diff --git a/.claude/skills b/.claude/skills new file mode 120000 index 00000000000..42c5394a18a --- /dev/null +++ b/.claude/skills @@ -0,0 +1 @@ +../skills \ No newline at end of file diff --git a/.cursor/BUGBOT.md b/.cursor/BUGBOT.md index 2f452276ebe..306b915063e 100644 --- a/.cursor/BUGBOT.md +++ b/.cursor/BUGBOT.md @@ -2,10 +2,16 @@ ## Changelog Enforcement -Any PR that modifies configuration structures or usage patterns must update `CHANGELOG.md`. This includes changes to config fields (added, removed, renamed, moved, or default value changes) in: +Any PR that introduces **breaking** configuration changes must update `CHANGELOG.md`. Breaking changes are those that require users to update existing configs: -- `src/prime_rl/*/config.py` -- `src/prime_rl/rl.py` -- `src/prime_rl/utils/config.py` +- **Renamed** config fields (old name no longer accepted) +- **Removed** config fields (field deleted or moved to a different path) +- **Moved** config fields (field relocated in the config hierarchy) -If such changes are detected without a corresponding `CHANGELOG.md` update, request that the author add an entry. +Additive changes (new fields with defaults, new optional features) and default value changes do **not** require a changelog entry. + +Config files live in: + +- `src/prime_rl/configs/` + +If breaking changes are detected without a corresponding `CHANGELOG.md` update, request that the author add an entry. diff --git a/.github/workflows/cpu_tests.yaml b/.github/workflows/cpu_tests.yaml index 2f1b1a803ea..0f80390ac8a 100644 --- a/.github/workflows/cpu_tests.yaml +++ b/.github/workflows/cpu_tests.yaml @@ -39,34 +39,58 @@ jobs: - name: Cleanup output_dir run: rm -rf /tmp/outputs - integration-tests: - name: Integration tests + slim-configs-install: + name: Slim install (prime-rl-configs only) runs-on: ubuntu-latest steps: - name: Checkout repository uses: actions/checkout@v4 - with: - submodules: true - name: Install uv uses: astral-sh/setup-uv@v5 with: enable-cache: true cache-dependency-glob: "uv.lock" - - name: Install dependencies - run: uv sync --all-extras --locked - - name: Run tests - env: - USERNAME_CI: CI_RUNNER - WANDB_API_KEY: ${{ secrets.WANDB_API_KEY }} - # Set WANDB_MODE to online only if WANDB_API_KEY is available, otherwise set to offline - # This is to allow running tests on forks without WANDB_API_KEY - WANDB_MODE: ${{ secrets.WANDB_API_KEY && 'online' || 'offline' }} - HF_TOKEN: ${{ secrets.HF_TOKEN }} - OPENAI_API_KEY: ${{ secrets.OPENAI_API_KEY }} - PRIME_API_KEY: ${{ secrets.PRIME_API_KEY }} - GITHUB_REF_NAME: ${{ github.ref_name }} - GITHUB_HEAD_REF: ${{ github.head_ref }} - PYTEST_OUTPUT_DIR: /tmp/outputs - run: PYTEST_OUTPUT_DIR=/tmp/outputs uv run pytest tests/integration -m "not gpu" - - name: Cleanup output_dir - run: rm -rf /tmp/outputs + - name: Build the slim wheel + run: uv build --wheel ./packages/prime-rl-configs --out-dir /tmp/slim-wheel + - name: Install slim wheel from a directory with no workspace context + # Build the wheel and install it from /tmp so uv can't discover the + # repo's root pyproject.toml [tool.uv.sources] and silently apply + # them. This mirrors what an external pip-only user gets. + run: | + uv venv --python 3.12 /tmp/slim + cd /tmp && uv pip install --python /tmp/slim/bin/python /tmp/slim-wheel/*.whl + - name: Verify configs import without heavy deps + run: | + /tmp/slim/bin/python - <<'PY' + import sys + + from prime_rl.configs.rl import RLConfig + from prime_rl.configs.sft import SFTConfig + from prime_rl.configs.inference import InferenceConfig + from prime_rl.configs.orchestrator import OrchestratorConfig + from prime_rl.configs.trainer import TrainerConfig + from prime_rl.configs.env_server import EnvServerConfig + from prime_rl.utils.config import BaseConfig, find_package_resource, rgetattr, rsetattr + from prime_rl.utils.validation import validate_shared_ckpt_config + + forbidden = ["torch", "transformers", "vllm", "wandb", "ring_flash_attn", + "verifiers", "prime", "datasets", "liger_kernel", "loguru"] + leaked = [m for m in forbidden if m in sys.modules] + if leaked: + raise SystemExit(f"slim install leaked heavy deps into sys.modules: {leaked}") + + # Slim wheel does not ship the templates/ data dir; resolver returns None. + assert find_package_resource("templates") is None, "templates should not be in slim wheel" + print("slim install OK — configs import with no heavy deps loaded") + PY + - name: Verify a real config (with [slurm]) parses on slim install + run: | + /tmp/slim/bin/python - <<'PY' + from pydantic_config import cli + from prime_rl.configs.rl import RLConfig + rl = cli(RLConfig, args=["@", "examples/Intellect-3.1/rl.toml"]) + assert rl.slurm is not None, "expected [slurm] in example config" + # On slim, template_path can't be auto-resolved (heavy wheel absent) — must stay None. + assert rl.slurm.template_path is None, f"unexpected template_path on slim: {rl.slurm.template_path}" + print("slim parse OK") + PY diff --git a/.github/workflows/devx_tag.yaml b/.github/workflows/devx_tag.yaml new file mode 100644 index 00000000000..d4891142a0e --- /dev/null +++ b/.github/workflows/devx_tag.yaml @@ -0,0 +1,94 @@ +name: DevX Tag + +on: + push: + branches: [main] + workflow_dispatch: + +permissions: + contents: write + +concurrency: + group: ${{ github.workflow }}-${{ github.ref }} + cancel-in-progress: false + +jobs: + create-dev-tag: + name: Create .dev tag + if: github.ref == 'refs/heads/main' + runs-on: ubuntu-latest + steps: + - name: Create next .dev tag from latest release + uses: actions/github-script@v7 + with: + script: | + const { owner, repo } = context.repo; + const sha = context.sha; + + let latestTag; + try { + const latestRelease = await github.rest.repos.getLatestRelease({ owner, repo }); + latestTag = latestRelease.data.tag_name; + } catch (error) { + if (error.status === 404) { + core.notice("No releases found. Skipping .dev tag creation."); + return; + } + throw error; + } + + const semverMatch = /^(v?)(\d+)\.(\d+)(?:\.(\d+))?$/.exec(latestTag); + if (!semverMatch) { + core.setFailed(`Latest release tag '${latestTag}' is not semver-like.`); + return; + } + + const [, versionPrefix, majorRaw, minorRaw, patchRaw] = semverMatch; + const major = Number.parseInt(majorRaw, 10); + const minor = Number.parseInt(minorRaw, 10); + const patch = Number.parseInt(patchRaw ?? "0", 10) + 1; + const baseTag = `${versionPrefix}${major}.${minor}.${patch}`; + + const prefix = `${baseTag}.dev`; + const escapeRegex = (value) => value.replace(/[.*+?^${}()|[\]\\]/g, "\\$&"); + const pattern = new RegExp(`^refs/tags/${escapeRegex(baseTag)}\\.dev(\\d+)$`); + + for (let attempt = 0; attempt < 5; attempt += 1) { + const refs = await github.paginate(github.rest.git.listMatchingRefs, { + owner, + repo, + ref: `tags/${prefix}`, + }); + + let maxVersion = 0; + for (const ref of refs) { + const match = pattern.exec(ref.ref); + if (match) { + const version = Number.parseInt(match[1], 10); + if (version > maxVersion) { + maxVersion = version; + } + } + } + + const tag = `${prefix}${maxVersion + 1}`; + + try { + await github.rest.git.createRef({ + owner, + repo, + ref: `refs/tags/${tag}`, + sha, + }); + core.notice(`Created tag ${tag} at ${sha}`); + return; + } catch (error) { + if (error.status === 422) { + core.warning(`Tag ${tag} already exists, retrying...`); + continue; + } + throw error; + } + } + + core.setFailed("Could not create a unique .dev tag after retries."); diff --git a/.github/workflows/gpu_tests.yaml b/.github/workflows/gpu_tests.yaml index 666ee5cc6ae..ed9e3908f7e 100644 --- a/.github/workflows/gpu_tests.yaml +++ b/.github/workflows/gpu_tests.yaml @@ -15,7 +15,7 @@ jobs: name: Unit tests runs-on: vm container: - image: pytorch/pytorch:2.5.1-cuda12.4-cudnn9-devel + image: pytorch/pytorch:2.10.0-cuda13.0-cudnn9-devel options: --gpus all timeout-minutes: 45 if: github.event_name == 'push' || github.event.pull_request.draft == false @@ -53,9 +53,9 @@ jobs: integration-tests: name: Integration tests (${{ matrix.test }}) - runs-on: vm + runs-on: ${{ matrix.runner }} container: - image: pytorch/pytorch:2.5.1-cuda12.4-cudnn9-devel + image: pytorch/pytorch:2.10.0-cuda13.0-cudnn9-devel options: --gpus all timeout-minutes: 60 if: github.event_name == 'push' || github.event.pull_request.draft == false @@ -64,11 +64,20 @@ jobs: matrix: include: - test: base - pytest_args: "tests/integration -m gpu --ignore=tests/integration/test_rl_lora.py --ignore=tests/integration/test_rl_multi_run_lora.py" + runner: vm + pytest_args: "tests/integration -m gpu --ignore=tests/integration/test_rl_lora.py --ignore=tests/integration/test_rl_multi_run_lora.py --ignore=tests/integration/test_benchmark_regression.py --ignore=tests/integration/test_alphabet_sort.py" - test: rl_lora + runner: vm pytest_args: "tests/integration/test_rl_lora.py -m gpu" + - test: alphabet_sort + runner: vm + pytest_args: "tests/integration/test_alphabet_sort.py -m gpu" - test: rl_multi_run + runner: 4xa6000 pytest_args: "tests/integration/test_rl_multi_run_lora.py -m gpu" + - test: benchmark_regression + runner: 4xa6000 + pytest_args: "tests/integration/test_benchmark_regression.py -m gpu" steps: - name: Install Git run: | @@ -96,7 +105,10 @@ jobs: HF_TOKEN: ${{ secrets.HF_TOKEN }} OPENAI_API_KEY: ${{ secrets.OPENAI_API_KEY }} PRIME_API_KEY: ${{ secrets.PRIME_API_KEY }} + PRIME_BASE_URL: ${{ secrets.PRIME_BASE_URL }} + PRIME_TEAM_ID: ${{ secrets.PRIME_TEAM_ID }} GITHUB_REF_NAME: ${{ github.ref_name }} GITHUB_HEAD_REF: ${{ github.head_ref }} PYTEST_OUTPUT_DIR: /tmp/outputs + VLLM_WORKER_MULTIPROC_METHOD: spawn run: uv run pytest ${{ matrix.pytest_args }} diff --git a/.github/workflows/nightly_tests.yaml b/.github/workflows/nightly_tests.yaml index 79b7f9c1927..c82a07011fa 100644 --- a/.github/workflows/nightly_tests.yaml +++ b/.github/workflows/nightly_tests.yaml @@ -56,8 +56,11 @@ jobs: needs: discover-experiments runs-on: research-cluster container: - image: pytorch/pytorch:2.5.1-cuda12.4-cudnn9-devel - options: --gpus all + image: pytorch/pytorch:2.10.0-cuda13.0-cudnn9-devel + # --shm-size: vLLM V1's DP engine/API-server IPC uses POSIX shm-backed + # MessageQueue (~160MB per queue); Docker's default 64MB /dev/shm triggers + # SIGBUS in DP engine cores on VLM runs (mm data overflows tmpfs). + options: --gpus all --shm-size=16g timeout-minutes: 1440 # 24 hours strategy: fail-fast: false diff --git a/.github/workflows/style.yaml b/.github/workflows/style.yaml index 120aebd3e46..39fa2d0cd31 100644 --- a/.github/workflows/style.yaml +++ b/.github/workflows/style.yaml @@ -20,3 +20,11 @@ jobs: submodules: true - name: Run ruff uses: astral-sh/ruff-action@v3 + with: + version: "0.13.0" + args: "check --config=pyproject.toml" + - name: Run ruff format (check) + uses: astral-sh/ruff-action@v3 + with: + version: "0.13.0" + args: "format --check --config=pyproject.toml" diff --git a/.github/workflows/sync-docs.yml b/.github/workflows/sync-docs.yml index 33c90147ab7..48d83308e5f 100644 --- a/.github/workflows/sync-docs.yml +++ b/.github/workflows/sync-docs.yml @@ -12,6 +12,39 @@ jobs: trigger-docs-sync: runs-on: ubuntu-latest steps: + - uses: actions/checkout@v4 + + - name: Validate mint.json pages exist + run: | + missing=() + pages=$(python3 -c " + import json, sys + data = json.load(open('docs/mint.json')) + pages = [] + def collect(obj): + if isinstance(obj, list): + for item in obj: collect(item) + elif isinstance(obj, dict): + if 'pages' in obj: collect(obj['pages']) + if 'groups' in obj: collect(obj['groups']) + elif isinstance(obj, str): + pages.append(obj) + collect(data.get('navigation', [])) + print('\n'.join(pages)) + ") + while IFS= read -r page; do + [ -z "$page" ] && continue + if [ ! -f "docs/${page}.md" ] && [ ! -f "docs/${page}.mdx" ]; then + missing+=("${page}") + fi + done <<< "$pages" + if [ ${#missing[@]} -gt 0 ]; then + echo "ERROR: The following pages listed in mint.json do not exist in docs/:" + printf ' - %s\n' "${missing[@]}" + exit 1 + fi + echo "All mint.json pages verified." + - name: Trigger public-docs workflow uses: peter-evans/repository-dispatch@v3 with: diff --git a/.gitignore b/.gitignore index c593bab3494..7eaa4976e61 100644 --- a/.gitignore +++ b/.gitignore @@ -5,6 +5,9 @@ configs/endpoints.py # Artifacts **/slurm +!src/prime_rl/slurm +!examples/slurm +!docs/slurm.md **/torchrun **/weights **/checkpoints @@ -196,3 +199,5 @@ debug_I2_zero_band .idea outputs/ + +third_party/ diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 6b89232fbeb..de1acda1be4 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -1,10 +1,10 @@ repos: - repo: https://github.com/astral-sh/ruff-pre-commit # Ruff version. - rev: v0.9.3 + rev: v0.13.0 hooks: # Run the linter. - - id: ruff + - id: ruff-check args: [ --fix, --config=pyproject.toml ] # Run the formatter. - id: ruff-format diff --git a/AGENTS.md b/AGENTS.md new file mode 100644 index 00000000000..55caa3db92c --- /dev/null +++ b/AGENTS.md @@ -0,0 +1,59 @@ +# AGENTS.md + +## Writing code + +- **Minimal try/except**: let errors propagate — silent failures hide bugs. Only catch exceptions for intentional fault tolerance (retries, robustness). +- **Targeted comments**: don't explain your work process or reference old code. Use targeted comments sparingly to clarify ambiguous logic. +- **Zen of Python**: remember the Zen of Python when writing code. +``` +Beautiful is better than ugly. +Explicit is better than implicit. +Simple is better than complex. +Complex is better than complicated. +Flat is better than nested. +Sparse is better than dense. +Readability counts. +Special cases aren't special enough to break the rules. +Although practicality beats purity. +Errors should never pass silently. +Unless explicitly silenced. +In the face of ambiguity, refuse the temptation to guess. +There should be one-- and preferably only one --obvious way to do it. +Although that way may not be obvious at first unless you're Dutch. +Now is better than never. +Although never is often better than *right* now. +If the implementation is hard to explain, it's a bad idea. +If the implementation is easy to explain, it may be a good idea. +Namespaces are one honking great idea -- let's do more of those! +``` + +## Running code + +- **Always use uv**: run code with `uv run` or `uv run `, never raw `python`. +- **Adding dependencies**: add to `pyproject.toml` and run `uv sync --all-extras` to install and lock them. +- **Git dependency pins**: when pinning git dependencies in `pyproject.toml`, always use a small (7-char) commit hash for the `rev` field. +- **Never edit `.venv/`**: the local virtual env is read-only. Edits there are silently overwritten by the next `uv sync` and don't propagate to teammates or CI. Reading files under `.venv/` to understand library behavior is fine; writing is not. To fix a dependency issue, update `pyproject.toml` (pin a fork via 7-char commit hash if needed), vendor the code into `src/`, or patch upstream. + +## Skills + +Skills live in `skills/` and are symlinked to `.claude/skills/`. They teach agents how to handle specific workflows (e.g. starting the inference server, writing configs). When you make changes to the codebase, check if any skills need to be updated to stay accurate. + +You are responsible for maintaining the skills folder. When a workflow fails and you fix it – whether with help from the user or through trial and error – you must update the skills to make implicit knowledge explicit. You are also responsible for keeping the skills up to date whenever you or anyone else modifies the code. + +## Testing + +Write tests as plain functions with pytest fixtures. Don't use class-based tests. + +- **Conservative test additions**: don't add new tests unless the user explicitly asks for them or it's clearly necessary. Editing existing tests is fine, but adding new test files or test functions should be the exception, not the default. +- **Test what matters**: only test code with clear, isolated logic — pure functions, abstract base classes, data transformations, well-defined algorithms. Don't test runtime-level code, framework glue, or anything that requires extensive mocking/patching just to get a test to pass. If you need to patch everything out to make it testable, it's probably not worth testing. + +## Git + +- **Branch prefixes**: use the following prefixes for branches: `feat/`, `fix/`, `chore/` + +## GitHub + +- **Draft PRs**: always create PRs as drafts (`gh pr create --draft`) to avoid triggering CI unnecessarily. +- **Pull requests**: do not include a "test plan" section in PR descriptions unless you actually ran tests to verify the changes or the user explicitly asked for one. +- **Keep PR descriptions in sync**: every time you push commits to a PR, also update the PR description (`gh pr edit --body-file ...`) so it reflects the current state of the branch — not just what was true when the PR was opened. Preserve any auto-generated blocks (e.g. ``). + diff --git a/CHANGELOG.md b/CHANGELOG.md index 9058ad1fde0..9e0de85b983 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,7 +1,50 @@ # Changelog -Documenting changes which affect configuration usage patterns (added/moved/removed/renamed fields, notable logic changes). +Documenting **breaking** configuration changes — renamed, removed, or moved fields that require users to update existing configs. +- **`orchestrator.advantage.length_penalty` → discriminated sub-config**: The scalar `length_penalty: Literal["tokens","turns"] | None` is replaced by a `LengthPenaltyConfig | None` discriminated on `type`. Token shaping now takes weighted completion + tool-response token costs. Migration: `length_penalty = "tokens"` becomes `[orchestrator.advantage.length_penalty]\ntype = "tokens"` (default weights `completion_weight = 1.0`, `tool_response_weight = 1.0` — total context). `length_penalty = "turns"` becomes `[orchestrator.advantage.length_penalty]\ntype = "turns"`. (2026-05-06) +- **`orchestrator.advantage.length_shaping` → `orchestrator.advantage.length_penalty`**: The boolean `length_shaping` flag has been replaced by `length_penalty: Literal["tokens", "turns"] | None` (default: `None`). `length_shaping = true` becomes `length_penalty = "tokens"`; `length_shaping = false` becomes `length_penalty = None`. The new `"turns"` option applies the same correctness-gated efficiency shaping using trajectory turn count instead of completion-token count. (2026-05-01) +- **`AdvantageInputs` API**: Replaced the `rewards`/`completion_lengths`/`num_turns` tensor fields with a single `rollouts: list[list[vf.RolloutOutput]]` (grouped by problem). Custom advantage functions can now access any rollout metadata. Existing custom advantages must update their signatures and extract per-rollout fields directly (e.g. `torch.tensor([[r["reward"] for r in g] for g in inputs.rollouts])`). (2026-05-01) +- **`orchestrator.teacher_rollout_model` now requires `orchestrator.use_sft_loss = true`**: External teacher rollout configs no longer rely on `trainer.loss.type = "sft"` to select SFT loss. Existing hard-distill configs must set `orchestrator.use_sft_loss = true` alongside `orchestrator.teacher_rollout_model`; the orchestrator validates the pair and stamps training samples with the per-batch SFT loss bool. (2026-04-26) +- **`model.attn = "eager"` (NEW option)**: Added `eager` as a valid value for the `model.attn` field. Required for GPT-OSS models on non-Hopper GPUs, since the only flash attention kernel GPT-OSS supports (`kernels-community/vllm-flash-attn3`) is Hopper-only. A clear error message is raised at model load time if GPT-OSS is used without `eager` on unsupported hardware. Also added `kernels` as a core dependency. (2026-04-05) +- **`[[orchestrator.env]]` → `[[orchestrator.train.env]]`**: Training environments and sampling are now configured under `[orchestrator.train]`. The old `[[orchestrator.env]]` and `[orchestrator.sampling]` paths are auto-translated with a deprecation warning and will be removed in a future release. (2026-04-09) +- **Per-env sampling overrides (NEW)**: Both `TrainEnvConfig` and `EvalEnvConfig` now accept a `[sampling]` section for per-env overrides. Unset fields inherit from the group-level sampling config (`[orchestrator.train.sampling]` or `[orchestrator.eval.sampling]`). (2026-04-09) +- **Per-env eval sampling overrides**: `EvalEnvConfig` now accepts a `[sampling]` section. Unset fields inherit from `[orchestrator.eval.sampling]`. All `EvalSamplingConfig` fields remain nullable (None = defer to inference server default). (2026-04-09) +- **`orchestrator.buffer.env_ratios` → per-env `orchestrator.train.env[].ratio`**: `buffer.env_ratios` has been removed. Set `ratio` on each `[[orchestrator.train.env]]` entry instead. Ratios must be all-or-nothing across envs (either all have a ratio or none do). (2026-04-05) +- **`orchestrator.val` removed**: The `[orchestrator.val]` config section (`ValConfig`) has been removed. Existing configs must delete this section. (2026-04-05) +- **`orchestrator.max_concurrent` removed**: Concurrency limiting via `max_concurrent` and the global semaphore have been removed. Existing configs must delete this field. (2026-04-05) +- **`orchestrator.buffer.hash_keys` default changed**: Default changed from `["task", "prompt"]` to `["env_name", "prompt"]`. The `task` field is no longer overridden by the orchestrator for env identification; `env_name` is used instead. Buffer checkpoints using the old default may not resume correctly. (2026-04-05) +- **`orchestrator.eval.env[].failed_rollouts` metric is now a ratio**: The `eval/{name}/failed_rollouts` metric now reports a ratio (0.0–1.0) instead of a raw count. Dashboards keying on this metric should be updated. (2026-04-05) +- **`orchestrator.sampling.temp_scheduler` removed**: Temperature scheduling (`TemperatureSchedulerConfig`) has been removed. `sampling.temperature` is now a required `float` (default `1.0`). Existing configs using `temp_scheduler` must replace it with a fixed `temperature` value. (2026-04-05) +- **`orchestrator.verification` removed**: The `[orchestrator.verification]` config section (`VerificationConfig`) has been removed. Rollout scoring is now always enabled. Existing configs must delete this section. (2026-04-05) +- **`orchestrator.eval.env[].skip_first` removed**: The `skip_first` field on `EvalEnvConfig` has been removed. Existing configs using this field must delete it. (2026-04-05) +- **`EvalSaveHFConfig` removed**: Unused config class removed. (2026-04-09) +- **`orchestrator.env[].max_total_completion_tokens` (NEW)**: Added `max_total_completion_tokens: int` to `EnvConfig` (default: `-1`, disabled). Limits the total completion tokens across all turns in a multi-turn rollout. Auto-populated into `extra_env_kwargs` and applied via the verifiers `set_max_total_completion_tokens()` setter. Works for both training and eval environments. (2026-04-06) +- **`sampling.max_tokens` → `sampling.max_completion_tokens`**: Renamed `max_tokens` to `max_completion_tokens` in both `SamplingConfig` and `EvalSamplingConfig` for consistency with OpenAI API naming. The old name is accepted with a deprecation warning. (2026-04-06) +- **`log.file` and `log.env_worker_logs` removed**: Removed `log.file` (from `LogConfig` and `SharedLogConfig`) and `log.env_worker_logs` (from `LogConfig`). Python file logging is replaced by deployment-level capture. Existing configs using these fields must delete them. Log paths unified: `.stdout` files renamed to `.log`, SLURM logs moved from `slurm/` to `logs/`. (2026-03-31) +- **`trainer.log.ranks_filter` (NEW)**: Added `ranks_filter: list[int]` to `TrainerLogConfig` (default: `[0]`). Controls which ranks appear in trainer console output via torchrun's `--local-ranks-filter`. (2026-03-31) +- **`wandb.log_extras.sample_ratio` / monitor sample logging defaults**: `wandb.log_extras.sample_ratio` is now actually applied to W&B sample-table logging via the shared monitor sampler (it was previously a no-op for WandB). Separately, the orchestrator no longer hard-caps sample logging to 8 rollouts before monitor-level sampling runs, so when monitor `sample_ratio` is `None`, monitors now receive and may log the full rollout batch for a step instead of at most 8 rollouts. This affects both W&B and Prime monitor sample logging behavior. (2026-03-27) +- **`[model.vlm].freeze_vision_encoder`**: Added to `VLMConfig`. Controls whether the vision encoder is frozen (default: `true`). When `false`, the vision encoder is trainable and FSDP-sharded per-block. Has no effect with LoRA. Moved from `[trainer.model].freeze_vision_encoder` into the VLM sub-config. (2026-03-25) +- **`loss_impl = "quack_fused"` (SFT)**: Added `quack_fused` option for the SFT `loss_impl` field. Uses quack-kernels for chunked linear + cross-entropy with CuTe DSL CUDA kernels, avoiding full logits materialization. Requires `quack-kernels` package. Does not support Gemma logit softcapping. Custom model impl (`model.impl = "custom"`) also gains quack RMSNorm acceleration on CUDA automatically. (2026-03-26) +- **`model.ac.targets`**: Removed the `mamba` selective activation checkpoint target and folded that behavior into the broader `linear_attn` target. Existing configs using `targets = ["mamba"]` will now fail validation at runtime with `ValueError`; switch them to `targets = ["linear_attn"]`. The `linear_attn` target description now covers supported token mixers outside the standard softmax-attention path, including NemotronH Mamba layers, Qwen3.5-MoE GatedDeltaNet layers, and AFMoE sliding-window attention layers. (2026-03-28) +- **`model.tp` (trainer `ModelConfig`)**: Removed from the trainer model config. Existing trainer configs must delete this field; it is no longer accepted. (2026-03-26) +- **`orchestrator.env[].num_workers`**: Added configurable env server worker count (`int | "auto"`, default: `"auto"`). When `"auto"`, scales based on concurrency (1 worker per 256 concurrent rollouts). Only used when the orchestrator spawns the env server (i.e. `address` is not set). (2026-03-25) +- **`[model.vlm]` (NEW — replaces auto-detection)**: VLM mode is now opt-in via a `[model.vlm]` sub-config with required `vision_encoder_attr` and `language_model_attr` fields. There is no auto-detection — if you train a VLM, you must add `[model.vlm]`. Existing multimodal configs need the new section. See `docs/multimodal.md` for the table of known model attrs. (2026-03-24) +- **`model.optimization_dtype` / `model.reduce_dtype` (VLM models, RL only)**: VLM dtype validation now only applies to RL training (`TrainerConfig`), not SFT. VLM models used with `sft` no longer require `optimization_dtype='bfloat16'` / `reduce_dtype='bfloat16'`. RL training still enforces both to match vLLM inference. (2026-03-24) +- **`model.ep_comm_backend`**: Added expert-parallel communication backend selection. Accepted values are `torch` (default, uses TorchTitan all-to-all collectives) and `deepep` (uses DeepEP custom kernels). This field only affects configurations with expert parallelism enabled via `model.ep > 1`. (2026-03-22) +- **`model.deepep_num_sms`**: Added DeepEP SM allocation control for intranode dispatch/combine kernels (default: `20`, minimum: `1`). This field is only used when `model.ep_comm_backend='deepep'`; it also determines the internode RDMA channel count (`num_channels = num_sms / 2`). (2026-03-22) +- **`model.deepep_token_chunk_size`**: Added optional token chunk size for DeepEP MoE pipelining (default: `None`, minimum: `1`). When set, DeepEP dispatch for chunk `i+1` overlaps expert compute for chunk `i`. This field is only used when `model.ep_comm_backend='deepep'`. (2026-03-22) +- **`model.optimization_dtype` / `model.reduce_dtype` (VLM models)**: Added validation that VLM models must use `optimization_dtype='bfloat16'` and `reduce_dtype='bfloat16'` to match vLLM inference. Previously valid configs with `float32` (the default) are now rejected for VLM model names. Set both fields to `"bfloat16"` when training VLMs. (2026-03-21) +- **`model.ac.mode`** and **`model.ac.targets`**: Added selective activation checkpointing configuration. `model.ac.mode` accepts `full` (default) or `selective`. When `selective`, `model.ac.targets` selects subcomponents to checkpoint. Supported public targets are currently `norm`, `attention_sdpa`, `mla_up_proj`, and `routed_experts`; runtime validation remains the source of truth. `model.ac.targets` defaults to `["norm"]`, and selective mode requires at least one target. (2026-03-20) +- **`orchestrator.advantage.length_shaping_alpha` → `orchestrator.advantage.length_shaping`**: Replaced GR³ (`length_shaping_alpha: float | None`) with correctness-gated efficiency shaping (`length_shaping: bool`, default: `false`). No hyperparameter needed, no `online_difficulty_filtering` requirement. Existing configs using `length_shaping_alpha` must switch to `length_shaping = true`. (2026-04-13) +- **`orchestrator.advantage.length_weighted_mean`**: Removed. The default advantage now always uses the plain per-problem mean baseline unless `orchestrator.advantage.length_shaping` is enabled. Existing configs must delete this field. (2026-03-19) +- **`prime_monitor.log_extras.sample_ratio`**: Added ratio-based rollout sampling (0.0–1.0, default: None). When set, caps the number of rollouts logged per step to `len(rollouts) * sample_ratio`. `None` preserves current behavior (log all rollouts). Interacts with existing `interval` gate which still runs first. (2026-03-12) +- **`client.connect_timeout`**: Added configurable TCP connect timeout for inference API requests (default: 30.0s). Previously hardcoded to 5.0s. Helps with vLLM or cluster flakiness (2026-03-11) +- **`model.fused_lm_head_token_chunk_size`**: Added as the fused LM-head chunking field for the token-chunked implementation. Unlike the removed `model.fused_lm_head_chunk_size`, this chunks over flattened sequence tokens rather than vocabulary rows. `model.fused_lm_head_chunk_size` is no longer accepted; switch configs to `model.fused_lm_head_token_chunk_size` explicitly. (2026-03-09) +- **`slurm.pre_run_command`**: Added optional shell command to run on the head node before starting the job. Useful for cleanup routines (e.g. killing stale processes, removing lock files). For all-nodes execution, wrap with `srun` in the command string (default: None) (2026-03-08) +- **`slurm.nodelist`**, **`slurm.exclude`**, **`slurm.account`**, **`slurm.time`**: Added common SLURM scheduling options (all default: None) (2026-03-08) +- **`orchestrator.verification.enabled`**: Added top-level rollout verification switch. `orchestrator.buffer.skip_verification` has been removed; use `verification.enabled = false` instead. When disabled, rewards are always 0 and reward-dependent buffer features (`online_difficulty_filtering`, `easy_threshold`, `hard_threshold`) must be unset (2026-03-03) +- **`client.dp_rank_count`**: Added configuration for data-parallel inference routing. When > 1, each `client.base_url` is expanded into `dp_rank_count` logical clients and pinned via `X-data-parallel-rank` to keep multi-turn rollouts on a consistent DP rank (default: 1) (2026-03-03) - **`model.lora`**: Moved from `model.experimental.lora` to `model.lora` (no longer experimental) (#1440, 2025-12-16) - Auto-set `api_server_count=1` on inference when LoRA is enabled, because vLLM doesn't support hotloading for multiple API servers (#1422, 2025-12-17) - **`inference.model.rope_scaling`**: Added RoPE scaling configuration passthrough to vLLM (#1447 2025-12-17) @@ -37,3 +80,71 @@ Documenting changes which affect configuration usage patterns (added/moved/remov - **`orchestrator.env.log`**: Removed. Use `orchestrator.log` for env worker logging instead (2026-01-15) - **`orchestrator.eval.retry.reraise`**: Changed default from `True` to `False`. When `False`, raises `tenacity.RetryError` after retries are exhausted instead of the original exception, allowing failed eval environments to be skipped with a warning (#1586, 2026-01-14) - **`model.ep`**: Expert parallelism now supported (with auto/custom impl only), changed from the old behaviour when `ep>1` was a no-op to a proper parallelization of the MoE layers. (#1595, 2026-01-15) +- **`orchestrator.reload_weights_on_start`**: Removed. The reload was a no-op in practice since vLLM servers already start with base weights, and LoRA runs skipped it. (#1829, 2026-02-19) +- **`orchestrator.client.elastic`**: Added elastic inference pool with DNS-based service discovery. Supports dynamic server scaling via any DNS hostname with multiple A records (Kubernetes headless services, Consul, Route53, etc.). Automatically syncs LoRA adapters on new servers and only exposes ready servers to workers (#1617, 2026-01-19) +- **`model.fused_lm_head_chunk_size`**: Replaced chunk size `int | None` setting with `int | Literal["auto", "disabled"]` setting. `auto` auto-sets to 2048 if possible. `disabled` explicitly disables chunked loss (use vanilla LM head). Default behaviour is to use `auto` for RL training and `disabled` for SFT training. (not changed from previous version) (#1649, 2026-01-23) +- **`client.skip_model_check`**: Added configuration to skip checking if the model is available in the inference pool. Useful for external APIs or API keys that don't support the /models endpoint (default: False) (#1543, 2026-01-06) +- **`orchestrator.sampling.temp_scheduler`**: Added optional temperature schedule configuration with linear and cosine schedules. Set either `sampling.temperature` (constant) or `sampling.temp_scheduler` (schedule), not both. Default remains 1.0 if neither is set. (2026-01-27) +- **`orchestrator.trajectory_strategy`**: Deprecated. Interleaving now automatically handles extension breaks by starting a new sample when the prefix doesn't match, achieving best-of-both behavior. The setting is ignored and interleaved mode is always used. (2026-01-30) +- **`model.impl`**: Removed `liger_kernel` model implementation from supported options. The Liger kernel dependency remains for SFT loss. (2026-01-30) +- **`log.json_logging`**: Added JSON structured logging option for log aggregation systems (Loki, Grafana, etc.). Outputs flat newline-delimited JSON with `timestamp`, `level`, `message`, `module`, `function`, `line` fields. Available on root `log`, `trainer.log`, and `orchestrator.log` (default: False) (2026-01-28) +- **`model.optim_cpu_offload`**: Added flag to offload optimizer states to CPU without moving parameters (default: False) (2026-01-31) +- **`orchestrator.tasks_per_minute`**: Added optional rate limiting for sandbox tasks per environment worker. Uses token bucket algorithm. If None (default), no rate limiting is applied (2026-02-02) +- **`model.cp`**: When `cp>1` with `attn="flash_attention_3"`, require `model.impl="custom"` (FA3 ring-attention kernel only in custom path) (2026-02-06) +- **`model.attn`**: Added `fa4` as an attention implementation option. Flash attention 4 is only supported with the custom implementation (#1726, 2026-02-06) +- **`inference.model.enable_prefix_caching`**: Added flag to enable prefix caching in vLLM. Passed to vLLM as `--enable-prefix-caching` (default: None) (2026-02-08) +- **`orchestrator.env.address`**: Added address field on `EnvConfig`. If set, connect to an environment server at this address; if None, spawn a server in a subprocess (2026-02-06) +- **`orchestrator.env.extra_env_kwargs`**: Added on `EnvConfig`. Extra kwargs passed to the env (e.g. seq_len, interleaved_rollouts, score_rollouts). Auto-populated by the orchestrator for training envs; generally not recommended for user override. Main use case is to match these kwargs when running an env in an isolated environment server (default: {}) (2026-02-06) +- **`OrchestratorConfig`**: Removed `workers_per_env`, `max_env_worker_restarts`, and `mask_env_responses` (2026-02-06) +- **`EvalSaveDiskConfig`**, **`EvalSaveConfig`**, **`RetryConfig`**, **`OnlineEvalConfig`**: Removed (2026-02-06) +- **`TemperatureScheduleConfig`**: Renamed to `TemperatureSchedulerConfig` (2026-02-06) +- **`optim.mu`**: Added Muon momentum (`mu`) config field (default: 0.95). Previously hardcoded to Muon class default. Also fixed `optim.betas1`/`optim.betas2` not being passed through to the Muon optimizer (2026-02-09) +- **`dump_config`**: Added `--dump-config ` flag to the `rl` command. When set, writes the resolved subconfigs (trainer, orchestrator, inference, teacher_inference) to the given directory and exits without starting any processes (2026-02-12) +- **`client.api_key_var`**: Changed default from "OPENAI_API_KEY" to "VLLM_API_KEY" (2026-02-12) +- **`orchestrator.filters`**: Added orchestrator-side rollout filters for detecting degenerate generations. Supports `[[filters]] type = "gibberish"` (rare tokens at high entropy) and `[[filters]] type = "repetition"` (high-confidence token streaks). Detected rollouts get reward zeroed and completion mask cleared (2026-02-13) +- **`inference.model.tool_call_parser`**: Changed default from `"hermes"` to auto-detection from model name. Uses `MODEL_TOOL_CALL_PARSER` dict to infer the correct vLLM tool call parser (e.g. Qwen3→`hermes`, GLM-4.5→`glm45`, GLM-4.7→`glm47`, MiniMax-M2→`minimax_m2`, INTELLECT-3→`hermes`). Unknown models default to `None`. Explicit values still take priority. (#1795, 2026-02-16) +- **`orchestrator.eval.cancel_inflight_rollouts_on_eval`**: Added flag to optionally cancel in-flight training rollouts before starting online evals. When enabled, avoids congestion by preventing training and eval rollouts from running simultaneously, but slows training as the rollout pipeline must refill after each eval (default: False) (2026-02-16) +- **`orchestrator.use_token_client`**: Added flag to use the token-in-token-out (TITO) client for training across all environments. When enabled, uses `openai_chat_completions_token` client type instead of `openai_chat_completions`. Only use when environments have linear history and the chat template has the extension property (default: False) (2026-02-21) +- **`model.cp` + AFMoE**: Context parallelism now works with AFMoE models via unified `substitute_ring_attn` which patches `_compute_attention` on both `FlashAttention` and `AfmoeFlashAttention` to use ring attention. Sliding window layers automatically get per-layer `window_size`; full attention layers default to `(-1, -1)`. Also plumbed `window_size` through the FA3 ring attention wrapper (`ring_fa3_varlen_func`). (2026-02-21) +- **`orchestrator.token_batch_size`** and **`orchestrator.max_inflight_rollouts`**: Added token-based batching via `token_batch_size` and explicit in-flight rollout control via `max_inflight_rollouts` (2026-02-23) +- **`orchestrator.batch_size`**: Now optional and mutually exclusive with `token_batch_size`. If neither is set, defaults to rollout mode with `batch_size=128` (2026-02-23) +- **`inference.enable_expert_parallel`**, **`inference.all2all_backend`**, and **`inference.enable_eplb`**: Added expert-parallel inference controls passed to vLLM as `--enable-expert-parallel`, `--all2all-backend`, and `--enable-eplb` (defaults: `False`, `"allgather_reducescatter"`, `False`) (2026-02-23) +- **`rl_slurm` / `sft_slurm` entrypoints**: Removed. SLURM submission is now handled by the unified `rl` and `sft` entrypoints. Add a `[slurm]` section to your config to submit via SLURM instead of running locally (2026-02-23) +- **`inference_gpu_ids` / `trainer_gpu_ids` / `teacher_gpu_ids`**: Removed from `RLConfig`. Replaced by `[deployment]` section with `type = "single_node"` (fields: `num_train_gpus`, `num_infer_gpus`, `num_teacher_gpus`) or `type = "multi_node"` (fields: `num_train_nodes`, `num_infer_nodes`, `num_teacher_nodes`, `nodes_per_fsdp_group`). Default is `single_node` with 1 train GPU and 1 infer GPU (2026-02-23) +- **`RLSLURMConfig`**: Removed. Fields `job_name`, `num_train_nodes`, `num_infer_nodes`, `gpus_per_node`, `nodes_per_fsdp_group`, `project_dir`, `slurm_template`, `dry_run` are now under `[slurm]` and `[deployment]` in the unified `RLConfig` (2026-02-23) +- **`SFTSLURMConfig`**: Removed. Fields `job_name`, `num_nodes`, `gpus_per_node`, `nodes_per_fsdp_group`, `project_dir`, `slurm_template`, `dry_run` are now under `[slurm]` and `[deployment]` in `SFTConfig` (2026-02-23) +- **`[deployment]` (RL)**: Added deployment configuration. `type = "single_node"` auto-derives contiguous GPU assignments from `num_infer_gpus`, `num_train_gpus`, `num_teacher_gpus`. `type = "multi_node"` requires `[slurm]` and uses `num_train_nodes`, `num_infer_nodes` (2026-02-23) +- **`[deployment]` (SFT)**: Added deployment configuration. `type = "single_node"` with `num_gpus` (default: 1). `type = "multi_node"` with `num_nodes`, `nodes_per_fsdp_group`, `hf_hub_offline` (2026-02-23) +- **`[slurm]` (RL)**: Added SLURM configuration with `job_name`, `project_dir`, `template_path`, `partition`, `dry_run`. When present, `uv run rl` generates and submits an sbatch script instead of running locally. Template is auto-selected based on deployment type (2026-02-23) +- **`[slurm]` (SFT)**: Added SLURM configuration with `job_name`, `project_dir`, `template_path`, `partition`, `dry_run`. When present, `uv run sft` generates and submits an sbatch script instead of running locally (2026-02-23) +- **`hf_hub_offline` (RL/SFT SLURM)**: Removed. `HF_HUB_OFFLINE=1` is now hardcoded in the multi-node SLURM templates (2026-02-23) +- **SLURM templates**: Moved from `src/prime_rl/slurm/` to `src/prime_rl/templates/` and renamed to `single_node_rl.sbatch.j2`, `multi_node_rl.sbatch.j2`, `single_node_sft.sbatch.j2`, `multi_node_sft.sbatch.j2` (2026-02-23) +- **Entrypoints**: Moved `rl` and `sft` entrypoints from `prime_rl.rl` / `prime_rl.sft` to `prime_rl.entrypoints.rl` / `prime_rl.entrypoints.sft`. No change to CLI usage (`uv run rl`, `uv run sft`) (2026-02-24) +- **`output_dir` (RL)**: Changed from `Path | None` (default `None`) to `Path` (default `Path("outputs")`). The SLURM-specific validation that rejected the default has been removed — `output_dir` now works the same for local and SLURM runs (2026-02-24) +- **`clean_output_dir`**: Added to `RLConfig` and `SFTConfig` (default: `False`). Training now raises `FileExistsError` when `output_dir` contains checkpoints from a previous run and not resuming. Set `clean_output_dir=true` to delete and start fresh, or set `ckpt.resume_step` to resume (2026-02-24) +- **`clean`**: Removed from `RLConfig`. The old `clean` flag (default: `True`) silently deleted logs, rollouts, and broadcasts on every local RL run. Superseded by the explicit `clean_output_dir` flag (2026-02-24) +- **Config consolidation**: All config modules moved into `prime_rl.configs` subpackage. `utils/config.py` + `transport/config.py` → `configs/shared.py`; `trainer/config.py` + `trainer/rl/config.py` → `configs/trainer.py`; `trainer/sft/config.py` → `configs/sft.py`; `orchestrator/config.py` → `configs/orchestrator.py`; `inference/config.py` → `configs/inference.py`; `rl_config.py` → `configs/rl.py`. Class renames: `SFTTrainerConfig` → `SFTConfig`, `RLTrainerConfig` → `TrainerConfig`. Component prefixes dropped from orchestrator and inference config classes (e.g. `OrchestratorCheckpointConfig` → `CheckpointConfig`). TypeAlias renames: dropped `Type` suffix (e.g. `LossConfigType` → `LossConfig`, `TransportConfigType` → `TransportConfig`), renamed `LossConfig` class → `DefaultLossConfig`. No TOML key changes. (2026-02-24) +- **`trainer.enable_router_replay`**: Added flag to enable router replay. If True, will return routed experts in the batch. This is only supported if `enable_return_routed_experts=True` in the inference config or pass `--enable-return-routed-experts` to vLLM server. This is only supported for custom models. (2026-02-22) +- **`inference.enable_return_routed_experts`**: Added flag to enable return routed experts. Passed to vLLM as `--enable-return-routed-experts` (2026-02-22) +- **`orchestrator.oversampling_factor`**: Added rollout-only over-sampling config that resolves `max_inflight_rollouts = int(batch_size * oversampling_factor)` when `max_inflight_rollouts` is unset. Cannot be used with `token_batch_size` or together with explicit `max_inflight_rollouts` (2026-02-25) +- **`sft.val`**: Added optional periodic SFT validation with `val/loss` and `val/num_tokens` logging. Configure via `sft.val.data` (validation dataset) and `sft.val.interval` (every N steps, default 50). Runs the full validation dataset each pass. (2026-02-26) +- **`model.fused_lm_head_chunk_size`**: Changed default value from 2048 to 8192 for RL training (2026-02-26) +- **`inference.data_parallel_size_local`** and **`inference.data_parallel_rpc_port`**: Added data-parallel node-local controls for vLLM, passed as `--data-parallel-size-local` and `--data-parallel-rpc-port` (defaults: `None`, `13345`) (2026-02-26) +- **`dump_config`**: Removed from `RLConfig`. Replaced by `dry_run` (see below) (2026-02-26) +- **`slurm.dry_run`**: Removed from `SlurmConfig`. Replaced by top-level `dry_run` (see below) (2026-02-26) +- **`dry_run`**: Added to `RLConfig` and `SFTConfig` (default: `False`). When set, validates the config, writes resolved subconfigs to `output_dir/configs/`, and exits without starting any processes. Works the same for both local and SLURM runs (2026-02-26) +- **Config output location**: Resolved subconfigs are now always written to `output_dir/configs/` instead of `.pydantic_config//`. This applies to both local and SLURM entrypoints, and for both single-node and multi-node deployments (2026-02-26) +- **SFT config filename**: The resolved SFT trainer config is now written as `sft.toml` instead of `trainer.toml` (2026-02-26) +- **Inference entrypoint**: Moved from `prime_rl.inference.server` to `prime_rl.entrypoints.inference`. No change to CLI usage (`uv run inference`) (2026-02-26) +- **`[slurm]` (inference)**: Added SLURM configuration to `InferenceConfig`. When present, `uv run inference` generates and submits an sbatch script instead of running locally. Each SLURM node runs an independent vLLM replica (no cross-node parallelism) (2026-02-26) +- **`[deployment]` (inference)**: Added deployment configuration. `type = "single_node"` (default) with `gpus_per_node`. `type = "multi_node"` with `num_nodes` and `gpus_per_node` — requires `[slurm]` (2026-02-26) +- **`inference.output_dir`**: Added directory for SLURM logs and generated scripts (default: `"outputs"`) (2026-02-26) +- **`inference.dry_run`**: Added flag (default: `False`). When set, validates config, writes resolved config to `output_dir/configs/`, and exits without starting inference or submitting SLURM jobs (2026-02-26) +- **`orchestrator.teacher_rollout_model`**: Added optional external rollout model configuration. When set, rollouts are generated from this endpoint/model instead of the student inference server. Accepts a `TeacherRolloutModelConfig` with `client` and `model` sub-fields, or `None` (default: `None`) (2026-02-26) +- **`trainer.loss.type = "sft"`**: Added SFT loss variant. Set `trainer.loss.type = "sft"` to use masked negative log-likelihood loss instead of the default RL loss (2026-02-26) +- **`trainer.loss.type = "custom"`**: Added custom loss variant. Set `trainer.loss.type = "custom"` with `import_path` (e.g. `"my_module.my_loss"`) and optional `kwargs` to use an external loss function (2026-02-26) +- **`trainer.loss` (default loss)**: Made IPO (DPPO-Binary TV variant ([arxiv](https://arxiv.org/pdf/2602.04879)) + Kimi-K2.5 KL ([Kimi-K2.5](https://arxiv.org/pdf/2602.02276))) the default loss. Removed `ratio_type`, `token_mask_low`, `token_mask_high`, `sequence_clip_high`, `geo_mask_low`, `geo_mask_high`, `sequence_mask_low`, `sequence_mask_high`. Added `ipo_mask_low` (default: 0.2) and `ipo_mask_high` (default: 0.2) for token-level probability-difference masking. Changed `kl_tau` default from `0.0` to `1e-3`. (2026-03-02) +- **Metrics logging overhaul**: All orchestrator metrics now follow a `{metric}/{scope}/{stat}` naming convention where scope is `all` (global) or an env name. Per-env breakdowns are always logged (previously only when >1 env). Key renames: `reward/mean` → `reward/all/mean`, `batch/solve_none` → `solve_none/all`, `val_reward/` → `val/reward/`, `metrics/{name}` → `metrics/{env}/{name}`, `stop_condition/{sc}` → `stop_condition/all/{sc}`, `error/mean` → `error/all/mean`. New per-env metrics: `solve_none/{env}`, `solve_all/{env}`, `effective_batch_size/{env}`, `stop_condition/{env}/generation_truncated`, `stop_condition/{env}/{sc}`, `error/{env}/mean`, plus per-env breakdowns for `seq_len`, `prefill_len`, `decode_len`, `is_truncated`, `samples_per_rollout`, `num_turns`, `generation_ms`, `scoring_ms`. Removed: `reward/std`, `reward/median`, per-error-type breakdown (`error/{error_type}`). Env name `"all"` is reserved and rejected at config validation. Env name validation now strips `@version` suffixes to match runtime behavior (e.g. `math@1.0` and `math@2.0` correctly detected as duplicates). Solve stats now use `example_id` grouping instead of index arithmetic (pre-existing bug fix). (2026-03-08) +- **`wandb.shared`** (experimental): Added shared W&B mode that logs trainer and orchestrator metrics to a single W&B run instead of two separate runs. Uses `wandb.Settings(mode="shared")` (requires wandb SDK >= 0.19.9). Enabled by default on the RL entrypoint. Disable with `--wandb.shared False`. Run ID is communicated via `WANDB_SHARED_RUN_ID` env var; process role via `WANDB_SHARED_LABEL`. Non-primary processes retry `wandb.init` up to 30 times waiting for the primary to create the run. Works with multi-node SLURM. (2026-03-18) +- **`wandb.id`**: Removed from `WandbConfig`. Run IDs are now managed internally via env vars for shared mode. (2026-03-18) +- **`trainer.ckpt.skip_optimizer`**: Added flag to skip loading optimizer states when resuming training. When set, only the model weights are loaded from the checkpoint. (2026-03-24) diff --git a/CLAUDE.md b/CLAUDE.md new file mode 100644 index 00000000000..43c994c2d36 --- /dev/null +++ b/CLAUDE.md @@ -0,0 +1 @@ +@AGENTS.md diff --git a/Dockerfile.cuda b/Dockerfile.cuda index cd587c941eb..f7e8fd7c59b 100644 --- a/Dockerfile.cuda +++ b/Dockerfile.cuda @@ -1,7 +1,7 @@ #inspired from https://github.com/astral-sh/uv-docker-example/blob/main/multistage.Dockerfile -# Build stage -FROM pytorch/pytorch:2.5.1-cuda12.4-cudnn9-devel AS builder +# Build stage - multi-arch base (supports amd64 + arm64) +FROM nvidia/cuda:12.8.1-cudnn-devel-ubuntu22.04 AS builder LABEL maintainer="prime intellect" LABEL repository="prime-rl" @@ -12,12 +12,16 @@ RUN echo "LC_ALL=en_US.UTF-8" >> /etc/environment ENV CUDA_HOME=/usr/local/cuda ENV PATH=$PATH:/usr/local/cuda/bin -# Install packages +# Install build tooling. +ARG DEBIAN_FRONTEND=noninteractive +ENV DEBIAN_FRONTEND=noninteractive +ENV TZ=Etc/UTC RUN apt-get update && apt-get install -y --no-install-recommends --force-yes \ build-essential \ curl \ sudo \ git \ + ninja-build \ && apt-get clean autoclean && rm -rf /var/lib/apt/lists/* /tmp/* /var/tmp/* # Download the latest installer @@ -38,12 +42,18 @@ COPY pyproject.toml /app/pyproject.toml COPY uv.lock /app/uv.lock COPY README.md /app/README.md COPY src/ /app/src/ +COPY packages/ /app/packages/ COPY configs /app/configs COPY examples /app/examples COPY benchmarks/scripts /app/benchmarks/scripts RUN --mount=type=cache,target=/app/.cache/uv \ - uv sync --all-extras --locked --no-dev + uv sync --extra flash-attn --extra flash-attn-3 --extra flash-attn-cute --extra envs --extra gpt-oss --group mamba-ssm --locked --no-dev + +# arm64: build flash-attn from source, fix namespace conflicts, apply workarounds +ARG TARGETARCH +COPY scripts/docker-arm64-post-install.sh /app/scripts/docker-arm64-post-install.sh +RUN if [ "$TARGETARCH" = "arm64" ]; then /app/scripts/docker-arm64-post-install.sh; fi FROM python:3.12-slim @@ -97,6 +107,10 @@ RUN rm /app/.venv/bin/python3.12 && ln -s /usr/local/bin/python /app/.venv/bin/p # Place executables in the environment at the front of the path ENV PATH="/app/.venv/bin:$PATH" +# HuggingFace Hub timeouts (defaults are 10s which causes issues on slow networks) +ENV HF_HUB_ETAG_TIMEOUT=500 +ENV HF_HUB_DOWNLOAD_TIMEOUT=300 + # Use entrypoint for setup (ulimit, etc) but default to sleep infinity for K8s ENTRYPOINT ["/app/docker-entrypoint.sh"] CMD ["sleep", "infinity"] diff --git a/README.md b/README.md index 537cc695eb8..562a51595fb 100644 --- a/README.md +++ b/README.md @@ -29,13 +29,37 @@ PRIME-RL: Async RL Training at Scale ## Overview -PRIME-RL is a framework for large-scale asynchronous reinforcement learning. It is designed to be easy-to-use and hackable, yet capable of scaling to 1000+ GPUs. Beyond that, here is why we think you might like it: +PRIME-RL is a framework for large-scale reinforcement learning. It is designed to be easy to use and hackable, yet capable of scaling to 1000+ GPUs. Here is what we think sets it apart: + +1. Fully asynchronous RL for high-throughput agentic training at scale. +2. Performant: built to train 1T+ MoE models on 1000+ GPUs with [FSDP2](https://docs.pytorch.org/tutorials/intermediate/FSDP_tutorial.html) for training and [vLLM](https://github.com/vllm-project/vllm) for inference, with FP8 inference, PD disaggregation, EP and CP parallelism, and more. +3. Native integration with [`verifiers`](https://github.com/PrimeIntellect-ai/verifiers) environments through the [Environments Hub](https://app.primeintellect.ai/dashboard/environments?ex_sort=most_stars), including built-in support for SWE and agentic environments. +4. End-to-end post-training: SFT, RL training, and evals. +5. Multi-node deployment with Slurm and Kubernetes support. +6. Multimodal support for VLMs such as Qwen3-VL. +7. Hackable, modular, and extensible by design. + + +## Models support + + +The trainer works with both Hugging Face and Prime custom `ModelForCausalLM` out of the box. For selected families (especially large MoE) we also ship highly optimized training code under `src/prime_rl/trainer/models/`, including expert parallelism (EP) for MoE layers and context parallelism (CP) for long sequences (see the table), and additional kernels like [quack-kernels](https://github.com/quack-kernels/quack-kernels). + +With `[model] impl = "auto"` (the default), the trainer selects that custom stack when the Hugging Face config type is registered. + +| Family | Example IDs | MoE | EP | CP | +|--------|-------------|-----|----|-----| +| GLM-5 (`glm_moe_dsa`) | `zai-org/GLM-5`, `zai-org/GLM-5-FP8` | yes | ✅ | ✅ | +| Qwen3 MoE (`qwen3_moe`) | `Qwen/Qwen3-30B-A3B`, … | yes | ✅ | ✅ | +| Qwen3.5 MoE (`qwen3_5_moe`) | `Qwen/Qwen3.5-35B-A3B`, … | yes | ✅ | ✅ | +| Qwen3 / Qwen3.5 VLMs | [multimodal.md](docs/multimodal.md) (`qwen3_vl`, `qwen3_5`, `qwen3_5_moe`) | MoE only on MoE VLMs | MoE only | ✅ | +| MiniMax M2 (`minimax_m2`) | `MiniMax/MiniMax-M2` | yes | ✅ | ✅ | +| Nemotron H (`nemotron_h`) | `nvidia/Nemotron-3-Nano-30B-A3B`, `nvidia/Nemotron-3-Super-120B-A12B`, … | yes | ✅ | ❌ | +| Trinity (`afmoe`) | `arcee-ai/Trinity-Mini`, … | yes | ✅ | ✅ | +| GLM-4 · GLM-4.5 MoE · INTELLECT-3 (`glm4_moe`) | `THUDM/GLM-4-9B-0414`, `zai-org/GLM-4.5-Air`, `zai-org/GLM-4.5`, `PrimeIntellect/INTELLECT-3`, … | yes | ✅ | ✅ | +| GPT-OSS (HF, MoE) | `openai/gpt-oss-20b`, `openai/gpt-oss-120b` | yes | ❌ | ✅ | +| Other HF causal LMs | Qwen3 dense, Mistral, … (`impl = "hf"`) | varies | ❌ | ✅ | -1. Integrates natively with [`verifiers`](https://github.com/PrimeIntellect-ai/verifiers) environments via the [Environments Hub](https://app.primeintellect.ai/dashboard/environments?ex_sort=most_stars) -2. Supports end-to-end post-training, including SFT and RL training and evals -3. Multi-node deployment with [FSDP2](https://docs.pytorch.org/tutorials/intermediate/FSDP_tutorial.html) training and [vLLM](https://github.com/vllm-project/vllm) inference backend -4. Designed for asynchronous agentic RL training at scale -5. Hackable, modular and extensible by nature ## Setup @@ -160,12 +184,27 @@ uv run hf auth login ## Training Examples We provide end-to-end training examples in the [`examples`](examples) directory to highlight features of the framework and guide you through the process of training your own models. + +### Basic Training: 1 to 8 GPUs + +Follow this guide to learn the basics of Prime-RL. You can train your own models on 1 to 8 GPUs. Ideal for getting started and exploring the capabilities of the framework. These guides cover most use cases -- single-turn, multi-turn, tool calling, etc. -- on toy environments and small models. + 1. [**Reverse Text**](examples/reverse_text/README.md): Train `Qwen3-0.6B` to reverse a small chunk of text. Demonstrates tiny-scale single-turn SFT and RL training. Can be trained on a single consumer GPU in a few minutes, and is ideal for getting started. 2. [**Wordle**](examples/wordle/README.md): Train `Qwen3-1.7B` to play Wordle. A fun example of multi-turn SFT and RL training. Can be trained on a 2-4 H100 GPUs in a few hours. Ideal for exploring the multi-turn training capabilities of the framework. 3. [**Alphabet Sort**](examples/alphabet_sort/README.md): Train `Qwen3-4B-Instruct-2507` to sort names alphabetically. Demonstrates multi-turn RL training via LoRA without SFT warmup. Can be trained on a single H100 GPU in just over an hour. Ideal for exploring LoRA-based training. 4. [**Wiki Search**](examples/wiki_search/README.md): Train `Qwen3-4B-Instruct-2507` to answer trivia questions by searching through a Wikipedia. Demonstrates multi-turn with web search tool use. +5. [**Hendrycks Sanity**](examples/hendrycks_sanity/README.md): Run a sanity check experiment on `DeepSeek-R1-Distill-Qwen-1.5B` using a filtered subset of MATH where the model already partially solves 20-80% of problems. Useful for algorithm ablations. + +### Advanced Training: 32 - 2048 GPUs: + +Follow this guide to train large models on hard reasoning and agentic / swe environments. +These guides are designed to be run from a Slurm cluster but can also be adapted to k8s deployments. -4. *More to come...* +1. [**Qwen 3 30B - A3B Math**](examples/qwen30b_math/README.md): Train `Qwen3-30B-A3B` to solve hard math problems. +2. [**Qwen 3 30B - A3B SWE**](examples/qwen30b_swe/README.md): Train `Qwen3-30B-A3B` to solve hard SWE problems. +3. [**Intellect-3.1**](examples/Intellect-3.1/README.md): Reproduce our `INTELLECT-3.1` training run. +4. [**MiniMax-M2.5 SWE**](examples/minimax_m2.5_swe/README.md): Train `MiniMax-M2.5` on agentic SWE tasks. +5. [**High-throughput GLM-5**](examples/glm5_pd_disag/README.md): Train `GLM-5` with PD disaggregation and FP8 inference on SWE. ## Docs @@ -179,7 +218,9 @@ Check out the [docs](docs) directory for in-depth guides on how to use PRIME-RL. - [**Checkpointing**](docs/checkpointing.md) - Saving and resuming training from checkpoints - [**Benchmarking**](docs/benchmarking.md) - Performance benchmarking and throughput measurement - [**Deployment**](docs/deployment.md) - Training deployment on single-GPU, multi-GPU, and multi-node clusters +- [**Memory Usage**](docs/memory_usage.md) - Techniques for reducing memory usage (activation checkpointing, offloading, EP, CP, LoRA, etc.) - [**Troubleshooting**](docs/troubleshooting.md) - Common issues and their solutions +- [**Multimodal**](docs/multimodal.md) - Training VLMs like Qwen3-VL ## Contributing diff --git a/benchmarks/baselines/benchmark-1xa6000-Qwen--Qwen3-0.6B-rl-full-1gpu-Recompute-flash_attention_2-65536-cp1-ep1.json b/benchmarks/baselines/benchmark-1xa6000-Qwen--Qwen3-0.6B-rl-full-1gpu-Recompute-flash_attention_2-65536-cp1-ep1.json index a558d9d60dc..de22864fe64 100644 --- a/benchmarks/baselines/benchmark-1xa6000-Qwen--Qwen3-0.6B-rl-full-1gpu-Recompute-flash_attention_2-65536-cp1-ep1.json +++ b/benchmarks/baselines/benchmark-1xa6000-Qwen--Qwen3-0.6B-rl-full-1gpu-Recompute-flash_attention_2-65536-cp1-ep1.json @@ -7,40 +7,42 @@ "seq_len": 65536, "ac": "Recompute", "attention": "flash_attention_2", - "output": "/tmp/benchmark_result.json", + "output": "benchmark_result.json", "dry_run": false, "timeout": 1500, "micro_batches": 2, - "docker_image": "primeintellect/prime-rl-jackmin@sha256:5a146f7dfdcdf6b0e90fd3f1ed0874d22a3a0641378dc8dad46ce213d02fa2e6", + "ep": 1, + "cp": 1, + "docker_image": null, "device_name": "NVIDIA RTX A6000", - "commit_sha": "unknown", - "timestamp": "2026-01-15T22:16:04.729569+00:00", + "commit_sha": "83c21ddac", + "timestamp": "2026-02-26T02:07:06+00:00", "success": true, "error_reason": null, - "time_taken": 278.55358465001336 + "time_taken": 276.12556480802596 }, "metrics": { "mfu": { - "mean": 17.440641920184348, - "std": 0.024251274789467057, - "min": 17.413318182049757, - "max": 17.459612680819426 + "mean": 16.41977165630161, + "std": 0.010586646038388221, + "min": 16.407648988464405, + "max": 16.42719598975201 }, "throughput": { - "mean": 2082.9043561040025, - "std": 2.8962859355307957, - "min": 2079.641131421904, - "max": 2085.1699997738997 + "mean": 1960.9836648020673, + "std": 1.2643440104327244, + "min": 1959.5358764831926, + "max": 1961.8703395331793 }, "step_time": { - "mean": 61.44155622298907, - "std": 0.33102868375328043, - "min": 61.21719340299023, - "max": 61.8217413329985 + "mean": 65.87, + "std": 0.13, + "min": 65.74, + "max": 66.0 }, "peak_memory": { - "gib": 19.501953125, - "pct": 41.141803507714734 + "gib": 31.806640625, + "pct": 67.03 } } -} \ No newline at end of file +} diff --git a/benchmarks/baselines/benchmark-4xa6000-Qwen--Qwen3-0.6B-rl-full-4gpu-Recompute-flash_attention_2-65536-cp1-ep1.json b/benchmarks/baselines/benchmark-4xa6000-Qwen--Qwen3-0.6B-rl-full-4gpu-Recompute-flash_attention_2-65536-cp1-ep1.json new file mode 100644 index 00000000000..b043c1c79e0 --- /dev/null +++ b/benchmarks/baselines/benchmark-4xa6000-Qwen--Qwen3-0.6B-rl-full-4gpu-Recompute-flash_attention_2-65536-cp1-ep1.json @@ -0,0 +1,48 @@ +{ + "config": { + "type": "rl", + "num_gpus": 4, + "model_name": "Qwen/Qwen3-0.6B", + "lora_rank": null, + "seq_len": 65536, + "ac": "Recompute", + "attention": "flash_attention_2", + "output": "benchmark_result.json", + "dry_run": false, + "timeout": 1500, + "micro_batches": 2, + "ep": 1, + "cp": 1, + "docker_image": null, + "device_name": "NVIDIA RTX A6000", + "commit_sha": "83c21ddac", + "timestamp": "2026-02-26T02:07:06+00:00", + "success": true, + "error_reason": null, + "time_taken": 291.91085167229176 + }, + "metrics": { + "mfu": { + "mean": 15.634260517326164, + "std": 0.009817863071768202, + "min": 15.622927380137481, + "max": 15.640172924064188 + }, + "throughput": { + "mean": 7468.685954343421, + "std": 4.690118598479809, + "min": 7463.271969943842, + "max": 7471.510386564679 + }, + "step_time": { + "mean": 68.34095695288852, + "std": 0.16054587555887065, + "min": 68.16114284005016, + "max": 68.46991707291454 + }, + "peak_memory": { + "gib": 23.25390625, + "pct": 49.05701678145735 + } + } +} \ No newline at end of file diff --git a/benchmarks/results/BENCHMARKS.md b/benchmarks/results/BENCHMARKS.md index 926f56dfa0d..77ec1caf560 100644 --- a/benchmarks/results/BENCHMARKS.md +++ b/benchmarks/results/BENCHMARKS.md @@ -2,7 +2,7 @@ Automated benchmark results for prime-rl using `--bench` flag. -**Last Updated:** 2026-01-16 01:16 UTC +**Last Updated:** 2026-02-01 00:43 UTC **Commit:** `unknown` **Docker Image:** `primeintellect/prime-rl-jackmin@sha256:5a146f7dfdcdf6b0e90fd3f1ed0874d22a3a0641378dc8dad46ce213d02fa2e6` @@ -12,87 +12,78 @@ Automated benchmark results for prime-rl using `--bench` flag. > :clock10: The Step Time shown is the time taken per micro batch. This differs from what gets displayed in the bench table which is the total step time. ## INTELLECT-3 -| Type | SeqLen | AC | Attn | Hardware | MFU | TPS | Step Time | Peak Mem | -|------|--------|----|----|----------|-----|-----|-----------|----------| -| RL LoRA(r=16) | 16384 | Recompute | FA3 | 8xH200 | 19.8% | 16.69k | 7.85s | 89.7 GiB | -| RL LoRA(r=16) | 16384 | Offload | FA3 | 8xH200 | 19.2% | 16.15k | 8.11s | 84.2 GiB | -| RL LoRA(r=16) | 16384 | Recompute | FA2 | 8xB200 | 2.5% | 4.78k | 27.45s | 90.4 GiB | -| RL LoRA(r=16) | 65536 | Offload | FA3 | 8xH200 | 25.1% | 9.76k | 53.73s | 128.8 GiB | -| RL LoRA(r=16) | 65536 | Offload | FA2 | 8xH200 | 16.2% | 6.30k | 83.23s | 130.3 GiB | -| RL LoRA(r=16) | 65536 | Recompute | FA2 | 8xB200 | 5.8% | 5.16k | 101.60s | 153.6 GiB | -| RL LoRA(r=16) | 65536 | Offload | FA2 | 8xB200 | 5.7% | 5.04k | 104.10s | 131.5 GiB | +| Type | SeqLen | AC | Attn | EP | CP | Hardware | MFU | TPS | Step Time | Peak Mem | +|------|--------|----|----|----|----|----------|-----|-----|-----------|----------| +| RL LoRA(r=16) | 16384 | Recompute | FA3 | 1 | 1 | 8xH200 | 19.8% | 16.69k | 7.85s | 89.7 GiB | +| RL LoRA(r=16) | 16384 | Offload | FA3 | 1 | 1 | 8xH200 | 19.2% | 16.15k | 8.11s | 84.2 GiB | +| RL LoRA(r=16) | 16384 | Recompute | FA2 | 1 | 1 | 8xB200 | 2.5% | 4.78k | 27.45s | 90.4 GiB | +| RL LoRA(r=16) | 65536 | Offload | FA3 | 1 | 1 | 8xH200 | 25.1% | 9.76k | 53.73s | 128.8 GiB | +| RL LoRA(r=16) | 65536 | Offload | FA2 | 1 | 1 | 8xH200 | 16.2% | 6.30k | 83.23s | 130.3 GiB | +| RL LoRA(r=16) | 65536 | Recompute | FA2 | 1 | 1 | 8xB200 | 5.8% | 5.16k | 101.60s | 153.6 GiB | +| RL LoRA(r=16) | 65536 | Offload | FA2 | 1 | 1 | 8xB200 | 5.7% | 5.04k | 104.10s | 131.5 GiB | ## Qwen3-0.6B -| Type | SeqLen | AC | Attn | Hardware | MFU | TPS | Step Time | Peak Mem | -|------|--------|----|----|----------|-----|-----|-----------|----------| -| RL Full | 16384 | Recompute | FA3 | 1xH100 HBM3 | 11.1% | 11.90k | 1.38s | 12.5 GiB | -| RL Full | 16384 | Recompute | FA2 | 1xA6000 | 10.9% | 3.69k | 4.44s | 12.5 GiB | -| RL Full | 65536 | Recompute | FA3 | 1xH100 HBM3 | 26.8% | 10.15k | 6.46s | 19.6 GiB | -| RL Full | 65536 | Offload | FA3 | 1xH100 HBM3 | 26.4% | 9.98k | 6.57s | 16.1 GiB | -| RL Full | 65536 | Recompute | FA2 | 1xA6000 | 17.4% | 2.08k | 31.46s | 19.5 GiB | -| RL LoRA(r=16) | 16384 | Recompute | FA3 | 1xH100 HBM3 | 10.9% | 13.34k | 1.23s | 5.1 GiB | -| RL LoRA(r=16) | 16384 | Recompute | FA2 | 1xA6000 | 6.3% | 2.41k | 6.79s | 5.1 GiB | -| RL LoRA(r=16) | 65536 | Recompute | FA3 | 1xH100 HBM3 | 25.2% | 9.97k | 6.58s | 11.5 GiB | -| RL LoRA(r=16) | 65536 | Offload | FA3 | 1xH100 HBM3 | 24.7% | 9.78k | 6.70s | 7.9 GiB | -| RL LoRA(r=16) | 65536 | Recompute | FA2 | 1xA6000 | 16.0% | 1.99k | 32.85s | 11.5 GiB | -| SFT Full | 8192 | Recompute | FA3 | 1xH100 HBM3 | 16.8% | 26.01k | 0.32s | 31.7 GiB | -| SFT Full | 8192 | Recompute | FA2 | 1xA6000 | 13.5% | 6.60k | 1.24s | 31.2 GiB | -| SFT Full | 16384 | Recompute | FA3 | 1xH100 HBM3 | 23.7% | 25.49k | 0.64s | 52.8 GiB | +| Type | SeqLen | AC | Attn | EP | CP | Hardware | MFU | TPS | Step Time | Peak Mem | +|------|--------|----|----|----|----|----------|-----|-----|-----------|----------| +| RL Full | 16384 | Recompute | FA3 | 1 | 1 | 1xH100 HBM3 | 11.1% | 11.90k | 1.38s | 12.5 GiB | +| RL Full | 16384 | Recompute | FA2 | 1 | 1 | 1xA6000 | 10.9% | 3.69k | 4.44s | 12.5 GiB | +| RL Full | 65536 | Recompute | FA3 | 1 | 1 | 1xH100 HBM3 | 26.8% | 10.15k | 6.46s | 19.6 GiB | +| RL Full | 65536 | Offload | FA3 | 1 | 1 | 1xH100 HBM3 | 26.4% | 9.98k | 6.57s | 16.1 GiB | +| RL Full | 65536 | Recompute | FA2 | 1 | 1 | 4xA6000 | 16.2% | 7.72k | 33.94s | 12.4 GiB | +| RL Full | 65536 | Recompute | FA2 | 1 | 1 | 1xA6000 | 17.1% | 2.04k | 32.12s | 19.5 GiB | +| RL LoRA(r=16) | 16384 | Recompute | FA3 | 1 | 1 | 1xH100 HBM3 | 10.9% | 13.34k | 1.23s | 5.1 GiB | +| RL LoRA(r=16) | 16384 | Recompute | FA2 | 1 | 1 | 1xA6000 | 6.3% | 2.41k | 6.79s | 5.1 GiB | +| RL LoRA(r=16) | 65536 | Recompute | FA3 | 1 | 1 | 1xH100 HBM3 | 25.2% | 9.97k | 6.58s | 11.5 GiB | +| RL LoRA(r=16) | 65536 | Offload | FA3 | 1 | 1 | 1xH100 HBM3 | 24.7% | 9.78k | 6.70s | 7.9 GiB | +| RL LoRA(r=16) | 65536 | Recompute | FA2 | 1 | 1 | 1xA6000 | 16.0% | 1.99k | 32.85s | 11.5 GiB | +| SFT Full | 8192 | Recompute | FA3 | 1 | 1 | 1xH100 HBM3 | 16.8% | 26.01k | 0.32s | 31.7 GiB | +| SFT Full | 8192 | Recompute | FA2 | 1 | 1 | 1xA6000 | 13.5% | 6.60k | 1.24s | 31.2 GiB | +| SFT Full | 16384 | Recompute | FA3 | 1 | 1 | 1xH100 HBM3 | 23.7% | 25.49k | 0.64s | 52.8 GiB | ## Qwen3-235B-A22B-Instruct-2507 -| Type | SeqLen | AC | Attn | Hardware | MFU | TPS | Step Time | Peak Mem | -|------|--------|----|----|----------|-----|-----|-----------|----------| -| RL LoRA(r=16) | 16384 | Recompute | FA2 | 8xB200 | 4.2% | 4.14k | 31.69s | 161.5 GiB | -| RL LoRA(r=16) | 16384 | Offload | FA2 | 8xB200 | 4.2% | 4.13k | 31.76s | 150.3 GiB | +| Type | SeqLen | AC | Attn | EP | CP | Hardware | MFU | TPS | Step Time | Peak Mem | +|------|--------|----|----|----|----|----------|-----|-----|-----------|----------| +| RL LoRA(r=16) | 16384 | Recompute | FA2 | 1 | 1 | 8xB200 | 4.2% | 4.14k | 31.69s | 161.5 GiB | +| RL LoRA(r=16) | 16384 | Offload | FA2 | 1 | 1 | 8xB200 | 4.2% | 4.13k | 31.76s | 150.3 GiB | ## Qwen3-30B-A3B-Instruct-2507 -| Type | SeqLen | AC | Attn | Hardware | MFU | TPS | Step Time | Peak Mem | -|------|--------|----|----|----------|-----|-----|-----------|----------| -| RL Full | 16384 | Recompute | FA3 | 8xH100 HBM3 | 2.9% | 6.11k | 21.44s | 74.6 GiB | -| RL Full | 16384 | Recompute | FA3 | 8xH200 | 2.8% | 5.94k | 22.06s | 74.6 GiB | -| RL Full | 65536 | Recompute | FA3 | 8xH200 | 15.4% | 12.75k | 41.13s | 105.4 GiB | -| RL LoRA(r=16) | 16384 | Recompute | FA3 | 8xH200 | 17.0% | 36.77k | 3.56s | 31.6 GiB | -| RL LoRA(r=16) | 16384 | Recompute | FA3 | 8xH100 HBM3 | 15.3% | 33.10k | 3.96s | 31.6 GiB | -| RL LoRA(r=16) | 16384 | Recompute | FA2 | 8xB200 | 1.9% | 9.27k | 14.14s | 32.0 GiB | -| RL LoRA(r=16) | 65536 | Recompute | FA3 | 8xH200 | 29.2% | 24.41k | 21.48s | 63.7 GiB | -| RL LoRA(r=16) | 65536 | Recompute | FA3 | 8xH100 HBM3 | 27.9% | 23.39k | 22.42s | 63.7 GiB | -| RL LoRA(r=16) | 65536 | Recompute | FA2 | 8xB200 | 6.9% | 13.23k | 39.64s | 64.8 GiB | -| SFT Full | 16384 | Recompute | FA3 | 8xH200 | 16.6% | 35.03k | 3.74s | 106.4 GiB | +| Type | SeqLen | AC | Attn | EP | CP | Hardware | MFU | TPS | Step Time | Peak Mem | +|------|--------|----|----|----|----|----------|-----|-----|-----------|----------| +| RL Full | 16384 | Recompute | FA3 | 1 | 1 | 8xH100 HBM3 | 2.9% | 6.11k | 21.44s | 74.6 GiB | +| RL Full | 16384 | Recompute | FA3 | 1 | 1 | 8xH200 | 2.8% | 5.94k | 22.06s | 74.6 GiB | +| RL Full | 65536 | Recompute | FA3 | 1 | 1 | 8xH200 | 15.4% | 12.75k | 41.13s | 105.4 GiB | +| RL LoRA(r=16) | 16384 | Recompute | FA3 | 1 | 1 | 8xH200 | 17.0% | 36.77k | 3.56s | 31.6 GiB | +| RL LoRA(r=16) | 16384 | Recompute | FA3 | 1 | 1 | 8xH100 HBM3 | 15.3% | 33.10k | 3.96s | 31.6 GiB | +| RL LoRA(r=16) | 16384 | Recompute | FA2 | 1 | 1 | 8xB200 | 1.9% | 9.27k | 14.14s | 32.0 GiB | +| RL LoRA(r=16) | 65536 | Recompute | FA3 | 1 | 1 | 8xH200 | 29.2% | 24.41k | 21.48s | 63.7 GiB | +| RL LoRA(r=16) | 65536 | Recompute | FA3 | 1 | 1 | 8xH100 HBM3 | 27.9% | 23.39k | 22.42s | 63.7 GiB | +| RL LoRA(r=16) | 65536 | Recompute | FA2 | 1 | 1 | 8xB200 | 6.9% | 13.23k | 39.64s | 64.8 GiB | +| SFT Full | 16384 | Recompute | FA3 | 1 | 1 | 8xH200 | 16.6% | 35.03k | 3.74s | 106.4 GiB | ## Qwen3-4B-Instruct-2507 -| Type | SeqLen | AC | Attn | Hardware | MFU | TPS | Step Time | Peak Mem | -|------|--------|----|----|----------|-----|-----|-----------|----------| -| RL Full | 16384 | Recompute | FA2 | 8xB200 | 6.5% | 27.54k | 4.76s | 17.0 GiB | -| RL Full | 16384 | Recompute | FA3 | 8xH200 | 14.7% | 27.44k | 4.78s | 17.1 GiB | -| RL Full | 16384 | Recompute | FA3 | 8xH100 HBM3 | 13.9% | 26.03k | 5.04s | 17.1 GiB | -| RL Full | 65536 | Recompute | FA3 | 8xH200 | 36.1% | 29.54k | 17.75s | 36.1 GiB | -| RL Full | 65536 | Recompute | FA3 | 8xH100 HBM3 | 35.4% | 28.97k | 18.10s | 36.1 GiB | -| RL LoRA(r=16) | 16384 | Recompute | FA3 | 8xH200 | 22.8% | 52.38k | 2.50s | 10.3 GiB | -| RL LoRA(r=16) | 16384 | Recompute | FA3 | 8xH100 HBM3 | 21.2% | 48.72k | 2.69s | 10.3 GiB | -| RL LoRA(r=16) | 16384 | Recompute | FA2 | 8xB200 | 8.8% | 46.08k | 2.84s | 10.2 GiB | -| RL LoRA(r=16) | 16384 | Recompute | FA3 | 1xH100 HBM3 | 21.7% | 6.23k | 2.63s | 23.3 GiB | -| RL LoRA(r=16) | 16384 | Offload | FA3 | 1xH100 HBM3 | 20.7% | 5.95k | 2.75s | 20.6 GiB | -| RL LoRA(r=16) | 16384 | Recompute | FA2 | 1xA6000 | 13.6% | 1.23k | 13.29s | 23.3 GiB | -| RL LoRA(r=16) | 65536 | Recompute | FA3 | 8xH200 | 37.5% | 33.43k | 15.68s | 28.8 GiB | -| RL LoRA(r=16) | 65536 | Recompute | FA3 | 8xH100 HBM3 | 36.8% | 32.80k | 15.98s | 28.8 GiB | -| RL LoRA(r=16) | 65536 | Recompute | FA2 | 8xB200 | 12.3% | 25.04k | 20.94s | 28.8 GiB | -| RL LoRA(r=16) | 65536 | Recompute | FA3 | 1xH100 HBM3 | 37.4% | 4.16k | 15.74s | 42.0 GiB | -| RL LoRA(r=16) | 65536 | Offload | FA3 | 1xH100 HBM3 | 36.3% | 4.04k | 16.21s | 31.1 GiB | -| SFT Full | 16384 | Recompute | FA2 | 8xB200 | 16.0% | 68.36k | 1.92s | 54.6 GiB | -| SFT Full | 16384 | Recompute | FA2 | 8xH200 | 28.4% | 53.14k | 2.47s | 54.6 GiB | -| SFT Full | 16384 | Recompute | FA2 | 8xH100 HBM3 | 26.6% | 49.72k | 2.64s | 54.6 GiB | -| SFT Full | 65536 | Recompute | FA2 | 8xB200 | 14.2% | 26.39k | 19.86s | 171.5 GiB | - -## Failed Benchmarks - -- **Qwen/Qwen3-30B-A3B-Instruct-2507** (RL Full) on 8xB200: Non-zero exit code: 1 -- **Qwen/Qwen3-0.6B** (SFT Full) on 1xH100 HBM3: [rank0]: torch.OutOfMemoryError: CUDA out of memory. Tried to allocate 37.09 GiB. GPU 0 has a total capacity of 79.18 GiB of which 34.48 GiB is free. Including non-PyTorch memory, this process has 44. -- **Qwen/Qwen3-235B-A22B-Instruct-2507** (RL LoRA(r=16)) on 8xB200: Non-zero exit code: 1 -- **Qwen/Qwen3-235B-A22B-Instruct-2507** (RL LoRA(r=16)) on 8xB200: [rank4]: torch.OutOfMemoryError: CUDA out of memory. Tried to allocate 8.00 GiB. GPU 4 has a total capacity of 178.36 GiB of which 7.20 GiB is free. Including non-PyTorch memory, this process has 171. -- **Qwen/Qwen3-30B-A3B-Instruct-2507** (SFT Full) on 8xB200: Non-zero exit code: 1 -- **Qwen/Qwen3-4B-Instruct-2507** (RL Full) on 8xB200: Non-zero exit code: 1 -- **Qwen/Qwen3-30B-A3B-Instruct-2507** (RL Full) on 8xB200: Non-zero exit code: 1 +| Type | SeqLen | AC | Attn | EP | CP | Hardware | MFU | TPS | Step Time | Peak Mem | +|------|--------|----|----|----|----|----------|-----|-----|-----------|----------| +| RL Full | 16384 | Recompute | FA2 | 1 | 1 | 8xB200 | 6.5% | 27.54k | 4.76s | 17.0 GiB | +| RL Full | 16384 | Recompute | FA3 | 1 | 1 | 8xH200 | 14.7% | 27.44k | 4.78s | 17.1 GiB | +| RL Full | 16384 | Recompute | FA3 | 1 | 1 | 8xH100 HBM3 | 13.9% | 26.03k | 5.04s | 17.1 GiB | +| RL Full | 65536 | Recompute | FA3 | 1 | 1 | 8xH200 | 36.1% | 29.54k | 17.75s | 36.1 GiB | +| RL Full | 65536 | Recompute | FA3 | 1 | 1 | 8xH100 HBM3 | 35.4% | 28.97k | 18.10s | 36.1 GiB | +| RL LoRA(r=16) | 16384 | Recompute | FA3 | 1 | 1 | 8xH200 | 22.8% | 52.38k | 2.50s | 10.3 GiB | +| RL LoRA(r=16) | 16384 | Recompute | FA3 | 1 | 1 | 8xH100 HBM3 | 21.2% | 48.72k | 2.69s | 10.3 GiB | +| RL LoRA(r=16) | 16384 | Recompute | FA2 | 1 | 1 | 8xB200 | 8.8% | 46.08k | 2.84s | 10.2 GiB | +| RL LoRA(r=16) | 16384 | Recompute | FA3 | 1 | 1 | 1xH100 HBM3 | 21.7% | 6.23k | 2.63s | 23.3 GiB | +| RL LoRA(r=16) | 16384 | Offload | FA3 | 1 | 1 | 1xH100 HBM3 | 20.7% | 5.95k | 2.75s | 20.6 GiB | +| RL LoRA(r=16) | 16384 | Recompute | FA2 | 1 | 1 | 1xA6000 | 13.6% | 1.23k | 13.29s | 23.3 GiB | +| RL LoRA(r=16) | 65536 | Recompute | FA3 | 1 | 1 | 8xH200 | 37.5% | 33.43k | 15.68s | 28.8 GiB | +| RL LoRA(r=16) | 65536 | Recompute | FA3 | 1 | 1 | 8xH100 HBM3 | 36.8% | 32.80k | 15.98s | 28.8 GiB | +| RL LoRA(r=16) | 65536 | Recompute | FA2 | 1 | 1 | 8xB200 | 12.3% | 25.04k | 20.94s | 28.8 GiB | +| RL LoRA(r=16) | 65536 | Recompute | FA3 | 1 | 1 | 1xH100 HBM3 | 37.4% | 4.16k | 15.74s | 42.0 GiB | +| RL LoRA(r=16) | 65536 | Offload | FA3 | 1 | 1 | 1xH100 HBM3 | 36.3% | 4.04k | 16.21s | 31.1 GiB | +| SFT Full | 16384 | Recompute | FA2 | 1 | 1 | 8xB200 | 16.0% | 68.36k | 1.92s | 54.6 GiB | +| SFT Full | 16384 | Recompute | FA2 | 1 | 1 | 8xH200 | 28.4% | 53.14k | 2.47s | 54.6 GiB | +| SFT Full | 16384 | Recompute | FA2 | 1 | 1 | 8xH100 HBM3 | 26.6% | 49.72k | 2.64s | 54.6 GiB | +| SFT Full | 65536 | Recompute | FA2 | 1 | 1 | 8xB200 | 14.2% | 26.39k | 19.86s | 171.5 GiB | diff --git a/benchmarks/scripts/aggregate_results.py b/benchmarks/scripts/aggregate_results.py index 5bfd9fed5a2..f7c05f64a4a 100644 --- a/benchmarks/scripts/aggregate_results.py +++ b/benchmarks/scripts/aggregate_results.py @@ -11,18 +11,19 @@ from pydantic import Field -from prime_rl.utils.pydantic_config import BaseSettings, parse_argv +from prime_rl.utils.config import BaseConfig, cli SHORTENED_ATTN_MAPPING = { "flash_attention_2": "FA2", "flash_attention_3": "FA3", + "flash_attention_4": "FA4", } # These words are stripped from the device name to get the short name DEVICE_NAME_STRIP_WORDS = ["NVIDIA", "RTX", "80GB", "40GB"] -class AggregateConfig(BaseSettings): +class AggregateConfig(BaseConfig): """Configuration for aggregating benchmark results.""" artifacts_dir: Annotated[Path, Field(description="Directory containing benchmark JSON artifacts")] @@ -196,7 +197,7 @@ def generate_markdown( def main(): - config = parse_argv(AggregateConfig) + config = cli(AggregateConfig) results = load_json_dir(config.artifacts_dir) print(f"Loaded {len(results)} benchmark results", file=sys.stderr) diff --git a/benchmarks/scripts/run_single_benchmark.py b/benchmarks/scripts/run_single_benchmark.py index 10f701e8112..3c8a5f95528 100644 --- a/benchmarks/scripts/run_single_benchmark.py +++ b/benchmarks/scripts/run_single_benchmark.py @@ -19,7 +19,7 @@ import torch from pydantic import Field -from prime_rl.utils.pydantic_config import BaseSettings, parse_argv +from prime_rl.utils.config import BaseConfig, cli MAX_LORAS = 4 @@ -46,7 +46,7 @@ def extract_oom_error_reason(output: str) -> str | None: return None -class BenchmarkConfig(BaseSettings): +class BenchmarkConfig(BaseConfig): """Configuration for running a single benchmark.""" type: Annotated[ @@ -62,12 +62,18 @@ class BenchmarkConfig(BaseSettings): seq_len: Annotated[int, Field(ge=1, description="Sequence length")] = 512 - ac: Annotated[Literal["Recompute", "Offload", "None"], Field(description="Activation checkpointing type")] = ( - "Recompute" - ) + ac: Annotated[ + Literal["Recompute", "Selective", "Offload"] | None, + Field(description="Activation checkpointing type"), + ] = "Recompute" + + selective_targets: Annotated[ + list[str] | None, + Field(description="Selective activation checkpoint targets when ac=Selective"), + ] = None attention: Annotated[ - Literal["sdpa", "flash_attention_2", "flash_attention_3"], + Literal["sdpa", "flash_attention_2", "flash_attention_3", "flash_attention_4"], Field(description="Attention implementation"), ] = "flash_attention_2" @@ -83,6 +89,11 @@ class BenchmarkConfig(BaseSettings): cp: Annotated[int, Field(ge=1, description="Context parallelism size (1 = no CP)")] = 1 + fused_lm_head_token_chunk_size: Annotated[ + int | None, + Field(description="Fused LM head token chunk size (None uses trainer default)"), + ] = None + docker_image: Annotated[str | None, Field(description="Docker image used for the benchmark")] = None # Metadata set by the script @@ -136,6 +147,10 @@ def build_command(config: BenchmarkConfig) -> list[str]: # Add activation checkpointing if enabled if config.ac == "Recompute": cmd.append("--model.ac") + elif config.ac == "Selective": + cmd.extend(["--model.ac", "--model.ac.mode", "selective"]) + if config.selective_targets: + cmd.extend(["--model.ac.targets", json.dumps(config.selective_targets)]) elif config.ac == "Offload": cmd.append("--model.ac-offloading") @@ -152,6 +167,10 @@ def build_command(config: BenchmarkConfig) -> list[str]: if config.cp > 1: cmd.extend(["--model.cp", str(config.cp)]) + # Fused LM head chunk size + if config.fused_lm_head_token_chunk_size is not None: + cmd.extend(["--model.fused-lm-head-token-chunk-size", str(config.fused_lm_head_token_chunk_size)]) + # Data configuration differs between RL and SFT if config.type.startswith("rl"): cmd.extend(["--data.fake.batch-size", str(config.micro_batches * config.num_gpus)]) @@ -239,7 +258,7 @@ def run_benchmark(config: BenchmarkConfig) -> None: def main(): - config = parse_argv(BenchmarkConfig) + config = cli(BenchmarkConfig) run_benchmark(config) diff --git a/configs/acereason_math/stage1.toml b/configs/acereason_math/stage1.toml index 8aac30c362e..9dbffe564ec 100644 --- a/configs/acereason_math/stage1.toml +++ b/configs/acereason_math/stage1.toml @@ -1,9 +1,10 @@ -inference_gpu_ids = [0,1,2,3,4,5] -trainer_gpu_ids = [6,7] - max_steps = 500 seq_len = 8192 +[deployment] +num_train_gpus = 2 +num_infer_gpus = 6 + [model] name = "deepseek-ai/DeepSeek-R1-Distill-Qwen-7B" @@ -18,11 +19,11 @@ name = "stage1" batch_size = 1024 rollouts_per_example = 8 -[orchestrator.sampling] +[orchestrator.train.sampling] temperature = 0.6 -max_tokens = 8192 +max_completion_tokens = 8192 -[[orchestrator.env]] +[[orchestrator.train.env]] id = "math-env" args = { dataset_name = "nvidia/AceReason-Math", dataset_subset = "default", question_key = "problem", math_verify_max_workers = 128, math_verify_timeout = 60 } diff --git a/configs/acereason_math/stage2.toml b/configs/acereason_math/stage2.toml index c773b494060..a3880c78cb5 100644 --- a/configs/acereason_math/stage2.toml +++ b/configs/acereason_math/stage2.toml @@ -1,9 +1,10 @@ -inference_gpu_ids = [0,1,2,3,4,5] -trainer_gpu_ids = [6,7] - max_steps = 1000 seq_len = 16384 +[deployment] +num_train_gpus = 2 +num_infer_gpus = 6 + [model] name = "deepseek-ai/DeepSeek-R1-Distill-Qwen-7B" @@ -19,11 +20,11 @@ name = "stage2" batch_size = 2048 rollouts_per_example = 16 -[orchestrator.sampling] +[orchestrator.train.sampling] temperature = 0.6 -max_tokens = 16384 +max_completion_tokens = 16384 -[[orchestrator.env]] +[[orchestrator.train.env]] id = "math-env" args = { dataset_name = "nvidia/AceReason-Math", dataset_subset = "default", question_key = "problem", math_verify_max_workers = 128, math_verify_timeout = 60 } diff --git a/configs/alphabet_sort/rl.toml b/configs/alphabet_sort/rl.toml index a8df0c2d98a..22daacdc5fa 100644 --- a/configs/alphabet_sort/rl.toml +++ b/configs/alphabet_sort/rl.toml @@ -13,14 +13,24 @@ name = "Qwen/Qwen3-4B-Instruct-2507" batch_size = 512 rollouts_per_example = 16 -[orchestrator.sampling] -max_tokens = 512 +[orchestrator.train.sampling] +max_completion_tokens = 512 -[[orchestrator.env]] -id = "primeintellect/alphabet-sort" +[[orchestrator.train.env]] +id = "alphabet-sort" name = "alphabet-sort" args = { min_turns = 2, max_turns = 2 } +[orchestrator.eval] +interval = 50 + +[[orchestrator.eval.env]] +id = "alphabet-sort" +name = "alphabet-sort" +num_examples = 50 +rollouts_per_example = 4 +args = { min_turns = 2, max_turns = 2 } + [trainer] [inference] \ No newline at end of file diff --git a/configs/ci/integration/alphabet_sort/start.toml b/configs/ci/integration/alphabet_sort/start.toml new file mode 100644 index 00000000000..50906f82f2c --- /dev/null +++ b/configs/ci/integration/alphabet_sort/start.toml @@ -0,0 +1,24 @@ +max_steps = 10 +seq_len = 2048 + +[ckpt] + +[model] +name = "PrimeIntellect/Qwen3-0.6B" + +[trainer.optim] +lr = 1e-5 + +[orchestrator] +batch_size = 128 +rollouts_per_example = 8 + +[orchestrator.train.sampling] +max_completion_tokens = 384 + +[[orchestrator.train.env]] +id = "alphabet-sort" +name = "alphabet-sort" +args = { min_turns = 2, max_turns = 2, min_names_per_turn = 1, max_names_per_turn = 3, similarity_power = 4, power_per_turn = false } + +[inference] diff --git a/configs/ci/integration/alphabet_sort_branch/start.toml b/configs/ci/integration/alphabet_sort_branch/start.toml new file mode 100644 index 00000000000..e3acaec5448 --- /dev/null +++ b/configs/ci/integration/alphabet_sort_branch/start.toml @@ -0,0 +1,29 @@ +max_steps = 10 +seq_len = 2048 + +[deployment] +type = "single_node" +num_infer_gpus = 1 +num_train_gpus = 1 + +[ckpt] + +[model] +name = "Qwen/Qwen3-0.6B" + +[trainer.optim] +lr = 1e-5 + +[orchestrator] +batch_size = 128 +rollouts_per_example = 8 + +[orchestrator.train.sampling] +max_completion_tokens = 1024 + +[[orchestrator.train.env]] +id = "primeintellect/alphabet-sort" +name = "alphabet-sort" +args = { min_turns = 2, max_turns = 2, min_names_per_turn = 1, max_names_per_turn = 3, similarity_power = 4, power_per_turn = false } + +[inference] diff --git a/configs/ci/integration/eval/multi_env.toml b/configs/ci/integration/eval/multi_env.toml deleted file mode 100644 index a47f91f96bb..00000000000 --- a/configs/ci/integration/eval/multi_env.toml +++ /dev/null @@ -1,15 +0,0 @@ -num_examples = 5 -rollouts_per_example = 3 - -[sampling] -max_tokens = 16 - -[[env]] -id = "primeintellect/math-env" -name = "gsm8k" -args = { dataset_name = "openai/gsm8k", dataset_subset = "main" } - -[[env]] -id = "primeintellect/math-env" -name = "hendrycks" -args = { dataset_name = "PrimeIntellect/Hendrycks-Math", dataset_subset = "default" } \ No newline at end of file diff --git a/configs/ci/integration/eval/single_env.toml b/configs/ci/integration/eval/single_env.toml deleted file mode 100644 index ab4e3f97e9a..00000000000 --- a/configs/ci/integration/eval/single_env.toml +++ /dev/null @@ -1,9 +0,0 @@ -num_examples = 5 -rollouts_per_example = 3 - -[sampling] -max_tokens = 16 - -[[env]] -id = "primeintellect/math-env" -args = { dataset_name = "openai/gsm8k", dataset_subset = "main" } diff --git a/configs/ci/integration/rl/resume.toml b/configs/ci/integration/rl/resume.toml index e19494c13bd..6a01d75b74e 100644 --- a/configs/ci/integration/rl/resume.toml +++ b/configs/ci/integration/rl/resume.toml @@ -1,6 +1,3 @@ -inference_gpu_ids = [0] -trainer_gpu_ids = [1] - max_steps = 25 seq_len = 2048 @@ -17,10 +14,10 @@ lr = 3e-6 batch_size = 128 rollouts_per_example = 16 -[orchestrator.sampling] -max_tokens = 128 +[orchestrator.train.sampling] +max_completion_tokens = 128 -[[orchestrator.env]] +[[orchestrator.train.env]] id = "reverse-text" [inference] \ No newline at end of file diff --git a/configs/ci/integration/rl/start.toml b/configs/ci/integration/rl/start.toml index 91e67d8ae0a..8abdf34059c 100644 --- a/configs/ci/integration/rl/start.toml +++ b/configs/ci/integration/rl/start.toml @@ -1,6 +1,3 @@ -inference_gpu_ids = [0] -trainer_gpu_ids = [1] - max_steps = 20 seq_len = 2048 @@ -16,10 +13,10 @@ lr = 3e-6 batch_size = 128 rollouts_per_example = 16 -[orchestrator.sampling] -max_tokens = 128 +[orchestrator.train.sampling] +max_completion_tokens = 128 -[[orchestrator.env]] +[[orchestrator.train.env]] id = "reverse-text" [inference] \ No newline at end of file diff --git a/configs/ci/integration/rl_lora/resume.toml b/configs/ci/integration/rl_lora/resume.toml index 3b4ccc4aebb..c4a17df3adf 100644 --- a/configs/ci/integration/rl_lora/resume.toml +++ b/configs/ci/integration/rl_lora/resume.toml @@ -1,6 +1,3 @@ -inference_gpu_ids = [0] -trainer_gpu_ids = [1] - max_steps = 25 seq_len = 2048 @@ -26,10 +23,10 @@ rollouts_per_example = 16 [orchestrator.model.lora] name = "r8-1e-4" -[orchestrator.sampling] -max_tokens = 128 +[orchestrator.train.sampling] +max_completion_tokens = 128 -[[orchestrator.env]] +[[orchestrator.train.env]] id = "reverse-text" [inference] diff --git a/configs/ci/integration/rl_lora/start.toml b/configs/ci/integration/rl_lora/start.toml index 03e43bd4657..d8f0ff09953 100644 --- a/configs/ci/integration/rl_lora/start.toml +++ b/configs/ci/integration/rl_lora/start.toml @@ -1,6 +1,3 @@ -inference_gpu_ids = [0] -trainer_gpu_ids = [1] - max_steps = 20 seq_len = 2048 @@ -25,10 +22,10 @@ rollouts_per_example = 16 [orchestrator.model.lora] name = "r8-1e-4" -[orchestrator.sampling] -max_tokens = 128 +[orchestrator.train.sampling] +max_completion_tokens = 128 -[[orchestrator.env]] +[[orchestrator.train.env]] id = "reverse-text" [inference] diff --git a/configs/ci/integration/rl_moe/start.toml b/configs/ci/integration/rl_moe/start.toml new file mode 100644 index 00000000000..067b53f5278 --- /dev/null +++ b/configs/ci/integration/rl_moe/start.toml @@ -0,0 +1,26 @@ +max_steps = 20 +seq_len = 2048 + +[ckpt] + +[model] +name = "samsja/mini-glm-moe" + +[trainer.optim] +lr = 3e-6 + +[orchestrator] +batch_size = 128 +rollouts_per_example = 16 + +[orchestrator.train.sampling] +max_completion_tokens = 128 + +[[orchestrator.train.env]] +id = "reverse-text" + +[inference] +gpu_memory_utilization = 0.7 + +[inference.model] +max_model_len = 2048 diff --git a/configs/ci/integration/rl_multi_run/orchestrator.toml b/configs/ci/integration/rl_multi_run/orchestrator.toml index 2d3c279308d..be34be66fbf 100644 --- a/configs/ci/integration/rl_multi_run/orchestrator.toml +++ b/configs/ci/integration/rl_multi_run/orchestrator.toml @@ -11,10 +11,10 @@ name = "PrimeIntellect/Qwen3-0.6B-Reverse-Text-SFT" [optim] lr = 3e-5 -[sampling] -max_tokens = 128 +[train.sampling] +max_completion_tokens = 128 -[[env]] +[[train.env]] id = "reverse-text" [ckpt] diff --git a/configs/ci/integration/rl_multi_run/trainer.toml b/configs/ci/integration/rl_multi_run/trainer.toml index cbe9bd6602e..38a683b273f 100644 --- a/configs/ci/integration/rl_multi_run/trainer.toml +++ b/configs/ci/integration/rl_multi_run/trainer.toml @@ -5,6 +5,7 @@ max_async_level = 5 [model] name = "PrimeIntellect/Qwen3-0.6B-Reverse-Text-SFT" seq_len = 65536 +fused_lm_head_token_chunk_size = 8192 [model.lora] rank = 8 diff --git a/configs/ci/integration/sft_lora/resume.toml b/configs/ci/integration/sft_lora/resume.toml new file mode 100644 index 00000000000..bcc9145e98f --- /dev/null +++ b/configs/ci/integration/sft_lora/resume.toml @@ -0,0 +1,30 @@ +max_steps = 20 + +[ckpt] +resume_step = 10 + +[ckpt.weights] +save_adapter_separately = true + +[model] +name = "PrimeIntellect/Qwen3-0.6B" + +[model.lora] +rank = 8 +target_modules = [ + "q_proj", + "k_proj", + "v_proj", + "o_proj", + "gate_proj", + "up_proj", + "down_proj", +] + +[data] +name = "PrimeIntellect/Reverse-Text-SFT" +batch_size = 4 +seq_len = 1024 + +[optim] +lr = 1.5e-5 diff --git a/configs/ci/integration/sft_lora/start.toml b/configs/ci/integration/sft_lora/start.toml new file mode 100644 index 00000000000..0d775bdfa65 --- /dev/null +++ b/configs/ci/integration/sft_lora/start.toml @@ -0,0 +1,29 @@ +max_steps = 10 + +[ckpt] + +[ckpt.weights] +save_adapter_separately = true + +[model] +name = "PrimeIntellect/Qwen3-0.6B" + +[model.lora] +rank = 8 +target_modules = [ + "q_proj", + "k_proj", + "v_proj", + "o_proj", + "gate_proj", + "up_proj", + "down_proj", +] + +[data] +name = "PrimeIntellect/Reverse-Text-SFT" +batch_size = 4 +seq_len = 1024 + +[optim] +lr = 1.5e-5 diff --git a/configs/ci/integration/synthesize/multi_env.toml b/configs/ci/integration/synthesize/multi_env.toml deleted file mode 100644 index a47f91f96bb..00000000000 --- a/configs/ci/integration/synthesize/multi_env.toml +++ /dev/null @@ -1,15 +0,0 @@ -num_examples = 5 -rollouts_per_example = 3 - -[sampling] -max_tokens = 16 - -[[env]] -id = "primeintellect/math-env" -name = "gsm8k" -args = { dataset_name = "openai/gsm8k", dataset_subset = "main" } - -[[env]] -id = "primeintellect/math-env" -name = "hendrycks" -args = { dataset_name = "PrimeIntellect/Hendrycks-Math", dataset_subset = "default" } \ No newline at end of file diff --git a/configs/ci/integration/synthesize/multi_turn_tool_call.toml b/configs/ci/integration/synthesize/multi_turn_tool_call.toml deleted file mode 100644 index bdeb900d7e5..00000000000 --- a/configs/ci/integration/synthesize/multi_turn_tool_call.toml +++ /dev/null @@ -1,8 +0,0 @@ -num_examples = 5 -rollouts_per_example = 3 - -[sampling] -max_tokens = 16 - -[[env]] -id = "primeintellect/wiki-search" \ No newline at end of file diff --git a/configs/ci/integration/synthesize/single_env.toml b/configs/ci/integration/synthesize/single_env.toml deleted file mode 100644 index ab4e3f97e9a..00000000000 --- a/configs/ci/integration/synthesize/single_env.toml +++ /dev/null @@ -1,9 +0,0 @@ -num_examples = 5 -rollouts_per_example = 3 - -[sampling] -max_tokens = 16 - -[[env]] -id = "primeintellect/math-env" -args = { dataset_name = "openai/gsm8k", dataset_subset = "main" } diff --git a/configs/ci/nightly/acereason_math.toml b/configs/ci/nightly/acereason_math.toml index a8647f20a7a..1a513e1755a 100644 --- a/configs/ci/nightly/acereason_math.toml +++ b/configs/ci/nightly/acereason_math.toml @@ -3,12 +3,13 @@ # - 300 instead of 500 steps # - No checkpointing # - No evals -inference_gpu_ids = [0,1,2,3] -trainer_gpu_ids = [4,5,6,7] - max_steps = 300 seq_len = 8192 +[deployment] +num_train_gpus = 4 +num_infer_gpus = 4 + [model] name = "deepseek-ai/DeepSeek-R1-Distill-Qwen-7B" @@ -16,12 +17,12 @@ name = "deepseek-ai/DeepSeek-R1-Distill-Qwen-7B" batch_size = 1024 rollouts_per_example = 8 -[orchestrator.sampling] +[orchestrator.train.sampling] temperature = 0.6 -max_tokens = 8192 +max_completion_tokens = 8192 -[[orchestrator.env]] -id = "primeintellect/math-env" +[[orchestrator.train.env]] +id = "math-env" name = "acereason-math" args = { dataset_name = "nvidia/AceReason-Math", dataset_subset = "default", question_key = "problem", math_verify_max_workers = 128, math_verify_timeout = 60 } diff --git a/configs/ci/nightly/hendrycks_math.toml b/configs/ci/nightly/hendrycks_math.toml deleted file mode 100644 index 5822223485c..00000000000 --- a/configs/ci/nightly/hendrycks_math.toml +++ /dev/null @@ -1,46 +0,0 @@ -inference_gpu_ids = [0,1,2,3] -trainer_gpu_ids = [4,5,6,7] - -max_steps = 300 -seq_len = 2048 - -[model] -name = "deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B" - -[orchestrator] -batch_size = 1024 -rollouts_per_example = 16 - -[orchestrator.sampling] -max_tokens = 2048 - -[[orchestrator.env]] -id = "primeintellect/math-env" -name = "hendrycks-math" -args = { dataset_name = "PrimeIntellect/Hendrycks-Math", dataset_subset = "default", math_verify_max_workers = 128, math_verify_timeout = 60 } - -[orchestrator.buffer] -easy_fraction = 0.0 -hard_fraction = 0.0 -easy_threshold = 1.0 -hard_threshold = 0.0 - -[orchestrator.val] -interval = 5 -num_examples = 128 - -[orchestrator.eval] -interval = 10 - -[orchestrator.eval.sampling] -max_tokens = 2048 - -[[orchestrator.eval.env]] -id = "primeintellect/math500" -name = "math500" -num_examples = 30 -rollouts_per_example = 4 - -[trainer] # Default trainer config - -[inference] diff --git a/configs/ci/nightly/multimodal_color_codeword.toml b/configs/ci/nightly/multimodal_color_codeword.toml new file mode 100644 index 00000000000..a90fdfb454b --- /dev/null +++ b/configs/ci/nightly/multimodal_color_codeword.toml @@ -0,0 +1,36 @@ +# Based on configs/multimodal/rl_color_codeword.toml but tuned for CI +max_steps = 25 +seq_len = 4096 + +[deployment] +num_train_gpus = 4 +num_infer_gpus = 4 + +[model] +name = "Qwen/Qwen3-VL-4B-Instruct" + +[model.vlm] +vision_encoder_attr = "model.visual" +language_model_attr = "model.language_model" + +[orchestrator] +batch_size = 256 +rollouts_per_example = 16 + +[orchestrator.train.sampling] +max_completion_tokens = 64 + +[[orchestrator.train.env]] +id = "color-codeword" +args = { images_per_turn = 1, max_turns = 3, num_examples = 1000, seed = 42 } + +[trainer] + +[trainer.model] +optimization_dtype = "bfloat16" +reduce_dtype = "bfloat16" + +[trainer.optim] +lr = 3e-6 + +[inference] diff --git a/configs/debug/eval.toml b/configs/debug/eval.toml deleted file mode 100644 index 1d00c580616..00000000000 --- a/configs/debug/eval.toml +++ /dev/null @@ -1,4 +0,0 @@ -use_tqdm = true - -[sampling] -max_tokens = 16 \ No newline at end of file diff --git a/configs/debug/eval/multi_env.toml b/configs/debug/eval/multi_env.toml deleted file mode 100644 index 7b9a364a22a..00000000000 --- a/configs/debug/eval/multi_env.toml +++ /dev/null @@ -1,13 +0,0 @@ -# Example of generating data in two single-turn environments (gsm8k and hendrycks) -num_examples = 16 -rollouts_per_example = 4 - -[[env]] -id = "primeintellect/math-env" -name = "gsm8k" -args = { dataset_name = "openai/gsm8k", dataset_subset = "main" } - -[[env]] -id = "primeintellect/math-env" -name = "hendrycks" -args = { dataset_name = "PrimeIntellect/Hendrycks-Math", dataset_subset = "default" } \ No newline at end of file diff --git a/configs/debug/eval/single_env.toml b/configs/debug/eval/single_env.toml deleted file mode 100644 index 80e37ef4bac..00000000000 --- a/configs/debug/eval/single_env.toml +++ /dev/null @@ -1,7 +0,0 @@ -# Example of generating data in a single single-turn environment (gsm8k) -num_examples = 16 -rollouts_per_example = 4 - -[[env]] -id = "primeintellect/math-env" -args = { dataset_name = "openai/gsm8k", dataset_subset = "main" } diff --git a/configs/debug/orch.toml b/configs/debug/orch.toml index 1195c3720f3..2e34d780baf 100644 --- a/configs/debug/orch.toml +++ b/configs/debug/orch.toml @@ -2,6 +2,6 @@ max_steps = 5 max_async_level = 5 batch_size = 16 -[sampling] -max_tokens = 16 +[train.sampling] +max_completion_tokens = 16 diff --git a/configs/debug/synthesize/multi_env.toml b/configs/debug/synthesize/multi_env.toml deleted file mode 100644 index 7b9a364a22a..00000000000 --- a/configs/debug/synthesize/multi_env.toml +++ /dev/null @@ -1,13 +0,0 @@ -# Example of generating data in two single-turn environments (gsm8k and hendrycks) -num_examples = 16 -rollouts_per_example = 4 - -[[env]] -id = "primeintellect/math-env" -name = "gsm8k" -args = { dataset_name = "openai/gsm8k", dataset_subset = "main" } - -[[env]] -id = "primeintellect/math-env" -name = "hendrycks" -args = { dataset_name = "PrimeIntellect/Hendrycks-Math", dataset_subset = "default" } \ No newline at end of file diff --git a/configs/debug/synthesize/multi_turn.toml b/configs/debug/synthesize/multi_turn.toml deleted file mode 100644 index b047c9a9f8b..00000000000 --- a/configs/debug/synthesize/multi_turn.toml +++ /dev/null @@ -1,7 +0,0 @@ -# Example of generating data in a multi-turn environment (alphabet-sort) -num_examples = 16 -rollouts_per_example = 4 - -[[env]] -id = "primeintellect/alphabet-sort" -args = { min_turns = 2, max_turns = 2 } \ No newline at end of file diff --git a/configs/debug/synthesize/multi_turn_tool_call.toml b/configs/debug/synthesize/multi_turn_tool_call.toml deleted file mode 100644 index af577cdb675..00000000000 --- a/configs/debug/synthesize/multi_turn_tool_call.toml +++ /dev/null @@ -1,6 +0,0 @@ -# Example of generating data in a multi-turn environment with tool calls -num_examples = 16 -rollouts_per_example = 4 - -[[env]] -id = "primeintellect/wiki-search" \ No newline at end of file diff --git a/configs/debug/synthesize/single_turn.toml b/configs/debug/synthesize/single_turn.toml deleted file mode 100644 index 80e37ef4bac..00000000000 --- a/configs/debug/synthesize/single_turn.toml +++ /dev/null @@ -1,7 +0,0 @@ -# Example of generating data in a single single-turn environment (gsm8k) -num_examples = 16 -rollouts_per_example = 4 - -[[env]] -id = "primeintellect/math-env" -args = { dataset_name = "openai/gsm8k", dataset_subset = "main" } diff --git a/configs/deepscaler/eval.toml b/configs/deepscaler/eval.toml deleted file mode 100644 index 24bfee8bb68..00000000000 --- a/configs/deepscaler/eval.toml +++ /dev/null @@ -1,12 +0,0 @@ -[model] -name = "deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B" - -[[env]] -id = "math500" -args = { use_think = true } -rollouts_per_example = 1 - -[[env]] -id = "aime2024" -args = { use_think = true } -rollouts_per_example = 16 diff --git a/configs/deepscaler/stage1.toml b/configs/deepscaler/stage1.toml index a0cbe4a9614..e27f7c6cba5 100644 --- a/configs/deepscaler/stage1.toml +++ b/configs/deepscaler/stage1.toml @@ -1,9 +1,10 @@ -inference_gpu_ids = [0,1,2,3,4,5] -trainer_gpu_ids = [6,7] - max_steps = 500 seq_len = 8192 +[deployment] +num_train_gpus = 2 +num_infer_gpus = 6 + [model] name = "deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B" @@ -18,11 +19,11 @@ interval = 100 batch_size = 1024 rollouts_per_example = 8 -[orchestrator.sampling] +[orchestrator.train.sampling] temperature = 0.6 -max_tokens = 8192 +max_completion_tokens = 8192 -[[orchestrator.env]] +[[orchestrator.train.env]] id = "math-env" name = "deepscaler" args = { dataset_name= "agentica-org/DeepScaleR-Preview-Dataset", dataset_subset = "default", question_key = "problem", answer_key = "solution", math_verify_max_workers = 128, math_verify_timeout = 60 } diff --git a/configs/deepscaler/stage2.toml b/configs/deepscaler/stage2.toml index 655c2510097..267c34ae78b 100644 --- a/configs/deepscaler/stage2.toml +++ b/configs/deepscaler/stage2.toml @@ -1,9 +1,10 @@ -inference_gpu_ids = [0,1,2,3,4,5] -trainer_gpu_ids = [6,7] - max_steps = 1000 seq_len = 16384 +[deployment] +num_train_gpus = 2 +num_infer_gpus = 6 + [model] name = "deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B" @@ -19,11 +20,11 @@ resume_step = 500 batch_size = 1024 rollouts_per_example = 8 -[orchestrator.sampling] +[orchestrator.train.sampling] temperature = 0.6 -max_tokens = 16384 +max_completion_tokens = 16384 -[[orchestrator.env]] +[[orchestrator.train.env]] id = "math-env" name = "deepscaler" args = { dataset_name= "agentica-org/DeepScaleR-Preview-Dataset", dataset_subset = "default", question_key = "problem", answer_key = "solution", math_verify_max_workers = 128, math_verify_timeout = 60 } diff --git a/configs/deepscaler/stage3.toml b/configs/deepscaler/stage3.toml index bd76c09531b..364a1cfa1db 100644 --- a/configs/deepscaler/stage3.toml +++ b/configs/deepscaler/stage3.toml @@ -1,9 +1,10 @@ -inference_gpu_ids = [0,1,2,3,4,5] -trainer_gpu_ids = [6,7] - max_steps = 1500 seq_len = 24576 +[deployment] +num_train_gpus = 2 +num_infer_gpus = 6 + [model] name = "deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B" @@ -19,11 +20,11 @@ resume_step = 1000 batch_size = 1024 rollouts_per_example = 8 -[orchestrator.sampling] +[orchestrator.train.sampling] temperature = 0.6 -max_tokens = 24576 +max_completion_tokens = 24576 -[[orchestrator.env]] +[[orchestrator.train.env]] id = "math-env" name = "deepscaler" args = { dataset_name= "agentica-org/DeepScaleR-Preview-Dataset", dataset_subset = "default", question_key = "problem", answer_key = "solution", math_verify_max_workers = 128, math_verify_timeout = 60 } diff --git a/configs/elastic/README.md b/configs/elastic/README.md new file mode 100644 index 00000000000..3772520eece --- /dev/null +++ b/configs/elastic/README.md @@ -0,0 +1,17 @@ +# Testing Elastic Inference + +To test the elastic inference pool without having the `rl` entrypoint start the inference server: + +1. **Run RL with the elastic config** (orchestrator discovers servers via DNS; it will find `localhost:8000` once the server is started): + + ```bash + uv run rl @ examples/alphabet_sort/rl_elastic.toml + ``` + + The config uses `[orchestrator.client.elastic]` with `hostname = "localhost"` and `port = 8000`. The orchestrator will wait for at least one ready server before starting rollouts. + +2. **Start the inference server manually** (in a separate terminal or on another machine): + + ```bash + uv run inference --model.name Qwen/Qwen3-4B-Instruct-2507 --enable-lora --max-lora-rank 32 + ``` \ No newline at end of file diff --git a/configs/elastic/rl.toml b/configs/elastic/rl.toml new file mode 100644 index 00000000000..387bfddbf2e --- /dev/null +++ b/configs/elastic/rl.toml @@ -0,0 +1,46 @@ +# Alphabet sort RL with elastic inference pool (no inference process started by rl). +# Start the inference server manually first, then run: uv run rl @ examples/alphabet_sort/rl_elastic.toml +# The orchestrator discovers servers via DNS (hostname) and only uses ready ones. +# Omitting [inference] means inference = None, so the rl entrypoint will not start an inference process. + +max_steps = 200 +seq_len = 2048 + +[ckpt] + +[model] +name = "Qwen/Qwen3-4B-Instruct-2507" + +[wandb] +project = "alphabet-sort" +name = "alphabet-sort-elastic" + +[trainer.model] +impl = "auto" + +[trainer.model.ac] +freq = 1 + +[trainer.model.lora] +rank = 32 +alpha = 64 + +[trainer.optim] +lr = 1e-5 + +[orchestrator] +batch_size = 512 +rollouts_per_example = 8 + +[orchestrator.train.sampling] +max_completion_tokens = 768 + +[orchestrator.client.elastic] +hostname = "localhost" +port = 8000 +sync_interval = 5.0 + +[[orchestrator.train.env]] +id = "alphabet-sort" +name = "alphabet-sort" +args = { min_turns = 3, max_turns = 5, power_per_turn = false } \ No newline at end of file diff --git a/configs/env_mix/env_mix.toml b/configs/env_mix/env_mix.toml index 7d482f4ff4d..64d8283fffd 100644 --- a/configs/env_mix/env_mix.toml +++ b/configs/env_mix/env_mix.toml @@ -1,9 +1,10 @@ -trainer_gpu_ids = [0,1,2,3] -inference_gpu_ids = [4,5,6,7] - max_steps = 300 seq_len = 8192 +[deployment] +num_train_gpus = 4 +num_infer_gpus = 4 + [wandb] project = "env-mix-debug" name = "env-mix" @@ -15,27 +16,27 @@ name = "Qwen/Qwen3-4B-Instruct-2507" batch_size = 512 rollouts_per_example = 16 -[[orchestrator.env]] +[[orchestrator.train.env]] id = "math-env" name = "math" args = { min_avg_reward = 0.1, rubric_max_workers = 128 } -[[orchestrator.env]] +[[orchestrator.train.env]] id = "code-env" name = "code" args = { min_avg_reward = 0.1 } -[[orchestrator.env]] +[[orchestrator.train.env]] id = "science-env" name = "science" args = { min_avg_reward = 0.1, pool_size = 128, rubric_max_workers = 128 } -[[orchestrator.env]] +[[orchestrator.train.env]] id = "logic-env" name = "logic" args = { min_avg_reward = 0.1 } -[[orchestrator.env]] +[[orchestrator.train.env]] id = "math-env" name = "math-python" args = { min_avg_reward = 0.1, sandbox_client_max_workers = 128, sandbox_timeout_minutes = 10, sandbox_memory_gb = 1, sandbox_disk_size_gb = 1, pip_install_packages = "numpy sympy", python_tool = true, rubric_max_workers = 128 } @@ -47,6 +48,5 @@ api_server_count = 8 gpu_memory_utilization = 0.8 [inference.model] -enable_auto_tool_choice = true tool_call_parser = "hermes" max_model_len = 8192 \ No newline at end of file diff --git a/configs/glm5_disagg_inference/inference.toml b/configs/glm5_disagg_inference/inference.toml new file mode 100644 index 00000000000..3576eb31a42 --- /dev/null +++ b/configs/glm5_disagg_inference/inference.toml @@ -0,0 +1,19 @@ +enable_expert_parallel = true +enable_eplb = false +use_deep_gemm = true +gpu_memory_utilization = 0.85 + +[model] +name = "zai-org/GLM-5-FP8" +tool_call_parser = "auto" +reasoning_parser = "glm45" + +[deployment] +type = "disaggregated" +num_prefill_nodes = 4 +num_decode_nodes = 2 +gpus_per_node = 8 + +[slurm] +job_name = "glm5" +partition = "cluster" diff --git a/configs/gsm8k/rl.toml b/configs/gsm8k/rl.toml index d313e1ab642..63d5d40a8ea 100644 --- a/configs/gsm8k/rl.toml +++ b/configs/gsm8k/rl.toml @@ -1,6 +1,3 @@ -inference_gpu_ids = [0] -trainer_gpu_ids = [1] - max_steps = 100 seq_len = 2048 @@ -15,11 +12,11 @@ name = "PrimeIntellect/Qwen3-0.6B" batch_size = 512 rollouts_per_example = 16 -[orchestrator.sampling] -max_tokens = 2048 +[orchestrator.train.sampling] +max_completion_tokens = 2048 -[[orchestrator.env]] -id = "primeintellect/math-env" +[[orchestrator.train.env]] +id = "math-env" name = "gsm8k" args = { dataset_name = "openai/gsm8k", dataset_subset = "main", math_verify_max_workers = 128, math_verify_timeout = 60 } diff --git a/configs/hendrycks_math/eval.toml b/configs/hendrycks_math/eval.toml deleted file mode 100644 index 34f572440d4..00000000000 --- a/configs/hendrycks_math/eval.toml +++ /dev/null @@ -1,5 +0,0 @@ -[sampling] -max_tokens = 2048 - -[[env]] -id = "math500" \ No newline at end of file diff --git a/configs/hendrycks_math/rl.toml b/configs/hendrycks_math/rl.toml index cdea9589fb9..acc1434ea95 100644 --- a/configs/hendrycks_math/rl.toml +++ b/configs/hendrycks_math/rl.toml @@ -1,6 +1,3 @@ -trainer_gpu_ids = [0] -inference_gpu_ids = [1] - max_steps = 500 seq_len = 2048 @@ -15,11 +12,11 @@ name = "Qwen/Qwen3-4B-Instruct-2507" batch_size = 512 rollouts_per_example = 16 -[orchestrator.sampling] -max_tokens = 2048 +[orchestrator.train.sampling] +max_completion_tokens = 2048 -[[orchestrator.env]] -id = "primeintellect/math-env" +[[orchestrator.train.env]] +id = "math-env" name = "hendrycks-math" args = { dataset_name = "PrimeIntellect/Hendrycks-Math", dataset_subset = "default", math_verify_max_workers = 128, math_verify_timeout = 60 } @@ -27,18 +24,14 @@ args = { dataset_name = "PrimeIntellect/Hendrycks-Math", dataset_subset = "defau easy_threshold = 1.0 hard_threshold = 0.0 -[orchestrator.val] -interval = 5 -num_examples = 128 - [orchestrator.eval] interval = 10 [orchestrator.eval.sampling] -max_tokens = 2048 +max_completion_tokens = 2048 [[orchestrator.eval.env]] -id = "primeintellect/math500" +id = "math500" num_examples = 30 rollouts_per_example = 4 diff --git a/configs/hendrycks_math/sanity.toml b/configs/hendrycks_math/sanity.toml new file mode 100644 index 00000000000..af95d038d45 --- /dev/null +++ b/configs/hendrycks_math/sanity.toml @@ -0,0 +1,39 @@ +max_steps = 3000 +seq_len = 8192 + +[wandb] +project = "hendrycks-math-debug" +name = "hendrycks-math-sanity" + +[model] +name = "deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B" + +[orchestrator] +batch_size = 512 +rollouts_per_example = 8 + +[[orchestrator.train.env]] +id = "math-env" +args = { dataset_name = "mikasenghaas/Sanity-Test-R1D-1.5B", dataset_subset = "default" } +name = "hendrycks-math" + +[orchestrator.eval] +interval = 50 + +[[orchestrator.eval.env]] +id = "aime2024" +name = "aime2024" +rollouts_per_example = 16 + +[trainer.model] +name = "deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B" + +[trainer.model.ac] + +[trainer.model.compile] + +[inference.model] +max_model_len = 8192 + +[log] +level = "debug" \ No newline at end of file diff --git a/configs/intellect_3/evals/api/deepseek.toml b/configs/intellect_3/evals/api/deepseek.toml deleted file mode 100644 index 4f896347b8e..00000000000 --- a/configs/intellect_3/evals/api/deepseek.toml +++ /dev/null @@ -1,5 +0,0 @@ -max_concurrent = 128 - -[client] -base_url = ["https://api.deepseek.com"] -api_key_var = "DEEPSEEK_API_KEY" \ No newline at end of file diff --git a/configs/intellect_3/evals/api/fireworks.toml b/configs/intellect_3/evals/api/fireworks.toml deleted file mode 100644 index e15afc6a89b..00000000000 --- a/configs/intellect_3/evals/api/fireworks.toml +++ /dev/null @@ -1,6 +0,0 @@ -max_concurrent = 128 - -[client] -base_url = ["https://api.fireworks.ai/inference/v1"] - -api_key_var = "FIREWORKS_API_KEY" \ No newline at end of file diff --git a/configs/intellect_3/evals/api/localhost.toml b/configs/intellect_3/evals/api/localhost.toml deleted file mode 100644 index aaa5e78a2f5..00000000000 --- a/configs/intellect_3/evals/api/localhost.toml +++ /dev/null @@ -1,4 +0,0 @@ -max_concurrent = 1024 - -[client] -base_url = [ "http://localhost:8000/v1" ] \ No newline at end of file diff --git a/configs/intellect_3/evals/api/openrouter/deepseek.toml b/configs/intellect_3/evals/api/openrouter/deepseek.toml deleted file mode 100644 index b0717c95f42..00000000000 --- a/configs/intellect_3/evals/api/openrouter/deepseek.toml +++ /dev/null @@ -1,8 +0,0 @@ -max_concurrent = 256 - -[client] -base_url = ["https://openrouter.ai/api/v1"] -api_key_var = "OPENROUTER_API_KEY" - -[client.extra_body] -provider = { order = ["deepseek", "novita/fp8", "atlas-cloud/fp8", "siliconflow/fp8"] } diff --git a/configs/intellect_3/evals/api/openrouter/zai.toml b/configs/intellect_3/evals/api/openrouter/zai.toml deleted file mode 100644 index 3c12a5448c3..00000000000 --- a/configs/intellect_3/evals/api/openrouter/zai.toml +++ /dev/null @@ -1,8 +0,0 @@ -max_concurrent = 128 - -[client] -base_url = ["https://openrouter.ai/api/v1"] -api_key_var = "OPENROUTER_API_KEY" - -[client.extra_body] -provider = { order = ["z-ai"], allow_fallbacks = false} \ No newline at end of file diff --git a/configs/intellect_3/evals/api/pinference.toml b/configs/intellect_3/evals/api/pinference.toml deleted file mode 100644 index 98ad9b18073..00000000000 --- a/configs/intellect_3/evals/api/pinference.toml +++ /dev/null @@ -1,8 +0,0 @@ -max_concurrent = 16 - -[model] -name = "PrimeIntellect/INTELLECT-3" - -[client] -base_url = ["https://dev-inference.pinference.ai/api/v1"] -api_key_var = "PRIME_API_KEY_DEV" diff --git a/configs/intellect_3/evals/api/together.toml b/configs/intellect_3/evals/api/together.toml deleted file mode 100644 index b0bbe913b83..00000000000 --- a/configs/intellect_3/evals/api/together.toml +++ /dev/null @@ -1,5 +0,0 @@ -max_concurrent = 128 - -[client] -base_url = ["https://api.together.ai/v1"] -api_key_var = "TOGETHER_API_KEY" \ No newline at end of file diff --git a/configs/intellect_3/evals/envs/aime2024.toml b/configs/intellect_3/evals/envs/aime2024.toml deleted file mode 100644 index cc0b564e1ff..00000000000 --- a/configs/intellect_3/evals/envs/aime2024.toml +++ /dev/null @@ -1,3 +0,0 @@ -[[env]] -id = "aime2024" -rollouts_per_example = 32 # 30 examples * 32 rollouts = 960 rollouts \ No newline at end of file diff --git a/configs/intellect_3/evals/envs/aime2025.toml b/configs/intellect_3/evals/envs/aime2025.toml deleted file mode 100644 index a62d658124a..00000000000 --- a/configs/intellect_3/evals/envs/aime2025.toml +++ /dev/null @@ -1,3 +0,0 @@ -[[env]] -id = "aime2025" -rollouts_per_example = 32 # 30 examples * 32 rollouts = 960 rollouts \ No newline at end of file diff --git a/configs/intellect_3/evals/envs/gpqa.toml b/configs/intellect_3/evals/envs/gpqa.toml deleted file mode 100644 index dd79325680c..00000000000 --- a/configs/intellect_3/evals/envs/gpqa.toml +++ /dev/null @@ -1,4 +0,0 @@ -[[env]] -id = "gpqa" -rollouts_per_example = 8 -args = { verifier = "regex" } diff --git a/configs/intellect_3/evals/envs/hle.toml b/configs/intellect_3/evals/envs/hle.toml deleted file mode 100644 index 183f59af364..00000000000 --- a/configs/intellect_3/evals/envs/hle.toml +++ /dev/null @@ -1,3 +0,0 @@ -[[env]] -id = "hle" -rollouts_per_example = 1 # 2158 examples * 1 rollouts = 2158 rollouts \ No newline at end of file diff --git a/configs/intellect_3/evals/envs/livecodebench.toml b/configs/intellect_3/evals/envs/livecodebench.toml deleted file mode 100644 index bf60dc9268f..00000000000 --- a/configs/intellect_3/evals/envs/livecodebench.toml +++ /dev/null @@ -1,4 +0,0 @@ -[[env]] -id = "livecodebench" -args = { pool_size = 128 } -rollouts_per_example = 2 # 454 examples * 2 rollouts = 908 rollouts \ No newline at end of file diff --git a/configs/intellect_3/evals/envs/mmlu_pro.toml b/configs/intellect_3/evals/envs/mmlu_pro.toml deleted file mode 100644 index 68c94c63bbd..00000000000 --- a/configs/intellect_3/evals/envs/mmlu_pro.toml +++ /dev/null @@ -1,3 +0,0 @@ -[[env]] -id = "mmlu-pro" -rollouts_per_example = 1 # 12K examples * 1 rollouts = 12000 rollouts \ No newline at end of file diff --git a/configs/intellect_3/evals/models/README.md b/configs/intellect_3/evals/models/README.md deleted file mode 100644 index c92b4a3ecda..00000000000 --- a/configs/intellect_3/evals/models/README.md +++ /dev/null @@ -1 +0,0 @@ -The model names in this directory match the names on OpenRouter. If you try hitting the model on a different provider, you might have to change the model name. \ No newline at end of file diff --git a/configs/intellect_3/evals/models/deepseek-r1-0528.toml b/configs/intellect_3/evals/models/deepseek-r1-0528.toml deleted file mode 100644 index 468b95a9c23..00000000000 --- a/configs/intellect_3/evals/models/deepseek-r1-0528.toml +++ /dev/null @@ -1,2 +0,0 @@ -[model] -name = "deepseek/deepseek-r1-0528" \ No newline at end of file diff --git a/configs/intellect_3/evals/models/deepseek-v3.2.toml b/configs/intellect_3/evals/models/deepseek-v3.2.toml deleted file mode 100644 index 28642f5a5c2..00000000000 --- a/configs/intellect_3/evals/models/deepseek-v3.2.toml +++ /dev/null @@ -1,2 +0,0 @@ -[model] -name = "deepseek/deepseek-v3.2-exp" \ No newline at end of file diff --git a/configs/intellect_3/evals/models/glm-4.5-air.toml b/configs/intellect_3/evals/models/glm-4.5-air.toml deleted file mode 100644 index c6fc7246177..00000000000 --- a/configs/intellect_3/evals/models/glm-4.5-air.toml +++ /dev/null @@ -1,2 +0,0 @@ -[model] -name = "z-ai/glm-4.5-air" \ No newline at end of file diff --git a/configs/intellect_3/evals/models/glm-4.6.toml b/configs/intellect_3/evals/models/glm-4.6.toml deleted file mode 100644 index 344b7016d17..00000000000 --- a/configs/intellect_3/evals/models/glm-4.6.toml +++ /dev/null @@ -1,2 +0,0 @@ -[model] -name = "z-ai/glm-4.6" \ No newline at end of file diff --git a/configs/intellect_3/evals/models/gpt-oss-120b-vertex.toml b/configs/intellect_3/evals/models/gpt-oss-120b-vertex.toml deleted file mode 100644 index d5543b6330e..00000000000 --- a/configs/intellect_3/evals/models/gpt-oss-120b-vertex.toml +++ /dev/null @@ -1,20 +0,0 @@ -max_concurrent = 64 -reasoning_field = "reasoning" - -[model] -name = "openai/gpt-oss-120b" - - -[client] -base_url = ["https://openrouter.ai/api/v1"] -api_key_var = "OPENROUTER_API_KEY" - -[sampling] -reasoning_effort = "high" -temperature = 1.0 -max_tokens = 122880 - -[sampling.extra_body.provider] -only = ["google-vertex"] -allow_fallbacks = false -require_parameters = true diff --git a/configs/intellect_3/evals/models/intellect-3.toml b/configs/intellect_3/evals/models/intellect-3.toml deleted file mode 100644 index db0f97f2d56..00000000000 --- a/configs/intellect_3/evals/models/intellect-3.toml +++ /dev/null @@ -1,2 +0,0 @@ -[model] -name = "PrimeIntellect/INTELLECT-3" \ No newline at end of file diff --git a/configs/math_group/rl.toml b/configs/math_group/rl.toml new file mode 100644 index 00000000000..965cb1cd1ca --- /dev/null +++ b/configs/math_group/rl.toml @@ -0,0 +1,45 @@ +max_steps = 500 + +[wandb] +project = "math-group-debug" +name = "math-group" + +[model] +name = "Qwen/Qwen3-4B-Instruct-2507" + +[orchestrator] +batch_size = 256 +rollouts_per_example = 8 + +[[orchestrator.train.env]] +id = "math-env" +name = "hendrycks-math" +args = { dataset_name = "PrimeIntellect/Hendrycks-Math", dataset_subset = "default" } +ratio = 0.5 + +[[orchestrator.train.env]] +id = "math-env" +name = "acereason-math" +args = { dataset_name = "nvidia/AceReason-Math", dataset_subset = "default", question_key = "problem" } +ratio = 0.5 + +[orchestrator.buffer] +easy_threshold = 1.0 +hard_threshold = 0.0 + +[orchestrator.eval] +interval = 50 + +[[orchestrator.eval.env]] +id = "aime2024" +name = "aime2024" +num_examples = 30 +rollouts_per_example = 4 + +[trainer.model] +seq_len = 4096 + +[inference.model] +max_model_len = 2048 + + diff --git a/configs/math_python/math_python.toml b/configs/math_python/math_python.toml index b88331abfb4..ae9ec700350 100644 --- a/configs/math_python/math_python.toml +++ b/configs/math_python/math_python.toml @@ -12,10 +12,10 @@ name = "Qwen/Qwen3-4B-Instruct-2507" batch_size = 512 rollouts_per_example = 16 -[orchestrator.sampling] -max_tokens = 512 +[orchestrator.train.sampling] +max_completion_tokens = 512 -[[orchestrator.env]] +[[orchestrator.train.env]] id = "math-python" args = { sandbox_client_max_workers = 128, sandbox_timeout_minutes = 10, sandbox_memory_gb = 1, sandbox_disk_size_gb = 1, pip_install_packages = "numpy sympy" } @@ -25,6 +25,5 @@ args = { sandbox_client_max_workers = 128, sandbox_timeout_minutes = 10, sandbox api_server_count = 4 [inference.model] -enable_auto_tool_choice = true tool_call_parser = "hermes" max_model_len = 4096 \ No newline at end of file diff --git a/configs/multi_reverse_text/rl.toml b/configs/multi_reverse_text/rl.toml index df995b5984c..06bead89465 100644 --- a/configs/multi_reverse_text/rl.toml +++ b/configs/multi_reverse_text/rl.toml @@ -8,19 +8,46 @@ name = "PrimeIntellect/Qwen3-0.6B-Reverse-Text-SFT" batch_size = 128 rollouts_per_example = 16 -[orchestrator.sampling] -max_tokens = 128 +# -- train envs -- -[[orchestrator.env]] +[orchestrator.train.sampling] +max_completion_tokens = 128 + +[[orchestrator.train.env]] id = "reverse-text" +name = "reverse-text-1" + +[orchestrator.train.env.sampling] +max_completion_tokens = 127 -[[orchestrator.env]] +[[orchestrator.train.env]] id = "reverse-text" name = "reverse-text-2" -[orchestrator.val] -interval = 1 +# -- eval envs -- + +[orchestrator.eval] +interval = 5 + +[orchestrator.eval.sampling] +max_completion_tokens = 512 + +[[orchestrator.eval.env]] +id = "reverse-text" +name = "eval-default" +num_examples = 32 +rollouts_per_example = 4 + +[[orchestrator.eval.env]] +id = "reverse-text" +name = "eval-custom" num_examples = 16 +rollouts_per_example = 2 +interval = 10 + +[orchestrator.eval.env.sampling] +max_completion_tokens = 256 +temperature = 0.5 [trainer.optim] lr = 3e-6 diff --git a/configs/multimodal/rl_color_codeword.toml b/configs/multimodal/rl_color_codeword.toml new file mode 100644 index 00000000000..46910ed8a08 --- /dev/null +++ b/configs/multimodal/rl_color_codeword.toml @@ -0,0 +1,35 @@ +max_steps = 15 +seq_len = 4096 + +[model] +name = "Qwen/Qwen3-VL-4B-Instruct" + +[model.vlm] +vision_encoder_attr = "model.visual" +language_model_attr = "model.language_model" + +[orchestrator] +batch_size = 256 +rollouts_per_example = 16 + + +[orchestrator.train.sampling] +max_completion_tokens = 64 + +[[orchestrator.train.env]] +id = "color-codeword" +args = { images_per_turn = 1, max_turns = 3, num_examples = 1000, seed = 42 } + +[trainer] + +[trainer.model] +optimization_dtype = "bfloat16" +reduce_dtype = "bfloat16" + +[trainer.optim] +lr = 3e-6 + +[inference] + +[inference.parallel] +dp = 1 diff --git a/configs/multimodal/rl_color_codeword_test.toml b/configs/multimodal/rl_color_codeword_test.toml new file mode 100644 index 00000000000..151bad0987f --- /dev/null +++ b/configs/multimodal/rl_color_codeword_test.toml @@ -0,0 +1,35 @@ +max_steps = 3 +seq_len = 2048 +output_dir = "outputs/rl_color_codeword_test" + +[model] +name = "Qwen/Qwen3-VL-4B-Instruct" + +[model.vlm] +vision_encoder_attr = "model.visual" +language_model_attr = "model.language_model" + +[orchestrator] +batch_size = 16 +rollouts_per_example = 2 + +[orchestrator.train.sampling] +max_completion_tokens = 32 + +[[orchestrator.train.env]] +id = "color-codeword" +args = { images_per_turn = 1, max_turns = 2, num_examples = 100, seed = 42 } + +[trainer] + +[trainer.model] +optimization_dtype = "bfloat16" +reduce_dtype = "bfloat16" + +[trainer.optim] +lr = 3e-6 + +[inference] + +[inference.parallel] +dp = 1 diff --git a/configs/nemotron_4node/rl.toml b/configs/nemotron_4node/rl.toml new file mode 100644 index 00000000000..0fda06a049d --- /dev/null +++ b/configs/nemotron_4node/rl.toml @@ -0,0 +1,76 @@ +output_dir = "/data/sami_debug/nemotron-debug" +clean_output_dir = true +max_steps = 5 +seq_len = 2048 + +[slurm] +job_name = "nemotron-4node" +pre_run_command = "uv sync --all-extras --group mamba-ssm" + +[deployment] +type = "multi_node" +num_train_nodes = 4 +num_infer_nodes = 4 + +[log] +level = "debug" + +[wandb] +project = "sami_debug_nemotron" +name = "nemotron-120b-math" +shared = true + +[weight_broadcast] +type = "nccl" + +[model] +name = "nvidia/NVIDIA-Nemotron-3-Super-120B-A12B-BF16" + +[trainer] +dist_timeout_seconds = 1800 + +[trainer.model] +impl = "custom" +attn = "flash_attention_2" +optim_cpu_offload = true +ep = 8 +fused_lm_head_token_chunk_size = 8192 + +[trainer.model.ac] +freq = 1 + +[trainer.model.ac_offloading] +max_inflight_activations = 5 + +[trainer.model.compile] + +[orchestrator] +batch_size = 128 +rollouts_per_example = 8 + +[orchestrator.train.sampling] +max_completion_tokens = 2048 + +[[orchestrator.train.env]] +id = "math-env" +name = "hendrycks-math" +args = { dataset_name = "PrimeIntellect/Hendrycks-Math", dataset_subset = "default", math_verify_max_workers = 128, math_verify_timeout = 60 } + +[orchestrator.buffer] +easy_threshold = 1.0 +hard_threshold = 0.0 + +[orchestrator.eval] +interval = 10 + +[orchestrator.eval.sampling] +max_completion_tokens = 2048 + +[[orchestrator.eval.env]] +id = "math500" +num_examples = 30 +rollouts_per_example = 4 + +[inference.parallel] +tp = 8 +dp = 1 diff --git a/configs/nemotron_debug/rl.toml b/configs/nemotron_debug/rl.toml new file mode 100644 index 00000000000..9e1e921b092 --- /dev/null +++ b/configs/nemotron_debug/rl.toml @@ -0,0 +1,75 @@ +output_dir = "outputs/rl-nemotron-debug" +clean_output_dir = true +max_steps = 500 +seq_len = 4096 + +[slurm] +job_name = "nemotron-sami-debug" + +[deployment] +type = "multi_node" +num_train_nodes = 2 +num_infer_nodes = 2 + +[log] +level = "debug" + +[wandb] +project = "nemotron_sami_debug" +name = "nemotron-120b-math" +shared = true + +[weight_broadcast] +type = "nccl" + +[model] +name = "nvidia/NVIDIA-Nemotron-3-Super-120B-A12B-BF16" + +[trainer.model] +impl = "custom" +attn = "flash_attention_2" +optim_cpu_offload = true +fused_lm_head_token_chunk_size = 8192 + +[trainer.model.debug] +random_init = true +num_layers = 8 + +[trainer.model.ac_offloading] +max_inflight_activations = 5 + +[trainer.model.ac] +freq = 1 + +[trainer.model.compile] + +[orchestrator] +batch_size = 512 +rollouts_per_example = 16 + +[orchestrator.train.sampling] +max_completion_tokens = 4096 + +[[orchestrator.train.env]] +id = "math-env" +name = "hendrycks-math" +args = { dataset_name = "PrimeIntellect/Hendrycks-Math", dataset_subset = "default", math_verify_max_workers = 128, math_verify_timeout = 60 } + +[orchestrator.buffer] +easy_threshold = 1.0 +hard_threshold = 0.0 + +[orchestrator.eval] +interval = 10 + +[orchestrator.eval.sampling] +max_completion_tokens = 4096 + +[[orchestrator.eval.env]] +id = "math500" +num_examples = 30 +rollouts_per_example = 4 + +[inference.parallel] +tp = 4 +dp = 4 diff --git a/configs/wiki_search/rl.toml b/configs/wiki_search/rl.toml index adf97dcb78d..d688ba21dc0 100644 --- a/configs/wiki_search/rl.toml +++ b/configs/wiki_search/rl.toml @@ -1,6 +1,3 @@ -inference_gpu_ids = [0] -trainer_gpu_ids = [1] - max_steps = 500 seq_len = 4096 @@ -33,15 +30,14 @@ batch_size = 512 rollouts_per_example = 16 oversampling_factor = 2.0 -[orchestrator.sampling] -max_tokens = 512 +[orchestrator.train.sampling] +max_completion_tokens = 512 [orchestrator.buffer] online_difficulty_filtering = true -[[orchestrator.env]] -id = "primeintellect/wiki-search" +[[orchestrator.train.env]] +id = "wiki-search" [inference.model] -enable_auto_tool_choice = true tool_call_parser = "hermes" diff --git a/docs/bring-your-own-algorithms.md b/docs/bring-your-own-algorithms.md new file mode 100644 index 00000000000..816d9c94c41 --- /dev/null +++ b/docs/bring-your-own-algorithms.md @@ -0,0 +1,143 @@ +# Bring Your Own Algorithms + +Prime-RL supports custom implementations for key algorithmic components, allowing you to experiment with different RL objectives and techniques. + +## 1. Custom Loss Functions + +The loss is computed **per-sequence** (per-sample). You provide a function that computes the loss for a single sequence, and the framework handles iteration and aggregation. + +### Interface + +```python +from prime_rl.trainer.rl.loss import LossInputs, LossOutputs + +def my_custom_loss(inputs: LossInputs, **kwargs) -> LossOutputs: + ... +``` + +#### LossInputs + +```python +@dataclass +class LossInputs: + trainer_logprobs: Float[Tensor, "seq"] # Log probs from current policy + inference_logprobs: Float[Tensor, "seq"] # Log probs from reference policy + teacher_logprobs: Float[Tensor, "seq"] | None # Optional teacher log probs + advantages: Float[Tensor, "seq"] # Per-token advantages + loss_mask: Bool[Tensor, "seq"] # Mask for valid tokens +``` + +#### LossOutputs + +```python +@dataclass +class LossOutputs: + loss: Float[Tensor, ""] # Scalar loss for this sequence + metrics: dict[str, Tensor] # Metrics to log +``` + +### Example: PPO Clipped Loss + +```python +import torch +from prime_rl.trainer.rl.loss import LossInputs, LossOutputs + +def ppo_clip_loss(inputs: LossInputs, clip_eps: float = 0.2) -> LossOutputs: + ratio = torch.exp(inputs.trainer_logprobs - inputs.inference_logprobs) + clipped_ratio = torch.clamp(ratio, 1 - clip_eps, 1 + clip_eps) + + surr1 = ratio * inputs.advantages + surr2 = clipped_ratio * inputs.advantages + + loss = -torch.min(surr1, surr2)[inputs.loss_mask].sum() + + return LossOutputs( + loss=loss, + metrics={"clip_frac": (ratio != clipped_ratio)[inputs.loss_mask].float().mean()}, + ) +``` + +### Configuration + +```toml +[loss] +type = "custom" +import_path = "my_module.ppo_clip_loss" +kwargs = { clip_eps = 0.2 } +``` + +--- + +## 2. Custom Advantage Functions + +Advantages are computed **per-example** (grouped by `rollouts_per_example`). You provide a function that computes advantages for a batch of examples. + +### Interface + +```python +from prime_rl.orchestrator.advantage import AdvantageInputs, AdvantageOutputs + +def my_custom_advantage(inputs: AdvantageInputs, **kwargs) -> AdvantageOutputs: + ... +``` + +#### AdvantageInputs + +```python +@dataclass +class AdvantageInputs: + # Rollouts grouped by problem: rollouts[i][j] is the j-th rollout for problem i. + rollouts: list[list[vf.RolloutOutput]] +``` + +Each `vf.RolloutOutput` carries the full rollout (`reward`, `trajectory`, etc.), so custom advantages can read any metadata they need (e.g. completion-token counts, turn counts, tool calls). + +#### AdvantageOutputs + +```python +@dataclass +class AdvantageOutputs: + advantages: Float[Tensor, "num_examples rollouts_per_example"] +``` + +### Example: Normalized Advantage + +```python +import torch +from prime_rl.orchestrator.advantage import AdvantageInputs, AdvantageOutputs + +def normalized_advantage(inputs: AdvantageInputs, eps: float = 1e-8) -> AdvantageOutputs: + """Normalize advantages to zero mean and unit variance per example.""" + rewards = torch.tensor([[r["reward"] for r in group] for group in inputs.rollouts]) + mean = rewards.mean(dim=1, keepdim=True) + std = rewards.std(dim=1, keepdim=True) + advantages = (rewards - mean) / (std + eps) + return AdvantageOutputs(advantages=advantages) +``` + +### Configuration + +```toml +[advantage] +type = "custom" +import_path = "my_module.normalized_advantage" +kwargs = { eps = 1e-8 } +``` + +--- + +## Default Implementations + +If no custom function is specified: + +- **Loss**: Uses `default_loss_fn` (masked importance sampling with KL against the inference policy, and optional masking strategies) +- **Advantage**: Uses `default_advantage_fn` (reward minus per-example baseline, a.k.a. DR-GRPO without std normalization) + +See `LossConfig` and `AdvantageConfig` for available parameters. + +## Tips + +- Your functions receive structured inputs via dataclasses with jaxtyping annotations +- Return metrics as scalars or 1D tensors - they'll be aggregated automatically +- Use the `loss_mask` / tensor shapes to handle variable-length sequences +- Test your custom functions with the provided test patterns before training diff --git a/docs/configs.md b/docs/configs.md index 06e8a018d20..5af0d996b02 100644 --- a/docs/configs.md +++ b/docs/configs.md @@ -30,4 +30,4 @@ name = "Qwen/Qwen-14B" PRIME_MODEL__NAME=Qwen/Qwen3-4B uv run ... @ qwen8b.toml @ qwen14b.toml --model.name Qwen/Qwen3-32B ``` -In this example, the CLI argument `--model.name Qwen/Qwen3-32B` will take precendence and the script will use `Qwen/Qwen3-32B` as the model name. If the CLI argument wasn't set, then the second config file would take precedence and the script would use `Qwen/Qwen-14B` as the model name. If the second config file wasn't set, then the first config file would take precedence and the script would use `Qwen/Qwen3-8B` as the model name. Finally, if the first config file wasn't set, then the environment variable would take precedence and the script would use `Qwen/Qwen-4B` as the model name. If the environment variable wasn't set, then the default value would be used and the script would use `Qwen/Qwen3-0.6B` as the model name. \ No newline at end of file +In this example, the CLI argument `--model.name Qwen/Qwen3-32B` will take precedence and the script will use `Qwen/Qwen3-32B` as the model name. If the CLI argument wasn't set, then the second config file would take precedence and the script would use `Qwen/Qwen-14B` as the model name. If the second config file wasn't set, then the first config file would take precedence and the script would use `Qwen/Qwen3-8B` as the model name. Finally, if the first config file wasn't set, then the environment variable would take precedence and the script would use `Qwen/Qwen-4B` as the model name. If the environment variable wasn't set, then the default value would be used and the script would use `Qwen/Qwen3-0.6B` as the model name. diff --git a/docs/deployment.md b/docs/deployment.md index 7f889d35c90..38d911b122c 100644 --- a/docs/deployment.md +++ b/docs/deployment.md @@ -83,11 +83,13 @@ uv run torchrun \ ### SLURM -TBD. +See the dedicated [SLURM guide](slurm.md). ## Inference -We rely on vLLMs multi-node deployment primitives and load balancing for multi-node deployments. Currently, vLLM supports multi-node data parallel deployment ([docs](https://docs.vllm.ai/en/v0.10.0/serving/data_parallel_deployment.html)). +For SLURM-based inference deployment, see the [SLURM guide](slurm.md#inference-examples). Each node runs an independent vLLM replica — no manual coordination needed. + +For manual multi-node deployment without SLURM, we rely on vLLM's multi-node data parallel deployment primitives ([docs](https://docs.vllm.ai/en/v0.10.0/serving/data_parallel_deployment.html)). First, decide which node will be your head node and find a reachable private IP address for it. If your nodes are not colocated, you will likely need to setup VPN (e.g. [Tailscale](https://tailscale.com)) for the nodes to reach each other. @@ -185,7 +187,7 @@ uv run rl \ --inference @ path/to/infer.toml \ ``` -You can configure to GPU IDs to use for the inference server and the trainer. For example, to run the inference server on GPUs IDs 0-5 with data parallelism and the trainer on GPUs IDs 6-7 +You can configure the GPU IDs to use for the inference server and the trainer. For example, to run the inference server on GPUs IDs 0-5 with data parallelism and the trainer on GPUs IDs 6-7 ```bash uv run rl \ @@ -237,9 +239,9 @@ uv run rl \ ### Multi-Node Training -> We currently require shared file system for multi-node RL training. +> We currently require a shared file system for multi-node RL training. -To faciliate multi-node RL training, ensure that all nodes have access to a shared file system and that the node that will run the inference server is reachable from the orchestrator via a private or public IP address. Then, set the following environment variables on all nodes: +To facilitate multi-node RL training, ensure that all nodes have access to a shared file system and that the node that will run the inference server is reachable from the orchestrator via a private or public IP address. Then, set the following environment variables on all nodes: ```bash # On all nodes @@ -281,7 +283,7 @@ Of course, you can further scale up the number of nodes used by the trainer and ### SLURM -TBD. +See the dedicated [SLURM guide](slurm.md). ## Kubernetes @@ -294,4 +296,4 @@ See the dedicated [Kubernetes guide](kubernetes.md) for complete documentation i - Component architecture - Scaling and distributed training - Configuration options -- Troubleshooting \ No newline at end of file +- Troubleshooting diff --git a/docs/disaggregated-inference.md b/docs/disaggregated-inference.md new file mode 100644 index 00000000000..65f5dacf840 --- /dev/null +++ b/docs/disaggregated-inference.md @@ -0,0 +1,91 @@ +# Disaggregated Prefill/Decode Inference + +Run MoE models with separate prefill and decode node groups for higher throughput. + +## Quick Start + +See [`configs/glm5_disagg_inference/inference.toml`](../configs/glm5_disagg_inference/inference.toml) for an example config. + +```bash +uv run inference @ configs/glm5_disagg_inference/inference.toml --output-dir /data/$USER/outputs +``` + +## Prefill/Decode Ratio + +| Workload | Recommended ratio (P:D) | Why | +|---|---|---| +| Agentic (SWE, Lean) | **3:1** | Long growing contexts → prefill-heavy | +| Non-agentic (math, chat) | **1:2** | Short prompts, long generations → decode-heavy | + +Monitor live queue depths: +```bash +curl -s http://:8100/metrics | grep num_requests_waiting +curl -s http://:8200/metrics | grep num_requests_waiting +``` + +If prefill has queued requests and decode has zero, add more prefill nodes (and vice versa). + +For historical averages (cumulative over the entire run), query the histogram metrics: +```bash +# Average queue time per request (seconds) +curl -s http://:/metrics | awk ' + /request_queue_time_seconds_sum\{/ { sum += $2 } + /request_queue_time_seconds_count\{/ { count += $2 } + END { if (count > 0) printf "avg queue: %.2fs (%d requests)\n", sum/count, count } +' + +# Average prefill/decode compute time +curl -s http://:/metrics | awk ' + /request_prefill_time_seconds_sum\{/ { ps += $2 } + /request_prefill_time_seconds_count\{/ { pc += $2 } + /request_decode_time_seconds_sum\{/ { ds += $2 } + /request_decode_time_seconds_count\{/ { dc += $2 } + END { + if (pc > 0) printf "avg prefill: %.2fs\n", ps/pc + if (dc > 0) printf "avg decode: %.2fs\n", ds/dc + } +' +``` + +Other useful metrics on the `/metrics` endpoint: +- `vllm:e2e_request_latency_seconds` — end-to-end latency +- `vllm:kv_cache_usage_perc` — KV cache memory pressure +- `vllm:nixl_xfer_time_seconds` — NIXL KV transfer duration +- `vllm:nixl_bytes_transferred` — bytes per KV transfer + +## UCX 1.19 + +NVSHMEM requires UCX >= 1.19 for multi-GPU CUDA support. Most clusters ship UCX 1.17 (via HPC-X), which causes `cuStreamCreate: invalid device context` errors during DeepEP internode dispatch. + +**Check your version:** +```bash +/opt/hpcx/ucx/bin/ucx_info -v | head -1 +# If < 1.19, you need to build from source +``` + +**Build UCX 1.19 (run once on a GPU node):** +```bash +salloc -N 1 --gres=gpu:1 bash -c 'bash scripts/install_nixl_from_source.sh' +``` + +This installs UCX 1.19 to `prime-rl/third_party/ucx/`. The sbatch template automatically adds it to `LD_LIBRARY_PATH`, overriding the system version. + +## Troubleshooting + +### `DeepEP error: timeout (dispatch CPU)` +NVSHMEM internode communication failing. Check: +1. UCX version >= 1.19? (`third_party/ucx/bin/ucx_info -v`) +2. NVSHMEM libs reachable at `/tmp/deepep_build/nvshmem/lib/`? If not: + ```bash + ssh 'mkdir -p /tmp/deepep_build/nvshmem && \ + ln -sfn /lib/python3.12/site-packages/nvidia/nvshmem/lib \ + /tmp/deepep_build/nvshmem/lib' + ``` +3. IBGDA driver enabled? `ssh 'cat /proc/driver/nvidia/params | grep EnableStreamMemOPs'` should show `1`. + +### Router healthy but requests hang +NIXL side channel not running on prefill. Check: +```bash +ssh 'ss -tlnp sport ge :5600 sport le :5608 | grep -c LISTEN' +# Should show 8 (one per DP rank). If 0, check logs for UCX/NVSHMEM errors. +``` diff --git a/docs/entrypoints.md b/docs/entrypoints.md index 0e1fb00b33f..d79a82faed9 100644 --- a/docs/entrypoints.md +++ b/docs/entrypoints.md @@ -8,7 +8,7 @@ The main usecase of PRIME-RL is RL training. Three main abstractions facilitate ### Orchestrator -The orchestrator is a lightweight CPU process that handles the core data and scheduling logic, serving as an intermediary between the trainer and inference service with bidirectional relays. In one direction, it collects rollouts from the inference server, assembles them into packed batches, and dispatches them to the trainer; in the other direction, it relays updated model weights from the trainer to the inference service. The orchestrator utilizes `verifiers` environments to abstract multi-turn rollout generation and scoring, leveraging async OpenAI-compatible inference clients. +The orchestrator is a lightweight CPU process that handles the core data and scheduling logic, serving as an intermediary between the trainer and inference service with bidirectional relays. In one direction, it collects rollouts from the inference server, assembles them into packed batches, and dispatches them to the trainer; in the other direction, it relays updated model weights from the trainer to the inference service. The orchestrator utilizes `verifiers` environments to abstract multi-turn rollout generation and scoring. Each training and evaluation environment is exposed as a `vf.EnvServer` as a sidecar to the orchestrator process (default) or as a standalone process (e.g. used in hosted training to run environments in containers). ### Trainer @@ -16,7 +16,7 @@ The trainer is responsible for producing an updated policy model given rollouts ### Inference -The inference service in its simplest form is a standard OpenAI-compatible server with a vLLM backend. The API specification is extended with two custom endpoints to enable updating the server with the latest policy: `update_weights` is used to reload model weights from a HF-compatible checkpoint on disk, and `reload_weights` is used to reset the weights to the base model in between experiments. Otherwise, we rely on vLLM's optimized kernels, parallelism strategies, and scheduling for fast rollout generation. Given the disaggregated nature of the service architecture, it can be directly extended to include multiple engines with a shared request pool, allowing operation across multiple clusters and straightforward integration of alternative inference engines (e.g. SGLang, Tokasaurus). We also heavily rely on native data parallelism in vLLM (also available in SGLang) for orchestrating the fleet of nodes dedicated to inference. +The inference service in its simplest form is a standard OpenAI-compatible server with a vLLM backend. The API specification is extended with a custom `update_weights` endpoint to reload model weights from a HF-compatible checkpoint on disk. Otherwise, we rely on vLLM's optimized kernels, parallelism strategies, and scheduling for fast rollout generation. Given the disaggregated nature of the service architecture, it can be directly extended to include multiple engines with a shared request pool, allowing operation across multiple clusters and straightforward integration of alternative inference engines (e.g. SGLang, Tokasaurus). We also heavily rely on native data parallelism in vLLM (also available in SGLang) for orchestrating the fleet of nodes dedicated to inference. ### RL @@ -50,7 +50,7 @@ For more details on multi-node deployment options, see the [deployment](deployme We provide a fairly straight-forward SFT trainer which is capable of fine-tuning any conversational model on multi-turn conversation with tool calling. It shares a lot of components with the RL trainer, such as the modeling code, parallelism techniques, checkpoint format, logger, etc. which ensures a seemless post-training workflow. -To start an SFT training, you need to prepare a dataset in [prompt-completion format](https://huggingface.co/docs/trl/en/dataset_formats#prompt-completion) (we do not support any other format). Single-turn fine-tuning should be compatible with the chat templates of most models. However, to properly handle loss masking, we require that the tokenizer's chat template satisfies a prefix property: the tokenization of any conversation prefix must be a prefix of the tokenization of the full conversation. For instance, tokenizing message 1 should yield a token sequence that forms a prefix of tokenizing messages 1 and 2, which in turn should be a prefix of tokenizing messages 1, 2, 3, and so forth. An example of a chat template that *does not* satisfy this property is Qwen3's chat template, as it strips away past think sections. +To start an SFT training, you need to prepare a conversational dataset in either [prompt-completion format](https://huggingface.co/docs/trl/en/dataset_formats#prompt-completion) or raw `messages` format. If `messages` is provided, the trainer interprets the full conversation as a single sample with an empty prompt and applies role-based loss masking across the whole chat. If both `messages` and `prompt` / `completion` are present, `messages` takes precedence. Single-turn fine-tuning should be compatible with the chat templates of most models. However, to properly handle loss masking, we require that the tokenizer's chat template satisfies a prefix property: the tokenization of any conversation prefix must be a prefix of the tokenization of the full conversation. For instance, tokenizing message 1 should yield a token sequence that forms a prefix of tokenizing messages 1 and 2, which in turn should be a prefix of tokenizing messages 1, 2, 3, and so forth. An example of a chat template that *does not* satisfy this property is Qwen3's chat template, as it strips away past think sections. On a single GPU, start the training with the `sft` entrypoint @@ -65,114 +65,3 @@ uv run torchrun --nproc-per-node 8 src/prime_rl/trainer/sft/train.py ... ``` For more details on multi-node deployment options, see the [deployment](deployment.md) documentation and see the [examples](examples) for concrete training configurations. To see all available configuration options, run `uv run sft --help`. - -## Evals - -You can eval any [verifiers](https://github.com/willccbb/verifiers) environment against API models, local models and checkpoints from an SFT or RL training using the `eval` entrypoint. - -> We recommned using the `vf-eval` entrypoint for debugging and developing an environment and the PRIME-RL eval entrypoint for production-grade workloads, as it has more advanced features such as multi-environment evaluation, saving results to the Environments Hub, and scales better to custom multi-node inference deployments. - -By default, the entrypoint uses the OpenAI API with `gpt-4.1-mini` to provide a smooth out-of-the-box experience without requiring a vLLM server. - -To check all available configuration options, run `uv run eval --help`. - -### API Models - -By default, the entrypoint uses the OpenAI API. First, set the API key as an environment variable: - -```bash -export OPENAI_API_KEY=... -``` - -### Single Environment - -```bash -uv run eval @ configs/debug/eval/single_env.toml -``` - -### Multiple Environments - -```bash -uv run eval @ configs/debug/eval/multi_env.toml -``` - -### Local Models - -To use a different API, or custom inference deployment, override the client and model configuration. For example, for a bare vLLM server, run: - -```bash -uv run eval \ - @ configs/debug/eval/local_model.toml \ - --client.base-url http://localhost:8000/v1 \ - --model.name -``` - -### Checkpoints - -To evaluate a training checkpoints, start an inference server with the model being the base model that you started training from and specify the directory containing the weight checkpoints with `--weights-dir`. - -```bash -uv run inference --model.name -``` - -```bash -uv run eval \ - @ configs/debug/eval/single_env.toml \ - --client.base-url http://localhost:8000/v1 \ - --model.name \ - --weights-dir outputs/weights -``` - -By default, this will evaluate the base model and all step checkpoints found in the weights directory. To skip evaling the base model, set `--no-eval-base` and to evaluate only specific steps, set `--steps` as a comma-separated list of integers representing the steps to evaluate. For example, - -## Synthetic Data - -You can generate synthetic data in any [verifiers](https://github.com/willccbb/verifiers) environment using the `synthesize` entrypoint. The entrypoint shares configuration with evals (`ModelConfig`, `ClientConfig`, `SamplingConfig`, and `EnvConfig`), making it easy to generate data using the same models and sampling settings as your evaluations. - -The synthesize entrypoint supports single-turn, multi-turn, and tool-calling environments, and can generate data across multiple environments in parallel. It automatically parses reasoning content from model responses (configurable via `--reasoning-field`, defaults to `reasoning_content`), and saves data in append mode using the same schema as verifiers, making it straightforward to convert to SFT datasets. The entrypoint is robust to errors during generation, scoring, or saving—failed groups are simply dropped. Environments specified in the format `{env_org}/{env_id}` are automatically installed, similar to the RL entrypoint. - -By default, the entrypoint uses the OpenAI API with `gpt-4.1-mini` (also the default for evals) to provide a smooth out-of-the-box experience without requiring a vLLM server. Make sure to set your API key as an environment variable - -```bash -export OPENAI_API_KEY=... -``` - -To check all available configuration options, run `uv run synthesize --help`. - -### Single-Turn Environments - -Generate synthetic data in a single-turn environment (e.g., `gsm8k`): - -```bash -uv run synthesize @ configs/debug/synthesize/single_turn.toml -``` - -By default, the entrypoint parses and saves the `reasoning_content` field from raw responses, as returned by vLLM when a reasoning parser is set. This behavior can be configured to handle different APIs: - -```bash -uv run synthesize @ configs/debug/synthesize/single_turn.toml --reasoning-field reasoning -``` - -### Multiple Environments - -Generate synthetic data across multiple environments (e.g., `gsm8k` and `hendrycks-math`): - -```bash -uv run synthesize @ configs/debug/synthesize/multi_env.toml -``` - -### Multi-Turn Environments - -Generate synthetic data in a multi-turn environment (e.g., `alphabet-sort`): - -```bash -uv run synthesize @ configs/debug/synthesize/multi_turn.toml -``` - -### Multi-Turn Environments with Tool Calls - -Generate synthetic data in a multi-turn environment with tool calls (e.g., `wiki-search`): - -```bash -uv run synthesize @ configs/debug/synthesize/multi_turn_tool_call.toml -``` diff --git a/docs/index.md b/docs/index.md index 8d19590974e..aa76871f8c7 100644 --- a/docs/index.md +++ b/docs/index.md @@ -7,7 +7,8 @@ This directory maintains the documentation for PRIME-RL. It is organized into th - [**Environments**](environments.md) - Installing and using verifiers environments from the Environments Hub - [**Async Training**](async.md) - Understanding asynchronous off-policy training and step semantics - [**Logging**](logging.md) - Logging with loguru, torchrun, and Weights & Biases -- [**Runs**](runs.md) - Multi-run training with the Runs object for concurrent LoRA adapters +- [**Platform Monitoring**](platform-monitoring.md) - Register runs on the Prime Intellect platform and stream training metrics +- [**MultiRunManager**](multi_run_manager.md) - Multi-run training with the MultiRunManager object for concurrent LoRA adapters - [**Checkpointing**](checkpointing.md) - Saving and resuming training from checkpoints - [**Benchmarking**](benchmarking.md) - Performance benchmarking and throughput measurement - [**Deployment**](deployment.md) - Training deployment on single-GPU, multi-GPU, and multi-node clusters diff --git a/docs/logging.md b/docs/logging.md index ee59b7ac035..cbd7e881f24 100644 --- a/docs/logging.md +++ b/docs/logging.md @@ -1,14 +1,6 @@ # Logging -prime-rl uses [loguru](https://loguru.readthedocs.io/en/stable/) for logging with a global logger pattern. Logs are written to both console and files under `{output_dir}/logs/`. For RL training, we recommend streaming file logs into tmux panes (as set up by `tmux.sh`). - -## tmux helper (`scripts/tmux.sh`) - -`scripts/tmux.sh` sets up a tmux session for RL runs with **three panes (one per subprocess)**: - -- **Trainer**: run `uv run rl ...` here -- **Orchestrator**: follows `{output_dir}/logs/orchestrator.stdout` -- **Inference**: follows `{output_dir}/logs/inference.stdout` +prime-rl uses [loguru](https://loguru.readthedocs.io/en/stable/) for logging with a global logger pattern. All logs are captured at the deployment level (stdout/stderr redirection for local, `tee` for SLURM) under `{output_dir}/logs/`. For RL training, we recommend streaming logs into tmux panes (as set up by `tmux.sh`). ## Logger Architecture @@ -20,7 +12,7 @@ We use a **singleton pattern** with a module-level global logger instance (`_LOG from prime_rl.utils.logger import setup_logger, get_logger # At entrypoint - call ONCE -logger = setup_logger("info", log_file=Path("output/logs/my.log")) +logger = setup_logger("info") # Anywhere else in codebase logger = get_logger() @@ -29,79 +21,66 @@ logger.info("Hello world") **How it works:** -1. **`setup_logger(log_level, log_file)`** - Initializes the global logger exactly once: - - Creates an isolated loguru `Logger` instance (not the default `loguru.logger`) to prevent third-party code from hijacking our logs - - Adds a stdout handler with colorized output - - Optionally adds a file handler (deletes existing file first) - - Raises `RuntimeError` if called twice +1. **`get_logger()`** - Returns the global logger instance. Always works — if `setup_logger` hasn't been called yet, it initializes a default logger automatically. Safe to call from any module at any time. -2. **`get_logger()`** - Returns the global logger instance: - - Raises `RuntimeError` if `setup_logger` hasn't been called yet - - Safe to call from any module after initialization +2. **`setup_logger(log_level)`** - Configures (or reconfigures) the global logger: + - Creates an isolated loguru `Logger` instance (not the default `loguru.logger`) to prevent third-party code from hijacking our logs + - Adds a stdout handler with colorized output (or JSON output if `json_logging=True`) + - Can be called multiple times — cleans up the previous logger before creating a new one 3. **`reset_logger()`** - Resets the global logger to `None`: - Used in subprocesses that inherit parent state (e.g., env workers) - Used in tests between test cases -## RL Log File Structure - -For RL training, logs are organized by component: +## Log File Structure -| Component | Log Path | Description | -|-----------|----------|-------------| -| **RL (parent)** | `logs/rl.log` | Main process that spawns subprocesses | -| **Inference** | `logs/inference.stdout` | vLLM inference server stdout/stderr | -| **Orchestrator** | `logs/orchestrator.log` | Rollout generation, buffer, scheduling | -| **Trainer** | `logs/trainer/rank_{N}.log` | Training process (one file per GPU rank) | -| **Env Workers** | `logs/env_workers/{env_name}/worker_{N}.log` | Per-environment worker logs (optional) | +Logs are captured at the deployment level — the entrypoint redirects subprocess stdout/stderr to files (local) or `tee` captures them (SLURM). The structure is consistent across deployment types: `logs/trainer.log` and `logs/inference.log` always exist, regardless of whether the run is local or multi-node SLURM. -## Per-Environment Worker Logging +### Local (single node) -Environment workers run in **separate subprocesses** to isolate event loop lag. Worker logging is controlled at the orchestrator level via `orchestrator.log`: - -```toml -[orchestrator.log] -level = "debug" # Log level for prime-rl logger -vf_level = "info" # Log level for verifiers library -env_worker_logs = true # Enable file logging for env workers ``` - -When `env_worker_logs = true`, logs are written to: -``` -output_dir/ -└── logs/ - └── env_workers/ - ├── {env_name_1}.log - ├── {env_name_2}.log +{output_dir}/logs/ +├── trainer.log # trainer stdout (rank 0 only) +├── orchestrator.log # orchestrator stdout +├── inference.log # vLLM inference server stdout +├── trainer/ +│ └── torchrun/ # per-rank stdout/stderr (all ranks) +└── envs/ + ├── train/{env_name}/ + │ ├── env_server.log + │ └── env_worker_{id}.log + └── eval/{env_name}/ └── ... ``` -All workers for an environment share the same log file. - -Set `env_worker_logs = false` to disable worker file logging (workers inherit parent process logging). - -## Torchrun - -For multi-node training with `torchrun`, all ranks log simultaneously. To filter to master rank only: +### SLURM multi-node -```bash -uv run torchrun \ - --local-ranks-filter 0 \ - --nproc-per-node 8 \ - ... +``` +{output_dir}/logs/ +├── trainer.log -> trainer/node_0.log (symlink) +├── inference.log -> inference/node_0.log (symlink) +├── orchestrator.log # orchestrator stdout +├── trainer/ +│ ├── node_0.log # per-node trainer output (rank 0 only) +│ ├── node_1.log +│ └── torchrun/ # per-rank stdout/stderr (all ranks) +├── inference/ +│ ├── node_0.log # per-node inference output +│ ├── node_1.log +│ └── router_0.log # vllm-router per replica +└── envs/ + └── ... ``` -You can also use torchrun's native log redirection: +Environment logs live under `logs/envs/train/{env_name}/` and `logs/envs/eval/{env_name}/`. Env log verbosity is controlled by `orchestrator.log.vf_level`. -```bash -uv run torchrun \ - --local-ranks-filter 0 \ - --nproc-per-node 8 \ - --log-dir outputs/torchrun \ - --redirects 3 \ - --tee 3 \ - ... -``` +Only rank 0 output is shown in `trainer.log`. Per-rank logs from all ranks are available under `logs/trainer/torchrun/{rdzv_id}/attempt_0/{rank}/{stdout,stderr}.log`, written by torchrun's `--log-dir`. + +## tmux helper (`scripts/tmux.sh`) -This writes to `outputs/torchrun/{rdzv_id}/attempt_0/{rank}/{stdout,stderr}.log`. +`scripts/tmux.sh` sets up a tmux session for RL runs with **four panes**: +- **Trainer**: follows `{output_dir}/logs/trainer.log` +- **Orchestrator**: follows `{output_dir}/logs/orchestrator.log` +- **Envs**: follows `{output_dir}/logs/envs/*/*/*.log` +- **Inference**: follows `{output_dir}/logs/inference.log` diff --git a/docs/memory_usage.md b/docs/memory_usage.md new file mode 100644 index 00000000000..b36c1172543 --- /dev/null +++ b/docs/memory_usage.md @@ -0,0 +1,132 @@ +# Reducing memory usage + +While most of our parallelism techniques in prime-rl are designed to scale training up (FSDP, EP, CP, ...), we also provide many tools to scale training down that allow training large MoE models on a limited amount of GPUs. + +These techniques target the trainer part of prime-rl. + + +## TLDR: config to use for maximum memory usage reduction with correct throughput + +```toml +[trainer.model] +impl = "custom" +attn = "flash_attention_2" +fused_lm_head_token_chunk_size = 1024 +ep = 8 +cp = 2 +optim_cpu_offload = true + +[trainer.model.compile] + +[trainer.model.ac] +freq = 1 + +[trainer.model.ac_offloading] +max_inflight_activations = 1 +``` + +## Activation checkpointing + +Activation checkpointing discards intermediate activations during the forward pass and recomputes them during the backward pass, trading compute for memory. + +To enable it, use: + +```toml +[trainer.model.ac] +freq = 1 +``` + +`freq` controls how often layers are checkpointed: every `freq` layers. Lower values yield lower memory usage (e.g. `freq = 1` checkpoints every layer). + +## Activation offloading + +Activation offloading offloads the activations to CPU to reduce the memory usage of the trainer. It can be used in combination with activation checkpointing. + +To enable it, use: + +```toml +[trainer.model.ac] +freq = 1 + +[trainer.model.ac_offloading] +max_inflight_activations = 5 +``` + +## Chunk loss + +Chunk loss splits the loss computation into smaller chunks to reduce the memory usage of the trainer. + +To enable it, use: + +```toml +[trainer.model] +fused_lm_head_token_chunk_size = auto +``` + + +## Expert parallelism + +While expert parallelism splits the weights of the experts across all GPUs like FSDP, using EP still reduces memory usage by reducing the communication size and therefore the FSDP buffer. + +EP is only available for models with MoE layers using the custom model implementation. + +``` +[trainer.model] +impl = "custom" +ep = 8 +``` + +## Context parallelism + +Context parallelism splits the context into smaller chunks to reduce the memory usage of the activations. We don't advise using CP across multiple nodes (i.e., increasing CP beyond 8). + +CP is only available for certain models and only with the custom model implementation. + +``` +[trainer.model] +impl = "custom" +cp = 2 +``` + +We recommend CP 2 or CP 4 for most 128K sequence length training runs. Can be pushed to 8. + + +## torch compile + +Enabling torch.compile can reduce the memory usage for certain model architectures, especially MoE with the custom model implementation. + +``` +[trainer.model.compile] +``` + +## CPU Optimizer offloading + +Offloading the optimizer states to CPU can reduce the memory usage of the trainer significantly, especially at low GPU counts where the optimizer states take a lot of memory as they won't be sharded enough. + +In RL, in contrast with pretraining, we end up with many gradient accumulation steps, so the cost of offloading the optimizer states is not as high as in pretraining, and indeed barely noticeable. + +``` +[trainer.optim] +optim_cpu_offload = true +``` + +## :warning: FSDP CPU offloading + +FSDP CPU offloading offloads the parameters, gradients, and optimizer states to CPU to reduce the memory usage of the trainer. + +This will make training significantly slower and is not recommended most of the time. + +``` +[trainer.model] +fsdp_cpu_offload = true +``` + +## :warning: Lora training + +LoRA training significantly reduces the memory usage of the trainer at the cost of smaller gradient updates. + +``` +[trainer.model.lora] +rank = 8 +``` + diff --git a/docs/mint.json b/docs/mint.json index da1104d9ec6..216fbe4fa48 100644 --- a/docs/mint.json +++ b/docs/mint.json @@ -10,7 +10,7 @@ "environments", "async", "logging", - "runs", + "multi_run_manager", "checkpointing", "benchmarking", "deployment", diff --git a/docs/multi_run_manager.md b/docs/multi_run_manager.md new file mode 100644 index 00000000000..bef6a6f566b --- /dev/null +++ b/docs/multi_run_manager.md @@ -0,0 +1,244 @@ +# MultiRunManager + +The `MultiRunManager` object is a global singleton that manages the parameters and components for multiple concurrent training runs within a single trainer process. +This allows multiple orchestrator deployments to share the same trainer. + +When `max_concurrent_runs > 1`, the trainer can train multiple runs in parallel. Each run: + +- Has its own LoRA adapter parameters +- Has its own optimizer and scheduler +- Saves its own checkpoints +- Tracks its own training progress (step, tokens, samples) +- Loads its own orchestrator configuration + +The `MultiRunManager` object provides: + +- **Bidirectional mapping** between run IDs (e.g., `run_abc123`) and run indices (0, 1, 2, ...) +- **Progress tracking** per run (step count, total tokens, total samples) +- **Configuration management** for orchestrator configs +- **Distributed synchronization** across ranks via the PyTorch distributed store +- **LoRA module registration** for multi-adapter parameter management +- **Creation hooks** for initializing per-run resources (optimizers, schedulers) +- **Run eviction** for removing runs that are misbehaving + +## **Initialization and run discovery** + +The `MultiRunManager` singleton is set up at the start of training: + +```python +from prime_rl.trainer.runs import setup_multi_run_manager, get_multi_run_manager + +# Initialize with output directory and max concurrent runs +setup_multi_run_manager(output_dir=Path("outputs/my-experiment"), max_runs=4) + +# Get the singleton instance anywhere in the codebase +multi_run_manager = get_multi_run_manager() +``` + +Each run's directory follows this structure: + +``` +{output_dir}/ +├── run_abc123/ +│ ├── control/ +│ │ ├── orch.toml # Orchestrator configuration +│ │ ├── config_validation_error.txt # Config validation errors (if any) +│ │ └── evicted.txt # Eviction reason (if evicted) +│ ├── checkpoints/ +│ │ └── step_100/ # Orchestrator checkpoints +│ ├── rollouts/ +│ │ └── step_100/ # Rollouts +│ └── broadcast/ +│ └── step_100/ # Broadcasted weights for inference +├── run_def456/ +│ └── ... +└── ... + +``` + +Runs are discovered by scanning the output directory for the pattern `run_*`. Each run must contain a valid orchestrator config at `{run_dir}/control/orch.toml` before they are added to the active runs otherwise they are ignored. When the maximum number of runs is reached, new `run_*` directories will not be picked up until old ones are deleted. + +```python +# Master rank scans for new/deleted runs +multi_run_manager.discover_runs() + +# All ranks synchronize state (must be called after discover_runs) +multi_run_manager.synchronize_state() +``` + +The `discover_runs()` method (master only): + +1. Scans the output directory for `run_*` directories +2. Filters out evicted runs (those with `control/evicted.txt`) +3. Detects new runs and deleted runs +4. Calls `forgotten_hook` for deleted runs (master only) +5. Loads and validates the orchestrator config for each new run +6. Updates internal mappings and data structures +7. Calls `discovered_hook` for new runs (master only) + +The `synchronize_state()` method (all ranks): + +1. Master broadcasts run state to all ranks via the distributed store +2. Non-master ranks catch up by calling internal `_delete_run_data` / `_create_run_data` +3. All ranks execute `deletion_hook` for deleted runs +4. All ranks execute `creation_hook` for new runs (e.g., optimizer setup, LoRA parameter reset) + +## Run Eviction + +The master proc on the trainer can evict a run using the `evict_run(idx: int, reason: str)` method. +This is useful when the trainer detects an issue with a run that requires it to be stopped (e.g., invalid data, resource constraints, or policy violations). + +```python +# Evict a run by its index (master only) +multi_run_manager.evict_run(idx=0, reason="Run exceeded memory limits") +``` + +The `evict_run()` method (master only): + +1. Writes the eviction reason to `{run_dir}/control/evicted.txt` +2. Logs a warning with the eviction details +3. The run is **not** immediately removed from the manager's data structures + +The eviction takes effect through two mechanisms: + +**On the trainer side:** +- The next `discover_runs()` call will filter out the evicted run (it checks for `evicted.txt`) +- The run will then be treated as deleted, triggering forgotten/deletion hooks +- The run index is returned to the unused pool + +**On the orchestrator side:** +- The orchestrator checks for `evicted.txt` at the start of each iteration in its main loop +- If found, it raises a `RuntimeError` with the eviction reason, causing the orchestrator to exit +- This surfaces the eviction reason to the user +- The orchestrator also self-evicts by writing `evicted.txt` if a training batch has no learning signal (all rollouts filtered out) on `MAX_EMPTY_BATCH_ATTEMPTS` (3) consecutive attempts + +## LoRA Module Registration + +LoRA modules register themselves with `MultiRunManager` for parameter management: + +```python +# In apply_lora_to_model() +lora_module = MultiLoRALinear( + base_layer=base_module, + rank=config.rank, + n_adapters=get_multi_run_manager().max_runs, + ... +) +lora_module.register_with_runs(get_multi_run_manager(), module_name) + +``` + +The `MultiRunManager` object then exposes: + +```python +# Get parameters for a specific run (used by optimizer creation) +multi_run_manager.get_named_parameters_for_run(idx) + +# Get state dict for a specific run (used by weight broadcast) +multi_run_manager.get_state_dict_for_run(idx) + +# Reset parameters for a new run +multi_run_manager.reset_run_parameters(idx) + +``` + +## Hooks + +The `MultiRunManager` object supports several types of hooks for different lifecycle events. +Deletion hooks are always called before creation hooks. + +```mermaid +flowchart TD + subgraph master["Rank 0 (Master)"] + discover["discover_runs()"] + forgotten["forgotten_hooks"] + validation["config_validation_hooks"] + discovered["discovered_hooks"] + + discover --> forgotten + forgotten --> validation + validation --> discovered + discovered --> discover + end + + subgraph rank1["Rank 1"] + wait1["waiting..."] + end + + subgraph rankN["Rank N"] + waitN["waiting..."] + end + + discovered --> barrier + wait1 --> barrier + waitN --> barrier + + barrier[["synchronize_state()"]] + + barrier --> deletion["deletion_hooks"] + deletion --> creation["creation_hooks"] + + style barrier fill:#fff9c4 +``` + +### Hook Registration + +```python +# These hooks are only called on the master as only master uses `discover_runs()` +# These hooks are thus only relevant to master only components (packer) +multi_run_manager.register_discovered_hook(callback) +multi_run_manager.register_forgotten_hook(callback) + +# These hooks are executed by all ranks in the order they were added during `synchronize_state()` +# This ensures DTensor creations and other distributed operations happen together +# Calling torch.dist.barrier() in a hook here should work +multi_run_manager.register_creation_hook(callback) +multi_run_manager.register_deletion_hook(callback) + +# These hooks validate the orchestrator config when runs are discovered: +multi_run_manager.register_config_validation_hook(callback) +``` + +The callback signatures: + +```python +def discovered_callback(idx: int, run_id: str, config: OrchestratorConfig) -> None: + """Called when a new run is discovered (master only). + + Args: + idx: The run's index (0 to max_runs-1) + run_id: The run's ID (e.g., "run_abc123") + config: The orchestrator config for the run + """ + # Example: Set the scaling factor for the run + multi_run_manager.scaling_factors[idx] = config.model.lora.alpha / config.model.lora.rank + +def forgotten_callback(idx: int, run_id: str) -> None: + """Called when a run is forgotten/removed (master only). + + Args: + idx: The run's index (0 to max_runs-1) + run_id: The run's ID (e.g., "run_abc123") + """ + pass + +def callback(idx: int, run_id: str) -> None: + """Called when a run is created/deleted. + + Args: + idx: The run's index (0 to max_runs-1) + run_id: The run's ID (e.g., "run_abc123") + """ + pass + +def config_validation_callback(config: OrchestratorConfig) -> tuple[bool, str]: + """Validate an orchestrator config. + + Args: + config: The orchestrator config to validate + + Returns: + (is_valid, error_message): If invalid, error_message is written to config dir + """ + return True, "" +``` diff --git a/docs/multimodal.md b/docs/multimodal.md new file mode 100644 index 00000000000..2d12cd324b7 --- /dev/null +++ b/docs/multimodal.md @@ -0,0 +1,62 @@ +# Multimodal (VLM) Support + +Prime-RL supports training vision-language models (VLMs) like Qwen3-VL. + +## VLM Configuration + +### Supported Models + +The built-in registry supports these model families out of the box: + +| Model Family | model_type | Vision Encoder | Language Model | +|-------------|------------|---------------|----------------| +| Qwen3-VL | `qwen3_vl` | `model.visual` | `model.language_model` | +| Qwen3.5 | `qwen3_5` | `model.visual` | `model.language_model` | +| Qwen3.5-MoE | `qwen3_5_moe` | `model.visual` | `model.language_model` | + +Enable VLM mode by adding a `[model.vlm]` section. Both fields are required — they tell prime-rl where the vision encoder and language model live on the model object: + +```toml +[model] +name = "Qwen/Qwen3-VL-4B-Instruct" + +[model.vlm] +vision_encoder_attr = "model.visual" +language_model_attr = "model.language_model" +``` + +For the registered models in the table above, use the attrs shown there. For custom VLMs, check your model's structure with `model.named_children()`. + +Both fields are dotted attribute paths resolved on the loaded model. A bad path raises a `ValueError` immediately — there are no silent fallbacks. + +The weight key prefix for NCCL broadcasting is derived automatically as `{language_model_attr}.layers.`. + +To add permanent support for a new model family, add an entry to `VLM_REGISTRY` in `src/prime_rl/utils/vlm.py`. + +## Current Limitations + +- **Vision encoder is frozen by default**: The vision encoder is frozen during training by default. Set `freeze_vision_encoder = false` in `[model.vlm]` to make it trainable. When unfrozen, the vision encoder is FSDP-sharded per-block for proper gradient flow. Note: this has no effect when using LoRA. + +- **No multimodal-safe truncation**: Token sequences are truncated to `seq_len`, but `pixel_values` and `image_grid_thw` are passed through unchanged. If a multimodal sample exceeds `seq_len`, image tokens can be dropped while image tensors still describe the full set of images. Ensure `seq_len` covers your longest VLM samples. + +- **Optimization dtype must be bfloat16**: Set `optimization_dtype = "bfloat16"` and `reduce_dtype = "bfloat16"` in your trainer config. + +- **Higher KL mismatch with multi-image inputs**: VLM training exhibits higher KL mismatch compared to text-only, especially with multiple images. + +- **Images are not logged**: The images the VLM sees during training are not logged to monitors. + +## How Multi-Turn VLM RL Training Works + +VLM training uses the same `interleave_rollout` path as text-only models. Multi-turn trajectory steps are merged into a single training sample wherever the extension property holds. + +Images are handled via a `VLMImageCache` built once per batch: + +1. **Extract**: Base64 images are decoded from trajectory step prompts into PIL images. +2. **Preprocess**: Images are processed through the HuggingFace image processor, producing `pixel_values` and `image_grid_thw`. +3. **Attach**: Each training sample receives the cumulative `pixel_values` up to its last merged step. + +Each multimodal sample becomes its own micro-batch during training (no packing) since image tensor sizes vary. + +## vLLM Configuration + +`VLLM_WORKER_MULTIPROC_METHOD=spawn` is required for VLM inference. This is set automatically when using `uv run rl @ ...`, but if you start the vLLM server yourself, make sure this environment variable is set. diff --git a/docs/on_policy_distillation.md b/docs/on_policy_distillation.md index d660579ccd3..093ccfec710 100644 --- a/docs/on_policy_distillation.md +++ b/docs/on_policy_distillation.md @@ -6,10 +6,11 @@ For more details, see [On-Policy Distillation](https://thinkingmachines.ai/blog/ ## Quick Start -Add `teacher_gpu_ids` and set `teacher_tau > 0`: +Add `num_teacher_gpus` to `[deployment]` and set `teacher_tau > 0`: ```toml -teacher_gpu_ids = [2, 3] +[deployment] +num_teacher_gpus = 2 [trainer.loss] teacher_tau = 0.5 @@ -18,7 +19,8 @@ teacher_tau = 0.5 This automatically starts a teacher inference server using the same model as inference. To use a different teacher model: ```toml -teacher_gpu_ids = [2, 3] +[deployment] +num_teacher_gpus = 2 [teacher_inference.model] name = "deepseek-ai/DeepSeek-R1-Distill-Qwen-7B" @@ -47,26 +49,68 @@ name = "deepseek-ai/DeepSeek-R1-Distill-Qwen-7B" For agentic environments where verification is expensive (code execution, tool use, multi-turn interactions), you can skip verification entirely and use only the teacher signal: ```toml -teacher_gpu_ids = [2, 3] +[deployment] +num_teacher_gpus = 2 [trainer.loss] teacher_tau = 1.0 adv_tau = 0.0 # Disable reward-based learning -[orchestrator.buffer] -skip_verification = true # Skip expensive verification +[orchestrator.verification] +enabled = false # Skip expensive verification ``` This runs pure on-policy distillation: the student learns to match the teacher without needing any reward signal. +## SFT Distillation ("Hard Distillation") From Teacher Rollouts + +Use this mode when you want to train from teacher-generated completions directly (hard distillation), without teacher token-level logprobs. + +```toml +[trainer.loss] +type = "sft" + +[orchestrator] +use_token_client = false + +[orchestrator.teacher_rollout_model.client] +base_url = ["https://your-openai-compatible-endpoint/v1"] +skip_model_check = true + +[orchestrator.teacher_rollout_model.model] +name = "teacher-model-name" +``` + +In this mode: +- Rollouts are generated from `orchestrator.teacher_rollout_model` +- The orchestrator uses text-level reconstruction with the student tokenizer +- The RL trainer optimizes masked NLL (`trainer.loss.type = "sft"`) +- Omit `[inference]` (no local inference server required) + +### Image Input (VLM) Support + +Yes, image input is supported in SFT/hard-distillation mode when the student model is multimodal (VLM). + +- Prompts can include OpenAI-style image items in `message.content`, e.g. `{"type": "image_url", "image_url": {"url": "data:image/..."}}` +- The orchestrator extracts and preprocesses images from trajectory prompts and attaches `pixel_values`/`image_grid_thw` to training samples +- No teacher token IDs/logprobs are required; reconstruction still happens from messages + +Notes: +- This path currently expects `data:image/...` payloads in message content +- The teacher rollout endpoint still needs to be able to handle the same multimodal prompts during generation + +Reference configs: +- `configs/alphabet_sort/sft_distill_hard_qwen4b_lora_prime_teacher.toml` +- `examples/alphabet_sort/sft_distill_hard.toml` + ## Parameters | Parameter | Default | Description | |-----------|---------|-------------| -| `teacher_gpu_ids` | `None` | GPUs for teacher server. Auto-starts server when set. | +| `deployment.num_teacher_gpus` | `None` | Number of GPUs for teacher server. Auto-starts server when set. | | `trainer.loss.teacher_tau` | `0.0` | Distillation strength. Set `> 0` to enable. | | `trainer.loss.adv_tau` | `1.0` | Weight for RL advantage signal. Set `0` for pure distillation. | -| `orchestrator.buffer.skip_verification` | `false` | Skip verification. Use with `adv_tau = 0`. | +| `orchestrator.verification.enabled` | `true` | Enable/disable verification. Set to `false` for pure distillation with `adv_tau = 0`. | ## Monitoring diff --git a/docs/platform-monitoring.md b/docs/platform-monitoring.md new file mode 100644 index 00000000000..31bcfe312bc --- /dev/null +++ b/docs/platform-monitoring.md @@ -0,0 +1,48 @@ +# Platform Monitoring + +Use `orchestrator.prime_monitor` to register a run on the Prime Intellect platform and stream training metrics, samples, and distributions. + +> **Internal-only for now:** external run registration is currently only enabled for internal / allowlisted teams. + +## Prerequisites + +You need a Prime API key with `rft:write` scope. + +Use the CLI: + +```bash +prime login +``` + +Or set an environment variable directly: + +```bash +export PRIME_API_KEY=pit_... +``` + +## Minimal config + +```toml +[orchestrator.prime_monitor] +run_name = "my-experiment" +``` + +You can also override from the CLI: + +```bash +uv run rl @ config.toml --orchestrator.prime_monitor.run_name "my-experiment" +``` + +## Troubleshooting + +### `API key not found` + +Set the env var from `api_key_var` or run: + +```bash +prime login +``` + +### `External training runs are not enabled for this team` + +Your team is not allowlisted yet. This feature is currently internal-only. diff --git a/docs/runs.md b/docs/runs.md deleted file mode 100644 index 7c94dfaf374..00000000000 --- a/docs/runs.md +++ /dev/null @@ -1,23 +0,0 @@ -# Runs - -The `Runs` object is a singleton that manages multiple concurrent training runs within a single trainer process. It is the central coordination point for multi-run RL training, enabling a single trainer to serve multiple orchestrator experiments simultaneously with separate LoRA adapters, optimizers, and schedulers. - -## Overview - -When `max_concurrent_runs > 1`, the trainer can train multiple runs in parallel. Each run: -- Has its own LoRA adapter weights -- Has its own optimizer and scheduler -- Tracks its own training progress (step, tokens, samples) -- Loads its own orchestrator configuration - -| Responsibility | Description | -|---------------|-------------| -| **Discovery** | Scans for `run_*` directories and loads configs | -| **Mapping** | Provides bidirectional run ID ↔ index mapping | -| **Progress** | Tracks per-run training step, tokens, samples | -| **Synchronization** | Keeps all ranks in sync via distributed store | -| **Hooks** | Enables lazy initialization of per-run resources | -| **LoRA Management** | Registers modules for multi-adapter parameter access | -| **State Access** | Provides per-run parameters and state dicts | - -This design enables efficient multi-tenant training where a single trainer can serve multiple experiments with independent adapter weights, optimizers, and learning rate schedules. diff --git a/docs/slurm.md b/docs/slurm.md new file mode 100644 index 00000000000..7b171d4894b --- /dev/null +++ b/docs/slurm.md @@ -0,0 +1,298 @@ +# SLURM + +The `rl`, `sft`, and `inference` entrypoints all have built-in SLURM support. Adding a `[slurm]` section to your config switches from local execution to SLURM job submission — no separate entrypoint needed. + +## Quick Start + +```bash +# Local run +uv run rl @ examples/reverse_text/rl.toml + +# SLURM run (same entrypoint, just add [slurm] to the config) +uv run rl @ examples/reverse_text/slurm_rl.toml +``` + +The SLURM config is a thin overlay that inherits from a base config and adds `[slurm]` + `[deployment]` sections: + +```toml +# examples/reverse_text/slurm_rl.toml +toml_files = ["rl.toml"] + +output_dir = "outputs/reverse-text-rl" + +[slurm] +job_name = "reverse-text-rl" +``` + +## How it works + +When `[slurm]` is present, the entrypoint: + +1. Resolves the full config +2. Renders a SLURM batch script from a Jinja2 template +3. Writes the script and resolved config to `{output_dir}/` +4. Submits via `sbatch` (or prints the script with `--slurm.dry-run`) + +For **single-node** jobs, the entire config is dumped to a TOML file and the template simply runs `uv run rl @` or `uv run sft @` on the allocated node. + +For **multi-node** jobs, sub-configs are written separately and `srun` dispatches processes across nodes. + +## Configuration + +### `[slurm]` — Job submission (shared between RL and SFT) + +| Field | Description | Default | +|---|---|---| +| `job_name` | SLURM job name | `"prime-rl"` | +| `project_dir` | Path to the project root on the cluster | `"."` | +| `template_path` | Path to a custom Jinja2 template | auto-selected | +| `partition` | SLURM partition | `"cluster"` | +| `nodelist` | Comma-separated list of specific nodes to run on (`--nodelist`) | `None` | +| `exclude` | Comma-separated list of nodes to exclude (`--exclude`) | `None` | +| `account` | SLURM account to charge (`--account`) | `None` | +| `time` | Maximum wall time, e.g. `"24:00:00"` (`--time`) | `None` | +| `pre_run_command` | Shell command to run on head node after env setup, before starting the job (e.g. cleanup) | `None` | + +### `[deployment]` — Node and GPU allocation + +**RL** uses a discriminated union with `type = "single_node"` (default) or `type = "multi_node"`: + +| Field | single_node | multi_node | +|---|---|---| +| `gpus_per_node` | Number of GPUs per node (default: 8) | Same | +| `num_train_gpus` | Training GPUs | — | +| `num_infer_gpus` | Inference GPUs | — | +| `num_train_nodes` | — | Training nodes | +| `num_infer_nodes` | — | Inference nodes | +| `nodes_per_fsdp_group` | — | Nodes per FSDP island (optional) | + +**SFT** follows the same pattern but only has training nodes: + +| Field | single_node | multi_node | +|---|---|---| +| `gpus_per_node` | Number of GPUs per node (default: 8) | Same | +| `num_gpus` | Number of GPUs (default: 1) | — | +| `num_nodes` | — | Training nodes (default: 2) | +| `nodes_per_fsdp_group` | — | Nodes per FSDP island (optional) | + +**Inference** runs independent vLLM replicas per node: + +| Field | single_node | multi_node | +|---|---|---| +| `gpus_per_node` | Number of GPUs per node (default: 8) | Same | +| `num_nodes` | — | Number of inference nodes (default: 1) | + +The SLURM template is auto-selected based on `deployment.type`. You can override it with `slurm.template_path`. + +### Constraints + +- `output_dir` should be explicitly set when using SLURM (defaults to `"outputs"`) +- Multi-node deployment requires `[slurm]` to be set + +--- + +## RL Examples + +### Single-node SLURM + +The simplest case: run on a single allocated node. No `[deployment]` needed — defaults to `single_node`. + +```toml +output_dir = "/shared/outputs/my-rl-run" + +[slurm] +job_name = "my-rl-run" +``` + +### Multi-node SLURM (Hendrycks Math) + +```toml +output_dir = "outputs/rl-math-moe" +max_steps = 500 +seq_len = 2048 + +[slurm] +job_name = "hendrycks-math-rl-moe" + +[deployment] +type = "multi_node" +num_train_nodes = 1 +num_infer_nodes = 1 + +[weight_broadcast] +type = "nccl" + +[model] +name = "Qwen/Qwen3-30B-A3B-Thinking-2507" + +[trainer.model] +impl = "custom" +attn = "flash_attention_3" +optim_cpu_offload = true + +[trainer.model.ac_offloading] +max_inflight_activations = 5 + +[trainer.model.ac] +freq = 1 + +[orchestrator] +batch_size = 512 +rollouts_per_example = 16 + +[orchestrator.sampling] +max_tokens = 2048 + +[[orchestrator.env]] +id = "math-env" +name = "hendrycks-math" +args = { dataset_name = "PrimeIntellect/Hendrycks-Math", dataset_subset = "default" } + +[inference.parallel] +tp = 4 +dp = 2 +``` + +See [`examples/hendrycks_math/rl.toml`](../examples/hendrycks_math/rl.toml) for the full example. + +--- + +## SFT Examples + +### Single-node SLURM + +```toml +output_dir = "/shared/outputs/my-sft-run" + +[slurm] +job_name = "my-sft-run" +``` + +### Multi-node SLURM (MoE SFT) + +```toml +output_dir = "outputs/sft-moe-math" +max_steps = 500 + +[slurm] +job_name = "sft-moe-math" + +[deployment] +type = "multi_node" +num_nodes = 2 + +[model] +name = "Qwen/Qwen3-30B-A3B-Thinking-2507" +impl = "custom" +attn = "flash_attention_3" +optim_cpu_offload = true + +[model.ac_offloading] +max_inflight_activations = 5 + +[model.ac] +freq = 1 + +[data] +type = "sft" +name = "PrimeIntellect/INTELLECT-3-SFT-10K" +subsets = ["default"] +splits = ["math"] +batch_size = 128 +seq_len = 8192 +``` + +See [`examples/hendrycks_math/sft.toml`](../examples/hendrycks_math/sft.toml) for the full example. + +--- + +## Inference Examples + +### Single-node SLURM + +Run a vLLM server on a single allocated node: + +```toml +output_dir = "/shared/outputs/my-inference" + +[model] +name = "Qwen/Qwen3-8B" + +[parallel] +tp = 8 + +[slurm] +job_name = "my-inference" +``` + +```bash +uv run inference @ inference_slurm.toml +``` + +### Multi-node SLURM + +Each node runs an independent vLLM replica. TP and DP must fit within a single node — there is no cross-node parallelism. + +```toml +output_dir = "/shared/outputs/my-inference" + +[model] +name = "PrimeIntellect/INTELLECT-3-RL-600" + +[parallel] +tp = 4 +dp = 2 + +[deployment] +type = "multi_node" +num_nodes = 4 + +[slurm] +job_name = "my-inference" +``` + +After submission, the SLURM template prints the inference URLs for all nodes (one per node). + +### Dry run + +Use `dry_run = true` to generate the sbatch script without submitting: + +```bash +uv run inference @ config.toml --dry-run true +``` + +--- + +## Custom SLURM Templates + +The default templates handle standard setups with InfiniBand detection, environment setup, and `srun`-based process dispatch. For advanced use cases (custom partitions, account settings, module loads, etc.), provide your own Jinja2 template: + +```bash +uv run rl @ my_config.toml --slurm.template-path path/to/my_template.sbatch.j2 +``` + +See [`src/prime_rl/templates/`](../src/prime_rl/templates/) for the default templates as a starting point. + +## Monitoring + +After submission, logs are available at: + +```bash +# All deployment types (trainer.log and inference.log are symlinks for multi-node) +tail -F {output_dir}/logs/trainer.log +tail -F {output_dir}/logs/orchestrator.log +tail -F {output_dir}/logs/inference.log + +# Multi-node: per-node logs +tail -F {output_dir}/logs/trainer/node_*.log +tail -F {output_dir}/logs/inference/node_*.log + +# Multi-node inference: per-replica router logs +tail -F {output_dir}/logs/inference/router_*.log +``` + +For convenience, a tmux launcher sets up a session with all log streams: + +```bash +bash scripts/tmux.sh my-rl-job /shared/outputs/my-rl-job +``` diff --git a/docs/testing-moe-at-small-scale.md b/docs/testing-moe-at-small-scale.md new file mode 100644 index 00000000000..ba2ca048f9f --- /dev/null +++ b/docs/testing-moe-at-small-scale.md @@ -0,0 +1,113 @@ +# Testing MoE at Small Scale + +When working on MoE architectures (GLM-4, Kimi, etc.), you can't iterate on a 100B+ parameter model locally. This guide shows how to create a small (~0.5B) MoE model with the same architecture, run SFT to warm it up, and run RL on it — all on 1-2 GPUs. + +The goal isn't performance. It's catching bugs in modeling code, state dict conversions, and training pipeline integration before running at scale. + +## Overview + +1. **Create + verify** a mini model with random weights and check HF <-> PrimeRL roundtrip +2. **SFT** to give it a non-trivial distribution +3. **RL** on reverse-text to validate the full pipeline + +## Prerequisites + +- At least 1 GPU for steps 1-2, 2 GPUs for step 3 (RL) +- Architecture presets are defined in `scripts/mini_moe.py` + +## Step 1: Create and verify the mini model + +```bash +uv run python scripts/mini_moe.py --arch glm4_moe --output-dir ./mini-glm-moe +``` + +This creates a ~543M parameter GLM-4 MoE (1024 hidden, 24 layers, 8 experts) with random weights, copies the tokenizer from the original GLM-4 model, then verifies that: +- Logits match between HF and PrimeRL implementations (`convert_to_prime`) +- The HF -> PrimeRL -> HF roundtrip is lossless (`convert_to_hf`) + +To re-run verification only (e.g. after a modeling code change): + +```bash +uv run python scripts/mini_moe.py --arch glm4_moe --output-dir ./mini-glm-moe --verify-only +``` + +## Step 2: SFT warmup + +Using the existing debug MoE SFT config with overrides for real data: + +```bash +uv run sft @ configs/debug/moe/sft/train.toml \ + --model.name ./mini-glm-moe \ + --data.name PrimeIntellect/Reverse-Text-SFT \ + --data.type null \ + --max_steps 200 \ + --optim.lr 1e-4 \ + --ckpt.weights +``` + +This fine-tunes on [PrimeIntellect/Reverse-Text-SFT](https://huggingface.co/datasets/PrimeIntellect/Reverse-Text-SFT) for 200 steps. Loss should drop from ~12 to ~2.5. The model won't be coherent, but it will have a non-trivial distribution so KL divergence is meaningful during RL. + +The latest weight checkpoint is saved under `outputs/weights/step_`. You can verify the roundtrip on it: + +```bash +uv run python scripts/mini_moe.py --arch glm4_moe --output-dir outputs/weights/step_200 --verify-only +``` + +A pre-built SFT'd model is available at [samsja/mini-glm-moe](https://huggingface.co/samsja/mini-glm-moe). + +## Step 3: RL (reverse-text) + +Requires 2 GPUs (one for inference, one for training). + +```bash +uv run rl @ configs/ci/integration/rl/start.toml \ + --model.name samsja/mini-glm-moe \ + --trainer.model.impl custom \ + --inference.gpu-memory-utilization 0.7 \ + --inference.model.max-model-len 2048 +``` + +Or to use the checkpoint from step 2: + +```bash +uv run rl @ configs/ci/integration/rl/start.toml \ + --model.name outputs/weights/step_200 \ + --trainer.model.impl custom \ + --inference.gpu-memory-utilization 0.7 \ + --inference.model.max-model-len 2048 +``` + +What to look for: +- **Training runs without crashing** — validates the full pipeline (inference server, orchestrator, trainer) +- **KL divergence is non-zero and finite** — confirms the reference model distribution is working +- **Loss is reasonable** — not NaN, not stuck at a constant value + +Don't expect the reward to go up meaningfully in 20 steps on a random model. + +## Adding a new architecture + +To test a new MoE architecture (e.g., Kimi2.5): + +1. Add modeling code under `src/prime_rl/trainer/models//` +2. Add a preset to `scripts/mini_moe.py` with the config class, small dimensions, HF model class, PrimeRL model class, and tokenizer source +3. Run steps 1-3 above with `--arch ` + +The preset defines the small config: + +```python +ARCH_PRESETS = { + "glm4_moe": { + "config_class": Glm4MoeConfig, + "config_kwargs": dict( + hidden_size=1024, + num_hidden_layers=24, + n_routed_experts=8, + # ... + ), + "hf_model_class": HFGlm4MoeForCausalLM, + "prime_model_class": PrimeRLGlm4MoeForCausalLM, + "tokenizer_source": "THUDM/GLM-4-9B-0414", + }, + # Add your new arch here +} +``` diff --git a/docs/trajectories.md b/docs/trajectories.md index 8d90e698a09..6e0b9d3e26f 100644 --- a/docs/trajectories.md +++ b/docs/trajectories.md @@ -2,29 +2,54 @@ Verifiers [v0.1.8](https://github.com/PrimeIntellect-ai/verifiers/releases/tag/v0.1.8) introduced trajectory-based rollouts, where each LLM request/response pair in a multi-turn interaction is recorded as an independent step. For details on the design decision, check the detailed [design document](https://github.com/PrimeIntellect-ai/verifiers/blob/main/notes/TRAJECTORIES.md) in the verifiers repository. -This allows two different training modes: -- **Interleaved Rollouts**: A multi-turn interaction is concatenated into a single training example (default, enabled via `--trajectory-strategy interleaved` on the orchestrator) -- **Branching Rollouts**: A multi-turn interaction is split into separate training examples (enabled via `--trajectory-strategy branching` on the orchestrator) +## Best-Effort Interleaved Rollouts -## Interleaved Rollouts +PRIME-RL uses a best-effort interleaving strategy that automatically merges consecutive trajectory steps when possible, and starts a new training sample when the extension property breaks. -Prior to verifiers v0.1.8, interleaved rollout was the only supported mode. Here, any multi-turn interaction is concatenated into a single training example. For example, a 3-turn conversation with alternating user and assistant messages will be concatenated into the following training example: +### The Extension Property + +A sequence of trajectory steps has the **extension property** when each successive step's prompt contains all previous prompts and completions as a prefix. When this holds: +- Multiple steps can be merged into a single training sample +- Compute scales as O(T) for a trajectory of length T + +When extension breaks (e.g., due to context compaction or thinking being stripped): +- A new training sample is started from that step +- Compute scales as O(T²) in the worst case (every step breaks extension) + +### How It Works -```txt -(U1,A1,U2,A2,U3,A3) ``` +5-step trajectory where extension breaks at step 4: + +Steps 1-3: extension holds → merged into Sample 1 +Step 4: extension breaks (e.g., thinking stripped from history) +Steps 4-5: extension holds → merged into Sample 2 + +Result: 2 training samples instead of 5 +``` + +This approach gives you the best of both worlds: +- When extension holds: O(T) compute, single merged sample +- When extension breaks: graceful fallback, no corrupted data +- Mixed scenarios: optimal merging where possible -This approach enforces a strict invariant: +### The Exact Prefix Invariant + +Interleaving enforces a strict invariant: > The prompt at turn $t$ must be the exact concatenation of prior messages exactly as the LLM originally generated them -We call this the "exact prefix" invariant. For example, at turn 2, the LLM should see U1,A1,U2 as the prompt, where U1 exactly matches the user message in turn 1 and A1 exactly matches the produced assistant message in turn 1. Any violation to this invariant will result in downstream problems when computing the importance sampling ratio during training. For example, assume that at turn 2 the prompt is U1,A1',U2 where A1' varies from A1. In this scenario it is not clear whether to add A1 or A1' to the interleaved rollout: +We call this the "exact prefix" invariant. For example, at turn 2, the LLM should see U1,A1,U2 as the prompt, where U1 exactly matches the user message in turn 1 and A1 exactly matches the produced assistant message in turn 1. Any violation to this invariant will result in downstream problems when computing the importance sampling ratio during training. + +For example, assume that at turn 2 the prompt is U1,A1',U2 where A1' varies from A1. In this scenario it is not clear whether to add A1 or A1' to the interleaved rollout: - If we add A1', the logprobs from turn 1 might be off because the inference LLM produced A1 but the trainer LLM is computing logprobs for A1' -- If we add A1, the logprobs from turn 2 might be off because the inference LLM is attending to A1' but the trainer LLM is attending to A1. +- If we add A1, the logprobs from turn 2 might be off because the inference LLM is attending to A1' but the trainer LLM is attending to A1 + +When the invariant is violated (extension breaks), PRIME-RL automatically starts a new training sample rather than producing corrupted data. ### Arbitrary Chat Templates -There exist chat templates which add, modify, or remove tokens across turns. One good example, is the chat template of the Qwen3-series of models, which strips thinking across user turns. +There exist chat templates which add, modify, or remove tokens across turns. One good example is the chat template of the Qwen3-series of models, which strips thinking across user turns. ```python from transformers import AutoTokenizer @@ -50,20 +75,22 @@ print(tokenizer.apply_chat_template(messages, tokenize=False)) # U2<|im_end|> ``` -The chat template automatically strips away past thinking section across user turn, which is often referred to as "interleaved thinking". Many chat templates, such as GLM or MiniMax, implement this approach, which makes training in multi-user turn environment (e.g. `alphabet-sort`) impossible when interleaving rollouts: Should we include the thinking section in the interleaved rollout or not? No matter what we do, we violate the exact prefix invariant, and will run into large logprob discrepancies between the trainer and inference. +The chat template automatically strips away past thinking sections across user turns, which is often referred to as "interleaved thinking". Many chat templates, such as GLM or MiniMax, implement this approach. -## Branching Rollouts +With best-effort interleaving, PRIME-RL handles this gracefully: when the thinking is stripped and the prefix no longer matches, a new training sample is started automatically. -When branching rollouts, a $n$-length multi-turn interaction will be split into $n$ separate training examples. For example, a 3-turn conversation with alternating user and assistant messages will be split into the following training examples: +### Discontinuous Trajectories by Design -```txt -1. U1,A1 -2. U1,A1,U2,A2 -3. U1,A1,U2,A2,U3,A3 -``` +Some multi-turn environments are intentionally discontinuous. For example, in a sub-agent calling scenario: + +1. Main agent receives a task and decides to delegate to a sub-agent +2. Sub-agent runs independently (possibly multiple turns with its own context) +3. Control returns to main agent with only the sub-agent's final result + +The main agent's trajectory is discontinuous because the sub-agent's internal conversation isn't part of its context. When the main agent resumes, its prompt doesn't extend the previous turn - it contains a summarized result instead. -This approach alleviates both issues that interleaved rollouts suffer from: it naturally handles any chat template, and avoids the retokenization issues all together. This is, because we *turn a single multi-turn interaction into many single-turn interactions*. In a sense this method is a safer default that works out of the box, but it comes at the (significant) cost of increased training time, because the number of training tokens now increases linearly with the number of turns. +Best-effort interleaving handles this naturally: each agent's contiguous turns get merged, but the handoff between agents starts a new sample. -## Verifiers vs. PRIME-RL +## Deprecated: Branching Mode -The verifiers trainer is using branching rollouts by default, while PRIME-RL supports two training modes, with interleaved rollouts being the default. It is an example of different design philosophies between the two projects. Verifiers is designed to be a general-purpose trainer that "just works" out of the box, while PRIME-RL is designed to be a high-performance production-scale trainer. \ No newline at end of file +The `--trajectory-strategy branching` option is deprecated. The best-effort interleaving strategy now handles all cases automatically, falling back to separate samples (equivalent to branching) when the extension property breaks. diff --git a/examples/Intellect-3.1/README.md b/examples/Intellect-3.1/README.md new file mode 100644 index 00000000000..b31828cb9b4 --- /dev/null +++ b/examples/Intellect-3.1/README.md @@ -0,0 +1,78 @@ +# Intellect-3.1 + +This example guides you through reproducing the [INTELLECT-3.1](https://huggingface.co/PrimeIntellect/INTELLECT-3.1) RL training on top of [GLM-4.5-Air-Base](https://huggingface.co/zai-org/GLM-4.5-Air-Base). + +## Requirements + +You need access to a Slurm cluster with at least 16 nodes to run this example. Each node must have a shared filesystem. In this guide we assume the NFS is mounted at `/shared`; you can change it to your own path. + +If you don't have access to that many nodes, you can downscale inference and training nodes to lower values and scale down the batch size accordingly. + +For training, going below 4 nodes might cause OOM. + +You also need to have prime-rl cloned on your cluster into the shared filesystem. + +```bash +git clone https://github.com/PrimeIntellect-ai/prime-rl.git /shared/prime-rl +cd /shared/prime-rl +uv sync --all-extras +``` + +You might also want to create a `.env` file inside the prime-rl directory to store environment variables used during training like W&B and Hugging Face tokens. The `.env` file is automatically sourced during training. + +```bash +touch .env +``` + +```bash +echo "WANDB_API_KEY=your_wandb_api_key" >> .env +echo "HUGGINGFACE_TOKEN=your_huggingface_token" >> .env +``` + +### sandbox + +The [mini-swe-agent-plus](https://github.com/PrimeIntellect-ai/sandbox-mini-swe-agent-plus) environment is configured to use Prime Intellect Sandboxes. You can find more information about the sandboxes [here](https://docs.primeintellect.ai/sandboxes/overview). + +You will need to create a sandbox account and add the credentials to the `.env` file. + +Alternatively, you can adapt the code of the environment to use your own sandbox implementation. + +## Tmux session + +We recommend using the tmux helper to start the run and look at the logs. + +From your Slurm head node: + +```bash +bash scripts/tmux.sh intellect-3.1 /shared/outputs/intellect-3.1 +``` + +You can then attach to it by doing `tmux attach -t intellect-3.1`. + +## Start the run + +Run the following command to start the RL training: + +PS: If using the tmux helper, you can run the command in the `Terminal` (window 0) pane and look at the logs in the `Logs` (window 1) pane. + +```bash +uv run rl @ examples/Intellect-3.1/rl.toml --output-dir /shared/outputs/intellect-3.1 +``` + +Output of the command: +``` +XXX:XX:XX INFO Wrote subconfigs to /shared/outputs/intellect-3.1/configs [rl.py::515] +XXX:XX:XX INFO Wrote SLURM script to /shared/outputs/intellect-3.1/rl.sbatch [rl.py::534] +XXX:XX:XX INFO Submitting: sbatch /shared/outputs/intellect-3.1/rl.sbatch [rl.py::540] +XXX:XX:XX SUCCESS Submitted batch job YYYY + +Logs: + Trainer: tail -F /shared/outputs/intellect-3.1/logs/trainer.log + Orchestrator: tail -F /shared/outputs/intellect-3.1/logs/orchestrator.log + Inference: tail -F /shared/outputs/intellect-3.1/logs/inference.log + Envs: tail -F /shared/outputs/intellect-3.1/logs/envs/*/*/*.log + Train: tail -F /shared/outputs/intellect-3.1/logs/envs/train/*/*.log + swe: tail -F /shared/outputs/intellect-3.1/logs/envs/train/swe/*.log +``` + + diff --git a/examples/Intellect-3.1/rl.toml b/examples/Intellect-3.1/rl.toml new file mode 100644 index 00000000000..8c4d1001e7e --- /dev/null +++ b/examples/Intellect-3.1/rl.toml @@ -0,0 +1,97 @@ +seq_len = 131072 + +[log] +level = "debug" + +[model] +name = "PrimeIntellect/INTELLECT-3-Base" + +[weight_broadcast] +type = "nccl" + +[deployment] +type = "multi_node" +num_train_nodes = 4 +num_infer_nodes = 12 + +[slurm] +job_name = "swe-qwen30b" + +[ckpt] +interval = 100 + +[wandb] +project = "intellect-3.1" + +[trainer.model] +impl = "custom" +attn = "flash_attention_3" +ep = 8 +cp = 2 +fused_lm_head_token_chunk_size = 1024 + +[trainer.model.compile] + +[trainer.model.ac_offloading] +max_inflight_activations = 5 + +[trainer.model.ac] +freq = 1 + +[trainer.optim] +type = "muon" +lr = 1e-6 +weight_decay = 0.01 + +[orchestrator] +batch_size = 2048 +oversampling_factor = 2 + +[[orchestrator.train.env]] +id = "mini-swe-agent-plus" +name = "swe" +ratio = 0.3 +args = { max_turns = 200, cpu_cores = 2, memory_gb = 4, disk_size_gb = 4, labels = ["mini-swe-agent-plus"], total_timeout_minutes = 720, sandbox_client_max_workers = 256, max_command_timeouts = 3, sandbox_command_timeout = 30} + +[[orchestrator.train.env]] +id = "deepdive" +name = "deepdive" +ratio = 0.2 +args = { finish_with_tool = true, open_max_workers = 128, cache_dir = "/tmp/i3_deepdive_cache_train" } + +[[orchestrator.train.env]] +id = "math-env" +name = "math" +ratio = 0.3 +args = { min_avg_reward = 0.0, max_avg_reward = 0.874} + +[[orchestrator.train.env]] +id = "logic-env" +name = "logic" +ratio = 0.2 +args = { min_avg_reward = 0.0, max_avg_reward = 0.874 } + +[[orchestrator.train.env]] +id = "code-env" +name = "code" +ratio = 0.2 +args = { pool_size = 512 } + +[orchestrator.buffer] +easy_threshold = 1.0 +online_difficulty_filtering = true +seed = 42 + +[orchestrator.eval] +interval = 25 + +[[orchestrator.eval.env]] +id = "mini-swe-agent-plus" +name = "swe-bench-verified-quick" +args = { dataset_name="PrimeIntellect/SWE-Bench-Verified-Quick", max_turns = 200, cpu_cores = 2, memory_gb = 4, disk_size_gb = 4, labels = ["mini-swe-agent-plus"], total_timeout_minutes = 720, sandbox_client_max_workers = 256, max_command_timeouts = 3, sandbox_command_timeout = 30} + +[[orchestrator.eval.env]] +id = "aime2025" + +[inference.parallel] +tp = 8 \ No newline at end of file diff --git a/examples/alphabet_sort/README.md b/examples/alphabet_sort/README.md index 436639f9a89..57eeef90dff 100644 --- a/examples/alphabet_sort/README.md +++ b/examples/alphabet_sort/README.md @@ -90,7 +90,8 @@ The correct alphabetical order by last name should be: Akman, Cakir, Durgun. The We train with LoRA (rank 32, alpha 64) for 100 steps. -![RL](rl/wandb.png) +![RL](rl.png) + *Check out the logs on [W&B](https://wandb.ai/primeintellect/alphabet-sort-4b-lora/workspace?nw=nwuserandrewpi).* ```bash # In the `Trainer` pane diff --git a/examples/alphabet_sort/rl/wandb.png b/examples/alphabet_sort/rl.png similarity index 100% rename from examples/alphabet_sort/rl/wandb.png rename to examples/alphabet_sort/rl.png diff --git a/examples/alphabet_sort/rl.toml b/examples/alphabet_sort/rl.toml index 10c6a65afe9..480cc9f04b8 100644 --- a/examples/alphabet_sort/rl.toml +++ b/examples/alphabet_sort/rl.toml @@ -12,7 +12,7 @@ project = "alphabet-sort" name = "alphabet-sort" [trainer.model] -impl = "liger_kernel" +impl = "auto" [trainer.model.ac] freq = 1 @@ -28,10 +28,10 @@ lr = 1e-5 batch_size = 512 rollouts_per_example = 8 -[orchestrator.sampling] -max_tokens = 768 +[orchestrator.train.sampling] +max_completion_tokens = 768 -[[orchestrator.env]] +[[orchestrator.train.env]] id = "primeintellect/alphabet-sort" name = "alphabet-sort" args = { min_turns = 3, max_turns = 5, power_per_turn = false } diff --git a/examples/alphabet_sort/sft_distill_hard.toml b/examples/alphabet_sort/sft_distill_hard.toml new file mode 100644 index 00000000000..35f29607a14 --- /dev/null +++ b/examples/alphabet_sort/sft_distill_hard.toml @@ -0,0 +1,49 @@ +max_steps = 24 +seq_len = 2048 + +[deployment] +type = "single_node" +gpus_per_node = 2 +num_train_gpus = 2 +num_infer_gpus = 0 + +[model] +name = "Qwen/Qwen3-4B-Instruct-2507" + +[wandb] +project = "alphabet-sort-sft-distill" +name = "qwen3-4b-hard-distill" + +[ckpt] + +[trainer.optim] +lr = 1e-5 + +[trainer.model.lora] +rank = 16 +alpha = 32 + +[trainer.ckpt.weights] +save_adapter_separately = true + +[orchestrator] +batch_size = 256 +rollouts_per_example = 4 +use_token_client = false +use_sft_loss = true + +[orchestrator.train.sampling] +max_completion_tokens = 512 +temperature = 0.7 + +[orchestrator.teacher_rollout_model.client] +base_url = ["https://api.pinference.ai/api/v1"] +api_key_var = "PRIME_API_KEY" + +[orchestrator.teacher_rollout_model.model] +name = "qwen/qwen3-235b-a22b-instruct-2507" + +[[orchestrator.train.env]] +id = "primeintellect/alphabet-sort" +name = "alphabet-sort" +args = { min_turns = 2, max_turns = 2, min_names_per_turn = 1, max_names_per_turn = 3, similarity_power = 4, power_per_turn = false } diff --git a/examples/alphabet_sort/slurm_rl.toml b/examples/alphabet_sort/slurm_rl.toml new file mode 100644 index 00000000000..3295bdb8631 --- /dev/null +++ b/examples/alphabet_sort/slurm_rl.toml @@ -0,0 +1,5 @@ +# Usage: uv run rl @ examples/alphabet_sort/rl.toml @ examples/alphabet_sort/slurm_rl.toml +output_dir = "outputs/alphabet-sort" + +[slurm] +job_name = "alphabet-sort" \ No newline at end of file diff --git a/examples/glm5_pd_disag/README.md b/examples/glm5_pd_disag/README.md new file mode 100644 index 00000000000..4fff3a47de9 --- /dev/null +++ b/examples/glm5_pd_disag/README.md @@ -0,0 +1,74 @@ +# GLM-5 PD Disaggregation + +This example guides you through deploying a highly efficient GLM-5 RL training with PD disaggregation and FP8 inference on agentic SWE tasks. + +## Requirements + +You need access to a Slurm cluster with at least 32 nodes to run this example. Each node must have a shared filesystem. In this guide we assume the NFS is mounted at `/shared`; you can change it to your own path. + +You also need to have prime-rl cloned on your cluster into the shared filesystem. + +```bash +git clone https://github.com/PrimeIntellect-ai/prime-rl.git /shared/prime-rl +cd /shared/prime-rl +uv sync --all-extras +``` + +You might also want to create a `.env` file inside the prime-rl directory to store environment variables used during training like W&B and Hugging Face tokens. The `.env` file is automatically sourced during training. + +```bash +touch .env +``` + +```bash +echo "WANDB_API_KEY=your_wandb_api_key" >> .env +echo "HUGGINGFACE_TOKEN=your_huggingface_token" >> .env +``` + +### sandbox + +The [mini-swe-agent-plus](https://github.com/PrimeIntellect-ai/sandbox-mini-swe-agent-plus) environment is configured to use Prime Intellect Sandboxes. You can find more information about the sandboxes [here](https://docs.primeintellect.ai/sandboxes/overview). + +You will need to create a sandbox account and add the credentials to the `.env` file. + +Alternatively, you can adapt the code of the environment to use your own sandbox implementation. + +## Tmux session + +We recommend using the tmux helper to start the run and look at the logs. + +From your Slurm head node: + +```bash +bash scripts/tmux.sh glm5-pd-disag /shared/outputs/glm5-pd-disag +``` + +You can then attach to it by doing `tmux attach -t glm5-pd-disag`. + +## Start the run + +Run the following command to start the RL training: + +PS: If using the tmux helper, you can run the command in the `Terminal` (window 0) pane and look at the logs in the `Logs` (window 1) pane. + +```bash +uv run rl @ examples/glm5_pd_disag/rl.toml --output-dir /shared/outputs/glm5-pd-disag +``` + +Output of the command: +``` +XXX:XX:XX INFO Wrote subconfigs to /shared/outputs/glm5-pd-disag/configs [rl.py::515] +XXX:XX:XX INFO Wrote SLURM script to /shared/outputs/glm5-pd-disag/rl.sbatch [rl.py::534] +XXX:XX:XX INFO Submitting: sbatch /shared/outputs/glm5-pd-disag/rl.sbatch [rl.py::540] +XXX:XX:XX SUCCESS Submitted batch job YYYY + +Logs: + Trainer: tail -F /shared/outputs/glm5-pd-disag/logs/trainer.log + Orchestrator: tail -F /shared/outputs/glm5-pd-disag/logs/orchestrator.log + Inference: tail -F /shared/outputs/glm5-pd-disag/logs/inference.log + Envs: tail -F /shared/outputs/glm5-pd-disag/logs/envs/*/*/*.log + Train: tail -F /shared/outputs/glm5-pd-disag/logs/envs/train/*/*.log + swe: tail -F /shared/outputs/glm5-pd-disag/logs/envs/train/swe/*.log +``` + + diff --git a/examples/glm5_pd_disag/rl.toml b/examples/glm5_pd_disag/rl.toml new file mode 100644 index 00000000000..ac7fd664726 --- /dev/null +++ b/examples/glm5_pd_disag/rl.toml @@ -0,0 +1,97 @@ +seq_len = 131072 + +[log] +level = "debug" + +[wandb] +project = "glm5-pd-disag" + +[model] +name = "zai-org/GLM-5" + +[weight_broadcast] +type = "nccl" +port = 29502 +timeout = 12000 +quantize_in_weight_transfer = true + +[deployment] +type = "multi_node" +num_train_nodes = 16 +num_infer_nodes = 8 +gpus_per_node = 8 +num_infer_replicas = 2 + +[inference.deployment] +type = "disaggregated" +num_prefill_nodes = 4 +num_decode_nodes = 4 + +[slurm] +job_name = "glm5-pd-disag" + + +[ckpt] +interval = 100 + +[trainer] +dist_timeout_seconds = 7200 + +[trainer.model] +impl = "custom" +attn = "flash_attention_2" # dsa is skipping flash attn anyway +fused_lm_head_token_chunk_size = 2048 +optim_cpu_offload = true +ep = 8 + +[trainer.model.ac] +freq = 1 + +[trainer.model.compile] + +[trainer.model.ac_offloading] +max_inflight_activations = 1 + +[trainer.ckpt] +skip_optimizer = true +skip_gather_master_weights = true + +[trainer.optim] +type = "adamw" +lr = 1e-6 +weight_decay = 0.1 + + +[orchestrator] +batch_size = 4096 +rollouts_per_example = 16 +oversampling_factor = 3 +max_off_policy_steps = 16 + +[orchestrator.eval] +interval = 20 + +[[orchestrator.eval.env]] +id = "mini-swe-agent-plus" +name = "swe-bench-verified-quick" +args = { dataset_name="PrimeIntellect/SWE-Bench-Verified-Quick", max_turns = 200, cpu_cores = 2, memory_gb = 4, disk_size_gb = 4, labels = ["mini-swe-agent-plus"], total_timeout_minutes = 720, sandbox_client_max_workers = 256, max_command_timeouts = 3, sandbox_command_timeout = 30} + +[[orchestrator.filters]] +type = "gibberish" +enforce = true + +[orchestrator.buffer] +easy_threshold = 1.0 +hard_threshold = 0.0 +seed = 42 + +[inference] +enable_expert_parallel = true +# we need <0.85 bc glm5 layers are too large for 0.85 +gpu_memory_utilization = 0.80 +enable_eplb = false +use_deep_gemm = true + +[inference.model] +name = "zai-org/GLM-5-FP8" +tool_call_parser = "glm47" diff --git a/examples/hendrycks_sanity/README.md b/examples/hendrycks_sanity/README.md new file mode 100644 index 00000000000..aa07c46cdf6 --- /dev/null +++ b/examples/hendrycks_sanity/README.md @@ -0,0 +1,25 @@ +# Hendrycks Sanity + +This example runs the Hendrycks Sanity Check experiment proposed in [Defeating the Training-Inference Mismatch](https://arxiv.org/abs/2510.26788). The sanity check tests whether an RL algorithm can reliably improve a model on problems it can *already partially solve*. The dataset is filtered from MATH to only include problems where the base model (`DeepSeek-R1-Distill-Qwen-1.5B`) solves 20-80% of the time across 40 rollouts. A reliable algorithm should push training accuracy on this "perfectible" subset above 95%. + +Because our trainer is asynchronous, we perform only one gradient step per batch (the inference engine generates the next batch while the trainer processes the current one). + +> This example runs on 8 GPUs (4 for inference, 4 for training). + +## Training + +Schedule training locally on a node with 8 GPUs + +```bash +uv run rl @ examples/hendrycks_sanity/rl.toml \ + --wandb.project your-project \ + --wandb.name your-run +``` + +Or schedule the same training via SLURM + +```bash +uv run rl @ examples/hendrycks_sanity/slurm_rl.toml \ + --wandb.project your-project \ + --wandb.name your-run +``` diff --git a/examples/hendrycks_sanity/rl.toml b/examples/hendrycks_sanity/rl.toml new file mode 100644 index 00000000000..925541abf1f --- /dev/null +++ b/examples/hendrycks_sanity/rl.toml @@ -0,0 +1,37 @@ +max_steps = 5000 +seq_len = 8192 + +[wandb] +project = "hendrycks-sanity" +name = "hendrycks-sanity" + +[deployment] +num_train_gpus = 4 +num_infer_gpus = 4 + +[model] +name = "deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B" + +[orchestrator] +batch_size = 512 +rollouts_per_example = 8 + +[[orchestrator.train.env]] +id = "math-env" # included in lock file +args = { dataset_name = "mikasenghaas/Sanity-Test-R1D-1.5B", dataset_subset = "default" } +name = "hendrycks-math" + +[orchestrator.eval] +interval = 50 + +[[orchestrator.eval.env]] +id = "primeintellect/aime2024" +name = "aime2024" +rollouts_per_example = 32 + +[trainer.model] +name = "deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B" +seq_len = 16384 + +[inference.model] +max_model_len = 8192 \ No newline at end of file diff --git a/examples/hendrycks_sanity/slurm_rl.toml b/examples/hendrycks_sanity/slurm_rl.toml new file mode 100644 index 00000000000..2fc4249d24f --- /dev/null +++ b/examples/hendrycks_sanity/slurm_rl.toml @@ -0,0 +1,5 @@ +# Usage: uv run rl @ examples/hendrycks_sanity/rl.toml @ examples/hendrycks_sanity/slurm_rl.toml +output_dir = "outputs/hendrycks-sanity" + +[slurm] +job_name = "hendrycks-sanity" \ No newline at end of file diff --git a/examples/minimax_m2.5_swe/README.md b/examples/minimax_m2.5_swe/README.md new file mode 100644 index 00000000000..3df08f3c8e6 --- /dev/null +++ b/examples/minimax_m2.5_swe/README.md @@ -0,0 +1,74 @@ +# MiniMax-M2.5 SWE + +This example guides you through RL training [MiniMax-M2.5](https://huggingface.co/MiniMaxAI/MiniMax-M2.5) on agentic SWE tasks. + +## Requirements + +You need access to a Slurm cluster with at least 16 nodes to run this example. Each node must have a shared filesystem. In this guide we assume the NFS is mounted at `/shared`; you can change it to your own path. + +You also need to have prime-rl cloned on your cluster into the shared filesystem. + +```bash +git clone https://github.com/PrimeIntellect-ai/prime-rl.git /shared/prime-rl +cd /shared/prime-rl +uv sync --all-extras +``` + +You might also want to create a `.env` file inside the prime-rl directory to store environment variables used during training like W&B and Hugging Face tokens. The `.env` file is automatically sourced during training. + +```bash +touch .env +``` + +```bash +echo "WANDB_API_KEY=your_wandb_api_key" >> .env +echo "HUGGINGFACE_TOKEN=your_huggingface_token" >> .env +``` + +### sandbox + +The [mini-swe-agent-plus](https://github.com/PrimeIntellect-ai/sandbox-mini-swe-agent-plus) environment is configured to use Prime Intellect Sandboxes. You can find more information about the sandboxes [here](https://docs.primeintellect.ai/sandboxes/overview). + +You will need to create a sandbox account and add the credentials to the `.env` file. + +Alternatively, you can adapt the code of the environment to use your own sandbox implementation. + +## Tmux session + +We recommend using the tmux helper to start the run and look at the logs. + +From your Slurm head node: + +```bash +bash scripts/tmux.sh minimax-swe /shared/outputs/minimax-swe +``` + +You can then attach to it by doing `tmux attach -t minimax-swe`. + +## Start the run + +Run the following command to start the RL training: + +PS: If using the tmux helper, you can run the command in the `Terminal` (window 0) pane and look at the logs in the `Logs` (window 1) pane. + +```bash +uv run rl @ examples/minimax_m2.5_swe/rl.toml --output-dir /shared/outputs/minimax-swe +``` + +Output of the command: +``` +XXX:XX:XX INFO Wrote subconfigs to /shared/outputs/minimax-swe/configs [rl.py::515] +XXX:XX:XX INFO Wrote SLURM script to /shared/outputs/minimax-swe/rl.sbatch [rl.py::534] +XXX:XX:XX INFO Submitting: sbatch /shared/outputs/minimax-swe/rl.sbatch [rl.py::540] +XXX:XX:XX SUCCESS Submitted batch job YYYY + +Logs: + Trainer: tail -F /shared/outputs/minimax-swe/logs/trainer.log + Orchestrator: tail -F /shared/outputs/minimax-swe/logs/orchestrator.log + Inference: tail -F /shared/outputs/minimax-swe/logs/inference.log + Envs: tail -F /shared/outputs/minimax-swe/logs/envs/*/*/*.log + Train: tail -F /shared/outputs/minimax-swe/logs/envs/train/*/*.log + swe: tail -F /shared/outputs/minimax-swe/logs/envs/train/swe/*.log +``` + + diff --git a/examples/minimax_m2.5_swe/rl.toml b/examples/minimax_m2.5_swe/rl.toml new file mode 100644 index 00000000000..cc9fd893121 --- /dev/null +++ b/examples/minimax_m2.5_swe/rl.toml @@ -0,0 +1,71 @@ +seq_len = 131072 + +[log] +level = "debug" + +[model] +name = "PrimeIntellect/MiniMax-M2.5-bf16" + +[weight_broadcast] +type = "nccl" + +[deployment] +type = "multi_node" +num_train_nodes = 4 +num_infer_nodes = 12 + +[slurm] +job_name = "swe-minimax" + +[ckpt] +interval = 100 + +[wandb] +project = "minimax-swe" + +[trainer.model] +impl = "custom" +attn = "flash_attention_3" +ep = 8 +cp = 4 +fused_lm_head_token_chunk_size = 1024 + +[trainer.model.compile] + +[trainer.model.ac_offloading] +max_inflight_activations = 5 + +[trainer.model.ac] +freq = 1 + +[trainer.optim] +type = "adamw" +lr = 1e-6 +weight_decay = 0.01 + +[orchestrator] +batch_size = 2048 +oversampling_factor = 2 +max_off_policy_steps = 16 + + +[[orchestrator.train.env]] +id = "mini-swe-agent-plus" +name = "swe" +args = { max_turns = 200, cpu_cores = 2, memory_gb = 4, disk_size_gb = 4, labels = ["mini-swe-agent-plus"], total_timeout_minutes = 720, sandbox_client_max_workers = 256, max_command_timeouts = 3, sandbox_command_timeout = 30} + +[orchestrator.eval] +interval = 25 + +[[orchestrator.eval.env]] +id = "mini-swe-agent-plus" +name = "swe-bench-verified-quick" +args = { dataset_name="PrimeIntellect/SWE-Bench-Verified-Quick", max_turns = 200, cpu_cores = 2, memory_gb = 4, disk_size_gb = 4, labels = ["mini-swe-agent-plus"], total_timeout_minutes = 720, sandbox_client_max_workers = 256, max_command_timeouts = 3, sandbox_command_timeout = 30} + +[orchestrator.buffer] +easy_threshold = 1.0 +hard_threshold = 0.0 +seed = 42 + +[inference.parallel] +tp = 8 \ No newline at end of file diff --git a/examples/multinode/rl.toml b/examples/multinode/rl.toml new file mode 100644 index 00000000000..6d2854077ac --- /dev/null +++ b/examples/multinode/rl.toml @@ -0,0 +1,55 @@ +output_dir = "outputs/rl-math-moe" +max_steps = 500 +seq_len = 2048 + +[slurm] +job_name = "hendrycks-math-rl-moe" + +[deployment] +type = "multi_node" +num_train_nodes = 1 +num_infer_nodes = 1 + +[log] +level = "debug" + +[wandb] +project = "hendrycks-math-debug" +name = "hendrycks-math" + +[weight_broadcast] +type = "nccl" + +[model] +name = "Qwen/Qwen3-30B-A3B-Thinking-2507" + +[trainer.model] +impl = "custom" +attn = "flash_attention_3" +optim_cpu_offload = true + +[trainer.model.ac_offloading] +max_inflight_activations = 5 + +[trainer.model.ac] +freq = 1 + +[orchestrator] +batch_size = 512 +rollouts_per_example = 16 + +[orchestrator.train.sampling] +max_completion_tokens = 2048 + +[[orchestrator.train.env]] +id = "math-env" +name = "hendrycks-math" +args = { dataset_name = "PrimeIntellect/Hendrycks-Math", dataset_subset = "default", math_verify_max_workers = 128, math_verify_timeout = 60 } + +[orchestrator.buffer] +easy_threshold = 1.0 +hard_threshold = 0.0 + +[inference.parallel] +tp = 4 +dp = 2 diff --git a/examples/multinode/sft.toml b/examples/multinode/sft.toml new file mode 100644 index 00000000000..be15363830a --- /dev/null +++ b/examples/multinode/sft.toml @@ -0,0 +1,35 @@ +output_dir = "outputs/sft-moe-math" +max_steps = 500 + +[slurm] +job_name = "sft-moe-math" + +[deployment] +type = "multi_node" +num_nodes = 2 + +[wandb] +project = "sft-moe-intellect" +name = "sft-moe-math" + +[model] +name = "Qwen/Qwen3-30B-A3B-Thinking-2507" +impl = "custom" +attn = "flash_attention_3" +optim_cpu_offload = true + +[model.compile] + +[model.ac_offloading] +max_inflight_activations = 5 + +[model.ac] +freq = 1 + +[data] +type = "sft" +name = "PrimeIntellect/INTELLECT-3-SFT-10K" +subsets = ["default"] +splits = ["math"] +batch_size = 128 +seq_len = 8192 diff --git a/examples/qwen30b_math/README.md b/examples/qwen30b_math/README.md new file mode 100644 index 00000000000..3fa5d64eb9e --- /dev/null +++ b/examples/qwen30b_math/README.md @@ -0,0 +1,66 @@ +# Qwen 3 30B - A3B Math + +This example guides you through RL training [Qwen3-30B-A3B-Thinking-2507](https://huggingface.co/Qwen/Qwen3-30B-A3B-Thinking-2507) on math tasks. + +## Requirements + +You need access to a Slurm cluster with at least 4 nodes to run this example. Each node must have a shared filesystem. In this guide we assume the NFS is mounted at `/shared`; you can change it to your own path. + +You also need to have prime-rl cloned on your cluster into the shared filesystem. + +```bash +git clone https://github.com/PrimeIntellect-ai/prime-rl.git /shared/prime-rl +cd /shared/prime-rl +uv sync --all-extras +``` + +You might also want to create a `.env` file inside the prime-rl directory to store environment variables used during training like W&B and Hugging Face tokens. The `.env` file is automatically sourced during training. + +```bash +touch .env +``` + +```bash +echo "WANDB_API_KEY=your_wandb_api_key" >> .env +echo "HUGGINGFACE_TOKEN=your_huggingface_token" >> .env +``` + +## Tmux session + +We recommend using the tmux helper to start the run and look at the logs. + +From your Slurm head node: + +```bash +bash scripts/tmux.sh qwen30b-math /shared/outputs/qwen30b-math +``` + +You can then attach to it by doing `tmux attach -t qwen30b-math`. + +## Start the run + +Run the following command to start the RL training: + +PS: If using the tmux helper, you can run the command in the `Terminal` (window 0) pane and look at the logs in the `Logs` (window 1) pane. + +```bash +uv run rl @ examples/qwen30b_math/rl.toml --output-dir /shared/outputs/qwen30b-math +``` + +Output of the command: +``` +XXX:XX:XX INFO Wrote subconfigs to /shared/outputs/qwen30b-math/configs [rl.py::515] +XXX:XX:XX INFO Wrote SLURM script to /shared/outputs/qwen30b-math/rl.sbatch [rl.py::534] +XXX:XX:XX INFO Submitting: sbatch /shared/outputs/qwen30b-math/rl.sbatch [rl.py::540] +XXX:XX:XX SUCCESS Submitted batch job YYYY + +Logs: + Trainer: tail -F /shared/outputs/qwen30b-math/logs/trainer.log + Orchestrator: tail -F /shared/outputs/qwen30b-math/logs/orchestrator.log + Inference: tail -F /shared/outputs/qwen30b-math/logs/inference.log + Envs: tail -F /shared/outputs/qwen30b-math/logs/envs/*/*/*.log + Train: tail -F /shared/outputs/qwen30b-math/logs/envs/train/*/*.log + math: tail -F /shared/outputs/qwen30b-math/logs/envs/train/math/*.log +``` + + diff --git a/examples/qwen30b_math/rl.toml b/examples/qwen30b_math/rl.toml new file mode 100644 index 00000000000..82edf7e0505 --- /dev/null +++ b/examples/qwen30b_math/rl.toml @@ -0,0 +1,67 @@ +seq_len = 32768 + +[log] +level = "debug" + +[model] +name = "Qwen/Qwen3-30B-A3B-Thinking-2507" + +[weight_broadcast] +type = "nccl" + +[deployment] +type = "multi_node" +num_train_nodes = 2 +num_infer_nodes = 2 + +[slurm] +job_name = "math-qwen30b" + +[ckpt] +interval = 100 + +[wandb] +project = "qwen30b-math" + +[trainer.model] +impl = "custom" +attn = "flash_attention_3" +ep = 8 + +[trainer.model.compile] + +[trainer.model.ac_offloading] +max_inflight_activations = 5 + +[trainer.model.ac] +freq = 1 + +[trainer.optim] +type = "adamw" +lr = 1e-6 + +[orchestrator] +batch_size = 512 +oversampling_factor = 2 +max_off_policy_steps = 8 + +[orchestrator.train.sampling] +max_completion_tokens = 32768 + +[[orchestrator.train.env]] +id = "math-env" +name = "math" + +[orchestrator.eval] +interval = 25 + +[[orchestrator.eval.env]] +id = "aime2025" + +[orchestrator.buffer] +easy_threshold = 1.0 +hard_threshold = 0.0 +seed = 42 + +[inference.parallel] +tp = 8 \ No newline at end of file diff --git a/examples/qwen30b_swe/README.md b/examples/qwen30b_swe/README.md new file mode 100644 index 00000000000..14a41cc004a --- /dev/null +++ b/examples/qwen30b_swe/README.md @@ -0,0 +1,74 @@ +# Qwen 3 30B - A3B SWE + +This example guides you through RL training [Qwen3-30B-A3B-Thinking-2507](https://huggingface.co/Qwen/Qwen3-30B-A3B-Thinking-2507) on agentic SWE tasks. + +## Requirements + +You need access to a Slurm cluster with at least 4 nodes to run this example. Each node must have a shared filesystem. In this guide we assume the NFS is mounted at `/shared`; you can change it to your own path. + +You also need to have prime-rl cloned on your cluster into the shared filesystem. + +```bash +git clone https://github.com/PrimeIntellect-ai/prime-rl.git /shared/prime-rl +cd /shared/prime-rl +uv sync --all-extras +``` + +You might also want to create a `.env` file inside the prime-rl directory to store environment variables used during training like W&B and Hugging Face tokens. The `.env` file is automatically sourced during training. + +```bash +touch .env +``` + +```bash +echo "WANDB_API_KEY=your_wandb_api_key" >> .env +echo "HUGGINGFACE_TOKEN=your_huggingface_token" >> .env +``` + +### sandbox + +The [mini-swe-agent-plus](https://github.com/PrimeIntellect-ai/sandbox-mini-swe-agent-plus) environment is configured to use Prime Intellect Sandboxes. You can find more information about the sandboxes [here](https://docs.primeintellect.ai/sandboxes/overview). + +You will need to create a sandbox account and add the credentials to the `.env` file. + +Alternatively, you can adapt the code of the environment to use your own sandbox implementation. + +## Tmux session + +We recommend using the tmux helper to start the run and look at the logs. + +From your Slurm head node: + +```bash +bash scripts/tmux.sh qwen30b-swe /shared/outputs/qwen30b-swe +``` + +You can then attach to it by doing `tmux attach -t qwen30b-swe`. + +## Start the run + +Run the following command to start the RL training: + +PS: If using the tmux helper, you can run the command in the `Terminal` (window 0) pane and look at the logs in the `Logs` (window 1) pane. + +```bash +uv run rl @ examples/qwen30b_swe/rl.toml --output-dir /shared/outputs/qwen30b-swe +``` + +Output of the command: +``` +XXX:XX:XX INFO Wrote subconfigs to /shared/outputs/qwen30b-swe/configs [rl.py::515] +XXX:XX:XX INFO Wrote SLURM script to /shared/outputs/qwen30b-swe/rl.sbatch [rl.py::534] +XXX:XX:XX INFO Submitting: sbatch /shared/outputs/qwen30b-swe/rl.sbatch [rl.py::540] +XXX:XX:XX SUCCESS Submitted batch job YYYY + +Logs: + Trainer: tail -F /shared/outputs/qwen30b-swe/logs/trainer.log + Orchestrator: tail -F /shared/outputs/qwen30b-swe/logs/orchestrator.log + Inference: tail -F /shared/outputs/qwen30b-swe/logs/inference.log + Envs: tail -F /shared/outputs/qwen30b-swe/logs/envs/*/*/*.log + Train: tail -F /shared/outputs/qwen30b-swe/logs/envs/train/*/*.log + swe: tail -F /shared/outputs/qwen30b-swe/logs/envs/train/swe/*.log +``` + + diff --git a/examples/qwen30b_swe/rl.toml b/examples/qwen30b_swe/rl.toml new file mode 100644 index 00000000000..ba5902b71b8 --- /dev/null +++ b/examples/qwen30b_swe/rl.toml @@ -0,0 +1,68 @@ +seq_len = 131072 + +[log] +level = "debug" + +[model] +name = "Qwen/Qwen3-30B-A3B-Thinking-2507" + +[weight_broadcast] +type = "nccl" + +[deployment] +type = "multi_node" +num_train_nodes = 2 +num_infer_nodes = 2 + +[slurm] +job_name = "swe-qwen30b" + +[ckpt] +interval = 100 + +[wandb] +project = "qwen30b-swe" + +[trainer.model] +impl = "custom" +attn = "flash_attention_3" +ep = 8 +cp = 2 + +[trainer.model.compile] + +[trainer.model.ac_offloading] +max_inflight_activations = 5 + +[trainer.model.ac] +freq = 1 + +[trainer.optim] +type = "adamw" +lr = 1e-6 + +[orchestrator] +batch_size = 512 +oversampling_factor = 2 +max_off_policy_steps = 16 + +[[orchestrator.train.env]] +id = "mini-swe-agent-plus" +name = "swe" +args = { max_turns = 200, cpu_cores = 2, memory_gb = 4, disk_size_gb = 4, labels = ["mini-swe-agent-plus"], total_timeout_minutes = 720, sandbox_client_max_workers = 256, max_command_timeouts = 3, sandbox_command_timeout = 30} + +[orchestrator.eval] +interval = 25 + +[[orchestrator.eval.env]] +id = "mini-swe-agent-plus" +name = "swe-bench-verified-quick" +args = { dataset_name="PrimeIntellect/SWE-Bench-Verified-Quick", max_turns = 200, cpu_cores = 2, memory_gb = 4, disk_size_gb = 4, labels = ["mini-swe-agent-plus"], total_timeout_minutes = 720, sandbox_client_max_workers = 256, max_command_timeouts = 3, sandbox_command_timeout = 30} + +[orchestrator.buffer] +easy_threshold = 1.0 +hard_threshold = 0.0 +seed = 42 + +[inference.parallel] +tp = 8 \ No newline at end of file diff --git a/examples/reverse_text/README.md b/examples/reverse_text/README.md index 4dde0352d0b..3890cd2b886 100644 --- a/examples/reverse_text/README.md +++ b/examples/reverse_text/README.md @@ -36,7 +36,7 @@ This is of course just a quick vibe check and no full-fledged evaluation, but we We will fine-tune `PrimeIntellect/Qwen3-0.6B` ([HF](https://huggingface.co/PrimeIntellect/Qwen3-0.6B)), which is a clone of `Qwen/Qwen3-0.6B` ([HF](https://huggingface.co/Qwen/Qwen3-0.6B)) with a chat template suitable for multi-turn RL, on `willcb/R1-reverse-wikipedia-paragraphs-v1-1000` ([HF](https://huggingface.co/datasets/willcb/R1-reverse-wikipedia-paragraphs-v1-1000)) which contains 1K examples of reversals of small paragraphs. -![SFT](sft/wandb.png) +![SFT](sft.png) *Check out the logs of the SFT run on [W&B](https://wandb.ai/primeintellect/examples?nw=s3p14m48jod).* To train on a single GPU, run @@ -72,7 +72,7 @@ We have uploaded the final model as [`PrimeIntellect/Qwen3-0.6B-Reverse-Text-SFT For the RL we will only do 20 steps at 8x16 rollouts, for a total batch size of 128 and sequence length 128. Because of the small context, training should be extremely quick. -![RL](rl/wandb.png) +![RL](rl.png) *Check out the logs of the RL run on [W&B](https://wandb.ai/primeintellect/examples?nw=yxjwjc556do).* ```bash diff --git a/examples/reverse_text/rl/wandb.png b/examples/reverse_text/rl.png similarity index 100% rename from examples/reverse_text/rl/wandb.png rename to examples/reverse_text/rl.png diff --git a/examples/reverse_text/rl.toml b/examples/reverse_text/rl.toml index c27fcbccef9..d19337b86c6 100644 --- a/examples/reverse_text/rl.toml +++ b/examples/reverse_text/rl.toml @@ -12,13 +12,15 @@ name = "reverse-text" batch_size = 128 rollouts_per_example = 16 -[orchestrator.sampling] -max_tokens = 128 +[orchestrator.train.sampling] +max_completion_tokens = 128 -[[orchestrator.env]] +[[orchestrator.train.env]] id = "reverse-text" [trainer.optim] lr = 3e-6 +[ckpt] # Checkpoint at the end of training + [inference] \ No newline at end of file diff --git a/examples/reverse_text/rl/infer.toml b/examples/reverse_text/rl/infer.toml deleted file mode 100644 index 7bc3175599b..00000000000 --- a/examples/reverse_text/rl/infer.toml +++ /dev/null @@ -1,2 +0,0 @@ -[model] -name = "PrimeIntellect/Qwen3-0.6B-Reverse-Text-SFT" \ No newline at end of file diff --git a/examples/reverse_text/rl/orch.toml b/examples/reverse_text/rl/orch.toml deleted file mode 100644 index 95d70b9cae2..00000000000 --- a/examples/reverse_text/rl/orch.toml +++ /dev/null @@ -1,13 +0,0 @@ -batch_size = 128 -rollouts_per_example = 16 -seq_len = 2048 -max_steps = 20 - -[model] -name = "PrimeIntellect/Qwen3-0.6B-Reverse-Text-SFT" - -[sampling] -max_tokens = 128 - -[[env]] -id = "reverse-text" diff --git a/examples/reverse_text/rl/train.toml b/examples/reverse_text/rl/train.toml deleted file mode 100644 index fc6652748f9..00000000000 --- a/examples/reverse_text/rl/train.toml +++ /dev/null @@ -1,7 +0,0 @@ -max_steps = 20 - -[model] -name = "PrimeIntellect/Qwen3-0.6B-Reverse-Text-SFT" - -[optim] -lr = 3e-6 \ No newline at end of file diff --git a/examples/reverse_text/sft/wandb.png b/examples/reverse_text/sft.png similarity index 100% rename from examples/reverse_text/sft/wandb.png rename to examples/reverse_text/sft.png diff --git a/examples/reverse_text/sft/train.toml b/examples/reverse_text/sft/train.toml deleted file mode 100644 index 9dd0d2582d3..00000000000 --- a/examples/reverse_text/sft/train.toml +++ /dev/null @@ -1,12 +0,0 @@ -max_steps = 100 - -[model] -name = "PrimeIntellect/Qwen3-0.6B" - -[data] -name = "willcb/R1-reverse-wikipedia-paragraphs-v1-1000" -seq_len = 4096 -batch_size = 32 - -[optim] -lr = 2e-5 \ No newline at end of file diff --git a/examples/reverse_text/slurm_rl.toml b/examples/reverse_text/slurm_rl.toml new file mode 100644 index 00000000000..cd3cce5ece8 --- /dev/null +++ b/examples/reverse_text/slurm_rl.toml @@ -0,0 +1,5 @@ +# Usage: uv run rl @ examples/reverse_text/rl.toml @ examples/reverse_text/slurm_rl.toml +output_dir = "outputs/reverse-text-rl" + +[slurm] +job_name = "reverse-text-rl" \ No newline at end of file diff --git a/examples/reverse_text/slurm_sft.toml b/examples/reverse_text/slurm_sft.toml new file mode 100644 index 00000000000..1c6056edcae --- /dev/null +++ b/examples/reverse_text/slurm_sft.toml @@ -0,0 +1,5 @@ +# Usage: uv run sft @ examples/reverse_text/sft.toml @ examples/reverse_text/slurm_sft.toml +output_dir = "outputs/reverse-text-sft" + +[slurm] +job_name = "reverse-text-sft" \ No newline at end of file diff --git a/examples/wiki_search/rl.toml b/examples/wiki_search/rl.toml index 5827b8d66e1..6abbb3d8156 100644 --- a/examples/wiki_search/rl.toml +++ b/examples/wiki_search/rl.toml @@ -1,9 +1,10 @@ -inference_gpu_ids = [0,1,2,3,4,5] -trainer_gpu_ids = [6,7] - max_steps = 200 seq_len = 4096 +[deployment] +num_train_gpus = 2 +num_infer_gpus = 6 + [model] name = "Qwen/Qwen3-4B-Instruct-2507" @@ -36,15 +37,16 @@ oversampling_factor = 2.0 [orchestrator.model.lora] name = "qwen3-4b-wiki-search" -[orchestrator.sampling] -max_tokens = 512 +[orchestrator.train.sampling] +max_completion_tokens = 512 [orchestrator.buffer] online_difficulty_filtering = true -[[orchestrator.env]] +[[orchestrator.train.env]] id = "primeintellect/wiki-search" +[ckpt] # Checkpoint at the end of training + [inference.model] -enable_auto_tool_choice = true tool_call_parser = "hermes" diff --git a/examples/wiki_search/slurm_rl.toml b/examples/wiki_search/slurm_rl.toml new file mode 100644 index 00000000000..2da75e17440 --- /dev/null +++ b/examples/wiki_search/slurm_rl.toml @@ -0,0 +1,5 @@ +# Usage: uv run rl @ examples/wiki_search/rl.toml @ examples/wiki_search/slurm_rl.toml +output_dir = "outputs/wiki-search" + +[slurm] +job_name = "wiki-search" \ No newline at end of file diff --git a/examples/wordle/rl.toml b/examples/wordle/rl.toml index 6dacdc27a92..123f4eb90a6 100644 --- a/examples/wordle/rl.toml +++ b/examples/wordle/rl.toml @@ -1,9 +1,10 @@ -inference_gpu_ids = [0,1,2,3,4,5] -trainer_gpu_ids = [6,7] - max_steps = 200 seq_len = 8192 +[deployment] +num_train_gpus = 2 +num_infer_gpus = 6 + [wandb] project = "wordle" name = "wordle" @@ -17,12 +18,12 @@ name = "PrimeIntellect/Qwen3-1.7B-Wordle-SFT" batch_size = 1024 rollouts_per_example = 16 -[[orchestrator.env]] +[[orchestrator.train.env]] id = "primeintellect/wordle" name = "wordle" -[orchestrator.sampling] -max_tokens = 1024 +[orchestrator.train.sampling] +max_completion_tokens = 1024 [trainer] # Default trainer config diff --git a/examples/wordle/slurm_rl.toml b/examples/wordle/slurm_rl.toml new file mode 100644 index 00000000000..85b66a52ce7 --- /dev/null +++ b/examples/wordle/slurm_rl.toml @@ -0,0 +1,5 @@ +# Usage: uv run rl @ examples/wordle/rl.toml @ examples/wordle/slurm_rl.toml +output_dir = "outputs/wordle-rl" + +[slurm] +job_name = "wordle-rl" \ No newline at end of file diff --git a/examples/wordle/slurm_sft.toml b/examples/wordle/slurm_sft.toml new file mode 100644 index 00000000000..1a5ace3e292 --- /dev/null +++ b/examples/wordle/slurm_sft.toml @@ -0,0 +1,5 @@ +# Usage: uv run sft @ examples/wordle/sft.toml @ examples/wordle/slurm_sft.toml +output_dir = "outputs/wordle-sft" + +[slurm] +job_name = "wordle-sft" \ No newline at end of file diff --git a/k8s/prime-rl/templates/deployment.yaml b/k8s/prime-rl/templates/deployment.yaml index 013d02ea7d0..1ed16b9f5e6 100644 --- a/k8s/prime-rl/templates/deployment.yaml +++ b/k8s/prime-rl/templates/deployment.yaml @@ -27,7 +27,7 @@ spec: {{- toYaml . | nindent 8 }} {{- end }} containers: - - name: prime-rl + - name: prime-rl-orchestrator image: "{{ .Values.image.repository }}:{{ .Values.image.tag }}" imagePullPolicy: {{ .Values.image.pullPolicy }} {{- if .Values.orchestrator.autoStart }} @@ -117,6 +117,7 @@ metadata: spec: serviceName: {{ .Release.Name }}-inference-headless replicas: {{ .Values.inference.replicas }} + podManagementPolicy: Parallel selector: matchLabels: {{- include "prime-rl.componentLabels" . | nindent 6 }} @@ -132,7 +133,7 @@ spec: runtimeClassName: {{ .Values.inference.runtimeClassName }} {{- end }} containers: - - name: prime-rl + - name: prime-rl-inference image: "{{ .Values.image.repository }}:{{ .Values.image.tag }}" imagePullPolicy: {{ .Values.image.pullPolicy }} {{- if .Values.inference.autoStart }} @@ -196,6 +197,29 @@ spec: memory: {{ .Values.inference.resources.limits.memory }} cpu: {{ .Values.inference.resources.limits.cpu }} {{- end }} + {{- if and .Values.inference.probes .Values.inference.probes.enabled }} + startupProbe: + httpGet: + path: /health + port: {{ .Values.inference.service.port }} + periodSeconds: {{ .Values.inference.probes.startup.periodSeconds }} + failureThreshold: {{ .Values.inference.probes.startup.failureThreshold }} + timeoutSeconds: {{ .Values.inference.probes.startup.timeoutSeconds }} + livenessProbe: + httpGet: + path: /liveness + port: {{ .Values.inference.service.port }} + periodSeconds: {{ .Values.inference.probes.liveness.periodSeconds }} + failureThreshold: {{ .Values.inference.probes.liveness.failureThreshold }} + timeoutSeconds: {{ .Values.inference.probes.liveness.timeoutSeconds }} + readinessProbe: + httpGet: + path: /health + port: {{ .Values.inference.service.port }} + periodSeconds: {{ .Values.inference.probes.readiness.periodSeconds }} + failureThreshold: {{ .Values.inference.probes.readiness.failureThreshold }} + timeoutSeconds: {{ .Values.inference.probes.readiness.timeoutSeconds }} + {{- end }} {{- if .Values.storage.enabled }} volumeMounts: - name: shared-data @@ -222,6 +246,7 @@ metadata: spec: serviceName: {{ .Release.Name }}-trainer-headless replicas: {{ .Values.trainer.replicas }} + podManagementPolicy: Parallel selector: matchLabels: {{- include "prime-rl.componentLabels" . | nindent 6 }} @@ -237,7 +262,7 @@ spec: runtimeClassName: {{ .Values.trainer.runtimeClassName }} {{- end }} containers: - - name: prime-rl + - name: prime-rl-trainer image: "{{ .Values.image.repository }}:{{ .Values.image.tag }}" imagePullPolicy: {{ .Values.image.pullPolicy }} {{- if .Values.trainer.autoStart }} @@ -318,6 +343,29 @@ spec: memory: {{ .Values.trainer.resources.limits.memory }} cpu: {{ .Values.trainer.resources.limits.cpu }} {{- end }} + {{- if and .Values.trainer.probes .Values.trainer.probes.enabled }} + startupProbe: + httpGet: + path: /health + port: {{ .Values.trainer.service.port }} + periodSeconds: {{ .Values.trainer.probes.startup.periodSeconds }} + failureThreshold: {{ .Values.trainer.probes.startup.failureThreshold }} + timeoutSeconds: {{ .Values.trainer.probes.startup.timeoutSeconds }} + livenessProbe: + httpGet: + path: /health + port: {{ .Values.trainer.service.port }} + periodSeconds: {{ .Values.trainer.probes.liveness.periodSeconds }} + failureThreshold: {{ .Values.trainer.probes.liveness.failureThreshold }} + timeoutSeconds: {{ .Values.trainer.probes.liveness.timeoutSeconds }} + readinessProbe: + httpGet: + path: /health + port: {{ .Values.trainer.service.port }} + periodSeconds: {{ .Values.trainer.probes.readiness.periodSeconds }} + failureThreshold: {{ .Values.trainer.probes.readiness.failureThreshold }} + timeoutSeconds: {{ .Values.trainer.probes.readiness.timeoutSeconds }} + {{- end }} {{- if .Values.storage.enabled }} volumeMounts: - name: shared-data diff --git a/k8s/prime-rl/values.yaml b/k8s/prime-rl/values.yaml index 5600d81736d..7fc3b489899 100644 --- a/k8s/prime-rl/values.yaml +++ b/k8s/prime-rl/values.yaml @@ -75,6 +75,23 @@ inference: runtimeClassName: nvidia + # Health probes for inference server (/health for startup/readiness, /liveness for engine liveness) + probes: + enabled: false + # Startup probe: allows long model loading time (up to 45 min by default) + startup: + periodSeconds: 10 + failureThreshold: 270 # 270 * 10s = 45 min max startup time + timeoutSeconds: 5 + liveness: + periodSeconds: 10 + failureThreshold: 3 + timeoutSeconds: 10 + readiness: + periodSeconds: 5 + failureThreshold: 2 + timeoutSeconds: 5 + # Trainer component trainer: enabled: true @@ -106,6 +123,23 @@ trainer: runtimeClassName: nvidia + # Health probes for trainer (requires metrics_server config) + probes: + enabled: false + # Startup probe: allows time for model loading + startup: + periodSeconds: 10 + failureThreshold: 60 # 60 * 10s = 10 min max startup time + timeoutSeconds: 30 + liveness: + periodSeconds: 30 + failureThreshold: 6 # 6 * 30s = 3 min before restart + timeoutSeconds: 30 + readiness: + periodSeconds: 10 + failureThreshold: 3 + timeoutSeconds: 30 + # Additional configuration config: # Example name (used for labeling) diff --git a/packages/prime-rl-configs/README.md b/packages/prime-rl-configs/README.md new file mode 100644 index 00000000000..be6a9563a66 --- /dev/null +++ b/packages/prime-rl-configs/README.md @@ -0,0 +1,26 @@ +# prime-rl-configs + +Slim config schema for [`prime-rl`](https://github.com/PrimeIntellect-ai/prime-rl), with no GPU or ML deps. + +`pip install prime-rl-configs` gives you `prime_rl.configs.*` (RL/SFT/inference/orchestrator/trainer/env-server schemas) without pulling in `torch`, `vllm`, `transformers`, `wandb`, etc. The full training stack lives in `prime-rl`, which depends on this package. + +## Install + +```sh +pip install git+https://github.com/PrimeIntellect-ai/prime-rl.git#subdirectory=packages/prime-rl-configs +``` + +## Usage + +The pip *distribution name* (`prime-rl-configs`) and the *import path* (`prime_rl.configs.*`) are different on purpose: this package contributes submodules to the shared `prime_rl` namespace. + +```python +from pydantic_config import cli +from prime_rl.configs.rl import RLConfig + +config = cli(RLConfig, args=["@", "path/to/rl.toml"]) +``` + +Other config classes live alongside `RLConfig` under `prime_rl.configs.*` (`sft`, `inference`, `orchestrator`, `trainer`, `env_server`). + +`import prime_rl` on its own succeeds but is empty — it's a [PEP 420](https://peps.python.org/pep-0420/) namespace package with no top-level attributes. Always import a submodule. diff --git a/packages/prime-rl-configs/pyproject.toml b/packages/prime-rl-configs/pyproject.toml new file mode 100644 index 00000000000..1d8f17f6f60 --- /dev/null +++ b/packages/prime-rl-configs/pyproject.toml @@ -0,0 +1,25 @@ +[project] +name = "prime-rl-configs" +version = "0.4.0" +description = "Slim config schema for prime-rl (no GPU/ML deps)." +readme = "README.md" +requires-python = "~=3.12.0" +dependencies = [ + "pydantic>=1.10.13", + # The PyPI package named `pydantic-config` is a different project that + # doesn't expose `BaseConfig` / `cli`. Pin to samsja's fork so plain + # `pip install prime-rl-configs` (no uv workspace context) resolves it. + "pydantic-config @ git+https://github.com/samsja/pydantic_config.git", + "tomli>=2.2.1", + "tomli-w>=1.2.0", +] + +[build-system] +requires = ["hatchling"] +build-backend = "hatchling.build" + +[tool.hatch.build.targets.wheel] +packages = ["src/prime_rl"] + +[tool.hatch.metadata] +allow-direct-references = true diff --git a/src/prime_rl/__init__.py b/packages/prime-rl-configs/src/prime_rl/configs/__init__.py similarity index 100% rename from src/prime_rl/__init__.py rename to packages/prime-rl-configs/src/prime_rl/configs/__init__.py diff --git a/packages/prime-rl-configs/src/prime_rl/configs/env_server.py b/packages/prime-rl-configs/src/prime_rl/configs/env_server.py new file mode 100644 index 00000000000..1e4e21614e2 --- /dev/null +++ b/packages/prime-rl-configs/src/prime_rl/configs/env_server.py @@ -0,0 +1,29 @@ +from pathlib import Path +from typing import Annotated + +from pydantic import Field, model_validator + +from prime_rl.configs.orchestrator import EnvConfig +from prime_rl.configs.shared import LogConfig +from prime_rl.utils.config import BaseConfig + + +class EnvServerConfig(BaseConfig): + """Configures an environment server.""" + + env: EnvConfig = EnvConfig() + log: LogConfig = LogConfig() + env_install_prerelease: bool = False + + output_dir: Annotated[ + Path, + Field( + description="Directory to write outputs to. Will be populated with checkpoints, weights, rollouts and logs as subdirectories. Should be set to a persistent directory with enough disk space. This value should be distinct across experiments running on a single node. See the README for more details." + ), + ] = Path("outputs") + + @model_validator(mode="after") + def validate_num_workers(self): + if self.env.num_workers == "auto": + self.env.num_workers = 1 + return self diff --git a/packages/prime-rl-configs/src/prime_rl/configs/inference.py b/packages/prime-rl-configs/src/prime_rl/configs/inference.py new file mode 100644 index 00000000000..805481de767 --- /dev/null +++ b/packages/prime-rl-configs/src/prime_rl/configs/inference.py @@ -0,0 +1,605 @@ +from argparse import Namespace +from pathlib import Path +from typing import Annotated, Any, Literal, TypeAlias + +from pydantic import BaseModel, ConfigDict, Field, model_validator +from pydantic_config import BaseConfig + +from prime_rl.configs.shared import BaseModelConfig, SlurmConfig +from prime_rl.utils.config import find_package_resource, rgetattr, rsetattr + +# TODO: Set thinking/ solution budget + + +class ServerConfig(BaseConfig): + """Configures the inference server.""" + + host: Annotated[str | None, Field(description="The host to bind to.")] = None + port: Annotated[int, Field(description="The port to bind to.")] = 8000 + liveness_timeout_seconds: Annotated[ + float, + Field( + gt=0, + description=( + "Timeout in seconds for the /liveness endpoint's internal vLLM worker RPC. " + "If Kubernetes liveness probes are enabled, keep the probe timeoutSeconds at least this high." + ), + ), + ] = 30.0 + + +class ParallelConfig(BaseConfig): + """Configures multi-node and multi-GPU setups through different types of parallelism (TP, DP, PP).""" + + tp: Annotated[ + int, + Field( + description="The tensor parallel size. It is passed to vLLM as `--tensor-parallel-size`", + ), + ] = 1 + + dp: Annotated[ + int, + Field( + ge=1, + description="The data parallel size. It is passed to vLLM as `--data-parallel-size`", + ), + ] = 1 + + def __str__(self) -> str: + return f"tp={self.tp} dp={self.dp}" + + +class ModelConfig(BaseModelConfig): + """Configures the inference model. Most arguments are passed directly to the vLLM LLM class (https://docs.vllm.ai/en/latest/api/vllm.LLM.html).""" + + dtype: Annotated[ + Literal["auto", "float16", "bfloat16", "float32"], + Field( + description="Data type for model weights and activations. If 'auto' will use FP16 precision for FP32 and FP16 models, and BF16 precision for BF16 models. Passed to vLLM as `--dtype`", + ), + ] = "auto" + + max_model_len: Annotated[ + int | None, + Field( + description="Maximum model context length. If None, will use the maximum context length from model config. Passed to vLLM as `--max-model-len`", + ), + ] = None + + enforce_eager: Annotated[ + bool, + Field( + description="Whether to enforce eager mode. If False, will use PyTorch eager and cuda graphs in hybrid for maximal performance. Passed to vLLM as `--enforce-eager`", + ), + ] = False + + trust_remote_code: Annotated[ + bool, + Field( + description="Whether to trust remote code. Passed to vLLM engine init", + ), + ] = False + + chat_template: Annotated[ + str | None, + Field( + description="Chat template to use. Can be a Jinja2 template string or a path to a template file. " + "Passed to vLLM as `--chat-template`. If None, uses the model's default.", + ), + ] = None + + tool_call_parser: Annotated[ + str | None, + Field( + description="The tool call parser to use. Passed to vLLM as `--tool-call-parser`. " + 'Set to "auto" to infer from the model name.', + ), + ] = "auto" + + reasoning_parser: Annotated[ + str | None, + Field( + description="Parser for extracting reasoning content from model outputs. Passed to vLLM as `--reasoning-parser`. Setting this enables reasoning mode.", + ), + ] = None + + rope_scaling: Annotated[ + dict[str, Any] | str | None, + Field( + description='RoPE scaling configuration as a dict. For YaRN, use: {rope_type="yarn", factor=4.0, original_max_position_embeddings=32768} or. Passed to vLLM as `--rope-scaling`.', + ), + ] = None + + +class WeightBroadcastConfig(BaseConfig): + """Configures weight broadcast settings.""" + + type: Annotated[Literal["nccl", "filesystem"], Field(description="The type of weight broadcast to use.")] = ( + "filesystem" + ) + + +class KVCacheOffloadConfig(BaseModel): + """CPU KV cache offloading for vLLM inference workers.""" + + model_config = ConfigDict(extra="forbid") + + cpu_bytes: Annotated[int, Field(gt=0, description="CPU bytes available for KV cache offloading per worker.")] = ( + 1_000_000_000 + ) + + +# Valid vLLM max_lora_rank values (from vllm/config/lora.py) +# TODO: on newer vLLM, can import via `get_args(vllm.config.lora.MaxLoRARanks)` +VALID_VLLM_LORA_RANKS = (8, 16, 32, 64, 128, 256, 320, 512) + +# vLLM all2all backend options for expert-parallel deployments. +All2AllBackend = Literal[ + "allgather_reducescatter", + "deepep_high_throughput", + "deepep_low_latency", + "flashinfer_nvlink_one_sided", + "flashinfer_nvlink_two_sided", +] + + +class BaseInferenceDeploymentConfig(BaseModel): + """Base deployment config for inference.""" + + model_config = ConfigDict(extra="forbid") + + gpus_per_node: Annotated[int, Field(description="Number of GPUs per node.")] = 8 + + +class SingleNodeInferenceDeploymentConfig(BaseInferenceDeploymentConfig): + """Configures a single-node inference deployment.""" + + type: Literal["single_node"] = "single_node" + + +class MultiNodeInferenceDeploymentConfig(BaseInferenceDeploymentConfig): + """Configures a multi-node inference deployment. Each node runs an independent vLLM replica.""" + + type: Literal["multi_node"] = "multi_node" + + num_nodes: Annotated[int, Field(ge=1, description="Number of inference nodes.")] = 2 + + router_port: Annotated[int, Field(description="Port for the vllm-router.")] = 8000 + backend_port: Annotated[int, Field(description="Port for vLLM backend instances.")] = 8100 + router_policy: Annotated[ + str, Field(description="Routing policy for the vllm-router (e.g. 'consistent_hash', 'round_robin').") + ] = "consistent_hash" + + +class DisaggregatedInferenceDeploymentConfig(BaseInferenceDeploymentConfig): + """Configures a disaggregated prefill/decode inference deployment. + + Each inference replica is split into separate prefill and decode node groups. + Requires NIXL for KV transfer and a vllm-router for request routing. + + Multi-replica support: set ``num_prefill_replicas`` / ``num_decode_replicas`` + to run multiple independent vLLM instances within the prefill / decode node + groups. For example, ``num_prefill_nodes=4, num_prefill_replicas=2`` creates + two prefill vLLM instances each spanning 2 nodes (EP16 with 8 GPUs/node). + """ + + type: Literal["disaggregated"] = "disaggregated" + + num_prefill_nodes: Annotated[int, Field(ge=1, description="Total number of prefill nodes.")] = 1 + num_decode_nodes: Annotated[int, Field(ge=1, description="Total number of decode nodes.")] = 1 + + num_prefill_replicas: Annotated[ + int, + Field( + ge=1, + description="Number of independent prefill vLLM instances. Must evenly divide num_prefill_nodes.", + ), + ] = 1 + num_decode_replicas: Annotated[ + int, + Field( + ge=1, + description="Number of independent decode vLLM instances. Must evenly divide num_decode_nodes.", + ), + ] = 1 + + router_port: Annotated[int, Field(description="Port for the vllm-router on each replica.")] = 8000 + prefill_port: Annotated[int, Field(description="Port for prefill vLLM instances.")] = 8100 + decode_port: Annotated[int, Field(description="Port for decode vLLM instances.")] = 8200 + router_policy: Annotated[ + str, Field(description="Routing policy for the vllm-router (e.g. 'consistent_hash', 'round_robin').") + ] = "consistent_hash" + + prefill_env_overrides: Annotated[ + dict[str, str], + Field(description="Extra environment variables exported only on prefill nodes."), + ] = {} + decode_env_overrides: Annotated[ + dict[str, str], + Field(description="Extra environment variables exported only on decode nodes."), + ] = {} + + @property + def num_nodes(self) -> int: + return self.num_prefill_nodes + self.num_decode_nodes + + @model_validator(mode="after") + def validate_replicas_divide_nodes(self): + if self.num_prefill_nodes % self.num_prefill_replicas != 0: + raise ValueError( + f"num_prefill_replicas ({self.num_prefill_replicas}) must evenly divide " + f"num_prefill_nodes ({self.num_prefill_nodes})" + ) + if self.num_decode_nodes % self.num_decode_replicas != 0: + raise ValueError( + f"num_decode_replicas ({self.num_decode_replicas}) must evenly divide " + f"num_decode_nodes ({self.num_decode_nodes})" + ) + return self + + +InferenceDeploymentConfig: TypeAlias = Annotated[ + SingleNodeInferenceDeploymentConfig | MultiNodeInferenceDeploymentConfig | DisaggregatedInferenceDeploymentConfig, + Field(discriminator="type"), +] + + +class InferenceExperimentalConfig(BaseConfig): + """Experimental features for inference.""" + + +class InferenceConfig(BaseConfig): + """Configures inference.""" + + # The server configuration + server: ServerConfig = ServerConfig() + + # The model configuration + model: ModelConfig = Field(default_factory=ModelConfig) + + # The parallel configuration + parallel: ParallelConfig = ParallelConfig() + + enable_lora: Annotated[ + bool, + Field( + description="Whether to enable LORA. Passed to vLLM as `--enable-lora`", + ), + ] = False + + max_loras: Annotated[ + int, + Field( + description="The maximum number of LoRAs to use. Passed to vLLM as `--max-loras`", + ), + ] = 8 + + # TODO: The default value is very high because our areal impl for lora isn't ideal + # We add a lora with the same name instead of changing weights inplace + # Because we dont cancel requests that are past max_async, these requests could be using a LoRA that gets unloaded which will crash the inference server + max_cpu_loras: Annotated[ + int, + Field( + description="The maximum number of LoRAs to use on CPU. Passed to vLLM as `--max-cpu-loras`", + ), + ] = 100 + + max_lora_rank: Annotated[ + int | None, + Field( + description="The maximum LoRA rank to use. Passed to vLLM as `--max-lora-rank`", + ), + ] = None + + lora_target_modules: Annotated[ + list[str] | None, + Field( + description="The target modules for LoRA. Passed to vLLM as `--lora-target-modules`.", + ), + ] = None + + enable_prefix_caching: Annotated[ + bool | None, + Field( + description="Whether to enable prefix caching. Passed to vLLM as `--enable-prefix-caching`", + ), + ] = None + + gpu_memory_utilization: Annotated[ + float, + Field( + description="The GPU memory utilization to use. Passed to vLLM as `--gpu-memory-utilization`", + ), + ] = 0.9 + + api_server_count: Annotated[ + int, + Field( + ge=0, + description="The number of API servers to use. Passed to vLLM as `--api-server-count`. Set to 0 for headless mode.", + ), + ] = 1 + + data_parallel_size_local: Annotated[ + int | None, + Field( + ge=1, + description="Number of data parallel replicas to run on this node. Passed to vLLM as `--data-parallel-size-local`.", + ), + ] = None + + data_parallel_rpc_port: Annotated[ + int, + Field( + ge=1, + le=65535, + description="RPC port for data parallel communication. Passed to vLLM as `--data-parallel-rpc-port`.", + ), + ] = 13345 + + seed: Annotated[ + int, + Field( + description="Seed the inference components. Passed to vLLM as `--seed`", + ), + ] = 0 + + enable_expert_parallel: Annotated[ + bool, + Field( + description="Enable expert parallelism for MoE models. Passed to vLLM as `--enable-expert-parallel`.", + ), + ] = False + + all2all_backend: Annotated[ + All2AllBackend, + Field( + description="All-to-all backend for expert parallel communication. Passed to vLLM as `--all2all-backend`.", + ), + ] = "allgather_reducescatter" + + enable_eplb: Annotated[ + bool, + Field( + description="Enable expert parallel load balancer (EPLB). Passed to vLLM as `--enable-eplb`.", + ), + ] = False + + enable_dbo: Annotated[ + bool, + Field( + description="Enable dual batch overlap (DBO). Passed to vLLM as `--enable-dbo`.", + ), + ] = False + + use_deep_gemm: Annotated[ + bool, + Field( + description="Force DeepGEMM FP8 kernels via VLLM_USE_DEEP_GEMM=1. Only works with per-tensor FP8 quantization (e.g. GLM-5-FP8).", + ), + ] = False + + weight_broadcast: Annotated[WeightBroadcastConfig, Field(description="The weight broadcast config.")] = ( + WeightBroadcastConfig() + ) + + kv_cache_offload: Annotated[ + KVCacheOffloadConfig | None, + Field( + description=( + "CPU KV cache offload config for inference workers. Standard inference uses vLLM's " + "OffloadingConnector. Disaggregated P/D deployments combine it with NIXL through " + "MultiConnector in the SLURM launcher." + ), + ), + ] = None + + enable_return_routed_experts: Annotated[ + bool, + Field( + description="Whether to enable return routed experts. Passed to vLLM as `--enable-return-routed-experts`", + ), + ] = False + + enable_fp32_lm_head: Annotated[ + bool, + Field( + description="Run the lm_head projection in fp32 via a native bf16xbf16 -> fp32 GEMM (`torch.mm` with `out_dtype=torch.float32`). Stabilizes logprob precision under FP8/bf16 inference, matching SGLang's `--enable-fp32-lm-head`. Implemented as a monkey-patch over vLLM's LogitsProcessor, activated by setting `additional_config[\"fp32_lm_head\"] = True` on the vLLM config.", + ), + ] = False + + vllm_extra: Annotated[ + dict[str, Any], + Field( + description="Extra arguments to pass to vLLM. These are applied as attributes on the vLLM namespace after config translation.", + ), + ] = {} + + # Launcher-only fields + + deployment: Annotated[ + InferenceDeploymentConfig, + Field( + description="Deployment configuration for inference.", + ), + ] = SingleNodeInferenceDeploymentConfig() + + slurm: Annotated[ + SlurmConfig | None, + Field( + description="SLURM configuration. If set, the run will be submitted as a SLURM job instead of running locally.", + ), + ] = None + + output_dir: Annotated[Path, Field(description="Directory for SLURM logs and generated scripts.")] = Path("outputs") + + dry_run: Annotated[bool, Field(description="Only validate and dump resolved configs and exit early.")] = False + + experimental: Annotated[ + InferenceExperimentalConfig, + Field(description="Experimental features for inference."), + ] = InferenceExperimentalConfig() + + @model_validator(mode="after") + def validate_multi_node_requires_slurm(self): + if self.deployment.type == "multi_node" and self.slurm is None: + raise ValueError("Must use SLURM for multi-node deployment.") + return self + + @model_validator(mode="after") + def auto_setup_kv_cache_offload(self): + if self.kv_cache_offload is not None: + if self.enable_prefix_caching is False: + raise ValueError("KV cache offloading requires inference.enable_prefix_caching to be true.") + if "enable_prefix_caching" not in self.model_fields_set: + self.enable_prefix_caching = True + + return self + + @model_validator(mode="after") + def auto_setup_disaggregated(self): + """Auto-configure inference for disaggregated P/D: enable EP and compute DP.""" + if self.deployment.type == "disaggregated": + if "enable_expert_parallel" not in self.model_fields_set: + self.enable_expert_parallel = True + if "enable_eplb" not in self.model_fields_set: + self.enable_eplb = False + gpus_per_node = self.deployment.gpus_per_node + tp = self.parallel.tp + dp_per_node = gpus_per_node // tp + if self.data_parallel_size_local is None: + self.data_parallel_size_local = dp_per_node + if self.parallel.dp == 1: + self.parallel.dp = dp_per_node + if self.api_server_count == 1: + self.api_server_count = dp_per_node + return self + + @model_validator(mode="after") + def auto_setup_slurm_template(self): + if self.slurm is not None and self.slurm.template_path is None: + templates_dir = find_package_resource("templates") + if templates_dir is not None: + self.slurm.template_path = templates_dir / "inference.sbatch.j2" + return self + + @model_validator(mode="after") + def auto_setup_max_lora_rank(self): + """Auto-setup max_lora_rank by rounding up to the nearest valid vLLM value. + + vLLM only accepts specific values for max_lora_rank: (1, 8, 16, 32, 64, 128, 256, 320, 512). + This validator ensures that any configured rank is rounded up to the minimum valid value + that can serve adapters of the requested rank. + """ + if self.max_lora_rank is not None: + original_rank = self.max_lora_rank + for valid_rank in VALID_VLLM_LORA_RANKS: + if valid_rank >= self.max_lora_rank: + self.max_lora_rank = valid_rank + break + else: + raise ValueError(f"max_lora_rank={original_rank} exceeds vLLM maximum of {VALID_VLLM_LORA_RANKS[-1]}") + return self + + @model_validator(mode="after") + def auto_setup_api_server_count(self): + """ + Ensures that we have at least as many API servers as data parallel + size. Unless LoRA is enabled, in which case only one API server is + supported (vLLM limitation). + """ + if self.vllm_extra.get("headless", False): + self.api_server_count = 0 + return self + + if "api_server_count" not in self.model_fields_set: + min_api_server_count = self.data_parallel_size_local or self.parallel.dp + if self.api_server_count < min_api_server_count: + self.api_server_count = min_api_server_count + + if self.enable_lora: + self.api_server_count = 1 # LoRA requires only one API server + return self + + def to_vllm(self) -> Namespace: + """Convert InferenceConfig to vLLM-compatible Namespace.""" + namespace = Namespace() + to_vllm = { + "server.host": "host", + "server.port": "port", + "server.liveness_timeout_seconds": "liveness_timeout_seconds", + "model.name": "model", + "model.dtype": "dtype", + "model.max_model_len": "max_model_len", + "model.enforce_eager": "enforce_eager", + "model.trust_remote_code": "trust_remote_code", + "model.chat_template": "chat_template", + "model.tool_call_parser": "tool_call_parser", + "model.reasoning_parser": "reasoning_parser", + "model.rope_scaling": "rope_scaling", + "parallel.tp": "tensor_parallel_size", + "parallel.dp": "data_parallel_size", + "data_parallel_size_local": "data_parallel_size_local", + "data_parallel_rpc_port": "data_parallel_rpc_port", + "enable_lora": "enable_lora", + "enable_prefix_caching": "enable_prefix_caching", + "max_loras": "max_loras", + "max_cpu_loras": "max_cpu_loras", + "max_lora_rank": "max_lora_rank", + "lora_target_modules": "lora_target_modules", + "gpu_memory_utilization": "gpu_memory_utilization", + "api_server_count": "api_server_count", + "enable_return_routed_experts": "enable_return_routed_experts", + "enable_expert_parallel": "enable_expert_parallel", + "all2all_backend": "all2all_backend", + "enable_eplb": "enable_eplb", + "enable_dbo": "enable_dbo", + "seed": "seed", + } + + for config_key, vllm_key in to_vllm.items(): + value = rgetattr(self, config_key.replace("-", "_")) + rsetattr(namespace, vllm_key, value) + + # Set `logprobs_mode` to `processed_logprobs` by default + rsetattr(namespace, "logprobs_mode", "processed_logprobs") + + if self.kv_cache_offload is not None: + rsetattr( + namespace, + "kv_transfer_config", + { + "kv_connector": "OffloadingConnector", + "kv_role": "kv_both", + "kv_connector_extra_config": { + "cpu_bytes_to_use": int(self.kv_cache_offload.cpu_bytes), + }, + }, + ) + + # Pass prime-rl-specific flags through vLLM's additional_config dict; + # workers read these via get_current_vllm_config().additional_config. + if self.enable_fp32_lm_head: + existing = getattr(namespace, "additional_config", None) or {} + existing["fp32_lm_head"] = True + rsetattr(namespace, "additional_config", existing) + + # Remove chat_template if not set (vLLM doesn't accept None) + if namespace.chat_template is None: + delattr(namespace, "chat_template") + + # Remove reasoning_parser if not set (vLLM doesn't accept None) + if namespace.reasoning_parser is None: + delattr(namespace, "reasoning_parser") + + # Remove lora_target_modules if not set (vLLM doesn't accept None) + if hasattr(namespace, "lora_target_modules") and namespace.lora_target_modules is None: + delattr(namespace, "lora_target_modules") + + # Remove rope_scaling if not set (vLLM doesn't accept None) + if hasattr(namespace, "rope_scaling"): + if namespace.rope_scaling is None: + delattr(namespace, "rope_scaling") + + return namespace diff --git a/packages/prime-rl-configs/src/prime_rl/configs/orchestrator.py b/packages/prime-rl-configs/src/prime_rl/configs/orchestrator.py new file mode 100644 index 00000000000..39118162273 --- /dev/null +++ b/packages/prime-rl-configs/src/prime_rl/configs/orchestrator.py @@ -0,0 +1,1343 @@ +import math +import warnings +from pathlib import Path +from typing import Annotated, Any, Literal, TypeAlias + +from pydantic import AliasChoices, BaseModel, ConfigDict, Field, model_validator + +from prime_rl.configs.shared import ( + BaseModelConfig, + ClientConfig, + FileSystemTransportConfig, + HeartbeatConfig, + LogConfig, + PrimeMonitorConfig, + RendererConfig, + TransportConfig, + WandbWithExtrasConfig, +) +from prime_rl.configs.trainer import TokenizerConfig +from prime_rl.utils.config import BaseConfig + + +class OptimizerConfig(BaseConfig): + """Per-run optimizer configuration for multi-run training.""" + + lr: Annotated[ + float, + Field( + ge=0, + description="Learning rate for this run.", + ), + ] = 1e-4 + + +class LoRAConfig(BaseConfig): + """Per-run LoRA configuration for multi-run training.""" + + name: Annotated[ + str | None, + Field( + description="Name of the LoRA adapter. If None, auto-generated from rank and alpha.", + ), + ] = None + + rank: Annotated[ + int | None, + Field( + ge=1, + description="LoRA rank for this run. Must be <= trainer's max rank. If None, uses trainer's rank.", + ), + ] = None + + alpha: Annotated[ + float | None, + Field( + ge=0, + description="LoRA alpha for this run. If None, uses trainer's alpha.", + ), + ] = None + + +class ModelConfig(BaseModelConfig): + """Extended model configuration with per-run LoRA settings.""" + + lora: Annotated[ + LoRAConfig | None, + Field( + description="LoRA configuration. If None, LoRA is not used.", + ), + ] = None + + +class TrainSamplingConfig(BaseConfig): + """Configures how tokens are sampled from the model for training. Largely follows the vLLM sampling parameters.""" + + temperature: Annotated[ + float, + Field( + ge=0, + description="Temperature for sampling.", + ), + ] = 1.0 + + repetition_penalty: Annotated[ + float, + Field( + ge=0, + description="Penalty for repeating tokens. Values > 1.0 discourage repetition, values < 1.0 encourage repetition, and 1.0 means no penalty.", + ), + ] = 1.0 + + max_completion_tokens: Annotated[ + int | None, + Field( + validation_alias=AliasChoices("max_completion_tokens", "max_tokens"), + description="Maximum number of output tokens to generate per turn. If None, will generate until maximum context length or EOS token is hit.", + ), + ] = None + + min_tokens: Annotated[ + int, + Field( + ge=0, + description="Minimum number of output tokens to generate per sequence.", + ), + ] = 0 + + seed: Annotated[ + int | None, + Field( + description="Random seed to use for sampling. If None, no seeding is used.", + ), + ] = None + + # Strictly speaking, extra_body is not a sampling parameter, but it is the + # easiest way to pass arbitrary extra parameters to the server via verifiers + extra_body: Annotated[ + dict[str, Any], + Field( + description="Extra body to pass with each request to the inference server. By default, it is set to an empty dictionary.", + ), + ] = {} + + def to_sampling_args(self) -> dict[str, Any]: + """Convert to OAI-compatible sampling args dict, omitting None values.""" + # Top-level OAI params + args: dict[str, Any] = { + "temperature": self.temperature, + "top_p": 1.0, + "logprobs": True, + } + if self.max_completion_tokens is not None: + args["max_completion_tokens"] = self.max_completion_tokens + if self.seed is not None: + args["seed"] = self.seed + + # vLLM extra_body params + extra_body = dict(self.extra_body) + if self.min_tokens > 0: + extra_body["min_tokens"] = self.min_tokens + if self.repetition_penalty != 1.0: + extra_body["repetition_penalty"] = self.repetition_penalty + if extra_body: + args["extra_body"] = extra_body + + return args + + @model_validator(mode="before") + @classmethod + def _deprecate_max_tokens(cls, data: Any) -> Any: + if isinstance(data, dict) and "max_tokens" in data and "max_completion_tokens" not in data: + warnings.warn( + "'max_tokens' is deprecated, use 'max_completion_tokens' instead. " + "Auto-translating for now, but this will be removed in a future release.", + FutureWarning, + stacklevel=2, + ) + return data + + +class EvalSamplingConfig(BaseConfig): + """Configures how tokens are sampled from the model for evaluation. + + All sampling fields default to None, meaning the inference server's own + default is used. Only explicitly set fields are forwarded. + """ + + temperature: Annotated[ + float | None, + Field(ge=0, description="Sampling temperature. None defers to the inference server default."), + ] = None + + repetition_penalty: Annotated[ + float | None, + Field(ge=0, description="Repetition penalty. None defers to the inference server default."), + ] = None + + top_p: Annotated[ + float | None, + Field(description="Nucleus sampling threshold. None defers to the inference server default."), + ] = None + + top_k: Annotated[ + int | None, + Field(description="Top-k sampling. None defers to the inference server default."), + ] = None + + min_p: Annotated[ + float | None, + Field(ge=0, description="Min-p sampling threshold. None defers to the inference server default."), + ] = None + + max_completion_tokens: Annotated[ + int | None, + Field( + validation_alias=AliasChoices("max_completion_tokens", "max_tokens"), + description="Maximum output tokens per turn. None defers to the inference server default.", + ), + ] = None + + min_tokens: Annotated[ + int | None, + Field(ge=0, description="Minimum output tokens per sequence. None defers to the inference server default."), + ] = None + + reasoning_effort: Annotated[ + Literal["minimal", "low", "medium", "high"] | None, + Field(description="Reasoning effort constraint for reasoning models."), + ] = None + + seed: Annotated[ + int | None, + Field(description="Random seed for sampling. None means no seeding."), + ] = None + + extra_body: Annotated[ + dict[str, Any], + Field(description="Extra body parameters forwarded to the inference server."), + ] = {} + + def to_sampling_args(self) -> dict[str, Any]: + """Convert to OAI-compatible sampling args dict. Only includes non-None fields.""" + args: dict[str, Any] = {} + if self.temperature is not None: + args["temperature"] = self.temperature + if self.top_p is not None: + args["top_p"] = self.top_p + if self.max_completion_tokens is not None: + args["max_completion_tokens"] = self.max_completion_tokens + if self.reasoning_effort is not None: + args["reasoning_effort"] = self.reasoning_effort + if self.seed is not None: + args["seed"] = self.seed + + extra_body = dict(self.extra_body) + if self.top_k is not None: + extra_body["top_k"] = self.top_k + if self.min_p is not None: + extra_body["min_p"] = self.min_p + if self.min_tokens is not None: + extra_body["min_tokens"] = self.min_tokens + if self.repetition_penalty is not None: + extra_body["repetition_penalty"] = self.repetition_penalty + if extra_body: + args["extra_body"] = extra_body + + return args + + @model_validator(mode="before") + @classmethod + def _deprecate_max_tokens(cls, data: Any) -> Any: + if isinstance(data, dict) and "max_tokens" in data and "max_completion_tokens" not in data: + warnings.warn( + "'max_tokens' is deprecated, use 'max_completion_tokens' instead. " + "Auto-translating for now, but this will be removed in a future release.", + FutureWarning, + stacklevel=2, + ) + return data + + +class EnvConfig(BaseConfig): + """Base environment configuration.""" + + id: Annotated[ + str, + Field( + description="Registered verifiers environment ID (e.g. 'math-env', 'primeintellect/math-env'). May include an @version suffix for installation." + ), + ] = "reverse-text" + + name: Annotated[ + str | None, + Field( + description="Display name for this environment in logs, metrics, and buffer keys. Defaults to the id (without @version). Must be unique across all envs in the same group." + ), + ] = None + + args: Annotated[ + dict, + Field( + description="Keyword arguments forwarded to vf.load_environment. See the environment's docstring for accepted args." + ), + ] = {} + + extra_env_kwargs: Annotated[ + dict[str, Any], + Field( + description=( + "Extra kwargs passed to an env (e.g. seq_len, max_total_completion_tokens). This field is auto-populated on the orchestrator for all envs. It is generally NOT recommended for this field to be overriden by the user. It's main use case is to match the extra_env_kwargs when running an env in an isolated environment server." + ), + ), + ] = {} + + address: Annotated[ + str | None, + Field( + description="ZMQ address of an external env server (e.g. 'tcp://host:5000'). When set, the orchestrator connects to this server instead of spawning one. When None (default), a subprocess env server is spawned automatically." + ), + ] = None + + num_workers: Annotated[ + int | Literal["auto"], + Field( + description="Number of worker processes for the spawned env server. 'auto' scales to 1 worker per 256 concurrent rollouts. Ignored when address is set (external server)." + ), + ] = "auto" + + ratio: Annotated[ + float | None, + Field( + gt=0, + description="Sampling weight for this environment in the buffer. When None for all envs, samples uniformly across all available problems. When set, must be set on all envs — values are relative weights normalized to probabilities (e.g. [1, 1] and [0.5, 0.5] are equivalent).", + ), + ] = None + + max_retries: Annotated[ + int, + Field(ge=0, description="Number of times the env server retries a failed rollout before returning an error."), + ] = 0 + + max_total_completion_tokens: Annotated[ + int, + Field( + description=( + "Maximum total completion tokens across all turns in a multi-turn rollout. " + "Set to -1 (default) to disable. Auto-populated into extra_env_kwargs." + ), + ), + ] = -1 + + timeout: Annotated[ + float | None, + Field( + validation_alias=AliasChoices("timeout", "timeout_seconds"), + description="Per-rollout wall-clock timeout in seconds. Set to None (default) to disable.", + ), + ] = None + + @property + def stripped_id(self) -> str: + """Environment ID without the @version suffix.""" + return self.id.split("@")[0] + + @property + def resolved_name(self) -> str: + return self.name or self.stripped_id + + @model_validator(mode="after") + def validate_env_name(self): + if self.resolved_name == "all": + raise ValueError( + 'Environment name "all" is reserved for global metric aggregation. Use a different name or id.' + ) + return self + + @model_validator(mode="after") + def resolve_max_total_completion_tokens(self): + self.extra_env_kwargs["max_total_completion_tokens"] = self.max_total_completion_tokens + return self + + @model_validator(mode="after") + def resolve_timeout(self): + if self.timeout is not None: + self.extra_env_kwargs["timeout_seconds"] = self.timeout + return self + + +class TrainEnvConfig(EnvConfig): + """Configures a training environment.""" + + sampling: Annotated[ + TrainSamplingConfig, + Field( + description="Per-env sampling overrides. Unset fields inherit from the group-level train sampling config.", + ), + ] = TrainSamplingConfig() + + +class EvalEnvConfig(EnvConfig): + """Configures an evaluation environment.""" + + sampling: Annotated[ + EvalSamplingConfig, + Field( + description="Per-env sampling overrides. Unset fields inherit from the group-level eval sampling config.", + ), + ] = EvalSamplingConfig() + + num_examples: Annotated[ + int, + Field( + description="Number of eval examples to sample from the dataset. Set to -1 to use all available examples." + ), + ] = -1 + + rollouts_per_example: Annotated[ + int, + Field( + ge=1, + description="Number of rollouts generated per example. Used for pass@k estimation (e.g. rollouts_per_example=8 enables pass@1 through pass@8).", + ), + ] = 1 + + interval: Annotated[ + int, + Field( + ge=1, + description="Per-env eval interval. If unset, inherits from the group-level eval interval.", + ), + ] = 100 + + +class TrainConfig(BaseConfig): + """Configures training environments and their shared sampling settings.""" + + env: list[TrainEnvConfig] = [TrainEnvConfig()] + + sampling: TrainSamplingConfig = TrainSamplingConfig() + + num_workers: Annotated[ + int | Literal["auto"], + Field( + description="Default number of worker processes for env servers. Can be overridden per env.", + ), + ] = "auto" + + max_retries: Annotated[ + int, + Field(ge=0, description="Default number of retries for failed rollouts. Can be overridden per env."), + ] = 0 + + @model_validator(mode="after") + def resolve_env_defaults(self): + """Resolve per-env overrides: inherit group-level sampling, num_workers, and max_retries.""" + group_sampling = self.sampling.model_dump() + for env in self.env: + if "sampling" not in env.model_fields_set: + env.sampling = TrainSamplingConfig(**group_sampling) + else: + merged = group_sampling | env.sampling.model_dump(exclude_unset=True) + env.sampling = TrainSamplingConfig(**merged) + if "num_workers" not in env.model_fields_set: + env.num_workers = self.num_workers + if "max_retries" not in env.model_fields_set: + env.max_retries = self.max_retries + return self + + @model_validator(mode="after") + def validate_unique_env_names(self): + env_names = [env.resolved_name for env in self.env] + duplicates = [n for n in env_names if env_names.count(n) > 1] + if duplicates: + raise ValueError( + f"Duplicate training environment names: {set(duplicates)}. Each env must have a unique name." + ) + return self + + @model_validator(mode="after") + def validate_env_ratios(self): + ratios = [env.ratio for env in self.env] + if all(r is None for r in ratios): + return self + if any(r is None for r in ratios): + raise ValueError("Either all envs must have a ratio or none of them. Got a mix of set and unset ratios.") + return self + + +class EvalConfig(BaseConfig): + """Configures evaluation using verifiers environments.""" + + env: list[EvalEnvConfig] = [EvalEnvConfig()] + + sampling: EvalSamplingConfig = Field( + default_factory=EvalSamplingConfig, + description="Shared sampling configuration for evals; can differ from training sampling.", + ) + + num_examples: Annotated[ + int, + Field( + description="Default number of eval examples per environment. Set to -1 to use all. Can be overridden per env." + ), + ] = -1 + + rollouts_per_example: Annotated[ + int, + Field(ge=1, description="Default number of rollouts per example. Can be overridden per env."), + ] = 1 + + num_workers: Annotated[ + int | Literal["auto"], + Field( + description="Default number of worker processes for env servers. Can be overridden per env.", + ), + ] = "auto" + + max_retries: Annotated[ + int, + Field(ge=0, description="Default number of retries for failed rollouts. Can be overridden per env."), + ] = 0 + + interval: Annotated[ + int, + Field( + ge=1, + description="Interval at which to evaluate the model.", + ), + ] = 100 + + @model_validator(mode="after") + def resolve_env_defaults(self): + """Resolve per-env overrides: inherit group-level sampling, num_workers, max_retries, num_examples, rollouts_per_example, and interval. Then resolve auto num_workers.""" + group_sampling = self.sampling.model_dump() + for env in self.env: + if "sampling" not in env.model_fields_set: + env.sampling = EvalSamplingConfig(**group_sampling) + else: + merged = group_sampling | env.sampling.model_dump(exclude_unset=True) + env.sampling = EvalSamplingConfig(**merged) + if "num_examples" not in env.model_fields_set: + env.num_examples = self.num_examples + if "rollouts_per_example" not in env.model_fields_set: + env.rollouts_per_example = self.rollouts_per_example + if "interval" not in env.model_fields_set: + env.interval = self.interval + if "num_workers" not in env.model_fields_set: + env.num_workers = self.num_workers + if "max_retries" not in env.model_fields_set: + env.max_retries = self.max_retries + # Resolve auto num_workers now that num_examples and rollouts_per_example are set + if env.num_workers == "auto": + if env.num_examples == -1: + env.num_workers = 4 + else: + max_concurrent = env.num_examples * env.rollouts_per_example + env.num_workers = max(1, math.ceil(max_concurrent / 256)) + return self + + @model_validator(mode="after") + def validate_unique_env_names(self): + env_names = [env.resolved_name for env in self.env] + duplicates = [n for n in env_names if env_names.count(n) > 1] + if duplicates: + raise ValueError( + f"Duplicate evaluation environment names: {set(duplicates)}. Each env must have a unique name." + ) + return self + + eval_base_model: Annotated[ + bool, + Field( + description="Whether to evaluate the base model we are training on.", + ), + ] = True + + skip_eval_on_resume: Annotated[ + bool, + Field( + validation_alias=AliasChoices("skip_eval_on_resume", "skip_eval_on_restart"), + description=( + "If True and resuming the orchestrator from a checkpoint, skip the (potentially redundant) " + "online eval that would otherwise run immediately at the resumed checkpoint step." + ), + ), + ] = True + + cancel_inflight_rollouts_on_eval: Annotated[ + bool, + Field( + description="Whether to cancel in-flight training rollouts before starting online evals. This is useful to avoid congestion (e.g. do not have training + eval rollouts happening at the same time) but leads to slower training steps as rollouts get cancelled and the pipeline has to fill up after each eval", + ), + ] = False + + +class CheckpointConfig(BaseConfig): + """Configures checkpointing the orchestrator.""" + + interval: Annotated[int | None, Field(ge=1, description="Interval at which to save the checkpoint.")] = None + + resume_step: Annotated[ + int | None, + Field( + ge=-1, + description="Step to resume orchestrator from. If None, will start from scratch. If -1, will restart from latest checkpoint available.", + ), + ] = None + + wait_for_weights_timeout: Annotated[ + int | None, + Field( + ge=1, + description="When resuming, wait up to this many seconds for the weight directory to appear. Useful when the orchestrator restarts while the trainer is still saving weights. If None (default), fail immediately if weights are not found.", + ), + ] = None + + keep_last: Annotated[ + int | None, + Field( + ge=1, + description="Keep at most this many recent step checkpoints on disk. If None, never clean old checkpoints based on recency.", + ), + ] = None + + keep_interval: Annotated[ + int | None, + Field( + ge=1, + description="Keep checkpoints at every N steps permanently (e.g., keep_interval=100 keeps step 100, 200, ...). If None, no interval-based keeping.", + ), + ] = None + + skip_progress: Annotated[ + bool, + Field( + description="Whether to skip loading the progress from checkpoint.", + ), + ] = False + + skip_buffer: Annotated[ + bool, + Field( + description="Whether to skip loading the buffer from checkpoint.", + ), + ] = False + + +class BufferConfig(BaseConfig): + """Configures the buffer for the orchestrator.""" + + seed: Annotated[ + int | None, + Field( + description="Random seed to use for the buffer. If set, the sampling from the buffer will be deterministic.", + ), + ] = None + + easy_threshold: Annotated[ + float | None, + Field( + description="Threshold for easy difficulty classification. If average reward >= this threshold, mark as easy.", + ), + ] = None + + hard_threshold: Annotated[ + float | None, + Field( + description="Threshold for hard difficulty classification. If average reward <= this threshold, mark as hard.", + ), + ] = None + + easy_fraction: Annotated[ + float, + Field( + ge=0, + le=1, + description="Fraction of easy problems to convert to normal when resuming or starting training. Only problems with difficulty 'normal' are sampled.", + ), + ] = 0.0 + + hard_fraction: Annotated[ + float, + Field( + ge=0, + le=1, + description="Fraction of hard problems to convert to normal when resuming or starting training. Only problems with difficulty 'normal' are sampled.", + ), + ] = 0.0 + + online_difficulty_filtering: Annotated[ + bool, + Field( + description="Whether to filter rollouts based on difficulty. If True, rollouts with average reward 0.0 or 1.0 are not added to the buffer.", + ), + ] = False + + hash_keys: Annotated[ + list[str], + Field( + min_length=1, + description="Keys to use for computing example hashes. Will be used to match examples from buffer checkpoints and determine buffer resume behavior.", + ), + ] = ["env_name", "prompt"] + + @model_validator(mode="after") + def validate_thresholds(self): + if self.easy_threshold is not None and self.hard_threshold is not None: + assert self.easy_threshold > self.hard_threshold, "easy_threshold must be greater than hard_threshold." + return self + + +class TokensLengthPenaltyConfig(BaseModel): + """Length penalty by weighted token cost. + + Effective cost = completion_weight * model_completion_tokens + + tool_response_weight * tool_response_tokens. + + Tool-response tokens are read from the rollout's harness metric + `*_total_tool_response_tokens` (e.g. `rlm_total_tool_response_tokens`); 0 if absent — + so for envs without tool accounting, only the completion term contributes regardless of weight. + """ + + model_config = ConfigDict(extra="forbid") + + type: Literal["tokens"] = "tokens" + completion_weight: Annotated[ + float, + Field( + ge=0, + allow_inf_nan=False, + description="Weight on model completion tokens. Finite and non-negative.", + ), + ] = 1.0 + tool_response_weight: Annotated[ + float, + Field( + ge=0, + allow_inf_nan=False, + description="Weight on tool-response tokens. Finite and non-negative.", + ), + ] = 1.0 + + +class TurnsLengthPenaltyConfig(BaseModel): + """Length penalty by trajectory turn count.""" + + model_config = ConfigDict(extra="forbid") + + type: Literal["turns"] = "turns" + + +LengthPenaltyConfig: TypeAlias = Annotated[ + TokensLengthPenaltyConfig | TurnsLengthPenaltyConfig, + Field(discriminator="type"), +] + + +class DefaultAdvantageConfig(BaseModel): + """Config for the default advantage.""" + + model_config = ConfigDict(extra="forbid") + + type: Literal["default"] = "default" + length_penalty: Annotated[ + LengthPenaltyConfig | None, + Field( + description=( + "Correctness-gated length penalty. `tokens` shapes by weighted token cost, " + "`turns` shapes by trajectory turn count, `None` disables shaping. In mixed " + "groups, lower-cost correct rollouts get amplified advantage (up to 2x), higher-cost correct " + "rollouts are unchanged, incorrect untouched. In all-correct groups, below-average-cost " + "rollouts get advantage in [0, 1], others get 0." + ) + ), + ] = None + + +class CustomAdvantageConfig(BaseModel): + """Config for a custom external advantage function.""" + + type: Literal["custom"] = "custom" + import_path: Annotated[ + str, Field(description="Import path to the advantage function (e.g., 'my_module.my_advantage')") + ] + kwargs: Annotated[ + dict[str, Any], Field(default_factory=dict, description="Kwargs to pass to the advantage function") + ] + + +AdvantageConfig: TypeAlias = Annotated[ + DefaultAdvantageConfig | CustomAdvantageConfig, + Field(discriminator="type"), +] + + +class GibberishFilterConfig(BaseModel): + """Flags rare tokens generated at high entropy (Section 5.2, https://arxiv.org/abs/2510.02387).""" + + model_config = ConfigDict(extra="forbid") + + type: Literal["gibberish"] = "gibberish" + enforce: Annotated[ + bool, + Field( + description="If True, skip detected rollouts entirely so they are not sent to the trainer. If False, only track detection metrics." + ), + ] = False + token_id_threshold: Annotated[ + int, + Field(description="Token IDs above this are candidates for gibberish. BPE tokens are sorted by merge order."), + ] = 100_000 + logprob_offset: Annotated[ + float, + Field(description="Offset from uniform distribution logprob. Threshold = -log(vocab_size) - logprob_offset."), + ] = 2.0 + + +class RepetitionFilterConfig(BaseModel): + """Flags rollouts where the model gets stuck in a repetition loop, emitting high-confidence tokens + for an extended stretch. A rollout is flagged when `window` consecutive tokens are each sampled + with probability above `prob_threshold`. (Section 3.2, https://arxiv.org/abs/2506.13585)""" + + model_config = ConfigDict(extra="forbid") + + type: Literal["repetition"] = "repetition" + enforce: Annotated[ + bool, + Field( + description="If True, skip detected rollouts entirely so they are not sent to the trainer. If False, only track detection metrics." + ), + ] = False + window: Annotated[ + int, + Field(ge=1, description="Number of consecutive high-probability steps before flagging."), + ] = 3_000 + prob_threshold: Annotated[ + float, + Field( + gt=0, + le=1, + description="Tokens sampled with probability above this are considered repetitive. Consecutive such tokens count toward the window.", + ), + ] = 0.99 + + +class ZeroAdvantageFilterConfig(BaseModel): + """Flags rollouts with zero advantage.""" + + model_config = ConfigDict(extra="forbid") + + type: Literal["zero_advantage"] = "zero_advantage" + enforce: Annotated[ + bool, + Field( + description="If True, skip detected rollouts entirely so they are not sent to the trainer. If False, only track detection metrics." + ), + ] = True + + +FilterConfig: TypeAlias = Annotated[ + GibberishFilterConfig | RepetitionFilterConfig | ZeroAdvantageFilterConfig, + Field(discriminator="type"), +] + + +class FileSystemWeightBroadcastConfig(BaseModel): + """Configures the filesystem weight broadcast.""" + + type: Literal["filesystem"] = "filesystem" + + +class NCCLWeightBroadcastConfig(BaseModel): + """Configures the NCCL weight broadcast.""" + + type: Literal["nccl"] = "nccl" + + host: Annotated[str, Field(description="The host to use for the NCCL broadcast.")] = "localhost" + port: Annotated[int, Field(description="The port to use for the NCCL broadcast.")] = 29501 + timeout: Annotated[int, Field(description="The timeout in seconds to use for the NCCL broadcast.")] = 1200 + quantize_in_weight_transfer: Annotated[ + bool, + Field(description="Use kernel-format FP8 quantized NCCL transfer for weight updates."), + ] = False + + inference_world_size: Annotated[ + int, + Field( + ge=1, + description="Total number of inference GPUs across all servers. Used by init_nccl_broadcast to compute per-server rank offsets.", + ), + ] = 1 + + +WeightBroadcastConfig: TypeAlias = Annotated[ + FileSystemWeightBroadcastConfig | NCCLWeightBroadcastConfig, Field(discriminator="type") +] + + +class OrchestratorExperimentalConfig(BaseConfig): + """Experimental features for the orchestrator.""" + + +class TeacherModelConfig(BaseConfig): + """Configures the teacher model for computing teacher logprobs (e.g. for distillation).""" + + client: Annotated[ + ClientConfig, + Field(description="The OAI client configuration for the teacher model."), + ] = ClientConfig() + + model: Annotated[ + ModelConfig, + Field(description="The model configuration for the teacher model."), + ] = ModelConfig() + + +class TeacherRolloutModelConfig(BaseConfig): + """Configures an external teacher model used to generate rollout text.""" + + client: Annotated[ + ClientConfig, + Field(description="The OAI client configuration for rollout generation."), + ] = ClientConfig() + + model: Annotated[ + ModelConfig, + Field(description="The model configuration for rollout generation."), + ] = ModelConfig() + + +class OrchestratorConfig(BaseConfig): + """Configures the orchestrator for RL training.""" + + # Training environments and sampling + train: TrainConfig = TrainConfig() + + # The OAI client configuration + client: ClientConfig = ClientConfig() + + # The model configuration + model: ModelConfig = ModelConfig() + + # The tokenizer configuration + tokenizer: TokenizerConfig = TokenizerConfig() + + # The renderer configuration (only used when use_renderer=True) + renderer: RendererConfig = RendererConfig() + + # The optimizer configuration (per-run LR for multi-run training) + optim: OptimizerConfig = OptimizerConfig() + + # The teacher model configuration (optional) + teacher_model: Annotated[ + TeacherModelConfig | None, + Field( + description="The teacher model configuration for computing teacher logprobs (e.g. for distillation). " + "If provided, teacher logprobs will be computed using the specified model. " + "If None, no teacher model will be used." + ), + ] = None + + # External teacher rollout model configuration (optional) + teacher_rollout_model: Annotated[ + TeacherRolloutModelConfig | None, + Field( + description=( + "Optional external teacher model used for rollout generation. " + "When set, rollouts are generated from this endpoint/model instead of the student inference server." + ), + ), + ] = None + + # When True, trainer uses SFT loss instead of RL loss (per-run override for hosted multi-tenant training) + use_sft_loss: Annotated[ + bool, + Field( + description=( + "When True, use SFT masked NLL loss instead of the trainer's configured RL loss. " + "Requires a teacher_rollout_model to be configured." + ), + ), + ] = False + + # The evaluation configuration + eval: EvalConfig | None = None + + # Data buffer configuration + buffer: BufferConfig = BufferConfig() + + # The advantage configuration + advantage: AdvantageConfig | None = DefaultAdvantageConfig() + + # Rollout filters (monitor by default, enforce optionally) + filters: list[FilterConfig] = [GibberishFilterConfig(), RepetitionFilterConfig(), ZeroAdvantageFilterConfig()] + + # The logging configuration + log: LogConfig = LogConfig() + + # The wandb configuration + wandb: WandbWithExtrasConfig | None = None + + # The prime monitor configuration + prime_monitor: PrimeMonitorConfig | None = None + + # Whether to collect inference server metrics (requires wandb) + collect_inference_metrics: bool = True + + # The checkpoint configuration + ckpt: CheckpointConfig | None = None + + weight_broadcast: WeightBroadcastConfig = FileSystemWeightBroadcastConfig() + + rollout_transport: TransportConfig = FileSystemTransportConfig() + + output_dir: Annotated[ + Path, + Field( + description="Directory to write outputs to. Will be populated with checkpoints, weights, rollouts and logs as subdirectories. Should be set to a persistent directory with enough disk space. This value should be distinct across experiments running on a single node. See the README for more details." + ), + ] = Path("outputs/run_default") + + tasks_per_minute: Annotated[ + int | None, + Field( + ge=1, + description="Rate limit for tasks per environment worker, in tasks per minute. Recommended for sandbox-backed environments to prevent sandbox-not-ready errors during autoscaling. When set to None, no rate limiting is applied. Note: with multiple workers, the effective total rate equals workers × this value.", + ), + ] = None + + batch_size: Annotated[ + int | None, + Field( + ge=1, + description="Number of samples to train on per step (rollout-based batching). Set this OR token_batch_size.", + ), + ] = None + + token_batch_size: Annotated[ + int | None, + Field( + ge=1, + description="Number of tokens to train on per step (token-based batching). Set this OR batch_size.", + ), + ] = None + + oversampling_factor: Annotated[ + float | None, + Field( + ge=1, + description=( + "Rollout-mode batching only. Multiplier used to derive max_inflight_rollouts from batch_size " + "when max_inflight_rollouts is unset." + ), + ), + ] = None + + max_inflight_rollouts: Annotated[ + int | None, + Field( + ge=1, + description=( + "Maximum number of rollouts to keep in-flight. Required for token-based batching. " + "If batch_size is set and this is unset, defaults to batch_size * oversampling_factor " + "(or batch_size when oversampling_factor is unset)." + ), + ), + ] = None + + rollouts_per_example: Annotated[ + int, + Field( + ge=1, + description="Number of output sequences to return per example during training.", + ), + ] = 1 + + seq_len: Annotated[ + int, + Field( + description="Sequence length to use for training. If a sample is shorter than this, it will be padded. If a sequence is longer than this, it will be truncated.", + ), + ] = 2048 + + # TODO(Mika): This should be automatic from the number of ZMQ connections + num_train_workers: Annotated[ + int, + Field(default=1, ge=1, description="Number of training workers to use for training."), + ] = 1 + + max_steps: Annotated[ + int | None, + Field( + description="Maximum number of training steps to run. If None, will run indefinitely.", + ), + ] = None + + max_off_policy_steps: Annotated[ + int, + Field( + ge=0, + description="Maximum number of policies that are allowed to generate a single rollout. Rollouts that are generated from more than `max_off_policy_steps` steps ahead of training will be discarded. Higher values yield better throughput, but lead to more off-policyness in training.", + ), + ] = 8 + + max_async_level: Annotated[ + int, + Field( + ge=0, + description="Maximum number of steps the inference can be ahead of training. If 0, will degenerate to synchronous on-policy RL. If >=1, training and inference will be overlapped.", + ), + ] = 1 + + strict_async_level: Annotated[ + bool, + Field( + description="Whether to strictly enforce the max async level. If True, will always ensure that the policy used for generating rollouts is exactly `max_async_level` steps ahead of training. If False, any policy that is at most `max_async_level` steps ahead of training is allowed, i.e. we always use the latest available policy.", + ), + ] = False + + bench: Annotated[ + bool, + Field( + description="Whether to run in benchmark mode. It will automatically set the maximum number of steps to run to 5, max async level to ~infinity and disable W&B.", + ), + ] = False + + seed: Annotated[int | None, Field(description="Random seed for the orchestrator.")] = 42 + + heartbeat: Annotated[ + HeartbeatConfig | None, Field(description="The heartbeat config for monitoring training progress.") + ] = None + + use_token_client: Annotated[ + bool, + Field( + description="Whether to use the token-in-token-out (TITO) client for training across all environments. " + "WARNING: Only use this if your environment has a linear history and the chat template has the extension " + "property (i.e. no tokens are ever removed or inserted by the chat template). Mutually exclusive with " + "``use_renderer``." + ), + ] = True + + use_renderer: Annotated[ + bool, + Field( + description="Whether to use the renderer client (client-side tokenization via the ``renderers`` package, " + "served by ``/v1/generate``). Mutually exclusive with ``use_token_client``. When True, the " + "``[orchestrator.renderer]`` block (name / tool_parser / reasoning_parser / pool_size) " + "applies; when False those fields must be left at their defaults. Not supported for VLMs — " + "VLMs must use the token client (TITO) so image preprocessing and chat templating stay server-side." + ), + ] = False + + env_install_prerelease: Annotated[ + bool, + Field( + description="Allow pre-release versions when installing environments (e.g. verifiers>=0.1.12.dev5). Passes --prerelease to prime env install." + ), + ] = False + + experimental: Annotated[ + OrchestratorExperimentalConfig, + Field(description="Experimental features for the orchestrator."), + ] = OrchestratorExperimentalConfig() + + @model_validator(mode="before") + @classmethod + def _env_to_train(cls, data: Any) -> Any: + """Allow [[env]] and [sampling] as shorthand for [train] with [[train.env]] and [train.sampling].""" + if not isinstance(data, dict): + return data + if "env" in data or "sampling" in data: + train = data.setdefault("train", {}) + if isinstance(train, dict): + if "env" in data: + warnings.warn( + "'[[orchestrator.env]]' is deprecated, use '[[orchestrator.train.env]]' instead. " + "Auto-translating for now, but this will be removed in a future release.", + FutureWarning, + stacklevel=2, + ) + train.setdefault("env", data.pop("env")) + if "sampling" in data: + warnings.warn( + "'[orchestrator.sampling]' is deprecated, use '[orchestrator.train.sampling]' instead. " + "Auto-translating for now, but this will be removed in a future release.", + FutureWarning, + stacklevel=2, + ) + train.setdefault("sampling", data.pop("sampling")) + return data + + @model_validator(mode="after") + def auto_setup_tokenizer(self): + if self.tokenizer.name is None: + self.tokenizer.name = self.model.name + if self.tokenizer.trust_remote_code is None: + self.tokenizer.trust_remote_code = self.model.trust_remote_code + return self + + @model_validator(mode="after") + def validate_unique_filter_types(self): + types = [f.type for f in self.filters] + if len(types) != len(set(types)): + raise ValueError(f"Duplicate filter types: {types}. Each filter type may only appear once.") + return self + + @model_validator(mode="after") + def validate_sft_distill_mode(self): + """Enforce the SFT hard distill invariants that involve only orchestrator fields. + + Runs at ``OrchestratorConfig`` level so hosted deployments (which load this + config standalone via the ``orchestrator`` entrypoint) get the same guarantees + as the combined ``rl`` entrypoint. + """ + has_teacher = self.teacher_rollout_model is not None + if self.use_sft_loss and not has_teacher: + raise ValueError( + "orchestrator.use_sft_loss = true requires orchestrator.teacher_rollout_model to be configured." + ) + if has_teacher and not self.use_sft_loss: + raise ValueError("orchestrator.teacher_rollout_model requires orchestrator.use_sft_loss = true.") + if has_teacher and self.use_token_client: + raise ValueError( + "orchestrator.use_token_client must be false when orchestrator.teacher_rollout_model is configured." + ) + if has_teacher and self.use_renderer: + raise ValueError( + "orchestrator.use_renderer must be false when orchestrator.teacher_rollout_model is configured " + "(external rollout uses MITO)." + ) + return self + + @model_validator(mode="after") + def validate_client_mode(self): + """The two client toggles select among three exclusive modes: + + - ``use_token_client=True`` + ``use_renderer=False`` → TITO (default) + - ``use_token_client=False`` + ``use_renderer=True`` → renderer + - ``use_token_client=False`` + ``use_renderer=False`` → MITO + + Both True is invalid: TITO and renderer are different wire protocols + (server-side templating vs client-side tokenization). + """ + if self.use_token_client and self.use_renderer: + raise ValueError( + "orchestrator.use_token_client and orchestrator.use_renderer are mutually exclusive. " + "Pick one: token client (TITO, server-side templating) or renderer client (client-side " + "tokenization)." + ) + return self + + @model_validator(mode="after") + def validate_renderer_vs_vlm(self): + """The renderer client takes plain message dicts and tokenizes + them client-side. VLMs need server-side image preprocessing and + chat templating, so they must use the token client (TITO) — fail + loudly when both are set.""" + if self.use_renderer and self.model.vlm is not None: + raise ValueError( + "orchestrator.use_renderer is not supported for VLMs. Use the token client " + "(``use_token_client=true``, the default) so image preprocessing and chat " + "templating stay on the inference server." + ) + return self + + @model_validator(mode="after") + def validate_renderer_args(self): + """``[orchestrator.renderer]`` knobs are only meaningful when + ``use_renderer=True``. Reject otherwise so callers don't silently + pass them and wonder why they're ignored.""" + if self.use_renderer: + return self + + renderer_args_set = [] + if self.renderer.name != "auto": + renderer_args_set.append(f"renderer.name={self.renderer.name!r}") + if self.renderer.tool_parser is not None: + renderer_args_set.append(f"renderer.tool_parser={self.renderer.tool_parser!r}") + if self.renderer.reasoning_parser is not None: + renderer_args_set.append(f"renderer.reasoning_parser={self.renderer.reasoning_parser!r}") + if self.renderer.pool_size is not None: + renderer_args_set.append(f"renderer.pool_size={self.renderer.pool_size!r}") + + if renderer_args_set: + raise ValueError( + "Renderer-specific args set without orchestrator.use_renderer=True: " + f"{', '.join(renderer_args_set)}. Either enable the renderer client or remove these knobs." + ) + return self + + @model_validator(mode="after") + def nccl_max_async_level(self): + if self.weight_broadcast.type == "nccl": + if not self.max_async_level == 1: + raise ValueError("max_async_level must be 1 for NCCL broadcast") + return self + + @model_validator(mode="after") + def resolve_batching(self): + has_rollout_batch = self.batch_size is not None + has_token_batch = self.token_batch_size is not None + + if has_rollout_batch and has_token_batch: + raise ValueError("Set exactly one of batch_size or token_batch_size") + + if not has_rollout_batch and not has_token_batch: + self.batch_size = 128 + + if has_token_batch: + if self.oversampling_factor is not None: + raise ValueError("oversampling_factor can only be set when batch_size is set") + if self.max_inflight_rollouts is None: + raise ValueError("max_inflight_rollouts must be set when token_batch_size is set") + else: + assert self.batch_size is not None + if self.batch_size % self.rollouts_per_example != 0: + raise ValueError("Batch size must be divisible by the number of samples per problem") + if self.max_inflight_rollouts is not None and self.oversampling_factor is not None: + expected_max_inflight_rollouts = int(self.batch_size * self.oversampling_factor) + if self.max_inflight_rollouts != expected_max_inflight_rollouts: + raise ValueError("max_inflight_rollouts conflicts with oversampling_factor * batch_size") + if self.max_inflight_rollouts is None: + oversampling_factor = self.oversampling_factor if self.oversampling_factor is not None else 1.0 + self.max_inflight_rollouts = int(self.batch_size * oversampling_factor) + + if self.max_inflight_rollouts is not None and self.max_inflight_rollouts < self.rollouts_per_example: + raise ValueError("max_inflight_rollouts must be at least the number of rollouts per example") + + # Resolve train env num_workers from max_inflight_rollouts + for env_cfg in self.train.env: + if env_cfg.num_workers == "auto": + assert self.max_inflight_rollouts is not None + env_cfg.num_workers = max(1, math.ceil(self.max_inflight_rollouts / 256)) + + return self + + @model_validator(mode="after") + def auto_setup_bench(self): + if self.bench: + self.max_steps = 4 # Run for 1 warmup step + 3 evaluation steps + self.max_async_level = int(1e9) # Never wait for RL weight checkpoints + + # Disable evaluation + self.eval = None + if self.wandb: + self.wandb.log_extras = None + if self.prime_monitor: + self.prime_monitor.log_extras = None + + return self + + @model_validator(mode="after") + def resolve_env_config(self): + """Populate extra_env_kwargs and vLLM sampling defaults from top-level fields.""" + is_vllm = self.teacher_rollout_model is None + for env in self.train.env: + env.extra_env_kwargs.update(max_seq_len=self.seq_len) + if is_vllm: + env.sampling.extra_body.setdefault("top_k", -1) + env.sampling.extra_body.setdefault("min_p", 0.0) + env.sampling.extra_body.setdefault("return_token_ids", True) + return self diff --git a/packages/prime-rl-configs/src/prime_rl/configs/rl.py b/packages/prime-rl-configs/src/prime_rl/configs/rl.py new file mode 100644 index 00000000000..a160af2c9fa --- /dev/null +++ b/packages/prime-rl-configs/src/prime_rl/configs/rl.py @@ -0,0 +1,986 @@ +import warnings +from pathlib import Path +from typing import Annotated, Literal, TypeAlias + +from pydantic import BaseModel, ConfigDict, Field, model_validator + +from prime_rl.configs.inference import InferenceConfig +from prime_rl.configs.inference import WeightBroadcastConfig as InferenceWeightBroadcastConfig +from prime_rl.configs.orchestrator import ( + CheckpointConfig as OrchestratorCheckpointConfig, +) +from prime_rl.configs.orchestrator import ( + FileSystemWeightBroadcastConfig as OrchestratorFileSystemWeightBroadcastConfig, +) +from prime_rl.configs.orchestrator import ( + NCCLWeightBroadcastConfig as OrchestratorNCCLWeightBroadcastConfig, +) +from prime_rl.configs.orchestrator import ( + OrchestratorConfig, +) +from prime_rl.configs.shared import ( + SlurmConfig, + VLMConfig, + WandbConfig, + WandbWithExtrasConfig, +) +from prime_rl.configs.trainer import ( + BenchConfig, + FakeDataLoaderConfig, + TokenizerConfig, + TrainerConfig, +) +from prime_rl.configs.trainer import ( + CheckpointConfig as TrainerCheckpointConfig, +) +from prime_rl.configs.trainer import ( + FileSystemWeightBroadcastConfig as TrainerFileSystemWeightBroadcastConfig, +) +from prime_rl.configs.trainer import ( + NCCLWeightBroadcastConfig as TrainerNCCLWeightBroadcastConfig, +) +from prime_rl.utils.config import BaseConfig, find_package_resource +from prime_rl.utils.validation import ( + validate_shared_ckpt_config, + validate_shared_max_async_level, + validate_shared_max_steps, + validate_shared_model_name, + validate_shared_output_dir, + validate_shared_tokenizer, + validate_shared_wandb_config, + validate_shared_weight_broadcast, +) + + +class RLExperimentalConfig(BaseConfig): + """Experimental features for RL training.""" + + +class SharedLogConfig(BaseConfig): + """Configures shared logging.""" + + level: Annotated[ + str | None, + Field( + description="The log level to use. When unset, the trainer and orchestrator log levels are used as-is (which themselves default to the PRIME_LOG_LEVEL env var if set, else 'info').", + ), + ] = None + + json_logging: Annotated[ + bool, + Field(description="Emit JSON logs (newline-delimited) for log aggregation (Loki, Grafana, etc.)."), + ] = False + + +class SharedWandbConfig(BaseConfig): + """Configures shared W&B configs.""" + + project: Annotated[str | None, Field(description="The W&B project to use.")] = "prime-rl" + + entity: Annotated[str | None, Field(description="The W&B entity to use.")] = None + + name: Annotated[str | None, Field(description="The W&B run name to use.")] = None + + group: Annotated[str | None, Field(description="The W&B group to use.")] = None + + tags: Annotated[list[str] | None, Field(description="The W&B tags to attach to the run.")] = None + + offline: Annotated[bool | None, Field(description="Whether to run W&B in offline mode.")] = False + + shared: Annotated[ + bool, + Field( + description="Use shared W&B mode to log trainer and orchestrator metrics to a single run. " + "Requires wandb SDK >= 0.19.9. Incompatible with offline mode.", + ), + ] = True + + @model_validator(mode="after") + def validate_shared_not_offline(self): + if self.shared and self.offline: + raise ValueError("W&B shared mode requires server connectivity and is incompatible with offline mode") + return self + + +class SharedCheckpointConfig(BaseConfig): + """Configures shared checkpoint configs.""" + + output_dir: Annotated[ + Path | None, + Field( + description="Override directory for checkpoints and weights. When set, checkpoints and weight snapshots are written here instead of under the trainer output_dir.", + ), + ] = None + + interval: Annotated[int | None, Field(description="The interval at which to save checkpoints.")] = None + + resume_step: Annotated[ + int | None, Field(description="The step to resume from. If None, will not resume from a checkpoint.") + ] = None + + keep_last: Annotated[ + int | None, + Field( + ge=1, + description="Keep at most this many recent step checkpoints on disk. If None, never clean old checkpoints based on recency.", + ), + ] = None + + keep_interval: Annotated[ + int | None, + Field( + ge=1, + description="Keep checkpoints at every N steps permanently (e.g., keep_interval=100 keeps step 100, 200, ...). If None, no interval-based keeping.", + ), + ] = None + + +class SharedModelConfig(BaseConfig): + """Configures shared model settings.""" + + name: Annotated[ + str, + Field(description="The name of the model to use."), + ] = "Qwen/Qwen3-0.6B" + + vlm: Annotated[ + "VLMConfig | None", + Field(description="VLM configuration. Set to enable vision-language model support."), + ] = None + + +class SharedWeightBroadcastConfig(BaseConfig): + """Configures shared weight broadcast settings.""" + + type: Annotated[Literal["nccl", "filesystem"], Field(description="The type of weight broadcast to use.")] = ( + "filesystem" + ) + + port: Annotated[int, Field(description="The port to use for NCCL weight broadcast.")] = 29501 + timeout: Annotated[int, Field(description="The timeout in seconds for NCCL weight broadcast.")] = 1200 + quantize_in_weight_transfer: Annotated[ + bool, + Field( + description=( + "Use kernel-format FP8 quantized NCCL transfer for weight updates. " + "When disabled, uses default HF checkpoint-format transfer." + ), + ), + ] = False + + +class BaseDeploymentConfig(BaseModel): + """Configures a base deployment.""" + + model_config = ConfigDict(extra="forbid") + + gpus_per_node: Annotated[int, Field(description="Number of GPUs per node.")] = 8 + + +class SingleNodeDeploymentConfig(BaseDeploymentConfig): + """Configures a single node deployment.""" + + type: Literal["single_node"] = "single_node" + + num_train_gpus: Annotated[int, Field(description="Number of training GPUs")] = 1 + num_infer_gpus: Annotated[int, Field(description="Number of inference GPUs")] = 1 + num_teacher_gpus: Annotated[int | None, Field(description="Number of teacher inference GPUs")] = None + + @model_validator(mode="after") + def validate_gpu_count(self): + total = self.num_train_gpus + self.num_infer_gpus + (self.num_teacher_gpus or 0) + if total > self.gpus_per_node: + raise ValueError( + f"Total GPU count ({total} = {self.num_train_gpus} train + {self.num_infer_gpus} infer" + f" + {self.num_teacher_gpus or 0} teacher) exceeds gpus_per_node ({self.gpus_per_node})." + ) + return self + + +class MultiNodeDeploymentConfig(BaseDeploymentConfig): + """Configures a multi node deployment.""" + + type: Literal["multi_node"] = "multi_node" + + num_train_nodes: Annotated[int, Field(description="Number of training nodes.")] + num_infer_nodes: Annotated[ + int, + Field( + ge=0, + description="Number of inference nodes per replica. Set to 0 to skip inference and orchestrator (requires fake data).", + ), + ] + num_infer_replicas: Annotated[ + int, + Field( + ge=1, + description="Number of independent inference replicas. Total inference nodes = num_infer_nodes * num_infer_replicas.", + ), + ] = 1 + num_teacher_nodes: Annotated[int | None, Field(description="Number of teacher inference nodes.")] = None + + nodes_per_fsdp_group: Annotated[ + int | None, + Field( + description="Number of training nodes per FSDP island. Auto-sets trainer.dp_replicate = num_train_nodes / nodes_per_fsdp_group." + ), + ] = None + + @property + def total_infer_nodes(self) -> int: + return self.num_infer_nodes * self.num_infer_replicas + + @model_validator(mode="after") + def teacher_inference_not_supported(self): + if self.num_teacher_nodes is not None: + raise ValueError("Teacher inference is not yet supported in multi node deployment.") + return self + + +DeploymentConfig: TypeAlias = Annotated[ + SingleNodeDeploymentConfig | MultiNodeDeploymentConfig, Field(discriminator="type") +] + + +class RLConfig(BaseConfig): + """Configures an RL training run.""" + + trainer: TrainerConfig + orchestrator: OrchestratorConfig + inference: Annotated[ + InferenceConfig | None, + Field( + description="The inference config. If None, the rl entrypoint will not start an inference server (useful for elastic inference pools or manually started servers)." + ), + ] = None + + teacher_inference: Annotated[ + InferenceConfig | None, + Field( + description="Teacher inference config. If None, will use the same config as inference or a default config. Only used when teacher GPUs or nodes are set." + ), + ] = None + + output_dir: Annotated[ + Path, + Field( + description="The directory to store the outputs. Should be set to a unique directory identifying the experiment." + ), + ] = Path("outputs") + + clean_output_dir: Annotated[ + bool, + Field( + description="If true, delete the output directory before starting training. Required to overwrite an output directory that contains checkpoints from a previous run when not resuming.", + ), + ] = False + + ### Shared configurations + + log: Annotated[ + SharedLogConfig, + Field( + description="Shared log configs. If None, will fallback to the log configs specified on submodule configs." + ), + ] = SharedLogConfig() + + ckpt: Annotated[ + SharedCheckpointConfig | None, + Field( + description="Shared checkpoint configs. If None, will fallback to the checkpoint configs specified on submodule configs." + ), + ] = None + + wandb: Annotated[ + SharedWandbConfig | None, + Field( + description="Shared W&B configs. If None, will fallback to the W&B configs specified on submodule configs." + ), + ] = None + + model: Annotated[ + SharedModelConfig | None, + Field( + description="Shared model configs. If None, will fallback to the model configs specified on submodule configs." + ), + ] = None + + tokenizer: Annotated[ + TokenizerConfig | None, + Field( + description="Shared tokenizer config. Propagated to trainer, orchestrator, and inference. " + "If None, each component uses its own tokenizer config (defaulting to model name).", + ), + ] = None + + max_steps: Annotated[ + int | None, + Field( + description="The maximum number of steps to train for. If None, will fallback to the max steps specified on submodule configs." + ), + ] = None + + max_model_len: Annotated[ + int | None, + Field( + description="The maximum model length to use. If None, will fallback to the max model length specified on submodule configs." + ), + ] = None + + seq_len: Annotated[ + int | None, + Field( + description="Shared sequence length. Propagates to trainer.model.seq_len and orchestrator.seq_len, " + "but only for those not explicitly set in the config. " + "Explicitly set per-component values always take precedence." + ), + ] = None + + max_async_level: Annotated[ + int | None, + Field( + description="The async level to use. If None, will fallback to the async level specified on submodule configs." + ), + ] = None + + weight_broadcast: Annotated[ + SharedWeightBroadcastConfig | None, Field(description="The weight broadcast config.") + ] = None + + bench: Annotated[ + bool, + Field( + description="Whether to run in benchmark mode. Automatically sets the trainer and orchestrator to benchmark mode and, if present, suffixes the W&B project with `-bench`.", + ), + ] = False + + deployment: DeploymentConfig = SingleNodeDeploymentConfig() + + slurm: Annotated[SlurmConfig | None, Field(description="SLURM configuration. If None, will run locally.")] = None + + dry_run: Annotated[bool, Field(description="Only validate and dump resolved configs and exit early.")] = False + + experimental: Annotated[ + RLExperimentalConfig, + Field(description="Experimental features for RL training."), + ] = RLExperimentalConfig() + + ### Validate configs (e.g. raise for unsupported (combinations of) configs) + + @model_validator(mode="after") + def validate_deployment(self): + if self.deployment.type == "multi_node": + if self.slurm is None: + raise ValueError("Must use SLURM for multi-node deployment.") + if self.deployment.num_infer_nodes > 0 and not self.inference: + raise ValueError("Must configure inference when using multi-node deployment with inference nodes.") + if self.deployment.num_infer_nodes == 0 and self.inference: + raise ValueError( + "Cannot configure inference with num_infer_nodes = 0. " + "Either set num_infer_nodes > 0 or remove the inference config." + ) + if self.deployment.num_infer_nodes == 0 and not self.trainer.data.fake and not self.bench: + raise ValueError( + "Must use fake data (trainer.data.fake or bench = true) when num_infer_nodes = 0, " + "since no orchestrator or inference server will be running." + ) + return self + + # TODO: fix this + @model_validator(mode="after") + def validate_no_teacher_in_multinode(self): + if self.deployment.type == "multi_node" and self.teacher_inference is not None: + raise ValueError( + "Teacher inference is not supported in multi-node deployment. " + "The SLURM template only handles inference and training nodes." + ) + return self + + @model_validator(mode="after") + def validate_enough_devices_for_nccl(self): + if self.deployment.type == "single_node": + if self.trainer.weight_broadcast.type == "nccl": + if self.deployment.num_train_gpus + self.deployment.num_infer_gpus < 2: + raise ValueError( + "NCCL weight broadcast requires at least 2 GPUs to build the broadcast process group." + ) + return self + + @model_validator(mode="after") + def validate_quantize_in_weight_transfer(self): + if self.weight_broadcast is None or not self.weight_broadcast.quantize_in_weight_transfer: + return self + + if self.weight_broadcast.type != "nccl": + raise ValueError("weight_broadcast.quantize_in_weight_transfer requires weight_broadcast.type = 'nccl'.") + + if self.inference is None: + raise ValueError("weight_broadcast.quantize_in_weight_transfer requires an inference config.") + + if self.trainer.model.impl != "custom": + raise ValueError("weight_broadcast.quantize_in_weight_transfer requires trainer.model.impl = 'custom'.") + + return self + + @model_validator(mode="after") + def validate_teacher_model(self): + if ( + self.trainer.loss.type == "default" and self.trainer.loss.teacher_tau > 0 + ) and not self.orchestrator.teacher_model: + raise ValueError( + "teacher_model must be configured when teacher_tau > 0. " + "Either set teacher_tau = 0, set deployment.num_teacher_gpus, or configure teacher_model manually." + ) + return self + + @model_validator(mode="after") + def validate_external_rollout_inference(self): + """Forbid configuring a local inference server when rollouts come from an external teacher. + + Orchestrator-only invariants (``use_sft_loss`` paired with ``teacher_rollout_model``, + and ``use_token_client`` coupling) live on ``OrchestratorConfig`` so the hosted + orchestrator entrypoint also enforces them. + """ + if self.orchestrator.teacher_rollout_model is not None and self.inference is not None: + raise ValueError( + "inference must be omitted when orchestrator.teacher_rollout_model is configured. " + "External rollout mode does not use the local inference server." + ) + return self + + ### Auto-setup and validate shared configs + + @model_validator(mode="after") + def auto_setup_output_dir(self): + """Auto-setup shared output directory for trainer and orchestrator.""" + self.trainer.output_dir = self.output_dir + self.orchestrator.output_dir = self.output_dir / "run_default" + + validate_shared_output_dir(self.trainer, self.orchestrator) + + return self + + @model_validator(mode="after") + def auto_setup_logs(self): + """Auto-setup shared log config for trainer and orchestrator.""" + if self.log is not None: + if self.log.level is not None: + self.trainer.log.level = self.log.level + self.orchestrator.log.level = self.log.level + self.trainer.log.json_logging = self.log.json_logging + self.orchestrator.log.json_logging = self.log.json_logging + + return self + + @model_validator(mode="after") + def auto_setup_ckpt(self): + """Auto-setup shared checkpoint config for trainer and orchestrator.""" + if self.ckpt is not None: + # Create checkpoint configs if not specified + if self.trainer.ckpt is None: + self.trainer.ckpt = TrainerCheckpointConfig() + if self.orchestrator.ckpt is None: + self.orchestrator.ckpt = OrchestratorCheckpointConfig() + + # If specified, override checkpoint output directory + if self.ckpt.output_dir is not None: + self.trainer.ckpt.output_dir = self.ckpt.output_dir + + # If specified, use the same ckpt interval + if self.ckpt.interval is not None: + self.trainer.ckpt.interval = self.ckpt.interval + self.orchestrator.ckpt.interval = self.ckpt.interval + + # If resuming training, ensure orchestrator resume from the same step + if self.ckpt.resume_step is not None: + self.trainer.ckpt.resume_step = self.ckpt.resume_step + self.orchestrator.ckpt.resume_step = self.ckpt.resume_step + + # If specified, propagate keep policy + if self.ckpt.keep_last is not None: + self.trainer.ckpt.keep_last = self.ckpt.keep_last + self.orchestrator.ckpt.keep_last = self.ckpt.keep_last + + if self.ckpt.keep_interval is not None: + self.trainer.ckpt.keep_interval = self.ckpt.keep_interval + self.orchestrator.ckpt.keep_interval = self.ckpt.keep_interval + + validate_shared_ckpt_config(self.trainer, self.orchestrator) + + return self + + @model_validator(mode="after") + def auto_setup_wandb(self): + """Auto-setup shared W&B config for trainer and orchestrator.""" + if self.wandb is not None: + if not self.trainer.wandb: + self.trainer.wandb = WandbConfig() + if not self.orchestrator.wandb: + self.orchestrator.wandb = WandbWithExtrasConfig() + + if self.wandb.project: + self.trainer.wandb.project = self.wandb.project + self.orchestrator.wandb.project = self.wandb.project + + if self.wandb.entity: + self.trainer.wandb.entity = self.wandb.entity + self.orchestrator.wandb.entity = self.wandb.entity + + if self.wandb.shared: + if self.wandb.name: + self.trainer.wandb.name = self.wandb.name + self.orchestrator.wandb.name = self.wandb.name + else: + if self.wandb.name: + self.trainer.wandb.name = f"{self.wandb.name}-trainer" + self.orchestrator.wandb.name = f"{self.wandb.name}-orchestrator" + + if self.wandb.group: + self.trainer.wandb.group = self.wandb.group + self.orchestrator.wandb.group = self.wandb.group + + if self.wandb.tags: + self.trainer.wandb.tags = self.wandb.tags.copy() + self.orchestrator.wandb.tags = self.wandb.tags.copy() + + if self.wandb.offline: + self.trainer.wandb.offline = self.wandb.offline + self.orchestrator.wandb.offline = self.wandb.offline + + validate_shared_wandb_config(self.trainer, self.orchestrator) + + if self.orchestrator.prime_monitor is not None and self.orchestrator.prime_monitor.run_name is None: + if self.wandb and self.wandb.name: + self.orchestrator.prime_monitor.run_name = self.wandb.name + + return self + + @model_validator(mode="after") + def auto_setup_model(self): + """Auto-setup shared model config for trainer, orchestrator, and inference.""" + if self.model is not None: + self.trainer.model.name = self.model.name + if self.inference is not None: + inference_model_explicitly_set = "name" in self.inference.model.model_fields_set + if not inference_model_explicitly_set: + self.inference.model.name = self.model.name + self.orchestrator.model.name = self.inference.model.name + else: + self.orchestrator.model.name = self.model.name + + if self.model.vlm is not None: + self.trainer.model.vlm = self.model.vlm + self.orchestrator.model.vlm = self.model.vlm + if self.inference is not None: + self.inference.model.vlm = self.model.vlm + + validate_shared_model_name(self.trainer, self.orchestrator, self.inference) + + return self + + @model_validator(mode="after") + def auto_setup_tokenizer(self): + """Auto-setup shared tokenizer config for trainer, orchestrator, and inference.""" + if self.tokenizer is not None: + # Shared tokenizer config: propagate to all components, then fill + # in name/trust_remote_code from model config where still unset. + self.trainer.tokenizer = self.tokenizer.model_copy() + self.orchestrator.tokenizer = self.tokenizer.model_copy() + for component in (self.trainer, self.orchestrator): + if component.tokenizer.name is None: + component.tokenizer.name = component.model.name + if component.tokenizer.trust_remote_code is None: + component.tokenizer.trust_remote_code = component.model.trust_remote_code + else: + # No shared tokenizer: re-derive from (now-correct) model names, + # since auto_setup_tokenizer on sub-configs already ran with defaults. + for component in (self.trainer, self.orchestrator): + component.tokenizer.name = component.model.name + component.tokenizer.trust_remote_code = component.model.trust_remote_code + + # Propagate chat_template to inference (vLLM --chat-template) + if self.inference is not None: + chat_template = self.trainer.tokenizer.chat_template + if chat_template is not None and self.inference.model.chat_template is None: + self.inference.model.chat_template = chat_template + + validate_shared_tokenizer(self.trainer, self.orchestrator, self.inference) + + return self + + @model_validator(mode="after") + def auto_setup_max_steps(self): + """Auto-setup shared max steps for trainer and orchestrator.""" + if self.max_steps is not None: + self.trainer.max_steps = self.max_steps + self.orchestrator.max_steps = self.max_steps + + validate_shared_max_steps(self.trainer, self.orchestrator) + + return self + + @model_validator(mode="after") + def auto_setup_async_level(self): + """Auto-setup shared async level for trainer and orchestrator.""" + if self.max_async_level is not None: + self.trainer.max_async_level = self.max_async_level + self.orchestrator.max_async_level = self.max_async_level + + validate_shared_max_async_level(self.trainer, self.orchestrator) + + return self + + @model_validator(mode="after") + def auto_setup_seq_len(self): + """Auto-setup shared seq_len for trainer and orchestrator. + + Only propagates to components that weren't explicitly set in the config. + Uses model_fields_set to detect explicit assignment. + """ + if self.seq_len is not None: + if "seq_len" not in self.trainer.model.model_fields_set: + self.trainer.model.seq_len = self.seq_len + if "seq_len" not in self.orchestrator.model_fields_set: + self.orchestrator.seq_len = self.seq_len + + if self.trainer.model.seq_len < self.orchestrator.seq_len: + raise ValueError( + f"Trainer model seq_len ({self.trainer.model.seq_len}) must be >= orchestrator seq_len ({self.orchestrator.seq_len}). " + f"The trainer needs to be able to handle sequences at least as long as those produced by the orchestrator." + ) + + return self + + @model_validator(mode="after") + def auto_setup_weight_broadcast(self): + """Auto-setup shared weight broadcast config for trainer, orchestrator, and inference.""" + if self.weight_broadcast is not None: + if self.weight_broadcast.type == "nccl": + inference_world_size = self.inference.parallel.dp * self.inference.parallel.tp if self.inference else 1 + self.trainer.weight_broadcast = TrainerNCCLWeightBroadcastConfig( + type=self.weight_broadcast.type, + inference_world_size=inference_world_size, + port=self.weight_broadcast.port, + timeout=self.weight_broadcast.timeout, + quantize_in_weight_transfer=self.weight_broadcast.quantize_in_weight_transfer, + ) + self.orchestrator.weight_broadcast = OrchestratorNCCLWeightBroadcastConfig( + type=self.weight_broadcast.type, + port=self.weight_broadcast.port, + timeout=self.weight_broadcast.timeout, + inference_world_size=inference_world_size, + quantize_in_weight_transfer=self.weight_broadcast.quantize_in_weight_transfer, + ) + elif self.weight_broadcast.type == "filesystem": + self.trainer.weight_broadcast = TrainerFileSystemWeightBroadcastConfig() + self.orchestrator.weight_broadcast = OrchestratorFileSystemWeightBroadcastConfig() + if self.inference is not None: + self.inference.weight_broadcast = InferenceWeightBroadcastConfig(type=self.weight_broadcast.type) + + validate_shared_weight_broadcast(self.trainer, self.orchestrator, self.inference) + + return self + + @model_validator(mode="after") + def validate_eplb_requires_quantized_weight_transfer(self): + if self.inference is None or not self.inference.enable_eplb: + return self + + # TODO(matej): check if weight reloading works itself before supporting EPLB without quantized transfer. + trainer_weight_broadcast = self.trainer.weight_broadcast + if trainer_weight_broadcast.type != "nccl" or not trainer_weight_broadcast.quantize_in_weight_transfer: + raise ValueError( + "inference.enable_eplb requires weight_broadcast.type = 'nccl' and " + "weight_broadcast.quantize_in_weight_transfer = true." + ) + + return self + + @model_validator(mode="after") + def auto_setup_bench(self): + if self.bench: + self.trainer.bench = BenchConfig() + self.orchestrator.bench = True + self.trainer.data.fake = FakeDataLoaderConfig( + batch_size=self.orchestrator.batch_size or 32, + ) + + trainer_bench_enabled = self.trainer.bench is not None + if trainer_bench_enabled != self.orchestrator.bench: + raise ValueError( + f"Trainer benchmark mode ({self.trainer.bench}) and orchestrator benchmark mode " + f"({self.orchestrator.bench}) must match. Use the top-level bench = true to set both." + ) + + return self + + @model_validator(mode="after") + def auto_setup_lora(self): + if self.trainer.model.lora is not None: + if self.trainer.weight_broadcast.type == "nccl": + raise ValueError("NCCL weight broadcast does not support LoRA yet.") + + if self.orchestrator.model.lora is None: + from prime_rl.configs.orchestrator import LoRAConfig + + self.orchestrator.model.lora = LoRAConfig() + + if ( + self.orchestrator.model.lora.rank is not None + and self.orchestrator.model.lora.rank != self.trainer.model.lora.rank + ): + raise ValueError( + f"orchestrator.model.lora.rank ({self.orchestrator.model.lora.rank}) conflicts with " + f"trainer.model.lora.rank ({self.trainer.model.lora.rank}). " + f"Remove orchestrator.model.lora.rank to inherit from trainer, or update trainer.model.lora.rank to match." + ) + + if ( + self.orchestrator.model.lora.alpha is not None + and self.orchestrator.model.lora.alpha != self.trainer.model.lora.alpha + ): + raise ValueError( + f"orchestrator.model.lora.alpha ({self.orchestrator.model.lora.alpha}) conflicts with " + f"trainer.model.lora.alpha ({self.trainer.model.lora.alpha}). " + f"Remove orchestrator.model.lora.alpha to inherit from trainer, or update trainer.model.lora.alpha to match." + ) + + if self.orchestrator.model.lora.rank is None: + self.orchestrator.model.lora.rank = self.trainer.model.lora.rank + + if self.orchestrator.model.lora.alpha is None: + self.orchestrator.model.lora.alpha = self.trainer.model.lora.alpha + + if self.orchestrator.model.lora.name is None: + self.orchestrator.model.lora.name = ( + f"r{self.orchestrator.model.lora.rank}-a{self.orchestrator.model.lora.alpha}" + ) + + if self.inference is not None: + self.inference.enable_lora = True + self.inference.max_lora_rank = self.trainer.model.lora.rank + else: + warnings.warn( + "LoRA is enabled, but inference is not configured. When manually starting the inference server, " + "make sure to set --enable_lora and --max-lora-rank.", + stacklevel=2, + ) + + return self + + @model_validator(mode="after") + def auto_setup_session_headers(self): + """Ensure X-Session-ID header is always set for sticky DP-aware routing at the inference router.""" + self.orchestrator.client.extra_headers_from_state.setdefault("X-Session-ID", "example_id") + return self + + @model_validator(mode="after") + def auto_setup_router_replay(self): + if self.trainer.enable_router_replay: + if self.inference is not None: + if self.inference.enable_return_routed_experts is False: + warnings.warn( + "Router replay is enabled, but inference.enable_return_routed_experts is False. Setting to True.", + stacklevel=2, + ) + self.inference.enable_return_routed_experts = True + else: + warnings.warn( + "Router replay is enabled, but inference is not configured. When manually starting the inference server, make sure to pass `--enable-return-routed-experts` to the vLLM server.", + stacklevel=2, + ) + return self + + @model_validator(mode="after") + def auto_setup_deployment(self): + if self.deployment.type == "single_node": # single-node + # set num_train_workers to the number of data replicas + non_data_parallel_size = self.trainer.model.cp + if self.deployment.num_train_gpus > 1: + self.orchestrator.num_train_workers = self.deployment.num_train_gpus // non_data_parallel_size + + # fill up inference capacity with dp ranks + if self.inference is not None: + num_infer_gpus = self.deployment.num_infer_gpus + if num_infer_gpus != self.inference.parallel.dp * self.inference.parallel.tp: + assert num_infer_gpus % self.inference.parallel.tp == 0, ( + "Number of inference GPUs must be divisible by the tensor parallel size" + ) + self.inference.parallel.dp = num_infer_gpus // self.inference.parallel.tp + # Ensure api_server_count matches DP so all workers are created. + # Without this, the NCCL broadcast group expects dp*tp workers + # but only api_server_count*tp exist, causing a deadlock. + dp = self.inference.parallel.dp + if self.inference.api_server_count < dp and not self.inference.enable_lora: + self.inference.api_server_count = dp + + elif self.deployment.type == "multi_node": # multi-node + self.orchestrator.num_train_workers = self.deployment.num_train_nodes * self.deployment.gpus_per_node + + if self.deployment.nodes_per_fsdp_group is not None: + if self.deployment.num_train_nodes % self.deployment.nodes_per_fsdp_group != 0: + raise ValueError( + f"deployment.num_train_nodes ({self.deployment.num_train_nodes}) must be divisible by " + f"deployment.nodes_per_fsdp_group ({self.deployment.nodes_per_fsdp_group})" + ) + self.trainer.model.dp_replicate = ( + self.deployment.num_train_nodes // self.deployment.nodes_per_fsdp_group + ) + + if ( + self.inference is not None + and self.inference.enable_expert_parallel + and self.inference.deployment.type != "disaggregated" + ): + inference_tp = self.inference.parallel.tp + if self.deployment.gpus_per_node % inference_tp != 0: + raise ValueError( + "deployment.gpus_per_node must be divisible by inference.parallel.tp " + "when inference.enable_expert_parallel is enabled in multi-node deployment." + ) + + inferred_dp_local = self.deployment.gpus_per_node // inference_tp + total_infer_gpus = self.deployment.num_infer_nodes * self.deployment.gpus_per_node + expected_global_world_size = self.inference.parallel.dp * inference_tp + if expected_global_world_size != total_infer_gpus: + raise ValueError( + "For multi-node expert parallel inference, inference.parallel.dp * inference.parallel.tp " + f"must match total inference GPUs ({total_infer_gpus}), got {expected_global_world_size}." + ) + + if self.inference.data_parallel_size_local is None: + self.inference.data_parallel_size_local = inferred_dp_local + elif self.inference.data_parallel_size_local != inferred_dp_local: + raise ValueError( + "inference.data_parallel_size_local must equal deployment.gpus_per_node / inference.parallel.tp " + f"({inferred_dp_local}) when inference.enable_expert_parallel is enabled in multi-node deployment." + ) + + if not self.inference.enable_lora and self.inference.api_server_count == self.inference.parallel.dp: + self.inference.api_server_count = inferred_dp_local + + # Auto-infer DP and api_server_count for standard multi-node inference. + # Without EP, vLLM only creates api_server_count * tp workers per node, + # not gpus_per_node workers. If DP isn't set, the broadcast group expects + # more workers than exist, deadlocking NCCL init. + if ( + self.inference is not None + and not self.inference.enable_expert_parallel + and self.inference.deployment.type != "disaggregated" + ): + dp_per_node = self.deployment.gpus_per_node // self.inference.parallel.tp + if self.inference.parallel.dp == 1 and dp_per_node > 1: + self.inference.parallel.dp = dp_per_node + if self.inference.data_parallel_size_local is None and dp_per_node > 1: + self.inference.data_parallel_size_local = dp_per_node + if self.inference.api_server_count == 1 and dp_per_node > 1: + self.inference.api_server_count = dp_per_node + + if self.weight_broadcast is not None and self.weight_broadcast.type == "nccl": + # Compute inference_world_size from actual worker count per server: + # each api_server runs tp workers that participate in collective_rpc. + api_server_count = self.inference.api_server_count if self.inference else 1 + tp = self.inference.parallel.tp if self.inference else 1 + total_infer_workers = self.deployment.total_infer_nodes * api_server_count * tp + assert self.trainer.weight_broadcast.type == "nccl" + self.trainer.weight_broadcast.host = "0.0.0.0" + self.trainer.weight_broadcast.inference_world_size = total_infer_workers + assert self.orchestrator.weight_broadcast.type == "nccl" + self.orchestrator.weight_broadcast.inference_world_size = total_infer_workers + + return self + + @model_validator(mode="after") + def auto_setup_disaggregated_inference(self): + """Auto-setup for disaggregated P/D inference within a multi-node deployment.""" + if self.inference is None or self.inference.deployment.type != "disaggregated": + return self + if self.deployment.type != "multi_node": + return self + + infer_deploy = self.inference.deployment + expected_infer_nodes = infer_deploy.num_prefill_nodes + infer_deploy.num_decode_nodes + if self.deployment.num_infer_nodes != expected_infer_nodes: + raise ValueError( + f"deployment.num_infer_nodes ({self.deployment.num_infer_nodes}) must equal " + f"inference.deployment.num_prefill_nodes ({infer_deploy.num_prefill_nodes}) + " + f"inference.deployment.num_decode_nodes ({infer_deploy.num_decode_nodes}) = {expected_infer_nodes}" + ) + + total_infer_gpus = self.deployment.total_infer_nodes * self.deployment.gpus_per_node + if self.weight_broadcast is not None and self.weight_broadcast.type == "nccl": + assert self.trainer.weight_broadcast.type == "nccl" + self.trainer.weight_broadcast.inference_world_size = total_infer_gpus + assert self.orchestrator.weight_broadcast.type == "nccl" + self.orchestrator.weight_broadcast.inference_world_size = total_infer_gpus + + return self + + @model_validator(mode="after") + def auto_setup_dp_rank_count(self): + """Auto-set orchestrator client dp_rank_count from inference DP size. + + Uses data_parallel_size_local (per-node DP) when set, since each base URL + points to a single node whose API server only knows about its local ranks. + Falls back to the global parallel.dp for single-node setups. + """ + if self.inference is not None and "dp_rank_count" not in self.orchestrator.client.model_fields_set: + self.orchestrator.client.dp_rank_count = ( + self.inference.data_parallel_size_local or self.inference.parallel.dp + ) + return self + + @model_validator(mode="after") + def auto_setup_teacher_inference(self): + """Auto-configure teacher inference server and orchestrator teacher_model client.""" + if self.deployment.type != "single_node": + return self + if self.deployment.num_teacher_gpus is None or self.deployment.num_teacher_gpus == 0: + return self + + import copy + + from prime_rl.configs.orchestrator import TeacherModelConfig + + if self.teacher_inference is None: + if self.inference is None: + self.teacher_inference = InferenceConfig() + else: + self.teacher_inference = copy.deepcopy(self.inference) + self.teacher_inference.server.port = (self.inference.server.port if self.inference else 8000) + 1 + elif self.inference is not None and self.teacher_inference.server.port == self.inference.server.port: + raise ValueError( + f"teacher_inference.server.port ({self.teacher_inference.server.port}) conflicts with " + f"inference.server.port ({self.inference.server.port}). " + "Either use different ports or let teacher_inference be auto-configured." + ) + + tp = self.teacher_inference.parallel.tp + num_teacher_gpus = self.deployment.num_teacher_gpus + if num_teacher_gpus != self.teacher_inference.parallel.dp * tp: + assert num_teacher_gpus % tp == 0, "Number of teacher GPUs must be divisible by tensor parallel size" + assert num_teacher_gpus > 0, "num_teacher_gpus cannot be zero" + self.teacher_inference.parallel.dp = num_teacher_gpus // tp + + if self.orchestrator.teacher_model is None: + self.orchestrator.teacher_model = TeacherModelConfig() + host = self.teacher_inference.server.host or "localhost" + port = self.teacher_inference.server.port + self.orchestrator.teacher_model.client.base_url = [f"http://{host}:{port}/v1"] + self.orchestrator.teacher_model.model.name = self.teacher_inference.model.name + + return self + + @model_validator(mode="after") + def auto_setup_slurm_template(self): + """Auto-setup the default single-node/multi-node SLURM template if no custom template is provided.""" + if self.slurm is not None and self.slurm.template_path is None: + templates_dir = find_package_resource("templates") + if templates_dir is not None: + if self.deployment.type == "single_node": + self.slurm.template_path = templates_dir / "single_node_rl.sbatch.j2" + else: + self.slurm.template_path = templates_dir / "multi_node_rl.sbatch.j2" + return self + + ### Warnings diff --git a/packages/prime-rl-configs/src/prime_rl/configs/sft.py b/packages/prime-rl-configs/src/prime_rl/configs/sft.py new file mode 100644 index 00000000000..868ae2d9736 --- /dev/null +++ b/packages/prime-rl-configs/src/prime_rl/configs/sft.py @@ -0,0 +1,434 @@ +import warnings +from pathlib import Path +from typing import Annotated, Literal, TypeAlias + +from pydantic import BaseModel, ConfigDict, Field, model_validator + +from prime_rl.configs.shared import ( + HeartbeatConfig, + SlurmConfig, + TrainerLogConfig, + WandbConfig, +) +from prime_rl.configs.trainer import ( + AdamWConfig, + BenchConfig, + CheckpointConfig, + ConstantSchedulerConfig, + GCConfig, + ModelConfig, + OptimizerConfig, + SchedulerConfig, + TokenizerConfig, +) +from prime_rl.utils.config import BaseConfig, find_package_resource + + +class BaseDataConfig(BaseModel): + """Base config for SFT data.""" + + batch_size: Annotated[int, Field(ge=1)] = 128 + seq_len: Annotated[int, Field(ge=1)] = 128 + pack_function: Literal["cat", "stack"] = "cat" + micro_batch_size: Annotated[int, Field(ge=1)] = 1 + + @model_validator(mode="after") + def validate_batch_size(self): + if self.batch_size % self.micro_batch_size != 0: + raise ValueError("Batch size must be divisible by micro batch size") + if self.batch_size < self.micro_batch_size: + raise ValueError("Batch size must be greater than or equal to micro batch size") + return self + + +class FakeDataConfig(BaseDataConfig): + """Configures fake data used for debugging.""" + + type: Literal["fake"] = "fake" + + length: Literal["fixed", "variable"] = "fixed" + input_ids: Literal["increasing", "random"] = "increasing" + + +class LossMaskConfig(BaseConfig): + """Configures which message types contribute to the loss. If True, the loss_mask will be True and the message type will contribute to the loss.""" + + system: Annotated[bool, Field(description="Whether system messages contribute to the loss.")] = False + user: Annotated[bool, Field(description="Whether user messages contribute to the loss.")] = False + assistant: Annotated[bool, Field(description="Whether assistant messages contribute to the loss.")] = True + tool: Annotated[bool, Field(description="Whether tool messages contribute to the loss.")] = False + + +class SFTDataConfig(BaseDataConfig): + """Configures the data used for training.""" + + type: Literal["sft"] = "sft" + + name: Annotated[str, Field(description="Name or path of the HF dataset to use.")] = ( + "PrimeIntellect/Reverse-Text-SFT" + ) + subsets: Annotated[list[str] | None, Field(description="Subsets to use from the HF dataset.")] = None + splits: Annotated[list[str] | None, Field(description="Splits to use from the HF dataset.")] = None + probabilities: Annotated[list[float] | None, Field(description="Probabilities to use for each subset/split.")] = ( + None + ) + stopping_strategy: Annotated[ + Literal["first_exhausted", "all_exhausted"], + Field(description=""), + ] = "all_exhausted" + shuffle: Annotated[bool, Field(description="Whether to shuffle the dataset at the beginning of each epoch.")] = True + seed: Annotated[ + int, + Field( + description="Random seed to use for shuffling the dataset. We also shuffle at the end of each epoch by adding epoch count to the seed." + ), + ] = 0 + + # Configuring + loss_mask: LossMaskConfig = LossMaskConfig() + + @model_validator(mode="after") + def validate_subsets_and_splits(self): + if self.subsets is not None or self.splits is not None: + if self.subsets is not None and self.splits is not None: + if len(self.subsets) != len(self.splits): + raise ValueError( + "Number of subsets must be equal to number of splits. Please specify which split to load for each subset." + ) + if self.subsets is not None and self.probabilities is not None: + if len(self.probabilities) != len(self.subsets): + raise ValueError( + "Number of probabilities must be equal to number of subsets. Please specify a probability for each subset." + ) + if self.splits is not None and self.probabilities is not None: + if len(self.probabilities) != len(self.splits): + raise ValueError( + "Number of probabilities must be equal to number of splits. Please specify a probability for each split." + ) + return self + + +class SFTValConfig(BaseConfig): + interval: Annotated[int, Field(ge=1, description="Run validation every N training steps.")] = 50 + eval_on_start: Annotated[bool, Field(description="Run validation before the first training step.")] = False + data: SFTDataConfig + + +DataConfig: TypeAlias = Annotated[FakeDataConfig | SFTDataConfig, Field(discriminator="type")] + + +class BaseDeploymentConfig(BaseModel): + """Base deployment config for SFT.""" + + model_config = ConfigDict(extra="forbid") + + gpus_per_node: Annotated[int, Field(description="Number of GPUs per node.")] = 8 + + +class SingleNodeDeploymentConfig(BaseDeploymentConfig): + """Configures a single-node SFT deployment.""" + + type: Literal["single_node"] = "single_node" + + num_gpus: Annotated[int, Field(description="Number of GPUs.")] = 1 + + @model_validator(mode="after") + def validate_gpu_count(self): + if self.num_gpus > self.gpus_per_node: + raise ValueError(f"num_gpus ({self.num_gpus}) exceeds gpus_per_node ({self.gpus_per_node}).") + return self + + +class MultiNodeDeploymentConfig(BaseDeploymentConfig): + """Configures a multi-node SFT deployment.""" + + type: Literal["multi_node"] = "multi_node" + + num_nodes: Annotated[int, Field(description="Number of training nodes.")] = 2 + + nodes_per_fsdp_group: Annotated[ + int | None, + Field( + description="Number of nodes per FSDP island. Auto-sets model.dp_replicate = num_nodes / nodes_per_fsdp_group." + ), + ] = None + + +SFTDeploymentConfig: TypeAlias = Annotated[ + SingleNodeDeploymentConfig | MultiNodeDeploymentConfig, Field(discriminator="type") +] + + +class SFTExperimentalConfig(BaseConfig): + """Experimental features for SFT training.""" + + +class SFTConfig(BaseConfig): + """Configures the SFT trainer""" + + # The model configuration + model: ModelConfig = ModelConfig() + + # The tokenizer configuration + tokenizer: TokenizerConfig = TokenizerConfig() + + # The data configuration + data: DataConfig = SFTDataConfig() + + # Optional validation configuration + val: SFTValConfig | None = None + + # The optimizer configuration + optim: OptimizerConfig = AdamWConfig() + + # The learning rate scheduler configuration + scheduler: SchedulerConfig = ConstantSchedulerConfig() + + # The checkpoint configuration + ckpt: CheckpointConfig | None = None + + # The logging configuration + log: TrainerLogConfig = TrainerLogConfig() + + # The wandb configuration + wandb: WandbConfig | None = None + + output_dir: Annotated[ + Path, + Field( + description="Directory to write outputs to. Will be populated with checkpoints and logs as subdirectories. Should be set to a persistent directory with enough disk space. This value should be distinct across experiments running on a single node. See the README for more details." + ), + ] = Path("outputs") + + clean_output_dir: Annotated[ + bool, + Field( + description="If true, delete the output directory before starting training. Required to overwrite an output directory that contains checkpoints from a previous run when not resuming.", + ), + ] = False + + matmul_precision: Annotated[ + Literal["highest", "high", "medium"], + Field( + description=( + "Precision for float32 matrix multiplications. " + "Use 'highest' for full FP32 (required on ROCm/AMD GPUs to avoid " + "catastrophic precision loss in softmax over large vocabularies). " + "Use 'high' to enable TF32 on NVIDIA GPUs for a speedup with minor " + "precision tradeoff. See torch.set_float32_matmul_precision docs." + ), + ), + ] = "high" + + max_steps: Annotated[ + int | None, + Field(description="Maximum number of steps to run training for. If None, will run indefinitely."), + ] = None + + memory_profiler_path: Annotated[Path | None, Field(description="Path to write memory profile to.")] = None + + bench: Annotated[ + BenchConfig | None, + Field( + description="Whether to run in benchmark mode. It will automatically set the maximum number of steps to run to 4 and use fake data.", + ), + ] = None + + gc: Annotated[ + GCConfig | None, + Field( + description="Garbage collection config. Disables automatic GC and runs deterministic collections every N steps to avoid stragglers. Set to null to use Python's default GC behavior.", + ), + ] = GCConfig() + + trace_path: Annotated[Path | None, Field(description="Path to write pytorch profiler trace to.")] = None + + dist_timeout_seconds: Annotated[ + int, + Field( + description="Timeout in seconds for torch distributed ops. Defaults to 600 seconds.", + ), + ] = 600 + + loss_impl: Annotated[ + Literal["liger", "torch", "liger_fused", "quack_fused"], + Field( + description="Implementation of the cross entropy loss function to use. " + "'liger_fused' fuses the lm_head projection with the CE loss to avoid materializing full logits. " + "'quack_fused' uses quack-kernels for chunked linear + CE with CuTe DSL CUDA kernels." + ), + ] = "torch" + + heartbeat: Annotated[ + HeartbeatConfig | None, Field(description="The heartbeat config for monitoring training progress.") + ] = None + + deployment: SFTDeploymentConfig = SingleNodeDeploymentConfig() + + slurm: Annotated[ + SlurmConfig | None, + Field( + description="SLURM configuration. If set, the run will be submitted as a SLURM job instead of running locally.", + ), + ] = None + + dry_run: Annotated[bool, Field(description="Only validate and dump resolved configs and exit early.")] = False + + experimental: Annotated[ + SFTExperimentalConfig, + Field(description="Experimental features for SFT training."), + ] = SFTExperimentalConfig() + + ### Pre-validation normalization + + @model_validator(mode="before") + @classmethod + def normalize_deployment(cls, data): + if not isinstance(data, dict): + return data + deployment = data.get("deployment") + if isinstance(deployment, dict) and deployment.get("type") == "multi_node": + for key in ("num_gpus",): + deployment.pop(key, None) + return data + + ### Validate configs (e.g. raise for unsupported (combinations of) configs) + + @model_validator(mode="after") + def deepep_disables_grad_clipping(self): + if self.model.ep_comm_backend == "deepep" and self.optim.max_norm is not None: + warnings.warn( + "Gradient clipping is not compatible with DeepEP. " + "Automatically setting optim.max_norm to None (disabled).", + stacklevel=1, + ) + self.optim.max_norm = None + return self + + @model_validator(mode="after") + def validate_deployment(self): + if self.deployment.type == "multi_node" and self.slurm is None: + raise ValueError("Must use SLURM for multi-node deployment.") + return self + + @model_validator(mode="after") + def validate_pack_function(self): + if self.model.cp > 1: + if self.data.pack_function != "cat": + raise ValueError("Packing function must be 'cat' when CP is enabled") + if self.val is not None and self.val.data.pack_function != "cat": + raise ValueError("Validation packing function must be 'cat' when CP is enabled") + return self + + @model_validator(mode="after") + def validate_cp_seq_len(self): + if self.model.cp > 1: + if self.data.seq_len % self.model.cp != 0: + raise ValueError("Sequence length must be divisible by CP degree") + if self.val is not None and self.val.data.seq_len % self.model.cp != 0: + raise ValueError("Validation sequence length must be divisible by CP degree") + return self + + @model_validator(mode="after") + def validate_cp_micro_batch_size(self): + if self.model.cp > 1: + if self.data.micro_batch_size != 1: + raise ValueError("Micro batch size must be 1 when CP is enabled") + if self.val is not None and self.val.data.micro_batch_size != 1: + raise ValueError("Validation micro batch size must be 1 when CP is enabled") + return self + + @model_validator(mode="after") + def validate_seq_len(self): + if self.data.pack_function == "stack" and self.data.seq_len % 256 != 0: + raise ValueError("The sequence length must be divisible by 256 when using pack function stack") + if self.val is not None and self.val.data.pack_function == "stack" and self.val.data.seq_len % 256 != 0: + raise ValueError("The validation sequence length must be divisible by 256 when using pack function stack") + return self + + @model_validator(mode="after") + def dont_do_massive_traces(self): + if self.trace_path: + if self.max_steps is None: + raise ValueError("Must specify max_steps when tracing") + if self.max_steps >= 10: + raise ValueError( + "Tracing more than 10 steps is not recommended as your trace will be massive. Remove this line if you really want to trace more steps." + ) + return self + + @model_validator(mode="after") + def validate_lora_adapter_saving(self): + if self.ckpt and self.ckpt.weights and self.ckpt.weights.save_adapter_separately: + lora_enabled = self.model and self.model.lora + if not lora_enabled: + raise ValueError( + "save_adapter_separately=True requires LoRA to be enabled. " + "Set model.lora or disable save_adapter_separately." + ) + return self + + @model_validator(mode="after") + def validate_opt_and_fsdp_offload(self): + if self.optim.type == "muon" and self.model.fsdp_cpu_offload: + raise ValueError("Muon optimizer does not support FSDP CPU offload") + return self + + @model_validator(mode="after") + def validate_and_disable_chunked_loss(self): + if isinstance(self.model.fused_lm_head_token_chunk_size, int): + raise ValueError( + "Chunked loss is not supported for SFT training yet, please set " + "`model.fused_lm_head_token_chunk_size` to 'disabled'" + ) + + self.model.fused_lm_head_token_chunk_size = "disabled" + return self + + @model_validator(mode="after") + def ep_only_with_custom_impl(self): + if self.model.ep > 1 and self.model.impl not in ("custom", "auto"): + raise ValueError("EP is only supported with the custom implementation or auto mode") + + return self + + ### Auto-setup and validate shared configs + + @model_validator(mode="after") + def auto_setup_bench(self): + if self.bench is not None: + self.max_steps = 4 # 1 Warmup + 3 Benchmark + if self.ckpt: # Do not checkpoint + self.ckpt = None + return self + + @model_validator(mode="after") + def auto_setup_tokenizer(self): + if self.tokenizer.name is None: + self.tokenizer.name = self.model.name + if self.tokenizer.trust_remote_code is None: + self.tokenizer.trust_remote_code = self.model.trust_remote_code + return self + + @model_validator(mode="after") + def auto_setup_deployment(self): + if self.deployment.type == "multi_node": + if self.deployment.nodes_per_fsdp_group is not None: + if self.deployment.num_nodes % self.deployment.nodes_per_fsdp_group != 0: + raise ValueError( + f"deployment.num_nodes ({self.deployment.num_nodes}) must be divisible by " + f"deployment.nodes_per_fsdp_group ({self.deployment.nodes_per_fsdp_group})" + ) + self.model.dp_replicate = self.deployment.num_nodes // self.deployment.nodes_per_fsdp_group + return self + + @model_validator(mode="after") + def auto_setup_slurm_template(self): + if self.slurm is not None and self.slurm.template_path is None: + templates_dir = find_package_resource("templates") + if templates_dir is not None: + if self.deployment.type == "single_node": + self.slurm.template_path = templates_dir / "single_node_sft.sbatch.j2" + else: + self.slurm.template_path = templates_dir / "multi_node_sft.sbatch.j2" + return self diff --git a/packages/prime-rl-configs/src/prime_rl/configs/shared.py b/packages/prime-rl-configs/src/prime_rl/configs/shared.py new file mode 100644 index 00000000000..d26c33d9a97 --- /dev/null +++ b/packages/prime-rl-configs/src/prime_rl/configs/shared.py @@ -0,0 +1,555 @@ +import os +from pathlib import Path +from typing import Annotated, Literal, TypeAlias + +from pydantic import BaseModel, Field, model_validator + +from prime_rl.utils.config import BaseConfig + + +class SlurmConfig(BaseConfig): + """Configures SLURM scheduling.""" + + job_name: Annotated[str, Field(description="The SLURM job name.")] = "prime-rl" + + project_dir: Annotated[ + Path, + Field(description="Path to the project root. Used to source .env, activate .venv, and run uv sync."), + ] = Path(".") + + template_path: Annotated[ + Path | None, + Field( + description="The path to the SLURM template file. If None, will use the default single-node/multi-node template." + ), + ] = None + + partition: Annotated[ + str, Field(description="The SLURM partition to use. Will be passed as #SBATCH --partition.") + ] = "cluster" + + nodelist: Annotated[ + str | None, + Field(description="Comma-separated list of specific nodes to run on. Passed as #SBATCH --nodelist."), + ] = None + + exclude: Annotated[ + str | None, + Field(description="Comma-separated list of nodes to exclude. Passed as #SBATCH --exclude."), + ] = None + + account: Annotated[ + str | None, + Field(description="SLURM account to charge. Passed as #SBATCH --account."), + ] = None + + time: Annotated[ + str | None, + Field(description="Maximum wall time (e.g. '24:00:00', '7-00:00:00'). Passed as #SBATCH --time."), + ] = None + + pre_run_command: Annotated[ + str | None, + Field( + description="Shell command to run on the head node before starting the job. " + "Runs after cd into project dir, .env sourcing, and venv activation. " + "Useful for cleanup routines like 'sudo pkill -f vllm'. " + "To run on all nodes, wrap with srun: 'srun bash -c \"pkill -f vllm || true\"'.", + ), + ] = None + + @property + def template_vars(self) -> dict: + """Common template variables for all SLURM templates.""" + return { + "job_name": self.job_name, + "project_dir": self.project_dir, + "partition": self.partition, + "nodelist": self.nodelist, + "exclude": self.exclude, + "account": self.account, + "time": self.time, + "pre_run_command": self.pre_run_command, + } + + @model_validator(mode="after") + def resolve_project_dir(self): + self.project_dir = self.project_dir.resolve() + return self + + +ServerType = Literal["vllm", "openai"] + + +class VLMConfig(BaseConfig): + """Configures vision-language model support. + + Presence of this config enables VLM mode. You must specify where the + vision encoder and language model live on the model object. + + Usage: + [model.vlm] + vision_encoder_attr = "model.visual" + language_model_attr = "model.language_model" + """ + + vision_encoder_attr: Annotated[ + str, + Field(description="Dotted attribute path to the vision encoder module (e.g. 'model.visual')."), + ] + + language_model_attr: Annotated[ + str, + Field(description="Dotted attribute path to the language model module (e.g. 'model.language_model')."), + ] + + freeze_vision_encoder: Annotated[ + bool, + Field( + description="Whether to freeze the vision encoder. When False, the vision encoder is trainable " + "and FSDP-sharded per-block. Has no effect with LoRA (LoRA freezes all non-adapter parameters).", + ), + ] = True + + +class BaseModelConfig(BaseConfig): + """Configures the model.""" + + name: Annotated[str, Field(description="Name or path of the HF model to use.")] = "Qwen/Qwen3-0.6B" + + trust_remote_code: Annotated[ + bool, + Field( + description="Whether to trust remote code for tokenizer initialization.", + ), + ] = False + + vlm: Annotated[ + "VLMConfig | None", + Field( + description="VLM configuration. Set this to enable vision-language model support.", + ), + ] = None + + +class RendererConfig(BaseConfig): + """Configures the client-side renderer (chat-template + response parsing). + + Only consumed when ``orchestrator.use_renderer = true``. The renderer + owns both directions: render messages → token ids on the client, and + parse model output tokens → structured ``content`` / ``reasoning_content`` + / ``tool_calls``. + """ + + name: Annotated[ + str, + Field( + description=( + "Renderer to use for chat template tokenization. " + "Options: 'auto' (detect from tokenizer), 'qwen3', 'qwen3_vl', " + "'qwen3.5', 'glm5', 'glm4.5', 'minimax-m2', 'deepseek_v3', " + "'kimi_k2', 'kimi_k25', 'nemotron3', 'gpt_oss', 'default'." + ), + ), + ] = "auto" + + tool_parser: Annotated[ + str | None, + Field( + description=( + "Name of a tool parser in renderers.parsers. Only consumed by " + "DefaultRenderer (model-specific renderers bake their own parsing " + "in). Options today: 'qwen3', 'qwen3.5', 'glm', 'deepseek_v3'." + ), + ), + ] = None + + reasoning_parser: Annotated[ + str | None, + Field( + description=( + "Name of a reasoning parser in renderers.parsers. Only consumed " + "by DefaultRenderer. Options today: 'think'." + ), + ), + ] = None + + pool_size: Annotated[ + int | None, + Field( + ge=1, + description=( + "Number of renderer slots shared across concurrent rollouts. " + "None keeps the verifiers default (1). Bump for long multi-turn " + "prompts where client-side jinja tokenization serializes." + ), + ), + ] = None + + +class ElasticConfig(BaseConfig): + """Configures elastic inference pool with DNS-based service discovery. + + Works with any DNS hostname that resolves to multiple IP addresses. + """ + + hostname: Annotated[ + str, + Field( + description="DNS hostname that resolves to inference server IPs.", + ), + ] + + port: Annotated[ + int, + Field( + description="Port that inference servers listen on.", + ), + ] = 8000 + + sync_interval: Annotated[ + float, + Field( + description="Interval in seconds between server discovery checks.", + ), + ] = 5.0 + + +class ClientConfig(BaseConfig): + """Configures the OAI client. + + Supports two modes: + - Static mode (default): Uses fixed base_url list + - Elastic mode: Uses DNS-based service discovery via hostname + + If elastic config is provided, base_url is ignored and servers are discovered dynamically. + """ + + timeout: Annotated[ + int, + Field( + description="Timeout in seconds. By default, it is set to 1200 seconds.", + ), + ] = 1200 + + connect_timeout: Annotated[ + float, + Field( + description="TCP connect timeout in seconds for inference API requests.", + ), + ] = 30.0 + + wait_for_ready_timeout: Annotated[ + int, + Field( + description="Timeout in seconds for waiting for the inference pool to become ready at startup. " + "Applies to both the static health check and elastic DNS-based discovery. Defaults to 1800 seconds.", + ), + ] = 1800 + + base_url: Annotated[ + list[str], + Field( + description="Base URLs to use for the OpenAI API. By default, it is set to a single server on localhost at port 8000 which matches the default local vLLM server configuration. If you specify more than one URL, the client will round-robin (chat) completion requests across all servers. Ignored if elastic config is provided.", + ), + ] = ["http://localhost:8000/v1"] + + api_key_var: Annotated[ + str, + Field( + description="Name of environment variable containing the API key to use for the inference API. Will parse using `os.getenv(client_config.api_key_var)`. Can be set to an arbitrary string if the inference server is not protected by an API key. If multiple URLs are specified, the same API key will be used for all servers.", + ), + ] = "VLLM_API_KEY" + + headers: Annotated[ + dict[str, str], + Field( + description="Headers to use for the OpenAI API. By default, it is set to an empty dictionary.", + ), + ] = {} + + extra_headers_from_state: Annotated[ + dict[str, str], + Field( + description="Maps HTTP header names to state field names. For each inference request, " + "the header value is dynamically read from the rollout state dict. " + 'e.g. {"X-Session-ID": "example_id"} enables sticky routing at the inference router.', + ), + ] = {} + + skip_model_check: Annotated[ + bool, + Field( + description="Whether to skip checking if the model is available in the inference pool. Useful for external APIs or API Keys that don't support the /models endpoint.", + ), + ] = False + + dp_rank_count: Annotated[ + int, + Field( + ge=1, + description=( + "Number of data-parallel ranks behind each base URL. When > 1, " + "each URL is expanded into dp_rank_count logical clients, each " + "pinned to a specific DP rank via the X-data-parallel-rank header. " + "This ensures all requests within a multi-turn rollout hit the same " + "DP engine, maximizing KV cache reuse. Auto-set from " + "inference.data_parallel_size_local (or inference.parallel.dp) " + "when using the RL entrypoint." + ), + ), + ] = 1 + + admin_base_url: Annotated[ + list[str] | None, + Field( + description="Separate base URLs for admin operations (weight updates, health checks). " + "When set, admin clients use these URLs instead of base_url, allowing weight " + "updates to bypass routers and hit each server directly. Used in disaggregated " + "P/D deployments where the inference router should not handle admin traffic.", + ), + ] = None + + elastic: Annotated[ + ElasticConfig | None, + Field( + description="Elastic inference pool configuration for DNS-based service discovery. If provided, base_url is ignored and inference servers are discovered dynamically via DNS.", + ), + ] = None + + router_url: Annotated[ + str | None, + Field( + description="URL of a vllm-router for load-aware inference routing. When set with elastic mode, inference requests go through the router while admin operations (weight updates, LoRA loading) still go directly to discovered pods.", + ), + ] = None + + @property + def is_elastic(self) -> bool: + """Check if elastic mode is enabled.""" + return self.elastic is not None + + +class LogConfig(BaseConfig): + """Configures the logger.""" + + level: Annotated[ + str, + Field( + default_factory=lambda: os.environ.get("PRIME_LOG_LEVEL", "info"), + description="Logging level for the process. Will determine the logging verbosity and format. Defaults to the PRIME_LOG_LEVEL env var if set, else 'info'.", + ), + ] + + vf_level: Annotated[ + str, + Field( + default_factory=lambda: os.environ.get("PRIME_VF_LOG_LEVEL", "info"), + description="Logging level for the verifiers package. Will determine the logging verbosity and format. Defaults to the PRIME_VF_LOG_LEVEL env var if set, else 'info'.", + ), + ] + + json_logging: Annotated[ + bool, + Field( + description="Emit JSON logs (newline-delimited) for log aggregation (Loki, Grafana, etc.).", + ), + ] = False + + log_data: Annotated[ + bool, + Field( + description="Whether to log the first data sample to the logger.", + ), + ] = False + + +class TrainerLogConfig(LogConfig): + """Trainer-specific log config.""" + + ranks_filter: Annotated[ + list[int], + Field(description="Which trainer ranks to show in console output. Passed to torchrun's --local-ranks-filter."), + ] = [0] + + +class LogExtrasConfig(BaseConfig): + """Configures extra logging for monitoring platforms.""" + + samples: Annotated[ + bool, + Field( + description="Whether to log prompt/response samples.", + ), + ] = True + + distributions: Annotated[ + bool, + Field( + description="Whether to log distributions (like rewards, advantages, etc.).", + ), + ] = True + + interval: Annotated[ + int, + Field( + ge=1, + description="Step interval at which to log extras.", + ), + ] = 10 + + sample_ratio: Annotated[ + float | None, + Field( + ge=0.0, + le=1.0, + description="Fraction of rollouts to log per step (0.0–1.0). " + "When set, the effective sample cap is len(rollouts) * sample_ratio. " + "1.0 = all rollouts, 0.5 = half, 0.0 = none. " + "None (default)", + ), + ] = None + + +class WandbConfig(BaseConfig): + """Configures logging to Weights and Biases.""" + + # Shared configs (May be overwritten by WandbConfig from `rl.py`) + project: Annotated[str, Field(description="The W&B project to log to.")] = "prime-rl" + + entity: Annotated[ + str | None, + Field( + description="The W&B entity to log to.", + ), + ] = None + + name: Annotated[ + str | None, + Field( + description="The W&B name to to use for logging.", + ), + ] = None + + group: Annotated[ + str | None, + Field( + description="The W&B group to use for logging.", + ), + ] = None + + tags: Annotated[ + list[str] | None, + Field( + description="The W&B tags to attach to the run.", + ), + ] = None + + offline: Annotated[bool, Field(description="Whether to run W&B in offline mode.")] = False + + +class WandbWithExtrasConfig(WandbConfig): + """Configures logging to Weights and Biases with extras.""" + + log_extras: Annotated[ + LogExtrasConfig | None, + Field( + description="Configuration for logging extras. If None, no extras are logged.", + ), + ] = LogExtrasConfig() + + +class PrimeMonitorConfig(BaseConfig): + """Configures logging to Prime Intellect API.""" + + base_url: Annotated[ + str, + Field( + description="The base URL for Prime Intellect monitoring API.", + ), + ] = "https://api.primeintellect.ai/api/v1/rft" + + api_key_var: Annotated[ + str, + Field( + description="Name of environment variable containing the API key for Prime Intellect API. Will parse using `os.getenv(config.api_key_var)`.", + ), + ] = "PRIME_API_KEY" + + log_extras: Annotated[ + LogExtrasConfig | None, + Field( + description="Configuration for logging extras. If None, no extras are logged.", + ), + ] = LogExtrasConfig() + + run_name: Annotated[ + str | None, + Field( + description="Name for the run shown on the platform. Defaults to the W&B run name if set, otherwise auto-generated by the platform.", + ), + ] = None + + team_id: Annotated[ + str | None, + Field( + description="Team ID to associate the run with.", + ), + ] = None + + frontend_url: Annotated[ + str | None, + Field( + description="Frontend base URL used for the dashboard link shown after registration. Defaults to the Prime CLI frontend URL when unset.", + ), + ] = None + + +class HeartbeatConfig(BaseConfig): + """Configures the heartbeat for BetterStack.""" + + url: Annotated[str, Field(description="The URL to send the heartbeat to.")] + + +class MetricsServerConfig(BaseConfig): + """Configures the Prometheus metrics server for trainer observability.""" + + port: Annotated[ + int, + Field( + ge=1, + le=65535, + description="Port to expose metrics and health endpoints. Defaults to 8000.", + ), + ] = 8000 + + host: Annotated[ + str, + Field( + description="Host to bind the server to. Defaults to 0.0.0.0.", + ), + ] = "0.0.0.0" + + +class BaseTransportConfig(BaseModel): + """Base configuration for transport.""" + + pass + + +class FileSystemTransportConfig(BaseTransportConfig): + """Configures filesystem-based transport for training examples.""" + + type: Literal["filesystem"] = "filesystem" + + +class ZMQTransportConfig(BaseTransportConfig): + """Configures ZMQ-based transport for training examples.""" + + type: Literal["zmq"] = "zmq" + host: Annotated[str, Field(description="The host address for ZMQ transport.")] = "localhost" + port: Annotated[int, Field(description="The base port for ZMQ transport.")] = 5555 + hwm: Annotated[int, Field(description="High water mark (max messages in queue) for ZMQ sockets.")] = 10 + + +TransportConfig: TypeAlias = Annotated[FileSystemTransportConfig | ZMQTransportConfig, Field(discriminator="type")] diff --git a/packages/prime-rl-configs/src/prime_rl/configs/trainer.py b/packages/prime-rl-configs/src/prime_rl/configs/trainer.py new file mode 100644 index 00000000000..a076d1e29c8 --- /dev/null +++ b/packages/prime-rl-configs/src/prime_rl/configs/trainer.py @@ -0,0 +1,1011 @@ +import warnings +from pathlib import Path +from typing import Annotated, Any, Literal, TypeAlias + +from pydantic import BaseModel, Field, model_validator + +from prime_rl.configs.shared import ( + BaseModelConfig, + FileSystemTransportConfig, + HeartbeatConfig, + MetricsServerConfig, + TrainerLogConfig, + TransportConfig, + WandbConfig, +) +from prime_rl.utils.config import BaseConfig + +# -- Shared trainer configs (used by both SFT and RL trainers) -- + +AttnImplementation: TypeAlias = Literal["eager", "sdpa", "flash_attention_2", "flash_attention_3", "fa4"] +EPCommBackend: TypeAlias = Literal["torch", "deepep"] + +# User-facing name -> internal name. Users set `flash_attention_4` in configs, +# which gets rewritten to `fa4` before pydantic validation. +# We use `fa4` internally because `flash_attention_*` triggers transformers +# to attempt installing a kernel from hub. +_ATTN_ALIASES = {"flash_attention_4": "fa4"} + + +class GCConfig(BaseConfig): + """Configures deterministic garbage collection to avoid stragglers in distributed training. + + Disables Python's automatic GC and runs manual collections every `freq` steps so all + ranks collect simultaneously, preventing one rank from stalling others. + """ + + interval: Annotated[ + int, + Field( + ge=1, + description="Run garbage collection every `interval` training steps.", + ), + ] = 50 + + +class ActivationCheckpointConfig(BaseConfig): + """Configures activation checkpointing.""" + + mode: Annotated[ + Literal["full", "selective"], + Field( + description="Whether to checkpoint whole transformer blocks (`full`) or selected subcomponents inside supported custom decoder layers (`selective`).", + ), + ] = "full" + + freq: Annotated[ + int, + Field( + ge=1, + description="Applies activation checkpointing to every `freq` layers. Defaults to 1.", + ), + ] = 1 + + targets: Annotated[ + list[str], + Field( + description="Selective checkpoint targets. `norm` checkpoints every norm module inside selected layers (decoder, attention, MLA, etc.). `attn_proj` checkpoints projection-side attention work outside the kernel, including input/output projections, attention-local norms, RoPE, gating, and model-specific MLA projection helpers where exposed. `mlp` checkpoints the entire dense MLP forward (not applicable to MoE layers). `mla_up_proj` checkpoints MLA Q/KV up-projection work where supported. `routed_experts` checkpoints routed expert compute in MoE layers (including LatentMoE). `linear_attn` checkpoints supported token mixers outside the standard softmax-attention path, including NemotronH Mamba layers, Qwen3.5-MoE GatedDeltaNet layers, and AFMoE sliding-window attention layers.", + ), + ] = ["norm"] + + @model_validator(mode="after") + def validate_selective_targets(self): + self.targets = list(dict.fromkeys(self.targets)) + if self.mode == "selective" and not self.targets: + raise ValueError("Selective activation checkpointing requires at least one target.") + return self + + +class ActivationOffloadingConfig(BaseConfig): + """Configures the activation offloading.""" + + pin_memory: Annotated[bool, Field(description="Whether to pin the offloaded activations to CPU memory.")] = True + + max_inflight_activations: Annotated[ + int, + Field( + ge=1, + description="The maximum number of activations to keep in while offloading further. (More activations means smoother overlap, but more gpu memory usage)", + ), + ] = 5 + + +class CompileConfig(BaseConfig): + """Configures model compilation.""" + + fullgraph: Annotated[ + bool, + Field(description="Whether to compile the transformer blocks with fullgraph."), + ] = False + + +class BenchConfig(BaseConfig): + """Configures benchmark mode.""" + + output_json: Annotated[ + Path | None, + Field(description="Path to write benchmark results as JSON. If not set, only prints to console."), + ] = None + + +class LoRAConfig(BaseConfig): + """Configuration for LoRA (Low-Rank Adaptation).""" + + rank: Annotated[ + int, + Field( + ge=1, + description="Rank of the low-rank decomposition matrices.", + ), + ] = 16 + + alpha: Annotated[ + float, + Field( + ge=0, + description="LoRA scaling parameter.", + ), + ] = 32.0 + + dropout: Annotated[ + float, + Field( + ge=0, + le=1, + description="LoRA dropout rate.", + ), + ] = 0.0 + + target_modules: Annotated[ + list[str], + Field( + description=( + "Module names or regex patterns for modules to apply LoRA to. Simple names (e.g., 'q_proj') " + "match any component in the module path. Regex patterns match anywhere in the name. " + "Names unknown to the current model are silently ignored, so defaults cover multiple architectures. " + "NemotronH note: 'experts' matches NonGatedGroupedExperts inside LatentMoE; 'fc1_latent_proj' and " + "'fc2_latent_proj' adapt the latent up/down projections. Add 'in_proj'/'out_proj' to also LoRA Mamba." + ), + ), + ] = [ + "q_proj", + "k_proj", + "v_proj", + "o_proj", + "gate_proj", + "up_proj", + "down_proj", + "experts", + "fc1_latent_proj", + "fc2_latent_proj", + ] + + modules_to_save: Annotated[ + list[str], + Field( + description="Module names or regex patterns for modules to keep fully trainable (not freeze). Simple names match any component in the module path. Regex patterns match anywhere in the name.", + ), + ] = [] + + +class DebugModelConfig(BaseConfig): + """Debugging feature around model and distributed training.""" + + num_layers: Annotated[ + int | None, + Field(description="The number of layers in the model."), + ] = None + + random_init: Annotated[ + bool, + Field( + description="Whether to random initialize the model.", + ), + ] = False + + force_balanced_routing: Annotated[ + bool, + Field( + description="If True, override MoE token-choice routing with a round-robin assignment so every expert receives an equal share of tokens. Intended for fake-data smoke tests where untrained routing would otherwise produce severe expert imbalance and OOM. Gating scores are still gathered from the override indices so the forward pass stays consistent.", + ), + ] = False + + +class ModelConfig(BaseModelConfig): + """Configures the model for training.""" + + seq_len: Annotated[int, Field(description="The sequence length to use for the model.")] = 2048 + + attn: Annotated[ + AttnImplementation, + Field( + description="The attention implementation to use. When CP is enabled, ring attention uses the matching kernel family (FA2 for flash_attention_2, FA3 for flash_attention_3, FA4 for fa4).", + ), + ] = "flash_attention_2" + + compile: Annotated[ + CompileConfig | None, + Field( + description="Whether to compile the model using `torch.compile`.", + ), + ] = None + + ac: Annotated[ + ActivationCheckpointConfig | None, + Field( + description="Whether to apply activation checkpointing to the model. If None, will not apply activation checkpointing.", + ), + ] = None + + ac_offloading: Annotated[ + ActivationOffloadingConfig | None, + Field( + description="Whether to apply activation offloading to the model. If None, will not apply activation offloading.", + ), + ] = None + + fsdp_cpu_offload: Annotated[ + bool, + Field( + description="Whether to enable FSDP CPU offloading for parameters, gradients, and optimizer states. When enabled, uses pinned memory for efficient CPU-GPU transfers.", + ), + ] = False + + optim_cpu_offload: Annotated[ + bool, + Field( + description="Whether to enable optimizer state CPU offloading. Unlike fsdp_cpu_offload, this only moves optimizer states (momentum, variance) to CPU, keeping weights on GPU. This avoids the H2D all-gather overhead while still saving GPU memory.", + ), + ] = False + + reshard_after_forward: Annotated[ + bool, Field(description="Whether to reshard the model after each forward pass.") + ] = True + + dp_replicate: Annotated[ + int, + Field( + description="The data parallel dim where model weights are replicated.", + ), + ] = 1 + + ep: Annotated[ + int, + Field( + description="The expert parallelism to use if the model has MoE layers. If 1, then no EP will be used.", + ), + ] = 1 + + ep_comm_backend: Annotated[ + EPCommBackend, + Field( + description=( + "Communication backend for expert parallelism. " + "`torch` uses TorchTitan all-to-all collectives and `deepep` uses DeepEP custom kernels." + ), + ), + ] = "torch" + + deepep_num_sms: Annotated[ + int, + Field( + ge=1, + description=( + "Number of SMs to allocate for DeepEP intranode dispatch/combine kernels. " + "Also determines internode RDMA channel count (num_channels = num_sms / 2). " + "Lower values leave more SMs for compute; higher values speed up dispatch/combine. " + "The optimal value depends on the EP degree and hardware." + "Only used when ep_comm_backend='deepep'." + ), + ), + ] = 20 + + deepep_token_chunk_size: Annotated[ + int | None, + Field( + ge=1, + description=( + "Optional token chunk size for DeepEP MoE pipelining. " + "When set, DeepEP dispatch for chunk i+1 is launched while experts compute chunk i. " + "Only used when ep_comm_backend='deepep'." + ), + ), + ] = None + + cp: Annotated[ + int, + Field( + description="The context parallelism size to use. If 1, then no CP will be used.", + ), + ] = 1 + + cp_style: Annotated[ + Literal["ring", "ulysses"], + Field( + description=( + "Context parallelism communication style. " + "'ring' uses ring-attention all-gather/reduce-scatter (current default, " + "requires custom kernels per attention type). " + "'ulysses' uses all-to-all to redistribute Q/K/V from sequence-sharded " + "to head-sharded, runs vanilla attention locally on the full sequence, " + "then all-to-all back. Works out-of-the-box with any attention kernel " + "(softmax FA, linear attention, mamba, etc.)." + ), + ), + ] = "ring" + + impl: Annotated[ + Literal["hf", "custom", "auto"], + Field( + description=( + "Model implementation to use. 'auto' (default) selects 'custom' if supported by the model, " + "otherwise 'hf'." + ), + ), + ] = "auto" + + optimization_dtype: Annotated[ + Literal["bfloat16", "float32"], + Field( + description="The dtype to use for the model optimization.", + ), + ] = "float32" + + reduce_dtype: Annotated[ + Literal["bfloat16", "float32"], + Field( + description="The dtype to use for the model reduce.", + ), + ] = "float32" + + moe_use_grouped_mm: Annotated[ + bool, + Field( + description="Whether to use grouped mm for the MoE layers. Require compute capability >= 9.0", + ), + ] = True + + fp8: Annotated[ + bool, + Field( + description="Whether to use FP8 training via DeepGEMM. Replaces nn.Linear layers with FP8 blockwise linear " + "and uses FP8 grouped GEMM for MoE experts. Requires SM90 (Hopper) GPUs and model.impl='custom'.", + ), + ] = False + + freeze_moe_router: Annotated[ + bool, + Field( + description="Whether to freeze the MoE router parameters during training.", + ), + ] = False + + lora: Annotated[ + LoRAConfig | None, + Field( + description="Whether to apply LoRA to the model. If None, will not apply LoRA.", + ), + ] = None + + debug: Annotated[ + DebugModelConfig, + Field( + description="Debugging feature around model and distributed training.", + ), + ] = DebugModelConfig() + + fused_lm_head_token_chunk_size: Annotated[ + int | Literal["auto", "disabled"], + Field( + description=( + "The flattened token chunk size to use for the fused LM head. " + "Three behaviors: " + "(1) int >= 1: explicitly set the number of tokens per LM-head chunk; " + "(2) 'auto': auto-enable (RL training auto-sets to 8192); " + "(3) 'disabled': explicitly disable fused LM head (use vanilla). " + "Explicitly setting an integer value for this feature isn't supported for SFT training." + ), + ), + ] = "disabled" + + @model_validator(mode="before") + @classmethod + def _normalize_attn_alias(cls, data): + """Rewrite user-facing `flash_attention_4` to internal `fa4` before validation.""" + if isinstance(data, dict) and data.get("attn") in _ATTN_ALIASES: + data["attn"] = _ATTN_ALIASES[data["attn"]] + return data + + @model_validator(mode="after") + def trust_remote_code_only_with_hf(self): + """Trust remote code only if the model is from HF.""" + if self.trust_remote_code: + if self.impl not in ("hf", "auto"): + raise ValueError("Trust remote code is only supported with the HF implementation or auto mode.") + return self + + @model_validator(mode="after") + def cp_only_with_flash_attn(self): + if self.cp > 1 and self.attn not in ["flash_attention_2", "flash_attention_3", "fa4"]: + raise ValueError("CP is only supported with flash attention 2, flash attention 3, or fa4") + if self.cp > 1 and self.attn in ("flash_attention_3", "fa4") and self.impl != "custom": + # Both ring and ulysses route FA3/FA4 through our custom FlashAttention class: + # ring patches `_compute_attention` with the ring kernel, ulysses patches it with + # the all-to-all wrapper around the FA3/FA4 kernel. The HF path patches + # `_flash_attention_forward` which only wraps FA2. + raise ValueError( + f"CP with {self.attn} requires model.impl='custom' " + "(FA3/FA4 paths are only implemented for the custom model attention class)" + ) + return self + + @model_validator(mode="after") + def ac_offloading_requires_ac(self): + """Automatically enable activation checkpointing when activation offloading is enabled.""" + if self.ac_offloading is not None and self.ac is None: + self.ac = ActivationCheckpointConfig() + return self + + @model_validator(mode="after") + def selective_ac_only_with_custom_impl(self): + if self.ac is not None and self.ac.mode == "selective" and self.impl not in ("custom", "auto"): + raise ValueError("Selective activation checkpointing requires model.impl='custom' or 'auto'") + return self + + @model_validator(mode="after") + def cpu_offload_mutual_exclusion(self): + if self.fsdp_cpu_offload and self.optim_cpu_offload: + raise ValueError("Cannot enable both fsdp_cpu_offload and optim_cpu_offload. Use one or the other.") + return self + + @model_validator(mode="after") + def flash_attention_4_only_with_custom_impl(self): + if self.attn == "fa4" and self.impl != "custom": + raise ValueError("Flash attention 4 is only supported with the custom implementation") + return self + + @model_validator(mode="after") + def fp8_only_with_custom_impl(self): + if self.fp8 and self.impl not in ("custom", "auto"): + raise ValueError("FP8 training is only supported with model.impl='custom' or 'auto'.") + return self + + @model_validator(mode="after") + def validate_ep_comm_backend(self): + if self.ep_comm_backend == "torch": + return self + + if self.ep <= 1: + raise ValueError(f"model.ep_comm_backend='{self.ep_comm_backend}' requires model.ep > 1.") + + return self + + +class TokenizerConfig(BaseConfig): + """Configuration for the tokenizer.""" + + name: Annotated[ + str | None, + Field(description="The name or path of the tokenizer to use. If None, will use the model's default tokenizer."), + ] = None + + trust_remote_code: Annotated[ + bool | None, + Field( + description="Whether to trust remote code for tokenizer initialization. If None, will use the model's default trust remote code setting.", + ), + ] = None + + chat_template: Annotated[ + str | None, + Field( + description="The chat template to use for the tokenizer. Can be a Jinja2 template string or a path to a template file. If None, will use the tokenizer's default chat template." + ), + ] = None + + +class ConstantSchedulerConfig(BaseModel): + """Configuration for constant learning rate scheduler.""" + + type: Literal["constant"] = "constant" + + +class LinearSchedulerConfig(BaseModel): + """Configuration for linear learning rate scheduler.""" + + type: Literal["linear"] = "linear" + + warmup_steps: Annotated[int, Field(ge=0, description="Number of warmup steps for the learning rate scheduler.")] = ( + 10 + ) + + decay_steps: Annotated[ + int, + Field( + ge=0, + description="Number of steps to decay the learning rate during the final portion of training.", + ), + ] = 10 + + min_lr: Annotated[float, Field(ge=0, description="Minimum learning rate to converge to.")] = 0.0 + + +class CosineSchedulerConfig(BaseModel): + """Configuration for cosine learning rate scheduler.""" + + type: Literal["cosine"] = "cosine" + + warmup_steps: Annotated[int, Field(ge=0, description="Number of warmup steps for the learning rate scheduler.")] = ( + 10 + ) + + min_lr: Annotated[float, Field(ge=0, description="Minimum learning rate to converge to.")] = 0.0 + + +SchedulerConfig: TypeAlias = Annotated[ + ConstantSchedulerConfig | LinearSchedulerConfig | CosineSchedulerConfig, Field(discriminator="type") +] + + +class BaseOptimizerConfig(BaseModel): + lr: Annotated[float, Field(ge=0)] = 1e-6 + weight_decay: Annotated[float, Field(ge=0)] = 0.01 + max_norm: Annotated[ + float | None, Field(ge=0, description="Maximum gradient norm to clip. If None, gradient clipping is disabled.") + ] = 1.0 + + +class SGDConfig(BaseOptimizerConfig): + type: Literal["sgd"] = "sgd" + nesterov: bool = True + momentum: float = 0.9 + + +class AdamWConfig(BaseOptimizerConfig): + type: Literal["adamw"] = "adamw" + + betas1: Annotated[float, Field(ge=0)] = 0.9 + betas2: Annotated[float, Field(ge=0)] = 0.999 + + +class MuonConfig(BaseOptimizerConfig): + type: Literal["muon"] = "muon" + + mu: Annotated[float, Field(ge=0, description="Momentum factor for the Muon algorithm.")] = 0.95 + betas1: Annotated[ + float, Field(ge=0, description="Beta1 for the AdamW/Lion sub-optimizer used on non-Muon params.") + ] = 0.9 + betas2: Annotated[ + float, Field(ge=0, description="Beta2 for the AdamW/Lion sub-optimizer used on non-Muon params.") + ] = 0.95 + + +class SignSGDConfig(BaseOptimizerConfig): + type: Literal["sign_sgd"] = "sign_sgd" + + +OptimizerConfig: TypeAlias = Annotated[ + SGDConfig | AdamWConfig | MuonConfig | SignSGDConfig, Field(discriminator="type") +] + + +class WeightCheckpointConfig(BaseConfig): + """Configures saving HF-compatible weight checkpoints.""" + + save_sharded: Annotated[ + bool, + Field( + description="Whether to save the weight checkpoint in sharded format.", + ), + ] = True + + save_format: Annotated[ + Literal["safetensors", "torch"], + Field( + description="The format to save the weight checkpoint in.", + ), + ] = "safetensors" + + save_adapter_separately: Annotated[ + bool, + Field( + description="Whether to save LoRA adapters separately before merging into full model weights.", + ), + ] = False + + +class CheckpointConfig(BaseConfig): + """Configures checkpointing the full model, optimizer and training state for resuming training.""" + + output_dir: Annotated[ + Path | None, + Field( + description="Override directory for checkpoints and weights. When set, checkpoints and weight snapshots are written here instead of under the trainer output_dir. Useful for writing large checkpoints to a separate storage volume.", + ), + ] = None + + interval: Annotated[ + int | None, + Field( + ge=1, + description="Interval at which to save the training checkpoint. If None, will only checkpoint at the end of training.", + ), + ] = None + + weights: WeightCheckpointConfig | None = WeightCheckpointConfig() + + skip_gather_master_weights: Annotated[ + bool, + Field( + description="When true, skip gathering and saving HF-compatible weight checkpoints. Useful for large models where the gather is expensive and only DCP checkpoints are needed.", + ), + ] = False + + weights_only: Annotated[ + bool, + Field( + description="When true, only save weight checkpoints (no optimizer/scheduler state). Much faster and smaller than full checkpoints, but cannot resume training.", + ), + ] = False + + resume_step: Annotated[ + int | None, + Field( + ge=-1, + description="Step to resume training from. If None, will start from scratch. If -1, will restart from latest checkpoint available.", + ), + ] = None + + keep_last: Annotated[ + int | None, + Field( + ge=1, + description="Keep at most this many recent step checkpoints on disk. If None, never clean old checkpoints based on recency.", + ), + ] = None + + keep_interval: Annotated[ + int | None, + Field( + ge=1, + description="Keep checkpoints at every N steps permanently (e.g., keep_interval=100 keeps step 100, 200, ...). If None, no interval-based keeping.", + ), + ] = None + + skip_progress: Annotated[ + bool, + Field( + description="Whether to skip loading the progress from checkpoint.", + ), + ] = False + + skip_scheduler: Annotated[ + bool, + Field( + description="Whether to skip loading the scheduler from checkpoint.", + ), + ] = False + + skip_dataloader: Annotated[ + bool, + Field( + description="Whether to skip loading the dataloader from checkpoint.", + ), + ] = False + + skip_optimizer: Annotated[ + bool, + Field( + description="Whether to skip loading the optimizer state from checkpoint.", + ), + ] = False + + +class DefaultLossConfig(BaseModel): + """Config for the default loss.""" + + type: Literal["default"] = "default" + + dppo_mask_low: Annotated[float, Field(ge=0, description="The low threshold for masking tokens.")] = 0.2 + dppo_mask_high: Annotated[float, Field(ge=0, description="The high threshold for masking tokens.")] = 0.2 + adv_tau: Annotated[float, Field(ge=0, description="The tau for advantages.")] = 1.0 + teacher_tau: Annotated[float, Field(ge=0, description="The tau for teacher logprobs.")] = 0.0 + kl_tau: Annotated[float, Field(ge=0, description="The tau for KL divergence.")] = 1e-3 + + +class SFTLossConfig(BaseModel): + """Config for SFT-style masked negative log-likelihood loss.""" + + type: Literal["sft"] = "sft" + + +class CustomLossConfig(BaseModel): + """Config for a custom external loss function.""" + + type: Literal["custom"] = "custom" + + import_path: Annotated[str, Field(description="Import path to the loss function (e.g., 'my_module.my_loss')")] + kwargs: Annotated[dict[str, Any], Field(default_factory=dict, description="Kwargs to pass to the loss function")] + + +LossConfig: TypeAlias = Annotated[DefaultLossConfig | SFTLossConfig | CustomLossConfig, Field(discriminator="type")] + + +class FakeDataLoaderConfig(BaseConfig): + """Configures a fake data loader sampling random micro batches for debugging.""" + + batch_size: Annotated[int, Field(ge=1)] = 2 + generate_samples: Annotated[ + bool, Field(description="Whether to generate separate samples and pack them into a single micro batch.") + ] = False + + +class DataLoaderConfig(BaseConfig): + """Configures the data loader used for training.""" + + fake: Annotated[FakeDataLoaderConfig | None, Field(description="Whether to use a fake data loader.")] = None + + +class BaseWeightBroadcastConfig(BaseModel): + """Configures the base weight broadcast.""" + + pass + + +class FileSystemWeightBroadcastConfig(BaseWeightBroadcastConfig): + """Configures the weight broadcast.""" + + type: Literal["filesystem"] = "filesystem" + save_sharded: Annotated[bool, Field(description="Whether to save the weight checkpoint in sharded format.")] = True + save_format: Annotated[ + Literal["safetensors", "torch"], Field(description="The format to save the weight checkpoint in.") + ] = "safetensors" + + +class NCCLWeightBroadcastConfig(BaseWeightBroadcastConfig): + """Configures the NCCL broadcast.""" + + type: Literal["nccl"] = "nccl" + host: Annotated[str, Field(description="The host to use for the NCCL broadcast.")] = "localhost" + port: Annotated[int, Field(description="The port to use for the NCCL broadcast.")] = 29501 + timeout: Annotated[int, Field(description="The timeout in seconds to use for the NCCL broadcast.")] = 1200 + # TODO: Should not be configurable, but auto-inferred + inference_world_size: Annotated[int, Field(description="The number of GPUs used for inference.")] = 1 + quantize_in_weight_transfer: Annotated[ + bool, + Field( + description=( + "Use kernel-format FP8 quantized NCCL transfer for weight updates. " + "When disabled, uses default HF checkpoint-format transfer." + ), + ), + ] = False + + +WeightBroadcastConfig: TypeAlias = Annotated[ + FileSystemWeightBroadcastConfig | NCCLWeightBroadcastConfig, Field(discriminator="type") +] + + +class TrainerExperimentalConfig(BaseConfig): + """Experimental features for the trainer.""" + + +class TrainerConfig(BaseConfig): + """Configures the RL trainer""" + + # The model configuration + model: ModelConfig = ModelConfig() + + # The tokenizer configuration + tokenizer: TokenizerConfig = TokenizerConfig() + + # The data configuration + data: DataLoaderConfig = DataLoaderConfig() + + # The loss configuration + loss: LossConfig = DefaultLossConfig() + + # The optimizer configuration + optim: OptimizerConfig = AdamWConfig() + + # The learning rate scheduler configuration + scheduler: SchedulerConfig = ConstantSchedulerConfig() + + # The checkpoint configuration + ckpt: CheckpointConfig | None = None + + weight_broadcast: WeightBroadcastConfig = FileSystemWeightBroadcastConfig() + + rollout_transport: TransportConfig = FileSystemTransportConfig() + + # The logging configuration + log: TrainerLogConfig = TrainerLogConfig() + + # The wandb configuration + wandb: WandbConfig | None = None + + output_dir: Annotated[ + Path, + Field( + description="Directory to write outputs to. Will be populated with checkpoints, weights, rollouts and logs as subdirectories. Should be set to a persistent directory with enough disk space. This value should be distinct across experiments running on a single node. See the README for more details." + ), + ] = Path("outputs") + + matmul_precision: Annotated[ + Literal["highest", "high", "medium"], + Field( + description=( + "Precision for float32 matrix multiplications. " + "Use 'highest' for full FP32 (required on ROCm/AMD GPUs to avoid " + "catastrophic precision loss in softmax over large vocabularies). " + "Use 'high' to enable TF32 on NVIDIA GPUs for a speedup with minor " + "precision tradeoff. See torch.set_float32_matmul_precision docs." + ), + ), + ] = "high" + + max_steps: Annotated[ + int | None, + Field( + description="Maximum number of steps to run training for. If None, will run indefinitely.", + ), + ] = None + + max_async_level: Annotated[ + int, + Field( + ge=0, + description="Maximum number of steps that inference can be ahead of training. Determines how 'off-policy' the inference engines can be. Higher values yield better throughput through async execution, but may yield lower performance. If 0, will be fully synchronous.", + ), + ] = 1 + + enable_router_replay: Annotated[ + bool, + Field( + description="Whether to enable router replay. If True, will return routed experts in the batch. This is only supported if `enable_return_routed_experts=True` in the inference config or pass `--enable-return-routed-experts` to vLLM server. This is only supported for custom models.", + ), + ] = False + + memory_profiler_path: Annotated[Path | None, Field(description="Path to write memory profile to.")] = None + + bench: Annotated[ + BenchConfig | None, + Field( + description="Whether to run in benchmark mode. It will automatically set the maximum number of steps to run to 4 and use fake data.", + ), + ] = None + + gc: Annotated[ + GCConfig | None, + Field( + description="Garbage collection config. Disables automatic GC and runs deterministic collections every N steps to avoid stragglers. Set to null to use Python's default GC behavior.", + ), + ] = GCConfig() + + trace_path: Annotated[Path | None, Field(description="Path to write pytorch profiler trace to.")] = None + + dist_timeout_seconds: Annotated[ + int, + Field( + description="Timeout in seconds for torch distributed ops. Defaults to 600 seconds.", + ), + ] = 600 + + heartbeat: Annotated[ + HeartbeatConfig | None, Field(description="The heartbeat config for monitoring training progress.") + ] = None + + metrics_server: Annotated[ + MetricsServerConfig | None, + Field(description="Prometheus metrics server config. If set, exposes /metrics endpoint for scraping."), + ] = None + + max_concurrent_runs: Annotated[ + int, + Field( + ge=1, + description="The maximum number of concurrent runs to allow. If 1, then only one run will be allowed at a time.", + ), + ] = 1 + + experimental: Annotated[ + TrainerExperimentalConfig, + Field(description="Experimental features for the trainer."), + ] = TrainerExperimentalConfig() + + @model_validator(mode="after") + def deepep_disables_grad_clipping(self): + if self.model.ep_comm_backend == "deepep" and self.optim.max_norm is not None: + warnings.warn( + "Gradient clipping is not compatible with DeepEP. " + "Automatically setting optim.max_norm to None (disabled).", + stacklevel=1, + ) + self.optim.max_norm = None + return self + + @model_validator(mode="after") + def vlms_require_bfloat16(self): + if self.model.vlm is not None and ( + self.model.optimization_dtype != "bfloat16" or self.model.reduce_dtype != "bfloat16" + ): + raise ValueError( + "VLM models must use optimization_dtype='bfloat16' and reduce_dtype='bfloat16' to match vLLM inference." + ) + return self + + @model_validator(mode="after") + def vlm_freeze_incompatible_with_lora(self): + if self.model.vlm is not None and not self.model.vlm.freeze_vision_encoder and self.model.lora is not None: + raise ValueError( + "freeze_vision_encoder=false is incompatible with LoRA. " + "LoRA freezes all non-adapter parameters including the vision encoder." + ) + return self + + @model_validator(mode="after") + def auto_setup_bench(self): + if self.bench is not None: + self.max_steps = 4 # 1 Warmup + 3 Benchmark + if not self.data.fake: + self.data.fake = FakeDataLoaderConfig() + if self.ckpt: # Do not checkpoint + self.ckpt = None + return self + + @model_validator(mode="after") + def dont_do_massive_traces(self): + if self.trace_path: + if self.max_steps is None: + raise ValueError("Must specify max_steps when tracing") + if self.max_steps >= 10: + raise ValueError( + "Tracing more than 10 steps is not recommended as your trace will be massive. Remove this line if you really want to trace more steps." + ) + return self + + @model_validator(mode="after") + def validate_lora_adapter_saving(self): + if self.ckpt and self.ckpt.weights and self.ckpt.weights.save_adapter_separately: + lora_enabled = self.model and self.model.lora + if not lora_enabled: + raise ValueError( + "save_adapter_separately=True requires LoRA to be enabled. " + "Set model.lora or disable save_adapter_separately." + ) + return self + + @model_validator(mode="after") + def validate_weight_broadcast_type(self): + if self.weight_broadcast.type == "nccl" and self.max_async_level != 1: + raise ValueError("NCCL weight broadcast only works with async level 1") + return self + + @model_validator(mode="after") + def validate_opt_and_fsdp_offload(self): + if self.optim.type == "muon" and self.model.fsdp_cpu_offload: + raise ValueError("Muon optimizer does not support FSDP CPU offload") + return self + + @model_validator(mode="after") + def validate_optim_cpu_offload_single_run(self): + if self.model.optim_cpu_offload and self.max_concurrent_runs > 1: + raise ValueError("Optimizer CPU offload is not supported with max_concurrent_runs > 1") + return self + + @model_validator(mode="after") + def validate_lora_broadcast(self): + if self.model.lora is not None and self.weight_broadcast.type == "nccl": + # TODO: Support this + raise ValueError("NCCL weight broadcast does not support LoRA yet.") + return self + + @model_validator(mode="after") + def auto_setup_tokenizer(self): + if self.tokenizer.name is None: + self.tokenizer.name = self.model.name + if self.tokenizer.trust_remote_code is None: + self.tokenizer.trust_remote_code = self.model.trust_remote_code + return self + + @model_validator(mode="after") + def auto_setup_fused_lm_head_token_chunk_size(self): + if self.model.fused_lm_head_token_chunk_size == "auto": + self.model.fused_lm_head_token_chunk_size = 8192 + + return self + + @model_validator(mode="after") + def ep_only_with_custom_impl(self): + if self.model.ep > 1 and self.model.impl not in ("custom", "auto"): + raise ValueError("EP is only supported with the custom implementation or auto mode") + + return self + + @model_validator(mode="after") + def router_replay_only_with_custom_impl(self): + if self.enable_router_replay and self.model.impl not in ("custom", "auto"): + raise ValueError("Router replay is only supported with the custom implementation or auto mode") + + return self diff --git a/packages/prime-rl-configs/src/prime_rl/utils/config.py b/packages/prime-rl-configs/src/prime_rl/utils/config.py new file mode 100644 index 00000000000..a76ead72f22 --- /dev/null +++ b/packages/prime-rl-configs/src/prime_rl/utils/config.py @@ -0,0 +1,56 @@ +from pathlib import Path +from typing import Any + +from pydantic import BaseModel +from pydantic_config import BaseConfig as BaseConfig # noqa: F401 +from pydantic_config import cli # noqa: F401 + + +def find_package_resource(subdir: str) -> Path | None: + """Find a directory contributed to the `prime_rl` namespace package by any installed wheel. + + Returns None if `subdir` is not present in any wheel — e.g. on a slim + `prime-rl-configs`-only install where `prime-rl`'s shipped resources + (templates, etc.) are absent. + """ + import prime_rl + + for p in prime_rl.__path__: + candidate = Path(p) / subdir + if candidate.is_dir(): + return candidate + return None + + +def rgetattr(obj: Any, attr_path: str) -> Any: + """Recursive getattr for dotted paths: rgetattr(cfg, "trainer.model.name").""" + current = obj + for attr in attr_path.split("."): + if not hasattr(current, attr): + raise AttributeError(f"'{type(current).__name__}' object has no attribute '{attr}'") + current = getattr(current, attr) + return current + + +def rsetattr(obj: Any, attr_path: str, value: Any) -> None: + """Recursive setattr for dotted paths: rsetattr(cfg, "trainer.model.name", "foo").""" + if "." not in attr_path: + return setattr(obj, attr_path, value) + parent_path, attr = attr_path.rsplit(".", 1) + setattr(rgetattr(obj, parent_path), attr, value) + + +def get_all_fields(model: BaseModel | type) -> list[str]: + if isinstance(model, BaseModel): + model_cls = model.__class__ + else: + model_cls = model + + fields = [] + for name, field in model_cls.model_fields.items(): + field_type = field.annotation + fields.append(name) + if field_type is not None and hasattr(field_type, "model_fields"): + sub_fields = get_all_fields(field_type) + fields.extend(f"{name}.{sub}" for sub in sub_fields) + return fields diff --git a/src/prime_rl/utils/validation.py b/packages/prime-rl-configs/src/prime_rl/utils/validation.py similarity index 61% rename from src/prime_rl/utils/validation.py rename to packages/prime-rl-configs/src/prime_rl/utils/validation.py index f05f4b4cfb0..9912cb47d79 100644 --- a/src/prime_rl/utils/validation.py +++ b/packages/prime-rl-configs/src/prime_rl/utils/validation.py @@ -2,13 +2,13 @@ from typing import Optional -from prime_rl.inference.config import InferenceConfig -from prime_rl.orchestrator.config import OrchestratorConfig -from prime_rl.trainer.rl.config import RLTrainerConfig +from prime_rl.configs.inference import InferenceConfig +from prime_rl.configs.orchestrator import OrchestratorConfig +from prime_rl.configs.trainer import TrainerConfig def validate_shared_ckpt_config( - trainer: RLTrainerConfig, + trainer: TrainerConfig, orchestrator: OrchestratorConfig, ) -> None: if trainer.ckpt and not orchestrator.ckpt: @@ -30,10 +30,19 @@ def validate_shared_ckpt_config( def validate_shared_model_name( - trainer: RLTrainerConfig, + trainer: TrainerConfig, orchestrator: OrchestratorConfig, inference: Optional[InferenceConfig] = None, ) -> None: + # Orchestrator must match inference (it queries the inference server) + if inference is not None: + if inference.model.name != orchestrator.model.name: + raise ValueError( + f"Inference model name ({inference.model.name}) and orchestrator model name ({orchestrator.model.name}) are not the same. " + "The orchestrator queries the inference server and must use the same model name." + ) + return + if trainer.model.name.startswith("Jackmin108/"): # The TT MoE models will have a different name on the orchestrator return if trainer.model.name != orchestrator.model.name: @@ -41,14 +50,9 @@ def validate_shared_model_name( f"Trainer model name ({trainer.model.name}) and orchestrator model name ({orchestrator.model.name}) are not the same. Please specify the same model name for both." ) - if inference and inference.model.name != orchestrator.model.name: - raise ValueError( - f"Inference model name ({inference.model.name}) and orchestrator model name ({orchestrator.model.name}. Please specify the same model name for both." - ) - def validate_shared_output_dir( - trainer: RLTrainerConfig, + trainer: TrainerConfig, orchestrator: OrchestratorConfig, ) -> None: if trainer.output_dir != orchestrator.output_dir.parent: @@ -58,9 +62,21 @@ def validate_shared_output_dir( def validate_shared_wandb_config( - trainer: RLTrainerConfig, + trainer: TrainerConfig, orchestrator: OrchestratorConfig, ) -> None: + if trainer.wandb and not orchestrator.wandb: + raise ValueError( + "Trainer W&B config is specified, but orchestrator W&B config is not. " + "This means only trainer metrics will be logged. Please specify [orchestrator.wandb] to log orchestrator metrics as well, " + "or use [wandb] to configure both at once." + ) + if orchestrator.wandb and not trainer.wandb: + raise ValueError( + "Orchestrator W&B config is specified, but trainer W&B config is not. " + "This means only orchestrator metrics will be logged. Please specify [trainer.wandb] to log trainer metrics as well, " + "or use [wandb] to configure both at once." + ) if trainer.wandb and orchestrator.wandb: if trainer.wandb.project != orchestrator.wandb.project: raise ValueError( @@ -69,7 +85,7 @@ def validate_shared_wandb_config( def validate_shared_max_steps( - trainer: RLTrainerConfig, + trainer: TrainerConfig, orchestrator: OrchestratorConfig, ) -> None: if trainer.max_steps != orchestrator.max_steps: @@ -79,7 +95,7 @@ def validate_shared_max_steps( def validate_shared_max_async_level( - trainer: RLTrainerConfig, + trainer: TrainerConfig, orchestrator: OrchestratorConfig, ) -> None: if trainer.max_async_level != orchestrator.max_async_level: @@ -88,8 +104,32 @@ def validate_shared_max_async_level( ) +def validate_shared_tokenizer( + trainer: TrainerConfig, + orchestrator: OrchestratorConfig, + inference: Optional[InferenceConfig] = None, +) -> None: + # Validate chat_template is consistent across all components. + # We only check chat_template (not name/trust_remote_code) because those + # are auto-derived from model names which may legitimately differ (e.g. + # when inference uses an FP8 quantized variant of the same model). + if trainer.tokenizer.chat_template != orchestrator.tokenizer.chat_template: + raise ValueError( + f"Trainer chat_template ({trainer.tokenizer.chat_template!r}) and orchestrator " + f"chat_template ({orchestrator.tokenizer.chat_template!r}) do not match. " + f"Use the shared [tokenizer] config to set chat_template for both." + ) + if inference is not None: + if trainer.tokenizer.chat_template != inference.model.chat_template: + raise ValueError( + f"Inference chat_template ({inference.model.chat_template!r}) does not match " + f"the shared tokenizer chat_template ({trainer.tokenizer.chat_template!r}). " + f"Use the shared [tokenizer] config to set chat_template for all components." + ) + + def validate_shared_weight_broadcast( - trainer: RLTrainerConfig, + trainer: TrainerConfig, orchestrator: OrchestratorConfig, inference: Optional[InferenceConfig] = None, ) -> None: diff --git a/pyproject.toml b/pyproject.toml index 301d1fefb65..83714bfdf2b 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,69 +1,74 @@ [project] name = "prime-rl" -version = "0.1.0" -description = "" +version = "0.4.0" +description = "Async RL Training at Scale" readme = "README.md" requires-python = "~=3.12.0" dependencies = [ + "prime-rl-configs", "beartype>=0.21.0", - "cydifflib>=1.2.0", "datasets>=4.0.0", "jaxtyping>=0.3.2", "liger-kernel>=0.5.10", "loguru>=0.7.3", - "numpy>=2.2.6", - "openai>=1.106.1", - "pydantic>=1.10.13", - "pydantic-settings>=2.10.1", - "pylatexenc>=2.10", - "tomli>=2.2.1", + "pyarrow>=21.0.0", + "numpy>=1.26.0", "torch>=2.9.0", - "transformers>=4.56.0", - "uvloop>=0.21.0", - "vllm==0.12.0", - "wandb>=0.20.1", - "lovely-tensors>=0.1.18", - "rich>=14.0.0", - "tomli-w>=1.2.0", - "textarena>=0.6.16", - "nltk>=3.9.1", - "math-verify>=0.8.0", + "torchvision", "torchdata>=0.11.0", - "accelerate>=1.10.1", - "blobfile>=3.0.0", - "reverse-text>=0.1.4", - "torchtitan", - "dion", - "reverse-text", - "verifiers>=0.1.8", - "prime-evals>=0.1.5", + "transformers", + "vllm>=0.17.0", + "wandb>=0.26.1", "ring-flash-attn>=0.1.8", - "prime>=0.5.0", - "tenacity>=8.2.0", + "flash-attn", + "prime>=0.5.73", "pyzmq>=27.1.0", + "aiolimiter>=1.2.1", + "tenacity>=8.2.0", + "openai>=1.106.1", + "rich>=14.0.0", + "setproctitle>=1.3.0", + "uvloop>=0.21.0", + "torchtitan", + "verifiers", + "renderers==0.1.6", + "dion", + "nvidia-ml-py>=12.575.51", ] [project.scripts] -rl = "prime_rl.rl:main" +rl = "prime_rl.entrypoints.rl:main" +sft = "prime_rl.entrypoints.sft:main" +inference = "prime_rl.entrypoints.inference:main" trainer = "prime_rl.trainer.rl.train:main" orchestrator = "prime_rl.orchestrator.orchestrator:main" -inference = "prime_rl.inference.server:main" -sft = "prime_rl.trainer.sft.train:main" -eval = "prime_rl.eval.eval:main" -synthesize = "prime_rl.synthesize.synthesize:main" +env-server = "prime_rl.orchestrator.env_server.env_server:main" + +[project.entry-points."vllm.general_plugins"] +prime_rl = "prime_rl.inference.patches:transformers_v5_compat" [project.optional-dependencies] -flash-attn = [ - "flash-attn @ https://github.com/mjun0812/flash-attention-prebuild-wheels/releases/download/v0.6.8/flash_attn-2.8.3+cu128torch2.9-cp312-cp312-linux_x86_64.whl", -] -flash-attn-3 = [ - "flash_attn_3 @ https://github.com/samsja/flash-attn-builds/releases/download/v0.1/flash_attn_3-3.0.0b1-cp39-abi3-linux_x86_64.whl", -] -metrics = [ - "prometheus_client>=0.20.0", +envs = [ + "reverse-text", + "alphabet-sort", + "wiki-search", + "math-python", + "math-env", + "code-env", + "science-env", + "logic-env", + "color-codeword", + "math500", + "aime2024", + "aime2025", + "mini_swe_agent_plus", + "deepdive" ] [dependency-groups] +mamba-ssm = [ + "mamba-ssm>=2.3.0", +] dev = [ "ipykernel>=6.29.5", "ipywidgets>=8.1.7", @@ -73,37 +78,96 @@ dev = [ ] +[tool.uv.workspace] +members = ["packages/*"] + [tool.uv] -no-build-isolation-package = ["flash-attn"] -prerelease = "allow" +exclude-newer = "7 days" +preview = true +environments = [ + "sys_platform == 'linux' and platform_machine == 'x86_64'", +] +override-dependencies = [ + "torch==2.10.0+rocm7.0", + "torchvision==0.25.0+rocm7.0", + "triton-rocm==3.6.0", + "transformers>=5.1.0.dev0", +] + +[tool.uv.exclude-newer-package] +# Self-vendored packages on our primeintellect index +reverse-text = false +alphabet-sort = false +wiki-search = false +math-python = false +math-env = false +code-env = false +science-env = false +logic-env = false +color-codeword = false +math500 = false +aime2024 = false +aime2025 = false +deepdive = false +mini_swe_agent_plus = false +# PrimeIntellect-published on PyPI (trusted publisher) +prime = false +prime-sandboxes = false +prime-tunnel = false +prime-evals = false +# Trusted git/URL sources pinned by rev or wheel URL (defensive: belt-and-suspenders) +verifiers = false +renderers = false +dion = false +pydantic-config = false [tool.uv.sources] -torch = { index = "pytorch-cu128" } -reverse-text = { index = "primeintellect" } -verifiers = { git = "https://github.com/PrimeIntellect-ai/verifiers.git", rev = "209774a" } +prime-rl-configs = { workspace = true } +torch = { index = "pytorch-rocm" } +torchvision = { index = "pytorch-rocm" } +triton-rocm = { index = "pytorch-rocm" } +vllm = { path = "/shared_silo/scratch/prime-rl/wheels/vllm-0.17.0+rocm700-cp312-cp312-linux_x86_64.whl" } +flash-attn = { path = "/shared_silo/scratch/prime-rl/wheels/flash_attn-2.8.3-cp312-cp312-linux_x86_64.whl" } +verifiers = { git = "https://github.com/PrimeIntellect-ai/verifiers.git", rev = "aa428f3" } torchtitan = { git = "https://github.com/pytorch/torchtitan", rev = "a1fdd7e" } -dion = { git = "https://github.com/samsja/dion.git", rev = "d891eebbd950a6ff11d9b5f806282e33df83df9d" } - -[tool.uv.extra-build-dependencies] -flash-attn = [{ requirement = "torch", match-runtime = true }] -flash-attn-3 = [{ requirement = "torch", match-runtime = true }] - -[tool.uv.extra-build-variables] -flash-attn = { FLASH_ATTENTION_SKIP_CUDA_BUILD = "TRUE" } +dion = { git = "https://github.com/samsja/dion.git", rev = "d891eeb" } +transformers = { git = "https://github.com/huggingface/transformers.git", rev = "c1c3424" } +pydantic-config = { git = "https://github.com/samsja/pydantic_config.git", branch = "main" } +reverse-text = { index = "primeintellect" } +alphabet-sort = { index = "primeintellect" } +wiki-search = { index = "primeintellect" } +math-python = { index = "primeintellect" } +math-env = { index = "primeintellect" } +code-env = { index = "primeintellect" } +science-env = { index = "primeintellect" } +logic-env = { index = "primeintellect" } +color-codeword = { index = "primeintellect" } +math500 = { index = "primeintellect" } +aime2024 = { index = "primeintellect" } +aime2025 = { index = "primeintellect" } +deepdive = { index = "primeintellect" } +mini_swe_agent_plus = { index = "primeintellect" } [[tool.uv.index]] name = "primeintellect" url = "https://hub.primeintellect.ai/primeintellect/simple/" +explicit = true [[tool.uv.index]] -name = "pytorch-cu128" -url = "https://download.pytorch.org/whl/test/cu128" +name = "pytorch-rocm" +url = "https://download.pytorch.org/whl/rocm7.0" explicit = true [tool.ruff.lint] select = ["F", "I"] ignore = ["F722", "F821"] # Need to ignore for jaxtyping (https://docs.kidger.site/jaxtyping/faq/) +[tool.ruff.lint.isort] +# `prime_rl` is split across two workspace wheels (prime-rl + prime-rl-configs). +# Without this, ruff classifies prime_rl.configs.* as third-party because it +# doesn't live under src/. +known-first-party = ["prime_rl"] + [tool.ruff] line-length = 120 @@ -118,5 +182,8 @@ markers = [ requires = ["hatchling"] build-backend = "hatchling.build" +[tool.hatch.build.targets.wheel] +packages = ["src/prime_rl"] + [tool.hatch.metadata] allow-direct-references = true diff --git a/scripts/docker-arm64-post-install.sh b/scripts/docker-arm64-post-install.sh new file mode 100755 index 00000000000..f02b3070b37 --- /dev/null +++ b/scripts/docker-arm64-post-install.sh @@ -0,0 +1,17 @@ +#!/bin/bash +# arm64 post-install fixups for Docker builds. +set -e + +echo "=== building flash-attn from source (sm_100 / GB200) ===" +# Run from /tmp so uv doesn't read pyproject.toml's [tool.uv.extra-build-variables] +# which sets FLASH_ATTENTION_SKIP_CUDA_BUILD=TRUE and prevents CUDA kernel compilation. +export TORCH_CUDA_ARCH_LIST="10.0" +export MAX_JOBS=4 +export FLASH_ATTENTION_FORCE_BUILD=TRUE +export FLASH_ATTENTION_SKIP_CUDA_BUILD=FALSE +(cd /tmp && uv pip install --python /app/.venv/bin/python \ + "flash-attn==2.8.3" --no-build-isolation --no-binary flash-attn --no-cache) + +echo "=== reinstalling flash-attn-cute (flash-attn overwrites it with a stub) ===" +uv pip install --reinstall --no-deps \ + "flash-attn-4 @ git+https://github.com/Dao-AILab/flash-attention.git@96bd151#subdirectory=flash_attn/cute" diff --git a/scripts/docker-entrypoint.sh b/scripts/docker-entrypoint.sh index ccb6585bef7..dfea2470480 100755 --- a/scripts/docker-entrypoint.sh +++ b/scripts/docker-entrypoint.sh @@ -4,5 +4,13 @@ set -e # Set higher ulimit for file descriptors to prevent API timeout issues ulimit -n 32000 2>/dev/null || echo "Warning: Could not set ulimit (may need --ulimit flag in docker run)" +# Allow runtime override of the verifiers package version. +# VERIFIERS_VERSION can be a git tag, branch, or commit hash. +if [ -n "$VERIFIERS_VERSION" ]; then + echo "Installing verifiers version: $VERIFIERS_VERSION" + uv pip install --reinstall-package verifiers \ + "verifiers @ git+https://github.com/PrimeIntellect-ai/verifiers.git@${VERIFIERS_VERSION}" +fi + # Execute the main command exec "$@" diff --git a/scripts/fix-flash-attn-cute.sh b/scripts/fix-flash-attn-cute.sh new file mode 100755 index 00000000000..20c092bc85d --- /dev/null +++ b/scripts/fix-flash-attn-cute.sh @@ -0,0 +1,23 @@ +#!/bin/bash +# Both `flash-attn` (FA2) and `flash-attn-cute` (FA4) ship a `flash_attn/cute/` +# sub-package. The one from `flash-attn` is a tiny stub, while the one from +# `flash-attn-cute` contains the real FA4 kernels (>1000 lines in interface.py). +# When both extras are installed, `uv sync` may install `flash-attn` *after* +# `flash-attn-cute`, causing the stub to overwrite the real module. +# +# This script reinstalls `flash-attn-cute` so the real module wins. +# Run it after `uv sync` if you have both flash-attn and flash-attn-cute extras enabled. + +set -e + +echo "Reinstalling flash-attn-cute to fix namespace conflict with flash-attn..." +uv pip install --reinstall --no-deps "flash-attn-4 @ git+https://github.com/Dao-AILab/flash-attention.git@96bd151#subdirectory=flash_attn/cute" + +# Verify installation +LINES=$(wc -l < "$(python -c 'import flash_attn.cute.interface as m; print(m.__file__)')") +if [ "$LINES" -gt 1000 ]; then + echo "Success: flash-attn-cute interface.py has $LINES lines (correct version)" +else + echo "Error: flash-attn-cute interface.py has only $LINES lines (wrong version)" + exit 1 +fi diff --git a/scripts/install.sh b/scripts/install.sh index 1be487c00cf..ab9b957d588 100755 --- a/scripts/install.sh +++ b/scripts/install.sh @@ -4,86 +4,118 @@ set -euo pipefail # Colors for output GREEN='\033[0;32m' YELLOW='\033[1;33m' +RED='\033[0;31m' NC='\033[0m' # No Color -log_info() { echo -e "${GREEN}[INFO]${NC} $*"; } -log_warn() { echo -e "${YELLOW}[WARN]${NC} $*"; } +log_info() { echo -e "${GREEN}[INFO]${NC} $*"; } +log_warn() { echo -e "${YELLOW}[WARN]${NC} $*"; } +log_error() { echo -e "${RED}[ERROR]${NC} $*" >&2; } REPO_ID="prime-rl" # Flag defaults (can be overridden via env) SKIP_CLONE=${SKIP_CLONE:-0} +assert_supported_platform() { + local os arch + os="$(uname -s)" + arch="$(uname -m)" + + if [ "$os" != "Linux" ]; then + log_error "Unsupported platform: ${os}/${arch}." + log_error "prime-rl currently supports Linux on x86_64 or aarch64 only." + exit 1 + fi + + case "$arch" in + x86_64 | aarch64) ;; + *) + log_error "Unsupported architecture: ${arch}." + log_error "prime-rl currently supports Linux on x86_64 or aarch64 only." + exit 1 + ;; + esac + + if ! command -v apt &>/dev/null; then + log_error "Unsupported Linux distribution: apt is required by this installer." + log_error "Use a Debian/Ubuntu-based Linux machine on x86_64 or aarch64, or install manually." + exit 1 + fi +} + has_ssh_access() { - # Probe SSH auth to GitHub without prompting; treat any nonzero as "no ssh" - # We try a quick ls-remote to avoid cloning on failure. - # Disable -e for the probe so the script doesn't exit on a failed test. - set +e - timeout 5s git ls-remote --heads git@github.com:PrimeIntellect-ai/${REPO_ID}.git >/dev/null 2>&1 - rc=$? - set -e - return $rc + # Probe SSH auth to GitHub without prompting; treat any nonzero as "no ssh" + # We try a quick ls-remote to avoid cloning on failure. + # Disable -e for the probe so the script doesn't exit on a failed test. + set +e + timeout 5s git ls-remote --heads git@github.com:PrimeIntellect-ai/${REPO_ID}.git >/dev/null 2>&1 + rc=$? + set -e + return $rc } ensure_known_hosts() { - # Make sure ~/.ssh exists with the right perms, then add GitHub host key. - mkdir -p "${HOME}/.ssh" - chmod 700 "${HOME}/.ssh" - # Use -H to hash hostnames; merge uniquely to avoid dupes. - if command -v ssh-keyscan >/dev/null 2>&1; then - ssh-keyscan -H github.com 2>/dev/null | sort -u \ - | tee -a "${HOME}/.ssh/known_hosts" >/dev/null + # Make sure ~/.ssh exists with the right perms, then add GitHub host key. + mkdir -p "${HOME}/.ssh" + chmod 700 "${HOME}/.ssh" + # Use -H to hash hostnames; merge uniquely to avoid dupes. + if command -v ssh-keyscan >/dev/null 2>&1; then + ssh-keyscan -H github.com 2>/dev/null | sort -u | + tee -a "${HOME}/.ssh/known_hosts" >/dev/null chmod 600 "${HOME}/.ssh/known_hosts" - fi + fi } main() { - # Ensure sudo exists - if ! command -v sudo &>/dev/null; then - apt update - apt install -y sudo - fi - - log_info "Installing base packages..." - sudo apt update && sudo apt install -y build-essential openssh-client curl git tmux htop nvtop - - log_info "Configuring SSH known_hosts for GitHub..." - ensure_known_hosts - - if [ "$SKIP_CLONE" -eq 1 ]; then - log_info "Skipping clone; assuming we are already inside the repo." + assert_supported_platform + + # Ensure sudo exists + if ! command -v sudo &>/dev/null; then + apt update + apt install -y sudo + fi + + log_info "Installing base packages..." + sudo apt update + sudo apt install -y build-essential openssh-client curl git tmux htop nvtop + + log_info "Configuring SSH known_hosts for GitHub..." + ensure_known_hosts + + if [ "$SKIP_CLONE" -eq 1 ]; then + log_info "Skipping clone; assuming we are already inside the repo." + else + log_info "Determining best way to clone (SSH vs HTTPS)..." + if has_ssh_access; then + log_info "SSH access to GitHub works. Cloning via SSH." + git clone git@github.com:PrimeIntellect-ai/${REPO_ID}.git else - log_info "Determining best way to clone (SSH vs HTTPS)..." - if has_ssh_access; then - log_info "SSH access to GitHub works. Cloning via SSH." - git clone git@github.com:PrimeIntellect-ai/${REPO_ID}.git - else - log_warn "SSH auth to GitHub not available. Cloning via HTTPS." - git clone https://github.com/PrimeIntellect-ai/${REPO_ID}.git - fi - - log_info "Entering project directory..." - cd ${REPO_ID} + log_warn "SSH auth to GitHub not available. Cloning via HTTPS." + git clone https://github.com/PrimeIntellect-ai/${REPO_ID}.git fi - log_info "Installing uv..." - curl -LsSf https://astral.sh/uv/install.sh | sh + log_info "Entering project directory..." + cd ${REPO_ID} + fi - log_info "Sourcing uv environment..." - if ! command -v uv &> /dev/null; then - source $HOME/.local/bin/env - fi + log_info "Installing uv..." + curl -LsSf https://astral.sh/uv/install.sh | sh + + log_info "Sourcing uv environment..." + if ! command -v uv &>/dev/null; then + source $HOME/.local/bin/env + fi - log_info "Installing prime..." - uv tool install prime + log_info "Installing prime..." + uv tool install prime - log_info "Syncing virtual environment..." - uv sync --all-extras + log_info "Syncing virtual environment..." + uv sync --all-extras - log_info "Installing pre-commit hooks..." - uv run pre-commit install + log_info "Installing pre-commit hooks..." + uv run pre-commit install - log_info "Installation completed!" + log_info "Installation completed!" } -main \ No newline at end of file +main diff --git a/scripts/install_deep_gemm.sh b/scripts/install_deep_gemm.sh new file mode 100755 index 00000000000..821bae0ffcf --- /dev/null +++ b/scripts/install_deep_gemm.sh @@ -0,0 +1,90 @@ +#!/usr/bin/env bash +# Build and install DeepGEMM from source for FP8 MoE inference + FP8 training. +# +# Pinned to d30fc36 — the commit our FP8 grouped-GEMM training path requires +# (newer than the 477618c that vLLM 0.17/0.19 used). The matching prebuilt +# wheel lives at tools/wheels/deep_gemm-2.3.0+d30fc36-cp312-cp312-linux_x86_64.whl +# and is what `pyproject.toml` consumes by default; this script is for +# rebuilding when torch / CUDA / vllm bumps invalidate that wheel. +# +# Requires CUDA 12.8+ and a Hopper/Blackwell GPU. +# +# Usage: +# bash scripts/install_deep_gemm.sh # build + install in-place +# bash scripts/install_deep_gemm.sh --wheel-dir tools/wheels # build wheel only +# +# Options: +# --ref REF DeepGEMM commit hash (default: d30fc36) +# --wheel-dir DIR Output wheel to DIR instead of installing + +set -euo pipefail + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +REPO_ROOT="$(dirname "$SCRIPT_DIR")" + +DEEPGEMM_GIT_REPO="https://github.com/deepseek-ai/DeepGEMM.git" +DEEPGEMM_GIT_REF="d30fc36c8f229f4f873b90a492f6e19e6e610923" +WHEEL_DIR="" + +while [[ $# -gt 0 ]]; do + case $1 in + --ref) DEEPGEMM_GIT_REF="$2"; shift 2 ;; + --wheel-dir) WHEEL_DIR="$2"; shift 2 ;; + *) echo "Unknown argument: $1" >&2; exit 1 ;; + esac +done + +# Skip if already installed at the right version +INSTALLED_VER=$(python -c "import deep_gemm; print(deep_gemm.__version__)" 2>/dev/null || echo "") +SHORT_REF="${DEEPGEMM_GIT_REF:0:7}" +if [[ "$INSTALLED_VER" == *"$SHORT_REF"* ]]; then + echo "DeepGEMM ${INSTALLED_VER} already installed, skipping." + exit 0 +fi + +# Auto-detect CUDA version +CUDA_VERSION=$(nvcc --version 2>/dev/null | grep -oP 'release \K[\d.]+' || echo "") +if [ -z "$CUDA_VERSION" ]; then + echo "ERROR: nvcc not found. CUDA toolkit required." >&2 + exit 1 +fi + +CUDA_MAJOR="${CUDA_VERSION%%.*}" +CUDA_MINOR="${CUDA_VERSION#*.}" +CUDA_MINOR="${CUDA_MINOR%%.*}" +if [ "$CUDA_MAJOR" -lt 12 ] || { [ "$CUDA_MAJOR" -eq 12 ] && [ "$CUDA_MINOR" -lt 8 ]; }; then + echo "Skipping DeepGEMM (requires CUDA 12.8+, got ${CUDA_VERSION})" + exit 0 +fi + +echo "================================================================" +echo " Building DeepGEMM (${SHORT_REF})" +echo " CUDA: ${CUDA_VERSION}" +echo "================================================================" + +TMPDIR=$(mktemp -d) +cleanup() { rm -rf "$TMPDIR"; } +trap cleanup EXIT + +git clone --recurse-submodules "$DEEPGEMM_GIT_REPO" "$TMPDIR/DeepGEMM" 2>&1 +cd "$TMPDIR/DeepGEMM" +git checkout "$DEEPGEMM_GIT_REF" +git submodule update --init --recursive + +cd "$REPO_ROOT" + +if [ -n "$WHEEL_DIR" ]; then + mkdir -p "$WHEEL_DIR" + uv build --no-build-isolation --wheel --out-dir "$WHEEL_DIR" "$TMPDIR/DeepGEMM" + echo "" + echo "Wheel built:" + ls -lh "$WHEEL_DIR"/deep_gemm*.whl +else + uv pip install --no-build-isolation --reinstall "$TMPDIR/DeepGEMM" + echo "" + echo "Installed: $(python -c 'import deep_gemm; print(deep_gemm.__version__)')" +fi + +echo "================================================================" +echo " DeepGEMM build complete" +echo "================================================================" diff --git a/scripts/install_ep_kernels.sh b/scripts/install_ep_kernels.sh new file mode 100755 index 00000000000..66071c3ec3a --- /dev/null +++ b/scripts/install_ep_kernels.sh @@ -0,0 +1,170 @@ +#!/usr/bin/env bash +# Install DeepEP kernels for expert-parallel disaggregated inference. +# +# Builds NVSHMEM + DeepEP from source and installs into the project venv. +# Required for disaggregated prefill/decode deployments with expert parallelism. +# +# Auto-detects the CUDA toolkit version that matches the installed PyTorch +# and looks for it at /usr/local/cuda-X.Y. Fails if not found. +# +# Prerequisites: +# - Matching CUDA toolkit installed (e.g. `sudo apt install cuda-toolkit-12-8`) +# - For multi-node: IBGDA driver configured (see --configure-drivers) +# +# Usage: +# bash scripts/install_ep_kernels.sh +# +# Options: +# --workspace DIR Build directory (default: ./ep_kernels_workspace) +# --deepep-ref REF DeepEP commit hash (default: 73b6ea4) +# --nvshmem-ver VER NVSHMEM version (default: 3.3.24) +# --configure-drivers Also configure IBGDA drivers (requires sudo, needs reboot) + +set -euo pipefail + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +REPO_ROOT="$(dirname "$SCRIPT_DIR")" + +DEEPEP_COMMIT_HASH="73b6ea4" +NVSHMEM_VER="3.3.24" +WORKSPACE="$REPO_ROOT/ep_kernels_workspace" +CONFIGURE_DRIVERS=0 + +while [[ $# -gt 0 ]]; do + case $1 in + --workspace) WORKSPACE="$2"; shift 2 ;; + --deepep-ref) DEEPEP_COMMIT_HASH="$2"; shift 2 ;; + --nvshmem-ver) NVSHMEM_VER="$2"; shift 2 ;; + --configure-drivers) CONFIGURE_DRIVERS=1; shift ;; + *) echo "Unknown argument: $1" >&2; exit 1 ;; + esac +done + +# Skip if DeepEP is already installed +if python -c "import deep_ep" 2>/dev/null; then + echo "DeepEP already installed, skipping." + exit 0 +fi + +# ── Auto-detect CUDA toolkit matching torch ─────────────────────────────────── +TORCH_CUDA_VER=$(python -c "import torch; print(torch.version.cuda)") +CUDA_MAJOR_MINOR=$(echo "$TORCH_CUDA_VER" | grep -oP '^\d+\.\d+') +CUDA_MAJOR=$(echo "$CUDA_MAJOR_MINOR" | cut -d. -f1) + +CUDA_HOME="/usr/local/cuda-${CUDA_MAJOR_MINOR}" +if [ ! -x "$CUDA_HOME/bin/nvcc" ]; then + echo "ERROR: Could not find CUDA toolkit matching torch (cuda ${TORCH_CUDA_VER}) at ${CUDA_HOME}" >&2 + echo "Install it with: sudo apt install cuda-toolkit-${CUDA_MAJOR_MINOR//./-}" >&2 + exit 1 +fi +export CUDA_HOME + +# Verify versions actually match +NVCC_VER=$("$CUDA_HOME/bin/nvcc" --version | grep -oP 'release \K[\d.]+') +echo "Torch CUDA: ${TORCH_CUDA_VER}, nvcc: ${NVCC_VER} (${CUDA_HOME})" + +# ── Auto-detect GPU architecture ────────────────────────────────────────────── +GPU_NAME=$(nvidia-smi --query-gpu=gpu_name --format=csv,noheader,nounits -i 0 2>/dev/null) +case "$GPU_NAME" in + *H100*|*H200*|*H800*) TORCH_CUDA_ARCH_LIST="9.0" ;; + *B100*|*B200*|*GB200*) TORCH_CUDA_ARCH_LIST="10.0" ;; + *) + echo "Could not auto-detect GPU arch from '$GPU_NAME'. Set TORCH_CUDA_ARCH_LIST manually." >&2 + echo " Hopper (H100/H200): TORCH_CUDA_ARCH_LIST=9.0" >&2 + echo " Blackwell (B200): TORCH_CUDA_ARCH_LIST=10.0" >&2 + exit 1 + ;; +esac +export TORCH_CUDA_ARCH_LIST + +echo "================================================================" +echo " Installing DeepEP kernels" +echo " CUDA: ${CUDA_HOME} (${NVCC_VER})" +echo " GPU: ${GPU_NAME} (arch ${TORCH_CUDA_ARCH_LIST})" +echo " DeepEP: ${DEEPEP_COMMIT_HASH}" +echo " NVSHMEM: ${NVSHMEM_VER}" +echo " Workspace: ${WORKSPACE}" +echo "================================================================" + +mkdir -p "$WORKSPACE" + +# ── Step 1: Install build deps ──────────────────────────────────────────────── +echo "" +echo "--- Installing build dependencies ---" +cd "$REPO_ROOT" +uv pip install cmake ninja + +# ── Step 2: Download and extract NVSHMEM ────────────────────────────────────── +echo "" +echo "--- Setting up NVSHMEM ${NVSHMEM_VER} ---" + +ARCH=$(uname -m) +case "${ARCH,,}" in + x86_64|amd64) NVSHMEM_SUBDIR="linux-x86_64" ;; + aarch64|arm64) NVSHMEM_SUBDIR="linux-sbsa" ;; + *) echo "Unsupported architecture: ${ARCH}" >&2; exit 1 ;; +esac + +NVSHMEM_DIR="$WORKSPACE/nvshmem" +if [ ! -d "$NVSHMEM_DIR/lib" ]; then + NVSHMEM_FILE="libnvshmem-${NVSHMEM_SUBDIR}-${NVSHMEM_VER}_cuda${CUDA_MAJOR}-archive.tar.xz" + NVSHMEM_URL="https://developer.download.nvidia.com/compute/nvshmem/redist/libnvshmem/${NVSHMEM_SUBDIR}/${NVSHMEM_FILE}" + + echo "Downloading ${NVSHMEM_URL}" + curl -fSL "${NVSHMEM_URL}" -o "$WORKSPACE/${NVSHMEM_FILE}" + tar -xf "$WORKSPACE/${NVSHMEM_FILE}" -C "$WORKSPACE" + mv "$WORKSPACE/${NVSHMEM_FILE%.tar.xz}" "$NVSHMEM_DIR" + rm -f "$WORKSPACE/${NVSHMEM_FILE}" + rm -rf "$NVSHMEM_DIR/lib/bin" "$NVSHMEM_DIR/lib/share" + echo "NVSHMEM extracted to ${NVSHMEM_DIR}" +else + echo "NVSHMEM already present at ${NVSHMEM_DIR}, skipping download" +fi + +export CMAKE_PREFIX_PATH="${NVSHMEM_DIR}/lib/cmake:${CMAKE_PREFIX_PATH:-}" +export NVSHMEM_DIR + +# ── Step 3: Build and install DeepEP ────────────────────────────────────────── +echo "" +echo "--- Building DeepEP (${DEEPEP_COMMIT_HASH}) ---" + +DEEPEP_DIR="$WORKSPACE/DeepEP" +if [ ! -d "$DEEPEP_DIR/.git" ]; then + git clone https://github.com/deepseek-ai/DeepEP "$DEEPEP_DIR" +fi + +cd "$DEEPEP_DIR" +git fetch origin +git checkout "$DEEPEP_COMMIT_HASH" + +WHEEL_DIR="$REPO_ROOT/deps" +mkdir -p "$WHEEL_DIR" +python setup.py bdist_wheel --dist-dir "$WHEEL_DIR" + +WHEEL=$(ls "$WHEEL_DIR"/deep_ep*.whl | head -1) +echo "" +echo "--- DeepEP wheel built at: $WHEEL ---" + +# ── Step 4 (optional): Configure IBGDA drivers ─────────────────────────────── +if [ "$CONFIGURE_DRIVERS" -eq 1 ]; then + echo "" + echo "--- Configuring IBGDA drivers (requires sudo) ---" + echo 'options nvidia NVreg_EnableStreamMemOPs=1 NVreg_RegistryDwords="PeerMappingOverride=1;"' | sudo tee -a /etc/modprobe.d/nvidia.conf + if command -v update-initramfs &> /dev/null; then + sudo update-initramfs -u + elif command -v dracut &> /dev/null; then + sudo dracut --force + else + echo "No supported initramfs update tool found." >&2 + exit 1 + fi + echo "" + echo "IBGDA configured. Please REBOOT the system to apply changes." +fi + +echo "" +echo "================================================================" +echo " DeepEP installation complete" +echo " NVSHMEM: ${NVSHMEM_DIR}" +echo " To verify: uv run python -c 'import deep_ep; print(deep_ep.__file__)'" +echo "================================================================" diff --git a/scripts/install_nixl_from_source.sh b/scripts/install_nixl_from_source.sh new file mode 100755 index 00000000000..dacb3d003d3 --- /dev/null +++ b/scripts/install_nixl_from_source.sh @@ -0,0 +1,71 @@ +#!/bin/bash +set -euo pipefail + +# Build UCX 1.19.x with CUDA + IB support, then build NIXL against it. +# Installs the unrepaired wheel (no auditwheel) so NIXL uses the UCX libs +# directly from third_party/ucx/ at runtime via LD_LIBRARY_PATH. + +SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)" +PROJECT_DIR="$(cd "$SCRIPT_DIR/.." && pwd)" +VENV_BIN="$PROJECT_DIR/.venv/bin" +PYTHON="$VENV_BIN/python" + +WORKSPACE="$PROJECT_DIR/nixl_workspace" +mkdir -p "$WORKSPACE" +UCX_SRC="$WORKSPACE/ucx_source" +UCX_INSTALL="$PROJECT_DIR/third_party/ucx" +NIXL_SRC="$WORKSPACE/nixl_source" +NIXL_VERSION="${NIXL_VERSION:-0.10.1}" +CUDA_PATH="${CUDA_HOME:-/usr/local/cuda}" +NPROC=$(nproc) + +export PATH="$VENV_BIN:$PATH" + +echo "=== Building UCX 1.19.x with CUDA + IB ===" +if [ ! -d "$UCX_SRC" ]; then + git clone https://github.com/openucx/ucx.git "$UCX_SRC" +fi +cd "$UCX_SRC" +git checkout v1.19.x + +if [ ! -f "$UCX_INSTALL/lib/libucs.so" ]; then + ./autogen.sh + ./configure \ + --prefix="$UCX_INSTALL" \ + --enable-shared \ + --disable-static \ + --disable-doxygen-doc \ + --enable-optimizations \ + --enable-cma \ + --enable-devel-headers \ + --enable-mt \ + --with-verbs \ + --with-cuda="$CUDA_PATH" \ + --with-ze=no + make -j"$NPROC" + make install + echo "=== UCX installed to $UCX_INSTALL ===" +else + echo "=== UCX already built, skipping ===" +fi + +echo "=== Building NIXL $NIXL_VERSION ===" +if [ ! -d "$NIXL_SRC" ]; then + git clone https://github.com/ai-dynamo/nixl.git "$NIXL_SRC" +else + cd "$NIXL_SRC" && git fetch --tags +fi +cd "$NIXL_SRC" +git checkout "$NIXL_VERSION" + +export PKG_CONFIG_PATH="$UCX_INSTALL/lib/pkgconfig" +export LD_LIBRARY_PATH="$UCX_INSTALL/lib:$UCX_INSTALL/lib/ucx:${LD_LIBRARY_PATH:-}" + +# Build and install directly (no auditwheel) so NIXL links to our UCX at runtime +WHEEL_DIR="$PROJECT_DIR/deps" +mkdir -p "$WHEEL_DIR" +uv pip install pip 2>/dev/null +"$PYTHON" -m pip wheel . --no-deps --wheel-dir="$WHEEL_DIR" + +WHEEL=$(ls "$WHEEL_DIR"/nixl*.whl | head -1) +echo "=== NIXL wheel built at: $WHEEL ===" diff --git a/scripts/mini_moe.py b/scripts/mini_moe.py new file mode 100644 index 00000000000..da23fdc277f --- /dev/null +++ b/scripts/mini_moe.py @@ -0,0 +1,274 @@ +"""Create and verify a mini MoE model for testing. + +Creates a small MoE model with random weights, saves it with a tokenizer, +and verifies the HF <-> PrimeRL weight conversion roundtrip. + +Usage: + # Create and verify + uv run python scripts/mini_moe.py --arch glm4_moe --output-dir ./mini-glm-moe + + # Verify only (on an existing checkpoint) + uv run python scripts/mini_moe.py --arch glm4_moe --output-dir ./mini-glm-moe --verify-only +""" + +import argparse +from pathlib import Path + +import torch +from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer +from transformers import Glm4MoeForCausalLM as HFGlm4MoeForCausalLM +from transformers.models.qwen3_5_moe.modeling_qwen3_5_moe import ( + Qwen3_5MoeForConditionalGeneration as HFQwen3_5MoeVLM, +) + +from prime_rl.trainer.models.glm4_moe import Glm4MoeConfig +from prime_rl.trainer.models.glm4_moe import Glm4MoeForCausalLM as PrimeRLGlm4MoeForCausalLM +from prime_rl.trainer.models.layers.lm_head import inject_prime_lm_head +from prime_rl.trainer.models.minimax_m2 import MiniMaxM2Config +from prime_rl.trainer.models.minimax_m2 import MiniMaxM2ForCausalLM as PrimeRLMiniMaxM2ForCausalLM +from prime_rl.trainer.models.qwen3_5_moe import Qwen3_5MoeForCausalLM as PrimeRLQwen3_5MoeVLM +from prime_rl.utils.logger import setup_logger +from prime_rl.utils.utils import default_dtype + +setup_logger("info") + + +def _qwen3_5_moe_vlm_config(): + """Build a tiny composite VLM config for Qwen3.5 MoE.""" + config = AutoConfig.from_pretrained("Qwen/Qwen3.5-35B-A3B", trust_remote_code=True, attn_implementation="sdpa") + config.use_cache = False + + tc = config.text_config + tc.vocab_size = 256 + tc.hidden_size = 256 + tc.num_hidden_layers = 2 + tc.num_attention_heads = 4 + tc.num_key_value_heads = 2 + tc.head_dim = 64 + tc.moe_intermediate_size = 128 + tc.shared_expert_intermediate_size = 128 + tc.num_experts = 4 + tc.num_experts_per_tok = 2 + tc.max_position_embeddings = 512 + tc.linear_key_head_dim = 32 + tc.linear_value_head_dim = 32 + tc.linear_num_key_heads = 4 + tc.linear_num_value_heads = 8 + tc.layer_types = ["full_attention", "linear_attention"] + tc.use_cache = False + + vc = config.vision_config + vc.depth = 2 + vc.hidden_size = 128 + vc.intermediate_size = 256 + vc.num_heads = 4 + vc.out_hidden_size = tc.hidden_size + + config.image_token_id = 250 + config.video_token_id = 251 + config.vision_start_token_id = 252 + config.vision_end_token_id = 253 + return config + + +ARCH_PRESETS = { + "glm4_moe": { + "config_class": Glm4MoeConfig, + "config_kwargs": dict( + vocab_size=151552, + hidden_size=1024, + intermediate_size=2048, + num_hidden_layers=24, + num_attention_heads=16, + num_key_value_heads=4, + hidden_act="silu", + max_position_embeddings=131072, + rms_norm_eps=1e-5, + rope_theta=1000000, + attention_bias=True, + partial_rotary_factor=0.5, + moe_intermediate_size=256, + n_routed_experts=8, + num_experts_per_tok=4, + n_shared_experts=1, + first_k_dense_replace=1, + norm_topk_prob=True, + use_qk_norm=False, + use_grouped_mm=False, + pad_token_id=151329, + eos_token_id=[151329, 151336, 151338], + ), + "hf_model_class": HFGlm4MoeForCausalLM, + "prime_model_class": PrimeRLGlm4MoeForCausalLM, + "tokenizer_source": "THUDM/GLM-4-9B-0414", + }, + "minimax_m2": { + "config_class": MiniMaxM2Config, + "config_kwargs": dict( + vocab_size=200064, + hidden_size=512, + intermediate_size=256, + num_hidden_layers=12, + num_attention_heads=8, + num_key_value_heads=4, + head_dim=64, + hidden_act="silu", + max_position_embeddings=4096, + rms_norm_eps=1e-6, + rope_theta=5000000, + rotary_dim=32, + num_local_experts=8, + num_experts_per_tok=4, + scoring_func="sigmoid", + use_routing_bias=True, + use_qk_norm=True, + qk_norm_type="per_layer", + use_grouped_mm=False, + auto_map={"AutoModelForCausalLM": "MiniMaxAI/MiniMax-M2.1--modeling_minimax_m2.MiniMaxM2ForCausalLM"}, + ), + "hf_model_class": None, # uses AutoModelForCausalLM with trust_remote_code + "prime_model_class": PrimeRLMiniMaxM2ForCausalLM, + "tokenizer_source": "MiniMaxAI/MiniMax-M2.1", + }, + "qwen3_5_moe_vlm": { + "config_fn": _qwen3_5_moe_vlm_config, + "hf_model_class": HFQwen3_5MoeVLM, + "prime_model_class": PrimeRLQwen3_5MoeVLM, + "tokenizer_source": "Qwen/Qwen3.5-35B-A3B", + "is_vlm": True, + }, + # glm_moe_dsa: HF implementation is incorrect, not supported here +} + + +def _create_hf_model(preset, config): + """Create an HF model from a preset and config.""" + hf_cls = preset["hf_model_class"] + if hf_cls is not None: + return hf_cls(config) + return AutoModelForCausalLM.from_config(config, trust_remote_code=True) + + +def _load_hf_model(preset, model_dir, config): + """Load an HF model from a preset and directory.""" + hf_cls = preset["hf_model_class"] + if hf_cls is not None: + return hf_cls.from_pretrained(str(model_dir), config=config) + return AutoModelForCausalLM.from_pretrained(str(model_dir), config=config, trust_remote_code=True) + + +def _create_hf_model_from_config(preset, config): + """Create an empty HF model from config (for roundtrip verification).""" + hf_cls = preset["hf_model_class"] + if hf_cls is not None: + return hf_cls._from_config(config) + return AutoModelForCausalLM.from_config(config, trust_remote_code=True) + + +def _build_config(preset): + """Build model config from preset (handles both config_class and config_fn styles).""" + if "config_fn" in preset: + return preset["config_fn"]() + return preset["config_class"](**preset["config_kwargs"]) + + +def create(arch: str, output_dir: Path) -> None: + preset = ARCH_PRESETS[arch] + config = _build_config(preset) + + text_config = getattr(config, "text_config", config) + print(f"Creating mini {arch} model...") + print(f" hidden_size={text_config.hidden_size}, layers={text_config.num_hidden_layers}") + + with torch.device("cpu"): + model = _create_hf_model(preset, config) + + param_count = sum(p.numel() for p in model.parameters()) + print(f" Parameters: {param_count / 1e6:.1f}M") + + print(f" Copying tokenizer from {preset['tokenizer_source']}...") + tokenizer = AutoTokenizer.from_pretrained(preset["tokenizer_source"], trust_remote_code=True) + + output_dir.mkdir(parents=True, exist_ok=True) + model.save_pretrained(output_dir) + tokenizer.save_pretrained(output_dir) + print(f" Saved to {output_dir}") + + +def verify(arch: str, model_dir: Path) -> None: + preset = ARCH_PRESETS[arch] + is_vlm = preset.get("is_vlm", False) + print(f"Verifying HF <-> PrimeRL roundtrip for {model_dir}...") + + trust_remote_code = preset["hf_model_class"] is None + config = AutoConfig.from_pretrained(str(model_dir), trust_remote_code=trust_remote_code) + config._attn_implementation = "sdpa" + if hasattr(config, "text_config"): + config.text_config._attn_implementation = "sdpa" + + text_config = getattr(config, "text_config", config) + vocab_size = text_config.vocab_size + + hf_model = _load_hf_model(preset, model_dir, config).to(device="cuda", dtype=torch.float32) + with torch.device("cuda"), default_dtype(torch.float32): + prime_model = preset["prime_model_class"]._from_config(config) + + with torch.no_grad(): + state_dict = hf_model.state_dict() + prime_model.convert_to_prime(state_dict) + prime_model.load_state_dict(state_dict) + + inject_prime_lm_head(prime_model, chunk_size=None) + + # Use tokens in safe range (avoid special VLM token IDs) + max_token = min(vocab_size, 200) if is_vlm else vocab_size + with torch.device("cuda"), default_dtype(torch.float32): + input_ids = torch.randint(0, max_token, (1, 64)) + position_ids = torch.arange(1, 65).unsqueeze(0) + + hf_output = hf_model(input_ids=input_ids, position_ids=position_ids) + prime_output = prime_model(input_ids, position_ids) + + if is_vlm: + # HF GatedDeltaNet has a dtype bug in float32 mode; just verify non-NaN output + assert not torch.isnan(prime_output["logits"]).any(), "PrimeRL VLM output contains NaN" + assert prime_output["logits"].shape == hf_output.logits.shape + print(" VLM forward pass verified (shape match, no NaN)") + else: + logits_diff = prime_output["logits"] - hf_output.logits + max_diff = logits_diff.abs().max().item() + print(f" HF vs PrimeRL max logits diff: {max_diff:.6f}") + assert max_diff < 0.1, f"HF vs PrimeRL logits mismatch: max diff {max_diff}" + + # Roundtrip weight conversion: HF -> PrimeRL -> HF + # Normalize both through the same roundtrip to handle expert format differences + prime_cls = preset["prime_model_class"] + with torch.no_grad(): + roundtrip_sd = prime_cls.convert_to_hf(dict(prime_model.state_dict())) + orig_sd = dict(hf_model.state_dict()) + prime_cls.convert_to_prime(orig_sd) + prime_cls.convert_to_hf(orig_sd) + + for key in orig_sd: + assert key in roundtrip_sd, f"Missing key after roundtrip: {key}" + assert torch.equal(orig_sd[key], roundtrip_sd[key]), f"Roundtrip mismatch at {key}" + print(" HF -> PrimeRL -> HF weight roundtrip verified") + + print(" Verification passed.") + + +def main(): + parser = argparse.ArgumentParser(description="Create and verify a mini MoE model") + parser.add_argument("--arch", choices=list(ARCH_PRESETS.keys()), required=True) + parser.add_argument("--output-dir", type=Path, required=True) + parser.add_argument("--verify-only", action="store_true", help="Skip creation, only verify an existing model") + args = parser.parse_args() + + if not args.verify_only: + create(args.arch, args.output_dir) + + verify(args.arch, args.output_dir) + + +if __name__ == "__main__": + main() diff --git a/scripts/tmux.sh b/scripts/tmux.sh index 7823305111b..3511d774584 100755 --- a/scripts/tmux.sh +++ b/scripts/tmux.sh @@ -49,54 +49,61 @@ if [[ ${#POSITIONAL[@]} -ge 2 ]]; then OUTPUT_DIR="${POSITIONAL[1]}" fi +LOG_DIR="${OUTPUT_DIR}/logs" + if tmux has-session -t "$SESSION_NAME" 2>/dev/null; then echo "Attaching to tmux session: $SESSION_NAME" exec tmux attach-session -t "$SESSION_NAME" -else - echo "Creating new tmux session: $SESSION_NAME" - - # Start new tmux session with first window - tmux new-session -d -s "$SESSION_NAME" -n "RL" - - # Window 1: RL - 3 vertical panes - tmux split-window -v -t "$SESSION_NAME:RL.0" - tmux split-window -v -t "$SESSION_NAME:RL.1" - tmux select-layout -t "$SESSION_NAME:RL" even-vertical - - # Pane titles - tmux select-pane -t "$SESSION_NAME:RL.0" -T "Trainer" - tmux select-pane -t "$SESSION_NAME:RL.1" -T "Orchestrator" - tmux select-pane -t "$SESSION_NAME:RL.2" -T "Inference" - - # Logs: Orchestrator - tmux send-keys -t "$SESSION_NAME:RL.1" \ - "echo \"Following orchestrator.stdout (tail -F; waits for rotate/create)...\"; tail -F \"${OUTPUT_DIR}/logs/orchestrator.stdout\" 2>/dev/null" \ - C-m - - # Logs: Inference - tmux send-keys -t "$SESSION_NAME:RL.2" \ - "echo \"Following inference.stdout (tail -F; waits for rotate/create)...\"; tail -F \"${OUTPUT_DIR}/logs/inference.stdout\" 2>/dev/null" \ - C-m - - # Window 2: Monitor - tmux new-window -t "$SESSION_NAME" -n "Monitor" - tmux split-window -h -t "$SESSION_NAME:Monitor" - tmux select-layout -t "$SESSION_NAME:Monitor" even-horizontal - - tmux select-pane -t "$SESSION_NAME:Monitor.0" -T "GPU" - tmux send-keys -t "$SESSION_NAME:Monitor.0" "nvtop" C-m - - tmux select-pane -t "$SESSION_NAME:Monitor.1" -T "CPU" - tmux send-keys -t "$SESSION_NAME:Monitor.1" "htop" C-m - - # Pane title styling - tmux set-option -t "$SESSION_NAME" -g pane-border-status top - tmux set-option -t "$SESSION_NAME" -g pane-border-format " #{pane_title} " - tmux set-window-option -t "$SESSION_NAME:RL" pane-border-status top - tmux set-window-option -t "$SESSION_NAME:Monitor" pane-border-status top - - # Focus trainer pane and attach - tmux select-window -t "$SESSION_NAME:RL" - tmux select-pane -t "$SESSION_NAME:RL.0" - exec tmux attach-session -t "$SESSION_NAME" -fi \ No newline at end of file +fi + +echo "Creating new tmux session: $SESSION_NAME" + +# Window 0: Launcher - empty shell +tmux new-session -d -s "$SESSION_NAME" -n "Launcher" + +# Window 1: Logs - 4 vertical panes +tmux new-window -t "$SESSION_NAME" -n "Logs" + +tmux split-window -v -t "$SESSION_NAME:Logs.0" +tmux split-window -v -t "$SESSION_NAME:Logs.1" +tmux split-window -v -t "$SESSION_NAME:Logs.2" +tmux select-layout -t "$SESSION_NAME:Logs" even-vertical + +tmux select-pane -t "$SESSION_NAME:Logs.0" -T "Trainer" +tmux select-pane -t "$SESSION_NAME:Logs.1" -T "Orchestrator" +tmux select-pane -t "$SESSION_NAME:Logs.2" -T "Envs" +tmux select-pane -t "$SESSION_NAME:Logs.3" -T "Inference" + +tmux send-keys -t "$SESSION_NAME:Logs.0" \ + "tail -F ${LOG_DIR}/trainer.log 2>/dev/null" C-m + +tmux send-keys -t "$SESSION_NAME:Logs.1" \ + "tail -F ${LOG_DIR}/orchestrator.log 2>/dev/null" C-m + +tmux send-keys -t "$SESSION_NAME:Logs.2" \ + "tail -F ${LOG_DIR}/envs/*/*/*.log 2>/dev/null" C-m + +tmux send-keys -t "$SESSION_NAME:Logs.3" \ + "tail -F ${LOG_DIR}/inference.log 2>/dev/null" C-m + +# Window 2: Claude Code with log context +tmux new-window -t "$SESSION_NAME" -n "Claude" +tmux send-keys -t "$SESSION_NAME:Claude" \ + "claude --permission-mode auto --append-system-prompt 'You are monitoring a prime-rl training run. The output directory is ${OUTPUT_DIR}. Log paths: + Trainer: ${LOG_DIR}/trainer.log + All nodes: ${LOG_DIR}/trainer/node_*.log + All ranks: ${LOG_DIR}/trainer/torchrun/*/*/*/*.log + Orchestrator: ${LOG_DIR}/orchestrator.log + Inference: ${LOG_DIR}/inference.log + Envs: ${LOG_DIR}/envs/*/*/*.log + Train envs: ${LOG_DIR}/envs/train/*/*.log +You are running inside tmux session \"${SESSION_NAME}\". The Launcher window (window 0) is where the user runs launch commands. You can read its contents with: tmux capture-pane -t ${SESSION_NAME}:Launcher -p +Help the user monitor and debug this run.'" C-m + +# Pane title styling +tmux set-option -t "$SESSION_NAME" -g pane-border-status top +tmux set-option -t "$SESSION_NAME" -g pane-border-format " #{pane_title} " + +# Focus launcher window and attach +tmux select-window -t "$SESSION_NAME:Launcher" +exec tmux attach-session -t "$SESSION_NAME" diff --git a/skills/config/SKILL.md b/skills/config/SKILL.md new file mode 100644 index 00000000000..e8dc13216c1 --- /dev/null +++ b/skills/config/SKILL.md @@ -0,0 +1,180 @@ +--- +name: config +description: How the prime-rl config system works — TOML files, CLI, config composition, and special patterns. Use when creating configs, debugging config errors, or overriding values via CLI. +--- + +# Config + +prime-rl uses `pydantic_config` (combines `tyro` and `pydantic`) for configuration. + +## Use configs + +Every entrypoint accepts TOML files via `@` syntax and CLI overrides to configure it. + +```bash +# Configure RL training with a TOML file +uv run rl @ examples/reverse_text/rl.toml + +# Override specific fields via CLI +uv run rl @ examples/reverse_text/rl.toml --max-steps 50 +``` + +Config resolve in the following order: + +1. CLI arguments +2. Config files (merged left-to-right) +3. Class defaults (lowest) + +## Compose configs + +Multiple config files are merged left-to-right (later files override earlier ones): + +```bash +uv run rl @ examples/reverse_text/rl.toml @ examples/reverse_text/slurm_rl.toml +``` + +Nested configs can be loaded for specific sections: + +```bash +uv run rl --model @ model.toml --data @ data.toml +``` + +Mixed composition works too: + +```bash +uv run rl @ base.toml --trainer @ trainer_override.toml --trainer.lr 1e-3 +``` + +Merging is deep — unset fields in the override are preserved from the base config. + +## Inspect & validate configs + +Use `--help` to see all available fields and their defaults. When combined with a config file, defaults reflect the TOML values: + +```bash +uv run rl --help # shows class defaults +uv run rl @ examples/reverse_text/rl.toml --help # shows defaults from TOML +``` + +Use `--dry-run` to validate and dump the fully resolved config: + +```bash +uv run rl @ examples/reverse_text/rl.toml --dry-run --output-dir /tmp/test +# Writes resolved TOML to /tmp/test/configs +``` + +## Naming + +CLI uses kebab-case (`--model.max-model-len`), TOML uses snake_case (`max_model_len`). Both refer to the same field. + +## General rules + +- **Fail early**: incompatible option combinations (e.g. CP requires flash attention, NCCL broadcast requires async level 1) should raise in `model_validator` at config resolution time, not at runtime. When adding new constraints, add a validator to the config class. +- **Deprecation**: when renaming or removing config fields, emit a deprecation warning with a clear migration path (e.g. "field X is deprecated, use Y instead"). Do not silently drop fields — help users update their configs. + +## Important patterns + +### Boolean fields + +```bash +uv run inference --model.enforce-eager # sets to true +uv run inference --model.no-enforce-eager # sets to false +``` + +In TOML, booleans must be explicit: + +```toml +[model] +enforce_eager = true +``` + +### None fields + +TOML has no null type. Use the string `"None"`: + +```toml +max_model_len = "None" +``` + +On the CLI, pass `None` as a plain string: + +```bash +uv run inference --model.max-model-len None +``` + +### List fields + +In TOML, use `[[double brackets]]` (array of tables) for lists of objects: + +```toml +[[orchestrator.env]] +id = "reverse-text" + +[[orchestrator.env]] +id = "math-env" +``` + +On the CLI, list items are indexed: `--env.0.id reverse-text --env.1.id math-env`. + +### Dict fields + +In TOML, use a section: + +```toml +[vllm_extra] +key1 = "value1" +key2 = 123 +``` + +On the CLI, pass as a JSON string: + +```bash +uv run inference --vllm-extra '{"key1": "value1", "key2": 123}' +``` + +### Discriminated unions + +Some config fields use discriminated unions (e.g. loss type, data type). Set the `type` field to select the variant: + +```toml +[trainer.loss] +type = "sft" + +[data] +type = "fake" +batch_size = 2 +``` + +On the CLI: + +```bash +uv run sft --data.type fake --data.batch-size 4 +``` + +If you wish to configure values of the default variant, you don't need to set the `type` field. + +### SFT hard distill override + +For hosted multi-tenant runs where the trainer image's `trainer.loss.type` is fixed, the orchestrator exposes a per-run override that forces SFT loss on every micro-batch without rebuilding the trainer. Set `orchestrator.use_sft_loss = true` alongside `orchestrator.teacher_rollout_model`; both must be configured together (the orchestrator validator enforces this). The orchestrator stamps each `TrainingSample.sft_loss = True`, which the trainer's `compute_loss` honors by dispatching to `sft_loss_fn` per batch — independent of the trainer's configured default loss. + +### Model fields + +For `BaseModel | None` fields (like `[ckpt]`, `[wandb]`, `[compile]`), a bare flag enables them with defaults: + +```bash +uv run rl @ config.toml --model.compile # enables compilation with defaults (fullgraph = false) +uv run rl @ config.toml --model.compile.fullgraph # enables compilation and sets nested field (fullgraph = true) +``` + +In TOML, an empty section header does the same: + +```toml +[ckpt] # enables checkpointing with defaults +``` + +## Key files + +- `src/prime_rl/utils/config.py` — re-exports `BaseConfig` and `cli` from pydantic_config +- `src/prime_rl/configs/` — all domain-specific config classes +- `configs/debug/` — minimal debug configs for testing +- `examples/` — full example configs for various tasks diff --git a/skills/entrypoints/SKILL.md b/skills/entrypoints/SKILL.md new file mode 100644 index 00000000000..f6589aca976 --- /dev/null +++ b/skills/entrypoints/SKILL.md @@ -0,0 +1,91 @@ +--- +name: entrypoints +description: All available prime-rl entrypoints — what they do, how to launch them, and which config class they use. Use when running commands, launching training, or starting servers. +--- + +# Entrypoints + +All entrypoints are run via `uv run ` and accept TOML configs via `@ path/to/config.toml` with CLI overrides. See the `config` skill for config system details. + +## `rl` — RL training + +Orchestrates the complete RL loop: launches inference server, orchestrator, and trainer as subprocesses. + +```bash +uv run rl @ examples/reverse_text/rl.toml +uv run rl @ examples/reverse_text/rl.toml @ examples/reverse_text/slurm_rl.toml # with SLURM +uv run rl @ examples/reverse_text/rl.toml --dry-run # generate scripts without running +``` + +- **Config:** `RLConfig` (`src/prime_rl/configs/rl.py`) +- **Entrypoint:** `src/prime_rl/entrypoints/rl.py` +- **SLURM:** yes — single-node and multi-node + +## `sft` — SFT training + +Trains a model on labeled data. Uses torchrun for distributed training. + +```bash +uv run sft @ examples/reverse_text/sft.toml +uv run sft @ examples/reverse_text/sft.toml --slurm # with SLURM +uv run sft @ examples/reverse_text/sft.toml --dry-run # generate scripts without running +``` + +The entrypoint launches torchrun internally — no need to call torchrun directly. + +- **Config:** `SFTConfig` (`src/prime_rl/configs/sft.py`) +- **Entrypoint:** `src/prime_rl/entrypoints/sft.py` +- **SLURM:** yes — single-node and multi-node + +## `inference` — Standalone inference server + +Launches a vLLM-based inference server with OpenAI-compatible API. + +```bash +uv run inference @ configs/debug/infer.toml +uv run inference --model.name Qwen/Qwen3-0.6B --model.enforce-eager +``` + +Always use the `inference` entrypoint — never `vllm serve` directly. + +Custom endpoints beyond standard OpenAI API: +- `/v1/chat/completions/tokens` — accepts token IDs as prompt input +- `/update_weights` — hot-reload model weights from the trainer +- `/load_lora_adapter` — load LoRA adapters at runtime +- `/init_broadcaster` — initialize weight broadcast for distributed training + +Check health with: +```bash +curl http://:/health +``` + +Check served models with: +```bash +curl http://:/v1/models +``` + +Test chat completions with: +```bash +curl http://localhost:8000/v1/chat/completions \ + -H "Content-Type: application/json" \ + -d '{"model": "Qwen/Qwen3-0.6B", "messages": [{"role": "user", "content": "Hi"}], "max_tokens": 50}' +``` + +- **Config:** `InferenceConfig` (`src/prime_rl/configs/inference.py`) +- **Entrypoint:** `src/prime_rl/entrypoints/inference.py` +- **SLURM:** yes — single-node, multi-node, and disaggregated deployments + +## Summary + +| Command | Purpose | SLURM | Typical use | +|---------|---------|-------|-------------| +| `rl` | Full RL pipeline | yes | Production RL training | +| `sft` | Supervised fine-tuning | yes | SFT training | +| `inference` | vLLM server | yes | Standalone inference or debugging | + +## Key directories + +- `src/prime_rl/entrypoints/` — top-level entrypoints (`rl`, `sft`, `inference`) +- `src/prime_rl/configs/` — all config classes +- `configs/debug/` — minimal configs for quick testing +- `examples/` — full example configs for various tasks diff --git a/skills/installation/SKILL.md b/skills/installation/SKILL.md new file mode 100644 index 00000000000..c0268e8a483 --- /dev/null +++ b/skills/installation/SKILL.md @@ -0,0 +1,73 @@ +--- +name: installation +description: How to install prime-rl and its optional dependencies. Use when setting up the project, installing extras like deep-gemm for FP8 models, or troubleshooting dependency issues. +--- + +# Installation + +## Basic +```bash +uv sync # core dependencies only +uv sync --group dev # dev tools: pytest, ruff, pre-commit +uv sync --all-extras # recommended: includes flash-attn, flash-attn-cute, etc. +``` + +## Advanced + +### Mamba-SSM (NemotronH models) + +For NemotronH (hybrid Mamba-Transformer-MoE) models, install `mamba-ssm` for Triton-based SSD kernels that match vLLM's precision: + +```bash +CUDA_HOME=/usr/local/cuda uv pip install mamba-ssm +``` + +Requires `nvcc` (CUDA toolkit). Without `mamba-ssm`, NemotronH falls back to HF's pure-PyTorch implementation which computes softplus in bf16, causing ~0.4 KL divergence vs vLLM. + +Note: do NOT install `causal-conv1d` unless your GPU architecture matches the compiled CUDA kernels. The code automatically falls back to PyTorch nn.Conv1d when it's absent. + +### FP8 inference with deep-gemm + +For certain models like GLM-5-FP8, you need `deep-gemm`. Install it via the `fp8-inference` dependency group: + +```bash +uv sync --group fp8-inference +``` + +This installs the pre-built `deep-gemm` wheel. No CUDA build step is needed. + +## Trainer DeepEP backend + +The trainer-side MoE `deepep` backend is optional and requires a local DeepEP build. + +Install using the provided script, which auto-detects CUDA toolkit and GPU architecture: + +```bash +bash scripts/install_ep_kernels.sh +``` + +The script downloads NVSHMEM, builds DeepEP from source, and places the wheel in `deps/`. It skips if DeepEP is already installed. Options: + +- `--workspace DIR` — build directory (default: `./ep_kernels_workspace`) +- `--deepep-ref REF` — DeepEP commit hash (default: `73b6ea4`) +- `--nvshmem-ver VER` — NVSHMEM version (default: `3.3.24`) +- `--configure-drivers` — configure IBGDA drivers for multi-node (requires sudo + reboot) + +Verify the install: + +```bash +uv run python -c 'import deep_ep; print(deep_ep.__file__)' +``` + +## Dev dependencies + +```bash +uv sync --group dev +``` + +Installs pytest, ruff, pre-commit, and other development tools. + +## Key files + +- `pyproject.toml` — all dependencies, extras, and dependency groups +- `uv.lock` — pinned lockfile (update with `uv sync --all-extras`) diff --git a/skills/monitor-run/SKILL.md b/skills/monitor-run/SKILL.md new file mode 100644 index 00000000000..2a8404dffb3 --- /dev/null +++ b/skills/monitor-run/SKILL.md @@ -0,0 +1,254 @@ +--- +name: monitor-run +description: How to monitor ongoing training runs — find output directories, check logs, diagnose performance, and inspect SLURM jobs. Use when asked to check on a run, debug training issues, or investigate performance. +--- + +# Monitor RL Run + +## Runbook + +### On launch + +Immediately gather context and write a summary of the run into `{output_dir}/STATUS.md`: + +1. Identify the output directory and read the resolved configs to understand the experiment (model, envs, hyperparameters, deployment details). +2. Make sure that the run started successfully and that all processes are alive. +3. Read the logs and note the current training step and health. + +### Recurring check-ins + +After the initial overview, schedule recurring check-ins. By default, check in every **1 hour** (the researcher can override this). + +At each check-in: + +1. Check that all processes are alive. +2. Read the logs — look for errors, warnings, hangs, or degraded performance +3. Note the current training step, key metrics, and checkpoint progress. +4. **Append an entry to `{output_dir}/STATUS.md`**: + +```markdown +## YYYY-MM-DD HH:MM UTC + +**Step**: {current_step} / {max_steps} +**Health**: {Healthy | Degraded | Down} + +**Progress**: reward/mean, seq_len, truncation rate, eval scores (if available), notable env-specific metrics. +**Stability**: entropy, mismatch KL, grad norm — flag any spikes or concerning trends. +**Performance**: trainer and orchestrator step times, who is waiting on whom, env server lag, inference pressure. + +**Notes**: (anything unusual — errors, restarts, hangs, etc. Omit if nothing notable.) +``` + +Always append — never overwrite previous entries. + +### Restarting a run + +**IMPORTANT**: Never restart a run unless you were explicitly instructed by the researcher. If you were given permission, make sure to ask the researcher for the exact command to resume a run and under what conditions a restart is necessary. + +**IMPORTANT**: Never run kill or launch commands directly from your shell. Instead, send them to the tmux **Launcher** window so the researcher can see exactly what was executed. + +Use `tmux send-keys` to dispatch commands to the Launcher pane: + +```bash +# Get your current tmux session name +SESSION=$(tmux display-message -p '#S') + +# Send a command to the Launcher window +tmux send-keys -t "$SESSION:Launcher" 'your command here' Enter +``` + +After a restart, verify that all processes are back up and healthy before resuming periodic check-ins. Check the process tree and tail the logs to confirm the run is making progress again. + +--- + +## Reference + +### Output directory and tmux session + +The output directory and tmux session name are typically provided by the researcher in the appended system prompt (see `scripts/tmux.sh` — the Claude window is launched with this context). If not provided, **ask the researcher** which output directory to monitor and which tmux session the run is in. + +The tmux session contains the **Launcher** window where the researcher runs launch commands — this is where you should send any restart commands (see [Restarting a run](#restarting-a-run)). + +Once you have the output directory, the resolved configs are at `{output_dir}/configs/`. + +### Configs + +The launcher writes resolved configs as TOML files to `{output_dir}/configs/`. Read `rl.toml` to get the full picture of the experiment (model, envs, hyperparameters, wandb, deployment). + +### Logs + +Logs are usually the most informative place to monitor a run. + +``` +{output_dir}/logs/ +├── trainer.log # trainer stdout (rank 0) +├── orchestrator.log # orchestrator stdout +├── inference.log # vLLM inference server stdout +├── trainer/ +│ ├── node_*.log # per-node logs (multi-node only) +│ └── torchrun/ # per-rank stdout/stderr (all ranks) +├── inference/ +│ ├── node_*.log # per-node logs (multi-node only) +│ └── router_0.log # vllm-router per replica (multi-node only) +└── envs/ + ├── train/{env_name}/ + │ ├── env_server.log + │ └── env_worker_{id}.log + └── eval/{env_name}/ + └── ... +``` + +Usually it's sufficient to tail `trainer.log`, `orchestrator.log`, and `inference.log`. For debugging, it may be necessary to check the per-node logs (`node_*.log`) or per-rank trainer logs under `torchrun/`. + +```bash +tail {output_dir}/logs/trainer.log # training progress — kl mismatch, entropy, grad norm, trainer step time +tail {output_dir}/logs/orchestrator.log # orchestrator — reward, rollouts, env execution, inference step time +tail {output_dir}/logs/inference.log # inference — completed HTTP requests, engine stats, OOM errors +tail {output_dir}/logs/envs/train/*/env_server.log # env server — aggregated stats across its workers (lag, task distribution) +tail {output_dir}/logs/envs/train/*/env_worker_*.log # env workers — individual env logs +``` + +All logs use loguru with the format `HH:mm:ss LEVEL message`. Log levels: `DEBUG`, `INFO`, `SUCCESS`, `WARNING`, `ERROR`. To scan for problems: + +```bash +grep -E "WARNING|ERROR" {output_dir}/logs/trainer.log +grep -E "WARNING|ERROR" {output_dir}/logs/orchestrator.log +grep -E "WARNING|ERROR" {output_dir}/logs/inference.log +grep -E "WARNING|ERROR" {output_dir}/logs/envs/train/*/env_server.log +grep -E "WARNING|ERROR" {output_dir}/logs/envs/train/*/env_worker_*.log +``` + +### Metrics + +All metrics below are logged to the console. The source column indicates which log file to check. + +#### Progress + +These tell you whether the model is learning and how the run is progressing. + +| Metric | Source | Description | +|--------|--------|-------------| +| `reward/{all,env}/mean` | orchestrator | mean training reward | +| `seq_len/{all,env}/mean` | orchestrator | average sequence length in tokens | +| `num_turns/{all,env}/mean` | orchestrator | average turns per rollout (ignore for single-turn envs) | +| `is_truncated/{all,env}/mean` | orchestrator | fraction of truncated rollouts | +| `empty_rollouts/{all,env}` | orchestrator | fraction of empty rollouts | +| `errored_rollouts/{all,env}` | orchestrator | fraction of errored rollouts | +| `metrics/{env}/{metric}` | orchestrator | env-specific metrics (e.g. pass rate for unit tests) | +| `eval/{env}/{avg@k,pass@k}` | orchestrator | eval scores (if eval is configured) | + +#### Stability + +These tell you whether training is healthy or diverging. + +| Metric | Source | Description | +|--------|--------|-------------| +| `mismatch_kl/mean` | trainer | KL divergence between trainer and (old) inference policy | +| `entropy/mean` | trainer | policy entropy | +| `optim/grad_norm` | trainer | gradient norm — spikes may precede divergence | + +#### Performance + +These tell you how fast the run is and where the bottlenecks are. Trainer and orchestrator step independently — compare their step times to identify who is waiting on whom. + +**Trainer** (in `trainer.log`): + +| Metric | Description | +|--------|-------------| +| `time/step` | total trainer step time | +| `time/wait_for_batch` | time waiting for orchestrator to deliver a batch — **high = orchestrator is the bottleneck** | +| `time/forward_backward` | forward/backward pass time | +| `time/broadcast_weights` | time broadcasting weights to inference | +| `time/save_ckpt` | checkpoint save time | +| `perf/throughput` | tokens/s throughput | +| `perf/mfu` | model FLOPs utilization % | + +**Orchestrator** (in `orchestrator.log`): + +| Metric | Description | +|--------|-------------| +| `time/step` | total orchestrator step time | +| `time/generate_completions` | rollout generation time | +| `time/wait_for_ckpt` | time waiting for trainer checkpoint — **high = trainer is the bottleneck** | +| `time/update_weights` | weight update time | +| `scheduler/async_level` | current async level | +| `scheduler/inflight_rollouts` | number of in-flight rollouts | + +**Env servers** (in `envs/train/{env_name}/env_server.log`): + +| Metric | Description | +|--------|-------------| +| event loop lag (min/mean/p90/p99/max) | server and worker lag stats, logged periodically | +| active task distribution | per-worker and per-env task counts | + +**Inference** (in `inference.log`): + +vLLM logs completed HTTP requests and occasionally engine stats. For live inference metrics, query the vLLM metrics endpoint directly: + +```bash +# Get vLLM Prometheus metrics (num running/queued reqs, KV cache usage, etc.) +curl -s http://localhost:8000/metrics | grep -E "num_requests|gpu_cache_usage" +``` + +Key vLLM metrics to watch: +- `vllm:num_requests_running` — requests currently being processed +- `vllm:num_requests_waiting` — requests queued waiting for KV cache space +- `vllm:gpu_cache_usage_perc` — KV cache pressure (approaching 1.0 = requests will queue) + +### Rollouts + +Plain-text rollouts (`verifiers` format) are saved every step alongside the binary training batch inside the run directory. For single-run (local) runs this is typically `{output_dir}/run_default`. + +``` +{output_dir}/{run_dir}/rollouts/ +└── step_{N}/ + ├── train_rollouts.jsonl # all train rollouts + ├── eval_rollouts.jsonl # all eval rollouts (only present when eval ran) + └── train_rollouts.bin # binary-encoded training batch (consumed by the trainer) +``` + +Each line in the `.jsonl` files is a JSON-serialized `vf.RolloutOutput` dict with fields like `example_id`, `task`, `prompt`, `completion`, `reward`, `trajectory`, `metrics`, etc. To inspect rollouts for a given step: + +```bash +# Count rollouts at a step +wc -l {output_dir}/{run_dir}/rollouts/step_42/train_rollouts.jsonl + +# Preview first rollout (pretty-printed) +head -1 {output_dir}/{run_dir}/rollouts/step_42/train_rollouts.jsonl | python -m json.tool + +# Extract rewards +jq '.reward' {output_dir}/{run_dir}/rollouts/step_42/train_rollouts.jsonl +``` + +### Errors and warnings + +As part of every check-in, grep all logs for `WARNING` and `ERROR` level messages. Pay special attention to env server and env worker logs — these are the most common source of issues since they run user-provided code. + +Common things to look for: +- **Env workers**: exceptions in environment execution, timeouts, sandbox errors, OOM kills +- **Orchestrator**: empty/errored rollout spikes, weight broadcast failures, checkpoint errors +- **Trainer**: NCCL/CUDA errors, OOM, NaN loss or gradients +- **Inference**: NCCL/CUDA errors, OOM, request timeouts + +A small number of warnings is normal (e.g. occasional env timeouts). Escalate to the researcher if you see errors that are persistent, increasing, or affect a large fraction of rollouts. + +### Processes + +All processes have custom process names, making them easy to identify in `ps`, `htop`, and `pstree`. Use this in case you need to debug a process that is not responding or behaving unexpectedly. + +``` +PRIME-RL::Launcher +├── PRIME-RL::Inference (vLLM server, GPU 0) +├── PRIME-RL::Orchestrator (CPU-only, data/scheduling) +│ ├── Verifiers::EnvServer (ZMQ env server per environment) +│ │ ├── Verifiers::EnvWorker0 +│ │ ├── Verifiers::EnvWorker1 +│ │ └── ... +│ └── ... +├── torchrun +│ └── PRIME-RL::Trainer (RL trainer, GPU 1+) +└── tail trainer.log +``` + +For multi-node runs, trainer and inference processes are distributed across separate nodes. Use `srun` or `ssh` to inspect processes on other nodes directly. + diff --git a/skills/release/SKILL.md b/skills/release/SKILL.md new file mode 100644 index 00000000000..79ac7cac35c --- /dev/null +++ b/skills/release/SKILL.md @@ -0,0 +1,18 @@ +--- +name: release +description: How to prepare and publish GitHub releases for prime-rl. Use when drafting release notes, tagging versions, or publishing releases. +--- + +# Releases + +## Preparing release notes + +1. **Style reference**: check the previous release (`gh release list --limit 1` then `gh release view `) to match the tone and formatting. +2. **Gather changes**: use `git log ..origin/main --oneline --no-merges` to list all commits since the last release. +3. **Check for new commits**: always `git fetch origin main` and re-check right before publishing, since PRs may have been merged while drafting. +4. **Structure**: organize notes into numbered highlight sections (`# 1.`, `# 2.`, ...), then `# Breaking Changes`, `# Bug Fixes`, and `# Misc`. +5. **Highlights**: group related PRs under a single highlight. Use `##` subsections when a highlight contains multiple items (e.g. Performance & Parallelism). Keep the top highlights for the most impactful user-facing features. +6. **Config examples**: when referencing TOML config, verify the exact field names against the actual code or docs — don't guess. +7. **Links**: use clickable links for docs (`[docs/foo.md](https://github.com/PrimeIntellect-ai/prime-rl/blob/main/docs/foo.md)`) and PR references (`[#1234](https://github.com/PrimeIntellect-ai/prime-rl/pull/1234)`). +8. **Contributors**: list all contributors ranked by number of commits, using their GitHub `@username`. Get usernames via the GitHub API, not git author names (which can be inconsistent). +9. **Draft first**: always create releases as `--draft` first, iterate on content, then publish. diff --git a/src/prime_rl/_compat.py b/src/prime_rl/_compat.py new file mode 100644 index 00000000000..268c38376b8 --- /dev/null +++ b/src/prime_rl/_compat.py @@ -0,0 +1,55 @@ +"""Compatibility shims for upstream regressions. + +Imported early (before any model code) by trainer and orchestrator entrypoints. +Each shim documents the upstream issue and removal condition. +""" + +# --------------------------------------------------------------------------- +# ring_flash_attn + transformers >= 5.4 +# +# ring_flash_attn 0.1.8 imports `is_flash_attn_greater_or_equal_2_10` from +# `transformers.modeling_flash_attention_utils`, removed in transformers 5.4. +# +# Upstream fix: https://github.com/zhuzilin/ring-flash-attention/pull/85 +# Remove once ring_flash_attn ships a fixed version. +# --------------------------------------------------------------------------- +import transformers.modeling_flash_attention_utils as _mfau + +if not hasattr(_mfau, "is_flash_attn_greater_or_equal_2_10"): + _mfau.is_flash_attn_greater_or_equal_2_10 = lambda: True + + +# --------------------------------------------------------------------------- +# transformers >= 5.5 hub_kernels offline regression +# +# lazy_load_kernel() resolves kernel versions via HfApi().list_repo_refs(), +# which raises OfflineModeIsEnabled when HF_HUB_OFFLINE=1 — even if the +# kernel is already cached locally. Only FileNotFoundError and AssertionError +# are caught; OfflineModeIsEnabled (a ConnectionError subclass) is not. +# +# This breaks Mamba-based models (NemotronH, Zamba2, Jamba, etc.) on SLURM +# worker nodes where HF_HUB_OFFLINE=1 is set. +# +# Fix: patch the except clause to also catch ConnectionError. +# Remove once huggingface/transformers fixes lazy_load_kernel offline handling. +# --------------------------------------------------------------------------- +try: + import transformers.integrations.hub_kernels as _hub_kernels +except ImportError: + _hub_kernels = None # transformers < 5.5, no patch needed + +if _hub_kernels is not None: + from huggingface_hub.errors import OfflineModeIsEnabled + + _original_lazy_load_kernel = _hub_kernels.lazy_load_kernel + + def _patched_lazy_load_kernel(kernel_name, mapping=_hub_kernels._KERNEL_MODULE_MAPPING): + try: + return _original_lazy_load_kernel(kernel_name, mapping) + except OfflineModeIsEnabled: + # Return None so NemotronH skips hub kernels; prime-rl's + # _patch_mamba2_use_triton_ssd() uses mamba_ssm directly. + mapping[kernel_name] = None + return None + + _hub_kernels.lazy_load_kernel = _patched_lazy_load_kernel diff --git a/src/prime_rl/entrypoints/inference.py b/src/prime_rl/entrypoints/inference.py new file mode 100644 index 00000000000..a14ae26e639 --- /dev/null +++ b/src/prime_rl/entrypoints/inference.py @@ -0,0 +1,152 @@ +import subprocess +import sys +from pathlib import Path + +import tomli_w + +from prime_rl.configs.inference import InferenceConfig +from prime_rl.utils.config import cli +from prime_rl.utils.logger import setup_logger +from prime_rl.utils.pathing import format_log_message, get_config_dir, get_log_dir +from prime_rl.utils.process import set_proc_title + +INFERENCE_TOML = "inference.toml" +INFERENCE_SBATCH = "inference.sbatch" + + +def write_config(config: InferenceConfig, output_dir: Path, exclude: set[str] | None = None) -> Path: + """Write resolved config to disk.""" + output_dir.mkdir(parents=True, exist_ok=True) + config_path = output_dir / INFERENCE_TOML + with open(config_path, "wb") as f: + tomli_w.dump(config.model_dump(exclude=exclude, exclude_none=True, mode="json"), f) + return config_path + + +def write_slurm_script(config: InferenceConfig, config_path: Path, script_path: Path) -> None: + """Write the SLURM script to disk.""" + from jinja2 import Environment, FileSystemLoader + + assert config.slurm is not None + assert config.slurm.template_path is not None + + env = Environment(loader=FileSystemLoader(config.slurm.template_path.parent), keep_trailing_newline=True) + template = env.get_template(config.slurm.template_path.name) + + is_disaggregated = config.deployment.type == "disaggregated" + dp_per_node = config.deployment.gpus_per_node // config.parallel.tp + + template_vars = dict( + **config.slurm.template_vars, + config_path=config_path, + output_dir=config.output_dir, + gpus_per_node=config.deployment.gpus_per_node, + dp_per_node=dp_per_node, + num_nodes=getattr(config.deployment, "num_nodes", 1), + port=config.server.port, + disaggregated=is_disaggregated, + kv_offload=config.kv_cache_offload is not None, + ) + + is_multi_node = config.deployment.type == "multi_node" + + if is_disaggregated: + template_vars.update( + num_prefill_nodes=config.deployment.num_prefill_nodes, + num_decode_nodes=config.deployment.num_decode_nodes, + num_prefill_replicas=config.deployment.num_prefill_replicas, + num_decode_replicas=config.deployment.num_decode_replicas, + prefill_port=config.deployment.prefill_port, + decode_port=config.deployment.decode_port, + router_port=config.deployment.router_port, + router_policy=config.deployment.router_policy, + data_parallel_rpc_port=config.data_parallel_rpc_port, + use_deep_gemm=config.use_deep_gemm, + prefill_env_overrides=config.deployment.prefill_env_overrides, + decode_env_overrides=config.deployment.decode_env_overrides, + kv_offload_cpu_bytes=int(config.kv_cache_offload.cpu_bytes) if config.kv_cache_offload else 0, + ) + elif is_multi_node: + template_vars.update( + router_port=config.deployment.router_port, + backend_port=config.deployment.backend_port, + router_policy=config.deployment.router_policy, + ) + + script = template.render(**template_vars) + + script_path.parent.mkdir(parents=True, exist_ok=True) + script_path.write_text(script) + + +def inference_slurm(config: InferenceConfig): + """Run inference via SLURM.""" + assert config.slurm is not None + + logger = setup_logger("info") + + config_dir = get_config_dir(config.output_dir) + exclude = ( + {"deployment", "slurm", "dry_run"} + if config.deployment.type in ("multi_node", "disaggregated") + else {"slurm", "dry_run"} + ) + config_path = write_config(config, config_dir, exclude=exclude) + logger.info(f"Wrote config to {config_path}") + + script_path = config.output_dir / INFERENCE_SBATCH + write_slurm_script(config, config_path, script_path) + logger.info(f"Wrote SLURM script to {script_path}") + + log_dir = get_log_dir(config.output_dir) + num_nodes = getattr(config.deployment, "num_nodes", 1) + log_message = format_log_message(log_dir=log_dir, inference=True, job_log=True, num_infer_nodes=num_nodes) + + if config.dry_run: + logger.success(f"Dry run complete. To submit manually:\n\n sbatch {script_path}\n\n{log_message}") + return + + logger.info(f"Submitting: sbatch {script_path}") + result = subprocess.run(["sbatch", str(script_path)], capture_output=True, text=True) + if result.returncode != 0: + logger.error(f"sbatch failed: {result.stderr.strip()}") + sys.exit(1) + + logger.success(f"{result.stdout.strip()}\n\n{log_message}") + + +def inference_local(config: InferenceConfig): + """Run inference locally.""" + from prime_rl.inference.server import setup_vllm_env + + logger = setup_logger("info") + + if config.dry_run: + logger.success("Dry run complete. To start inference locally, remove --dry-run from your command.") + return + + host = config.server.host or "0.0.0.0" + port = config.server.port + logger.info(f"Starting inference on http://{host}:{port}/v1\n") + + setup_vllm_env(config) + + from prime_rl.inference.vllm.server import server # pyright: ignore + + server(config, vllm_extra=config.vllm_extra) + + +def inference(config: InferenceConfig): + if config.slurm is not None: + inference_slurm(config) + else: + inference_local(config) + + +def main(): + set_proc_title("Inference") + inference(cli(InferenceConfig)) + + +if __name__ == "__main__": + main() diff --git a/src/prime_rl/entrypoints/rl.py b/src/prime_rl/entrypoints/rl.py new file mode 100644 index 00000000000..b740f58ba29 --- /dev/null +++ b/src/prime_rl/entrypoints/rl.py @@ -0,0 +1,588 @@ +import json +import os +import signal +import subprocess +import sys +import time +import uuid +from pathlib import Path +from subprocess import Popen +from threading import Event, Thread + +import pynvml +import tomli_w + +import prime_rl._compat # noqa: F401 — patch ring_flash_attn compat before transitive import +from prime_rl.configs.rl import RLConfig +from prime_rl.trainer.model import pre_download_model +from prime_rl.utils.config import cli +from prime_rl.utils.logger import get_logger, setup_logger +from prime_rl.utils.pathing import ( + clean_future_steps, + format_log_message, + get_ckpt_dir, + resolve_latest_ckpt_step, + validate_output_dir, +) +from prime_rl.utils.process import cleanup_processes, cleanup_threads, monitor_process, set_proc_title +from prime_rl.utils.utils import ( + get_free_port, + get_log_dir, +) + +RL_TOML = "rl.toml" +RL_SBATCH = "rl.sbatch" + +TRAINER_TOML = "trainer.toml" +ORCHESTRATOR_TOML = "orchestrator.toml" +INFERENCE_TOML = "inference.toml" +TEACHER_INFERENCE_TOML = "teacher_inference.toml" + + +def get_physical_gpu_ids() -> list[int]: + """Return physical GPU IDs visible to the launcher.""" + raw_visible = os.environ.get("CUDA_VISIBLE_DEVICES") + if raw_visible is None: + pynvml.nvmlInit() + return list(range(pynvml.nvmlDeviceGetCount())) + return [int(token.strip()) for token in raw_visible.split(",") if token.strip()] + + +def write_config(config: RLConfig, output_dir: Path, exclude: set[str] | None = None) -> None: + """Write resolved config to disk, excluding launcher-only fields.""" + output_dir.mkdir(parents=True, exist_ok=True) + config_dict = config.model_dump(exclude=exclude, exclude_none=True, mode="json") + with open(output_dir / RL_TOML, "wb") as f: + tomli_w.dump(config_dict, f) + + +def write_subconfigs(config: RLConfig, output_dir: Path) -> None: + """Write resolved subconfigs to disk as TOML files.""" + output_dir.mkdir(parents=True, exist_ok=True) + + with open(output_dir / TRAINER_TOML, "wb") as f: + tomli_w.dump(config.trainer.model_dump(exclude_none=True, mode="json"), f) + + with open(output_dir / ORCHESTRATOR_TOML, "wb") as f: + tomli_w.dump(config.orchestrator.model_dump(exclude_none=True, mode="json"), f) + + if config.inference is not None: + # Exclude launcher-only fields that are not needed by the vLLM server + exclude_inference = {"deployment", "slurm", "output_dir", "dry_run"} + with open(output_dir / INFERENCE_TOML, "wb") as f: + tomli_w.dump(config.inference.model_dump(exclude=exclude_inference, exclude_none=True, mode="json"), f) + + teacher_inference = getattr(config, "teacher_inference", None) + if teacher_inference is not None: + with open(output_dir / TEACHER_INFERENCE_TOML, "wb") as f: + tomli_w.dump(teacher_inference.model_dump(exclude_none=True, mode="json"), f) + + +def check_gpus_available(gpu_ids: list[int]) -> None: + """Raise error if there are existing processes on the specified GPUs.""" + pynvml.nvmlInit() + + occupied = [] + for gpu_id in gpu_ids: + handle = pynvml.nvmlDeviceGetHandleByIndex(gpu_id) + processes = pynvml.nvmlDeviceGetComputeRunningProcesses(handle) + if processes: + pids = [p.pid for p in processes] + occupied.append((gpu_id, pids)) + + if occupied: + msg = "Existing processes found on GPUs:\n" + for gpu_id, pids in occupied: + msg += f" GPU {gpu_id}: PIDs {pids}\n" + msg += "Kill these processes or use different GPUs." + raise RuntimeError(msg) + + +def rl_local(config: RLConfig): + assert config.deployment.type == "single_node" + + logger = setup_logger( + config.log.level or os.environ.get("PRIME_LOG_LEVEL", "info"), + json_logging=config.log.json_logging, + ) + + config_dir = config.output_dir / "configs" + write_subconfigs(config, config_dir) + logger.info(f"Wrote subconfigs to {config_dir}") + + if config.dry_run: + logger.success("Dry run complete. To start an RL run locally, remove --dry-run from your command.") + return + + # Derive launcher-local GPU IDs from deployment config + gpu_offset = 0 + num_infer_gpus = config.deployment.num_infer_gpus if config.inference is not None else 0 + infer_local_gpu_ids = list(range(gpu_offset, gpu_offset + num_infer_gpus)) + gpu_offset += num_infer_gpus + trainer_local_gpu_ids = list(range(gpu_offset, gpu_offset + config.deployment.num_train_gpus)) + gpu_offset += config.deployment.num_train_gpus + num_teacher_gpus = config.deployment.num_teacher_gpus or 0 + teacher_local_gpu_ids = list(range(gpu_offset, gpu_offset + num_teacher_gpus)) if num_teacher_gpus > 0 else [] + + total_requested_gpus = num_infer_gpus + config.deployment.num_train_gpus + num_teacher_gpus + physical_gpu_ids = get_physical_gpu_ids() + if total_requested_gpus > len(physical_gpu_ids): + raise ValueError( + f"Requested {total_requested_gpus} GPUs via deployment settings, but only " + f"{len(physical_gpu_ids)} physical GPU(s) are available: {physical_gpu_ids}" + ) + physical_gpu_mapping = {local_id: physical_gpu_ids[local_id] for local_id in range(total_requested_gpus)} + logger.info(f"Using local->physical GPU mapping: {physical_gpu_mapping}") + + infer_gpu_ids = [physical_gpu_mapping[local_gpu_id] for local_gpu_id in infer_local_gpu_ids] + trainer_gpu_ids = [physical_gpu_mapping[local_gpu_id] for local_gpu_id in trainer_local_gpu_ids] + teacher_gpu_ids = [physical_gpu_mapping[local_gpu_id] for local_gpu_id in teacher_local_gpu_ids] + + start_command = sys.argv + logger.info("Starting RL run") + logger.debug(f"RL start command: {' '.join(start_command)}") + + # Build shared W&B env vars for subprocesses + wandb_shared_env: dict[str, str] = {} + if config.wandb and config.wandb.shared: + wandb_shared_env["WANDB_SHARED_MODE"] = "1" + wandb_shared_env["WANDB_SHARED_RUN_ID"] = os.environ.get("WANDB_SHARED_RUN_ID", uuid.uuid4().hex) + + # Check for existing processes on GPUs + all_gpu_ids = list(set(infer_gpu_ids + trainer_gpu_ids + teacher_gpu_ids)) + check_gpus_available(all_gpu_ids) + + # Validate client port matches inference server port + if config.inference is not None and not config.orchestrator.client.is_elastic: + from urllib.parse import urlparse + + base_url = config.orchestrator.client.base_url[0] + parsed = urlparse(base_url) + client_port = parsed.port + expected_port = config.inference.server.port + if client_port != expected_port: + raise ValueError( + f"orchestrator.client.base_url port ({client_port}) does not match " + f"inference.server.port ({expected_port}). " + f"Update the base_url to use port {expected_port} to match the inference server." + ) + + # Prepare paths to communicate with the trainer + log_dir = get_log_dir(config.output_dir) + log_dir.mkdir(parents=True, exist_ok=True) + + # Start processes + processes: list[Popen] = [] + monitor_threads: list[Thread] = [] + error_queue: list[Exception] = [] + stop_events: dict[str, Event] = {} + + def sigterm_handler(signum, frame): + logger.warning("Received SIGTERM, terminating all processes...") + cleanup_threads(monitor_threads) + cleanup_processes(processes) + sys.exit(1) + + signal.signal(signal.SIGTERM, sigterm_handler) + + try: + # Optionally, start inference process + if config.inference: + inference_cmd = ["inference", "@", (config_dir / INFERENCE_TOML).as_posix()] + logger.info(f"Starting inference on GPU(s) {' '.join(map(str, infer_gpu_ids))}") + logger.debug(f"Inference start command: {' '.join(inference_cmd)}") + # If we don't log stdout, the server hangs + with open(log_dir / "inference.log", "w") as log_file: + inference_process = Popen( + inference_cmd, + env={ + **os.environ, + "CUDA_VISIBLE_DEVICES": ",".join(map(str, infer_gpu_ids)), + }, + stdout=log_file, + stderr=log_file, + ) + processes.append(inference_process) + + # Start monitoring thread + stop_event = Event() + stop_events["inference"] = stop_event + monitor_thread = Thread( + target=monitor_process, + args=(inference_process, stop_event, error_queue, "inference"), + daemon=True, + ) + monitor_thread.start() + monitor_threads.append(monitor_thread) + else: + if config.orchestrator.teacher_rollout_model is None: + logger.warning( + "No inference config specified, skipping starting inference server. Make sure your inference server is running." + ) + else: + logger.info( + "No inference config specified, using orchestrator.teacher_rollout_model for rollout generation." + ) + + # Optionally, start teacher inference process + if config.teacher_inference: + if not teacher_gpu_ids: + raise ValueError( + "teacher_inference is configured but deployment.num_teacher_gpus is not set. " + "Either set deployment.num_teacher_gpus to start a teacher inference server, " + "or omit teacher_inference and configure orchestrator.teacher_model to use an existing server." + ) + + teacher_inference_cmd = ["inference", "@", (config_dir / TEACHER_INFERENCE_TOML).as_posix()] + logger.info(f"Starting teacher inference process on GPU(s) {' '.join(map(str, teacher_gpu_ids))}") + logger.debug(f"Teacher inference start command: {' '.join(teacher_inference_cmd)}") + with open(log_dir / "teacher_inference.log", "w") as log_file: + teacher_inference_process = Popen( + teacher_inference_cmd, + env={ + **os.environ, + "CUDA_VISIBLE_DEVICES": ",".join(map(str, teacher_gpu_ids)), + }, + stdout=log_file, + stderr=log_file, + ) + processes.append(teacher_inference_process) + + # Start monitoring thread + stop_event = Event() + stop_events["teacher_inference"] = stop_event + monitor_thread = Thread( + target=monitor_process, + args=(teacher_inference_process, stop_event, error_queue, "teacher_inference"), + daemon=True, + ) + monitor_thread.start() + monitor_threads.append(monitor_thread) + elif ( + config.trainer.loss.type == "default" and config.trainer.loss.teacher_tau > 0 + ) or config.orchestrator.teacher_model: + logger.warning( + "No teacher_inference config specified, skipping starting teacher inference server. " + "Is your teacher inference server running? Make sure orchestrator.teacher_model is configured." + ) + + # Start orchestrator process + orchestrator_cmd = [ + "orchestrator", + "@", + (config_dir / ORCHESTRATOR_TOML).as_posix(), + ] + logger.info("Starting orchestrator process") + logger.debug(f"Orchestrator start command: {' '.join(orchestrator_cmd)}") + with open(log_dir / "orchestrator.log", "w") as log_file: + orchestrator_process = Popen( + orchestrator_cmd, + stdout=log_file, + stderr=log_file, + env={ + **os.environ, + **wandb_shared_env, + "WANDB_SHARED_LABEL": "orchestrator", + "LOGURU_FORCE_COLORS": "1", + "WANDB_PROGRAM": "uv run rl", + "WANDB_ARGS": json.dumps(start_command), + }, + ) + processes.append(orchestrator_process) + + # Start monitoring thread + stop_event = Event() + stop_events["orchestrator"] = stop_event + monitor_thread = Thread( + target=monitor_process, + args=(orchestrator_process, stop_event, error_queue, "orchestrator"), + daemon=True, + ) + monitor_thread.start() + monitor_threads.append(monitor_thread) + + # Start training process + trainer_cmd = [ + "torchrun", + "--role=trainer", + f"--rdzv-endpoint=localhost:{get_free_port()}", + f"--rdzv-id={uuid.uuid4().hex}", + # Pipe all logs to file, and only master rank logs to stdout + f"--log-dir={log_dir / 'trainer' / 'torchrun'}", + f"--local-ranks-filter={','.join(map(str, config.trainer.log.ranks_filter))}", + "--redirect=3", + "--tee=3", + f"--nproc-per-node={len(trainer_gpu_ids)}", + "-m", + "prime_rl.trainer.rl.train", + "@", + (config_dir / TRAINER_TOML).as_posix(), + ] + logger.info(f"Starting trainer on GPU(s) {' '.join(map(str, trainer_gpu_ids))}") + logger.debug(f"Training start command: {' '.join(trainer_cmd)}") + with open(log_dir / "trainer.log", "w") as log_file: + trainer_process = Popen( + trainer_cmd, + env={ + **os.environ, + **wandb_shared_env, + "WANDB_SHARED_LABEL": "trainer", + "CUDA_VISIBLE_DEVICES": ",".join(map(str, trainer_gpu_ids)), + "PYTHONUNBUFFERED": "1", + "PYTORCH_CUDA_ALLOC_CONF": "expandable_segments:True", + "LOGURU_FORCE_COLORS": "1", + "WANDB_PROGRAM": "uv run rl", + "WANDB_ARGS": json.dumps(start_command), + }, + stdout=log_file, + stderr=log_file, + ) + processes.append(trainer_process) + + # Start monitoring thread + stop_event = Event() + stop_events["trainer"] = stop_event + monitor_thread = Thread( + target=monitor_process, args=(trainer_process, stop_event, error_queue, "trainer"), daemon=True + ) + monitor_thread.start() + monitor_threads.append(monitor_thread) + + # Monitor all processes for failures + logger.success("Startup complete. Showing trainer logs...") + + tail_process = Popen( + f"tail -F '{log_dir / 'trainer.log'}' | sed -u 's/^\\[[a-zA-Z]*[0-9]*\\]://'", + shell=True, + ) + processes.append(tail_process) + + # Check for errors from monitor threads + while not (stop_events["orchestrator"].is_set() and stop_events["trainer"].is_set()): + if error_queue: + error = error_queue[0] + logger.error(f"Error: {error}") + logger.error("Terminating all processes...") + cleanup_threads(monitor_threads) + cleanup_processes(processes) + sys.exit(1) + + # Small delay to avoid busy waiting + time.sleep(1) + + # Check if any critical process failed + if orchestrator_process.returncode != 0: + logger.error(f"Orchestrator failed with exit code {orchestrator_process.returncode}") + cleanup_threads(monitor_threads) + cleanup_processes(processes) + sys.exit(1) + + if trainer_process.returncode != 0: + logger.error(f"Trainer failed with exit code {trainer_process.returncode}") + cleanup_threads(monitor_threads) + cleanup_processes(processes) + sys.exit(1) + + logger.success("RL training finished!") + + # Cleanup threads and processes + cleanup_threads(monitor_threads) + cleanup_processes(processes) + + except KeyboardInterrupt: + logger.warning("Received interrupt signal, terminating all processes...") + cleanup_threads(monitor_threads) + cleanup_processes(processes) + sys.exit(1) + except Exception as e: + logger.error(f"Error occurred: {e}") + cleanup_threads(monitor_threads) + cleanup_processes(processes) + raise + + +def write_slurm_script(config: RLConfig, config_dir: Path, script_path: Path) -> None: + """Write the SLURM script to disk.""" + from jinja2 import Environment, FileSystemLoader + + assert config.slurm is not None + assert config.slurm.template_path is not None + + env = Environment(loader=FileSystemLoader(config.slurm.template_path.parent), keep_trailing_newline=True) + template = env.get_template(config.slurm.template_path.name) + + if config.deployment.type == "single_node": + script = template.render( + **config.slurm.template_vars, + config_path=config_dir / RL_TOML, + output_dir=config.output_dir, + gpus_per_node=config.deployment.gpus_per_node, + ) + elif config.inference is not None and config.inference.deployment.type == "disaggregated": + infer_deploy = config.inference.deployment + + script = template.render( + **config.slurm.template_vars, + is_disaggregated=True, + config_dir=config_dir, + output_dir=config.output_dir, + orchestrator_output_dir=config.orchestrator.output_dir, + num_train_nodes=config.deployment.num_train_nodes, + num_infer_nodes=infer_deploy.num_nodes * config.deployment.num_infer_replicas, + nodes_per_infer_replica=infer_deploy.num_nodes, + num_infer_replicas=config.deployment.num_infer_replicas, + num_prefill_nodes=infer_deploy.num_prefill_nodes, + num_decode_nodes=infer_deploy.num_decode_nodes, + num_prefill_replicas=infer_deploy.num_prefill_replicas, + num_decode_replicas=infer_deploy.num_decode_replicas, + gpus_per_node=config.deployment.gpus_per_node, + router_port=infer_deploy.router_port, + prefill_port=infer_deploy.prefill_port, + decode_port=infer_deploy.decode_port, + inference_tp=config.inference.parallel.tp, + inference_data_parallel_rpc_port=config.inference.data_parallel_rpc_port, + use_deep_gemm=config.inference.use_deep_gemm, + prefill_env_overrides=infer_deploy.prefill_env_overrides, + decode_env_overrides=infer_deploy.decode_env_overrides, + dp_per_node=config.deployment.gpus_per_node // config.inference.parallel.tp, + kv_offload=config.inference.kv_cache_offload is not None, + kv_offload_cpu_bytes=int(config.inference.kv_cache_offload.cpu_bytes) + if config.inference.kv_cache_offload + else 0, + use_nccl_broadcast=config.weight_broadcast is not None and config.weight_broadcast.type == "nccl", + wandb_shared=config.wandb is not None and config.wandb.shared, + ranks_filter=",".join(map(str, config.trainer.log.ranks_filter)), + ) + else: + script = template.render( + **config.slurm.template_vars, + is_disaggregated=False, + config_dir=config_dir, # TODO: should prob have each subconfig path separately + output_dir=config.output_dir, + orchestrator_output_dir=config.orchestrator.output_dir, + num_train_nodes=config.deployment.num_train_nodes, + num_infer_nodes=config.deployment.total_infer_nodes, + nodes_per_infer_replica=config.deployment.num_infer_nodes, + num_infer_replicas=config.deployment.num_infer_replicas, + num_teacher_nodes=config.deployment.num_teacher_nodes, + gpus_per_node=config.deployment.gpus_per_node, + router_port=getattr(config.inference.deployment, "router_port", 8000) if config.inference else 8000, + backend_port=getattr(config.inference.deployment, "backend_port", 8100) if config.inference else 8100, + inference_tp=config.inference.parallel.tp if config.inference else 1, + inference_enable_expert_parallel=config.inference.enable_expert_parallel if config.inference else False, + inference_data_parallel_rpc_port=config.inference.data_parallel_rpc_port if config.inference else 29600, + dp_per_node=(config.deployment.gpus_per_node // config.inference.parallel.tp) if config.inference else 1, + kv_offload=config.inference is not None and config.inference.kv_cache_offload is not None, + use_nccl_broadcast=config.weight_broadcast is not None and config.weight_broadcast.type == "nccl", + wandb_shared=config.wandb is not None and config.wandb.shared, + ranks_filter=",".join(map(str, config.trainer.log.ranks_filter)), + ) + + script_path.parent.mkdir(parents=True, exist_ok=True) + script_path.write_text(script) + + +def rl_slurm(config: RLConfig): + assert config.slurm is not None + + logger = setup_logger( + config.log.level or os.environ.get("PRIME_LOG_LEVEL", "info"), json_logging=config.log.json_logging + ) + + config_dir = config.output_dir / "configs" + log_dir = get_log_dir(config.output_dir) + + if config.deployment.type == "single_node": + write_config(config, config_dir, exclude={"slurm", "dry_run", "clean_output_dir"}) + logger.info(f"Wrote config to {config_dir / RL_TOML}") + + train_env_names = [env.resolved_name for env in config.orchestrator.train.env] + eval_env_names = [env.resolved_name for env in config.orchestrator.eval.env] if config.orchestrator.eval else [] + + log_message = format_log_message( + log_dir=log_dir, + trainer=True, + orchestrator=True, + inference=True, + train_env_names=train_env_names, + eval_env_names=eval_env_names, + ) + else: + write_subconfigs(config, config_dir) + logger.info(f"Wrote subconfigs to {config_dir}") + + train_env_names = [env.resolved_name for env in config.orchestrator.train.env] + eval_env_names = [env.resolved_name for env in config.orchestrator.eval.env] if config.orchestrator.eval else [] + + has_infer = config.deployment.num_infer_nodes > 0 + log_message = format_log_message( + log_dir=log_dir, + trainer=True, + orchestrator=has_infer, + inference=has_infer, + train_env_names=train_env_names, + eval_env_names=eval_env_names, + num_train_nodes=config.deployment.num_train_nodes, + num_infer_nodes=config.deployment.total_infer_nodes if has_infer else 0, + ) + + script_path = config.output_dir / RL_SBATCH + write_slurm_script(config, config_dir, script_path) + logger.info(f"Wrote SLURM script to {script_path}") + + if config.dry_run: + logger.success(f"Dry run complete. To submit manually:\n\n sbatch {script_path}\n\n{log_message}") + return + + logger.info(f"Submitting: sbatch {script_path}") + result = subprocess.run(["sbatch", str(script_path)], capture_output=True, text=True) + if result.returncode != 0: + logger.error(f"sbatch failed: {result.stderr.strip()}") + sys.exit(1) + + logger.success(f"{result.stdout.strip()}\n\n{log_message}") + + +def rl(config: RLConfig): + resuming = config.ckpt is not None and config.ckpt.resume_step is not None + clean = config.clean_output_dir and not os.environ.get("NEVER_CLEAN_OUTPUT_DIR") + ckpt_output_dir = config.ckpt.output_dir if config.ckpt else None + validate_output_dir(config.output_dir, resuming=resuming, clean=clean, ckpt_output_dir=ckpt_output_dir) + config.output_dir.mkdir(parents=True, exist_ok=True) + if ckpt_output_dir is not None: + ckpt_output_dir.mkdir(parents=True, exist_ok=True) + + # Clean stale rollouts and broadcasts. When resuming, anything past the resume + # step is stale. When training from scratch, every existing step directory is + # stale — without this, a fresh run in a dirty output_dir would pick up rollouts + # from a previous run and the orchestrator would see a negative async level. + resume_step: int | None = None + if resuming: + resume_step = config.ckpt.resume_step + if resume_step == -1: + ckpt_base = ckpt_output_dir if ckpt_output_dir is not None else config.output_dir + resume_step = resolve_latest_ckpt_step(get_ckpt_dir(ckpt_base)) + + if resume_step is not None: + get_logger().info(f"Resuming from step {resume_step}, cleaning future rollouts and broadcasts") + clean_future_steps(config.output_dir, resume_step) + else: + get_logger().info("Training from scratch, cleaning any stale rollouts and broadcasts") + clean_future_steps(config.output_dir, -1) + + if not config.dry_run: + pre_download_model(config.trainer.model.name) + + if config.slurm is not None: + rl_slurm(config) + else: + rl_local(config) + + +def main(): + set_proc_title("Launcher") + rl(cli(RLConfig)) + + +if __name__ == "__main__": + main() diff --git a/src/prime_rl/entrypoints/sft.py b/src/prime_rl/entrypoints/sft.py new file mode 100644 index 00000000000..64b5ccb1912 --- /dev/null +++ b/src/prime_rl/entrypoints/sft.py @@ -0,0 +1,216 @@ +import os +import subprocess +import sys +import uuid +from pathlib import Path +from subprocess import Popen +from threading import Event, Thread + +import tomli_w + +import prime_rl._compat # noqa: F401 — patch ring_flash_attn compat before transitive import +from prime_rl.configs.sft import SFTConfig +from prime_rl.trainer.model import pre_download_model +from prime_rl.utils.config import cli +from prime_rl.utils.logger import setup_logger +from prime_rl.utils.pathing import format_log_message, get_config_dir, get_log_dir, validate_output_dir +from prime_rl.utils.process import cleanup_processes, cleanup_threads, monitor_process, set_proc_title +from prime_rl.utils.utils import get_free_port + +SFT_TOML = "sft.toml" +SFT_SBATCH = "sft.sbatch" + + +def write_config(config: SFTConfig, config_path: Path, exclude: set[str] | None = None) -> None: + """Write resolved config to disk, excluding launcher-only fields.""" + config_path.parent.mkdir(parents=True, exist_ok=True) + config_dict = config.model_dump(exclude=exclude, exclude_none=True, mode="json") + with open(config_path, "wb") as f: + tomli_w.dump(config_dict, f) + + +def write_slurm_script(config: SFTConfig, config_path: Path, script_path: Path) -> None: + """Write the SLURM script to disk.""" + from jinja2 import Environment, FileSystemLoader + + assert config.slurm is not None + assert config.slurm.template_path is not None + + env = Environment(loader=FileSystemLoader(config.slurm.template_path.parent), keep_trailing_newline=True) + template = env.get_template(config.slurm.template_path.name) + + if config.deployment.type == "single_node": + script = template.render( + **config.slurm.template_vars, + config_path=config_path, + output_dir=config.output_dir, + gpus_per_node=config.deployment.gpus_per_node, + ) + else: + script = template.render( + **config.slurm.template_vars, + config_path=config_path, + output_dir=config.output_dir, + num_nodes=config.deployment.num_nodes, + gpus_per_node=config.deployment.gpus_per_node, + ranks_filter=",".join(map(str, config.log.ranks_filter)), + ) + + script_path.parent.mkdir(parents=True, exist_ok=True) + script_path.write_text(script) + + +def sft_slurm(config: SFTConfig): + """Run SFT training via SLURM.""" + assert config.slurm is not None + + logger = setup_logger(config.log.level or "info", json_logging=config.log.json_logging) + + config_dir = get_config_dir(config.output_dir) + config_path = config_dir / SFT_TOML + exclude = ( + {"deployment", "slurm", "dry_run", "clean_output_dir"} + if config.deployment.type == "multi_node" + else {"slurm", "dry_run", "clean_output_dir"} + ) + write_config(config, config_path, exclude=exclude) + logger.info(f"Wrote config to {config_path}") + + script_path = config.output_dir / SFT_SBATCH + write_slurm_script(config, config_path, script_path) + logger.info(f"Wrote SLURM script to {script_path}") + + log_dir = get_log_dir(config.output_dir) + num_nodes = config.deployment.num_nodes if config.deployment.type == "multi_node" else 1 + log_message = format_log_message(log_dir=log_dir, trainer=True, num_train_nodes=num_nodes) + + if config.dry_run: + logger.success(f"Dry run complete. To submit manually:\n\n sbatch {script_path}\n\n{log_message}") + return + + logger.info(f"Submitting: sbatch {script_path}") + result = subprocess.run(["sbatch", str(script_path)], capture_output=True, text=True) + if result.returncode != 0: + logger.error(f"sbatch failed: {result.stderr.strip()}") + sys.exit(1) + + logger.success(f"{result.stdout.strip()}\n\n{log_message}") + + +def sft_local(config: SFTConfig): + """Run SFT training locally with process monitoring and cleanup.""" + assert config.deployment.type == "single_node" + + logger = setup_logger(config.log.level or "info", json_logging=config.log.json_logging) + + config_dir = get_config_dir(config.output_dir) + config_path = config_dir / SFT_TOML + write_config(config, config_path) + logger.info(f"Wrote config to {config_path}") + + if config.dry_run: + logger.success("Dry run complete. To start an SFT run locally, remove --dry-run from your command.") + return + + log_dir = config.output_dir / "logs" + log_dir.mkdir(parents=True, exist_ok=True) + + trainer_cmd = [ + "torchrun", + "--role=trainer", + f"--rdzv-endpoint=localhost:{get_free_port()}", + f"--rdzv-id={uuid.uuid4().hex}", + f"--log-dir={config.output_dir / 'logs' / 'trainer' / 'torchrun'}", + f"--local-ranks-filter={','.join(map(str, config.log.ranks_filter))}", + "--redirect=3", + "--tee=3", + f"--nproc-per-node={config.deployment.num_gpus}", + "-m", + "prime_rl.trainer.sft.train", + "@", + (config_dir / SFT_TOML).as_posix(), + ] + + logger.info(f"Starting SFT trainer with {config.deployment.num_gpus} GPU(s)") + logger.debug(f"Trainer command: {' '.join(trainer_cmd)}") + + processes: list[Popen] = [] + monitor_threads: list[Thread] = [] + error_queue: list[Exception] = [] + + try: + with open(log_dir / "trainer.log", "w") as log_file: + trainer_process = Popen( + trainer_cmd, + env={ + **os.environ, + "PYTHONUNBUFFERED": "1", + "PYTORCH_CUDA_ALLOC_CONF": "expandable_segments:True", + }, + stdout=log_file, + stderr=log_file, + ) + processes.append(trainer_process) + + stop_event = Event() + monitor_thread = Thread( + target=monitor_process, + args=(trainer_process, stop_event, error_queue, "trainer"), + daemon=True, + ) + monitor_thread.start() + monitor_threads.append(monitor_thread) + + logger.success("Startup complete. Showing trainer logs...") + tail_process = Popen( + f"tail -F '{log_dir / 'trainer.log'}' | sed -u 's/^\\[[a-zA-Z]*[0-9]*\\]://'", + shell=True, + ) + processes.append(tail_process) + + stop_event.wait() + + if trainer_process.returncode != 0: + logger.error(f"Trainer failed with exit code {trainer_process.returncode}") + cleanup_threads(monitor_threads) + cleanup_processes(processes) + sys.exit(1) + + logger.success("SFT training finished!") + cleanup_threads(monitor_threads) + cleanup_processes(processes) + + except KeyboardInterrupt: + logger.warning("Received interrupt signal, terminating all processes...") + cleanup_threads(monitor_threads) + cleanup_processes(processes) + sys.exit(1) + except Exception as e: + logger.error(f"Error occurred: {e}") + cleanup_threads(monitor_threads) + cleanup_processes(processes) + raise + + +def sft(config: SFTConfig): + resuming = config.ckpt is not None and config.ckpt.resume_step is not None + clean = config.clean_output_dir and not os.environ.get("NEVER_CLEAN_OUTPUT_DIR") + validate_output_dir(config.output_dir, resuming=resuming, clean=clean) + config.output_dir.mkdir(parents=True, exist_ok=True) + + if not config.dry_run: + pre_download_model(config.model.name) + + if config.slurm is not None: + sft_slurm(config) + else: + sft_local(config) + + +def main(): + set_proc_title("SFT") + sft(cli(SFTConfig)) + + +if __name__ == "__main__": + main() diff --git a/src/prime_rl/eval/config.py b/src/prime_rl/eval/config.py deleted file mode 100644 index 27ecce85483..00000000000 --- a/src/prime_rl/eval/config.py +++ /dev/null @@ -1,135 +0,0 @@ -from pathlib import Path -from typing import Annotated - -from pydantic import Field, model_validator - -from prime_rl.orchestrator.config import EvalConfig -from prime_rl.utils.config import ClientConfig, LogConfig, ModelConfig, WandbConfig -from prime_rl.utils.pydantic_config import BaseSettings - - -class OfflineEvalConfig(EvalConfig, BaseSettings): - """Configures evaluation.""" - - # The client configuration - client: ClientConfig = ClientConfig(timeout=36000, base_url=["https://api.openai.com/v1"]) - - # The model configuration - model: ModelConfig = ModelConfig(name="gpt-4.1-mini") - - # The wandb configuration - wandb: WandbConfig | None = None - - # The logging configuration - log: LogConfig = LogConfig() - - reasoning_field: Annotated[ - str, - Field( - description="The field in the raw model response that contains the reasoning content. Defaults to 'reasoning_content', which is the default for vLLM when serving a model with a reasoning parser. Other APIs (e.g. DeepSeek, GLM, etc.) may use different field names.", - ), - ] = "reasoning_content" - - output_dir: Annotated[ - Path, - Field( - description="Directory to write outputs to. Will be populated with artifacts such as reports and HF datasets as subdirectories. Should be set to a persistent directory with enough disk space." - ), - ] = Path("outputs") - - weights_dir: Annotated[ - Path | None, - Field( - description="Directory to load weight checkpoints (searches for `{weights_dir}/step_{x}`) generated during a training run (RL/ SFT). If set, will automatically eval all checkpoints found, including the base model. If None, will only eval the base model.", - ), - ] = None - - steps: Annotated[ - list[int] | None, - Field( - description="Steps to evaluate. If None, will evaluate all steps found in the weights directory. If set, will only evaluate the specified steps. If any of the specified steps are not found in the weights directory, will raise an error.", - ), - ] = None - - eval_base: Annotated[ - bool, - Field( - description="Whether to evaluate the base model. If True, will evaluate the base model before evaluating the checkpoints.", - ), - ] = True - - use_tqdm: Annotated[ - bool, - Field( - description="Whether to use tqdm to display progress bars during model generation.", - ), - ] = False - - max_concurrent: Annotated[ - int | None, - Field( - description="Maximum number of concurrent rollouts to generate and score. Will create a global semaphore and pass to verifiers Environment. If None, will not limit concurrency.", - ), - ] = None - - resume_path: Annotated[ - Path | None, - Field( - description="Path to the directory containing results.jsonl to resume from. When set, will read existing example_ids from the results.jsonl file, filter the dataset to exclude already-evaluated examples, and append new results to the existing file. Can be an absolute path or relative to the current working directory.", - ), - ] = None - - watcher: Annotated[ - bool, - Field( - description=( - "If True, watch `weights_dir` for newly-created stable checkpoints (folders named `step_{x}` that contain a `STABLE` file) " - "and evaluate them as they appear, instead of immediately iterating over all existing checkpoints." - ), - ), - ] = False - - @model_validator(mode="after") - def validate_eval_base(self): - if self.weights_dir is None and not self.eval_base: - raise ValueError( - "You should either evaluate the base model and/or checkpoints. Set `--eval-base` or specify a weight checkpoint directory with `--weights-dir`." - ) - return self - - @model_validator(mode="after") - def validate_resume_stream(self): - """Enforce per_rollout when resume_path or save.stream is enabled, as both require per-rollout scheduling.""" - if self.resume_path is not None: - if not self.save.stream: - raise ValueError( - "resume_path requires save.stream. Streaming saves are required for resume functionality." - ) - return self - - @model_validator(mode="after") - def validate_resume_multiple_environments(self): - """Prevent resume_path when multiple environments are configured, as they would share the same results.jsonl file.""" - if self.resume_path is not None and len(self.env) > 1: - raise ValueError( - f"resume_path cannot be used when evaluating multiple environments ({len(self.env)} environments configured). " - f"All environments will read from and write to the same file ({self.resume_path / 'results.jsonl'}), " - f"which may cause environments to incorrectly skip rollouts based on other environments' data. " - f"Consider evaluating environments separately." - ) - return self - - @model_validator(mode="after") - def validate_resume_multiple_checkpoints(self): - """Prevent resume_path when multiple checkpoints will be evaluated, as they would share the same results.jsonl file.""" - if self.resume_path is not None and self.weights_dir is not None: - # When weights_dir is set, multiple checkpoint evaluations will occur (at least 1 checkpoint, - # and possibly the base model if eval_base=True). All would write to the same results.jsonl file. - raise ValueError( - f"resume_path cannot be used when evaluating multiple checkpoints (weights_dir is set). " - f"When resume_path is set, all checkpoint evaluations (including base model if eval_base=True) " - f"would write to the same results.jsonl file ({self.resume_path / 'results.jsonl'}), " - f"causing later checkpoints to incorrectly skip examples that were already evaluated for earlier checkpoints. " - f"Consider evaluating checkpoints separately, each with its own resume_path, or remove resume_path to use checkpoint-specific output directories." - ) - return self diff --git a/src/prime_rl/eval/envs.py b/src/prime_rl/eval/envs.py deleted file mode 100644 index 50fff6e6a34..00000000000 --- a/src/prime_rl/eval/envs.py +++ /dev/null @@ -1,46 +0,0 @@ -from typing import TYPE_CHECKING, Any, List - -from prime_rl.utils.envs import ( - _ENV_PARSERS as _BASE_ENV_PARSERS, - get_env_value, - get_dir, - set_defaults, -) - -if TYPE_CHECKING: - # Enable type checking for shared envs - # ruff: noqa - from prime_rl.utils.envs import * - - # HF - HF_HUB_CACHE: str - HF_HUB_DISABLE_PROGRESS_BARS: str - HF_HUB_ETAG_TIMEOUT: int - - # vLLM - VLLM_CONFIGURE_LOGGING: int - - -_EVAL_ENV_PARSERS = { - "HF_HUB_CACHE": str, - "HF_HUB_DISABLE_PROGRESS_BARS": str, - "HF_HUB_ETAG_TIMEOUT": int, - "VLLM_CONFIGURE_LOGGING": int, - **_BASE_ENV_PARSERS, -} - -_EVAL_ENV_DEFAULTS = { - "HF_HUB_DISABLE_PROGRESS_BARS": "1", # Disable HF progress bars - "HF_HUB_ETAG_TIMEOUT": "500", # Set request timeout to 500s to avoid model download issues - "VLLM_CONFIGURE_LOGGING": "0", # Disable vLLM logging -} - -set_defaults(_EVAL_ENV_DEFAULTS) - - -def __getattr__(name: str) -> Any: - return get_env_value(_EVAL_ENV_PARSERS, name) - - -def __dir__() -> List[str]: - return get_dir(_EVAL_ENV_PARSERS) diff --git a/src/prime_rl/eval/eval.py b/src/prime_rl/eval/eval.py deleted file mode 100644 index ffb1d37aaa5..00000000000 --- a/src/prime_rl/eval/eval.py +++ /dev/null @@ -1,160 +0,0 @@ -import asyncio - -from prime_rl.eval.config import OfflineEvalConfig -from prime_rl.eval.utils import run_evals -from prime_rl.orchestrator.utils import set_semaphore -from prime_rl.utils.client import ( - check_has_model, - check_health, - reload_weights, - setup_admin_clients, - setup_clients, - setup_evals_client, - update_weights, -) -from prime_rl.utils.logger import intercept_verifiers_logging, setup_logger -from prime_rl.utils.monitor import setup_monitor -from prime_rl.utils.pydantic_config import parse_argv -from prime_rl.utils.utils import clean_exit, get_env_ids_to_install, get_step_path, install_env - - -@clean_exit -async def eval(config: OfflineEvalConfig): - # Initialize the logger - logger = setup_logger( - config.log.level, log_file=config.output_dir / "logs" / "eval.log" if config.log.file else None - ) - intercept_verifiers_logging(level=config.log.vf_level) - - env_names = [env.name or env.id for env in config.env] - logger.info(f"Starting evals for {config.model.name} in environments {', '.join(env_names)}") - logger.info(f"Using sampling config {config.sampling}") - - # Install environments - env_ids_to_install = get_env_ids_to_install(config.env) - for env_id in env_ids_to_install: - install_env(env_id) - - # Initialize the monitor - logger.info(f"Initializing monitor ({config.wandb})") - setup_monitor( - config=config.wandb, - output_dir=None, - run_config=config, - ) - - # Setup clients - logger.info( - f"Initializing OpenAI client (base_url={', '.join(config.client.base_url)}, api_key_var={config.client.api_key_var}, headers={config.client.headers})" - ) - clients = setup_clients(config.client) - admin_clients = setup_admin_clients(config.client) - evals_client = setup_evals_client() - - # Check health of the client - logger.info("Waiting for inference pool to be ready") - await check_health(admin_clients) - await check_has_model(clients, config.model.name) - logger.success(f"Inference pool is healthy and serves {config.model.name}") - - # Reset weights to base model to allow reusing inference server across runs - logger.info("Resetting weights to base model") - await reload_weights(admin_clients) - - # Run benchmarks on base model - await set_semaphore(config.max_concurrent or -1) - if config.eval_base: - logger.info(f"Evaluating model {config.model.name}") - await run_evals( - clients=clients, - eval_config=config, - model_config=config.model, - sampling_config=config.sampling, - evals_client=evals_client, - reasoning_field=config.reasoning_field, - output_dir=config.output_dir, - ckpt_step=0, - resume_path=config.resume_path, - ) - - # If specified, evaluate all checkpoints found in the weights directory - if config.weights_dir is not None: - logger.info(f"Evaluating weight checkpoints in {config.weights_dir}") - # Only include checkpoints that have a STABLE file (indicating save completed) - ckpt_steps = sorted( - [ - int(step_path.name.split("_")[-1]) - for step_path in config.weights_dir.glob("step_*") - if (step_path / "STABLE").exists() - ] - ) - logger.info(f"Found {len(ckpt_steps)} stable weight checkpoints (steps: {', '.join(map(str, ckpt_steps))})") - - if config.steps is not None: - missing_steps = [step for step in config.steps if step not in ckpt_steps] - if missing_steps: - raise ValueError( - f"Step {missing_steps[0]} not found (or not stable) in weights directory {config.weights_dir}" - ) - - # Filter the steps to evaluate - if config.steps is not None: - ckpt_steps = [step for step in ckpt_steps if step in config.steps] - - async def _eval_ckpt(ckpt_step: int) -> bool: - try: - logger.info(f"Evaluating model {config.model.name} at checkpoint {ckpt_step}") - await update_weights(admin_clients, get_step_path(config.weights_dir, ckpt_step)) - - await run_evals( - clients=clients, - eval_config=config, - model_config=config.model, - sampling_config=config.sampling, - evals_client=evals_client, - reasoning_field=config.reasoning_field, - output_dir=config.output_dir, - ckpt_step=ckpt_step, - resume_path=config.resume_path, - ) - return True - except Exception: - logger.exception(f"Checkpoint eval failed (ckpt_step={ckpt_step})") - return False - - logger.info(f"Evaluating {len(ckpt_steps)} weight checkpoints (steps: {', '.join(map(str, ckpt_steps))})") - for ckpt_step in ckpt_steps: - await _eval_ckpt(ckpt_step) - - if config.watcher: - # Keep watcher behavior consistent with previous versions: after evaluating all currently-stable - # checkpoints, watch for new ones and evaluate them as they appear. - already_evaluated_ckpt_steps = list(ckpt_steps) - while True: - # Only include checkpoints that have a STABLE file (indicating save completed) - all_ckpt_steps = sorted( - [ - int(step_path.name.split("_")[-1]) - for step_path in config.weights_dir.glob("step_*") - if (step_path / "STABLE").exists() - ] - ) - new_ckpt_steps = [step for step in all_ckpt_steps if step not in already_evaluated_ckpt_steps] - if len(new_ckpt_steps) > 0: - logger.info(f"New checkpoints to evaluate: {', '.join(map(str, new_ckpt_steps))}") - for ckpt_step in new_ckpt_steps: - await _eval_ckpt(ckpt_step) - already_evaluated_ckpt_steps.append(ckpt_step) - else: - logger.info("No new checkpoints to evaluate, waiting for 10 seconds") - await asyncio.sleep(10) - - logger.success("Eval finished!") - - -def main(): - asyncio.run(eval(parse_argv(OfflineEvalConfig))) - - -if __name__ == "__main__": - main() diff --git a/src/prime_rl/eval/registry.py b/src/prime_rl/eval/registry.py deleted file mode 100644 index 8133660c149..00000000000 --- a/src/prime_rl/eval/registry.py +++ /dev/null @@ -1,29 +0,0 @@ -from typing import Literal - -from datasets import Dataset, load_dataset - -Benchmark = Literal["math500", "aime24", "aime25", "livecodebench-v5"] - -_BENCHMARKS_DATASET_NAMES: dict[Benchmark, str] = { - "math500": "PrimeIntellect/MATH-500", - "aime24": "PrimeIntellect/AIME-24", - "aime25": "PrimeIntellect/AIME-25", - "livecodebench-v5": "PrimeIntellect/LiveCodeBench-v5", -} - -_BENCHMARK_DISPLAY_NAMES: dict[Benchmark, str] = { - "math500": "MATH-500", - "aime24": "AIME-24", - "aime25": "AIME-25", - "livecodebench-v5": "LiveCodeBench-V5", -} - - -def get_benchmark_dataset(name: Benchmark) -> Dataset: - dataset = load_dataset(_BENCHMARKS_DATASET_NAMES[name], split="train") - assert isinstance(dataset, Dataset) - return dataset - - -def get_benchmark_display_name(name: Benchmark) -> str: - return _BENCHMARK_DISPLAY_NAMES[name] diff --git a/src/prime_rl/eval/utils.py b/src/prime_rl/eval/utils.py deleted file mode 100644 index 4ccb63fe847..00000000000 --- a/src/prime_rl/eval/utils.py +++ /dev/null @@ -1,864 +0,0 @@ -import asyncio -import json -import re -import time -from copy import deepcopy -from itertools import cycle -from multiprocessing import Process -from pathlib import Path -from typing import Any - -import numpy as np -import pandas as pd -import verifiers as vf -from huggingface_hub import whoami -from openai import AsyncOpenAI, BadRequestError -from prime_evals import AsyncEvalsClient -from tenacity import RetryCallState, RetryError, retry, stop_after_attempt, wait_exponential -from tqdm.asyncio import tqdm -from verifiers import load_environment -from verifiers.envs.environment import get_results_path -from verifiers.utils.eval_utils import get_hf_hub_dataset_name, make_dataset, sanitize_metadata, save_to_disk - -from prime_rl.eval.config import OfflineEvalConfig -from prime_rl.orchestrator.config import EvalConfig, EvalSamplingConfig, EvalSaveConfig, ModelConfig, RetryConfig -from prime_rl.orchestrator.utils import set_semaphore -from prime_rl.synthesize.utils import merge_reasoning_content, save_result -from prime_rl.utils.client import setup_clients, setup_evals_client -from prime_rl.utils.config import ClientConfig -from prime_rl.utils.logger import get_logger, intercept_verifiers_logging, reset_logger, setup_logger -from prime_rl.utils.monitor import get_monitor -from prime_rl.utils.utils import capitalize, get_eval_dir, get_step_path -from prime_rl.utils.vf import ( - generate_group, - generate_rollout, - get_completion_len, - get_is_truncated, - to_serializable_state, -) - -WRITE_LOCK = asyncio.Lock() - - -def read_existing_rollout_ids(results_file: Path) -> set[tuple[int, int]]: - """Read existing (example_id, rollout_idx) pairs from results file for resume.""" - rollout_ids: set[tuple[int, int]] = set() - with open(results_file, "r") as f: - for line in f: - result = json.loads(line) - rollout_ids.add((result["example_id"], result["rollout_idx"])) - return rollout_ids - - -def read_existing_results(results_file: Path) -> pd.DataFrame: - results = [] - with open(results_file, "r") as f: - for line in f: - result = json.loads(line) - info = result.get("info", {}) - prime_info = info.get("prime_rl", {}) - rollout_status = prime_info.get("rollout_status", "ok") - results.append( - { - "example_id": result["example_id"], - "reward": result["reward"], - "completion_len": result["completion_len"], - "is_truncated": result["is_truncated"], - "rollout_status": rollout_status, - } - ) - - return pd.DataFrame(results) - - -def read_completed_example_ids(results_file: Path, rollouts_per_example: int) -> set[int]: - """Read example_ids that have complete groups (all rollouts_per_example completed) for group-based resume.""" - example_id_counts: dict[int, int] = {} - with open(results_file, "r") as f: - for line in f: - result = json.loads(line) - example_id = result["example_id"] - example_id_counts[example_id] = example_id_counts.get(example_id, 0) + 1 - - # Return example_ids that have exactly rollouts_per_example results - completed_example_ids = { - example_id for example_id, count in example_id_counts.items() if count == rollouts_per_example - } - return completed_example_ids - - -def compute_pass_at_k(rewards: list[int]) -> dict[str, float]: - total_attempts = len(rewards) - k = total_attempts // 2 - - if k == 0: - return {"pass@1": float(any(reward == 1.0 for reward in rewards))} - - num_trials = 100 - pass_rates = [] - - for _ in range(num_trials): - sampled_rewards = np.random.choice(rewards, size=k, replace=False) - pass_rate = float(any(reward == 1.0 for reward in sampled_rewards)) - pass_rates.append(pass_rate) - - return {f"pass@{k}": float(np.mean(pass_rates))} - - -def prepare_sampling_args(sampling_config: EvalSamplingConfig) -> dict[str, Any]: - """Prepare sampling args for the client.""" - # Initialize sampling args - sampling_args: dict[str, Any] = {} - - # Apply sampling arguments, if specified - if sampling_config.temperature is not None: - sampling_args["temperature"] = sampling_config.temperature - if sampling_config.max_tokens is not None: - sampling_args["max_tokens"] = sampling_config.max_tokens - if sampling_config.top_p is not None: - sampling_args["top_p"] = sampling_config.top_p - if sampling_config.reasoning_effort is not None: - sampling_args["reasoning_effort"] = sampling_config.reasoning_effort - - extra_body: dict[str, Any] = sampling_config.extra_body.copy() - - # Apply vLLM-specific sampling arguments, if specified - if sampling_config.top_k is not None: - extra_body["top_k"] = sampling_config.top_k - if sampling_config.min_p is not None: - extra_body["min_p"] = sampling_config.min_p - if sampling_config.min_tokens is not None: - extra_body["min_tokens"] = sampling_config.min_tokens - if sampling_config.repetition_penalty is not None: - extra_body["repetition_penalty"] = sampling_config.repetition_penalty - - sampling_args["extra_body"] = extra_body - - return sampling_args - - -# TODO: Move to verifiers to avoid code drift -def make_result(state: vf.State, reasoning_field: str, rollout_idx: int) -> dict: - """Translates a finished rollout state to a synthetic dataset row.""" - completion = merge_reasoning_content(state["completion"], state["trajectory"], reasoning_field) - result_dict = { - "example_id": state["example_id"], - "rollout_idx": rollout_idx, - "prompt": state["prompt"], - "completion": completion, - "task": state["task"], - "reward": state["reward"], - "generation_ms": state["timing"]["generation_ms"], - "scoring_ms": state["timing"]["scoring_ms"], - "total_ms": state["timing"]["total_ms"], - "info": state.get("info", {}), - "answer": state.get("answer", ""), - "completion_len": get_completion_len(state), - "is_truncated": get_is_truncated(state), - } - for metric_name, metric_value in state["metrics"].items(): - result_dict[metric_name] = metric_value - - result_dict["oai_tools"] = json.dumps(state["oai_tools"]) - - return result_dict - - -async def make_and_save_result(state: vf.State, save_file: Path, reasoning_field: str, rollout_idx: int): - """Translates and saves a finished rollout state to a synthetic dataset row.""" - result_dict = await asyncio.to_thread(make_result, state, reasoning_field, rollout_idx) - await save_result(result_dict, save_file) - - -def log_retry_attempt(retry_state: RetryCallState) -> None: - """Log retry attempts at WARNING level using the global logger.""" - logger = get_logger() - exception = retry_state.outcome.exception() - wait_time = retry_state.next_action.sleep - logger.warning( - f"Retrying {retry_state.fn.__name__} in {wait_time:.1f} seconds as it raised {exception.__class__.__name__}: {exception}" - ) - - -def parse_and_calculate_max_tokens(error_message: str) -> int | None: - """ - Example error message: - "This endpoint's maximum context length is 131072 tokens. However, you requested - about 131419 tokens (347 of text input, 131072 in the output)." - """ - context_match = re.search(r"maximum context length is (\d+) tokens", error_message) - prompt_match = re.search(r"(\d+) of text input", error_message) - - if context_match and prompt_match: - context_length = int(context_match.group(1)) - prompt_tokens = int(prompt_match.group(1)) - max_tokens = context_length - prompt_tokens - if max_tokens < 1: - return None - return max_tokens - return None - - -def _make_no_response_state( - *, - example: dict, - rollout_idx: int, - env_name_or_id: str, - error: Exception, -) -> dict: - """Create a vf.State-like dict for a rollout that never produced a model response.""" - # NOTE: Keep completion/trajectory empty so merge_reasoning_content() doesn't assert. - return { - "example_id": example.get("example_id"), - "prompt": example.get("prompt"), - "completion": [], - "task": example.get("task"), - "reward": None, - "timing": {"generation_ms": 0.0, "scoring_ms": 0.0, "total_ms": 0.0}, - "info": { - "prime_rl": { - "rollout_status": "no_response", - "env": env_name_or_id, - "rollout_idx": rollout_idx, - "error": repr(error), - } - }, - "answer": example.get("answer", ""), - "metrics": {}, - "trajectory": [], - "oai_tools": [], - } - - -def _is_valid_reward(reward: Any) -> bool: - if reward is None: - return False - if isinstance(reward, float) and np.isnan(reward): - return False - return isinstance(reward, (int, float)) - - -def _format_avg_reward(rewards: list) -> str: - valid = [r for r in rewards if _is_valid_reward(r)] - if not valid: - return "N/A" - return f"{(sum(valid) / len(valid)):.4f}" - - -async def generate_and_save_rollout( - client: AsyncOpenAI, - env: vf.Environment, - model_name: str, - example: dict, - rollout_idx: int, - env_name_or_id: str, - sampling_args: dict, - save_file: Path | None, - reasoning_field: str, - retry_config: RetryConfig, - pbar: tqdm, - rewards_accumulator: list, - rewards_lock: asyncio.Lock, -) -> vf.State: - """Generate and optionally save a single rollout, updating progress bar per-rollout.""" - logger = get_logger() - _sampling_args = deepcopy(sampling_args) - - @retry( - stop=stop_after_attempt(retry_config.max_attempts), - wait=wait_exponential( - multiplier=retry_config.wait_multiplier, min=retry_config.wait_min, max=retry_config.wait_max - ), - before_sleep=log_retry_attempt, - reraise=retry_config.reraise, - ) - async def _generate_rollout( - client: AsyncOpenAI, - env: vf.Environment, - model_name: str, - example: dict, - sampling_args: dict, - ) -> vf.State: - """Asynchronously generate and score a single rollout.""" - logger = get_logger() - try: - state = await generate_rollout(client, env, model_name, example, sampling_args) - # Re-raise infrastructure errors to trigger retry - if state.get("error") and isinstance(state["error"], vf.InfraError): - raise state["error"] - return state - except BadRequestError as e: - # Check if this is a context length error and retry with adjusted max_tokens - error_message = str(e) - new_max_tokens = parse_and_calculate_max_tokens(error_message) - - if new_max_tokens is not None: - logger.warning(f"Context length error: reducing max_tokens to {new_max_tokens}.") - sampling_args["max_tokens"] = new_max_tokens - state = await generate_rollout(client, env, model_name, example, sampling_args) - if state.get("error") and isinstance(state["error"], vf.InfraError): - raise state["error"] - return state - raise - - try: - state = await _generate_rollout(client, env, model_name, example, _sampling_args) - except (RetryError, Exception) as e: - # IMPORTANT: do not raise here (asyncio.gather would cancel all other rollouts). - logger.exception( - f"Rollout generation failed after retries; recording no_response " - f"(env={env_name_or_id}, example_id={example.get('example_id')}, rollout_idx={rollout_idx}): {repr(e)}" - ) - failed_state = _make_no_response_state( - example=example, rollout_idx=rollout_idx, env_name_or_id=env_name_or_id, error=e - ) - - # Save placeholder row if streaming saves enabled (supports resume + complete accounting). - if save_file is not None: - result_dict = await asyncio.to_thread(make_result, failed_state, reasoning_field, rollout_idx) - await save_result(result_dict, save_file) - - # Update progress + running average without biasing downward. - async with rewards_lock: - rewards_accumulator.append(None) - pbar.set_postfix({"Avg Reward": _format_avg_reward(rewards_accumulator)}) - pbar.update(1) - - return failed_state # type: ignore[return-value] - - # Save with rollout_idx if streaming saves enabled. - if save_file is not None: - result_dict = await asyncio.to_thread(make_result, state, reasoning_field, rollout_idx) - await save_result(result_dict, save_file) - - # Update running average immediately. - async with rewards_lock: - rewards_accumulator.append(state["reward"]) - pbar.set_postfix({"Avg Reward": _format_avg_reward(rewards_accumulator)}) - pbar.update(1) - - return state - - -async def generate_and_save_group( - client: AsyncOpenAI, - env: vf.Environment, - model_name: str, - example: dict, - rollouts_per_example: int, - env_name_or_id: str, - sampling_args: dict, - save_file: Path | None, - reasoning_field: str, - retry_config: RetryConfig, - pbar: tqdm, - rewards_accumulator: list, - rewards_lock: asyncio.Lock, -) -> list[vf.State]: - """Generate a group of rollouts, save results, and update progress (group-level).""" - logger = get_logger() - _sampling_args = deepcopy(sampling_args) - - @retry( - stop=stop_after_attempt(retry_config.max_attempts), - wait=wait_exponential( - multiplier=retry_config.wait_multiplier, min=retry_config.wait_min, max=retry_config.wait_max - ), - before_sleep=log_retry_attempt, - reraise=retry_config.reraise, - ) - async def _generate_group( - client: AsyncOpenAI, - env: vf.Environment, - model_name: str, - example: dict, - rollouts_per_example: int, - sampling_args: dict, - ) -> list[vf.State]: - """Asynchronously generate and score a group of rollouts.""" - logger = get_logger() - try: - states = await generate_group(client, env, model_name, example, rollouts_per_example, sampling_args) - # Re-raise infrastructure errors to trigger retry (check all states in group) - for state in states: - if state.get("error") and isinstance(state["error"], vf.InfraError): - raise state["error"] - return states - except BadRequestError as e: - error_message = str(e) - new_max_tokens = parse_and_calculate_max_tokens(error_message) - - if new_max_tokens is not None: - logger.warning(f"Context length error: reducing max_tokens to {new_max_tokens}.") - sampling_args["max_tokens"] = new_max_tokens - states = await generate_group(client, env, model_name, example, rollouts_per_example, sampling_args) - for state in states: - if state.get("error") and isinstance(state["error"], vf.InfraError): - raise state["error"] - return states - raise - - try: - states = await _generate_group(client, env, model_name, example, rollouts_per_example, _sampling_args) - except (RetryError, Exception) as e: - # IMPORTANT: do not raise here (asyncio.gather would cancel other groups). - logger.exception( - f"Group generation failed after retries; recording no_response rollouts " - f"(env={env_name_or_id}, example_id={example.get('example_id')}): {repr(e)}" - ) - failed_states = [ - _make_no_response_state( - example=example, rollout_idx=rollout_idx, env_name_or_id=env_name_or_id, error=e - ) - for rollout_idx in range(rollouts_per_example) - ] - - if save_file is not None: - await asyncio.gather( - *[ - make_and_save_result(state, save_file, reasoning_field, rollout_idx) - for rollout_idx, state in enumerate(failed_states) - ] - ) - - async with rewards_lock: - rewards_accumulator.extend([None] * rollouts_per_example) - pbar.set_postfix({"Avg Reward": _format_avg_reward(rewards_accumulator)}) - pbar.update(rollouts_per_example) - - return failed_states # type: ignore[return-value] - - # Save results for the generated group. - if save_file is not None: - await asyncio.gather( - *[ - make_and_save_result(state, save_file, reasoning_field, rollout_idx) - for rollout_idx, state in enumerate(states) - ] - ) - - # Update running average after group completes. - async with rewards_lock: - for state in states: - rewards_accumulator.append(state["reward"]) - pbar.set_postfix({"Avg Reward": _format_avg_reward(rewards_accumulator)}) - pbar.update(rollouts_per_example) - - return states - - -async def run_eval( - clients: list[AsyncOpenAI], - env_id: str, - env_name: str | None, - env_args: dict, - num_examples: int, - rollouts_per_example: int, - reasoning_field: str, - output_dir: Path, - ckpt_step: int, - model_config: ModelConfig, - sampling_config: EvalSamplingConfig, - save_config: EvalSaveConfig, - retry_config: RetryConfig, - evals_client: AsyncEvalsClient, - per_rollout: bool = False, - step: int | None = None, - resume_path: Path | None = None, -) -> None: - # Get the logger - logger = get_logger() - monitor = get_monitor() - eval_start_time = time.perf_counter() - - # Load the eval environment - env_name_or_id = env_name or env_id - env = load_environment(env_id, **env_args) - dataset = env.get_eval_dataset(n=num_examples) - sampling_args = prepare_sampling_args(sampling_config) - - # Determine streaming save path - if save_config.stream: - if resume_path is not None: - # resume_path points to a directory containing results.jsonl - path_to_save = resume_path / "results.jsonl" - else: - base_path = get_results_path(env_name_or_id, model_config.name, base_path=output_dir) - path_to_save = base_path / "results.jsonl" - path_to_save.parent.mkdir(parents=True, exist_ok=True) - - # Create shared structure for tracking rewards - rewards_accumulator: list = [] - rewards_lock = asyncio.Lock() - examples = dataset.to_list() - - if per_rollout: - # Per-rollout scheduling: enables live progress updates and per-rollout resume - all_rollouts: list[tuple[dict, int]] = [ - (example, rollout_idx) for example in examples for rollout_idx in range(rollouts_per_example) - ] - - # Filter out already-completed rollouts on resume - if resume_path is not None: - existing_rollout_ids = read_existing_rollout_ids(path_to_save) - original_count = len(all_rollouts) - all_rollouts = [ - (ex, ridx) for ex, ridx in all_rollouts if (ex["example_id"], ridx) not in existing_rollout_ids - ] - skipped_count = original_count - len(all_rollouts) - logger.info( - f"Resuming from {path_to_save}: skipping {skipped_count} already-completed rollouts, " - f"{len(all_rollouts)} remaining" - ) - # Populate rewards_accumulator with existing rewards - existing_results_df = read_existing_results(path_to_save) - rewards_accumulator.extend(existing_results_df.reward.tolist()) - - total_rollouts = len(all_rollouts) - logger.info( - f"Evaluating {env_name_or_id} ({num_examples=}, {rollouts_per_example=}) " - f"{'with default args' if env_args == {} else f'with args {env_args}'} and extra_body {sampling_args['extra_body']}\n" - f"{'Saving results to ' + str(path_to_save) if save_config.stream else 'Results will be saved at end of evaluation'}" - ) - - pbar = tqdm(total=total_rollouts, desc="Evaluating") - if rewards_accumulator: - pbar.set_postfix({"Avg Reward": _format_avg_reward(rewards_accumulator)}) - - all_states = await asyncio.gather( - *[ - generate_and_save_rollout( - client, - env, - model_config.name, - example, - rollout_idx, - env_name_or_id, - sampling_args, - path_to_save if save_config.stream else None, - reasoning_field, - retry_config, - pbar, - rewards_accumulator, - rewards_lock, - ) - for client, (example, rollout_idx) in zip(cycle(clients), all_rollouts) - ] - ) - else: - # Group-based scheduling: preserves group-based rubric scoring - # Filter out completed example_ids on resume - if resume_path is not None and path_to_save is not None and path_to_save.exists(): - completed_example_ids = read_completed_example_ids(path_to_save, rollouts_per_example) - original_count = len(examples) - examples = [ex for ex in examples if ex["example_id"] not in completed_example_ids] - skipped_count = original_count - len(examples) - logger.info( - f"Resuming from {path_to_save}: skipping {skipped_count} already-completed example_ids " - f"({len(completed_example_ids)} complete groups), {len(examples)} remaining" - ) - # Populate rewards_accumulator with existing rewards for metrics - existing_results_df = read_existing_results(path_to_save) - rewards_accumulator.extend(existing_results_df.reward.tolist()) - - total_rollouts = len(examples) * rollouts_per_example - logger.info( - f"Evaluating {env_name_or_id} ({num_examples=}, {rollouts_per_example=}) " - f"{'with default args' if env_args == {} else f'with args {env_args}'} and extra_body {sampling_args['extra_body']}\n" - f"{'Saving results to ' + str(path_to_save) if save_config.stream else 'Results will be saved at end of evaluation'}" - ) - - pbar = tqdm(total=total_rollouts, desc="Evaluating") - if rewards_accumulator: - pbar.set_postfix({"Avg Reward": _format_avg_reward(rewards_accumulator)}) - - group_states_list = await asyncio.gather( - *[ - generate_and_save_group( - client, - env, - model_config.name, - example, - rollouts_per_example, - env_name_or_id, - sampling_args, - path_to_save if save_config.stream else None, - reasoning_field, - retry_config, - pbar, - rewards_accumulator, - rewards_lock, - ) - for client, example in zip(cycle(clients), examples) - ] - ) - all_states = [state for group in group_states_list for state in group] - - k = rollouts_per_example - new_results_df = pd.DataFrame( - { - "example_id": [state["example_id"] for state in all_states], - "reward": [state["reward"] for state in all_states], - "completion_len": [get_completion_len(state) for state in all_states], - "is_truncated": [get_is_truncated(state) for state in all_states], - "rollout_status": [ - (state.get("info", {}) or {}).get("prime_rl", {}).get("rollout_status", "ok") - for state in all_states - ], - } - ) - - # If resuming, combine with existing results for accurate metrics - if resume_path is not None: - existing_results_df = read_existing_results(path_to_save) - results_df = pd.concat([existing_results_df, new_results_df], ignore_index=True) - logger.info( - f"Combined existing ({len(existing_results_df)}) and new ({len(new_results_df)}) results for metrics" - ) - else: - results_df = new_results_df - - unique_rewards = results_df.reward.dropna().unique() - could_be_binary = set(unique_rewards).issubset({0.0, 1.0}) - if could_be_binary: - pass_at_k = ( - results_df.groupby("example_id") - .apply(lambda x: compute_pass_at_k(x.reward.dropna()), include_groups=False) - .apply(pd.Series) - ) - else: - pass_at_k = None - logger.warning("Skipping computing pass@k rates because the task rewards appear to be non-binary") - - # Log statistics to console - eval_time = time.perf_counter() - eval_start_time - no_response_rate = float((results_df.rollout_status == "no_response").mean()) if "rollout_status" in results_df else 0.0 - message = f"Evaluated {env_name_or_id} in {eval_time:.2f}s (Avg@{k}={results_df.reward.mean():.4f}" - if could_be_binary: - assert pass_at_k is not None - for pass_rate, pass_rate_score in pd.Series(pass_at_k.mean()).items(): - message += f", {capitalize(str(pass_rate))}: {pass_rate_score:.4f}" - message += ( - f", No-response: {no_response_rate * 100:.1f}%" - f", Completion Length: {results_df.completion_len.mean():.2f} (±{results_df.completion_len.std():.2f}, ∈[{results_df.completion_len.min():.2f}, {results_df.completion_len.max():.2f}])" - f", Truncated: {results_df.is_truncated.mean() * 100:.1f}%)" - ) - logger.success(message) - - # Log statistics to monitor - eval_metrics = { - f"avg@{k}": results_df.reward.mean(), - "no_response/pct": no_response_rate * 100.0, - "no_response/count": int((results_df.rollout_status == "no_response").sum()) if "rollout_status" in results_df else 0, - "completion_len/avg": results_df.completion_len.mean().item(), - "completion_len/max": results_df.completion_len.max().item(), - "completion_len/min": results_df.completion_len.min().item(), - "is_truncated/mean": results_df.is_truncated.mean().item(), - "time": eval_time, - } - if could_be_binary: - assert pass_at_k is not None - eval_metrics.update(pd.Series(pass_at_k.mean()).to_dict()) - eval_metrics = {**{f"eval/{env_name_or_id}/{k}": v for k, v in eval_metrics.items()}} - eval_metrics.update({"progress/ckpt_step": ckpt_step, "step": step or ckpt_step}) - monitor.log(eval_metrics, step=None) - - # Save results - if save_config.disk is not None or save_config.hf is not None or save_config.env_hub: - outputs = env._prepare_rollout_results( - all_states=[to_serializable_state(state) for state in all_states], # type: ignore - model=model_config.name, - client=clients[0], # We use the first client - state_columns=None, - results_path=None, - gen_sampling_args=sampling_args, - start_time=eval_start_time, - ) - dataset = make_dataset(outputs) - metadata_dict = sanitize_metadata(outputs["metadata"]) - - if save_config.disk is not None: - is_online = step is not None - default_save_path = ( - get_step_path(get_eval_dir(output_dir), ckpt_step) / env_name_or_id - if is_online - else outputs["metadata"]["path_to_save"] - ) - save_path = save_config.disk.path or default_save_path - save_to_disk(dataset, metadata_dict, save_path) - logger.info(f"Saved eval results for {env_name_or_id} to disk ({save_path})") - - if save_config.hf is not None: - dataset_name = save_config.hf.dataset_name or get_hf_hub_dataset_name(outputs) - dataset_subset = save_config.hf.dataset_subset or env.env_id - dataset_split = save_config.hf.dataset_split or "evals" - dataset.push_to_hub(dataset_name, dataset_subset, split=dataset_split, private=save_config.hf.private) - default_org = whoami().get("name", "") - repo_name = dataset_name if "/" in dataset_name else f"{default_org}/{dataset_name}" - logger.info( - f"Pushed {'private' if save_config.hf.private else 'public'} eval results for {env_name_or_id} to HF Hub (https://huggingface.co/datasets/{repo_name})" - ) - - if save_config.env_hub: - eval_name = f"{env_id}--{model_config.name.replace('/', '--')}" - - # Create evaluation for environment - create_response = await evals_client.create_evaluation( - name=eval_name, - environments=[{"id": env_id}], - model_name=model_config.name, - framework="verifiers", - metadata=metadata_dict, - metrics=eval_metrics, - ) - - eval_id = create_response.get("evaluation_id") - assert eval_id is not None - - # Push samples - await evals_client.push_samples(eval_id, dataset.to_list()) - - # Finalize evaluation - await evals_client.finalize_evaluation(eval_id, metrics=eval_metrics) - - logger.info( - f"Pushed eval results for {env_id} to Environments Hub (https://app.primeintellect.ai/dashboard/evaluations/{eval_id})" - ) - - -async def run_evals( - clients: list[AsyncOpenAI], - eval_config: EvalConfig | OfflineEvalConfig, - model_config: ModelConfig, - sampling_config: EvalSamplingConfig, - evals_client: AsyncEvalsClient, - reasoning_field: str, - output_dir: Path, - ckpt_step: int, - step: int | None = None, - resume_path: Path | None = None, -): - logger = get_logger() - - async def run_eval_safe(env): - """Wrapper to catch RetryError and skip env.""" - try: - await run_eval( - clients=clients, - env_id=env.id, - env_name=env.name, - env_args=env.args, - num_examples=env.num_examples or eval_config.num_examples, - reasoning_field=reasoning_field, - rollouts_per_example=env.rollouts_per_example or eval_config.rollouts_per_example, - output_dir=output_dir, - model_config=model_config, - sampling_config=sampling_config, - save_config=eval_config.save, - retry_config=eval_config.retry, - evals_client=evals_client, - per_rollout=eval_config.per_rollout, - ckpt_step=ckpt_step, - step=step, - resume_path=resume_path, - ) - return True - except RetryError as e: - env_name = env.name or env.id - logger.warning( - f"Env '{env_name}' exhausted {eval_config.retry.max_attempts} retry attempts, skipping. " - f"Last error: {e.last_attempt.exception()}" - ) - return None - - results = await asyncio.gather(*[run_eval_safe(env) for env in eval_config.env]) - - skipped = sum(1 for r in results if r is None) - if skipped > 0: - logger.warning(f"Skipped {skipped}/{len(eval_config.env)} environments due to retry exhaustion") - - -def _run_evals_in_subprocess( - client_config: ClientConfig, - eval_config: EvalConfig, - model_config: ModelConfig, - sampling_config: EvalSamplingConfig, - reasoning_field: str, - output_dir: str, - ckpt_step: int, - step: int | None, - max_concurrent: int, -): - """Entry point for eval subprocess. Creates its own event loop and clients.""" - # Setup logger for subprocess (reset first since we inherit parent's global state when forked) - reset_logger() - logger = setup_logger("info") - intercept_verifiers_logging(level="info") - logger.info(f"Eval subprocess started for checkpoint step {ckpt_step}") - - # Create fresh clients in subprocess - clients = setup_clients(client_config) - evals_client = setup_evals_client() - - async def _run(): - await set_semaphore(max_concurrent) - await run_evals( - clients=clients, - eval_config=eval_config, - model_config=model_config, - sampling_config=sampling_config, - evals_client=evals_client, - reasoning_field=reasoning_field, - output_dir=Path(output_dir), - ckpt_step=ckpt_step, - step=step, - ) - - asyncio.run(_run()) - logger.info(f"Eval subprocess finished for checkpoint step {ckpt_step}") - - -async def run_evals_subprocess( - client_config: ClientConfig, - eval_config: EvalConfig | OfflineEvalConfig, - model_config: ModelConfig, - sampling_config: EvalSamplingConfig, - reasoning_field: str, - output_dir: Path, - ckpt_step: int, - step: int | None = None, - max_concurrent: int = -1, -): - """Run evals in a separate subprocess to isolate the event loop. - - This prevents eval's async work from causing event loop lag in the - orchestrator's main loop. The orchestrator blocks until eval completes. - """ - logger = get_logger() - logger.info(f"Spawning eval subprocess for checkpoint step {ckpt_step}") - - process = Process( - target=_run_evals_in_subprocess, - args=( - client_config, - eval_config, - model_config, - sampling_config, - reasoning_field, - str(output_dir), - ckpt_step, - step, - max_concurrent, - ), - daemon=True, - ) - process.start() - - # Wait for process to complete without blocking the event loop - while process.is_alive(): - await asyncio.sleep(0.5) - - if process.exitcode != 0: - logger.error(f"Eval subprocess failed with exit code {process.exitcode}") - else: - logger.success(f"Eval subprocess completed for checkpoint step {ckpt_step}") diff --git a/src/prime_rl/inference/config.py b/src/prime_rl/inference/config.py deleted file mode 100644 index 2a982feb549..00000000000 --- a/src/prime_rl/inference/config.py +++ /dev/null @@ -1,276 +0,0 @@ -import os -from argparse import Namespace -from typing import Annotated, Any, Literal - -from pydantic import Field, model_validator - -from prime_rl.utils.pydantic_config import BaseConfig, BaseSettings, get_all_fields -from prime_rl.utils.utils import rgetattr, rsetattr - -# TODO: Set thinking/ solution budget - - -class ServerConfig(BaseConfig): - """Configures the inference server.""" - - host: Annotated[str | None, Field(description="The host to bind to.")] = None - port: Annotated[int, Field(description="The port to bind to.")] = 8000 - - -class ParallelConfig(BaseConfig): - """Configures multi-node and multi-GPU setups through different types of parallelism (TP, DP, PP).""" - - tp: Annotated[ - int, - Field( - description="The tensor parallel size. It is passed to vLLM as `--tensor-parallel-size`", - ), - ] = 1 - - dp: Annotated[ - int, - Field( - ge=1, - description="The data parallel size. It is passed to vLLM as `--data-parallel-size`", - ), - ] = 1 - - def __str__(self) -> str: - return f"tp={self.tp} dp={self.dp}" - - -class ModelConfig(BaseConfig): - """Configures the inference model. Most arguments are passed directly to the vLLM LLM class (https://docs.vllm.ai/en/latest/api/vllm.LLM.html).""" - - name: Annotated[ - str, - Field( - description="Name or path of the HF model to use.", - ), - ] = "Qwen/Qwen3-0.6B" - - dtype: Annotated[ - Literal["auto", "float16", "bfloat16", "float32"], - Field( - description="Data type for model weights and activations. If 'auto' will use FP16 precision for FP32 and FP16 models, and BF16 precision for BF16 models. Passed to vLLM as `--dtype`", - ), - ] = "auto" - - max_model_len: Annotated[ - int | None, - Field( - description="Maximum model context length. If None, will use the maximum context length from model config. Passed to vLLM as `--max-model-len`", - ), - ] = None - - enforce_eager: Annotated[ - bool, - Field( - description="Whether to enforce eager mode. If False, will use PyTorch eager and cuda graphs in hybrid for maximal performance. Passed to vLLM as `--enforce-eager`", - ), - ] = False - - trust_remote_code: Annotated[ - bool, - Field( - description="Whether to trust remote code. Passed to vLLM engine init", - ), - ] = False - - enable_auto_tool_choice: Annotated[ - bool, - Field( - description="Whether to enable auto tool choice. Passed to vLLM as `--enable-auto-tool-choice`", - ), - ] = False - - tool_call_parser: Annotated[ - str, - Field( - description="The tool call parser to use. Passed to vLLM as `--tool-call-parser`", - ), - ] = "hermes" - - reasoning_parser: Annotated[ - str | None, - Field( - description="Parser for extracting reasoning content from model outputs. Passed to vLLM as `--reasoning-parser`. Setting this enables reasoning mode.", - ), - ] = None - - rope_scaling: Annotated[ - dict[str, Any] | str | None, - Field( - description='RoPE scaling configuration as a dict. For YaRN, use: {rope_type="yarn", factor=4.0, original_max_position_embeddings=32768} or. Passed to vLLM as `--rope-scaling`.', - ), - ] = None - - -class WeightBroadcastConfig(BaseSettings): - """Configures weight broadcast settings.""" - - type: Annotated[Literal["nccl", "filesystem"], Field(description="The type of weight broadcast to use.")] = ( - "filesystem" - ) - - -# Valid vLLM max_lora_rank values (from vllm/config/lora.py) -# TODO: on newer vLLM, can import via `get_args(vllm.config.lora.MaxLoRARanks)` -VALID_VLLM_LORA_RANKS = (8, 16, 32, 64, 128, 256, 320, 512) - - -class InferenceConfig(BaseSettings): - """Configures inference.""" - - # The server configuration - server: ServerConfig = ServerConfig() - - # The model configuration - model: ModelConfig = ModelConfig() - - # The parallel configuration - parallel: ParallelConfig = ParallelConfig() - - enable_lora: Annotated[ - bool, - Field( - description="Whether to enable LORA. Passed to vLLM as `--enable-lora`", - ), - ] = False - - max_loras: Annotated[ - int, - Field( - description="The maximum number of LoRAs to use. Passed to vLLM as `--max-loras`", - ), - ] = 8 - - # TODO: The default value is very high because our areal impl for lora isn't ideal - # We add a lora with the same name instead of changing weights inplace - # Because we dont cancel requests that are past max_async, these requests could be using a LoRA that gets unloaded which will crash the inference server - max_cpu_loras: Annotated[ - int, - Field( - description="The maximum number of LoRAs to use on CPU. Passed to vLLM as `--max-cpu-loras`", - ), - ] = 100 - - max_lora_rank: Annotated[ - int | None, - Field( - description="The maximum LoRA rank to use. Passed to vLLM as `--max-lora-rank`", - ), - ] = None - - gpu_memory_utilization: Annotated[ - float, - Field( - description="The GPU memory utilization to use. Passed to vLLM as `--gpu-memory-utilization`", - ), - ] = 0.9 - - api_server_count: Annotated[ - int, - Field( - ge=1, - description="The number of API servers to use. Passed to vLLM as `--api-server-count`", - ), - ] = 1 - - seed: Annotated[ - int, - Field( - description="Seed the inference components. Passed to vLLM as `--seed`", - ), - ] = 0 - - weight_broadcast: Annotated[WeightBroadcastConfig, Field(description="The weight broadcast config.")] = ( - WeightBroadcastConfig() - ) - - @model_validator(mode="after") - def nccl_and_dp(self): - if self.weight_broadcast.type == "nccl" and self.parallel.dp != 1: - raise ValueError("NCCL broadcast backend requires data parallel size to be 1") - return self - - @model_validator(mode="after") - def auto_setup_dynamic_lora_updating(self): - if self.enable_lora: - os.environ["VLLM_ALLOW_RUNTIME_LORA_UPDATING"] = "True" - return self - - @model_validator(mode="after") - def round_up_max_lora_rank(self): - """Round up max_lora_rank to the nearest valid vLLM value. - - vLLM only accepts specific values for max_lora_rank: (1, 8, 16, 32, 64, 128, 256, 320, 512). - This validator ensures that any configured rank is rounded up to the minimum valid value - that can serve adapters of the requested rank. - """ - if self.max_lora_rank is not None: - original_rank = self.max_lora_rank - for valid_rank in VALID_VLLM_LORA_RANKS: - if valid_rank >= self.max_lora_rank: - self.max_lora_rank = valid_rank - break - else: - raise ValueError(f"max_lora_rank={original_rank} exceeds vLLM maximum of {VALID_VLLM_LORA_RANKS[-1]}") - return self - - @model_validator(mode="after") - def auto_setup_api_server_count(self): - """ - Ensures that we have at least as many API servers as data parallel - size. Unless LoRA is enabled, in which case only one API server is - supported (vLLM limitation). - """ - if self.api_server_count < self.parallel.dp: - self.api_server_count = self.parallel.dp - - if self.enable_lora: - self.api_server_count = 1 # LoRA requires only one API server - return self - - def to_vllm(self) -> Namespace: - """Convert InferenceConfig to vLLM-compatible Namespace.""" - namespace = Namespace() - to_vllm = { - "server.host": "host", - "server.port": "port", - "model.name": "model", - "model.dtype": "dtype", - "model.max_model_len": "max_model_len", - "model.enforce_eager": "enforce_eager", - "model.trust_remote_code": "trust_remote_code", - "model.enable_auto_tool_choice": "enable_auto_tool_choice", - "model.tool_call_parser": "tool_call_parser", - "model.reasoning_parser": "reasoning_parser", - "model.rope_scaling": "rope_scaling", - "parallel.tp": "tensor_parallel_size", - "parallel.dp": "data_parallel_size", - "enable_lora": "enable_lora", - "max_loras": "max_loras", - "max_cpu_loras": "max_cpu_loras", - "max_lora_rank": "max_lora_rank", - "gpu_memory_utilization": "gpu_memory_utilization", - "api_server_count": "api_server_count", - } - - for key in get_all_fields(self): - value = rgetattr(self, key.replace("-", "_")) - rsetattr(namespace, to_vllm.get(key, key), value) - - # Set `logprobs_mode` to `processed_logprobs` by default - rsetattr(namespace, "logprobs_mode", "processed_logprobs") - - # Remove reasoning_parser if not set (vLLM doesn't accept None) - if namespace.reasoning_parser is None: - delattr(namespace, "reasoning_parser") - - # Remove rope_scaling if not set (vLLM doesn't accept None) - if hasattr(namespace, "rope_scaling"): - if namespace.rope_scaling is None: - delattr(namespace, "rope_scaling") - - return namespace diff --git a/src/prime_rl/inference/patches.py b/src/prime_rl/inference/patches.py index b56e1cffd51..c33f3ccf423 100644 --- a/src/prime_rl/inference/patches.py +++ b/src/prime_rl/inference/patches.py @@ -1,48 +1,527 @@ -# Monkeypatch PrometheusStatLogger to avoid NotImplementedError for LoRA in DP mode -def monkey_patch_prometheus_stat_logger_for_lora_in_dp_mode(): - from vllm.v1.metrics import loggers as vllm_metrics_loggers +import torch +from vllm.triton_utils import tl, triton - _original_prometheus_stat_logger_init = vllm_metrics_loggers.PrometheusStatLogger.__init__ - def _patched_prometheus_stat_logger_init(self, vllm_config, engine_indexes=None): - """Patched init that temporarily disables lora_config to skip the DP mode check.""" - original_lora_config = vllm_config.lora_config - vllm_config.lora_config = None - try: - _original_prometheus_stat_logger_init(self, vllm_config, engine_indexes) - finally: - vllm_config.lora_config = original_lora_config - # Re-initialize LoRA metrics if needed (after the DP check is bypassed) - if original_lora_config is not None: - self.labelname_max_lora = "max_lora" - self.labelname_waiting_lora_adapters = "waiting_lora_adapters" - self.labelname_running_lora_adapters = "running_lora_adapters" - self.max_lora = original_lora_config.max_loras - self.gauge_lora_info = vllm_metrics_loggers.PrometheusStatLogger._gauge_cls( - name="vllm:lora_requests_info", - documentation="Running stats on lora requests.", - multiprocess_mode="sum", - labelnames=[ - self.labelname_max_lora, - self.labelname_waiting_lora_adapters, - self.labelname_running_lora_adapters, - ], +def transformers_v5_compat(): + """vLLM general plugin: patch transformers v5 config attrs that vLLM still expects. + + Registered as a ``vllm.general_plugins`` entry-point so it runs automatically + in every vLLM process, including spawned workers. + """ + from transformers import Qwen3VLMoeTextConfig + + if not hasattr(Qwen3VLMoeTextConfig, "tie_word_embeddings"): + Qwen3VLMoeTextConfig.tie_word_embeddings = False + + _patch_qwen35_lora() + _patch_lora_key_prefix() + monkey_patch_deep_gemm_ep_scatter() + monkey_patch_deep_gemm_silu_mul_quant_int64() + monkey_patch_dp_engine_core_pause_resume_deadlock() + + +@triton.jit +def _apply_expert_map_triton(expert_id, expert_map): + if expert_id != -1: + expert_id = tl.load(expert_map + expert_id).to(expert_id.dtype) + return expert_id + + +@triton.jit +def _fwd_kernel_ep_scatter_2_int64( + total_token_num, + expert_start_loc, + recv_x, + recv_x_stride0, + recv_x_stride1, + recv_x_scale, + recv_x_scale_stride0, + recv_x_scale_stride1, + recv_topk, + recv_topk_stride0, + recv_topk_stride1, + output_tensor, + output_tensor_stride0, + output_tensor_stride1, + output_tensor_scale, + output_tensor_scale_stride0, + output_tensor_scale_stride1, + output_index, + output_index_stride0, + output_index_stride1, + topk_num: tl.constexpr, + expert_map, + HAS_EXPERT_MAP: tl.constexpr, + HIDDEN_SIZE: tl.constexpr, + HIDDEN_SIZE_PAD: tl.constexpr, + SCALE_HIDDEN_SIZE: tl.constexpr, + SCALE_HIDDEN_SIZE_PAD: tl.constexpr, +): + start_token_id = tl.program_id(0) + grid_num = tl.num_programs(0) + + offset_in = tl.arange(0, HIDDEN_SIZE_PAD) + mask = offset_in < HIDDEN_SIZE + + offset_in_s = tl.arange(0, SCALE_HIDDEN_SIZE_PAD) + mask_s = offset_in_s < SCALE_HIDDEN_SIZE + + output_tensor_stride0 = output_tensor_stride0.to(tl.int64) + + for token_id in range(start_token_id, total_token_num, grid_num): + to_copy = tl.load(recv_x + token_id * recv_x_stride0 + offset_in, mask=mask) + to_copy_s = tl.load( + recv_x_scale + token_id * recv_x_scale_stride0 + offset_in_s, + mask=mask_s, + ) + + for topk_index in tl.range(0, topk_num, 1, num_stages=4): + expert_id = tl.load(recv_topk + token_id * recv_topk_stride0 + topk_index) + + if HAS_EXPERT_MAP: + expert_id = _apply_expert_map_triton(expert_id, expert_map) + + if expert_id >= 0: + dest_token_index = tl.atomic_add(expert_start_loc + expert_id, 1) + dest_token_index_i64 = dest_token_index.to(tl.int64) + + tl.store( + output_index + token_id * output_index_stride0 + topk_index, + dest_token_index, + ) + + output_tensor_ptr = output_tensor + dest_token_index_i64 * output_tensor_stride0 + output_tensor_scale_ptr = output_tensor_scale + dest_token_index * output_tensor_scale_stride0 + tl.store(output_tensor_ptr + offset_in, to_copy, mask=mask) + tl.store(output_tensor_scale_ptr + offset_in_s, to_copy_s, mask=mask_s) + + +def _triton_ep_scatter_int64( + recv_x: torch.Tensor, + recv_x_scale: torch.Tensor, + recv_topk: torch.Tensor, + num_recv_tokens_per_expert: torch.Tensor, + expert_map: torch.Tensor | None, + expert_start_loc: torch.Tensor, + output_tensor: torch.Tensor, + output_tensor_scale: torch.Tensor, + m_indices: torch.Tensor, + output_index: torch.Tensor, +) -> None: + from vllm.model_executor.layers.fused_moe import deep_gemm_utils + + block_e = 128 + num_warps = 8 + num_experts = num_recv_tokens_per_expert.shape[0] + hidden_size = recv_x.shape[1] + + assert m_indices.shape[0] % block_e == 0 + assert expert_start_loc.shape[0] == num_experts + + deep_gemm_utils._fwd_kernel_ep_scatter_1[(num_experts,)]( + num_recv_tokens_per_expert, + expert_start_loc, + m_indices, + num_experts=num_experts, + num_warps=num_warps, + BLOCK_E=block_e, + BLOCK_EXPERT_NUM=triton.next_power_of_2(num_experts), + ) + + grid = min(recv_topk.shape[0], 1024 * 8) + _fwd_kernel_ep_scatter_2_int64[(grid,)]( + recv_topk.shape[0], + expert_start_loc, + recv_x, + recv_x.stride(0), + recv_x.stride(1), + recv_x_scale, + recv_x_scale.stride(0), + recv_x_scale.stride(1), + recv_topk, + recv_topk.stride(0), + recv_topk.stride(1), + output_tensor, + output_tensor.stride(0), + output_tensor.stride(1), + output_tensor_scale, + output_tensor_scale.stride(0), + output_tensor_scale.stride(1), + output_index, + output_index.stride(0), + output_index.stride(1), + topk_num=recv_topk.shape[1], + expert_map=expert_map, + HAS_EXPERT_MAP=expert_map is not None, + num_warps=num_warps, + HIDDEN_SIZE=hidden_size, + HIDDEN_SIZE_PAD=triton.next_power_of_2(hidden_size), + SCALE_HIDDEN_SIZE=recv_x_scale.shape[1], + SCALE_HIDDEN_SIZE_PAD=triton.next_power_of_2(recv_x_scale.shape[1]), + ) + + +def monkey_patch_deep_gemm_ep_scatter(): + # Temporary local carry of the upstream fix while it is under review: + # issue: https://github.com/vllm-project/vllm/issues/39211 + # PR: https://github.com/vllm-project/vllm/pull/39213 + from vllm.logger import init_logger + from vllm.model_executor.layers.fused_moe import deep_gemm_utils + + logger = init_logger(__name__) + + deep_gemm_utils.ep_scatter = torch.no_grad()(_triton_ep_scatter_int64) + logger.warning("Enabled int64-addressing Triton patch for vLLM DeepGEMM ep_scatter.") + + +@triton.jit +def _silu_mul_per_token_group_quant_fp8_colmajor_int64_kernel( + y_ptr, + y_q_ptr, + y_s_ptr, + M: tl.int64, + N: tl.int64, + y_s_col_stride: tl.int64, + eps, + fp8_min: tl.constexpr, + fp8_max: tl.constexpr, + use_ue8m0: tl.constexpr, + GROUP_SIZE: tl.constexpr, + BLOCK_M: tl.constexpr, + BLOCK_N: tl.constexpr, +): + pid_m = tl.program_id(0) + pid_n = tl.program_id(1) + N_2 = N // 2 + + m_offset = (pid_m * BLOCK_M).to(tl.int64) + n_offset = (pid_n * BLOCK_N).to(tl.int64) + if m_offset >= M: + return + + offs_n = tl.arange(0, BLOCK_N).to(tl.int64) + offs_m = tl.arange(0, BLOCK_M).to(tl.int64) + + base_y_ptr = y_ptr + m_offset * N + n_offset + act_in_ptrs = base_y_ptr + offs_m[:, None] * N + offs_n[None, :] + + act_in = tl.load(act_in_ptrs) + mul_in = tl.load(act_in_ptrs + N_2) + + act_in = act_in.to(tl.float32) + one_f32 = tl.cast(1, tl.float32) + silu_out = (act_in / (one_f32 + tl.exp(-act_in))).to(y_ptr.dtype.element_ty) + y = (silu_out * mul_in).to(tl.float32) + + absmax = tl.maximum(tl.max(tl.abs(y), axis=1), eps) + scale_raw = absmax * (1.0 / fp8_max) + y_s = tl.math.exp2(tl.ceil(tl.log2(scale_raw))) if use_ue8m0 else scale_raw + y_s = tl.reshape(y_s, (BLOCK_M, 1)) + y_q = tl.clamp(y / y_s, fp8_min, fp8_max).to(y_q_ptr.dtype.element_ty) + + base_y_q_ptr = y_q_ptr + m_offset * N_2 + n_offset + y_q_ptrs = base_y_q_ptr + offs_m[:, None] * N_2 + offs_n[None, :] + tl.store(y_q_ptrs, y_q) + + group_id = n_offset // GROUP_SIZE + base_y_s_ptr = y_s_ptr + group_id * y_s_col_stride + m_offset + y_s_ptrs = base_y_s_ptr + offs_m + y_s = tl.reshape(y_s, (BLOCK_M,)) + tl.store(y_s_ptrs, y_s) + + +def _silu_mul_per_token_group_quant_fp8_colmajor_int64( + input: torch.Tensor, + output: torch.Tensor | None = None, + use_ue8m0: bool | None = None, + eps: float = 1e-10, +): + from vllm.platforms import current_platform + from vllm.utils.deep_gemm import is_deep_gemm_e8m0_used + + group_size = 128 + assert input.ndim == 2 + if output is not None: + assert output.ndim == 2 + assert input.size(0) % group_size == 0 + assert input.size(1) % (group_size * 2) == 0 + + if use_ue8m0 is None: + use_ue8m0 = is_deep_gemm_e8m0_used() + + M, N = input.size() + N_2 = N // 2 + + fp8_dtype = current_platform.fp8_dtype() + if output is None: + output = torch.empty((M, N_2), dtype=fp8_dtype, device=input.device) + + output_scales = torch.empty(((N_2 // group_size), M), dtype=torch.float32, device=input.device).transpose(0, 1) + + block_m = 8 + block_n = group_size + assert M % block_m == 0 + assert N_2 % block_n == 0 + + finfo = torch.finfo(fp8_dtype) + fp8_min = -224.0 if current_platform.is_fp8_fnuz() else finfo.min + fp8_max = 224.0 if current_platform.is_fp8_fnuz() else finfo.max + + grid = (M // block_m, N_2 // block_n) + _silu_mul_per_token_group_quant_fp8_colmajor_int64_kernel[grid]( + input, + output, + output_scales, + M, + N, + output_scales.stride(-1), + eps, + fp8_min, + fp8_max, + use_ue8m0, + group_size, + block_m, + block_n, + ) + + return output, output_scales + + +def monkey_patch_deep_gemm_silu_mul_quant_int64(): + # Temporary local carry for large DeepGEMM profile shapes whose row offsets + # exceed signed int32 address arithmetic in vLLM's Triton kernel. + import sys + + from vllm.logger import init_logger + from vllm.model_executor.layers.quantization.utils import fp8_utils + + logger = init_logger(__name__) + + fp8_utils.silu_mul_per_token_group_quant_fp8_colmajor = _silu_mul_per_token_group_quant_fp8_colmajor_int64 + + deep_gemm_moe_module = sys.modules.get("vllm.model_executor.layers.fused_moe.experts.deep_gemm_moe") + if deep_gemm_moe_module is not None: + deep_gemm_moe_module.silu_mul_per_token_group_quant_fp8_colmajor = ( + _silu_mul_per_token_group_quant_fp8_colmajor_int64 + ) + + logger.warning("Enabled int64-addressing Triton patch for vLLM DeepGEMM SiLU/mul FP8 quant.") + + +def _patch_qwen35_lora(): + """Fix Qwen3.5 LoRA: align packed_modules_mapping with output_sizes. + + Qwen3.5's GDN layers use create_qkvz_proj with 4 output_sizes (q, k, v, z) + but packed_modules_mapping only lists 2 entries, causing an IndexError + during LoRA initialization. + + Also generalizes MergedColumnParallelLinearWithLoRA.can_replace_layer + to accept any number of packed modules (not just 2), and generalizes + MergedColumnParallelLinearWithShardedLoRA.slice_lora_a to handle N + subloras instead of the hardcoded 2 (needed for fully_sharded_loras=True). + + Upstream: https://github.com/vllm-project/vllm/issues/36372 + """ + from vllm.lora.layers.column_parallel_linear import ( + MergedColumnParallelLinearWithLoRA, + MergedColumnParallelLinearWithShardedLoRA, + ) + from vllm.model_executor.models.qwen3_5 import ( + Qwen3_5ForCausalLMBase, + Qwen3_5ForConditionalGeneration, + Qwen3_5MoeForConditionalGeneration, + ) + + qkvz_fix = ["in_proj_q", "in_proj_k", "in_proj_v", "in_proj_z"] + + Qwen3_5ForCausalLMBase.packed_modules_mapping["in_proj_qkvz"] = qkvz_fix + Qwen3_5ForConditionalGeneration.packed_modules_mapping["in_proj_qkvz"] = qkvz_fix + + Qwen3_5MoeForConditionalGeneration.is_3d_moe_weight = False + + from vllm.lora.layers.utils import _not_fully_sharded_can_replace + + @classmethod + @_not_fully_sharded_can_replace + def can_replace_layer(cls, source_layer, lora_config, packed_modules_list, model_config=None): + from vllm.model_executor.layers.linear import MergedColumnParallelLinear + + return type(source_layer) is MergedColumnParallelLinear and len(packed_modules_list) == len( + source_layer.output_sizes + ) + + MergedColumnParallelLinearWithLoRA.can_replace_layer = can_replace_layer + + def slice_lora_a(self, lora_a): + output_shard_size = self.lora_a_stacked[0].shape[2] + output_start_idx = self.tp_rank * output_shard_size + return [ + a[output_start_idx : output_start_idx + output_shard_size, :] if a is not None else None for a in lora_a + ] + + MergedColumnParallelLinearWithShardedLoRA.slice_lora_a = slice_lora_a + + +def _patch_lora_key_prefix(): + """Patch vLLM's LoRA loading to handle keys without base_model.model. prefix. + + This is a copy of the upstream patch: https://github.com/vllm-project/vllm/pull/38522 + We can remove this patch once that PR makes it into a release. + """ + from vllm.lora.lora_model import ( + LoRAModel, + PEFTHelper, + TensorizerConfig, + WeightsMapper, + get_lora_id, + is_base_embedding_weights, + os, + parse_fine_tuned_lora_name, + safetensors, + ) + + def _patched_from_local_checkpoint( + cls, + lora_dir: str, + expected_lora_modules: set[str], + peft_helper: PEFTHelper, + *, + lora_model_id: int | None = None, + device: str = "cuda", + dtype: torch.dtype | None = None, + model_vocab_size: int | None = None, + weights_mapper: WeightsMapper | None = None, + tensorizer_config_dict: dict | None = None, + skip_prefixes: list[str] | None = None, + ) -> "LoRAModel": + """Create a LoRAModel from a local checkpoint. + + Args: + lora_dir: The local path that has lora data. + expected_lora_modules: Name of modules that are expected to be + replaced by lora. + peft_helper: Loaded lora configuration information. + lora_model_id: LoRA model id. If not given, automatically set by + a global counter. + device: Device where the lora model is loaded. + dtype: dtype of the lora model weights. + skip_prefixes: List of module name prefixes to skip during loading. + Models can define this to skip modules not used in inference + (e.g., MTP layers). Format: ["mtp."] + + Returns: + Loaded LoRA Model. + """ + lora_tensor_path = os.path.join(lora_dir, "adapter_model.safetensors") + lora_bin_file_path = os.path.join(lora_dir, "adapter_model.bin") + lora_pt_file_path = os.path.join(lora_dir, "adapter_model.pt") + + tensors: dict[str, torch.Tensor] = {} + unexpected_modules: list[list[str] | str] = [] + + def check_unexpected_modules(modules: dict): + for lora_module in modules.keys(): # noqa + if is_base_embedding_weights(lora_module): + continue + # Handle PEFT file format where experts.base_layer is the + # gate_up_proj and experts is the down_proj + if "base_layer" in lora_module: + continue + # Skip modules based on model-defined prefixes + if skip_prefixes and cls._should_skip_module(lora_module, skip_prefixes): + continue + module_name, _ = parse_fine_tuned_lora_name(lora_module, weights_mapper) + # Case for expert lora weights. + # For standard MoE models the name ends in "...experts", + # so expert_idx+1 yields "experts" which is in + # expected_lora_modules. + # For Qwen 3.5 MoE (and similar models) the expert index + # is embedded: "...experts.N.down_proj". Taking everything + # after ".experts" gives "experts.N.down_proj" which is + # never in the expected set even though "down_proj" is. + # Qwen3-30B-A3B goes the other way: the expected set + # contains the fully-qualified per-expert name + # ("experts.N.down_proj") but not the bare suffix. + # Accept either form. + if ".experts" in module_name: + expert_suffix = module_name.split(".")[-1] + experts_qualified = "experts" + module_name.split(".experts", 1)[-1] + if expert_suffix not in expected_lora_modules and experts_qualified not in expected_lora_modules: + unexpected_modules.append(module_name) + + elif module_name.rsplit(".", 1)[-1] not in expected_lora_modules: + unexpected_modules.append(module_name) + + if unexpected_modules: + raise ValueError( + f"While loading {lora_dir}, expected" + f" target modules in {expected_lora_modules}" + f" but received {unexpected_modules}." + f" Please verify that the loaded LoRA module is correct" + ) + + if tensorizer_config_dict: + from tensorizer import TensorDeserializer + + tensorizer_config = TensorizerConfig(**tensorizer_config_dict) + lora_tensor_path = os.path.join(tensorizer_config.tensorizer_dir, "adapter_model.tensors") + tensorizer_args = tensorizer_config._construct_tensorizer_args() + tensors = TensorDeserializer( + lora_tensor_path, + dtype=tensorizer_config.dtype, + **tensorizer_args.deserialization_kwargs, ) + check_unexpected_modules(tensors) + + elif os.path.isfile(lora_tensor_path): + # Find unexpected modules. + # Use safetensor key as a source of truth to find expected modules. + # in peft if you have target_modules A, B, C and C does not exist + # in the model it won’t error and model will be trained with A, B + # loraified. C won’t exist in the safetensor but it will exist in + # the target_modules of the adapter_config.json. + unexpected_modules = [] + with safetensors.safe_open(lora_tensor_path, framework="pt") as f: # type: ignore + # Load tensors if there are only expected modules. + check_unexpected_modules(f) + for module in f.keys(): # noqa + tensors[module] = f.get_tensor(module) + elif os.path.isfile(lora_bin_file_path) or os.path.isfile(lora_pt_file_path): + lora_file_path = lora_bin_file_path if os.path.isfile(lora_bin_file_path) else lora_pt_file_path + tensors = torch.load(lora_file_path, map_location=device, weights_only=True) + check_unexpected_modules(tensors) + else: + raise ValueError(f"{lora_dir} doesn't contain tensors") + + return cls.from_lora_tensors( + lora_model_id=get_lora_id() if lora_model_id is None else lora_model_id, + tensors=tensors, + peft_helper=peft_helper, + device=device, + dtype=dtype, + model_vocab_size=model_vocab_size, + weights_mapper=weights_mapper, + skip_prefixes=skip_prefixes, + ) - vllm_metrics_loggers.PrometheusStatLogger.__init__ = _patched_prometheus_stat_logger_init + LoRAModel.from_local_checkpoint = classmethod(_patched_from_local_checkpoint) # Monkeypatch LoadLoRAAdapter to allow loading the same adapter multiple times +# TODO: may be removable if we pass load_inplace=True (supported since vLLM 0.18, PR #31326) def monkey_patch_load_lora_adapter(): - from vllm.entrypoints.openai.serving_models import ( - ErrorResponse, - HTTPStatus, - LoadLoRAAdapterRequest, - LoRARequest, + from http import HTTPStatus + + from vllm.entrypoints.openai.engine.protocol import ErrorResponse + from vllm.entrypoints.openai.models.serving import ( OpenAIServingModels, create_error_response, - logger, ) + from vllm.entrypoints.serve.lora.protocol import LoadLoRAAdapterRequest + from vllm.logger import init_logger + from vllm.lora.request import LoRARequest + + logger = init_logger(__name__) async def _patched_load_lora_adapter( self: OpenAIServingModels, request: LoadLoRAAdapterRequest, base_model_name: str | None = None @@ -84,6 +563,7 @@ async def _patched_load_lora_adapter( # Monkeypatch LRUCacheWorkerLoRAManager to allow loading adapter inplace without doing it every request +# TODO: may be removable if we pass load_inplace=True (supported since vLLM 0.18, PR #31326) def monkey_patch_LRUCacheWorkerLoRAManager(): from vllm.lora.worker_manager import LoRARequest, LRUCacheLoRAModelManager, LRUCacheWorkerLoRAManager @@ -137,3 +617,395 @@ def _patched_add_adapter( LRUCacheWorkerLoRAManager._apply_adapters = _patched__apply_adapters LRUCacheWorkerLoRAManager.add_adapter = _patched_add_adapter + + +# Monkeypatch WorkerLoRAManager._load_adapter to skip the per-module regex +# warning loop. On wide MoE models (Qwen3.5-35B-A3B) it spends minutes +# recompiling regex patterns inside is_supported_lora_module — purely to emit +# logger.warning_once about modules that will be ignored. Adapter validity is +# already enforced by from_local_checkpoint, so dropping the warnings is safe. +def monkey_patch_skip_lora_module_warnings(): + from vllm.exceptions import LoRAAdapterNotFoundError + from vllm.lora.lora_model import LoRAModel + from vllm.lora.peft_helper import PEFTHelper + from vllm.lora.request import LoRARequest + from vllm.lora.utils import get_adapter_absolute_path + from vllm.lora.worker_manager import WorkerLoRAManager + + def _patched_load_adapter(self: WorkerLoRAManager, lora_request: LoRARequest) -> LoRAModel: + try: + supported_lora_modules = self._adapter_manager.supported_lora_modules + packed_modules_mapping = self._adapter_manager.packed_modules_mapping + expected_lora_lst: list[str] = [] + for module in supported_lora_modules: + if module in packed_modules_mapping: + expected_lora_lst.extend(packed_modules_mapping[module]) + else: + expected_lora_lst.append(module) + if module == "experts": + expected_lora_lst.append(module) + expected_lora_modules = set(expected_lora_lst) + lora_path = get_adapter_absolute_path(lora_request.lora_path) + + peft_helper = PEFTHelper.from_local_dir( + lora_path, + self.max_position_embeddings, + lora_request.tensorizer_config_dict, + ) + peft_helper.validate_legal(self.lora_config) + + model = self._adapter_manager.model + hf_to_vllm_mapper = getattr(model, "hf_to_vllm_mapper", None) + lora_skip_prefixes = getattr(model, "lora_skip_prefixes", None) + + lora = self._lora_model_cls.from_local_checkpoint( + lora_path, + expected_lora_modules, + peft_helper=peft_helper, + lora_model_id=lora_request.lora_int_id, + device="cpu", + dtype=self.lora_config.lora_dtype, + model_vocab_size=self.vocab_size, + tensorizer_config_dict=lora_request.tensorizer_config_dict, + weights_mapper=hf_to_vllm_mapper, + skip_prefixes=lora_skip_prefixes, + ) + except FileNotFoundError as e: + raise LoRAAdapterNotFoundError(lora_request.lora_name, lora_request.lora_path) from e + + return lora + + WorkerLoRAManager._load_adapter = _patched_load_adapter + + +# Monkeypatch TokenizeParams to fix overly conservative validation +def monkey_patch_tokenize_params_validation(): + """ + Patch TokenizeParams validation to only reject requests where the prompt + itself exceeds max_model_len, not where prompt + max_tokens > max_model_len. + + Original behavior: + - Rejects if prompt_len > (max_model_len - max_tokens) + + Patched behavior: + - Only rejects if prompt_len > max_model_len + - Lets the engine naturally cap generation at max_model_len + """ + from vllm.exceptions import VLLMValidationError + from vllm.renderers.params import TokenizeParams + + def _patched_token_len_check(self, tokenizer, tokens): + """Only validate that prompt fits in max_model_len, not prompt+max_tokens""" + if self.max_total_tokens is not None and len(tokens) > self.max_total_tokens: + raise VLLMValidationError( + f"The prompt is {len(tokens)} tokens, which exceeds the " + f"model's maximum context length of {self.max_total_tokens} tokens. " + f"Please reduce the length of the input prompt.", + parameter="input_tokens", + value=len(tokens), + ) + return tokens + + def _patched_text_len_check(self, tokenizer, text): + """Only validate text length against max_model_len, not max_input_tokens""" + if self.max_total_tokens is None or tokenizer is None: + return text + + if self.truncate_prompt_tokens is None: + max_chars = self.max_total_tokens * tokenizer.max_chars_per_token + if len(text) > max_chars: + raise VLLMValidationError( + f"You passed {len(text)} input characters. " + f"However, the model's context length is only " + f"{self.max_total_tokens} tokens " + f"(at most {max_chars} characters). " + f"Please reduce the length of the input prompt.", + parameter="input_text", + value=len(text), + ) + return text + + def _patched_get_encode_kwargs(self): + """Use max_total_tokens (max_model_len) instead of max_input_tokens for HF tokenizer truncation. + + The original uses max_input_tokens (= max_model_len - max_tokens) + 1, which causes HuggingFace's + tokenizer.encode() to left-truncate prompts before _token_len_check even runs. + """ + max_length = self.truncate_prompt_tokens + if max_length is not None and max_length < 0: + max_length = self.max_total_tokens + elif max_length is None and self.max_total_tokens is not None: + max_length = self.max_total_tokens + 1 + + return dict( + truncation=max_length is not None, + max_length=max_length, + add_special_tokens=self.add_special_tokens, + ) + + TokenizeParams._token_len_check = _patched_token_len_check + TokenizeParams._text_len_check = _patched_text_len_check + TokenizeParams.get_encode_kwargs = _patched_get_encode_kwargs + + +def monkey_patch_minimax_m2_for_lora(): + """Patch vLLM's MiniMaxM2 model for LoRA compatibility. + + These patches are only needed when using LoRA with MiniMax M2 but are safe + to apply unconditionally (verified with non-LoRA runs). We apply them at + import time because the worker __init__ runs before the vLLM config is + available, so we can't check if LoRA is enabled. + + Problem 1 — Gate dtype mismatch: + vLLM's MiniMaxM2MoE creates the gate (router) with params_dtype=float32 + and casts inputs to float32. When LoRA is enabled, vLLM wraps ALL + ReplicatedLinear layers (including the gate) with LoRA support. Even + though our adapter has no gate LoRA weights, the LoRA Triton kernel + still runs for all wrapped layers when any adapter is active — and it + asserts inputs are float16/bfloat16. Qwen3 MoE doesn't have this + problem because its gate uses the model dtype. + Fix: recreate the gate in model dtype and remove the float32 cast. + FusedMoE already has router_logits_dtype=float32, so routing precision + is preserved inside the expert dispatch. + + Problem 2 — Adapter key naming mismatch: + PrimeRL saves adapter keys using its internal naming convention + (mlp.experts.{j}.gate_proj/down_proj/up_proj), which matches Qwen3 MoE + but not MiniMax M2. vLLM's MiniMax M2 model expects HF-style keys + (block_sparse_moe.experts.{j}.w1/w2/w3). For full model weights this + is handled by vLLM's load_weights(), but LoRA adapters are loaded + through a separate path (LoRAModel.from_local_checkpoint) that doesn't + have model-specific key translation. + Fix: set hf_to_vllm_mapper on the model class so vLLM remaps adapter + keys during LoRA loading. This attribute is only read by _load_adapter + in the LoRA worker manager — it has no effect without LoRA. + """ + from vllm.model_executor.models.minimax_m2 import MiniMaxM2ForCausalLM, MiniMaxM2MoE + from vllm.model_executor.models.utils import WeightsMapper + + # --- Gate dtype fix (only matters with LoRA, safe without) --- + _original_init = MiniMaxM2MoE.__init__ + + def _patched_init(self, config, quant_config=None, prefix=""): + _original_init(self, config, quant_config, prefix) + from vllm.model_executor.layers.linear import ReplicatedLinear + + self.gate = ReplicatedLinear( + config.hidden_size, + config.num_local_experts, + bias=False, + quant_config=None, + prefix=f"{prefix}.gate", + ) + + def _patched_forward(self, hidden_states): + from vllm.distributed import tensor_model_parallel_all_reduce + + num_tokens, hidden_dim = hidden_states.shape + hidden_states = hidden_states.view(-1, hidden_dim) + router_logits, _ = self.gate(hidden_states) + final_hidden_states = self.experts(hidden_states=hidden_states, router_logits=router_logits) + if self.tp_size > 1: + final_hidden_states = tensor_model_parallel_all_reduce(final_hidden_states) + return final_hidden_states.view(num_tokens, hidden_dim) + + MiniMaxM2MoE.__init__ = _patched_init + MiniMaxM2MoE.forward = _patched_forward + + # --- Adapter key remapping (only read by vLLM's LoRA adapter loader) --- + MiniMaxM2ForCausalLM.hf_to_vllm_mapper = WeightsMapper( + orig_to_new_substr={ + ".mlp.experts.": ".block_sparse_moe.experts.", + ".gate_proj.": ".w1.", + ".down_proj.": ".w2.", + ".up_proj.": ".w3.", + }, + ) + + +def monkey_patch_harmony_stop_token_propagation(): + """Fix: vLLM doesn't merge harmony stop tokens into per-request SamplingParams. + + The harmony mode sets stop_token_ids (including <|call|> and <|return|>) in + default_sampling_params at server init, but ChatCompletionRequest.to_sampling_params() + ignores them, using only self.stop_token_ids (which defaults to []). + + Upstream: https://github.com/vllm-project/vllm/issues/22519 + """ + from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest + + _original_to_sampling_params = ChatCompletionRequest.to_sampling_params + + def _patched_to_sampling_params(self, max_tokens, default_sampling_params): + params = _original_to_sampling_params(self, max_tokens, default_sampling_params) + # Merge harmony stop tokens from default_sampling_params + default_stop_ids = default_sampling_params.get("stop_token_ids", []) + if default_stop_ids: + existing = set(params.stop_token_ids or []) + merged = list(existing | set(default_stop_ids)) + params.stop_token_ids = merged + return params + + ChatCompletionRequest.to_sampling_params = _patched_to_sampling_params + + +def monkey_patch_dp_engine_core_pause_resume_deadlock(): + """Fix DP pause/resume deadlocks around weight updates. + + Bug 1 (job 3756): while paused, START_DP_WAVE can wake idle ranks into the + DP loop. Those ranks then run dummy batches and hit DP collectives while + other ranks are still in NCCL weight transfer. + + Bug 2 (jobs 3769/3771): resume ties the DP running state to local + unfinished requests, but the DP wave state is global. Ranks with no local + work still need to re-enter the loop so they can participate in the same + DP collectives as ranks that are resuming remote-KV or decode work. + + Fix: + - ignore START_DP_WAVE wakeups while paused + - on resume, wake every DP rank and force an immediate global unfinished + sync instead of waiting for the normal 32-step cadence + + This keeps the upstream pause-side fix from + https://github.com/vllm-project/vllm/pull/37024 and extends it with the + resume-side wave-state fix. + """ + from vllm.config import ParallelConfig + from vllm.v1.core.sched.interface import PauseState + from vllm.v1.engine import EngineCoreOutputs, EngineCoreRequestType + from vllm.v1.engine.core import DPEngineCoreProc, EngineCore, EngineCoreProc + from vllm.v1.request import Request + + _base_add_request = EngineCore.add_request + _base_handle_client_request = EngineCoreProc._handle_client_request + _base_resume_scheduler = DPEngineCoreProc.resume_scheduler + + def _patched_add_request(self, request: Request, request_wave: int = 0): + _base_add_request(self, request, request_wave) + if self.has_coordinator and request_wave != self.current_wave: + if request_wave > self.current_wave: + self.current_wave = request_wave + elif not self.engines_running and self.scheduler.pause_state == PauseState.UNPAUSED: + self.engines_running = True + self.output_queue.put_nowait((-1, EngineCoreOutputs(start_wave=self.current_wave))) + + def _patched_handle_client_request(self, request_type, request): + if request_type == EngineCoreRequestType.START_DP_WAVE: + new_wave, exclude_eng_index = request + if exclude_eng_index != self.engine_index and new_wave >= self.current_wave: + self.current_wave = new_wave + if not self.engines_running and self.scheduler.pause_state == PauseState.UNPAUSED: + self.engines_running = True + else: + _base_handle_client_request(self, request_type, request) + + def _patched_resume_scheduler(self): + was_paused = self.scheduler.pause_state != PauseState.UNPAUSED + _base_resume_scheduler(self) + if was_paused: + self.engines_running = True + self._force_dp_running_state_sync = True + + def _patched_has_global_unfinished_reqs(self, local_unfinished: bool) -> bool: + self.step_counter += 1 + if getattr(self, "_force_dp_running_state_sync", False): + self._force_dp_running_state_sync = False + return ParallelConfig.has_unfinished_dp(self.dp_group, local_unfinished) + if self.step_counter % 32 != 0: + return True + return ParallelConfig.has_unfinished_dp(self.dp_group, local_unfinished) + + DPEngineCoreProc.add_request = _patched_add_request + DPEngineCoreProc._handle_client_request = _patched_handle_client_request + DPEngineCoreProc.resume_scheduler = _patched_resume_scheduler + DPEngineCoreProc._has_global_unfinished_reqs = _patched_has_global_unfinished_reqs + + +def monkey_patch_no_moe_lora(): + """This disables LoRA for MoE layers and makes them pick better kernels. + + Otherwise, the oracle will always try to pick TritonExperts. + For blackwells, we want TRTLLMFlashInfer. + """ + from vllm.model_executor.layers.fused_moe.config import FusedMoEConfig, logger + + def _patched__post_init__(self: FusedMoEConfig): + if self.dp_size > 1: + logger.debug_once("Using FusedMoEConfig::max_num_tokens=%d", self.max_num_tokens) + + assert self.max_num_tokens > 0 + + if self.router_logits_dtype is None: + self.router_logits_dtype = self.in_dtype + + if self.hidden_dim_unpadded is None: + self.hidden_dim_unpadded = self.hidden_dim + if self.intermediate_size_per_partition_unpadded is None: + self.intermediate_size_per_partition_unpadded = self.intermediate_size_per_partition + + # Disable LoRA for MoE layers + self.is_lora_enabled = False + + FusedMoEConfig.__post_init__ = _patched__post_init__ + + +def monkey_patch_fp32_lm_head(): + """Run the lm_head projection in fp32, via a native bf16xbf16 -> fp32 GEMM. + + Uses ``torch.mm(..., out_dtype=torch.float32)`` (PyTorch >= 2.10) so the + matmul accumulates and emits fp32 directly without zero-padding the bf16 + operands or maintaining a separate fp32 weight copy. This avoids the + epilogue truncation to bf16 that `F.linear(bf16, bf16)` does, which is + where lm_head precision actually leaks before the sampler's softmax. + + Activated by setting ``additional_config["fp32_lm_head"] = True`` on the + vLLM namespace; the launcher does this when ``inference.enable_fp32_lm_head`` + is set. The flag is captured once on ``LogitsProcessor.__init__`` (where + vLLM guarantees a ``set_current_vllm_config()`` context) and stored on the + instance — reading it from ``_get_logits`` during serving doesn't work + because vLLM doesn't keep the context set during forwards. + + Tracks vllm-project/vllm#24567 (which uses the operand-upcast approach). + Per @Jackmin801 on PR #2438, native ``out_dtype=fp32`` mm is more efficient + and just as correct. + """ + import torch + from vllm.config import get_current_vllm_config + from vllm.logger import init_logger + from vllm.model_executor.layers.logits_processor import LogitsProcessor + + logger = init_logger(__name__) + + _original_init = LogitsProcessor.__init__ + _original_get_logits = LogitsProcessor._get_logits + + def _patched_init(self, *args, **kwargs): + _original_init(self, *args, **kwargs) + vllm_config = get_current_vllm_config() + additional_config = vllm_config.additional_config or {} + self._fp32_lm_head_enabled = additional_config.get("fp32_lm_head", False) + if self._fp32_lm_head_enabled: + logger.warning("fp32 lm_head ENABLED for this LogitsProcessor instance.") + + def _patched_get_logits(self, hidden_states, lm_head, embedding_bias): + if not getattr(self, "_fp32_lm_head_enabled", False): + return _original_get_logits(self, hidden_states, lm_head, embedding_bias) + + # Native bf16xbf16 -> fp32 GEMM. torch.mm requires 2D inputs; vLLM v1's + # generative path passes 2D [num_tokens, hidden_size] hidden_states, but + # flatten defensively in case some future caller passes 3D. + flat = hidden_states.reshape(-1, hidden_states.shape[-1]) + logits = torch.mm(flat, lm_head.weight.t(), out_dtype=torch.float32) + if embedding_bias is not None: + logits = logits + embedding_bias.to(torch.float32) + if hidden_states.dim() > 2: + logits = logits.reshape(*hidden_states.shape[:-1], -1) + + logits = self._gather_logits(logits) + if logits is not None: + logits = logits[..., : self.org_vocab_size] + return logits + + LogitsProcessor.__init__ = _patched_init + LogitsProcessor._get_logits = _patched_get_logits + logger.info("Installed fp32 lm_head patch (native out_dtype=fp32 mm).") diff --git a/src/prime_rl/inference/server.py b/src/prime_rl/inference/server.py index 6ddff80764b..322f0145eb0 100644 --- a/src/prime_rl/inference/server.py +++ b/src/prime_rl/inference/server.py @@ -1,14 +1,27 @@ -from prime_rl.inference.config import InferenceConfig -from prime_rl.utils.pydantic_config import parse_argv +import os + +from prime_rl.configs.inference import InferenceConfig +from prime_rl.utils.config import cli + + +def setup_vllm_env(config: InferenceConfig): + """Set vLLM environment variables based on config. Must be called before importing vLLM.""" + + # spawn is more robust in vLLM nightlies and Qwen3-VL (fork can deadlock with multithreaded processes) + os.environ.setdefault("VLLM_WORKER_MULTIPROC_METHOD", "spawn") + + if config.enable_lora: + os.environ["VLLM_ALLOW_RUNTIME_LORA_UPDATING"] = "True" def main(): - config = parse_argv(InferenceConfig, allow_extras=True) + config = cli(InferenceConfig) + setup_vllm_env(config) # We import here to be able to set environment variables before importing vLLM from prime_rl.inference.vllm.server import server # pyright: ignore - server(config, vllm_args=config.get_unknown_args()) + server(config, vllm_extra=config.vllm_extra) if __name__ == "__main__": diff --git a/src/prime_rl/inference/vllm/server.py b/src/prime_rl/inference/vllm/server.py index e05237c4000..53ae22c1045 100644 --- a/src/prime_rl/inference/vllm/server.py +++ b/src/prime_rl/inference/vllm/server.py @@ -1,56 +1,176 @@ +import asyncio from argparse import Namespace from http import HTTPStatus +from typing import Any +import uvloop +from fastapi import APIRouter, Depends, Request from fastapi.responses import JSONResponse, StreamingResponse -from vllm.entrypoints.chat_utils import load_chat_template -from vllm.entrypoints.cli.serve import run_api_server_worker_proc -from vllm.entrypoints.logger import RequestLogger +from starlette.datastructures import State +from vllm.engine.protocol import EngineClient +from vllm.entrypoints.openai.api_server import init_app_state +from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionResponse +from vllm.entrypoints.openai.cli_args import make_arg_parser, validate_parsed_serve_args +from vllm.entrypoints.openai.engine.protocol import ErrorResponse +from vllm.entrypoints.openai.engine.serving import OpenAIServing +from vllm.entrypoints.openai.models.serving import OpenAIServingModels from vllm.entrypoints.openai.utils import validate_json_request -from vllm.entrypoints.openai.protocol import ChatCompletionResponse, ErrorResponse +from vllm.entrypoints.serve.lora.protocol import LoadLoRAAdapterRequest from vllm.entrypoints.utils import load_aware_call, with_cancellation +from vllm.logger import init_logger +from vllm.utils.argparse_utils import FlexibleArgumentParser + +from prime_rl.configs.inference import InferenceConfig +from prime_rl.utils.logger import get_logger + +MODEL_TOOL_CALL_PARSER: dict[str, str] = { + # GLM-4.5 + "zai-org/GLM-4.5": "glm45", + "zai-org/GLM-4.5-FP8": "glm45", + "zai-org/GLM-4.5-Base": "glm45", + "zai-org/GLM-4.5-Air": "glm45", + "zai-org/GLM-4.5-Air-FP8": "glm45", + "zai-org/GLM-4.5-Air-Base": "glm45", + "zai-org/GLM-4.5V": "glm45", + "zai-org/GLM-4.5V-FP8": "glm45", + # GLM-4.7 + "zai-org/GLM-4.7": "glm47", + "zai-org/GLM-4.7-FP8": "glm47", + "zai-org/GLM-4.7-Flash": "glm47", + # GLM-5 + "zai-org/GLM-5": "glm47", + "zai-org/GLM-5-FP8": "glm47", + # GLM-5.1 + "zai-org/GLM-5.1": "glm47", + "zai-org/GLM-5.1-FP8": "glm47", + # MiniMax M2 + "MiniMaxAI/MiniMax-M2": "minimax_m2", + "MiniMaxAI/MiniMax-M2.1": "minimax_m2", + "MiniMaxAI/MiniMax-M2.5": "minimax_m2", + # INTELLECT-3 + "PrimeIntellect/INTELLECT-3": "hermes", + "PrimeIntellect/INTELLECT-3-FP8": "hermes", + "PrimeIntellect/INTELLECT-3.1": "hermes", + # Qwen3 dense + "Qwen/Qwen3-0.6B": "hermes", + "Qwen/Qwen3-0.6B-Base": "hermes", + "Qwen/Qwen3-0.6B-FP8": "hermes", + "Qwen/Qwen3-1.7B": "hermes", + "Qwen/Qwen3-1.7B-Base": "hermes", + "Qwen/Qwen3-1.7B-FP8": "hermes", + "Qwen/Qwen3-4B": "hermes", + "Qwen/Qwen3-4B-Base": "hermes", + "Qwen/Qwen3-4B-FP8": "hermes", + "Qwen/Qwen3-8B": "hermes", + "Qwen/Qwen3-8B-Base": "hermes", + "Qwen/Qwen3-8B-FP8": "hermes", + "Qwen/Qwen3-14B": "hermes", + "Qwen/Qwen3-14B-Base": "hermes", + "Qwen/Qwen3-14B-FP8": "hermes", + "Qwen/Qwen3-32B": "hermes", + "Qwen/Qwen3-32B-FP8": "hermes", + # Qwen3 MoE + "Qwen/Qwen3-30B-A3B": "hermes", + "Qwen/Qwen3-30B-A3B-Base": "hermes", + "Qwen/Qwen3-30B-A3B-FP8": "hermes", + "Qwen/Qwen3-235B-A22B": "hermes", + "Qwen/Qwen3-235B-A22B-FP8": "hermes", + # Qwen3 2507 + "Qwen/Qwen3-4B-Instruct-2507": "hermes", + "Qwen/Qwen3-4B-Thinking-2507": "hermes", + "Qwen/Qwen3-4B-Instruct-2507-FP8": "hermes", + "Qwen/Qwen3-4B-Thinking-2507-FP8": "hermes", + "Qwen/Qwen3-30B-A3B-Instruct-2507": "hermes", + "Qwen/Qwen3-30B-A3B-Thinking-2507": "hermes", + "Qwen/Qwen3-30B-A3B-Instruct-2507-FP8": "hermes", + "Qwen/Qwen3-30B-A3B-Thinking-2507-FP8": "hermes", + "Qwen/Qwen3-235B-A22B-Instruct-2507": "hermes", + "Qwen/Qwen3-235B-A22B-Thinking-2507": "hermes", + "Qwen/Qwen3-235B-A22B-Instruct-2507-FP8": "hermes", + "Qwen/Qwen3-235B-A22B-Thinking-2507-FP8": "hermes", + # Qwen3-Next + "Qwen/Qwen3-Next-80B-A3B-Instruct": "hermes", + "Qwen/Qwen3-Next-80B-A3B-Thinking": "hermes", + "Qwen/Qwen3-Next-80B-A3B-Instruct-FP8": "hermes", + "Qwen/Qwen3-Next-80B-A3B-Thinking-FP8": "hermes", + # Qwen3-Coder + "Qwen/Qwen3-Coder-480B-A35B-Instruct": "hermes", + "Qwen/Qwen3-Coder-480B-A35B-Instruct-FP8": "hermes", + "Qwen/Qwen3-Coder-30B-A3B-Instruct": "hermes", + "Qwen/Qwen3-Coder-30B-A3B-Instruct-FP8": "hermes", + # Qwen3-Coder-Next + "Qwen/Qwen3-Coder-Next": "hermes", + "Qwen/Qwen3-Coder-Next-Base": "hermes", + "Qwen/Qwen3-Coder-Next-FP8": "hermes", + # Qwen3.5 dense (uses qwen3_coder tool format, not hermes) + "Qwen/Qwen3.5-0.8B": "qwen3_coder", + "Qwen/Qwen3.5-0.8B-Base": "qwen3_coder", + "Qwen/Qwen3.5-2B": "qwen3_coder", + "Qwen/Qwen3.5-2B-Base": "qwen3_coder", + "Qwen/Qwen3.5-4B": "qwen3_coder", + "Qwen/Qwen3.5-4B-Base": "qwen3_coder", + "Qwen/Qwen3.5-9B": "qwen3_coder", + "Qwen/Qwen3.5-9B-Base": "qwen3_coder", + "Qwen/Qwen3.5-27B": "qwen3_coder", + "Qwen/Qwen3.5-27B-FP8": "qwen3_coder", + # Qwen3.5 MoE (uses qwen3_coder tool format, not hermes) + "Qwen/Qwen3.5-35B-A3B": "qwen3_coder", + "Qwen/Qwen3.5-35B-A3B-Base": "qwen3_coder", + "Qwen/Qwen3.5-35B-A3B-FP8": "qwen3_coder", + "Qwen/Qwen3.5-122B-A10B": "qwen3_coder", + "Qwen/Qwen3.5-122B-A10B-FP8": "qwen3_coder", + "Qwen/Qwen3.5-397B-A17B": "qwen3_coder", + "Qwen/Qwen3.5-397B-A17B-FP8": "qwen3_coder", + # NemotronH + "nvidia/NVIDIA-Nemotron-3-Super-120B-A12B-BF16": "qwen3_coder", + "nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16": "qwen3_coder", +} -from fastapi.responses import JSONResponse, StreamingResponse -from vllm.entrypoints.chat_utils import load_chat_template -from vllm.entrypoints.logger import RequestLogger -from vllm.entrypoints.openai.protocol import ChatCompletionRequest, ChatCompletionResponse, ErrorResponse -from vllm.entrypoints.utils import load_aware_call, with_cancellation +def resolve_tool_call_parser(model_name: str, tool_call_parser: str | None) -> str | None: + """Resolve tool_call_parser from model name if set to "auto".""" + if tool_call_parser == "auto": + return MODEL_TOOL_CALL_PARSER.get(model_name) + return tool_call_parser + + +logger = get_logger() from prime_rl.inference.patches import ( - monkey_patch_prometheus_stat_logger_for_lora_in_dp_mode, + monkey_patch_harmony_stop_token_propagation, monkey_patch_load_lora_adapter, + monkey_patch_tokenize_params_validation, ) from prime_rl.inference.vllm.serving_chat_with_tokens import ( ChatCompletionRequestWithTokens, OpenAIServingChatWithTokens, ) -# NOTE: Monkeypatch PrometheusStatLogger to avoid NotImplementedError for LoRA in DP mode -monkey_patch_prometheus_stat_logger_for_lora_in_dp_mode() +# NOTE: Fix harmony stop token propagation for GPT-OSS models +# Upstream issue still open: https://github.com/vllm-project/vllm/issues/22519 +monkey_patch_harmony_stop_token_propagation() # NOTE: Monkeypatch LoadLoRAAdapter to allow loading the same adapter multiple times +# May be removable if we pass load_inplace=True (supported since vLLM 0.18, PR #31326) monkey_patch_load_lora_adapter() +# NOTE: Monkeypatch TokenizeParams to fix overly conservative validation +# Still needed in vLLM 0.20 — upstream rejects prompt_len > max_model_len - max_tokens +monkey_patch_tokenize_params_validation() -# ruff: noqa -import vllm.entrypoints.openai.api_server +logger = init_logger("vllm.entrypoints.openai.api_server") -import uvloop -import vllm.envs as envs -from fastapi import Request -from fastapi import Depends, HTTPException, Request -from starlette.datastructures import State -from vllm.engine.protocol import EngineClient -from vllm.entrypoints.openai.api_server import ( - router, - engine_client, - base, - init_app_state, -) -from vllm.entrypoints.openai.cli_args import make_arg_parser, validate_parsed_serve_args -from vllm.logger import init_logger -from vllm.utils.argparse_utils import FlexibleArgumentParser +# Create our own router for custom endpoints +router = APIRouter() -from prime_rl.inference.config import InferenceConfig -logger = init_logger("vllm.entrypoints.openai.api_server") +def engine_client(request: Request) -> EngineClient: + return request.app.state.engine_client + + +def base(request: Request) -> OpenAIServing: + return request.app.state.openai_serving_tokenization + + +def models(request: Request) -> OpenAIServingModels: + return request.app.state.openai_serving_models WORKER_EXTENSION_CLS = { @@ -63,35 +183,6 @@ def chat_with_tokens(request: Request) -> OpenAIServingChatWithTokens | None: return request.app.state.openai_serving_chat_with_tokens -@router.post("/update_weights") -async def update_weights(request: Request): - data = await request.json() - await engine_client(request).collective_rpc("update_weights", args=(data.get("weight_dir"),)) - return {"status": "ok"} - - -@router.post("/reload_weights") -async def reload_weights(request: Request): - await engine_client(request).collective_rpc("reload_weights") - return {"status": "ok"} - - -@router.post("/init_broadcaster") -async def init_broadcaster(request: Request): - data = await request.json() - host = data.get("host") - port = data.get("port") - timeout = data.get("timeout") - # Support both legacy and new field names - server_rank = data.get("server_rank") - num_inference_server = data.get("num_inference_server") - await engine_client(request).collective_rpc( - "init_broadcaster", - args=(host, port, server_rank, num_inference_server, timeout), - ) - return {"status": "ok"} - - @router.post( "/v1/chat/completions/tokens", dependencies=[Depends(validate_json_request)], @@ -108,10 +199,7 @@ async def _chat_with_tokens(request: ChatCompletionRequestWithTokens, raw_reques handler = chat_with_tokens(raw_request) if handler is None: return base(raw_request).create_error_response(message="The model does not support Chat Completions API") - try: - generator = await handler.create_chat_completion_with_tokens(request, raw_request) - except Exception as e: - raise HTTPException(status_code=HTTPStatus.INTERNAL_SERVER_ERROR.value, detail=str(e)) from e + generator = await handler.create_chat_completion_with_tokens(request, raw_request) if isinstance(generator, ErrorResponse): return JSONResponse(content=generator.model_dump(), status_code=generator.error.code) @@ -121,86 +209,167 @@ async def _chat_with_tokens(request: ChatCompletionRequestWithTokens, raw_reques return StreamingResponse(content=generator, media_type="text/event-stream") -async def custom_init_app_state(engine_client: EngineClient, state: State, args: Namespace): +@router.post("/pause") +async def pause(request: Request): + await engine_client(request).pause_generation(mode="keep", clear_cache=False) + return {"status": "paused"} + + +@router.post("/resume") +async def resume(request: Request): + await engine_client(request).resume_generation() + return {"status": "resumed"} + + +@router.post("/update_weights") +async def update_weights(request: Request): + data = await request.json() + await engine_client(request).collective_rpc("update_weights_from_path", args=(data.get("weight_dir"),)) + return {"status": "ok"} + + +@router.post("/load_lora_adapter") +async def load_lora_adapter(lora_request: LoadLoRAAdapterRequest, raw_request: Request): + """Wrapper around vLLM's /v1/load_lora_adapter.""" + handler = models(raw_request) + response = await handler.load_lora_adapter(lora_request) + if isinstance(response, ErrorResponse): + return JSONResponse(content=response.model_dump(), status_code=response.error.code) + return {"status": "ok"} + + +@router.get("/liveness") +async def liveness(raw_request: Request): + """Check that the engine event loop can service a no-op worker RPC.""" + try: + await asyncio.wait_for( + engine_client(raw_request).collective_rpc("liveness_probe"), + timeout=raw_request.app.state.liveness_timeout_seconds, + ) + except asyncio.TimeoutError: + return JSONResponse({"status": "engine_unresponsive"}, status_code=503) + return {"status": "ok"} + + +@router.post("/init_broadcaster") +async def init_broadcaster(request: Request): + data = await request.json() + host = data.get("host") + port = data.get("port") + timeout = data.get("timeout") + rank_offset = data.get("rank_offset") + inference_world_size = data.get("inference_world_size") + quantize_in_weight_transfer = data.get("quantize_in_weight_transfer", False) + await engine_client(request).collective_rpc( + "init_broadcaster", + args=(host, port, rank_offset, inference_world_size, timeout, quantize_in_weight_transfer), + ) + return {"status": "ok"} + + +async def custom_init_app_state( + engine_client: EngineClient, + state: State, + args: Namespace, + supported_tasks: tuple, +): """ Modifies init_app_state: - 1. Set up the custom OpenAIServingChatWithTokens state. - 2. Monkey-patch to allow updating lora adapters in-place. + 1. Call the original init_app_state to set up standard state, including + vLLM 0.20's ``serving_tokens`` for ``/inference/v1/generate``. + 2. Replace ``serving_chat`` with our ``OpenAIServingChatWithTokens`` wrapper + so the ``/v1/chat/completions/tokens`` (TITO) endpoint can stream + token IDs alongside the rendered chat completion. + 3. Replace ``serving_tokens`` with ``PrimeRlServingTokens`` so DP-rank + routing and ``routed_experts`` export survive the migration off the + legacy ``/v1/generate`` endpoint. """ - # Setup the regular app state first (in-place) - await init_app_state(engine_client, state, args) - - # NOTE: Initialize the custom OpenAIServingChatWithTokens state here - # TODO: Here, we repeat some calls done in init_app_state to be able to - # correctly set up the OpenAIServingChatWithTokens state, which is a bit - # brittle, and could probably be made nicer - if args.enable_log_requests: - request_logger = RequestLogger(max_log_len=args.max_log_len) + await init_app_state(engine_client, state, args, supported_tasks) + + state.reset_prefix_cache_after_update = getattr(args, "reset_prefix_cache_after_update", True) + state.liveness_timeout_seconds = args.liveness_timeout_seconds + + # TITO: server-side chat templating + token IDs. + if "generate" in supported_tasks and state.openai_serving_chat is not None: + original_chat = state.openai_serving_chat + serving_chat = object.__new__(OpenAIServingChatWithTokens) + serving_chat.__dict__.update(original_chat.__dict__) + state.openai_serving_chat = serving_chat + state.openai_serving_chat_with_tokens = serving_chat else: - request_logger = None - - supported_tasks = await engine_client.get_supported_tasks() - resolved_chat_template = load_chat_template(args.chat_template) - - state.openai_serving_chat_with_tokens = ( - OpenAIServingChatWithTokens( - engine_client, - state.openai_serving_models, - args.response_role, - request_logger=request_logger, - chat_template=resolved_chat_template, - chat_template_content_format=args.chat_template_content_format, - trust_request_chat_template=args.trust_request_chat_template, - return_tokens_as_token_ids=args.return_tokens_as_token_ids, - enable_auto_tools=args.enable_auto_tool_choice, - exclude_tools_when_tool_choice_none=args.exclude_tools_when_tool_choice_none, - tool_parser=args.tool_call_parser, - reasoning_parser=args.structured_outputs_config.reasoning_parser, - enable_prompt_tokens_details=args.enable_prompt_tokens_details, - enable_force_include_usage=args.enable_force_include_usage, - enable_log_outputs=args.enable_log_outputs, - log_error_stack=args.log_error_stack, - ) - if "generate" in supported_tasks - else None - ) + state.openai_serving_chat_with_tokens = None + + # Swap in our ServingTokens subclass for /inference/v1/generate so the + # X-data-parallel-rank header and routed_experts response field — both + # used by prime-RL's renderer / router-replay paths — keep working. + if "generate" in supported_tasks and state.serving_tokens is not None: + from prime_rl.inference.vllm.serving_tokens import PrimeRlServingTokens + + upstream = state.serving_tokens + prime_serving = object.__new__(PrimeRlServingTokens) + prime_serving.__dict__.update(upstream.__dict__) + state.serving_tokens = prime_serving + + +import vllm.entrypoints.openai.api_server +import vllm.v1.utils +from vllm.entrypoints.openai.api_server import build_app as _original_build_app +from vllm.v1.utils import run_api_server_worker_proc as _original_run_api_server_worker_proc + + +def custom_build_app(args: Namespace, supported_tasks: tuple, model_config=None): + """ + Wrap build_app to include our custom router. + """ + app = _original_build_app(args, supported_tasks, model_config) + app.include_router(router) + return app def custom_run_api_server_worker_proc(listen_address, sock, args, client_config=None, **uvicorn_kwargs) -> None: """ - Modifies run_api_server_worker_proc: - 1. Re-import our module to ensure monkey patches are applied in child processes + Re-import our module in child processes so monkey patches (custom routes, + custom init_app_state) are applied in multi-API-server mode. """ - # NOTE: This hack ensures that monkey patches are applied in child processes - # to make our custom routes work in multi-API-server settings. import prime_rl.inference.vllm.server # noqa: F401 - run_api_server_worker_proc(listen_address, sock, args, client_config, **uvicorn_kwargs) - + _original_run_api_server_worker_proc(listen_address, sock, args, client_config, **uvicorn_kwargs) -import vllm.entrypoints.openai.api_server -import vllm.entrypoints.cli.serve -# Also monkey patch run_api_server_worker_proc for multi-api-server mode -# This is needed because worker processes spawned by run_multi_api_server -# re-import modules and would otherwise use the original run_server_worker vllm.entrypoints.openai.api_server.init_app_state = custom_init_app_state -vllm.entrypoints.cli.serve.run_api_server_worker_proc = custom_run_api_server_worker_proc +vllm.entrypoints.openai.api_server.build_app = custom_build_app +vllm.v1.utils.run_api_server_worker_proc = custom_run_api_server_worker_proc # Adapted from vllm/entrypoints/cli/serve.py # Only difference we do some config translation (i.e. pass populated namespace # to `parse_args`) and additional arg validation -def server(config: InferenceConfig, vllm_args: list[str]): - from vllm.entrypoints.openai.api_server import run_server +def server(config: InferenceConfig, vllm_extra: dict[str, Any] | None = None): + import os + from vllm.entrypoints.cli.serve import run_headless, run_multi_api_server + from vllm.entrypoints.openai.api_server import run_server + + # Signal worker processes to disable LoRA on MoE layers when LoRA targets don't include experts + if config.lora_target_modules and not any("expert" in m for m in config.lora_target_modules): + os.environ["PRIME_NO_MOE_LORA"] = "1" + + namespace = config.to_vllm() + if vllm_extra: + for key, value in vllm_extra.items(): + setattr(namespace, key, value) parser = FlexibleArgumentParser(description="vLLM OpenAI-Compatible RESTful API server.") parser = make_arg_parser(parser) - args = parser.parse_args(args=vllm_args, namespace=config.to_vllm()) + args = parser.parse_args(args=[], namespace=namespace) assert args is not None validate_parsed_serve_args(args) + args.tool_call_parser = resolve_tool_call_parser(args.model, args.tool_call_parser) + args.enable_auto_tool_choice = args.tool_call_parser is not None + if args.tool_call_parser is not None: + logger.info(f"Using tool_call_parser='{args.tool_call_parser}' for model '{args.model}'") + # Set the worker extension class based on the broadcast backend args.worker_extension_cls = WORKER_EXTENSION_CLS[config.weight_broadcast.type] diff --git a/src/prime_rl/inference/vllm/serving_chat_with_tokens.py b/src/prime_rl/inference/vllm/serving_chat_with_tokens.py index 666db4ffe17..fae9465fbef 100644 --- a/src/prime_rl/inference/vllm/serving_chat_with_tokens.py +++ b/src/prime_rl/inference/vllm/serving_chat_with_tokens.py @@ -1,38 +1,81 @@ -from collections.abc import AsyncGenerator +from collections.abc import AsyncGenerator, AsyncIterator from typing import ClassVar, Optional, Union -import jinja2 from fastapi import Request from pydantic import Field -from vllm.entrypoints.openai.protocol import ( - ChatCompletionRequest, - ChatCompletionResponse, - ErrorResponse, - RequestResponseMetadata, -) -from vllm.entrypoints.openai.serving_chat import OpenAIServingChat +from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest, ChatCompletionResponse +from vllm.entrypoints.openai.chat_completion.serving import OpenAIServingChat +from vllm.entrypoints.openai.engine.protocol import ErrorResponse, RequestResponseMetadata +from vllm.entrypoints.openai.engine.serving import GenerationError from vllm.entrypoints.utils import get_max_tokens +from vllm.exceptions import VLLMValidationError from vllm.logger import init_logger from vllm.outputs import RequestOutput +from vllm.reasoning import ReasoningParser from vllm.sampling_params import BeamSearchParams, SamplingParams -from vllm.tokenizers.mistral import ( - MistralTokenizer, - maybe_serialize_tool_calls, - truncate_tool_call_ids, - validate_request_params, -) -from vllm.v1.sample.logits_processor import validate_logits_processors_parameters + +from prime_rl.inference.vllm.serving_tokens import _RoutedExpertsCaptureBase logger = init_logger(__name__) +class _RoutedExpertsCapture(_RoutedExpertsCaptureBase): + """Chat-endpoint variant: mutates choices in-place because + ``ChatCompletionResponseChoice`` is ``extra='allow'``, so an extra + ``routed_experts`` attribute survives serialization.""" + + def post_process(self, response: ChatCompletionResponse) -> None: + for choice in response.choices: + if choice.index in self.routed_experts: + choice.routed_experts = self.routed_experts[choice.index] + + class ChatCompletionRequestWithTokens(ChatCompletionRequest): field_names: ClassVar[Optional[set[str]]] = None tokens: list[int] = Field(description=("Prompt tokens to use for the request.")) class OpenAIServingChatWithTokens(OpenAIServingChat): - """OpenAI-compatible generate API that allows token-in.""" + """OpenAI-compatible generate API that allows token-in and routed experts capture.""" + + async def chat_completion_full_generator( + self, + request: ChatCompletionRequest, + result_generator: AsyncIterator[RequestOutput], + request_id: str, + model_name: str, + conversation, + tokenizer, + request_metadata: RequestResponseMetadata, + reasoning_parser: ReasoningParser | None = None, + ) -> ErrorResponse | ChatCompletionResponse: + # We need to override the full_generator to be able to capture the routed experts + # By default, VLLM does not save the routed experts into ChatCompletionResponse.choices, so we need to capture them manually + # How this works: + # 1. We create a custom generator that encapsulates the original result_generator in self._generator + # 2. We override it's __aiter__ method to also capture the routed experts as an extra field in ChatCompletionResponse.choices + # 3. We override the full_generator method to use the custom generator instead of the original one if expert routing is enabled + if self.model_config.enable_return_routed_experts: + capture = _RoutedExpertsCapture(result_generator) + result_generator = capture + else: + capture = None + + response = await super().chat_completion_full_generator( + request, + result_generator, + request_id, + model_name, + conversation, + tokenizer, + request_metadata, + reasoning_parser, + ) + + if capture and isinstance(response, ChatCompletionResponse): + capture.post_process(response) + + return response async def create_chat_completion_with_tokens( self, @@ -40,102 +83,40 @@ async def create_chat_completion_with_tokens( raw_request: Optional[Request] = None, ) -> Union[AsyncGenerator[str, None], ChatCompletionResponse, ErrorResponse]: """ - Copy of OpenAIServingChat.create_chat_completion, adapted to use prompt - ids directly via ChatCompletionRequestWithTokens. - """ - error_check_ret = await self._check_model(request) - if error_check_ret is not None: - logger.error("Error with model %s", error_check_ret) - return error_check_ret - - # If the engine is dead, raise the engine's DEAD_ERROR. - # This is required for the streaming case, where we return a - # success status before we actually start generating text :). - if self.engine_client.errored: - raise self.engine_client.dead_error + Chat Completion API similar to OpenAI's API. + See https://platform.openai.com/docs/api-reference/chat/create + for the API specification. This API mimics the OpenAI + Chat Completion API. + """ + # Streaming response + tokenizer = self.renderer.tokenizer + assert tokenizer is not None + reasoning_parser: ReasoningParser | None = None try: - lora_request = self._maybe_get_adapters(request, supports_default_mm_loras=True) - - model_name = self.models.model_name(lora_request) - - tokenizer = await self.engine_client.get_tokenizer() - - tool_parser = self.tool_parser - - if isinstance(tokenizer, MistralTokenizer): - # because of issues with pydantic we need to potentially - # re-serialize the tool_calls field of the request - # for more info: see comment in `maybe_serialize_tool_calls` - maybe_serialize_tool_calls(request) - truncate_tool_call_ids(request) - validate_request_params(request) - - if ( - request.tool_choice == "auto" - and not (self.enable_auto_tools and tool_parser is not None) - and not isinstance(tokenizer, MistralTokenizer) - and not self.use_harmony - ): - # for hf tokenizers, "auto" tools requires - # --enable-auto-tool-choice and --tool-call-parser - return self.create_error_response( - '"auto" tool choice requires --enable-auto-tool-choice and --tool-call-parser to be set' + if self.reasoning_parser_cls: + # Pass the same chat template kwargs as used in tokenization + chat_template_kwargs = self._prepare_extra_chat_template_kwargs( + request.chat_template_kwargs, + self.default_chat_template_kwargs, ) - - if request.tools is None or (request.tool_choice == "none" and self.exclude_tools_when_tool_choice_none): - tool_dicts = None - else: - tool_dicts = [tool.model_dump() for tool in request.tools] - - if not self.use_harmony: - # Common case. - error_check_ret = self._validate_chat_template( - request_chat_template=request.chat_template, - chat_template_kwargs=request.chat_template_kwargs, - trust_request_chat_template=self.trust_request_chat_template, - ) - if error_check_ret is not None: - return error_check_ret - ( - conversation, - request_prompts, - engine_prompts, - ) = await self._preprocess_chat( - request, + reasoning_parser = self.reasoning_parser_cls( tokenizer, - request.messages, - chat_template=request.chat_template or self.chat_template, - chat_template_content_format=self.chat_template_content_format, - add_generation_prompt=request.add_generation_prompt, - continue_final_message=request.continue_final_message, - tool_dicts=tool_dicts, - documents=request.documents, - chat_template_kwargs=request.chat_template_kwargs, - tool_parser=tool_parser, - add_special_tokens=request.add_special_tokens, + chat_template_kwargs=chat_template_kwargs, # type: ignore[call-arg] ) - else: - # For GPT-OSS. - ( - conversation, - request_prompts, - engine_prompts, - ) = self._make_request_with_harmony(request) - except (ValueError, TypeError, RuntimeError, jinja2.TemplateError) as e: - logger.exception("Error in preprocessing prompt inputs") - return self.create_error_response(f"{e} {e.__cause__}") - - # In-place override the engine_prompts with the tokens from the request - assert len(engine_prompts) == 1 - if engine_prompts[0]["prompt_token_ids"] != request.tokens: - logger.warning( - "Prompt tokens provided in request do not match the engine prompt tokens. This may happen due to retokenization discrepancies in multi-turn conversations. Since you are using the /v1/chat/completions/tokens endpoint, we assume you want this behavior and use the provided prompt tokens. If this is undesired, use the standard /v1/chat/completions endpoint instead." - ) - logger.debug(f"engine_prompt_tokens:\n{engine_prompts[0]['prompt_token_ids']}") - logger.debug(f"request_tokens:\n{request.tokens}") + except RuntimeError as e: + logger.exception("Error in reasoning parser creation.") + return self.create_error_response(str(e)) + result = await self.render_chat_request(request) + if isinstance(result, ErrorResponse): + return result + + conversation, engine_prompts = result - engine_prompts[0]["prompt_token_ids"] = request.tokens + # We override prompt tokens directly. + # VLM conversations use MITO (message-based) instead of TITO, so + # multi_modal_data is not expected here. + engine_prompts[0]["prompt_token_ids"] = request.tokens # type: ignore request_id = f"chatcmpl-{self._base_request_id(raw_request, request.request_id)}" @@ -143,26 +124,45 @@ async def create_chat_completion_with_tokens( if raw_request: raw_request.state.request_metadata = request_metadata + try: + lora_request = self._maybe_get_adapters(request, supports_default_mm_loras=True) + + model_name = self.models.model_name(lora_request) + except (ValueError, TypeError, RuntimeError) as e: + logger.exception("Error preparing request components") + return self.create_error_response(e) + # Extract data_parallel_rank from header (router can inject it) data_parallel_rank = self._get_data_parallel_rank(raw_request) # Schedule the request and get the result generator. + max_model_len = self.model_config.max_model_len generators: list[AsyncGenerator[RequestOutput, None]] = [] try: for i, engine_prompt in enumerate(engine_prompts): - prompt_text, _, _ = self._get_prompt_components(request_prompts[i]) + prompt_token_ids = self._extract_prompt_components(engine_prompt).token_ids + # If we are creating sub requests for multiple prompts, ensure that they # have unique request ids. sub_request_id = request_id if len(engine_prompts) == 1 else f"{request_id}_{i}" - if self.default_sampling_params is None: - self.default_sampling_params = {} + prompt_len = self._extract_prompt_len(engine_prompt) + if prompt_len >= max_model_len: + raise VLLMValidationError( + f"This model's maximum context length is " + f"{max_model_len} tokens. However, your request has " + f"{prompt_len} input tokens. Please reduce the length of " + "the input messages.", + parameter="input_tokens", + value=prompt_len, + ) max_tokens = get_max_tokens( - max_model_len=self.max_model_len, - request=request, - input_length=len(engine_prompt["prompt_token_ids"]), - default_sampling_params=self.default_sampling_params, + max_model_len, + request.max_completion_tokens if request.max_completion_tokens is not None else request.max_tokens, + self._extract_prompt_len(engine_prompt), + self.default_sampling_params, + self.override_max_tokens, ) sampling_params: SamplingParams | BeamSearchParams @@ -171,17 +171,12 @@ async def create_chat_completion_with_tokens( else: sampling_params = request.to_sampling_params( max_tokens, - self.model_config.logits_processor_pattern, self.default_sampling_params, ) - validate_logits_processors_parameters( - self.logits_processors, - sampling_params, - ) self._log_inputs( sub_request_id, - request_prompts[i], + engine_prompt, params=sampling_params, lora_request=lora_request, ) @@ -197,36 +192,28 @@ async def create_chat_completion_with_tokens( trace_headers=trace_headers, ) else: - engine_request, tokenization_kwargs = await self._process_inputs( - sub_request_id, - engine_prompt, - sampling_params, - lora_request=lora_request, - trace_headers=trace_headers, - priority=request.priority, + reasoning_ended = ( + reasoning_parser.is_reasoning_end(prompt_token_ids or []) if reasoning_parser else None ) generator = self.engine_client.generate( - engine_request, + engine_prompt, sampling_params, sub_request_id, lora_request=lora_request, trace_headers=trace_headers, priority=request.priority, - prompt_text=prompt_text, - tokenization_kwargs=tokenization_kwargs, data_parallel_rank=data_parallel_rank, + reasoning_ended=reasoning_ended, ) generators.append(generator) except ValueError as e: - # TODO: Use a vllm-specific Validation Error - return self.create_error_response(str(e)) + return self.create_error_response(e) assert len(generators) == 1 (result_generator,) = generators - # Streaming response if request.stream: return self.chat_completion_stream_generator( request, @@ -236,6 +223,7 @@ async def create_chat_completion_with_tokens( conversation, tokenizer, request_metadata, + reasoning_parser, ) try: @@ -247,7 +235,9 @@ async def create_chat_completion_with_tokens( conversation, tokenizer, request_metadata, + reasoning_parser, ) + except GenerationError: + raise # Let FastAPI's global generation_error_handler handle it except ValueError as e: - # TODO: Use a vllm-specific Validation Error - return self.create_error_response(str(e)) + return self.create_error_response(e) diff --git a/src/prime_rl/inference/vllm/serving_tokens.py b/src/prime_rl/inference/vllm/serving_tokens.py new file mode 100644 index 00000000000..9fe08335915 --- /dev/null +++ b/src/prime_rl/inference/vllm/serving_tokens.py @@ -0,0 +1,305 @@ +"""Prime-RL extensions to vLLM's `/inference/v1/generate` handler. + +vLLM 0.20 ships a generic tokens-in / tokens-out handler at +``vllm.entrypoints.serve.disagg.serving.ServingTokens`` that already covers +prefix-cache salting, lora dispatch, multimodal features, prompt logprobs and +priority. Three prime-RL features are not in the upstream protocol though, so +we subclass it to add them back: + +1. ``data_parallel_rank`` routing — read from the ``X-data-parallel-rank`` + header and forwarded to ``engine_client.generate``. The DP-replicated + inference servers prime-RL runs need this to target a specific replica. + +2. ``routed_experts`` per-token export — when the engine emits routing + decisions (``enable_return_routed_experts``), surface them on each choice. + This is what the trainer's router-replay path consumes. + +3. Server-side ``max_tokens`` defaulting — ``ServingTokens`` hands the + client-supplied ``SamplingParams`` to the engine verbatim, and + ``SamplingParams.max_tokens`` defaults to ``16`` (a dataclass-level + default that predates the OpenAI-compat layer). Every other vLLM + endpoint masks this server-side via + ``vllm.entrypoints.utils.get_max_tokens`` (see e.g. + ``OpenAIServingChat`` at ``serving.py:284``); the disagg endpoint + skips that path. Mirror it here so callers that omit ``max_tokens`` + don't silently truncate at 16. Drop once vLLM patches upstream. + +Everything else (request/response schema, sampling params, error handling) +delegates to upstream so we track future vLLM changes for free. +""" + +from __future__ import annotations + +import base64 +from collections.abc import AsyncGenerator +from functools import cached_property + +import numpy as np +from fastapi import Request +from pydantic import Field +from vllm.entrypoints.openai.engine.protocol import ErrorResponse, RequestResponseMetadata +from vllm.entrypoints.serve.disagg.protocol import ( + GenerateRequest, + GenerateResponse, + GenerateResponseChoice, +) +from vllm.entrypoints.serve.disagg.serving import ServingTokens +from vllm.entrypoints.utils import get_max_tokens +from vllm.outputs import RequestOutput +from vllm.sampling_params import RequestOutputKind, SamplingParams + + +class PrimeRlGenerateResponseChoice(GenerateResponseChoice): + routed_experts: dict | None = Field( + default=None, + description=( + "Per-token expert routing decisions (base85-encoded int32 array + shape). " + "Populated only when the engine was launched with " + "``enable_return_routed_experts=True``; otherwise ``None``." + ), + ) + + +class PrimeRlGenerateResponse(GenerateResponse): + choices: list[PrimeRlGenerateResponseChoice] + + +def encode_routed_experts(arr: np.ndarray) -> dict: + return { + "data": base64.b85encode(arr.tobytes()).decode("ascii"), + "shape": list(arr.shape), + } + + +class _RoutedExpertsCaptureBase: + """Wraps the engine result generator and accumulates a + ``{output_index: encoded_experts}`` map as outputs stream. Subclasses + implement ``post_process`` to fold the captured map into the response + in whatever shape the endpoint returns (in-place vs rebuilt).""" + + def __init__(self, generator: AsyncGenerator[RequestOutput, None]): + self._generator = generator + self.routed_experts: dict[int, dict] = {} + + async def __aiter__(self): + async for request_output in self._generator: + for output in request_output.outputs: + if output.routed_experts is not None: + self.routed_experts[output.index] = encode_routed_experts(output.routed_experts) + yield request_output + + +class _RoutedExpertsCapture(_RoutedExpertsCaptureBase): + """Generate-endpoint variant: rebuilds the response with + ``PrimeRlGenerateResponseChoice`` because upstream's + ``GenerateResponseChoice`` isn't ``extra='allow'``, so an attribute + set after construction wouldn't survive serialization.""" + + def post_process(self, response: GenerateResponse) -> PrimeRlGenerateResponse: + new_choices = [ + PrimeRlGenerateResponseChoice( + **choice.model_dump(), + routed_experts=self.routed_experts.get(choice.index), + ) + for choice in response.choices + ] + return PrimeRlGenerateResponse( + request_id=response.request_id, + choices=new_choices, + prompt_logprobs=response.prompt_logprobs, + kv_transfer_params=response.kv_transfer_params, + ) + + +async def _client_set_max_tokens(raw_request: Request | None) -> bool: + """Whether the inbound JSON body carried ``sampling_params.max_tokens``. + + ``GenerateRequest.sampling_params`` is parsed into a ``SamplingParams`` + instance, which means an unset ``max_tokens`` is indistinguishable from + an explicit ``max_tokens=16`` once the request reaches the handler — + both surface as ``sampling_params.max_tokens == 16``. We re-read the + cached body to recover that distinction. When we can't (no raw_request, + non-JSON body, or read error), pessimistically assume the client did + set it so we never clobber an explicit value. + """ + if raw_request is None: + return True + try: + body = await raw_request.json() + except Exception: + return True + if not isinstance(body, dict): + return True + sp = body.get("sampling_params") + return isinstance(sp, dict) and "max_tokens" in sp + + +class PrimeRlServingTokens(ServingTokens): + """ServingTokens + DP-rank routing + routed_experts export + max_tokens defaulting.""" + + @cached_property + def _max_tokens_defaults(self) -> tuple[dict, int | None]: + """Server-side ``max_tokens`` defaulting inputs, mirroring ``OpenAIServingChat``. + + Computed lazily because ``custom_init_app_state`` swaps in this + subclass via ``object.__new__`` + ``__dict__.update`` (so our + ``__init__`` never runs). + """ + diff = self.model_config.get_diff_sampling_param() + mc = self.model_config + override = ( + diff.get("max_tokens") + if mc.generation_config not in ("auto", "vllm") + # Upstream uses ``getattr(..., {})`` directly. Defensive ``or {}`` + # in case a downstream caller ever sets the attribute to ``None`` + # (``getattr``'s default only fires when the attribute is missing, + # not when it exists with a ``None`` value). + else (getattr(mc, "override_generation_config", None) or {}).get("max_new_tokens") + ) + return diff, override + + async def serve_tokens( + self, + request: GenerateRequest, + raw_request: Request | None = None, + ) -> PrimeRlGenerateResponse | ErrorResponse | AsyncGenerator[str, None]: + # Mirrors upstream ``ServingTokens.serve_tokens`` (vllm 0.20). Diffs: + # (a) inject ``data_parallel_rank`` from the inbound header into + # ``engine_client.generate``; (b) default ``sampling_params.max_tokens`` + # to ``max_model_len - prompt_len`` when the caller didn't set it; and + # (c) dispatch to our overridden response builder so ``routed_experts`` + # makes it into the JSON. + error_check_ret = await self._check_model(request) + if error_check_ret is not None: + return error_check_ret + + if self.engine_client.errored: + raise self.engine_client.dead_error + + lora_request = self._maybe_get_adapters(request, supports_default_mm_loras=True) + model_name = self.models.model_name(lora_request) + + request_id = f"generate-tokens-{self._base_request_id(raw_request, request.request_id)}" + request_metadata = RequestResponseMetadata(request_id=request_id) + if raw_request: + raw_request.state.request_metadata = request_metadata + + # Build the engine input — features-aware (MM) or text-only fallback. + # Identical to upstream so we keep tracking it. + if features := request.features: + from vllm.entrypoints.serve.disagg.mm_serde import decode_mm_kwargs_item + from vllm.inputs import mm_input + from vllm.multimodal.inputs import ( + MultiModalKwargsItem, + PlaceholderRange, + ) + + mm_placeholders = { + modality: [PlaceholderRange(offset=p.offset, length=p.length) for p in ranges] + for modality, ranges in features.mm_placeholders.items() + } + mm_kwargs: dict[str, list[MultiModalKwargsItem | None]] = {} + if features.kwargs_data is not None: + for modality, items in features.kwargs_data.items(): + mm_kwargs[modality] = [decode_mm_kwargs_item(item) if item is not None else None for item in items] + else: + for modality, hashes in features.mm_hashes.items(): + mm_kwargs[modality] = [None] * len(hashes) + engine_input = mm_input( + prompt_token_ids=request.token_ids, + mm_kwargs=mm_kwargs, # type: ignore[arg-type] + mm_hashes=features.mm_hashes, + mm_placeholders=mm_placeholders, + cache_salt=request.cache_salt, + ) + else: + (engine_input,) = await self.openai_serving_render.preprocess_completion( + request, + prompt_input=request.token_ids, + prompt_embeds=None, + skip_mm_cache=True, + ) + + sampling_params: SamplingParams = request.sampling_params + + # Server-side ``max_tokens`` defaulting — see module docstring. + # Mirrors ``OpenAIServingChat`` (vllm/entrypoints/openai/chat_completion/ + # serving.py:284) so callers that omit ``max_tokens`` don't get capped + # at vLLM's 16-token ``SamplingParams`` default. + if not await _client_set_max_tokens(raw_request): + diff_sp, override = self._max_tokens_defaults + sampling_params.max_tokens = get_max_tokens( + max_model_len=self.model_config.max_model_len, + max_tokens=None, + input_length=len(request.token_ids), + default_sampling_params=diff_sp, + override_max_tokens=override, + ) + + if self.force_no_detokenize: + sampling_params.detokenize = False + if request.stream: + sampling_params.output_kind = RequestOutputKind.DELTA + + self._log_inputs( + request_id, + engine_input, + params=sampling_params, + lora_request=lora_request, + ) + + trace_headers = None if raw_request is None else await self._get_trace_headers(raw_request.headers) + + result_generator = self.engine_client.generate( + engine_input, + sampling_params, + request_id, + lora_request=lora_request, + trace_headers=trace_headers, + priority=request.priority, + data_parallel_rank=self._get_data_parallel_rank(raw_request), + ) + + if request.stream: + # Streaming path: defer to upstream — prime-RL's renderer client + # only consumes the full response, so adding routed_experts to the + # streaming choice schema is unnecessary churn. + return self.serve_tokens_stream_generator( + request, + result_generator, + request_id, + model_name, + request_metadata, + ) + + return await self.serve_tokens_full_generator( + request, result_generator, request_id, model_name, request_metadata + ) + + async def serve_tokens_full_generator( # type: ignore[override] + self, + request: GenerateRequest, + result_generator: AsyncGenerator[RequestOutput, None], + request_id: str, + model_name: str, + request_metadata: RequestResponseMetadata, + ) -> ErrorResponse | GenerateResponse: + # Mirror serving_chat_with_tokens: wrap the result generator to capture + # routed_experts as it streams, defer the rest to upstream, then post- + # process the response into our PrimeRlGenerateResponse subclass so the + # encoded experts surface in the JSON. Skipping the wrapper when the + # engine isn't producing routed experts keeps us a no-op subclass on + # the common path. + capture: _RoutedExpertsCapture | None = None + if self.model_config.enable_return_routed_experts: + capture = _RoutedExpertsCapture(result_generator) + result_generator = capture # type: ignore[assignment] + + response = await super().serve_tokens_full_generator( + request, result_generator, request_id, model_name, request_metadata + ) + + if capture is not None and isinstance(response, GenerateResponse): + response = capture.post_process(response) + + return response diff --git a/src/prime_rl/inference/vllm/worker/__init__.py b/src/prime_rl/inference/vllm/worker/__init__.py index a4c2ebd0d2f..85108925381 100644 --- a/src/prime_rl/inference/vllm/worker/__init__.py +++ b/src/prime_rl/inference/vllm/worker/__init__.py @@ -1,4 +1,29 @@ -from prime_rl.inference.patches import monkey_patch_LRUCacheWorkerLoRAManager +import logging +import os + +from prime_rl.inference.patches import ( + monkey_patch_fp32_lm_head, + monkey_patch_LRUCacheWorkerLoRAManager, + monkey_patch_minimax_m2_for_lora, + monkey_patch_no_moe_lora, + monkey_patch_skip_lora_module_warnings, +) + +logger = logging.getLogger(__name__) # Monkeypatch LRUCacheWorkerLoRAManager to allow loading adapter inplace without doing it every request monkey_patch_LRUCacheWorkerLoRAManager() +# Skip the per-module regex warning loop in WorkerLoRAManager._load_adapter +# (minutes-long stall on wide MoE models like Qwen3.5-35B-A3B) +monkey_patch_skip_lora_module_warnings() +# Monkeypatch MiniMaxM2 MoE gate dtype and adapter key mapping for LoRA compatibility +monkey_patch_minimax_m2_for_lora() +# Disable LoRA on MoE layers so vLLM picks better kernels (e.g. TRTLLMFlashInfer on Blackwell) +if os.environ.get("PRIME_NO_MOE_LORA") == "1": + logger.info("PRIME_NO_MOE_LORA=1: disabling LoRA on MoE layers") + monkey_patch_no_moe_lora() +else: + logger.info("PRIME_NO_MOE_LORA=0: no patch applied") + +# Install fp32 lm_head patch; self-gates on additional_config["fp32_lm_head"] at call time +monkey_patch_fp32_lm_head() diff --git a/src/prime_rl/inference/vllm/worker/filesystem.py b/src/prime_rl/inference/vllm/worker/filesystem.py index fb2900c446d..1049d17b0c0 100644 --- a/src/prime_rl/inference/vllm/worker/filesystem.py +++ b/src/prime_rl/inference/vllm/worker/filesystem.py @@ -2,7 +2,8 @@ from torch.nn import Module from vllm.model_executor.model_loader import DefaultModelLoader, get_model_loader -from vllm.model_executor.model_loader.utils import process_weights_after_loading + +from prime_rl.inference.vllm.worker.weight_transfer import load_weights_checkpoint_layerwise # This is to get type hints for the Worker class but not actually extend it at runtime as this is required by vLLM worker extension if TYPE_CHECKING: @@ -20,11 +21,19 @@ def init_broadcaster(self) -> None: """Initialize the broadcaster.""" ... - def update_weights(self, weight_path: str) -> None: + def liveness_probe(self) -> None: + """No-op RPC used by the API server liveness endpoint.""" + return None + + def update_weights_from_path(self, weight_path: str) -> None: """Update weights from a specified path in shared filesystem containing a HF-compatible checkpoint.""" # Get vLLM model runner and model + # When enforce_eager=True, model isn't wrapped by torch.compile so no .runnable attr model_runner = self.model_runner - model = model_runner.model.runnable + if hasattr(model_runner.model, "runnable"): + model = model_runner.model.runnable + else: + model = model_runner.model assert isinstance(model, Module) # Get vLLM model loader @@ -38,8 +47,9 @@ def update_weights(self, weight_path: str) -> None: allow_patterns_overrides=getattr(model, "allow_patterns_overrides", None), ) weights_iterator = model_loader._get_weights_iterator(local_source) - model.load_weights(weights_iterator) # type: ignore - - # Process weights after loading (important for some models) - device = next(model.parameters()).device - process_weights_after_loading(model, self.model_runner.model_config, device) + load_weights_checkpoint_layerwise( + model, + weights_iterator, + self.model_runner.model_config, + self.vllm_config, + ) diff --git a/src/prime_rl/inference/vllm/worker/nccl.py b/src/prime_rl/inference/vllm/worker/nccl.py index 27dd54d1c80..705790e4091 100644 --- a/src/prime_rl/inference/vllm/worker/nccl.py +++ b/src/prime_rl/inference/vllm/worker/nccl.py @@ -4,10 +4,15 @@ import torch from torch.nn import Module from vllm.distributed.device_communicators.pynccl import PyNcclCommunicator -from vllm.distributed.parallel_state import get_tp_group from vllm.distributed.utils import StatelessProcessGroup from vllm.logger import init_logger -from vllm.model_executor.model_loader.utils import process_weights_after_loading + +from prime_rl.inference.vllm.worker.weight_transfer import ( + load_weights_checkpoint_layerwise, + load_weights_kernel, + update_mla_absorbed_weights, +) +from prime_rl.utils.nccl import disable_nccl_p2p_if_unavailable # This is to get type hints for the Worker class but not actually extend it at runtime as this is required by vLLM worker extension if TYPE_CHECKING: @@ -67,6 +72,7 @@ def __init__( timeout: int, ): logger.info(f"Initializing NCCL broadcast receiver ({host}:{port}, rank={rank}, world_size={world_size})") + disable_nccl_p2p_if_unavailable() pg = StatelessProcessGroup.create(host=host, port=port, rank=rank, world_size=world_size, store_timeout=timeout) self.communicator = PyNcclCommunicator(pg, device=device) @@ -86,35 +92,65 @@ def receive_state_dict(self): class NCCLWeightUpdateWorker(Worker): """vLLM worker extension for updating weights in-place using NCCL.""" - def init_broadcaster(self, host: str, port: int, server_rank: int, num_inference_server: int, timeout: int) -> None: - """Initialize the NCCL broadcast receiver.""" - tp_size = get_tp_group().world_size - tp_rank = get_tp_group().rank - global_rank_inference = (server_rank * tp_size) + tp_rank - global_inference_world_size = num_inference_server * tp_size + def init_broadcaster( + self, + host: str, + port: int, + rank_offset: int, + inference_world_size: int, + timeout: int, + quantize_in_weight_transfer: bool = False, + ) -> None: + """Initialize the NCCL broadcast receiver. + + Args: + rank_offset: Starting GPU offset for this server in the global inference group. + inference_world_size: Total number of inference GPUs across all servers. + """ + self.quantize_in_weight_transfer = quantize_in_weight_transfer + # Use the worker's device index directly as the local rank. + # The previous dp_group-based computation broke in vLLM v1 multiprocess + # DP mode where each worker is a separate process with a singleton + # DP group (rank_in_group is always 0). + local_rank = self.device.index + global_rank_inference = rank_offset + local_rank logger.info( - f"Worker [tp={tp_rank} server_rank={server_rank}] -> [global_rank={global_rank_inference} global_world_size={global_inference_world_size}]" + f"Worker [local_rank={local_rank} rank_offset={rank_offset}] " + f"-> [global_rank={global_rank_inference} inference_world_size={inference_world_size}]" ) self.nccl_broadcast_receiver = NCCLWeightBroadcastReceiver( host=host, port=port, rank=global_rank_inference + 1, # +1 as the trainer broadcaster is on rank 0 - world_size=global_inference_world_size + 1, # +1 as the trainer broadcaster is on rank 0 + world_size=inference_world_size + 1, # +1 as the trainer broadcaster is on rank 0 device=self.device, timeout=timeout, ) - def update_weights(self, weight_dir: str) -> None: + def liveness_probe(self) -> None: + """No-op RPC used by the API server liveness endpoint.""" + return None + + def update_weights_from_path(self, weight_dir: str) -> None: """Update weights with the nccl communicator.""" model_runner = self.model_runner - model = model_runner.model.runnable + if hasattr(model_runner.model, "runnable"): + model = model_runner.model.runnable + else: + model = model_runner.model assert isinstance(model, Module) state_iter = self.nccl_broadcast_receiver.receive_state_dict() - model.load_weights(state_iter) # type: ignore - - # # Process weights after loading (important for some models) - device = next(model.parameters()).device - process_weights_after_loading(model, self.model_runner.model_config, device) + if self.quantize_in_weight_transfer: + load_weights_kernel(model, state_iter) + update_mla_absorbed_weights(model) + return + + load_weights_checkpoint_layerwise( + model, + state_iter, + self.model_runner.model_config, + self.vllm_config, + ) diff --git a/src/prime_rl/inference/vllm/worker/weight_transfer.py b/src/prime_rl/inference/vllm/worker/weight_transfer.py new file mode 100644 index 00000000000..7b448f021a7 --- /dev/null +++ b/src/prime_rl/inference/vllm/worker/weight_transfer.py @@ -0,0 +1,152 @@ +from typing import Generator, Iterable + +import torch +from torch.nn import Module +from vllm.config import set_current_vllm_config +from vllm.logger import init_logger +from vllm.model_executor.model_loader.reload import finalize_layerwise_reload, initialize_layerwise_reload + +logger = init_logger("vllm.inference.vllm.worker_weight_transfer") + + +def load_weights_checkpoint_layerwise( + model: Module, + state_iter: Iterable[tuple[str, torch.Tensor]], + model_config, + vllm_config, +) -> None: + logger.info("Reloading checkpoint-format weights with vLLM layerwise processing") + device = next(model.parameters()).device + with torch.device(device), set_current_vllm_config(vllm_config): + initialize_layerwise_reload(model) + model.load_weights(state_iter) # type: ignore + finalize_layerwise_reload(model, model_config) + + +def _invert_logical_to_physical_map(logical_to_physical_map: torch.Tensor, num_physical_experts: int) -> torch.Tensor: + """Build a physical expert -> logical expert map from vLLM EPLB state.""" + physical_to_logical = torch.full( + (num_physical_experts,), + -1, + dtype=torch.long, + device=logical_to_physical_map.device, + ) + logical_indices = torch.arange( + logical_to_physical_map.shape[0], + dtype=torch.long, + device=logical_to_physical_map.device, + )[:, None].expand_as(logical_to_physical_map) + physical_indices = logical_to_physical_map.to(torch.long) + invalid = (physical_indices < -1) | (physical_indices >= num_physical_experts) + if invalid.any(): + invalid_indices = physical_indices[invalid].unique().tolist() + raise ValueError(f"EPLB maps to invalid physical experts: {invalid_indices}") + + valid = physical_indices >= 0 + physical_to_logical[physical_indices[valid]] = logical_indices[valid] + return physical_to_logical + + +def _build_expert_source_indices(module) -> torch.Tensor | None: + if module._expert_map is None: + return None + + physical_indices = torch.where(module._expert_map >= 0)[0] + local_indices = module._expert_map[physical_indices] + physical_indices = physical_indices[local_indices.argsort()] + + eplb_layer_state = getattr(module, "eplb_state", None) + logical_to_physical_map = getattr(eplb_layer_state, "logical_to_physical_map", None) + if logical_to_physical_map is None: + return physical_indices + + physical_to_logical = _invert_logical_to_physical_map(logical_to_physical_map, module.global_num_experts) + logical_indices = physical_to_logical[physical_indices.to(physical_to_logical.device)] + if (logical_indices < 0).any(): + missing = physical_indices[(logical_indices < 0).to(physical_indices.device)].tolist() + raise ValueError(f"EPLB has no logical mapping for local physical experts: {missing}") + + return logical_indices.to(physical_indices.device) + + +def build_expert_map(model: Module) -> dict[str, torch.Tensor]: + """Map FusedMoE module names to source expert indices local to this worker.""" + from vllm.model_executor.layers.fused_moe.layer import FusedMoE + + source_indices_by_module: dict[str, torch.Tensor] = {} + for module_name, module in model.named_modules(): + if not isinstance(module, FusedMoE): + continue + source_indices = _build_expert_source_indices(module) + if source_indices is None: + continue + source_indices_by_module[module_name] = source_indices + return source_indices_by_module + + +@torch.no_grad() +def load_weights_kernel(model: Module, state_iter: Generator[tuple[str, torch.Tensor], None, None]) -> None: + """Load vLLM kernel-format tensors using in-place copy_ updates.""" + params = dict(model.named_parameters()) + expert_source_indices = build_expert_map(model) + + loaded = 0 + skipped: list[str] = [] + shape_mismatches: list[str] = [] + + for name, tensor in state_iter: + if name not in params: + skipped.append(name) + continue + + param = params[name] + if param.shape != tensor.shape: + for module_name, source_indices in expert_source_indices.items(): + if not name.startswith(f"{module_name}."): + continue + tensor = tensor[source_indices.to(tensor.device)] + break + + if param.shape != tensor.shape: + shape_mismatches.append(f"{name}: param={list(param.shape)} != received={list(tensor.shape)}") + continue + + param.copy_(tensor) + loaded += 1 + + if shape_mismatches: + raise ValueError(f"Kernel weight transfer had {len(shape_mismatches)} shape mismatches: {shape_mismatches}") + if skipped: + raise ValueError(f"Kernel weight transfer skipped {len(skipped)} weights not found in model: {skipped}") + logger.debug(f"Kernel weight transfer copied {loaded} weights in-place") + + +@torch.no_grad() +def update_mla_absorbed_weights(model: Module) -> None: + """Recompute MLA absorbed KV weights after in-place kv_b_proj updates.""" + from vllm.model_executor.layers.quantization.utils.quant_utils import get_and_maybe_dequant_weights + + for name, module in model.named_modules(): + has_absorbed_weights = hasattr(module, "W_UV") or hasattr(module, "W_UK_T") + if not has_absorbed_weights or not hasattr(module, "kv_b_proj"): + continue + + if hasattr(module, "W_UV"): + out_dtype = module.W_UV.dtype + else: + out_dtype = torch.bfloat16 + + kv_b_proj_weight = get_and_maybe_dequant_weights(module.kv_b_proj, out_dtype=out_dtype).T + kv_b_proj_weight = kv_b_proj_weight.view( + module.kv_lora_rank, + module.num_heads, + module.qk_nope_head_dim + module.v_head_dim, + ) + w_uk, w_uv = kv_b_proj_weight.split([module.qk_nope_head_dim, module.v_head_dim], dim=-1) + + if hasattr(module, "W_UV"): + module.W_UV.copy_(w_uv.transpose(0, 1)) + if hasattr(module, "W_UK_T"): + module.W_UK_T.copy_(w_uk.permute(1, 2, 0)) + + logger.debug(f"Updated MLA absorbed weights for module {name}") diff --git a/src/prime_rl/orchestrator/advantage.py b/src/prime_rl/orchestrator/advantage.py index b9774a67670..63b1d503250 100644 --- a/src/prime_rl/orchestrator/advantage.py +++ b/src/prime_rl/orchestrator/advantage.py @@ -1,29 +1,166 @@ +from dataclasses import dataclass +from typing import Callable + import torch +import verifiers as vf +from jaxtyping import Float +from torch import Tensor + +from prime_rl.configs.orchestrator import ( + AdvantageConfig, + CustomAdvantageConfig, + LengthPenaltyConfig, + TokensLengthPenaltyConfig, + TurnsLengthPenaltyConfig, +) +from prime_rl.orchestrator.vf_utils import get_model_completion_len, get_num_turns, get_tool_response_len +from prime_rl.utils.utils import import_object + + +@dataclass +class AdvantageInputs: + """Inputs for advantage computation. + + `rollouts` is grouped by problem: `rollouts[i][j]` is the j-th rollout for problem i, + so `len(rollouts) == num_problems` and `len(rollouts[0]) == rollouts_per_example`. + """ + + rollouts: list[list[vf.RolloutOutput]] + + +@dataclass +class AdvantageOutputs: + """Outputs from advantage computation.""" + + advantages: Float[Tensor, "num_problems rollouts_per_example"] + + +AdvantageFn = Callable[..., AdvantageOutputs] +"""Type for an advantage function. + +Expected signature: + def my_advantage(inputs: AdvantageInputs, **kwargs) -> AdvantageOutputs: + ... +""" + + +def default_advantage_fn( + inputs: AdvantageInputs, + length_penalty: LengthPenaltyConfig | None = None, +) -> AdvantageOutputs: + """Default GRPO advantage: reward minus per-problem baseline. + + `length_penalty` enables correctness-gated efficiency shaping over a per-rollout + cost: tokens (weighted completion + tool-response) or trajectory turn count. + """ + rewards = torch.tensor([[r["reward"] for r in group] for group in inputs.rollouts], dtype=torch.float32) + + if isinstance(length_penalty, TokensLengthPenaltyConfig): + w_c = length_penalty.completion_weight + w_t = length_penalty.tool_response_weight + costs = torch.tensor( + [ + [w_c * get_model_completion_len(r) + w_t * get_tool_response_len(r) for r in group] + for group in inputs.rollouts + ], + dtype=rewards.dtype, + ) + return AdvantageOutputs(advantages=_efficiency_shaping(rewards, costs)) + if isinstance(length_penalty, TurnsLengthPenaltyConfig): + costs = torch.tensor( + [[get_num_turns(r) for r in group] for group in inputs.rollouts], + dtype=rewards.dtype, + ) + return AdvantageOutputs(advantages=_efficiency_shaping(rewards, costs)) + + baseline = rewards.mean(dim=1, keepdim=True) + return AdvantageOutputs(advantages=rewards - baseline) -from prime_rl.orchestrator.config import AdvantageConfig + +def _efficiency_shaping( + rewards: Float[Tensor, "num_problems rollouts_per_example"], + costs: Float[Tensor, "num_problems rollouts_per_example"], +) -> Float[Tensor, "num_problems rollouts_per_example"]: + """Correctness-gated efficiency shaping with bounded advantages. + + Shapes rewards with a bounded efficiency bonus before standard GRPO subtraction, + preserving zero-mean advantages per group. `costs` is a per-rollout cost (e.g., + completion length in tokens or number of turns). + + Correct rollouts get reward amplified by up to 2x based on relative efficiency. + Incorrect rollouts are untouched. Lower-cost correct rollouts get higher advantage. + """ + max_reward = rewards.max(dim=1, keepdim=True).values + correct_mask = rewards >= max_reward + num_correct = correct_mask.sum(dim=1, keepdim=True) + + # No shaping when max reward is 0 — no correct rollouts to differentiate + has_correct = max_reward > 0 + + # Mean cost of correct rollouts per problem + correct_costs = costs * correct_mask + mean_correct_cost = correct_costs.sum(dim=1, keepdim=True) / num_correct.clamp(min=1) + + # Bounded efficiency bonus: [0, 1], positive for below-average cost, zero for above. + # When mean_correct_cost is 0 (e.g. tool-only shaping with no harness metric, or + # all-zero turn counts), no rollouts can be differentiated — fall back to no bonus. + has_cost = mean_correct_cost > 0 + safe_mean = torch.where(has_cost, mean_correct_cost, torch.ones_like(mean_correct_cost)) + bonus = (1 - costs / safe_mean).clamp(0, 1) * has_cost + + # Shape rewards: correct rollouts amplified by up to 2x, incorrect untouched + shaped_rewards = rewards * (1 + bonus * correct_mask) + baseline = shaped_rewards.mean(dim=1, keepdim=True) + + shaped = shaped_rewards - baseline + unshaped = rewards - rewards.mean(dim=1, keepdim=True) + return torch.where(has_correct, shaped, unshaped) + + +def setup_advantage_fn(config: AdvantageConfig) -> AdvantageFn: + """Setup advantage function from config.""" + if isinstance(config, CustomAdvantageConfig): + custom_fn = import_object(config.import_path) + kwargs = config.kwargs + + def advantage_fn(inputs: AdvantageInputs) -> AdvantageOutputs: + return custom_fn(inputs, **kwargs) + + return advantage_fn + + def advantage_fn(inputs: AdvantageInputs) -> AdvantageOutputs: + return default_advantage_fn(inputs, length_penalty=config.length_penalty) + + return advantage_fn def compute_advantages( - rewards: list[float], - completion_lengths: list[int], + rollouts: list[vf.RolloutOutput], samples_per_problem: int, advantage_config: AdvantageConfig | None, -) -> list[float]: +) -> None: """ - Computes advantages from a flattened list of rewards, grouped by problem. + Computes advantages from rollouts, grouped by problem. + Stores advantages in-place on the rollouts. Args: - rewards: Flattened list of rewards where first `samples_per_problem` rewards are for the first problem - completion_lengths: List of completion lengths for each reward + rollouts: List of rollouts to store advantages on samples_per_problem: Number of samples (and thus, rewards) per problem - advantage_config: Configuration for advantage computation + advantage_config: Configuration for advantage computation (DefaultAdvantageConfig or CustomAdvantageConfig) """ + rewards = [r["reward"] for r in rollouts] + if not advantage_config: - return rewards - rewards = torch.tensor(rewards).view(-1, samples_per_problem) - lengths = torch.tensor(completion_lengths).view(-1, samples_per_problem) - if advantage_config.length_weighted_mean: - baseline = (rewards * lengths).sum(dim=1, keepdim=True) / lengths.sum(dim=1, keepdim=True) - else: - baseline = rewards.mean(dim=1, keepdim=True) - return (rewards - baseline).flatten().tolist() + for rollout, reward in zip(rollouts, rewards): + rollout["advantage"] = reward + return + + advantage_fn = setup_advantage_fn(advantage_config) + grouped = [rollouts[i : i + samples_per_problem] for i in range(0, len(rollouts), samples_per_problem)] + inputs = AdvantageInputs(rollouts=grouped) + + result = advantage_fn(inputs) + advantages = result.advantages.flatten().tolist() + + for rollout, advantage in zip(rollouts, advantages): + rollout["advantage"] = advantage diff --git a/src/prime_rl/orchestrator/buffer.py b/src/prime_rl/orchestrator/buffer.py index 5c22cfb789d..0217faea761 100644 --- a/src/prime_rl/orchestrator/buffer.py +++ b/src/prime_rl/orchestrator/buffer.py @@ -1,293 +1,323 @@ +from __future__ import annotations + import hashlib import json import random from collections import defaultdict from functools import partial from pathlib import Path -from typing import cast +from typing import TYPE_CHECKING, cast import verifiers as vf -from datasets import Dataset +from verifiers.utils.save_utils import make_serializable -from prime_rl.orchestrator.config import BufferConfig +from prime_rl.configs.orchestrator import BufferConfig from prime_rl.utils.logger import get_logger from prime_rl.utils.utils import format_num, mean, mean_normalize -from prime_rl.utils.vf import from_serializable_state, to_serializable_state +if TYPE_CHECKING: + from prime_rl.orchestrator.envs import TrainEnv, TrainEnvs -class Buffer: - """A buffer for storing rollouts and metadata.""" - - POOLS = ["easy", "normal", "hard"] - - def __init__( - self, - dataset: Dataset, - env_names: list[str], - buffer_config: BufferConfig, - ): - self.dataset = dataset - self.env_names = env_names - self.config = buffer_config - self.logger = get_logger() - if self.config.seed is not None: - random.seed(self.config.seed) - - # Basic assertions - assert "example_id" in self.dataset.column_names, "The dataset must contain a `example_id` column." - assert "prompt" in self.dataset.column_names, "The dataset must contain a `prompt` column." - assert "task" in self.dataset.column_names, "The dataset must contain a `task` column." - assert len(self.dataset) > 0, "The dataset must contain at least one example." - assert isinstance(self.dataset["example_id"][0], int), "The `example_id` column must be of type int." - assert len(set(self.dataset["example_id"])) == len(self.dataset), "The `example_id` column must be unique." - assert set(self.dataset["task"]) == set(self.env_names), "The `task` column must contain all environment names." - - # Initialize example buffer (env_name -> (example_id -> example)) - self.example_buffer: dict[str, dict[int, dict]] = defaultdict(dict) - for example in map(partial(cast, dict), self.dataset): - self.example_buffer[example["task"]][example["example_id"]] = example - assert len(self.example_buffer) == len(self.env_names) - self.logger.debug( - f"Initialized buffer with {format_num(len(self.dataset), precision=0)} example(s) in {len(self.env_names)} environment(s)" - ) +POOLS = ["easy", "normal", "hard"] - if self.config.env_ratios is not None: - # Convert ratios to probabilities - env_ratio = mean_normalize(self.config.env_ratios) - self.env_probs = {env_name: ratio for env_name, ratio in zip(self.env_names, env_ratio)} - self.logger.debug( - f"Sampling buffer according to provided environment ratios ({', '.join(f'{k}={v:.2f}' for k, v in self.env_probs.items())})" - ) - else: - # Count examples per environment to sample according to natural env distribution - env_counts = [len(self.example_buffer[env_name]) for env_name in self.env_names] - env_ratio = mean_normalize(env_counts) - self.env_probs = {env_name: ratio for env_name, ratio in zip(self.env_names, env_ratio)} - self.logger.debug( - f"Sampling buffer according to natural environment distribution ({', '.join(f'{k}={v:.2f}' for k, v in self.env_probs.items())})" - ) - # Initialize buffers for easy/ hard examples +class _EnvBuffer: + """Manages examples and difficulty pools for a single env.""" + + def __init__(self, env: TrainEnv, config: BufferConfig): + self.env_name = env.name + self.config = config + + dataset = env.get_dataset(seed=config.seed) + if "example_id" not in dataset.column_names: + dataset = dataset.map(lambda ex, idx: {**ex, "example_id": idx}, with_indices=True) + + assert len(dataset) > 0, f"Dataset for {env.name} must contain at least one example." + assert "example_id" in dataset.column_names, f"Dataset for {env.name} must contain an `example_id` column." + assert "prompt" in dataset.column_names, f"Dataset for {env.name} must contain a `prompt` column." + + self.examples: dict[int, dict] = {} + for example in map(partial(cast, dict), dataset): + example["env_name"] = env.name + self.examples[example["example_id"]] = example + self.easy_examples: list[dict] = [] self.hard_examples: list[dict] = [] - # Initialize rollout buffer (flat list of rollouts) - self.rollout_buffer: list[vf.State] = [] - self.reset_step_metrics() + @property + def num_normal(self) -> int: + return len(self.examples) + + @property + def num_total(self) -> int: + return self.num_normal + len(self.easy_examples) + len(self.hard_examples) + + def sample_example(self) -> dict: + key = random.choice(tuple(self.examples)) + return self.examples[key] + def get_example_hash(self, example: dict) -> str: - """Returns a hash of the example based on hash keys.""" hash_keys = [key for key in self.config.hash_keys if key in example] assert hash_keys, "No hashable keys found in example." return hashlib.sha256(json.dumps([example[key] for key in hash_keys]).encode()).hexdigest() - def save(self, path: Path) -> None: - """Saves pool assignments and rollout buffer.""" - path.mkdir(parents=True, exist_ok=True) + def update_pools(self, example_id: int, avg_reward: float) -> str: + """Assign example to pool based on reward. Returns pool name.""" + if self.config.easy_threshold is not None and avg_reward >= self.config.easy_threshold: + pool = "easy" + elif self.config.hard_threshold is not None and avg_reward <= self.config.hard_threshold: + pool = "hard" + else: + pool = "normal" - def write_jsonl(lst: list[dict], path: Path) -> None: - with open(path, "w") as f: - for item in lst: - f.write(json.dumps(item) + "\n") + if pool != "normal" and example_id in self.examples: + example = self.examples.pop(example_id) + target = self.easy_examples if pool == "easy" else self.hard_examples + target.append(example) - write_jsonl(self.easy_examples, path / "easy_examples.jsonl") - write_jsonl(self.hard_examples, path / "hard_examples.jsonl") + self.num_examples_per_step[pool] += 1 + return pool - serializable_rollouts = [to_serializable_state(rollout) for rollout in self.rollout_buffer] - write_jsonl(serializable_rollouts, path / "rollout_buffer.jsonl") + def reset_step_metrics(self) -> None: + zero = lambda: {p: 0 for p in POOLS} + self.num_examples_per_step = zero() + self.num_rollouts_per_step = zero() - def load(self, path: Path) -> None: - """Loads pool assignments and rollouts.""" + def get_metrics(self) -> dict[str, float]: + metrics = {} + num_examples = sum(self.num_examples_per_step.values()) + num_rollouts = sum(self.num_rollouts_per_step.values()) - def read_jsonl(path: Path) -> list[dict]: - with open(path, "r") as f: - return [json.loads(line) for line in f] + for pool in ["easy", "hard"]: + if num_examples: + metrics[f"evicted_examples/{self.env_name}/{pool}"] = self.num_examples_per_step[pool] / num_examples + if num_rollouts: + metrics[f"filtered_rollouts/{self.env_name}/{pool}"] = self.num_rollouts_per_step[pool] / num_rollouts - saved_easy_examples = read_jsonl(path / "easy_examples.jsonl") - saved_hard_examples = read_jsonl(path / "hard_examples.jsonl") - saved_rollout_buffer = [ - from_serializable_state(rollout) for rollout in read_jsonl(path / "rollout_buffer.jsonl") - ] - - if any(saved_easy_examples) or any(saved_hard_examples) or any(saved_rollout_buffer): - # Build hash lookup for example buffer (env -> (example_hash -> example_id)) - example_hash_lookup = defaultdict(dict) - all_hashes = set() - for env in self.example_buffer: - for example_id, example in self.example_buffer[env].items(): - example_hash = self.get_example_hash(example) - if example_hash in all_hashes: - self.logger.warning( - f"Duplicate example hash found based on hash_keys={self.config.hash_keys}. Overwriting with latest example. This may cause unexpected behavior when resuming the buffer." - ) - example_hash_lookup[env][example_hash] = example_id - all_hashes.add(example_hash) - - def move_saved_pool(saved_examples: list[dict], target_pool: list[dict]) -> int: - """Moves saved examples to the target pool from example buffer based on hash lookup.""" - num_moved = 0 - for example in saved_examples: - example_hash = self.get_example_hash(example) - for env in example_hash_lookup: - if example_hash in example_hash_lookup[env]: - example_id = example_hash_lookup[env][example_hash] - example = self.example_buffer[env].pop(example_id, None) - if example is not None: - target_pool.append(example) - num_moved += 1 - break - return num_moved - - if any(saved_easy_examples): - num_moved = move_saved_pool(saved_easy_examples, self.easy_examples) - self.logger.debug( - f"Loaded {num_moved}/{len(saved_easy_examples)} example(s) to easy pool from checkpoint." - ) - if num_moved != len(saved_easy_examples): - num_not_moved = len(saved_easy_examples) - num_moved - self.logger.warning( - f"Could not move {num_not_moved} example(s) from checkpoint to easy pool. This usually means you resumed with an env mix that does not contain all previous examples." - ) + pool_counts = [len(self.easy_examples), self.num_normal, len(self.hard_examples)] + pool_ratios = mean_normalize(pool_counts) + for pool, ratio in zip(POOLS, pool_ratios): + metrics[f"pool/{self.env_name}/{pool}"] = ratio - if any(saved_hard_examples): - num_moved = move_saved_pool(saved_hard_examples, self.hard_examples) - self.logger.debug( - f"Moved {num_moved}/{len(saved_hard_examples)} example(s) to hard pool from checkpoint." - ) - if num_moved != len(saved_hard_examples): - num_not_moved = len(saved_hard_examples) - num_moved - self.logger.warning( - f"Could not move {num_not_moved} example(s) from checkpoint to hard pool. This usually means you resumed with an env mix that does not contain all previous examples." - ) + self.reset_step_metrics() + return metrics - if any(saved_rollout_buffer): - # Extend rollout buffer, but only include rollouts for which the example still exists in the example buffer - valid_saved_rollouts = [ - rollout for rollout in saved_rollout_buffer if rollout["task"] in self.env_names - ] - self.rollout_buffer.extend(valid_saved_rollouts) - self.logger.debug(f"Loaded {len(valid_saved_rollouts)} rollout(s) from checkpoint.") - - # Load rollouts, filtering out removed environments and problems - def convert_examples_to_normal(examples: list[dict], fraction: float) -> int: - """Moves a fraction of examples from the given pool back to normal.""" - if fraction <= 0.0 or not examples: - return 0 - num_moved = round(len(examples) * fraction) - if num_moved <= 0: - return 0 - for _ in range(num_moved): - example = random.choice(examples) - env_name = example["task"] - example_id = example["example_id"] - examples.remove(example) - self.example_buffer[env_name][example_id] = example - return num_moved - - num_easy_examples = len(self.easy_examples) - num_moved = convert_examples_to_normal(self.easy_examples, self.config.easy_fraction) - self.logger.debug(f"Converted {num_moved}/{num_easy_examples} example(s) back to normal from easy pool.") - num_hard_examples = len(self.hard_examples) - num_moved = convert_examples_to_normal(self.hard_examples, self.config.hard_fraction) - self.logger.debug(f"Converted {num_moved}/{num_hard_examples} example(s) back to normal from hard pool.") - else: - self.logger.debug("No easy/ hard examples or rollouts found in checkpoint") - def sample_examples(self, n: int) -> list[dict]: - """Samples n examples from the buffer, respecting env ratios.""" +class Buffer: + """Manages multiple Buffers with env-ratio-aware sampling.""" - non_empty_envs = [env for env, examples in self.example_buffer.items() if examples] + def __init__(self, envs: TrainEnvs, config: BufferConfig): + self.config = config + self.logger = get_logger() - if not non_empty_envs: - raise ValueError("No environments left with examples.") + if config.seed is not None: + random.seed(config.seed) + + self.env_buffers: dict[str, _EnvBuffer] = {} + for env in envs: + self.env_buffers[env.name] = _EnvBuffer(env, config) + self.env_names = envs.names + + total = sum(eb.num_total for eb in self.env_buffers.values()) + self.logger.debug( + f"Initialized buffer with {format_num(total, precision=0)} example(s) " + f"in {len(self.env_names)} environment(s)" + ) + + env_ratios = [env.config.ratio for env in envs] + if any(r is not None for r in env_ratios): + env_ratio = mean_normalize(env_ratios) + self.env_probs = dict(zip(self.env_names, env_ratio)) + self.logger.debug( + f"Sampling buffer according to provided environment ratios " + f"({', '.join(f'{k}={v:.2f}' for k, v in self.env_probs.items())})" + ) + else: + env_counts = [self.env_buffers[name].num_normal for name in self.env_names] + env_ratio = mean_normalize(env_counts) + self.env_probs = dict(zip(self.env_names, env_ratio)) + self.logger.debug( + f"Sampling buffer according to natural environment distribution " + f"({', '.join(f'{k}={v:.2f}' for k, v in self.env_probs.items())})" + ) - non_empty_env_probs = [self.env_probs[env] for env in non_empty_envs] - sampled_examples = [] - for sampled_env in random.choices(non_empty_envs, weights=non_empty_env_probs, k=n): - sampled_example = random.choice(list(self.example_buffer[sampled_env].values())) - sampled_examples.append(sampled_example) + self.rollout_buffer: list[vf.RolloutOutput] = [] - return sampled_examples + def sample_examples(self, n: int) -> list[dict]: + """Samples n examples across envs, respecting env ratios.""" + non_empty = [name for name, eb in self.env_buffers.items() if eb.examples] + if not non_empty: + raise ValueError("No environments left with examples.") - def update(self, rollouts: list[vf.State]): - """Updates the buffer state with completed rollouts.""" + weights = [self.env_probs[name] for name in non_empty] + return [self.env_buffers[name].sample_example() for name in random.choices(non_empty, weights=weights, k=n)] + def update(self, rollouts: list[vf.RolloutOutput]): + """Updates buffer state with completed rollouts.""" rollouts_by_example = defaultdict(list) for rollout in rollouts: - rollouts_by_example[rollout["example_id"]].append(rollout) + rollouts_by_example[(rollout["env_name"], rollout["example_id"])].append(rollout) - for example_id, example_rollouts in rollouts_by_example.items(): + for (env_name, example_id), example_rollouts in rollouts_by_example.items(): + eb = self.env_buffers[env_name] avg_reward = mean([r["reward"] for r in example_rollouts]) - env_name = example_rollouts[0]["task"] - - if self.config.easy_threshold is not None and avg_reward >= self.config.easy_threshold: - pool = "easy" - elif self.config.hard_threshold is not None and avg_reward <= self.config.hard_threshold: - pool = "hard" - else: - pool = "normal" - - if pool != "normal" and example_id in self.example_buffer[env_name]: - example = self.example_buffer[env_name].pop(example_id) - target_pool = self.easy_examples if pool == "easy" else self.hard_examples - target_pool.append(example) + eb.update_pools(example_id, avg_reward) - self.num_examples_per_step[env_name][pool] += 1 if self.config.online_difficulty_filtering: if avg_reward == 0.0: - self.num_rollouts_per_step[env_name]["hard"] += len(example_rollouts) + eb.num_rollouts_per_step["hard"] += len(example_rollouts) continue elif avg_reward == 1.0: - self.num_rollouts_per_step[env_name]["easy"] += len(example_rollouts) + eb.num_rollouts_per_step["easy"] += len(example_rollouts) continue - self.num_rollouts_per_step[env_name]["normal"] += len(example_rollouts) + eb.num_rollouts_per_step["normal"] += len(example_rollouts) self.rollout_buffer.extend(example_rollouts) - def sample_rollouts(self, n: int) -> list[vf.State]: + def sample_rollouts(self, n: int) -> list[vf.RolloutOutput]: """Samples the latest n rollouts from the buffer.""" n = min(n, len(self.rollout_buffer)) - sampled_rollouts = self.rollout_buffer[-n:] + sampled = self.rollout_buffer[-n:] self.rollout_buffer = self.rollout_buffer[:-n] - return sampled_rollouts + return sampled - def reset_step_metrics(self) -> None: - """Reset per-step metrics (called after get_metrics).""" - zero_per_pool = lambda: {p: 0 for p in self.POOLS} - # num examples per env per step per pool (env_name -> (pool -> num_examples)) - self.num_examples_per_step = {env: zero_per_pool() for env in self.env_names} - # num rollouts per env per step per pool (env_name -> (pool -> num_rollouts)) - self.num_rollouts_per_step = {env: zero_per_pool() for env in self.env_names} + def save(self, path: Path) -> None: + """Saves pool assignments and rollout buffer.""" + path.mkdir(parents=True, exist_ok=True) - def get_metrics(self) -> dict[str, float]: - """Returns the buffer metrics for the current step.""" + def write_jsonl(lst: list, filepath: Path) -> None: + with open(filepath, "w") as f: + for item in lst: + f.write(json.dumps(item, default=make_serializable) + "\n") - metrics = {} + all_easy = [ex for eb in self.env_buffers.values() for ex in eb.easy_examples] + all_hard = [ex for eb in self.env_buffers.values() for ex in eb.hard_examples] + write_jsonl(all_easy, path / "easy_examples.jsonl") + write_jsonl(all_hard, path / "hard_examples.jsonl") + write_jsonl(self.rollout_buffer, path / "rollout_buffer.jsonl") - # sum over envs (e.g. log globally) - num_examples_per_step_per_pool = { - pool: sum(self.num_examples_per_step[env][pool] for env in self.env_names) for pool in self.POOLS - } - num_rollouts_per_step_per_pool = { - pool: sum(self.num_rollouts_per_step[env][pool] for env in self.env_names) for pool in self.POOLS - } - num_examples_per_step = sum(num_examples_per_step_per_pool.values()) - num_rollouts_per_step = sum(num_rollouts_per_step_per_pool.values()) + def load(self, path: Path) -> None: + """Loads pool assignments and rollouts from checkpoint.""" - for pool in ["easy", "hard"]: - if num_examples_per_step: - metrics[f"evicted_examples/{pool}"] = num_examples_per_step_per_pool[pool] / num_examples_per_step - if num_rollouts_per_step: - metrics[f"filtered_rollouts/{pool}"] = num_rollouts_per_step_per_pool[pool] / num_rollouts_per_step + def read_jsonl(filepath: Path) -> list[dict]: + with open(filepath, "r") as f: + return [json.loads(line) for line in f] - total_normal = sum(len(self.example_buffer[env]) for env in self.env_names) - pool_counts = [len(self.easy_examples), total_normal, len(self.hard_examples)] - pool_ratios = mean_normalize(pool_counts) - for pool, pool_ratio in zip(self.POOLS, pool_ratios): - metrics[f"pool/{pool}"] = pool_ratio + saved_easy = read_jsonl(path / "easy_examples.jsonl") + saved_hard = read_jsonl(path / "hard_examples.jsonl") + saved_rollouts = cast(list[vf.RolloutOutput], read_jsonl(path / "rollout_buffer.jsonl")) - self.reset_step_metrics() + if not any(saved_easy) and not any(saved_hard) and not any(saved_rollouts): + self.logger.debug("No easy/ hard examples or rollouts found in checkpoint") + return + + # Build hash lookup across all env buffers: env -> (hash -> example_id) + hash_lookup: dict[str, dict[str, int]] = defaultdict(dict) + all_hashes: set[str] = set() + for env_name, eb in self.env_buffers.items(): + for example_id, example in eb.examples.items(): + h = eb.get_example_hash(example) + if h in all_hashes: + self.logger.warning( + f"Duplicate example hash found based on hash_keys={self.config.hash_keys}. " + "Overwriting with latest example. This may cause unexpected behavior when resuming the buffer." + ) + hash_lookup[env_name][h] = example_id + all_hashes.add(h) + + def move_saved_pool(saved_examples: list[dict], pool_name: str) -> int: + num_moved = 0 + for example in saved_examples: + # Use any env buffer to compute hash (hash_keys are config-level) + first_eb = next(iter(self.env_buffers.values())) + h = first_eb.get_example_hash(example) + for env_name, env_hashes in hash_lookup.items(): + if h in env_hashes: + example_id = env_hashes[h] + eb = self.env_buffers[env_name] + matched = eb.examples.pop(example_id, None) + if matched is not None: + target = eb.easy_examples if pool_name == "easy" else eb.hard_examples + target.append(matched) + num_moved += 1 + break + return num_moved + + if any(saved_easy): + num_moved = move_saved_pool(saved_easy, "easy") + self.logger.debug(f"Loaded {num_moved}/{len(saved_easy)} example(s) to easy pool from checkpoint.") + if num_moved != len(saved_easy): + self.logger.warning( + f"Could not move {len(saved_easy) - num_moved} example(s) from checkpoint to easy pool. " + "This usually means you resumed with an env mix that does not contain all previous examples." + ) + + if any(saved_hard): + num_moved = move_saved_pool(saved_hard, "hard") + self.logger.debug(f"Moved {num_moved}/{len(saved_hard)} example(s) to hard pool from checkpoint.") + if num_moved != len(saved_hard): + self.logger.warning( + f"Could not move {len(saved_hard) - num_moved} example(s) from checkpoint to hard pool. " + "This usually means you resumed with an env mix that does not contain all previous examples." + ) + + if any(saved_rollouts): + valid = [r for r in saved_rollouts if r.get("env_name") in self.env_names] + self.rollout_buffer.extend(valid) + self.logger.debug(f"Loaded {len(valid)} rollout(s) from checkpoint.") + + def convert_to_normal(eb: _EnvBuffer, pool: list[dict], fraction: float) -> int: + if fraction <= 0.0 or not pool: + return 0 + num_to_move = round(len(pool) * fraction) + if num_to_move <= 0: + return 0 + for _ in range(num_to_move): + example = random.choice(pool) + pool.remove(example) + eb.examples[example["example_id"]] = example + return num_to_move + + for eb in self.env_buffers.values(): + n_easy = len(eb.easy_examples) + moved = convert_to_normal(eb, eb.easy_examples, self.config.easy_fraction) + self.logger.debug(f"Converted {moved}/{n_easy} example(s) back to normal from easy pool ({eb.env_name}).") + n_hard = len(eb.hard_examples) + moved = convert_to_normal(eb, eb.hard_examples, self.config.hard_fraction) + self.logger.debug(f"Converted {moved}/{n_hard} example(s) back to normal from hard pool ({eb.env_name}).") + + def get_metrics(self) -> dict[str, float]: + metrics = {} + + # Aggregate cross-env totals + total_examples_per_pool = {p: 0 for p in POOLS} + total_rollouts_per_pool = {p: 0 for p in POOLS} + for eb in self.env_buffers.values(): + for p in POOLS: + total_examples_per_pool[p] += eb.num_examples_per_step[p] + total_rollouts_per_pool[p] += eb.num_rollouts_per_step[p] + + total_examples = sum(total_examples_per_pool.values()) + total_rollouts = sum(total_rollouts_per_pool.values()) + + for pool in ["easy", "hard"]: + if total_examples: + metrics[f"evicted_examples/{pool}"] = total_examples_per_pool[pool] / total_examples + if total_rollouts: + metrics[f"filtered_rollouts/{pool}"] = total_rollouts_per_pool[pool] / total_rollouts + + total_normal = sum(eb.num_normal for eb in self.env_buffers.values()) + total_easy = sum(len(eb.easy_examples) for eb in self.env_buffers.values()) + total_hard = sum(len(eb.hard_examples) for eb in self.env_buffers.values()) + pool_ratios = mean_normalize([total_easy, total_normal, total_hard]) + for pool, ratio in zip(POOLS, pool_ratios): + metrics[f"pool/{pool}"] = ratio + + # Per-env metrics + for eb in self.env_buffers.values(): + metrics.update(eb.get_metrics()) return metrics diff --git a/src/prime_rl/orchestrator/ckpt.py b/src/prime_rl/orchestrator/ckpt.py index c5a32caea6c..19277e30632 100644 --- a/src/prime_rl/orchestrator/ckpt.py +++ b/src/prime_rl/orchestrator/ckpt.py @@ -4,8 +4,8 @@ import torch +from prime_rl.configs.orchestrator import CheckpointConfig from prime_rl.orchestrator.buffer import Buffer -from prime_rl.orchestrator.config import CheckpointConfig from prime_rl.utils.logger import get_logger from prime_rl.utils.utils import get_ckpt_dir, get_step_path @@ -29,7 +29,7 @@ def __init__(self, output_dir: Path, config: CheckpointConfig): def get_ckpt_path(self, step: int) -> Path: return get_step_path(self.ckpt_dir, step) / "orchestrator" - def _save_to_path( + def save_to_path( self, ckpt_path: Path, progress: Progress, @@ -47,7 +47,7 @@ def _save_to_path( self.logger.debug(f"Orchestrator checkpoint saved in {time.perf_counter() - start_time:.2f} seconds") - def _load_from_path(self, ckpt_path: Path, progress: Progress, buffer: Buffer) -> None: + def load_from_path(self, ckpt_path: Path, progress: Progress, buffer: Buffer) -> None: """Loads a checkpoint from a given path in-place.""" self.logger.debug(f"Loading checkpoint from {ckpt_path}") start_time = time.perf_counter() @@ -76,7 +76,7 @@ def load(self, progress: Progress, buffer: Buffer, step: int) -> None: ckpt_path = self.get_ckpt_path(step) if not ckpt_path.exists(): raise FileNotFoundError(f"Checkpoint not found at {ckpt_path}") - self._load_from_path(ckpt_path, progress, buffer) + self.load_from_path(ckpt_path, progress, buffer) def save( self, @@ -87,7 +87,7 @@ def save( """Saves the full checkpoint state for a specified step.""" ckpt_path = self.get_ckpt_path(step) ckpt_path.mkdir(parents=True, exist_ok=True) - self._save_to_path(ckpt_path, progress, buffer) + self.save_to_path(ckpt_path, progress, buffer) def setup_ckpt_manager(output_dir: Path, config: CheckpointConfig | None) -> CheckpointManager | None: diff --git a/src/prime_rl/orchestrator/config.py b/src/prime_rl/orchestrator/config.py deleted file mode 100644 index dd71850b971..00000000000 --- a/src/prime_rl/orchestrator/config.py +++ /dev/null @@ -1,809 +0,0 @@ -from pathlib import Path -from typing import Annotated, Any, Literal, TypeAlias - -from pydantic import AliasChoices, BaseModel, Field, model_validator - -from prime_rl.transport.config import FileSystemTransportConfig, TransportConfigType -from prime_rl.utils.config import ( - ClientConfig, - HeartbeatConfig, - LogConfig, - PrimeMonitorConfig, - WandbWithExtrasConfig, -) -from prime_rl.utils.config import ( - ModelConfig as BaseModelConfig, -) -from prime_rl.utils.pydantic_config import BaseConfig, BaseSettings - - -class OptimizerConfig(BaseConfig): - """Per-run optimizer configuration for multi-run training.""" - - lr: Annotated[ - float, - Field( - ge=0, - description="Learning rate for this run.", - ), - ] = 1e-4 - - -class LoRAConfig(BaseConfig): - """Per-run LoRA configuration for multi-run training.""" - - name: Annotated[ - str | None, - Field( - description="Name of the LoRA adapter. If None, auto-generated from rank and alpha.", - ), - ] = None - - rank: Annotated[ - int | None, - Field( - ge=1, - description="LoRA rank for this run. Must be <= trainer's max rank. If None, uses trainer's rank.", - ), - ] = None - - alpha: Annotated[ - float, - Field( - ge=0, - description="LoRA alpha for this run.", - ), - ] = 32.0 - - -class ModelConfig(BaseModelConfig): - """Extended model configuration with per-run LoRA settings.""" - - lora: Annotated[ - LoRAConfig | None, - Field( - description="LoRA configuration. If None, LoRA is not used.", - ), - ] = None - - -class SamplingConfig(BaseConfig): - """Configures how tokens are sampled from the model for training. Largely follows the vLLM sampling parameters.""" - - temperature: Annotated[ - float, - Field( - ge=0, - description="Scales the output probability distribution. Lower values => more deterministic, higher values => more random. If 0, will sample greedily.", - ), - ] = 1.0 - - repetition_penalty: Annotated[ - float, - Field( - ge=0, - description="Penalty for repeating tokens. Values > 1.0 discourage repetition, values < 1.0 encourage repetition, and 1.0 means no penalty.", - ), - ] = 1.0 - - max_tokens: Annotated[ - int | None, - Field( - description="Maximum number of output tokens to generate per turn. If None, will generate until maximum context length or EOS token is hit.", - ), - ] = None - - min_tokens: Annotated[ - int, - Field( - ge=0, - description="Minimum number of output tokens to generate per sequence.", - ), - ] = 0 - - seed: Annotated[ - int | None, - Field( - description="Random seed to use for sampling. If None, no seeding is used.", - ), - ] = None - - # Strictly speaking, extra_body is not a sampling parameter, but it is the - # easiest way to pass arbitrary extra parameters to the server via verifiers - extra_body: Annotated[ - dict[str, Any], - Field( - description="Extra body to pass with each request to the inference server. By default, it is set to an empty dictionary.", - ), - ] = {} - - -class EvalSamplingConfig(BaseConfig): - """Configures how tokens are sampled from the model for evaluation. Largely follows the vLLM sampling parameters.""" - - temperature: Annotated[ - float | None, - Field( - ge=0, - description="Scales the output probability distribution. Lower values => more deterministic, higher values => more random. If 0, will sample greedily. Defaults to None, which means we fall back to the inference server's default value.", - ), - ] = None - - repetition_penalty: Annotated[ - float | None, - Field( - ge=0, - description="Penalty for repeating tokens. Values > 1.0 discourage repetition, values < 1.0 encourage repetition, and 1.0 means no penalty. Defaults to None, which means we fall back to the inference server's default value.", - ), - ] = None - - top_p: Annotated[ - float | None, - Field( - description="Cumulative probability of the top tokens to consider. If 1, all tokens are considered. Defaults to None, which means we fall back to the inference server's default value.", - ), - ] = None - - top_k: Annotated[ - int | None, - Field( - description="Number of top tokens to consider. If -1, all tokens are considered. Defaults to None, which means we fall back to the inference server's default value.", - ), - ] = None - - min_p: Annotated[ - float | None, - Field( - description="Minimum probability for a token to be considered, relative to the probability of the most likely token. If 0, all tokens are considered. Defaults to None, which means we fall back to the inference server's default value.", - ), - ] = None - - max_tokens: Annotated[ - int | None, - Field( - description="Maximum number of output tokens to generate per turn. If None, will generate until maximum context length or EOS token is hit.", - ), - ] = None - - min_tokens: Annotated[ - int | None, - Field( - description="Minimum number of output tokens to generate per sequence. Defaults to None, which means we fall back to the inference server's default value.", - ), - ] = None - - reasoning_effort: Annotated[ - Literal["minimal", "low", "medium", "high"] | None, - Field( - description="Constrains effort on reasoning for reasoning models. Currently supported values are minimal, low, medium, and high. Defaults to None, which means we fall back to the inference server's default value.", - ), - ] = None - - seed: Annotated[ - int | None, - Field( - description="Random seed to use for sampling. If None, no seeding is used. Defaults to None, which means we fall back to the inference server's default value.", - ), - ] = None - - # Strictly speaking, extra_body is not a sampling parameter, but it is the - # easiest way to pass arbitrary extra parameters to the server via verifiers - extra_body: Annotated[ - dict[str, Any], - Field( - description="Extra body to use for the OpenAI API. By default, it is set to an empty dictionary.", - ), - ] = {} - - -class EvalSaveDiskConfig(BaseConfig): - """Configures how to save the eval results to disk.""" - - path: Annotated[ - Path | None, - Field( - description="The path to save the eval results to. If None, will default to /evals// for online evals and the verifiers default for offline evals." - ), - ] = None - - -class EvalSaveHFConfig(BaseConfig): - """Configures how to save the eval results to HF.""" - - dataset_name: Annotated[ - str | None, - Field( - description="The name of the HF dataset to save the eval results to. If None, will auto-generate a name." - ), - ] = None - - dataset_subset: Annotated[ - str | None, - Field( - description="The subset name of the HF dataset to save the evaluation results. If None, will default to the environment ID.", - ), - ] = None - - dataset_split: Annotated[ - str | None, - Field( - description="The split name of the HF dataset to save the evaluation results. If None, will default to 'evals'.", - ), - ] = None - - private: Annotated[ - bool, - Field(description="Whether to save the eval results to a private HF dataset."), - ] = False - - -class EvalSaveConfig(BaseConfig): - disk: EvalSaveDiskConfig | None = None - hf: EvalSaveHFConfig | None = None - env_hub: Annotated[ - bool, - Field( - description="Whether to push eval results to Prime Environment Hub. Automatically pushes all evaluated environments. Requires PRIME_API_KEY and authorization for the environments." - ), - ] = False - stream: Annotated[ - bool, - Field( - description="Whether to save results incrementally as rollouts complete.", - ), - ] = False - - -class RetryConfig(BaseConfig): - """Configures retry behavior for rollout generation.""" - - max_attempts: Annotated[ - int, - Field( - ge=1, - description="Maximum number of retry attempts.", - ), - ] = 10 - - wait_multiplier: Annotated[ - float, - Field( - ge=0, - description="Multiplier for exponential backoff wait time.", - ), - ] = 1.0 - - wait_min: Annotated[ - float, - Field( - ge=0, - description="Minimum wait time in seconds between retries.", - ), - ] = 1.0 - - wait_max: Annotated[ - float, - Field( - ge=0, - description="Maximum wait time in seconds between retries.", - ), - ] = 60.0 - - reraise: Annotated[ - bool, - Field( - description="Whether to reraise the original exception after all retries are exhausted. If False, raises tenacity.RetryError instead.", - ), - ] = False - - -class EnvConfig(BaseConfig): - """Configures an environment for training.""" - - id: Annotated[str, Field(description="ID of the environment to use.")] = "reverse-text" - args: Annotated[dict, Field(description="Arguments to pass to the environment.")] = {} - name: Annotated[str | None, Field(description="Name of the environment to use.")] = None - - -class EvalEnvConfig(EnvConfig): - """Configures an environment for evaluation.""" - - num_examples: Annotated[ - int | None, - Field( - description="Number of examples to evaluate per environment. If not set, will use 'num_examples' from main config." - ), - ] = None - rollouts_per_example: Annotated[ - int | None, - Field( - description="Number of samples to generate per example for each environment. If not set, will use 'rollouts_per_example' from main config." - ), - ] = None - - skip_first: Annotated[ - int, - Field( - description="Number of examples to skip from the beginning of the dataset.", - ), - ] = 0 - - -class ValConfig(BaseConfig): - """Configures the validation of the model.""" - - num_examples: Annotated[ - int, Field(ge=1, description="Number of examples to use for validation. If -1, will use all examples.") - ] = 16 - rollouts_per_example: Annotated[ - int, Field(ge=1, description="Number of samples to generate per example for validation.") - ] = 1 - interval: Annotated[int, Field(description="Interval at which to validate the model.")] = 10 - - -class EvalConfig(BaseConfig): - """Configures evaluation using verifiers environments.""" - - env: list[EvalEnvConfig] = [EvalEnvConfig()] - sampling: EvalSamplingConfig = Field( - default_factory=EvalSamplingConfig, - description="Shared sampling configuration for evals; can differ from training sampling.", - ) - save: EvalSaveConfig = Field( - default_factory=EvalSaveConfig, - description="Configures how to save the eval results.", - ) - retry: RetryConfig = Field( - default_factory=RetryConfig, - description="Configures retry behavior for rollout generation.", - ) - num_examples: Annotated[int, Field(description="Number of examples to evaluate per environment.")] = -1 - rollouts_per_example: Annotated[ - int, Field(ge=1, description="Number of samples to generate per example for each environment.") - ] = 1 - reasoning_field: Annotated[ - str, - Field( - description="The field in the raw model response that contains the reasoning content. Defaults to 'reasoning_content', which is the default for vLLM when serving a model with a reasoning parser. Other APIs (e.g. DeepSeek, GLM, etc.) may use different field names.", - ), - ] = "reasoning_content" - per_rollout: Annotated[ - bool, - Field( - description="Schedule rollouts individually instead of as groups. Enables live progress updates and per-rollout resume, but incompatible with group-based rubrics.", - ), - ] = False - - -class OnlineEvalConfig(EvalConfig): - """Configures online evaluation.""" - - interval: Annotated[ - int, - Field( - ge=1, - description="Interval at which to evaluate the model.", - ), - ] = 100 - - eval_base_model: Annotated[ - bool, - Field( - description="Whether to evaluate the base model we are training on.", - ), - ] = True - - skip_eval_on_resume: Annotated[ - bool, - Field( - validation_alias=AliasChoices("skip_eval_on_resume", "skip_eval_on_restart"), - description=( - "If True and resuming the orchestrator from a checkpoint, skip the (potentially redundant) " - "online eval that would otherwise run immediately at the resumed checkpoint step." - ), - ), - ] = True - - -class CheckpointConfig(BaseConfig): - """Configures checkpointing the orchestrator.""" - - interval: Annotated[int | None, Field(ge=1, description="Interval at which to save the checkpoint.")] = None - - resume_step: Annotated[ - int | None, - Field( - ge=-1, - description="Step to resume orchestrator from. If None, will start from scratch. If -1, will restart from latest checkpoint available.", - ), - ] = None - - keep_last: Annotated[ - int | None, - Field( - ge=1, - description="Keep at most this many recent step checkpoints on disk. If None, never clean old checkpoints based on recency.", - ), - ] = None - - keep_interval: Annotated[ - int | None, - Field( - ge=1, - description="Keep checkpoints at every N steps permanently (e.g., keep_interval=100 keeps step 100, 200, ...). If None, no interval-based keeping.", - ), - ] = None - - skip_progress: Annotated[ - bool, - Field( - description="Whether to skip loading the progress from checkpoint.", - ), - ] = False - - skip_buffer: Annotated[ - bool, - Field( - description="Whether to skip loading the buffer from checkpoint.", - ), - ] = False - - -class BufferConfig(BaseConfig): - """Configures the buffer for the orchestrator.""" - - seed: Annotated[ - int | None, - Field( - description="Random seed to use for the buffer. If set, the sampling from the buffer will be deterministic.", - ), - ] = None - - env_ratios: Annotated[ - list[float] | None, - Field( - description=( - "Ratios for sampling from each environment. " - "If None, samples uniformly across all available problems (not environments)." - ), - ), - ] = None - - easy_threshold: Annotated[ - float | None, - Field( - description="Threshold for easy difficulty classification. If average reward >= this threshold, mark as easy.", - ), - ] = None - - hard_threshold: Annotated[ - float | None, - Field( - description="Threshold for hard difficulty classification. If average reward <= this threshold, mark as hard.", - ), - ] = None - - easy_fraction: Annotated[ - float, - Field( - ge=0, - le=1, - description="Fraction of easy problems to convert to normal when resuming or starting training. Only problems with difficulty 'normal' are sampled.", - ), - ] = 0.0 - - hard_fraction: Annotated[ - float, - Field( - ge=0, - le=1, - description="Fraction of hard problems to convert to normal when resuming or starting training. Only problems with difficulty 'normal' are sampled.", - ), - ] = 0.0 - - online_difficulty_filtering: Annotated[ - bool, - Field( - description="Whether to filter rollouts based on difficulty. If True, rollouts with average reward 0.0 or 1.0 are not added to the buffer.", - ), - ] = False - - hash_keys: Annotated[ - list[str], - Field( - min_length=1, - description="Keys to use for computing example hashes. Will be used to match examples from buffer checkpoints and determine buffer resume behavior.", - ), - ] = ["task", "prompt"] - - skip_verification: Annotated[ - bool, - Field( - description=( - "Whether to skip verification of rollouts using the environment's rubric. " - "If True, rewards are always set to 0, online_difficulty_filtering is disabled, " - "and easy/hard thresholds are not used." - ), - ), - ] = False - - @model_validator(mode="after") - def validate_thresholds(self): - if self.easy_threshold is not None and self.hard_threshold is not None: - assert self.easy_threshold > self.hard_threshold, "easy_threshold must be greater than hard_threshold." - return self - - @model_validator(mode="after") - def validate_env_ratios(self): - if self.env_ratios is not None: - assert all(ratio > 0 for ratio in self.env_ratios), "All env_ratios must be positive." - return self - - @model_validator(mode="after") - def validate_skip_verification(self): - """Validate that skip_verification is not used with reward-dependent features.""" - if self.skip_verification: - if self.online_difficulty_filtering: - raise ValueError( - "skip_verification cannot be True when online_difficulty_filtering is True. " - "These features depend on rewards which are disabled when skip_verification=True." - ) - if self.easy_threshold is not None: - raise ValueError( - "skip_verification cannot be True when easy_threshold is set. " - "Easy threshold depends on rewards which are disabled when skip_verification=True." - ) - if self.hard_threshold is not None: - raise ValueError( - "skip_verification cannot be True when hard_threshold is set. " - "Hard threshold depends on rewards which are disabled when skip_verification=True." - ) - return self - - -class AdvantageConfig(BaseConfig): - length_weighted_mean: bool = False - - -class FileSystemWeightBroadcastConfig(BaseModel): - """Configures the filesystem weight broadcast.""" - - type: Literal["filesystem"] = "filesystem" - - -class NCCLWeightBroadcastConfig(BaseModel): - """Configures the NCCL weight broadcast.""" - - type: Literal["nccl"] = "nccl" - - host: Annotated[str, Field(description="The host to use for the NCCL broadcast.")] = "localhost" - port: Annotated[int, Field(description="The port to use for the NCCL broadcast.")] = 29501 - timeout: Annotated[int, Field(description="The timeout in seconds to use for the NCCL broadcast.")] = 1200 - - -WeightBroadcastConfigType: TypeAlias = FileSystemWeightBroadcastConfig | NCCLWeightBroadcastConfig - - -class TeacherModelConfig(BaseConfig): - """Configures the teacher model for computing teacher logprobs (e.g. for distillation).""" - - client: Annotated[ - ClientConfig, - Field(description="The OAI client configuration for the teacher model."), - ] = ClientConfig() - - model: Annotated[ - ModelConfig, - Field(description="The model configuration for the teacher model."), - ] = ModelConfig() - - -class OrchestratorConfig(BaseSettings): - """Configures the orchestrator for RL training.""" - - # The OAI client configuration - client: ClientConfig = ClientConfig() - - # The model configuration - model: ModelConfig = ModelConfig() - - # The optimizer configuration (per-run LR for multi-run training) - optim: OptimizerConfig = OptimizerConfig() - - # The teacher model configuration (optional) - teacher_model: Annotated[ - TeacherModelConfig | None, - Field( - description="The teacher model configuration for computing teacher logprobs (e.g. for distillation). " - "If provided, teacher logprobs will be computed using the specified model. " - "If None, no teacher model will be used." - ), - ] = None - - # The sampling configuration - sampling: SamplingConfig = SamplingConfig() - - # The environment configuration - env: list[EnvConfig] = [EnvConfig()] - - # The evaluation configuration - eval: OnlineEvalConfig | None = None - - # Data buffer configuration - buffer: BufferConfig = BufferConfig() - - # The advantage configuration - advantage: AdvantageConfig | None = AdvantageConfig() - - # The logging configuration - log: LogConfig = LogConfig() - - # The wandb configuration - wandb: WandbWithExtrasConfig | None = None - - # The prime monitor configuration - prime_monitor: PrimeMonitorConfig | None = None - - # The checkpoint configuration - ckpt: CheckpointConfig | None = None - - # The validation configuration - val: ValConfig | None = None - - weight_broadcast: Annotated[WeightBroadcastConfigType, Field(discriminator="type")] = ( - FileSystemWeightBroadcastConfig() - ) - - rollout_transport: Annotated[TransportConfigType, Field(discriminator="type")] = FileSystemTransportConfig() - - trajectory_strategy: Annotated[ - Literal["interleaved", "branching"], - Field( - description="Strategy to use for building training examples from multi-turn rollouts. If interleaved, will try to concatenate consecutive trajectory steps into a single training example. If branching, will create a separate training example for each trajectory step." - ), - ] = "interleaved" - - output_dir: Annotated[ - Path, - Field( - description="Directory to write outputs to. Will be populated with checkpoints, weights, rollouts and logs as subdirectories. Should be set to a persistent directory with enough disk space. This value should be distinct across experiments running on a single node. See the README for more details." - ), - ] = Path("outputs/run_default") - - max_concurrent: Annotated[ - int | None, - Field( - description="Maximum number of concurrent rollouts to generate and score. Will create a global semaphore and pass to verifiers Environment. If None, will not limit concurrency.", - ), - ] = None - - workers_per_env: Annotated[ - int, - Field( - ge=1, - description="Number of worker subprocesses to spawn per environment. Multiple workers enable isolation of event loop lag - if one worker slows down, others continue at full speed. Uses least-pending routing to distribute load.", - ), - ] = 1 - - max_env_worker_restarts: Annotated[ - int, - Field( - ge=-1, - description="Maximum number of automatic restarts for an environment worker that dies unexpectedly. After this limit, the orchestrator will crash. Set to -1 for unlimited restarts.", - ), - ] = 5 - - batch_size: Annotated[int, Field(ge=1, description="Number of samples to train on per step.")] = 128 - - oversampling_factor: Annotated[ - float, - Field( - ge=1, - description="Factor by which to oversample the batch. Will lead to more in-flight group rollout requests at the same time.", - ), - ] = 1.0 - - rollouts_per_example: Annotated[ - int, - Field( - ge=1, - description="Number of output sequences to return per example during training.", - ), - ] = 1 - - seq_len: Annotated[ - int, - Field( - description="Sequence length to use for training. If a sample is shorter than this, it will be padded. If a sequence is longer than this, it will be truncated.", - ), - ] = 2048 - - mask_env_responses: Annotated[ - bool, - Field( - description="Whether to mask environment responses from the loss.", - ), - ] = True - - # TODO(Mika): This should be automatic from the number of ZMQ connections - num_train_workers: Annotated[ - int, - Field(default=1, ge=1, description="Number of training workers to use for training."), - ] = 1 - - max_steps: Annotated[ - int | None, - Field( - description="Maximum number of training steps to run. If None, will run indefinitely.", - ), - ] = None - - max_off_policy_steps: Annotated[ - int, - Field( - ge=0, - description="Maximum number of policies that are allowed to generate a single rollout. Rollouts that are generated from more than `max_off_policy_steps` steps ahead of training will be discarded. Higher values yield better throughput, but lead to more off-policyness in training.", - ), - ] = 8 - - max_async_level: Annotated[ - int, - Field( - ge=0, - description="Maximum number of steps the inference can be ahead of training. If 0, will degenerate to synchronous on-policy RL. If >=1, training and inference will be overlapped.", - ), - ] = 1 - - strict_async_level: Annotated[ - bool, - Field( - description="Whether to strictly enforce the max async level. If True, will always ensure that the policy used for generating rollouts is exactly `max_async_level` steps ahead of training. If False, any policy that is at most `max_async_level` steps ahead of training is allowed, i.e. we always use the latest available policy.", - ), - ] = False - - bench: Annotated[ - bool, - Field( - description="Whether to run in benchmark mode. It will automatically set the maximum number of steps to run to 5, max async level to ~infinity and disable W&B.", - ), - ] = False - - seed: Annotated[int | None, Field(description="Random seed for the orchestrator.")] = 42 - - heartbeat: Annotated[ - HeartbeatConfig | None, Field(description="The heartbeat config for monitoring training progress.") - ] = None - - @model_validator(mode="after") - def nccl_max_async_level(self): - if self.weight_broadcast.type == "nccl": - if not self.max_async_level == 1: - raise ValueError("max_async_level must be 1 for NCCL broadcast") - return self - - @model_validator(mode="after") - def validate_batch_size(self): - if self.batch_size % self.rollouts_per_example != 0: - raise ValueError("Batch size must be divisible by the number of samples per problem") - return self - - @model_validator(mode="after") - def validate_env_ratios(self): - if self.buffer.env_ratios is not None: - assert len(self.buffer.env_ratios) == len(self.env), "env_ratios length must match number of environments" - return self - - @model_validator(mode="after") - def auto_setup_bench(self): - if self.bench: - self.max_steps = 4 # Run for 1 warmup step + 3 evaluation steps - self.max_async_level = int(1e9) # Never wait for RL weight checkpoints - - # Disable evaluation - self.eval = None - if self.wandb: - self.wandb.log_extras = None - if self.prime_monitor: - self.prime_monitor.log_extras = None - - return self diff --git a/src/prime_rl/orchestrator/env_server/env_server.py b/src/prime_rl/orchestrator/env_server/env_server.py new file mode 100644 index 00000000000..8ff057b9d97 --- /dev/null +++ b/src/prime_rl/orchestrator/env_server/env_server.py @@ -0,0 +1,45 @@ +import asyncio + +from verifiers.serve import ZMQEnvServer + +from prime_rl.configs.env_server import EnvServerConfig +from prime_rl.utils.config import cli +from prime_rl.utils.logger import setup_logger +from prime_rl.utils.pathing import get_log_dir +from prime_rl.utils.process import set_proc_title +from prime_rl.utils.utils import clean_exit, get_env_ids_to_install, install_env + + +@clean_exit +def run_server(config: EnvServerConfig): + setup_logger(config.log.level, json_logging=config.log.json_logging) + + # install environment if not already installed + env_ids_to_install = set() + env_ids_to_install.update(get_env_ids_to_install([config.env])) + for env_id in env_ids_to_install: + install_env(env_id, prerelease=config.env_install_prerelease) + + log_dir = (get_log_dir(config.output_dir) / config.env.resolved_name).as_posix() + + server = ZMQEnvServer( + env_id=config.env.stripped_id, + env_args=config.env.args, + extra_env_kwargs=config.env.extra_env_kwargs, + log_level=config.log.level, + log_dir=log_dir, + json_logging=config.log.json_logging, + num_workers=config.env.num_workers, + **{"address": config.env.address} if config.env.address is not None else {}, + ) + asyncio.run(server.run()) + + +def main(): + """Main entry-point for env-server. Run using `uv run env-server`""" + set_proc_title("EnvServer") + run_server(cli(EnvServerConfig)) + + +if __name__ == "__main__": + main() diff --git a/src/prime_rl/orchestrator/env_worker.py b/src/prime_rl/orchestrator/env_worker.py deleted file mode 100644 index 963ce54f2a6..00000000000 --- a/src/prime_rl/orchestrator/env_worker.py +++ /dev/null @@ -1,431 +0,0 @@ -""" -Environment worker subprocess. - -Runs environment rollouts in a separate process to isolate event loop lag. -""" - -import asyncio -import queue -import uuid -from dataclasses import dataclass -from itertools import cycle -from multiprocessing import Process, Queue -from pathlib import Path - -import verifiers as vf -from openai import AsyncOpenAI - -from prime_rl.utils.client import setup_clients -from prime_rl.utils.config import ClientConfig -from prime_rl.utils.logger import get_logger, intercept_verifiers_logging, reset_logger, setup_logger - - -class WorkerDiedError(Exception): - """Raised when a worker subprocess dies unexpectedly.""" - - pass - - -@dataclass -class RolloutRequest: - """Request to generate rollouts for an example.""" - - request_id: str - example_id: int - rollouts_per_example: int - model_name: str # Model name to use for this request (may change for LoRA) - - -@dataclass -class RolloutResponse: - """Response containing rollout results.""" - - request_id: str - results: list[dict] # Simplified state dicts - lag_metrics: dict | None = None # Event loop lag metrics from worker - - -def extract_result(state: vf.State) -> dict: - """Extract only the fields needed from vf.State for IPC. - - The extracted dict must contain all fields needed by: - - Buffer.update(): example_id, task, reward - - orchestrator metrics: reward, is_truncated, error, timing, metrics, trajectory - - interleave_rollout/branch_rollout: trajectory[*]["tokens"] with all token fields - """ - # Get trajectory with tokens (needed for training) - trajectory = [] - for step in state.get("trajectory", []): - traj_step = { - "prompt": step.get("prompt"), - "completion": step.get("completion"), - # tokens dict contains: prompt_ids, prompt_mask, completion_ids, - # completion_mask, completion_logprobs, is_truncated - "tokens": step.get("tokens"), - } - trajectory.append(traj_step) - - return { - # Required by buffer - "example_id": state.get("example_id"), - "task": state.get("task"), - "reward": state.get("reward"), - # Required by orchestrator metrics - "is_truncated": state.get("is_truncated", False), - "error": type(state["error"]).__name__ if state.get("error") else None, - "timing": dict(state.get("timing", {})), - "metrics": state.get("metrics", {}), - # Required for training examples - "prompt": state.get("prompt"), - "completion": state.get("completion"), - "trajectory": trajectory, - } - - -async def process_request( - request: RolloutRequest, - env: vf.Environment, - client_cycle: cycle, - semaphore: asyncio.Semaphore, - example_lookup: dict[int, dict], - sampling_args: dict, -) -> RolloutResponse: - """Process a single rollout request.""" - client = next(client_cycle) - example = example_lookup[request.example_id] - group_inputs = [vf.RolloutInput(**example) for _ in range(request.rollouts_per_example)] - - states = await env.run_group( - group_inputs=group_inputs, - client=client, - model=request.model_name, - gen_sampling_args=sampling_args, - gen_sem=semaphore, - score_sem=semaphore, - ) - - results = [extract_result(state) for state in states] - return RolloutResponse(request_id=request.request_id, results=results) - - -async def worker_loop( - request_queue: Queue, - response_queue: Queue, - env: vf.Environment, - clients: list[AsyncOpenAI], - max_concurrent: int, - env_id: str, - example_lookup: dict[int, dict], - sampling_args: dict, -): - """Main async loop for processing requests.""" - from prime_rl.orchestrator.event_loop_lag import EventLoopLagMonitor - - client_cycle = cycle(clients) - semaphore = asyncio.Semaphore(max_concurrent) if max_concurrent > 0 else asyncio.Semaphore(10000) - - # Start event loop lag monitor for this worker - lag_monitor = EventLoopLagMonitor(interval=0.1) # More frequent sampling for workers - lag_monitor_task = asyncio.create_task(lag_monitor.run()) - - # Track in-flight tasks - pending_tasks: dict[asyncio.Task, str] = {} - - def check_for_requests(): - """Non-blocking check for new requests.""" - while True: - try: - request = request_queue.get_nowait() - except queue.Empty: - break - if request is None: # Shutdown signal - return False - task = asyncio.create_task( - process_request(request, env, client_cycle, semaphore, example_lookup, sampling_args) - ) - pending_tasks[task] = request.request_id - return True - - try: - while True: - # Check for new requests - if not check_for_requests(): - break - - if not pending_tasks: - # No pending tasks, wait a bit for new requests - await asyncio.sleep(0.01) - continue - - # Wait for at least one task to complete - done, _ = await asyncio.wait(pending_tasks.keys(), timeout=0.1, return_when=asyncio.FIRST_COMPLETED) - - for task in done: - pending_tasks.pop(task) - response = task.result() - # Attach lag metrics to response - response.lag_metrics = lag_monitor.get_metrics() - response_queue.put(response) - finally: - # Cleanup - lag_monitor_task.cancel() - for task in pending_tasks: - task.cancel() - - -def worker_main( - request_queue: Queue, - response_queue: Queue, - env_id: str, - env_args: dict, - client_config_dict: dict, - seq_len: int, - interleaved_rollouts: bool, - max_concurrent: int, - example_lookup: dict[int, dict], - sampling_args: dict, - log_level: str, - vf_log_level: str, - log_file: str | None, - worker_name: str | None = None, -): - """Main entry point for worker process.""" - # Reset logger inherited from parent process, then setup fresh logger for this worker - if log_file: - reset_logger() - setup_logger(log_level, log_file=Path(log_file), append=True, tag=worker_name) - intercept_verifiers_logging(level=vf_log_level) - - # Load environment - env = vf.load_environment(env_id, **env_args) - env.set_max_seq_len(seq_len) - env.set_interleaved_rollouts(interleaved_rollouts) - - # Create clients - client_config = ClientConfig(**client_config_dict) - clients = setup_clients(client_config) - - # Run async loop - asyncio.run( - worker_loop( - request_queue, - response_queue, - env, - clients, - max_concurrent, - env_id, - example_lookup, - sampling_args, - ) - ) - - -class EnvWorker: - """Manages a worker subprocess for an environment.""" - - def __init__( - self, - env_id: str, - env_args: dict, - client_config: ClientConfig, - model_name: str, - seq_len: int, - interleaved_rollouts: bool, - max_concurrent: int, - example_lookup: dict[int, dict], - sampling_args: dict, - worker_name: str | None = None, - log_level: str = "warn", - vf_log_level: str = "warn", - log_file: str | None = None, - max_restarts: int = 5, - ): - self.env_id = env_id - self.env_args = env_args - self.client_config = client_config - self.model_name = model_name - self.seq_len = seq_len - self.interleaved_rollouts = interleaved_rollouts - self.max_concurrent = max_concurrent - self.example_lookup = example_lookup - self.sampling_args = sampling_args - self.worker_name = worker_name or env_id - - self.log_level = log_level - self.vf_log_level = vf_log_level - self.log_file = log_file - self.max_restarts = max_restarts - - self.request_queue: Queue = Queue() - self.response_queue: Queue = Queue() - self.process: Process | None = None - - # Track pending requests for response matching - self.pending_futures: dict[str, asyncio.Future] = {} - - # Track latest lag metrics from this worker - self.latest_lag_metrics: dict = {} - - # Track intentional shutdown to avoid false error on clean stop - self._stopping = False - # Track if worker died unexpectedly (prevents scheduler from routing to dead worker) - self._dead = False - # Track restart count to prevent infinite restart loops - self._restart_count = 0 - # Track fatal error when max restarts exceeded (orchestrator should crash) - self._fatal_error: Exception | None = None - # Track successful responses since last restart (to reset restart count) - self._responses_since_restart = 0 - - def start(self): - """Start the worker process.""" - self.process = Process( - target=worker_main, - args=( - self.request_queue, - self.response_queue, - self.env_id, - self.env_args, - self.client_config.model_dump(), - self.seq_len, - self.interleaved_rollouts, - self.max_concurrent, - self.example_lookup, - self.sampling_args, - self.log_level, - self.vf_log_level, - self.log_file, - self.worker_name, - ), - daemon=True, - ) - self.process.start() - self._stopping = False # Reset after process is alive to avoid race condition - self._dead = False # Reset in case of restart - - def stop(self): - """Stop the worker process.""" - self._stopping = True - if self.process and self.process.is_alive(): - self.request_queue.put(None) # Shutdown signal - self.process.join(timeout=5) - if self.process.is_alive(): - self.process.terminate() - - def _restart(self): - """Restart the worker process after unexpected death.""" - # Clean up old process if it exists - if self.process is not None: - if self.process.is_alive(): - self.process.terminate() - # Always join to reap zombie process, even if already dead - self.process.join(timeout=5) - self.process.close() - - # Clear queues to avoid stale data (drain without blocking) - while True: - try: - self.request_queue.get_nowait() - except queue.Empty: - break - while True: - try: - self.response_queue.get_nowait() - except queue.Empty: - break - - # Start fresh process - self.start() - - async def submit_request( - self, - example_id: int, - rollouts_per_example: int, - ) -> tuple[asyncio.Future, str]: - """Submit a rollout request and return a (future, request_id) tuple.""" - request_id = uuid.uuid4().hex - request = RolloutRequest( - request_id=request_id, - example_id=example_id, - rollouts_per_example=rollouts_per_example, - model_name=self.model_name, - ) - - loop = asyncio.get_event_loop() - future = loop.create_future() - self.pending_futures[request_id] = future - - self.request_queue.put(request) - return future, request_id - - async def collect_responses(self): - """Background task to collect responses and resolve futures.""" - logger = get_logger() - while True: - # Drain queue first to salvage any responses before checking for dead worker - while True: - try: - response: RolloutResponse = self.response_queue.get_nowait() - except queue.Empty: - break - # Store latest lag metrics from worker - if response.lag_metrics: - self.latest_lag_metrics = response.lag_metrics - if response.request_id in self.pending_futures: - future = self.pending_futures.pop(response.request_id) - # Check if future was cancelled (e.g., by update_policy) - if not future.done(): - future.set_result(response.results) - # Track successful responses; reset restart count after stable operation - self._responses_since_restart += 1 - if self._responses_since_restart >= 10 and self._restart_count > 0: - logger.debug(f"Worker '{self.worker_name}' stable after {self._responses_since_restart} responses, resetting restart count") - self._restart_count = 0 - self._responses_since_restart = 0 - - # Check if worker process died unexpectedly (but not during intentional shutdown) - if self.process and not self.process.is_alive() and not self._stopping: - exit_code = self.process.exitcode - error = WorkerDiedError(f"Worker '{self.worker_name}' died unexpectedly (exit code: {exit_code})") - # Mark worker as dead so scheduler won't route new requests here - self._dead = True - # Fail remaining pending futures so callers don't hang indefinitely - for future in self.pending_futures.values(): - if not future.done(): - future.set_exception(error) - self.pending_futures.clear() - - # Check if we've exceeded max restarts (-1 means unlimited) - self._restart_count += 1 - if self.max_restarts >= 0 and self._restart_count > self.max_restarts: - logger.error( - f"Worker '{self.worker_name}' died {self._restart_count} times, exceeding max restarts ({self.max_restarts}). Giving up." - ) - # Store fatal error so orchestrator can detect and crash - self._fatal_error = error - raise error - - # Log warning and restart the worker automatically - restart_info = f"{self._restart_count}/{self.max_restarts}" if self.max_restarts >= 0 else f"{self._restart_count}" - logger.warning( - f"Worker '{self.worker_name}' died unexpectedly (exit code: {exit_code}). " - f"Restarting worker automatically ({restart_info}). In-flight requests will be rescheduled." - ) - self._responses_since_restart = 0 # Reset on restart - self._restart() - - await asyncio.sleep(0.01) - - def update_model_name(self, model_name: str): - """Update the model name for future requests.""" - self.model_name = model_name - - @property - def pending_count(self) -> int: - """Number of pending requests for this worker. - - Returns a large number if the worker is dead to prevent scheduler from selecting it. - """ - if self._dead: - return 999999 # Effectively infinite - scheduler will pick other workers - return len(self.pending_futures) diff --git a/src/prime_rl/orchestrator/envs.py b/src/prime_rl/orchestrator/envs.py index 2416a0c1f0b..1de52994a99 100644 --- a/src/prime_rl/orchestrator/envs.py +++ b/src/prime_rl/orchestrator/envs.py @@ -1,35 +1,398 @@ -from typing import TYPE_CHECKING, Any +from __future__ import annotations -from prime_rl.utils.envs import _ENV_PARSERS as _BASE_ENV_PARSERS, get_env_value, get_dir, set_defaults +import asyncio +import atexit +import multiprocessing as mp +import time +from collections.abc import Awaitable, Callable, Iterator, Sequence +from multiprocessing.process import BaseProcess +from pathlib import Path +from typing import Generic, TypeVar -if TYPE_CHECKING: - # Enable type checking for shared envs - # ruff: noqa - from prime_rl.utils.envs import * +import pandas as pd +import verifiers as vf +from verifiers.serve import ZMQEnvClient, ZMQEnvServer +from verifiers.utils.serve_utils import get_free_port - # vLLM - VLLM_CONFIGURE_LOGGING: int +from prime_rl.configs.orchestrator import EnvConfig, EvalEnvConfig, TrainEnvConfig +from prime_rl.orchestrator.eval_utils import compute_pass_at_k +from prime_rl.orchestrator.vf_utils import get_completion_len +from prime_rl.utils.logger import ProgressTracker, get_logger +from prime_rl.utils.monitor import get_monitor +from prime_rl.utils.utils import capitalize - # tqdm - TQDM_DISABLE: int +REQUIRED_STATE_COLUMNS = ["trajectory", "sampling_args"] -_ORCHESTRATOR_ENV_PARSERS = { - "VLLM_CONFIGURE_LOGGING": int, - "TQDM_DISABLE": int, - **_BASE_ENV_PARSERS, -} +class Env: + """Wraps a vf.Environment - only exposes features used in PRIME-RL.""" -_ORCHESTRATOR_ENV_DEFAULTS = { - "VLLM_CONFIGURE_LOGGING": "0", -} + def __init__(self, config: EnvConfig): + self.config = config + self.sampling_args: dict = {} -set_defaults(_ORCHESTRATOR_ENV_DEFAULTS) + get_logger().info(f"Initializing {config.resolved_name} ({config})") + self._env: vf.Environment = vf.load_environment(config.stripped_id, **config.args) + self._env_client: ZMQEnvClient | None = None + self._env_server_process: BaseProcess | None = None + @property + def name(self) -> str: + return self.config.resolved_name -def __getattr__(name: str) -> Any: - return get_env_value(_ORCHESTRATOR_ENV_PARSERS, name) + @property + def env(self) -> vf.Environment: + return self._env + @property + def env_client(self) -> ZMQEnvClient: + if not self._env_client: + raise RuntimeError( + f"Env {self.name} has no env client connected. Call connect() first to connect to an env server." + ) + return self._env_client -def __dir__() -> list[str]: - return get_dir(_ORCHESTRATOR_ENV_PARSERS) + @property + def requires_group_scoring(self) -> bool: + return any(self.env.rubric._is_group_func(func) for func in self.env.rubric._get_reward_funcs()) + + async def start( + self, + log_dir: Path, + log_level: str | None = None, + json_logging: bool = False, + ) -> None: + """Spawn an env server (if needed) and connect to it.""" + if self.config.address is None: + address = self._spawn(log_dir=log_dir, log_level=log_level, json_logging=json_logging) + else: + address = self.config.address + get_logger().debug(f"Connecting {self.name} to env server {address}") + self._env_client = ZMQEnvClient(address=address, name=self.name) + await self.env_client.wait_for_server_startup() + + def _spawn( + self, + log_dir: Path, + log_level: str | None = None, + json_logging: bool = False, + ) -> str: + assert isinstance(self.config.num_workers, int), ( + f"num_workers must be resolved before spawn, got {self.config.num_workers!r}" + ) + num_workers = self.config.num_workers + address = f"tcp://127.0.0.1:{get_free_port()}" + get_logger().debug(f"Spawning env server {self.name} ({address=}, {num_workers=})") + process = mp.get_context("spawn").Process( + target=ZMQEnvServer.run_server, + args=( + self.config.stripped_id, + self.config.args, + self.config.extra_env_kwargs, + log_level, + (log_dir / self.name).as_posix(), + ), + kwargs=dict( + address=address, + json_logging=json_logging, + console_logging=False, + num_workers=num_workers, + ), + daemon=False, + ) + process.start() + self._env_server_process = process + return address + + def _sampling_args_with_salt(self, cache_salt: str) -> dict: + sampling_args = {**self.sampling_args} + extra_body = {**sampling_args.get("extra_body", {}), "cache_salt": cache_salt} + sampling_args["extra_body"] = extra_body + return sampling_args + + async def run_rollout( + self, + client: vf.ClientConfig, + example: dict, + model_name: str, + cache_salt: str, + ) -> vf.RolloutOutput: + """Run a single rollout for an example.""" + return await self.env.run_rollout( + vf.RolloutInput(**example), + client=client, + model=model_name, + sampling_args=self._sampling_args_with_salt(cache_salt), + max_retries=self.config.max_retries, + state_columns=REQUIRED_STATE_COLUMNS, + env_client=self.env_client, + ) + + async def run_group( + self, + client: vf.ClientConfig, + example: dict, + model_name: str, + rollouts_per_example: int, + cache_salt: str, + ) -> list[vf.RolloutOutput]: + """Run a group of rollouts for an example. Required for group-scoring envs.""" + return await self.env.run_group( + [vf.RolloutInput(**example) for _ in range(rollouts_per_example)], + client=client, + model=model_name, + sampling_args=self._sampling_args_with_salt(cache_salt), + max_retries=self.config.max_retries, + state_columns=REQUIRED_STATE_COLUMNS, + env_client=self.env_client, + ) + + def shutdown(self) -> None: + if self._env_server_process is None: + return + self._env_server_process.terminate() + self._env_server_process = None + + +class TrainEnv(Env): + config: TrainEnvConfig + + def __init__(self, config: TrainEnvConfig): + super().__init__(config) + self.sampling_args = config.sampling.to_sampling_args() + + def get_dataset(self, seed: int | None = None): + return self.env.get_dataset(seed=seed) + + +class EvalEnv(Env): + config: EvalEnvConfig + + def __init__(self, config: EvalEnvConfig): + super().__init__(config) + self.sampling_args = config.sampling.to_sampling_args() + self.examples = self.env.get_eval_dataset(n=config.num_examples).to_list() + + async def evaluate( + self, + model_name: str, + get_client: Callable[[], Awaitable[vf.ClientConfig]], + ckpt_step: int, + step: int, + cache_salt: str, + ) -> list[vf.RolloutOutput]: + num_examples = len(self.examples) + rollouts_per_example = self.config.rollouts_per_example + get_logger().info(f"Evaluating {self.name} ({num_examples=}, {rollouts_per_example=})") + total_rollouts = num_examples * rollouts_per_example + pbar = ProgressTracker(total=total_rollouts, desc=f"Evaluating {self.name}") + eval_start = time.perf_counter() + + if self.requires_group_scoring: + + async def run_with_progress(example: dict) -> list[vf.RolloutOutput] | None: + """Run rollouts_per_example rollouts as a scored group for one example.""" + try: + client = await get_client() + outputs = await self.run_group( + client=client, + example=example, + model_name=model_name, + rollouts_per_example=rollouts_per_example, + cache_salt=cache_salt, + ) + pbar.update(rollouts_per_example) + return outputs + except Exception as e: + get_logger().warning(f"Group failed: {e}") + pbar.update(rollouts_per_example) + return None + + coros = [run_with_progress(example) for example in self.examples] + + else: + + async def run_with_progress(example: dict) -> list[vf.RolloutOutput] | None: + """Run a single rollout for one example.""" + try: + client = await get_client() + output = await self.run_rollout( + client=client, example=example, model_name=model_name, cache_salt=cache_salt + ) + pbar.update(1) + return [output] + except Exception as e: + get_logger().warning(f"Rollout failed: {e}") + pbar.update(1) + return None + + coros = [run_with_progress(example) for example in self.examples for _ in range(rollouts_per_example)] + + try: + results = await asyncio.gather(*coros) + finally: + pbar.close() + + successful_outputs = [o for group in results if group is not None for o in group] + failed_count = total_rollouts - len(successful_outputs) + eval_time = time.perf_counter() - eval_start + + if failed_count: + get_logger().warning( + f"{failed_count}/{total_rollouts} ({failed_count / total_rollouts * 100:.1f}%) rollouts failed" + ) + + if not successful_outputs: + get_logger().warning(f"All rollouts failed for {self.name}, skipping logging metrics") + get_monitor().log( + { + f"eval/{self.name}/failed_rollouts": failed_count / total_rollouts, + "progress/ckpt_step": ckpt_step, + "step": step, + }, + step=step, + ) + return [] + + # Log metrics + monitor = get_monitor() + + rows = [ + { + "example_id": o["example_id"], + "reward": o["reward"], + "completion_len": get_completion_len(o), + "is_truncated": o["is_truncated"], + "has_error": o.get("error") is not None, + "no_response": not o.get("completion"), + } + for o in successful_outputs + ] + results_df = pd.DataFrame(rows) + + unique_rewards = results_df.reward.dropna().unique() + could_be_binary = set(unique_rewards).issubset({0.0, 1.0}) + if could_be_binary: + pass_at_k = ( + results_df.groupby("example_id") + .apply(lambda x: compute_pass_at_k(x.reward.dropna()), include_groups=False) + .apply(pd.Series) + ) + else: + pass_at_k = None + get_logger().warning("Skipping computing pass@k rates because the task rewards appear to be non-binary") + + message = ( + f"Evaluated {self.name} in {eval_time:.2f}s (Avg@{rollouts_per_example}={results_df.reward.mean():.4f}" + ) + if could_be_binary: + assert pass_at_k is not None + for pass_rate, pass_rate_score in pd.Series(pass_at_k.mean()).items(): + message += f", {capitalize(str(pass_rate))}: {pass_rate_score:.4f}" + + message += ( + f", No-response: {results_df.no_response.mean() * 100:.1f}%" + f", Completion Length: {results_df.completion_len.mean():.2f} (±{results_df.completion_len.std():.2f}, ∈[{results_df.completion_len.min():.2f}, {results_df.completion_len.max():.2f}])" + f", Truncated: {results_df.is_truncated.mean() * 100:.1f}%)" + ) + get_logger().success(message) + + eval_metrics = { + f"avg@{rollouts_per_example}": float(results_df.reward.mean()), + "no_response/mean": float(results_df.no_response.mean()), + "no_response/count": int(results_df.no_response.sum()), + "completion_len/mean": results_df.completion_len.mean().item(), + "completion_len/max": results_df.completion_len.max().item(), + "completion_len/min": results_df.completion_len.min().item(), + "is_truncated/mean": results_df.is_truncated.mean().item(), + "failed_rollouts": failed_count / total_rollouts, + "time": eval_time, + } + if could_be_binary: + assert pass_at_k is not None + eval_metrics.update(pd.Series(pass_at_k.mean()).to_dict()) + eval_metrics = {f"eval/{self.name}/{key}": v for key, v in eval_metrics.items()} + eval_metrics.update({"progress/ckpt_step": ckpt_step, "step": step}) + monitor.log(eval_metrics, step=step) + monitor.log_eval_samples(successful_outputs, env_name=self.name, step=step) + + return successful_outputs + + +EnvT = TypeVar("EnvT", bound=Env) + + +class Envs(Generic[EnvT]): + """Base container for a set of Env instances.""" + + _envs: dict[str, EnvT] + + @property + def names(self) -> list[str]: + return list(self._envs.keys()) + + @property + def configs(self) -> list[EnvConfig]: + return [env.config for env in self._envs.values()] + + def get(self, name: str) -> EnvT: + return self._envs[name] + + def __iter__(self) -> Iterator[EnvT]: + return iter(self._envs.values()) + + def __len__(self) -> int: + return len(self._envs) + + async def start( + self, + log_dir: Path, + log_level: str | None = None, + json_logging: bool = False, + ) -> None: + """Spawn env servers (where needed) and connect env clients one at a time. + + Serialized to avoid a TOCTOU port race: get_free_port() only holds the port + until it returns, so parallel spawns can hand the same port to two children. + """ + for env in self: + await env.start(log_dir=log_dir, log_level=log_level, json_logging=json_logging) + atexit.register(self.shutdown) + + def shutdown(self) -> None: + """Terminate all spawned env server processes in parallel.""" + processes = [env._env_server_process for env in self if env._env_server_process is not None] + if not processes: + return + logger = get_logger() + logger.info(f"Shutting down {len(processes)} env server(s), waiting for sandbox cleanup...") + for p in processes: + p.terminate() + for p in processes: + p.join(timeout=25) + if p.is_alive(): + logger.warning(f"Env server {p.pid} did not exit after 25s, force killing") + p.kill() + p.join(timeout=5) + for env in self: + env._env_server_process = None + + +class TrainEnvs(Envs[TrainEnv]): + """Collection of training environments.""" + + def __init__(self, configs: Sequence[TrainEnvConfig]): + self._envs: dict[str, TrainEnv] = {} + for config in configs: + env = TrainEnv(config) + self._envs[env.name] = env + + +class EvalEnvs(Envs[EvalEnv]): + """Collection of evaluation environments.""" + + def __init__(self, configs: Sequence[EvalEnvConfig]): + self._envs: dict[str, EvalEnv] = {} + for config in configs: + env = EvalEnv(config) + self._envs[env.name] = env diff --git a/src/prime_rl/orchestrator/eval_utils.py b/src/prime_rl/orchestrator/eval_utils.py new file mode 100644 index 00000000000..75118d63fa3 --- /dev/null +++ b/src/prime_rl/orchestrator/eval_utils.py @@ -0,0 +1,44 @@ +import numpy as np + + +def compute_eval_ckpt_step( + ckpt_step: int, + prev_ckpt_step: int, + last_eval_step: int, + interval: int, + eval_base_model: bool = True, +) -> int | None: + """Determine which checkpoint step (if any) should trigger an eval. + + Handles the case where ckpt_step jumps over interval boundaries by finding + the highest interval-aligned step in (prev_ckpt_step, ckpt_step]. + + Returns the interval step to eval at, or None if no eval should run. + """ + if ckpt_step <= prev_ckpt_step: + return None + highest_interval_step = (ckpt_step // interval) * interval + if highest_interval_step > prev_ckpt_step and highest_interval_step > last_eval_step: + if highest_interval_step == 0: + if ckpt_step == 0 and eval_base_model: + return 0 + else: + return highest_interval_step + return None + + +def _pass_at_k(n: int, c: int, k: int) -> float: + """Unbiased estimator of pass@k (Chen et al., 2021). + + Computes 1 - C(n-c, k) / C(n, k) in a numerically stable way. + """ + if n - c < k: + return 1.0 + return 1.0 - np.prod(1.0 - k / np.arange(n - c + 1, n + 1)) + + +def compute_pass_at_k(rewards: list[float]) -> dict[str, float]: + n = len(rewards) + c = sum(r == 1.0 for r in rewards) + ks = [2**i for i in range(n.bit_length())] + return {f"pass@{k}": _pass_at_k(n, c, k) for k in ks} diff --git a/src/prime_rl/orchestrator/filters.py b/src/prime_rl/orchestrator/filters.py new file mode 100644 index 00000000000..b2921d22b1d --- /dev/null +++ b/src/prime_rl/orchestrator/filters.py @@ -0,0 +1,188 @@ +"""Orchestrator-side rollout filters for detecting degenerate generations. + +Filters run after rollouts complete, inspecting token IDs and logprobs to +detect gibberish or repetition. Detection metrics are always tracked. +When enforce=True, detected rollouts are skipped entirely during training and +are not sent to the trainer. Reward is kept as-is for baseline calculation. +""" + +import math +from dataclasses import dataclass +from typing import Protocol + +import verifiers as vf + +from prime_rl.configs.orchestrator import FilterConfig +from prime_rl.utils.logger import get_logger + + +@dataclass +class FilterResult: + detected: bool + detection_index: int | None = None + + +class RolloutFilter(Protocol): + name: str + enforce: bool + + def check(self, rollout: vf.RolloutOutput) -> FilterResult: ... + + +@dataclass +class GibberishFilter: + """Flags rollouts containing rare tokens generated at high entropy. + + A token is flagged when both: + - id(token) > token_id_threshold (rare BPE token) + - logprob(token) < -log(vocab_size) - logprob_offset (high entropy) + + References: + Section 5.2, https://arxiv.org/abs/2510.02387 + """ + + name: str + token_id_threshold: int + logprob_threshold: float + enforce: bool = False + + def check(self, rollout: vf.RolloutOutput) -> FilterResult: + global_idx = 0 + for step in rollout["trajectory"]: + tokens = step["tokens"] + if tokens is None: + continue + for token_id, logprob in zip(tokens["completion_ids"], tokens["completion_logprobs"]): + if token_id > self.token_id_threshold and logprob < self.logprob_threshold: + return FilterResult(detected=True, detection_index=global_idx) + global_idx += 1 + return FilterResult(detected=False) + + +@dataclass +class RepetitionFilter: + """Flags rollouts with pathological repetition loops. + + Counts consecutive tokens where logprob > log(prob_threshold), indicating + the model is generating with very high confidence. When the streak reaches + the window size, the rollout is flagged. + + References: + Section 3.2, https://arxiv.org/abs/2506.13585 + """ + + name: str + window: int + logprob_threshold: float + enforce: bool = False + + def check(self, rollout: vf.RolloutOutput) -> FilterResult: + consecutive = 0 + global_idx = 0 + for step in rollout["trajectory"]: + tokens = step["tokens"] + if tokens is None: + continue + for logprob in tokens["completion_logprobs"]: + if logprob > self.logprob_threshold: + consecutive += 1 + else: + consecutive = 0 + if consecutive >= self.window: + return FilterResult(detected=True, detection_index=global_idx) + global_idx += 1 + return FilterResult(detected=False) + + +@dataclass +class ZeroAdvantageFilter: + """Flags rollouts with zero advantage. + + This filter is applied after advantages are computed and checks if the + rollout's advantage field is zero. + """ + + name: str + enforce: bool = True + + def check(self, rollout: vf.RolloutOutput) -> FilterResult: + advantage = rollout.get("advantage") + if advantage is not None and advantage == 0.0: + return FilterResult(detected=True) + return FilterResult(detected=False) + + +def setup_filter(config: FilterConfig, vocab_size: int) -> RolloutFilter: + """Create a RolloutFilter from a filter config.""" + if config.type == "gibberish": + return GibberishFilter( + name="gibberish", + token_id_threshold=config.token_id_threshold, + logprob_threshold=-math.log(vocab_size) - config.logprob_offset, + enforce=config.enforce, + ) + elif config.type == "repetition": + return RepetitionFilter( + name="repetition", + window=config.window, + logprob_threshold=math.log(config.prob_threshold), + enforce=config.enforce, + ) + elif config.type == "zero_advantage": + return ZeroAdvantageFilter( + name="zero_advantage", + enforce=config.enforce, + ) + raise ValueError(f"Unknown filter type: {config.type}") + + +def setup_filters(configs: list[FilterConfig], vocab_size: int) -> list[RolloutFilter]: + """Create RolloutFilters from a list of filter configs.""" + filters = [setup_filter(config, vocab_size) for config in configs] + if filters: + get_logger().info(f"Configured {len(filters)} rollout filter(s):") + for config, filt in zip(configs, filters): + mode = "Enforcing" if filt.enforce else "Monitoring" + params = ", ".join(f"{k}={v}" for k, v in config.model_dump().items()) + get_logger().info(f" {mode} {filt.name} filter ({params})") + return filters + + +def apply_filters(filters: list[RolloutFilter], rollouts: list[vf.RolloutOutput]) -> None: + """Flag rollouts in-place with per-filter detection and drop decision. + + Each rollout gets a `filters` dict with per-filter detection booleans and + an `is_filtered` bool that is True iff an enforcing filter detected it. + First matching filter wins per rollout (no double-counting). Reward and + trajectory tokens are left untouched so the rollout can still contribute + to baseline calculations and metric aggregation. + """ + for rollout in rollouts: + rollout["filters"] = {f.name: False for f in filters} + rollout["is_filtered"] = False + + if not filters: + return + + counts: dict[str, int] = {f.name: 0 for f in filters} + total_detected = 0 + total_enforced = 0 + + for rollout in rollouts: + for filt in filters: + result = filt.check(rollout) + if result.detected: + counts[filt.name] += 1 + total_detected += 1 + rollout["filters"][filt.name] = True + if filt.enforce: + rollout["is_filtered"] = True + total_enforced += 1 + break + + if total_detected > 0: + enforced_msg = f", enforced {total_enforced}" if total_enforced > 0 else "" + get_logger().info( + f"Detected {total_detected}/{len(rollouts)} rollouts " + f"({', '.join(f'{name}={c}' for name, c in counts.items() if c > 0)})" + enforced_msg + ) diff --git a/src/prime_rl/orchestrator/inference_metrics.py b/src/prime_rl/orchestrator/inference_metrics.py new file mode 100644 index 00000000000..7e756622f84 --- /dev/null +++ b/src/prime_rl/orchestrator/inference_metrics.py @@ -0,0 +1,236 @@ +from __future__ import annotations + +import asyncio +import time +from collections import deque + +import wandb +from httpx import AsyncClient +from prometheus_client.parser import text_string_to_metric_families + +from prime_rl.utils.logger import get_logger + +POLL_INTERVAL = 5.0 +WINDOW_SIZE = 20 + +# Gauge metrics: collected as instantaneous values +GAUGE_METRICS = { + "vllm:num_requests_running", + "vllm:num_requests_waiting", + "vllm:gpu_cache_usage_perc", + "vllm:gpu_prefix_cache_hit_rate", +} + +# Counter metrics: converted to per-second rates via delta/dt +COUNTER_METRICS = { + "vllm:prompt_tokens", + "vllm:generation_tokens", + "vllm:request_success", +} + +COUNTER_RATE_NAMES = { + "vllm:prompt_tokens": "prefill_throughput_tps", + "vllm:generation_tokens": "decode_throughput_tps", + "vllm:request_success": "completed_requests_per_s", +} + +# Histogram metrics: converted to average latency per interval +HISTOGRAM_METRICS = { + "vllm:nixl_xfer_time_seconds", +} + +_COUNTER_TOTAL_TO_NAME = {f"{name}_total": name for name in COUNTER_METRICS} + +# Gauges where we log both max and mean across engines (to show imbalance) +_DUAL_AGG_GAUGES = {"vllm:gpu_cache_usage_perc", "vllm:gpu_prefix_cache_hit_rate"} + + +def parse_prometheus_text(text: str) -> tuple[dict[str, float], dict[str, float], dict[str, tuple[float, float]]]: + """Parse Prometheus exposition format into (gauges, counters, histograms). + + For gauge metrics, returns the per-server aggregate (sum for queue sizes, + max for per-engine metrics like kv cache within a single server). + """ + gauges: dict[str, float] = {} + counters: dict[str, float] = {} + histograms: dict[str, tuple[float, float]] = {} + + for family in text_string_to_metric_families(text): + if family.type == "gauge" and family.name in GAUGE_METRICS: + for sample in family.samples: + if family.name in _DUAL_AGG_GAUGES: + gauges[family.name] = max(gauges.get(family.name, 0.0), sample.value) + else: + gauges[family.name] = gauges.get(family.name, 0.0) + sample.value + + elif family.type == "counter" and family.name in COUNTER_METRICS: + for sample in family.samples: + counters[family.name] = counters.get(family.name, 0.0) + sample.value + + elif family.name in _COUNTER_TOTAL_TO_NAME: + canonical = _COUNTER_TOTAL_TO_NAME[family.name] + for sample in family.samples: + counters[canonical] = counters.get(canonical, 0.0) + sample.value + + elif family.type == "histogram" and family.name in HISTOGRAM_METRICS: + h_sum = 0.0 + h_count = 0.0 + for sample in family.samples: + if sample.name.endswith("_sum"): + h_sum += sample.value + elif sample.name.endswith("_count"): + h_count += sample.value + histograms[family.name] = (h_sum, h_count) + + return gauges, counters, histograms + + +class InferenceMetricsCollector: + """Polls vLLM Prometheus /metrics and logs aggregated values to W&B. + + Runs independently of training steps on a time-based axis. + """ + + def __init__(self, admin_clients: list[AsyncClient]): + self.admin_clients = admin_clients + self.logger = get_logger() + self._gauge_history: dict[str, deque[float]] = {} + self._rate_history: dict[str, deque[float]] = {} + self._prev_counters: dict[str, tuple[float, float]] = {} + self._prev_histograms: dict[str, tuple[float, float, float]] = {} + self._task: asyncio.Task | None = None + + async def start(self): + wandb.define_metric("inference/*", step_metric="_timestamp") + + async def poll_loop(): + while True: + try: + await self._collect_and_log() + except Exception as e: + self.logger.debug(f"Inference metrics poll failed: {e!r}") + await asyncio.sleep(POLL_INTERVAL) + + self._task = asyncio.create_task(poll_loop()) + + async def _collect_and_log(self): + now = time.monotonic() + + async def fetch(client: AsyncClient) -> str | None: + try: + response = await client.get("/metrics", timeout=5.0) + response.raise_for_status() + return response.text + except Exception as e: + self.logger.debug(f"Failed to fetch metrics from {client.base_url}: {e!r}") + return None + + results = await asyncio.gather(*[fetch(client) for client in self.admin_clients]) + + # For dual-agg gauges, collect per-server values to compute both max and mean + dual_agg_values: dict[str, list[float]] = {} + agg_sum_gauges: dict[str, float] = {} + agg_counters: dict[str, float] = {} + agg_histograms: dict[str, tuple[float, float]] = {} + n_servers = 0 + + for text in results: + if text is None: + continue + n_servers += 1 + gauges, counters, histograms = parse_prometheus_text(text) + + for name, value in gauges.items(): + if name in _DUAL_AGG_GAUGES: + dual_agg_values.setdefault(name, []).append(value) + else: + agg_sum_gauges[name] = agg_sum_gauges.get(name, 0.0) + value + + for name, value in counters.items(): + agg_counters[name] = agg_counters.get(name, 0.0) + value + + for name, (h_sum, h_count) in histograms.items(): + prev = agg_histograms.get(name, (0.0, 0.0)) + agg_histograms[name] = (prev[0] + h_sum, prev[1] + h_count) + + if n_servers == 0: + return + + # Update gauge history — sum gauges + for name, value in agg_sum_gauges.items(): + short = name.removeprefix("vllm:") + if short not in self._gauge_history: + self._gauge_history[short] = deque(maxlen=WINDOW_SIZE) + self._gauge_history[short].append(value) + + # Update gauge history — dual-agg gauges (max + mean across engines) + for name, values in dual_agg_values.items(): + short = name.removeprefix("vllm:") + max_key = f"{short}_max" + mean_key = f"{short}_mean" + if max_key not in self._gauge_history: + self._gauge_history[max_key] = deque(maxlen=WINDOW_SIZE) + if mean_key not in self._gauge_history: + self._gauge_history[mean_key] = deque(maxlen=WINDOW_SIZE) + self._gauge_history[max_key].append(max(values)) + self._gauge_history[mean_key].append(sum(values) / len(values)) + + # Compute rates from counters + for name, value in agg_counters.items(): + rate_name = COUNTER_RATE_NAMES[name] + prev = self._prev_counters.get(name) + self._prev_counters[name] = (now, value) + if prev is None: + continue + prev_time, prev_value = prev + dt = now - prev_time + if dt <= 0: + continue + delta = value - prev_value + if delta < 0: + continue + rate = delta / dt + if rate_name not in self._rate_history: + self._rate_history[rate_name] = deque(maxlen=WINDOW_SIZE) + self._rate_history[rate_name].append(rate) + + # Compute average histogram latency + for name, (h_sum, h_count) in agg_histograms.items(): + short = name.removeprefix("vllm:") + rate_name = f"{short}_avg_ms" + prev = self._prev_histograms.get(name) + self._prev_histograms[name] = (now, h_sum, h_count) + if prev is None: + continue + prev_time, prev_sum, prev_count = prev + d_sum = h_sum - prev_sum + d_count = h_count - prev_count + if d_count < 0 or d_sum < 0: + continue + if d_count > 0: + avg_ms = (d_sum / d_count) * 1000.0 + if rate_name not in self._rate_history: + self._rate_history[rate_name] = deque(maxlen=WINDOW_SIZE) + self._rate_history[rate_name].append(avg_ms) + + # Build smoothed metrics dict + metrics: dict[str, float] = {} + for short, values in self._gauge_history.items(): + if values: + metrics[f"inference/{short}"] = sum(values) / len(values) + for rate_name, values in self._rate_history.items(): + if values: + metrics[f"inference/{rate_name}"] = sum(values) / len(values) + + if metrics: + metrics["_timestamp"] = time.time() + wandb.log(metrics) + + async def stop(self): + if self._task is not None: + self._task.cancel() + try: + await self._task + except asyncio.CancelledError: + pass + self._task = None diff --git a/src/prime_rl/orchestrator/orchestrator.py b/src/prime_rl/orchestrator/orchestrator.py index 5afc9272880..bc1128ebc7c 100644 --- a/src/prime_rl/orchestrator/orchestrator.py +++ b/src/prime_rl/orchestrator/orchestrator.py @@ -1,14 +1,25 @@ import asyncio -import random +import gc +import os import time import tomli_w +import prime_rl._compat # noqa: F401 — patch ring_flash_attn compat before transitive import from prime_rl.orchestrator.advantage import compute_advantages +from prime_rl.orchestrator.eval_utils import compute_eval_ckpt_step from prime_rl.orchestrator.event_loop_lag import EventLoopLagMonitor +from prime_rl.orchestrator.inference_metrics import InferenceMetricsCollector from prime_rl.orchestrator.patches import monkey_patch_chat_completion_logprobs, monkey_patch_oai_iterable_types -from prime_rl.orchestrator.trajectories import branch_rollout, interleave_rollout +from prime_rl.orchestrator.trajectories import ( + build_vlm_image_cache, + interleave_rollout, + offload_images_to_disk, + pretokenize_rollout_trajectory, +) from prime_rl.transport import TrainingBatch, TrainingSample, setup_training_batch_sender +from prime_rl.utils.pathing import get_log_dir, get_rollout_dir, get_step_path +from prime_rl.utils.usage_reporter import UsageReporter # This monkey patch is necessary to avoid Pydantic validating fields using typing.Iterable (e.g. in multimodal or tool call messages) lazily which leads to tokenization errors, for more info see https://github.com/PrimeIntellect-ai/prime-rl/pull/1249 monkey_patch_oai_iterable_types() @@ -18,36 +29,40 @@ monkey_patch_chat_completion_logprobs() # Import environment before any other imports + import pandas as pd import verifiers as vf -from loguru import logger -from transformers import AutoTokenizer +from renderers.base import create_renderer +from transformers import AutoProcessor -from prime_rl.eval.utils import run_evals_subprocess +from prime_rl.configs.orchestrator import OrchestratorConfig from prime_rl.orchestrator.buffer import Buffer from prime_rl.orchestrator.ckpt import Progress, setup_ckpt_manager -from prime_rl.orchestrator.config import BufferConfig, OrchestratorConfig +from prime_rl.orchestrator.envs import EvalEnv, EvalEnvs, TrainEnvs +from prime_rl.orchestrator.filters import apply_filters, setup_filters from prime_rl.orchestrator.scheduler import Scheduler from prime_rl.orchestrator.utils import ( compute_teacher_logprobs, - get_sampling_args, get_weight_dir, print_benchmark, - set_semaphore, + set_default_executor, + setup_external_rollout_model, +) +from prime_rl.orchestrator.vf_utils import ( + get_seq_len, + intercept_vf_logging, + save_rollouts, ) +from prime_rl.trainer.model import setup_tokenizer from prime_rl.utils.client import ( - check_has_model, - check_health, init_nccl_broadcast, - reload_weights, - setup_admin_clients, - setup_clients, - update_weights, + setup_inference_pool, ) +from prime_rl.utils.config import cli from prime_rl.utils.heartbeat import Heartbeat -from prime_rl.utils.logger import intercept_verifiers_logging, setup_logger +from prime_rl.utils.logger import setup_logger from prime_rl.utils.monitor import setup_monitor -from prime_rl.utils.pydantic_config import parse_argv +from prime_rl.utils.process import set_proc_title from prime_rl.utils.utils import ( clean_exit, get_env_ids_to_install, @@ -55,18 +70,30 @@ resolve_latest_ckpt_step, to_col_format, ) -from prime_rl.utils.vf import generate_batch, get_completion_len, get_prompt_len, get_seq_len + +# Hard wall-clock budget for the orchestrator's post-training cleanup. If the +# graceful shutdown sequence (scheduler / inference pool / env teardown) is +# still running after this many seconds, we force-exit the process so the run +# pod terminates instead of sitting wedged forever. The training checkpoint +# and artifacts are persisted *before* this point, so a forced exit is safe. +SHUTDOWN_TIMEOUT_S = 300 + +# Maximum number of times to attempt generating a training batch when all +# rollouts are filtered out. After this many attempts, the orchestrator crashes +# rather than silently skipping training steps. +MAX_EMPTY_BATCH_ATTEMPTS = 3 @clean_exit -@logger.catch(reraise=True) async def orchestrate(config: OrchestratorConfig): # Initialize the logger logger = setup_logger( - config.log.level, log_file=config.output_dir / "logs" / "orchestrator.log" if config.log.file else None + config.log.level, + json_logging=config.log.json_logging, ) - intercept_verifiers_logging(level=config.log.vf_level) + intercept_vf_logging(logger="verifiers.serve", level="WARN") # show logs from env clients logger.info("Starting orchestrator") + set_default_executor() event_loop_lag_monitor = EventLoopLagMonitor() event_loop_lag_monitor_task = asyncio.create_task(event_loop_lag_monitor.run()) @@ -76,45 +103,60 @@ async def orchestrate(config: OrchestratorConfig): logger.warning(f"Running in benchmark mode (max_steps={config.max_steps})") # Save configs to output directory - config_dir = config.output_dir / "configs" + config_dir = config.output_dir / "control" config_dir.mkdir(parents=True, exist_ok=True) with open(config_dir / "orch.toml", "wb") as f: tomli_w.dump(config.model_dump(exclude_none=True, mode="json"), f) # Install environments env_ids_to_install = set() - env_ids_to_install.update(get_env_ids_to_install(config.env)) + env_ids_to_install.update(get_env_ids_to_install(config.train.env)) if config.eval is not None: env_ids_to_install.update(get_env_ids_to_install(config.eval.env)) for env_id in env_ids_to_install: - install_env(env_id) + install_env(env_id, prerelease=config.env_install_prerelease) - # Setup client - logger.info( - f"Initializing OpenAI client (base_url={', '.join(config.client.base_url)}, api_key_var={config.client.api_key_var}, headers={config.client.headers})" - ) - clients = setup_clients(config.client) - admin_clients = setup_admin_clients(config.client) + # Setup rollout inference pool (handles both static and elastic modes) + rollout_client_config, rollout_model_name, enable_policy_updates = setup_external_rollout_model(config, logger) - # Setup teacher model client if configured - teacher_clients = None - teacher_admin_clients = None - teacher_model_name = None + # Setup teacher inference pool if configured if config.teacher_model: logger.info( - f"Initializing teacher OpenAI client (base_url={', '.join(config.teacher_model.client.base_url)}, " + f"Initializing teacher inference pool (base_url={', '.join(config.teacher_model.client.base_url)}, " f"model={config.teacher_model.model.name})" ) - teacher_clients = setup_clients(config.teacher_model.client) - teacher_admin_clients = setup_admin_clients(config.teacher_model.client) - teacher_model_name = config.teacher_model.model.name + teacher_inference_pool = await setup_inference_pool( + config.teacher_model.client, + model_name=config.teacher_model.model.name, + train_client_type="openai_chat_completions", + ) + else: + teacher_inference_pool = None + + # Check if this is a vision-language model (used throughout for VLM-specific paths) + is_vlm = config.model.vlm is not None - # Load tokenizer - logger.info(f"Initializing tokenizer for {config.model.name}") - tokenizer = AutoTokenizer.from_pretrained(config.model.name, trust_remote_code=config.model.trust_remote_code) + # Load tokenizer and processor (processor only for VLM models) + logger.info(f"Initializing tokenizer ({config.tokenizer})") + tokenizer = setup_tokenizer(config.tokenizer) - # Setup monitor + processor = None + if is_vlm: + logger.info(f"Loading VLM processor for {config.model.name}") + processor = AutoProcessor.from_pretrained( + config.model.name, trust_remote_code=config.model.trust_remote_code, use_fast=True + ) + + renderer, inference_pool = await setup_rollout_inference_pool( + config=config, + rollout_client_config=rollout_client_config, + rollout_model_name=rollout_model_name, + tokenizer=tokenizer, + logger=logger, + ) + + # Setup monitor (may register the run and set RUN_ID in the environment) logger.info(f"Initializing monitor (wandb={config.wandb}, prime_monitor={config.prime_monitor})") monitor = setup_monitor( wandb_config=config.wandb, @@ -122,41 +164,62 @@ async def orchestrate(config: OrchestratorConfig): output_dir=config.output_dir, tokenizer=tokenizer, run_config=config, + keep_full_history=config.bench, ) + # Read run_id AFTER setup_monitor so that newly registered runs are captured + run_id = os.getenv("RUN_ID", "") + + # Usage reporter requires BOTH the base URL and the API key. Activating + # with only one set used to crash every POST inside httpx (None header + # value), so we now gate construction on both being present and log a + # clear warning when half-configured. + usage_base_url = os.environ.get("PI_USAGE_BASE_URL") + usage_api_key = os.environ.get("PI_USAGE_API_KEY") + if usage_base_url and usage_api_key: + usage_reporter = UsageReporter() + else: + if usage_base_url and not usage_api_key: + logger.warning("PI_USAGE_BASE_URL is set but PI_USAGE_API_KEY is missing; usage reporting disabled.") + usage_reporter = None + # Setup heartbeat (only on rank 0, orchestrator is single process) heart = None if config.heartbeat is not None: logger.info("Initializing heartbeat") heart = Heartbeat(config.heartbeat.url) - # Load environment and extract dataset - logger.info( - f"Loading {len(config.env)} training environment(s) ({', '.join(env.name or env.id for env in config.env)})" - ) - env = vf.EnvGroup( - envs=[vf.load_environment(env.id, **env.args) for env in config.env], - env_names=[env.name or env.id for env in config.env], - map_kwargs=dict(writer_batch_size=1), # Set defensively to not error on map operations on large datasets + # Build rollout filters + rollout_filters = setup_filters(config.filters, vocab_size=tokenizer.vocab_size) + + # Load environments + logger.info("Loading training environments") + train_envs = TrainEnvs(config.train.env) + logger.info(f"Loaded {len(train_envs)} training environment(s) ({', '.join(train_envs.names)})") + + await train_envs.start( + log_dir=get_log_dir(config.output_dir.parent) / "envs" / "train", + log_level=config.log.vf_level, + json_logging=config.log.json_logging, ) - env.set_max_seq_len(config.seq_len) - if config.trajectory_strategy == "interleaved": - logger.info("Using token prompts in environment to avoid retokenization discrepancies in multi-turn rollouts") - env.set_interleaved_rollouts(True) - if config.buffer.skip_verification: - logger.info("Skipping verification (rewards will be set to 0)") - env.set_score_rollouts(False) + logger.success("Train environment(s) ready") + + eval_envs: EvalEnvs | None = None + if config.eval: + logger.info("Loading eval environment(s)") + eval_envs = EvalEnvs(config.eval.env) + logger.info(f"Loaded {len(eval_envs)} eval environment(s) ({', '.join(eval_envs.names)})") + + await eval_envs.start( + log_dir=get_log_dir(config.output_dir.parent) / "envs" / "eval", + log_level=config.log.vf_level, + json_logging=config.log.json_logging, + ) + logger.success("Eval environment(s) ready") # Setup buffer logger.info(f"Setting up buffer ({config.buffer})") - train_dataset = env.get_dataset(seed=config.buffer.seed) - buffer = Buffer(train_dataset, env.env_names, config.buffer) - if config.val is not None: - val_buffer_config = BufferConfig(env_ratios=config.buffer.env_ratios) - val_dataset = env.get_eval_dataset(seed=val_buffer_config.seed) - val_buffer = Buffer(val_dataset, env.env_names, val_buffer_config) - else: - val_buffer = None + buffer = Buffer(train_envs, config.buffer) # Get checkpoint manager logger.info(f"Initializing checkpoint manager ({config.ckpt})") @@ -169,55 +232,66 @@ async def orchestrate(config: OrchestratorConfig): else: checkpoint_step = config.ckpt.resume_step - # Setup scheduler (uses subprocess workers for env execution) scheduler = Scheduler( - admin_clients=admin_clients, - client_config=config.client, - env_configs=config.env, + train_envs=train_envs, buffer=buffer, - config=config, - oversampling_factor=config.oversampling_factor, + inference_pool=inference_pool, + max_inflight_rollouts=config.max_inflight_rollouts, max_async_level=config.max_async_level, max_off_policy_steps=config.max_off_policy_steps, strict_async_level=config.strict_async_level, + tasks_per_minute=config.tasks_per_minute, + enable_policy_updates=enable_policy_updates, lora_name=config.model.lora.name if config.model.lora else None, - output_dir=config.output_dir, + config=config, ) + scheduler.model_name = rollout_model_name - if checkpoint_step is not None and config.model.lora is not None: + if checkpoint_step is not None and config.model.lora is not None and enable_policy_updates: + assert config.model.lora.name is not None scheduler.model_name = config.model.lora.name - for workers in scheduler.workers.values(): - for worker in workers: - worker.model_name = config.model.lora.name - - await scheduler.start() - # Check health of the client + # Check health of the inference pool logger.info("Waiting for inference pool to be ready") - await check_health(admin_clients) - await check_has_model(clients, config.model.name) + await inference_pool.wait_for_ready(rollout_model_name) + logger.success("Inference pool ready") + # Start inference metrics collector (requires W&B) + inference_metrics_collector = None + if config.wandb is not None and config.collect_inference_metrics: + inference_metrics_collector = InferenceMetricsCollector(inference_pool.admin_clients) + await inference_metrics_collector.start() + # Check health of teacher inference server if configured - if teacher_admin_clients is not None: + if config.teacher_model and teacher_inference_pool: logger.info("Waiting for teacher inference pool to be ready") - await check_health(teacher_admin_clients) - await check_has_model(teacher_clients, teacher_model_name) + await teacher_inference_pool.wait_for_ready(config.teacher_model.model.name) logger.success("Teacher inference pool ready") # Set up weight broadcast backend - logger.info(f"Initializing weight broadcast ({config.weight_broadcast})") - if config.weight_broadcast.type == "nccl": - await init_nccl_broadcast( - admin_clients, config.weight_broadcast.host, config.weight_broadcast.port, config.weight_broadcast.timeout - ) + if enable_policy_updates: + logger.info(f"Initializing weight broadcast ({config.weight_broadcast})") + if config.weight_broadcast.type == "nccl": + await init_nccl_broadcast( + inference_pool.admin_clients, + config.weight_broadcast.host, + config.weight_broadcast.port, + config.weight_broadcast.timeout, + inference_world_size=config.weight_broadcast.inference_world_size, + quantize_in_weight_transfer=config.weight_broadcast.quantize_in_weight_transfer, + ) + else: + logger.info("Skipping weight broadcast initialization (SFT distillation mode)") # Setup training batch sender for sending training examples to trainer logger.info(f"Initializing training batch sender ({config.rollout_transport})") training_batch_sender = setup_training_batch_sender(config.output_dir, config.rollout_transport) - # Track last online eval checkpoint step for this process - last_eval_step = -1 + # Track last online eval checkpoint step per eval env + last_eval_steps: dict[str, int] = {env.name: -1 for env in eval_envs} if eval_envs else {} + # Track previous ckpt_step to detect when ckpt_step jumps over eval interval boundaries + prev_ckpt_step = -1 # Reset weights to base model if starting from scratch progress = Progress() @@ -227,37 +301,39 @@ async def orchestrate(config: OrchestratorConfig): logger.info(f"Resuming training from checkpoint step {checkpoint_step}") scheduler.ckpt_step = progress.step # Always resume from the latest checkpoint if config.eval and config.eval.skip_eval_on_resume: - last_eval_step = scheduler.ckpt_step + prev_ckpt_step = scheduler.ckpt_step + last_eval_steps = {name: scheduler.ckpt_step for name in last_eval_steps} logger.info(f"Skipping online eval on resume (ckpt_step={scheduler.ckpt_step})") - - await update_weights( - admin_clients, - get_weight_dir(config.output_dir, scheduler.ckpt_step), - lora_name=config.model.lora.name if config.model.lora else None, - ) + else: + # Allow eval at resumed step by setting prev_ckpt_step one behind + prev_ckpt_step = scheduler.ckpt_step - 1 + + if enable_policy_updates: + # In NCCL mode, skip existence check - weights are broadcasted, not stored on disk + check_exists = config.weight_broadcast.type != "nccl" + wait_timeout = config.ckpt.wait_for_weights_timeout if config.ckpt else None + weights_path = get_weight_dir( + config.output_dir, scheduler.ckpt_step, check_exists=check_exists, wait_timeout=wait_timeout + ) + lora_name = config.model.lora.name if config.model.lora else None + await inference_pool.update_weights(weights_path, lora_name=lora_name, step=scheduler.ckpt_step) else: - logger.info("Training from scratch. Resetting weights to base model") - if config.model.lora is None: - await reload_weights(admin_clients) + logger.info("Training from scratch") # Iterate over dataset in batches - max_steps = config.max_steps or int(1e9) - logger.info(f"Starting orchestrator loop (max_steps={max_steps or 'infinite'})") + logger.info(f"Starting orchestrator loop (max_steps={config.max_steps or 'infinite'})") is_first_step = True - await set_semaphore(config.max_concurrent or -1) - - # Start update policy loop - update_policy_task = asyncio.create_task(scheduler.update_policy_loop()) while True: - # Check if update_policy_task has failed and propagate the exception - if update_policy_task.done(): - # End all other tasks - for task in asyncio.all_tasks(): - task.cancel() - update_policy_task.result() # Raises if the task failed - # Capture ckpt_step once for consistency (it's updated by update_policy_loop concurrently) - ckpt_step = scheduler.ckpt_step + # Check if this run has been evicted by the trainer + evicted_path = config.output_dir / "control" / "evicted.txt" + if evicted_path.exists(): + reason = evicted_path.read_text().strip() + raise RuntimeError(f"Run evicted by trainer: {reason}") + + # Capture ckpt_step once for consistency (it's updated inside the scheduler) + ckpt_step = scheduler.ckpt_step if enable_policy_updates else progress.step + scheduler.ckpt_step = ckpt_step # Save checkpoint (if we are at an interval step and not at the first or last step) is_last_step = config.max_steps is not None and progress.step == config.max_steps - 1 @@ -280,94 +356,229 @@ async def orchestrate(config: OrchestratorConfig): logger.info(f"Starting orchestrator step {progress.step}") step_start_time = time.perf_counter() - # Run evals BEFORE training (blocking, in subprocess to isolate event loop) - # This ensures weights don't change during eval and eval doesn't cause event loop lag - if ( - config.eval - and ckpt_step % config.eval.interval == 0 - and ckpt_step > last_eval_step - and ((ckpt_step == 0 and config.eval.eval_base_model) or ckpt_step > 0) - ): - last_eval_step = ckpt_step - logger.info(f"Running evals for checkpoint step {ckpt_step} (blocking, subprocess)") + # Run evals BEFORE training (blocking). Weight updates are paused via + # scheduler.checkpoint_ready during eval to ensure consistent weights. + # Each eval env has its own interval, so we check each independently. + envs_to_eval: list[EvalEnv] = [] + if config.eval: + assert eval_envs is not None + for eval_env in eval_envs: + eval_ckpt_step = compute_eval_ckpt_step( + ckpt_step=ckpt_step, + prev_ckpt_step=prev_ckpt_step, + last_eval_step=last_eval_steps[eval_env.name], + interval=eval_env.config.interval, + eval_base_model=config.eval.eval_base_model, + ) + if eval_ckpt_step is not None: + last_eval_steps[eval_env.name] = ckpt_step + envs_to_eval.append(eval_env) + + if envs_to_eval: + env_names = ", ".join(e.name for e in envs_to_eval) + logger.info(f"Running evals at {ckpt_step=} for {env_names}") - # Pause weight updates during eval + # Pause weight updates and re-scheduling of training rollouts during eval + # to avoid evaluating across different checkpoints and avoid congestion scheduler.checkpoint_ready.clear() - await run_evals_subprocess( - client_config=config.client, - eval_config=config.eval, - model_config=config.model, - sampling_config=config.eval.sampling, - reasoning_field=config.eval.reasoning_field, - output_dir=config.output_dir, - ckpt_step=ckpt_step, - step=progress.step, - max_concurrent=config.max_concurrent or -1, + # For heavy eval workloads, it might be necessary additionally cancel in-flight training rollouts + if config.eval.cancel_inflight_rollouts_on_eval: + logger.info("Cancelling in-flight training rollouts before starting evals to avoid congestion.") + await scheduler.cancel_inflight_rollouts() + + eval_results = await asyncio.gather( + *[ + eval_env.evaluate( + model_name=scheduler.model_name, + get_client=inference_pool.get_eval_client, + ckpt_step=ckpt_step, + step=progress.step, + cache_salt=str(ckpt_step), + ) + for eval_env in envs_to_eval + ] ) + # Save eval rollouts to disk (fire-and-forget background thread) + eval_rollouts = [o for outputs in eval_results for o in outputs] + if eval_rollouts: + step_path = get_step_path(get_rollout_dir(config.output_dir), progress.step) + await asyncio.to_thread( + save_rollouts, eval_rollouts, step_path / "eval_rollouts.jsonl", exclude_keys={"trajectory"} + ) + # Resume weight updates scheduler.checkpoint_ready.set() - # Schedule generating the training batch - generate_completions_start_time = time.perf_counter() - train_task = asyncio.create_task(scheduler.generate_batch(step=progress.step)) - - # Schedule running validation at the specified interval - if val_buffer and config.val and progress.step % config.val.interval == 0: - logger.info(f"Running validation for step {progress.step}") - val_examples = val_buffer.sample_examples(config.val.num_examples) - val_task = asyncio.create_task( - generate_batch( - clients=clients, - env=env, - model_name=config.model.name, - examples=val_examples, - rollouts_per_example=config.val.rollouts_per_example, - sampling_args=get_sampling_args(config.sampling), - pbar_description="Generating rollouts (val)", + # Update prev_ckpt_step for next iteration + prev_ckpt_step = ckpt_step + + # Schedule generating the training batch. Retry on empty-after-filter + # batches so the trainer never receives an empty batch. + generate_completions_time = 0.0 + train_rollouts: list[vf.RolloutOutput] = [] + num_rollouts = 0 + num_unique_examples = 0 + n_trainable = 0 + for attempt in range(MAX_EMPTY_BATCH_ATTEMPTS): + train_rollouts = await scheduler.generate_batch(step=progress.step) + generate_completions_time += scheduler.last_batch_generation_time + + # Compute advantages (in-place) + num_rollouts = len(train_rollouts) + num_unique_examples = len({(r["env_name"], r["example_id"]) for r in train_rollouts}) + compute_advantages(train_rollouts, config.rollouts_per_example, config.advantage) + + # Apply rollout filters — sets rollout["filters"] and rollout["is_filtered"] + apply_filters(rollout_filters, train_rollouts) + + n_trainable = sum(1 for r in train_rollouts if not r["is_filtered"]) + if n_trainable > 0: + break + + if attempt == MAX_EMPTY_BATCH_ATTEMPTS - 1: + logger.error( + f"Attempt {attempt + 1}/{MAX_EMPTY_BATCH_ATTEMPTS} at step {progress.step} " + f"filtered out all {num_rollouts} rollouts - crashing orchestrator" ) + reason = ( + f"All {num_rollouts} rollouts were filtered out on " + f"{MAX_EMPTY_BATCH_ATTEMPTS} consecutive attempts at step {progress.step}" + ) + evicted_path = config.output_dir / "control" / "evicted.txt" + evicted_path.parent.mkdir(parents=True, exist_ok=True) + evicted_path.write_text(reason) + raise RuntimeError(reason) + + logger.warning( + f"Attempt {attempt + 1}/{MAX_EMPTY_BATCH_ATTEMPTS} at step {progress.step} " + f"filtered out all {num_rollouts} rollouts - retrying batch generation" + ) + + trainable_ratio = n_trainable / num_rollouts + if trainable_ratio <= 0.1: + logger.warning( + f"Only {n_trainable}/{num_rollouts} rollouts in the batch are trainable " + f"({trainable_ratio:.1%}) - this can mean the tasks are too easy or too hard for the " + "model, consider reviewing the task difficulty of your environment(s)" + ) + + # Save train rollouts to disk (fire-and-forget background thread) + step_path = get_step_path(get_rollout_dir(config.output_dir), progress.step) + await asyncio.to_thread( + save_rollouts, train_rollouts, step_path / "train_rollouts.jsonl", exclude_keys={"trajectory"} + ) + + # VLM: offload base64 images to disk immediately to free memory + if is_vlm: + offload_start = time.perf_counter() + num_offloaded = offload_images_to_disk(train_rollouts, config.output_dir) + if num_offloaded: + logger.info( + f"VLM offloaded {num_offloaded} unique images to disk in {time.perf_counter() - offload_start:.2f}s" + ) + + # Convert rollouts to training samples + parallel_preprocess_start = time.perf_counter() + + # Stage 1: pretokenize + (for VLM) build image cache concurrently. + # Pretokenize is a no-op when the renderer client already populated + # `tokens` on each trajectory step, but the fallback-tokenizer path + # and image-cache build are both CPU-heavy. Running them on threads + # and awaiting a single gather lets whichever finishes first free + # the event loop immediately and, with max_async_level >= 2, overlaps + # this whole stage with inference for the next batch. + async def _pretokenize_all() -> None: + await asyncio.gather( + *( + asyncio.to_thread( + pretokenize_rollout_trajectory, + rollout, + tokenizer, + processor=processor, + renderer=renderer, + ) + for rollout in train_rollouts + ) + ) + + if is_vlm: + mm_token_type_ids_mapping = {} + if hasattr(processor, "image_token_id") and processor.image_token_id is not None: + mm_token_type_ids_mapping[processor.image_token_id] = 1 + if hasattr(processor, "video_token_id") and processor.video_token_id is not None: + mm_token_type_ids_mapping[processor.video_token_id] = 2 + _, vlm_cache = await asyncio.gather( + _pretokenize_all(), + asyncio.to_thread(build_vlm_image_cache, train_rollouts, processor), + ) + logger.info( + f"VLM timing: extract={vlm_cache.extract_time:.2f}s, preprocess={vlm_cache.preprocess_time:.2f}s " + f"({vlm_cache.num_unique_images} unique images from {vlm_cache.num_unique_examples} examples)" ) else: - val_task = asyncio.create_task(asyncio.sleep(0)) # Dummy task - - # Await train rollouts, process results and write batch to disk to consume by trainer - await train_task - generate_completions_time = time.perf_counter() - generate_completions_start_time - train_rollouts = train_task.result() - - # Compute advantages - rewards = [rollout["reward"] for rollout in train_rollouts] - completion_lens = [get_completion_len(rollout) for rollout in train_rollouts] - advantages = compute_advantages( - rewards, - completion_lens, - config.rollouts_per_example, - config.advantage, + await _pretokenize_all() + vlm_cache = None + mm_token_type_ids_mapping = None + + # Process rollouts in parallel + def process_rollout(rollout: vf.RolloutOutput, rollout_idx: int) -> list[TrainingSample] | None: + return interleave_rollout( + rollout, + vlm_cache=vlm_cache, + cache_key=rollout_idx, + mm_token_type_ids_mapping=mm_token_type_ids_mapping, + ) + + results = await asyncio.gather( + *(asyncio.to_thread(process_rollout, r, rollout_idx) for rollout_idx, r in enumerate(train_rollouts)) ) - # Update and sample rollouts from the buffer - make_train_example = interleave_rollout if config.trajectory_strategy == "interleaved" else branch_rollout + # Collect results and assign advantages. Metrics are computed over all + # rollouts; only non-filtered samples are sent to the trainer. train_examples: list[TrainingSample] = [] - for train_rollout, advantage in zip(train_rollouts, advantages): - train_example = make_train_example(train_rollout) - if train_example is not None: - for te in train_example: - te.advantage = advantage - te.reward = train_rollout["reward"] - train_examples.extend(train_example) + rollout_prefill_lens: list[int] = [] + rollout_decode_lens: list[int] = [] + rollout_samples_per_rollout: list[int] = [] + num_prefill_tokens = 0 + num_decode_tokens = 0 + for rollout, samples in zip(train_rollouts, results): + rollout_prefill_tokens = 0 + rollout_decode_tokens = 0 + if samples is None: + samples = [] + rollout_samples_per_rollout.append(len(samples)) + for sample in samples: + sample.advantage = rollout["advantage"] + sample.reward = rollout["reward"] + if config.use_sft_loss: + sample.sft_loss = True + sample_decode_tokens = sum(sample.completion_mask) + sample_prefill_tokens = len(sample.prompt_ids) + len(sample.completion_mask) - sample_decode_tokens + rollout_decode_tokens += sample_decode_tokens + rollout_prefill_tokens += sample_prefill_tokens + if not rollout["is_filtered"]: + train_examples.append(sample) + rollout_prefill_lens.append(rollout_prefill_tokens) + rollout_decode_lens.append(rollout_decode_tokens) + num_prefill_tokens += rollout_prefill_tokens + num_decode_tokens += rollout_decode_tokens + + parallel_preprocess_time = time.perf_counter() - parallel_preprocess_start logger.debug( - f"Converted {len(train_rollouts)} training rollouts to {len(train_examples)} training examples using {config.trajectory_strategy} strategy" + f"Converted {len(train_rollouts)} rollouts ({num_unique_examples} unique examples) " + f"to {len(train_examples)} training examples" ) # Compute teacher logprobs if teacher model is configured teacher_logprobs_time = 0 - if config.teacher_model is not None: + if config.teacher_model and teacher_inference_pool: logger.info(f"Computing teacher logprobs for {len(train_examples)} training examples") teacher_logprobs_start_time = time.perf_counter() teacher_logprobs_list = await compute_teacher_logprobs( - clients=teacher_clients, - model_name=teacher_model_name, + clients=teacher_inference_pool.train_clients, + model_name=config.teacher_model.model.name, samples=train_examples, ) for train_example, teacher_logprobs in zip(train_examples, teacher_logprobs_list): @@ -377,193 +588,245 @@ async def orchestrate(config: OrchestratorConfig): training_batch = TrainingBatch( examples=train_examples, - temperature=config.sampling.temperature, step=progress.step, ) - assert len(training_batch.examples) != 0, "Step with no samples is not allowed" + training_batch_sender.send(training_batch) - # Await and process val results - await val_task - val_outputs = val_task.result() + step_time = time.perf_counter() - step_start_time # Gather metrics in dataframes results_df = pd.DataFrame( { "example_id": [rollout["example_id"] for rollout in train_rollouts], - "task": [rollout["task"] for rollout in train_rollouts], + "env_name": [rollout["env_name"] for rollout in train_rollouts], "reward": [rollout["reward"] for rollout in train_rollouts], "is_truncated": [rollout["is_truncated"] for rollout in train_rollouts], - "error": [rollout["error"] for rollout in train_rollouts], - "completion_len": [get_completion_len(rollout) for rollout in train_rollouts], - "prompt_len": [get_prompt_len(rollout) for rollout in train_rollouts], + "is_filtered": [rollout["is_filtered"] for rollout in train_rollouts], + "stop_condition": [rollout.get("stop_condition") for rollout in train_rollouts], "seq_len": [get_seq_len(rollout) for rollout in train_rollouts], + "prefill_len": rollout_prefill_lens, + "decode_len": rollout_decode_lens, + "samples_per_rollout": rollout_samples_per_rollout, "num_turns": [len(rollout["trajectory"]) for rollout in train_rollouts], - "generation_ms": [rollout["timing"]["generation_ms"] for rollout in train_rollouts], - "scoring_ms": [rollout["timing"]["scoring_ms"] for rollout in train_rollouts], } ) - # Gather individual reward function metrics + # Separate DataFrames for env reward function metrics, filter flags, and per-rollout timings + # to avoid column name collisions metrics_df = pd.DataFrame([rollout["metrics"] for rollout in train_rollouts]) - - val_results_df = ( - pd.DataFrame( + filter_df = pd.DataFrame([rollout["filters"] for rollout in train_rollouts]) + timing_df = pd.DataFrame( + [ { - "example_id": [rollout["input"]["example_id"] for rollout in val_outputs], - "task": [rollout["input"]["task"] for rollout in val_outputs], - "reward": [rollout["reward"] for rollout in val_outputs], + "total": rollout["timing"]["total"], + "setup": rollout["timing"]["setup"]["duration"], + "generation": rollout["timing"]["generation"]["duration"], + "model": rollout["timing"]["model"]["duration"], + "env": rollout["timing"]["env"]["duration"], + "scoring": rollout["timing"]["scoring"]["duration"], + "overhead": rollout["timing"]["overhead"], } - ) - if val_outputs is not None - else None + for rollout in train_rollouts + ] ) - # Update progress metrics and throughput + # Update progress metrics num_tokens = int(results_df.seq_len.sum()) progress.total_tokens += num_tokens - progress.total_samples += config.batch_size - progress.total_problems += config.batch_size // config.rollouts_per_example - throughput = num_tokens / generate_completions_time - - # Compute solve all and none tensors - solve_all = ( - results_df.groupby("example_id") - .apply(lambda x: x.reward.sum() == config.rollouts_per_example, include_groups=False) - .mean() - ) - solve_none = results_df.groupby("example_id").apply(lambda x: x.reward.sum() == 0, include_groups=False).mean() - effective_batch_size = 1 - solve_none - solve_all + progress.total_samples += num_rollouts + progress.total_problems += num_unique_examples - step_time = time.perf_counter() - step_start_time + def compute_solve_rates(df): + """Compute solve_none, solve_all, effective_batch_size for a set of rollouts.""" + reward_per_problem = df.groupby(["env_name", "example_id"]).reward.sum() + solve_none = (reward_per_problem == 0).mean() + solve_all = (reward_per_problem == config.rollouts_per_example).mean() + return solve_none, solve_all, 1 - solve_none - solve_all + + # Group by (env_name, example_id) to average across rollouts within each problem + by_example = results_df.groupby(["env_name", "example_id"]) + + solve_none, solve_all, effective_batch_size = compute_solve_rates(results_df) to_log = { # Progress metrics "progress/tokens": num_tokens, - "progress/samples": config.batch_size, - "progress/problems": config.batch_size // config.rollouts_per_example, + "progress/prefill_tokens": num_prefill_tokens, + "progress/decode_tokens": num_decode_tokens, + "progress/samples": num_rollouts, + "progress/problems": num_unique_examples, "progress/total_tokens": progress.total_tokens, "progress/total_samples": progress.total_samples, "progress/total_problems": progress.total_problems, "progress/ckpt_step": ckpt_step, # Shared W&B axis # Sequence length metrics - "seq_len/mean": results_df.groupby("example_id").seq_len.mean().mean(), - "seq_len/max": results_df.groupby("example_id").seq_len.mean().max(), - "seq_len/min": results_df.groupby("example_id").seq_len.mean().min(), - "prompt_len/mean": results_df.groupby("example_id").prompt_len.mean().mean(), - "prompt_len/max": results_df.groupby("example_id").prompt_len.mean().max(), - "prompt_len/min": results_df.groupby("example_id").prompt_len.mean().min(), - "completion_len/mean": results_df.groupby("example_id").completion_len.mean().mean(), - "completion_len/max": results_df.groupby("example_id").completion_len.mean().max(), - "completion_len/min": results_df.groupby("example_id").completion_len.mean().min(), - "is_truncated/mean": results_df.groupby("example_id").is_truncated.mean().mean(), - "is_truncated/max": results_df.groupby("example_id").is_truncated.mean().max(), - "is_truncated/min": results_df.groupby("example_id").is_truncated.mean().min(), - # Turn metrics - "num_turns/mean": results_df.groupby("example_id").num_turns.mean().mean(), - "num_turns/max": results_df.groupby("example_id").num_turns.mean().max(), - "num_turns/min": results_df.groupby("example_id").num_turns.mean().min(), - # Verifier timing metrics - "generation_ms/mean": results_df.groupby("example_id").generation_ms.mean().mean(), - "generation_ms/max": results_df.groupby("example_id").generation_ms.mean().max(), - "generation_ms/min": results_df.groupby("example_id").generation_ms.mean().min(), - "scoring_ms/mean": results_df.groupby("example_id").scoring_ms.mean().mean(), - "scoring_ms/max": results_df.groupby("example_id").scoring_ms.mean().max(), - "scoring_ms/min": results_df.groupby("example_id").scoring_ms.mean().min(), - # Performance metrics - "perf/throughput": throughput, - # Train reward - "reward/mean": results_df.reward.mean(), - # Batch metrics - "batch/solve_none": solve_none, - "batch/solve_all": solve_all, - "batch/effective_batch_size": effective_batch_size, - # Error metrics - "error/mean": (~results_df.error.isna()).mean(), + "seq_len/all/mean": by_example.seq_len.mean().mean(), + "seq_len/all/max": by_example.seq_len.mean().max(), + "seq_len/all/min": by_example.seq_len.mean().min(), + "prefill_len/all/mean": by_example.prefill_len.mean().mean(), + "prefill_len/all/max": by_example.prefill_len.mean().max(), + "prefill_len/all/min": by_example.prefill_len.mean().min(), + "decode_len/all/mean": by_example.decode_len.mean().mean(), + "decode_len/all/max": by_example.decode_len.mean().max(), + "decode_len/all/min": by_example.decode_len.mean().min(), + "is_truncated/all/mean": by_example.is_truncated.mean().mean(), + "is_truncated/all/max": by_example.is_truncated.mean().max(), + "stop_condition/all/generation_truncated": ( + results_df.is_truncated & (results_df.stop_condition != "prompt_too_long") + ).mean(), + **{ + f"stop_condition/all/{sc}": rate + for sc, rate in results_df.stop_condition.dropna().value_counts(normalize=True).items() + }, + "samples_per_rollout/all/mean": by_example.samples_per_rollout.mean().mean(), + "samples_per_rollout/all/max": by_example.samples_per_rollout.mean().max(), + "samples_per_rollout/all/min": by_example.samples_per_rollout.mean().min(), + "num_turns/all/mean": by_example.num_turns.mean().mean(), + "num_turns/all/max": by_example.num_turns.mean().max(), + "num_turns/all/min": by_example.num_turns.mean().min(), **{ - f"error/{error}": error_rate - for error, error_rate in results_df.error.dropna().value_counts(normalize=True).items() + f"timing/all/{key}/{stat}": getattr( + timing_df[key].groupby([results_df.env_name, results_df.example_id]).mean(), + stat, + )() + for key in timing_df.columns + for stat in ("mean", "max", "min") }, - # Env metrics - **{f"metrics/{metric}": metrics_df[metric].mean() for metric in metrics_df.columns}, + # Train reward + "reward/all/mean": by_example.reward.mean().mean(), + "reward/all/max": by_example.reward.mean().max(), + "reward/all/min": by_example.reward.mean().min(), + # Solve / batch metrics + "solve_none/all": solve_none, + "solve_all/all": solve_all, + "effective_batch_size/all": effective_batch_size, + **{f"batch/{env}": r for env, r in results_df.env_name.value_counts(normalize=True).items()}, # Time metrics "time/step": step_time, "time/generate_completions": generate_completions_time, "time/teacher_logprobs": teacher_logprobs_time, "time/save_ckpt": save_ckpt_time, + "time/parallel_preprocess": parallel_preprocess_time, # Scheduler metrics **scheduler.get_metrics(), # Buffer metrics **buffer.get_metrics(), # Event loop lag metrics **event_loop_lag_monitor.get_metrics(), + # Rollout filter metrics (detection rate per filter + overall drop rate) + "filters/all/is_filtered": results_df.is_filtered.astype(float).mean(), + **{f"filters/all/{name}": filter_df[name].astype(float).mean() for name in filter_df.columns}, # W&B axis "step": progress.step, } - # If more than one env, add per-env metrics - if results_df.task.nunique() > 1: - per_env_reward = results_df.groupby("task").reward.mean().to_dict() - to_log.update({f"reward/{env}": reward for env, reward in per_env_reward.items()}) - - per_env_ratio = results_df.task.value_counts(normalize=True).to_dict() - to_log.update({f"batch/{env}": ratio for env, ratio in per_env_ratio.items()}) - - # Optionally, add val metrics - if val_results_df is not None: - to_log.update({"val_reward/mean": val_results_df.reward.mean()}) - - if val_results_df.task.nunique() > 1: - per_env_reward = val_results_df.groupby("task").reward.mean().to_dict() - to_log.update({f"val_reward/{env}": reward for env, reward in per_env_reward.items()}) - - per_env_ratio = val_results_df.task.value_counts(normalize=True).to_dict() - to_log.update({f"val_batch/{env}": ratio for env, ratio in per_env_ratio.items()}) + # Per-env metrics + per_env_columns = [ + "seq_len", + "prefill_len", + "decode_len", + "is_truncated", + "samples_per_rollout", + "num_turns", + ] + + for env, env_df in results_df.groupby("env_name"): + env_by_example = env_df.groupby("example_id") + for col in per_env_columns: + to_log[f"{col}/{env}/mean"] = env_by_example[col].mean().mean() + to_log[f"{col}/{env}/max"] = env_by_example[col].mean().max() + if col != "is_truncated": + to_log[f"{col}/{env}/min"] = env_by_example[col].mean().min() + env_timing_df = timing_df.loc[env_df.index] + for key in timing_df.columns: + per_example = env_timing_df.groupby(env_df["example_id"])[key].mean() + to_log[f"timing/{env}/{key}/mean"] = per_example.mean() + to_log[f"timing/{env}/{key}/max"] = per_example.max() + to_log[f"timing/{env}/{key}/min"] = per_example.min() + to_log[f"reward/{env}/mean"] = env_by_example.reward.mean().mean() + to_log[f"reward/{env}/max"] = env_by_example.reward.mean().max() + to_log[f"reward/{env}/min"] = env_by_example.reward.mean().min() + solve_none, solve_all, effective_batch_size = compute_solve_rates(env_df) + to_log[f"solve_none/{env}"] = solve_none + to_log[f"solve_all/{env}"] = solve_all + to_log[f"effective_batch_size/{env}"] = effective_batch_size + to_log[f"stop_condition/{env}/generation_truncated"] = ( + env_df.is_truncated & (env_df.stop_condition != "prompt_too_long") + ).mean() + for sc, rate in env_df.stop_condition.dropna().value_counts(normalize=True).items(): + to_log[f"stop_condition/{env}/{sc}"] = rate + env_metrics_df = metrics_df.loc[env_df.index] + for metric in metrics_df.columns: + to_log[f"metrics/{env}/{metric}"] = env_metrics_df.groupby(env_df["example_id"])[metric].mean().mean() + to_log[f"filters/{env}/is_filtered"] = env_df.is_filtered.astype(float).mean() + env_filter_df = filter_df.loc[env_df.index] + for name in filter_df.columns: + to_log[f"filters/{env}/{name}"] = env_filter_df[name].astype(float).mean() # Log metrics to monitor(s) - monitor.log(to_log) + monitor.log(to_log, step=progress.step) - # Log samples to monitor(s) if enabled - subset_train_rollouts = random.sample(train_rollouts, min(8, len(train_rollouts))) - monitor.log_samples(subset_train_rollouts, step=progress.step) + # Log samples to monitor(s) if enabled. + monitor.log_samples(train_rollouts, step=progress.step) # Log distributions (rewards, advantages) if enabled monitor.log_distributions( distributions={ - "rewards": rewards, - "advantages": advantages, + "rewards": [r["reward"] for r in train_rollouts], + "advantages": [r["advantage"] for r in train_rollouts], }, step=progress.step, ) - step_message = f"Step {progress.step} | Time: {step_time:.2f}s | Reward: {results_df.reward.mean():.4f} |{f' Val. Reward: {val_results_df.reward.mean():.4f} |' if val_results_df is not None else ''} Throughput: {throughput:.1f} tokens/s | Seq. Length: {results_df.groupby('example_id').seq_len.mean().mean():.1f} tokens/sample | Async Level: {scheduler.async_level} | Max. Off-Policy Level: {scheduler.max_off_policy_level}" + if usage_reporter and run_id: + usage_reporter.report_training_usage( + run_id=run_id, + step=progress.step, + tokens=num_prefill_tokens + num_decode_tokens, + ) + + reward_mean = by_example.reward.mean().mean() + step_message = f"Step {progress.step} | Time: {step_time:.2f}s | Reward: {reward_mean:.4f} | Seq. Length: {by_example.seq_len.mean().mean():.1f} tokens/sample | Max. Off-Policy Level: {scheduler.max_off_policy_level}" logger.success(step_message) # Increment step progress.step += 1 is_first_step = False + # Free large per-step objects to prevent memory accumulation + del train_rollouts, train_examples, training_batch, vlm_cache + del results_df, metrics_df + gc.collect() + event_loop_lag_monitor.reset() # Send heartbeat if configured if heart is not None: heart.beat() - if config.eval: - logger.info("Running final evals (subprocess)") - await run_evals_subprocess( - client_config=config.client, - eval_config=config.eval, - model_config=config.model, - sampling_config=config.eval.sampling, - reasoning_field=config.eval.reasoning_field, - output_dir=config.output_dir, - ckpt_step=scheduler.ckpt_step, - step=progress.step, - max_concurrent=config.max_concurrent or -1, + if config.eval and eval_envs is not None: + logger.info("Running final evals") + eval_results = await asyncio.gather( + *[ + eval_env.evaluate( + model_name=scheduler.model_name, + get_client=inference_pool.get_eval_client, + ckpt_step=ckpt_step, + step=progress.step, + cache_salt=str(ckpt_step), + ) + for eval_env in eval_envs + ] ) - # Log final (immutable) samples and distributions to monitor(s) - monitor.log_final_samples() + # Save final eval rollouts to disk + eval_rollouts = [o for outputs in eval_results for o in outputs] + if eval_rollouts: + step_path = get_step_path(get_rollout_dir(config.output_dir), progress.step) + await asyncio.to_thread( + save_rollouts, eval_rollouts, step_path / "eval_rollouts.jsonl", exclude_keys={"trajectory"} + ) + monitor.save_final_summary() # Write final checkpoint @@ -571,14 +834,44 @@ async def orchestrate(config: OrchestratorConfig): logger.info("Writing final checkpoint") ckpt_manager.save(progress, buffer, step=progress.step) - # Close training batch sender - training_batch_sender.close() + # Bounded best-effort cleanup. Each await below may block on a remote peer + # (env-server ZMQ recv, inference admin httpx aclose, etc.). The outer + # asyncio.wait gives the whole sequence a single deadline; if anything + # wedges past SHUTDOWN_TIMEOUT_S we force-exit the process. Individual + # awaits intentionally do NOT have their own timeouts — asyncio.wait_for + # would itself hang on an uncancellable await, which is exactly the + # failure mode we're guarding against. + async def _graceful_shutdown() -> None: + training_batch_sender.close() + await scheduler.stop() + if inference_metrics_collector is not None: + await inference_metrics_collector.stop() + await inference_pool.stop() + if teacher_inference_pool is not None: + await teacher_inference_pool.stop() + event_loop_lag_monitor_task.cancel() + # Shutdown env processes (also registered as atexit handler for crash safety) + train_envs.shutdown() + if eval_envs is not None: + eval_envs.shutdown() + + shutdown_task = asyncio.create_task(_graceful_shutdown()) + _, pending = await asyncio.wait({shutdown_task}, timeout=SHUTDOWN_TIMEOUT_S) + + if pending: + logger.warning( + f"Orchestrator shutdown did not complete within {SHUTDOWN_TIMEOUT_S}s; " + "forcing process exit. Training artifacts are already persisted." + ) + os._exit(0) - # Stop env workers - await scheduler.stop() + # asyncio.wait swallows task exceptions; re-raise so a fast cleanup + # failure surfaces the same way as it did when each step was awaited + # directly. + await shutdown_task - # Cancel event loop lag monitor task - event_loop_lag_monitor_task.cancel() + if usage_reporter: + usage_reporter.close() logger.success("Orchestrator finished.") @@ -589,8 +882,77 @@ async def orchestrate(config: OrchestratorConfig): def main(): """Main entry-point for orchestrator. Run using `uv run orchestrator`""" + set_proc_title("Orchestrator") + asyncio.run(orchestrate(cli(OrchestratorConfig))) + + +async def setup_rollout_inference_pool( + *, + config: OrchestratorConfig, + rollout_client_config, + rollout_model_name: str, + tokenizer, + logger, +): + """Set up rollout inference. + + Routing policy is driven by ``config.use_token_client`` and + ``config.use_renderer`` (mutually exclusive — config-level validators + block both being True): + + - external teacher rollout → MITO (``openai_chat_completions``), + forced regardless of the toggles (config-level validator + rejects ``use_token_client`` / ``use_renderer`` in that case) + - ``use_renderer=True`` → renderer client (``/v1/generate``). + Not allowed for VLMs (validated at config time). + - ``use_token_client=True`` → TITO + (``openai_chat_completions_token``, ``/v1/chat/completions/tokens``). + Default. VLMs land here too. + - both False → MITO (``openai_chat_completions``). + """ + if config.teacher_rollout_model is not None: + logger.info("Using external rollout model (MITO) without renderer client") + inference_pool = await setup_inference_pool( + rollout_client_config, + model_name=rollout_model_name, + train_client_type="openai_chat_completions", + eval_client_type="openai_chat_completions", + ) + return None, inference_pool + + if config.use_renderer: + renderer = create_renderer( + tokenizer, + renderer=config.renderer.name, + tool_parser=config.renderer.tool_parser, + reasoning_parser=config.renderer.reasoning_parser, + ) + logger.info(f"Initialized {type(renderer).__name__} for {config.model.name}") + inference_pool = await setup_inference_pool( + rollout_client_config, + model_name=rollout_model_name, + train_client_type="renderer", + eval_client_type="openai_chat_completions", + renderer_name=config.renderer.name, + tool_parser=config.renderer.tool_parser, + reasoning_parser=config.renderer.reasoning_parser, + renderer_pool_size=config.renderer.pool_size, + ) + logger.info("Using direct renderer rollout client") + return renderer, inference_pool - asyncio.run(orchestrate(parse_argv(OrchestratorConfig))) + train_client_type = "openai_chat_completions_token" if config.use_token_client else "openai_chat_completions" + if config.use_token_client: + logger.info("Using token client (TITO) for rollouts — server-side templating, /v1/chat/completions/tokens") + else: + logger.info("Using MITO (openai_chat_completions) for rollouts") + inference_pool = await setup_inference_pool( + rollout_client_config, + model_name=rollout_model_name, + train_client_type=train_client_type, + eval_client_type="openai_chat_completions", + ) + return None, inference_pool if __name__ == "__main__": diff --git a/src/prime_rl/orchestrator/scheduler.py b/src/prime_rl/orchestrator/scheduler.py index 4092360fbde..a8427b69f72 100644 --- a/src/prime_rl/orchestrator/scheduler.py +++ b/src/prime_rl/orchestrator/scheduler.py @@ -1,25 +1,20 @@ -""" -Scheduler that runs environments in subprocesses. - -Isolates event loop lag from environment execution. -""" +from __future__ import annotations import asyncio import time -from pathlib import Path -from typing import NamedTuple +from collections import Counter, defaultdict +from dataclasses import dataclass, field -from httpx import AsyncClient -from tqdm import tqdm +import verifiers as vf +from aiolimiter import AsyncLimiter +from prime_rl.configs.orchestrator import OrchestratorConfig from prime_rl.orchestrator.buffer import Buffer -from prime_rl.orchestrator.config import EnvConfig, OrchestratorConfig -from prime_rl.orchestrator.env_worker import EnvWorker, WorkerDiedError -from prime_rl.orchestrator.utils import get_sampling_args -from prime_rl.utils.client import update_weights -from prime_rl.utils.config import ClientConfig -from prime_rl.utils.logger import get_logger -from prime_rl.utils.pathing import get_env_worker_log_file +from prime_rl.orchestrator.envs import TrainEnvs +from prime_rl.orchestrator.vf_utils import get_seq_len +from prime_rl.utils.async_utils import safe_cancel, safe_cancel_all +from prime_rl.utils.client import InferencePool +from prime_rl.utils.logger import ProgressTracker, get_logger from prime_rl.utils.utils import ( get_broadcast_dir, get_latest_ckpt_step, @@ -28,19 +23,32 @@ ) -class InflightRolloutInfo(NamedTuple): - """Metadata for an in-flight group rollout request.""" +@dataclass +class InflightRequest: + """Metadata for an in-flight request.""" off_policy_steps: int - worker: EnvWorker - request_id: str + client_config: vf.ClientConfig + env_name: str + group_id: int | None = None + rollout_count: int = 1 -class Scheduler: - """Asynchronously schedules group rollout requests using subprocess workers. +@dataclass +class GroupState: + """Tracks the state of a rollout group (one example × N rollouts).""" + + example: dict + rollouts_to_schedule: int + completed_rollouts: list[vf.RolloutOutput] = field(default_factory=list) + pinned_client: vf.ClientConfig | None = None - Runs environment execution in separate processes to isolate event loop lag - from the main orchestrator process. + +class Scheduler: + """ + Asynchronously manages scheduling of rollout requests and policy updates. + Keeps a constant number of rollouts in-flight (continuous batching) and + updates the policy as soon as it becomes available. References: - AReal: https://arxiv.org/abs/2505.24298v1 @@ -49,291 +57,462 @@ class Scheduler: def __init__( self, - admin_clients: list[AsyncClient], - client_config: ClientConfig, - env_configs: list[EnvConfig], + train_envs: TrainEnvs, + inference_pool: InferencePool, buffer: Buffer, config: OrchestratorConfig, - oversampling_factor: float, + max_inflight_rollouts: int, max_async_level: int, max_off_policy_steps: int, strict_async_level: bool, + tasks_per_minute: int | None, + enable_policy_updates: bool = True, lora_name: str | None = None, - output_dir: Path | None = None, ): self.logger = get_logger() - self.admin_clients = admin_clients - self.client_config = client_config + if tasks_per_minute is not None: + self.rate_limiter = AsyncLimiter(max_rate=tasks_per_minute, time_period=60) + else: + self.rate_limiter = None + self.train_envs = train_envs self.buffer = buffer self.config = config self.batch_size = config.batch_size + self.token_batch_size = config.token_batch_size self.rollouts_per_example = config.rollouts_per_example - self.seq_len = config.seq_len - self.problems_per_batch = int(oversampling_factor * self.batch_size // self.rollouts_per_example) + self.max_inflight_rollouts = max_inflight_rollouts self.max_async_level = max_async_level self.max_off_policy_steps = max_off_policy_steps self.strict_async_level = strict_async_level + self.enable_policy_updates = enable_policy_updates self.lora_name = lora_name - self.sampling_args = get_sampling_args(config.sampling) self.model_name = self.config.model.name + self.json_logging = config.log.json_logging - # Build example lookup dicts per env (example_id -> example) - self.example_lookups: dict[str, dict[int, dict]] = {} - for env_config in env_configs: - env_name = env_config.name or env_config.id - self.example_lookups[env_name] = buffer.example_buffer[env_name].copy() - - # Create workers - multiple per env - self.workers_per_env = config.workers_per_env or 1 - self.workers: dict[str, list[EnvWorker]] = {} - self.env_names: list[str] = [] - for env_config in env_configs: - env_name = env_config.name or env_config.id - self.env_names.append(env_name) - self.workers[env_name] = [] - - # Setup log file if env worker file logging is enabled (all workers share one file) - env_log_file = None - if config.log.env_worker_logs and output_dir is not None: - env_log_file = get_env_worker_log_file(output_dir, env_name) - env_log_file.parent.mkdir(parents=True, exist_ok=True) - - for worker_idx in range(self.workers_per_env): - worker = EnvWorker( - env_id=env_config.id, - env_args=env_config.args, - client_config=client_config, - model_name=self.model_name, - seq_len=config.seq_len, - interleaved_rollouts=config.trajectory_strategy == "interleaved", - max_concurrent=config.max_concurrent or -1, - example_lookup=self.example_lookups[env_name], - sampling_args=self.sampling_args, - worker_name=f"{env_name}_{worker_idx}", - log_level=config.log.level, - vf_log_level=config.log.vf_level, - log_file=str(env_log_file) if env_log_file else None, - max_restarts=config.max_env_worker_restarts, - ) - self.workers[env_name].append(worker) + # Inference pool - used for admin operations (adapter sync) and metrics + self.inference_pool = inference_pool + + group_scoring_envs = [env.name for env in train_envs if env.requires_group_scoring] + if group_scoring_envs: + self.logger.info(f"Group rollout scoring active for env(s): {', '.join(group_scoring_envs)}") - # Track in-flight requests: future -> info - self.inflight_group_rollouts: dict[asyncio.Future, InflightRolloutInfo] = {} + # Track in-flight requests: task -> info + self.inflight_requests: dict[asyncio.Task, InflightRequest] = {} + + # Track in-progress groups while rollouts are generated independently. + self.next_group_id = 0 + self.groups: dict[int, GroupState] = {} self.step, self.ckpt_step = 0, 0 self.checkpoint_ready = asyncio.Event() self.checkpoint_ready.set() self.update_weights_time, self.wait_for_ckpt_time = 0, 0 + self.update_policy_task: asyncio.Task | None = None + self.inflight_policy_update_task: asyncio.Task | None = None + self.policy_update_lock = asyncio.Lock() self.cancelled_rollouts_count = 0 + self.empty_rollouts_by_env: dict[str, int] = defaultdict(int) + self.errored_rollouts_by_env: dict[str, int] = defaultdict(int) + self.total_rollouts_by_env: dict[str, int] = defaultdict(int) + self.last_batch_generation_time = 0.0 - # Background tasks - self._response_collectors: list[asyncio.Task] = [] - - async def start(self): - """Start all workers and response collectors.""" - total_workers = sum(len(workers) for workers in self.workers.values()) - self.logger.info(f"Starting {total_workers} env worker(s) ({self.workers_per_env} per env)") - for workers in self.workers.values(): - for worker in workers: - worker.start() - # Start response collector for each worker - task = asyncio.create_task(worker.collect_responses()) - self._response_collectors.append(task) - - async def stop(self): - """Stop all workers and collectors.""" - for task in self._response_collectors: - task.cancel() - for workers in self.workers.values(): - for worker in workers: - worker.stop() - - async def schedule_group_rollout(self): - """Asynchronously schedules a group rollout request.""" - example = self.buffer.sample_examples(n=1)[0] - - # Route to worker for this example's environment - task = example["task"] - workers = self.workers[task] - worker = min(workers, key=lambda w: w.pending_count) + @property + def uses_token_batching(self) -> bool: + return self.token_batch_size is not None - future, request_id = await worker.submit_request( - example_id=example["example_id"], - rollouts_per_example=self.config.rollouts_per_example, + @property + def batch_target(self) -> int: + if self.uses_token_batching: + assert self.token_batch_size is not None + return self.token_batch_size + assert self.batch_size is not None + return self.batch_size + + def get_batch_progress_increment(self, rollouts: list[vf.RolloutOutput]) -> int: + if self.uses_token_batching: + return sum(get_seq_len(rollout) for rollout in rollouts) + return len(rollouts) + + def finalize_batch_rollouts(self, rollouts: list[vf.RolloutOutput]) -> list[vf.RolloutOutput]: + if self.batch_size is None: + return rollouts + return rollouts[: self.batch_size] + + async def cancel_inflight_rollouts(self): + """Cancel all in-flight rollout requests.""" + count = sum(info.rollout_count for info in self.inflight_requests.values()) + await safe_cancel_all(list(self.inflight_requests)) + self.inflight_requests.clear() + self.groups.clear() + self.cancelled_rollouts_count += count + + @staticmethod + def _client_identity(c: vf.ClientConfig) -> tuple[str, str | None]: + return ( + c.api_base_url, + c.extra_headers.get("X-data-parallel-rank"), ) - self.inflight_group_rollouts[future] = InflightRolloutInfo( + async def _select_least_loaded_client(self) -> vf.ClientConfig: + """Select the client with the fewest in-flight tasks. + + Uses (api_base_url, dp_rank) as identity rather than client_idx so that + load tracking survives elastic pool refreshes (which reassign indices). + """ + clients = self.inference_pool.train_clients + while not clients: + await asyncio.sleep(1) + clients = self.inference_pool.train_clients + inflight = Counter(self._client_identity(info.client_config) for info in self.inflight_requests.values()) + return min(clients, key=lambda c: inflight[self._client_identity(c)]) + + async def drop_group(self, group_id: int) -> int: + """Drop a group and cancel any remaining in-flight rollouts for it. Returns the number of cancelled rollouts.""" + tasks_to_cancel = [] + rollout_count = 0 + for task, info in list(self.inflight_requests.items()): + if info.group_id != group_id: + continue + self.inflight_requests.pop(task, None) + tasks_to_cancel.append(task) + rollout_count += info.rollout_count + self.groups.pop(group_id, None) + await safe_cancel_all(tasks_to_cancel) + return rollout_count + + async def schedule_rollout(self, group_id: int): + """Asynchronously schedules a rollout request (or a group request for group-scoring envs).""" + if self.rate_limiter: + await self.rate_limiter.acquire() + group = self.groups.get(group_id) + if group is None or group.rollouts_to_schedule <= 0: + return + + if group.pinned_client is not None: + client_config = group.pinned_client + else: + client_config = await self._select_least_loaded_client() + if group_id not in self.groups: + return + group.pinned_client = client_config + + env_name = group.example["env_name"] + env = self.train_envs.get(env_name) + + cache_salt = str(self.ckpt_step) + if env.requires_group_scoring: + rollout_count = group.rollouts_to_schedule + group.rollouts_to_schedule = 0 + task = asyncio.create_task( + env.run_group( + client=client_config, + example=group.example, + model_name=self.model_name, + rollouts_per_example=rollout_count, + cache_salt=cache_salt, + ) + ) + else: + rollout_count = 1 + group.rollouts_to_schedule -= 1 + task = asyncio.create_task( + env.run_rollout( + client=client_config, + example=group.example, + model_name=self.model_name, + cache_salt=cache_salt, + ) + ) + self.inflight_requests[task] = InflightRequest( off_policy_steps=0, - worker=worker, - request_id=request_id, + client_config=client_config, + env_name=env_name, + group_id=group_id, + rollout_count=rollout_count, ) + @property + def inflight_rollout_count(self) -> int: + return sum(info.rollout_count for info in self.inflight_requests.values()) + + @property + def inflight_sample_count(self) -> int: + pending = sum(g.rollouts_to_schedule for g in self.groups.values()) + return self.inflight_rollout_count + pending + + async def _schedule_next_request(self) -> bool: + remaining_capacity = self.max_inflight_rollouts - self.inflight_rollout_count + + if remaining_capacity <= 0: + return False + + for group_id, group in self.groups.items(): + if group.rollouts_to_schedule <= 0: + continue + env = self.train_envs.get(group.example["env_name"]) + cost = group.rollouts_to_schedule if env.requires_group_scoring else 1 + if cost <= remaining_capacity: + await self.schedule_rollout(group_id=group_id) + return True + + if remaining_capacity < self.rollouts_per_example: + return False + + example = self.buffer.sample_examples(n=1)[0] + group_id = self.next_group_id + self.next_group_id += 1 + self.groups[group_id] = GroupState(example=example, rollouts_to_schedule=self.rollouts_per_example) + await self.schedule_rollout(group_id=group_id) + return True + + async def _fill_inflight_requests(self) -> None: + while await self._schedule_next_request(): + pass + async def update_policy_loop(self): """Continuously checks for new policy checkpoints.""" while True: - await self.update_policy() + await self.maybe_update_policy() await asyncio.sleep(1) - async def update_policy(self): - """Updates the policy to the latest available checkpoint. Aborts rollout requests that are older than the max retention steps.""" + def _compute_next_ckpt_step(self) -> int: latest_ckpt_step = get_latest_ckpt_step(get_broadcast_dir(self.config.output_dir)) or 0 async_away_ckpt_step = max(self.step - self.max_async_level, 0) - next_ckpt_step = ( - async_away_ckpt_step if self.strict_async_level else max(async_away_ckpt_step, latest_ckpt_step) - ) + if self.strict_async_level: + return async_away_ckpt_step + return max(async_away_ckpt_step, latest_ckpt_step) - if next_ckpt_step > self.ckpt_step: - if next_ckpt_step == async_away_ckpt_step: - self.logger.info( - f"Orchestrator paused: waiting for trainer process to complete checkpoint {next_ckpt_step} " - f"(>{self.max_async_level} step(s) ahead). Training is progressing normally." - ) - self.checkpoint_ready.clear() - wait_for_ckpt_start_time = time.perf_counter() - await wait_for_path(get_step_path(get_broadcast_dir(self.config.output_dir), next_ckpt_step) / "STABLE") - self.wait_for_ckpt_time = time.perf_counter() - wait_for_ckpt_start_time - self.logger.debug(f"Waited for checkpoint {next_ckpt_step} for {self.wait_for_ckpt_time:.2f}s") - - self.logger.debug( - f"Got new policy with step {next_ckpt_step}. Updating weights and cancelling old rollout requests." + async def _apply_policy_update(self, next_ckpt_step: int) -> None: + async_away_ckpt_step = max(self.step - self.max_async_level, 0) + if next_ckpt_step == async_away_ckpt_step: + self.logger.info( + f"Orchestrator paused: waiting for trainer process to complete checkpoint {next_ckpt_step} " + f"(>{self.max_async_level} step(s) ahead). Training is progressing normally." ) - - # Update weights on inference servers - update_weights_start_time = time.perf_counter() - await update_weights( - self.admin_clients, - get_step_path(get_broadcast_dir(self.config.output_dir), next_ckpt_step), - lora_name=self.lora_name, + self.checkpoint_ready.clear() + wait_for_ckpt_start_time = time.perf_counter() + await wait_for_path(get_step_path(get_broadcast_dir(self.config.output_dir), next_ckpt_step) / "STABLE") + self.wait_for_ckpt_time = time.perf_counter() - wait_for_ckpt_start_time + self.logger.info( + f"Orchestrator resumed: checkpoint {next_ckpt_step} ready (after {self.wait_for_ckpt_time:.2f}s)" ) - self.update_weights_time = time.perf_counter() - update_weights_start_time - self.logger.debug(f"Updated weights to step {next_ckpt_step} in {self.update_weights_time:.2f}s") - if self.lora_name is not None: - self.model_name = self.lora_name + self.logger.debug( + f"Got new policy with step {next_ckpt_step}. Updating weights and cancelling old rollout requests." + ) - # Update model name on all workers - for workers in self.workers.values(): - for worker in workers: - worker.update_model_name(self.model_name) + update_weights_start_time = time.perf_counter() + weights_path = get_step_path(get_broadcast_dir(self.config.output_dir), next_ckpt_step) + await self.inference_pool.update_weights(weights_path, lora_name=self.lora_name, step=next_ckpt_step) + self.update_weights_time = time.perf_counter() - update_weights_start_time + self.logger.debug(f"Updated weights to step {next_ckpt_step} in {self.update_weights_time:.2f}s") - self.checkpoint_ready.set() + self.ckpt_step = next_ckpt_step + if self.lora_name is not None: + self.model_name = self.lora_name + self.inference_pool.update_model_name(self.model_name) - # Handle off-policy tracking - cancel old requests - futures_to_remove = [] - futures_to_update = [] - - for future, info in self.inflight_group_rollouts.items(): - if info.off_policy_steps > self.max_off_policy_steps: - if not future.done(): - future.cancel() - futures_to_remove.append((future, info.worker)) - else: - futures_to_update.append((future, info.off_policy_steps + 1, info.worker, info.request_id)) - - # Remove cancelled - for future, worker in futures_to_remove: - self.inflight_group_rollouts.pop(future, None) - self.cancelled_rollouts_count += len(futures_to_remove) - - # Update off-policy steps for remaining - for future, off_policy_steps, worker, request_id in futures_to_update: - if future in self.inflight_group_rollouts: - self.inflight_group_rollouts[future] = InflightRolloutInfo( - off_policy_steps=off_policy_steps, - worker=worker, - request_id=request_id, - ) - - if len(futures_to_remove) > 0: - self.logger.warning( - f"Cancelled {len(futures_to_remove)} old rollout requests (will refill naturally). Consider increasing max_off_policy_steps to avoid this." - ) + self.checkpoint_ready.set() + await self._update_off_policy() - self.ckpt_step = next_ckpt_step + async def _get_or_start_policy_update_task(self, next_ckpt_step: int) -> asyncio.Task: + async with self.policy_update_lock: + task = self.inflight_policy_update_task + if task is not None and not task.done(): + return task - async def generate_batch(self, step: int) -> list[dict]: - """Generate a batch of rollouts using workers. + task = asyncio.create_task(self._apply_policy_update(next_ckpt_step)) + self.inflight_policy_update_task = task - Returns list of result dicts (not vf.State, since those stay in workers). - """ + def _clear_inflight_policy_update(done_task: asyncio.Task) -> None: + if self.inflight_policy_update_task is done_task: + self.inflight_policy_update_task = None + + task.add_done_callback(_clear_inflight_policy_update) + return task + + async def maybe_update_policy(self): + """Updates the policy to the latest available checkpoint. Aborts rollout requests that are older than the max retention steps.""" + if not self.enable_policy_updates: + self.ckpt_step = self.step + self.checkpoint_ready.set() + return + + while True: + next_ckpt_step = self._compute_next_ckpt_step() + if next_ckpt_step <= self.ckpt_step: + return + + task = await self._get_or_start_policy_update_task(next_ckpt_step) + await asyncio.shield(task) + + async def _update_off_policy(self) -> None: + stale_group_ids = { + info.group_id + for info in self.inflight_requests.values() + if info.group_id is not None and info.off_policy_steps >= self.max_off_policy_steps + } + tasks_to_increment = [ + task + for task, info in list(self.inflight_requests.items()) + if info.group_id is None or info.group_id not in stale_group_ids + ] + + counts = await asyncio.gather(*(self.drop_group(gid) for gid in stale_group_ids)) + removed = sum(counts) + for task in tasks_to_increment: + info = self.inflight_requests.get(task) + if info is None: + continue + info.off_policy_steps += 1 + + self.cancelled_rollouts_count += removed + if removed: + self.logger.warning( + f"Cancelled {removed} old rollout requests (will refill naturally). " + f"Consider increasing max_off_policy_steps to avoid this." + ) + + async def generate_batch(self, step: int) -> list[vf.RolloutOutput]: + """Continuously generates a batch of rollouts.""" self.step = step - # Schedule initial tasks + if self.enable_policy_updates: + # Cancel the previous update policy task to avoid concurrent updates + if self.update_policy_task is not None: + await safe_cancel(self.update_policy_task) + + # Manually check the async barrier before starting the step, then re-create the update policy loop + # This ensures that we respect max_async_level, while still listening for policy updates mid-step + await self.maybe_update_policy() + self.update_policy_task = asyncio.create_task(self.update_policy_loop()) + else: + self.ckpt_step = step + self.checkpoint_ready.set() + + batch_start_time = time.perf_counter() + self.logger.debug("Starting to generate batch rollouts") - while len(self.inflight_group_rollouts) < self.problems_per_batch: - await self.schedule_group_rollout() - batch_rollouts: list[dict] = [] - pbar = tqdm(total=self.config.batch_size, desc="Generating rollouts (train)") + batch_rollouts: list[vf.RolloutOutput] = [] + batch_progress = 0 + pbar = ProgressTracker( + total=self.batch_target, desc="Generating rollouts (train)", json_logging=self.json_logging, step=step + ) - while len(batch_rollouts) < self.config.batch_size: - # Wait for at least one future to complete - done, _ = await asyncio.wait( - self.inflight_group_rollouts.keys(), + while batch_progress < self.batch_target: + await self._fill_inflight_requests() + inflight_tasks = list(self.inflight_requests.keys()) + + finished_tasks, _ = await asyncio.wait( + inflight_tasks, return_when=asyncio.FIRST_COMPLETED, ) - await self.checkpoint_ready.wait() - for finished_future in done: - if len(batch_rollouts) >= self.config.batch_size: - batch_rollouts = batch_rollouts[: self.config.batch_size] + for finished_task in finished_tasks: + if batch_progress >= self.batch_target: break - # Safely pop the future from tracking - if self.inflight_group_rollouts.pop(finished_future, None) is None: + rollout_info = self.inflight_requests.pop(finished_task, None) + if rollout_info is None: continue - try: - group_results: list[dict] = finished_future.result() - - # Update buffer with results - self.buffer.update(group_results) - accepted_rollouts = self.buffer.sample_rollouts(n=self.config.rollouts_per_example) + group_id = rollout_info.group_id + env_name = rollout_info.env_name - batch_rollouts.extend(accepted_rollouts) - pbar.update(len(accepted_rollouts)) + try: + group = self.groups.get(group_id) + if group is None: + continue + + env = self.train_envs.get(env_name) + result = finished_task.result() + rollouts: list[vf.RolloutOutput] = result if isinstance(result, list) else [result] + self.total_rollouts_by_env[env_name] += len(rollouts) + + # Check for empty/errored rollouts and reschedule + valid_rollouts = [] + has_failures = False + for rollout in rollouts: + if rollout["error"] is not None: + self.errored_rollouts_by_env[env_name] += 1 + has_failures = True + self.logger.warning( + f"Rollout error in group {group_id} ({env_name}), re-scheduling " + f"({len(group.completed_rollouts)}/{self.rollouts_per_example} complete): " + f"{rollout['error']['error_chain_repr']}" + ) + elif len(rollout["trajectory"]) == 0: + self.empty_rollouts_by_env[env_name] += 1 + has_failures = True + self.logger.warning( + f"Empty trajectory in group {group_id} ({env_name}), re-scheduling " + f"({len(group.completed_rollouts)}/{self.rollouts_per_example} complete)" + ) + else: + rollout["env_name"] = env_name + valid_rollouts.append(rollout) + + if has_failures and env.requires_group_scoring: + # Group scoring requires all rollouts — discard partial results, reschedule full group + group.completed_rollouts.clear() + group.rollouts_to_schedule = self.rollouts_per_example + continue + + # For individual scoring, reschedule only the failed ones + group.rollouts_to_schedule += len(rollouts) - len(valid_rollouts) + group.completed_rollouts.extend(valid_rollouts) + if len(group.completed_rollouts) < self.rollouts_per_example: + continue + completed_rollouts = self.groups.pop(group_id).completed_rollouts except asyncio.CancelledError: - pass # Request was cancelled, will be rescheduled - except WorkerDiedError as e: - # Worker died - check if it will auto-restart or if orchestrator should crash - self.logger.warning(f"Rollout lost due to worker death: {e}") - # Check if any worker has permanently failed (exceeded max restarts) - self._check_for_fatal_worker_errors() + if group_id is not None: + await self.drop_group(group_id) + continue except Exception as e: self.logger.warning(f"Rollout failed: {e}") + if group_id is not None: + await self.drop_group(group_id) + continue + + self.buffer.update(completed_rollouts) + accepted_rollouts = self.buffer.sample_rollouts(n=self.rollouts_per_example) - await self.schedule_group_rollout() + batch_rollouts.extend(accepted_rollouts) + progress_increment = self.get_batch_progress_increment(accepted_rollouts) + batch_progress += progress_increment + pbar.update(progress_increment) + await self._fill_inflight_requests() + + batch_rollouts = self.finalize_batch_rollouts(batch_rollouts) pbar.close() + self.last_batch_generation_time = time.perf_counter() - batch_start_time return batch_rollouts - def _check_for_fatal_worker_errors(self): - """Check if any worker has permanently failed and raise the error to crash orchestrator.""" - for workers in self.workers.values(): - for worker in workers: - if worker._fatal_error is not None: - self.logger.error(f"Worker '{worker.worker_name}' permanently failed, crashing orchestrator") - raise worker._fatal_error + async def stop(self) -> None: + await self.cancel_inflight_rollouts() + if self.update_policy_task is not None: + await safe_cancel(self.update_policy_task) + self.update_policy_task = None + if self.inflight_policy_update_task is not None: + await safe_cancel(self.inflight_policy_update_task) + self.inflight_policy_update_task = None @property def max_off_policy_level(self) -> int: - if not self.inflight_group_rollouts: - return 0 - return max(info.off_policy_steps for info in self.inflight_group_rollouts.values()) - - @property - def min_off_policy_level(self) -> int: - if not self.inflight_group_rollouts: + steps = [info.off_policy_steps for info in self.inflight_requests.values()] + if not steps: return 0 - return min(info.off_policy_steps for info in self.inflight_group_rollouts.values()) + return max(steps) @property def mean_off_policy_level(self) -> float: - if not self.inflight_group_rollouts: + steps = [info.off_policy_steps for info in self.inflight_requests.values()] + if not steps: return 0 - steps = [info.off_policy_steps for info in self.inflight_group_rollouts.values()] return sum(steps) / len(steps) @property @@ -341,26 +520,35 @@ def async_level(self) -> int: return self.step - self.ckpt_step def get_metrics(self) -> dict[str, float]: + total_rollouts = sum(self.total_rollouts_by_env.values()) metrics = { "time/wait_for_ckpt": self.wait_for_ckpt_time, "time/update_weights": self.update_weights_time, - "batch/async_level": self.async_level, - "batch/off_policy_level/max": self.max_off_policy_level, - "batch/off_policy_level/mean": self.mean_off_policy_level, - "batch/off_policy_level/min": self.min_off_policy_level, - "batch/cancelled_rollouts": self.cancelled_rollouts_count, + "scheduler/async_level": self.async_level, + "scheduler/inflight_rollouts": self.inflight_rollout_count, + "scheduler/inflight_samples": self.inflight_sample_count, + "scheduler/cancelled_rollouts": self.cancelled_rollouts_count, + "empty_rollouts/all": sum(self.empty_rollouts_by_env.values()) / max(total_rollouts, 1), + "errored_rollouts/all": sum(self.errored_rollouts_by_env.values()) / max(total_rollouts, 1), + "off_policy_level/all/max": self.max_off_policy_level, + "off_policy_level/all/mean": self.mean_off_policy_level, } + for env_name in self.total_rollouts_by_env: + env_total = max(self.total_rollouts_by_env[env_name], 1) + metrics[f"empty_rollouts/{env_name}"] = self.empty_rollouts_by_env.get(env_name, 0) / env_total + metrics[f"errored_rollouts/{env_name}"] = self.errored_rollouts_by_env.get(env_name, 0) / env_total + by_env: dict[str, list[int]] = {} + for info in self.inflight_requests.values(): + by_env.setdefault(info.env_name, []).append(info.off_policy_steps) + for env_name, steps in by_env.items(): + metrics[f"off_policy_level/{env_name}/max"] = max(steps) + metrics[f"off_policy_level/{env_name}/mean"] = sum(steps) / len(steps) self.cancelled_rollouts_count = 0 + self.empty_rollouts_by_env.clear() + self.errored_rollouts_by_env.clear() + self.total_rollouts_by_env.clear() - # Add per-worker lag metrics and pending counts - for workers in self.workers.values(): - for worker in workers: - worker_key = worker.worker_name - # Track pending count per worker (useful for debugging load balancing) - metrics[f"worker/{worker_key}/pending"] = worker.pending_count - if worker.latest_lag_metrics: - for metric_name, value in worker.latest_lag_metrics.items(): - # e.g. "worker_lag/env_0/max" - metrics[f"worker_lag/{worker_key}/{metric_name.split('/')[-1]}"] = value + # Add inference pool metrics (e.g. elastic pool server counts) + metrics.update(self.inference_pool.get_metrics()) return metrics diff --git a/src/prime_rl/orchestrator/trajectories.py b/src/prime_rl/orchestrator/trajectories.py index 3a107a623f1..3a45ee9ada5 100644 --- a/src/prime_rl/orchestrator/trajectories.py +++ b/src/prime_rl/orchestrator/trajectories.py @@ -1,106 +1,846 @@ -from copy import deepcopy +import base64 +import hashlib +import time +from concurrent.futures import ThreadPoolExecutor +from io import BytesIO +from pathlib import Path +from typing import Any +import torch import verifiers as vf +from PIL import Image +from transformers.tokenization_utils import PreTrainedTokenizer from prime_rl.transport import TrainingSample +from prime_rl.utils.chat_template import ( + common_prefix_len, + deserialize_tool_calls, + normalize_messages, + render_messages, + strip_message_content, +) from prime_rl.utils.logger import get_logger +# We use list() instead of deepcopy() for flat lists (token IDs, logprobs) - safe because +# primitives are immutable. pixel_values/image_grid_thw are not mutated after creation. -def interleave_rollout(state: vf.State) -> list[TrainingSample] | None: + +def _align_routed_experts( + routed_experts: list[list[list[int]]] | None, + expected_len: int, +) -> list[list[list[int]]] | None: + """Align routed_experts length with the expected token count. + + VLLM's capturer uses `num_tokens - 1` slot mappings because the final + generated token was never fed as input to a forward pass and has no + routing decision. Append zero-filled entries for the missing positions. """ - Convert vf.State to a *single* trainable rollout by interleaving the trajectory. + if routed_experts is None or not routed_experts: + return routed_experts + deficit = expected_len - len(routed_experts) + if deficit <= 0: + return routed_experts + num_layers = len(routed_experts[0]) + topk = len(routed_experts[0][0]) + zero_entry = [[0] * topk for _ in range(num_layers)] + return routed_experts + [zero_entry for _ in range(deficit)] + + +def _common_prefix_len(a: list[int], b: list[int]) -> int: + return common_prefix_len(a, b) + + +def _normalize_messages(messages: Any, default_role: str) -> list[dict[str, Any]]: + return normalize_messages(messages, default_role) + + +def _deserialize_tool_calls(messages: list[dict[str, Any]]) -> list[dict[str, Any]]: + return deserialize_tool_calls(messages) + + +def _strip_message_content(messages: list[dict[str, Any]]) -> list[dict[str, Any]]: + return strip_message_content(messages) + - NOTE: - - This requires that consecutive trajectory steps share token prefixes (incremental tokenization) - - This approach is suceptible to introduce subtle difference due to re-tokenization in multi-turn environments. +def _render_messages( + tokenizer: PreTrainedTokenizer, + messages: list[dict[str, Any]], + add_generation_prompt: bool = False, + tools: list[dict[str, Any]] | None = None, + processor=None, +) -> list[int]: + return render_messages( + tokenizer, + messages, + add_generation_prompt=add_generation_prompt, + tools=tools, + processor=processor, + ) + + +def _prepare_messages_for_processor(messages: list[dict[str, Any]]) -> list[dict[str, Any]]: + """Convert messages to the format expected by the VLM processor. + + - Converts image_url items to image items with loaded PIL Images + - Strips extra fields (e.g. image_url on text items) that confuse the processor + - Ensures all message content is in list format (processor requires this) """ - logger = get_logger() + prepared = [] + for msg in messages: + content = msg.get("content") + if isinstance(content, str): + prepared.append({**msg, "content": [{"type": "text", "text": content}]}) + continue - trajectory = state["trajectory"] - if len(trajectory) == 0: - logger.warning(f"No trajectory steps for example {state['example_id']}. Skipping rollout.") - return None + if not isinstance(content, list): + prepared.append(msg) + continue - has_error = state["error"] is not None + new_content = [] + for item in content: + if item.get("type") == "image_url": + url = item.get("image_url", {}).get("url", "") + if url.startswith(_FILE_URL_PREFIX): + img = _load_file_image(url) + elif url.startswith("data:image"): + b64_data = url.split(",", 1)[1] + img = Image.open(BytesIO(base64.b64decode(b64_data))) + else: + new_content.append(item) + continue + new_content.append({"type": "image", "image": img}) + elif item.get("type") == "text": + new_content.append({"type": "text", "text": item.get("text", "")}) + else: + new_content.append(item) + prepared.append({**msg, "content": new_content}) - # Initialize the rollout with prompt and completion from first trajectory step - first_step = trajectory[0] - if has_error: - completion_mask = [False] * len(first_step["tokens"]["completion_mask"]) - else: - completion_mask = [bool(i) for i in first_step["tokens"]["completion_mask"]] - interleaved_rollout = TrainingSample( - prompt_ids=deepcopy(first_step["tokens"]["prompt_ids"]), - prompt_mask=[bool(i) for i in first_step["tokens"]["prompt_mask"]], - completion_ids=deepcopy(first_step["tokens"]["completion_ids"]), - completion_mask=completion_mask, - completion_logprobs=deepcopy(first_step["tokens"]["completion_logprobs"]), - teacher_logprobs=None, # Populated at the end after full sequence length is known if teacher model is configured - advantage=None, + return prepared + + +def _tokenize_step_from_messages( + step: vf.TrajectoryStep, + tokenizer: PreTrainedTokenizer, + tools: list[dict[str, Any]] | None = None, + processor=None, +) -> dict[str, Any]: + prompt = _normalize_messages(step.get("prompt"), default_role="user") + completion = _normalize_messages(step.get("completion"), default_role="assistant") + + prompt = _strip_message_content(_deserialize_tool_calls(prompt)) + completion = _strip_message_content(_deserialize_tool_calls(completion)) + + assert all(m.get("role") == "assistant" for m in completion), ( + "Expected all completion messages to be assistant role for SFT distillation, " + f"got roles: {[m.get('role') for m in completion]}" ) - # Interleave all other trajectory steps into completion - prefix_tokens = first_step["tokens"]["prompt_ids"] + first_step["tokens"]["completion_ids"] - for step_idx, step in enumerate(trajectory[1:], start=2): - tokens = step["tokens"] - assert tokens is not None - prev_trajectory_and_new_prompt_ids = tokens["prompt_ids"] + if processor is not None: + prompt = _prepare_messages_for_processor(prompt) + completion = _prepare_messages_for_processor(completion) - # Incremental tokenization assumption - if not prefix_tokens == prev_trajectory_and_new_prompt_ids[: len(prefix_tokens)]: - logger.warning( - f"Found mismatch in prefix tokens for example {state['example_id']} at trajectory step {step_idx}" - ) + all_messages = prompt + completion + prompt_has_assistant_completion = len(completion) > 0 and completion[0].get("role") == "assistant" + prompt_ids = _render_messages( + tokenizer, + prompt, + add_generation_prompt=prompt_has_assistant_completion, + tools=tools, + processor=processor, + ) + full_ids = _render_messages( + tokenizer, + all_messages, + tools=tools, + processor=processor, + ) - # Extend the completion with the new prompt - prompt_ids = deepcopy(prev_trajectory_and_new_prompt_ids[len(prefix_tokens) :]) - interleaved_rollout.completion_ids.extend(prompt_ids) - interleaved_rollout.completion_mask.extend([False] * len(prompt_ids)) - interleaved_rollout.completion_logprobs.extend([0.0] * len(prompt_ids)) + split_idx = _common_prefix_len(prompt_ids, full_ids) + original_prompt_len = len(prompt_ids) - # Extend the completion with the new completion tokens - completion_ids = deepcopy(tokens["completion_ids"]) - completion_logprobs = deepcopy(tokens["completion_logprobs"]) - interleaved_rollout.completion_ids.extend(completion_ids) - if has_error: - interleaved_rollout.completion_mask.extend([False] * len(tokens["completion_mask"])) + prompt_ids = full_ids[:split_idx] + completion_ids = full_ids[split_idx:] + completion_mask = [True] * len(completion_ids) + completion_logprobs = [0.0] * len(completion_ids) + + return { + "prompt_ids": prompt_ids, + "prompt_mask": [False] * len(prompt_ids), + "completion_ids": completion_ids, + "completion_mask": completion_mask, + "completion_logprobs": completion_logprobs, + "routed_experts": None, + "prompt_prefix_len": split_idx, + "original_prompt_len": original_prompt_len, + } + + +def _convert_tools_to_oai_format(tool_defs: list) -> list[dict[str, Any]] | None: + """Convert verifiers Tool objects or dicts to OAI function-calling format.""" + if not tool_defs: + return None + + def _get(tool: Any, key: str) -> Any: + if isinstance(tool, dict): + return tool.get(key) + return getattr(tool, key, None) + + return [ + { + "type": "function", + "function": { + "name": _get(tool, "name"), + "description": _get(tool, "description"), + "parameters": _get(tool, "parameters"), + **({} if _get(tool, "strict") is None else {"strict": _get(tool, "strict")}), + }, + } + for tool in tool_defs + ] + + +def _tokenize_step_with_renderer( + step: vf.TrajectoryStep, + renderer, + tools: list[dict[str, Any]] | None = None, +) -> dict[str, Any]: + """Tokenize a trajectory step using a Renderer.""" + from renderers.base import build_trajectory_step + + prompt = _normalize_messages(step.get("prompt"), default_role="user") + completion = _normalize_messages(step.get("completion"), default_role="assistant") + prompt = _strip_message_content(_deserialize_tool_calls(prompt)) + completion = _strip_message_content(_deserialize_tool_calls(completion)) + return build_trajectory_step(renderer, prompt, completion, tools=tools) + + +def pretokenize_rollout_trajectory( + output: vf.RolloutOutput, + tokenizer: PreTrainedTokenizer, + processor=None, + renderer=None, +) -> bool: + """Populate missing step tokens from prompt/completion messages. + + When a renderer is provided, uses it for tokenization (faster, deterministic). + Otherwise falls back to the tokenizer + apply_chat_template path. + """ + logger = get_logger() + tools = _convert_tools_to_oai_format(output.get("tool_defs", [])) + + for step_idx, step in enumerate(output["trajectory"]): + if step["tokens"] is not None: + continue + + if renderer is not None: + step["tokens"] = _tokenize_step_with_renderer(step, renderer, tools=tools) else: - interleaved_rollout.completion_mask.extend([bool(i) for i in tokens["completion_mask"]]) - interleaved_rollout.completion_logprobs.extend(completion_logprobs) + reconstructed = _tokenize_step_from_messages(step, tokenizer, tools=tools, processor=processor) + if reconstructed["prompt_prefix_len"] < reconstructed["original_prompt_len"]: + logger.debug( + f"Prompt tokenization was non-prefix for example {output['example_id']} step {step_idx}. " + f"Using longest common prefix length {reconstructed['prompt_prefix_len']} " + f"(original prompt had {reconstructed['original_prompt_len']} tokens)." + ) + reconstructed.pop("prompt_prefix_len") + reconstructed.pop("original_prompt_len") + step["tokens"] = reconstructed + + return True + - # New prefix is the current prompt and completion ids concatenated - prefix_tokens = tokens["prompt_ids"] + tokens["completion_ids"] +def interleave_rollout( + output: vf.RolloutOutput, + vlm_cache: "VLMImageCache | None" = None, + cache_key: int | None = None, + mm_token_type_ids_mapping: dict[int, int] | None = None, +) -> list[TrainingSample] | None: + """ + Convert vf.RolloutOutput to trainable rollouts by interleaving trajectory steps + where the extension property holds. + + When consecutive steps share token prefixes (extension property), they are + merged into a single sample. When extension breaks (e.g., due to context + compaction or a change in control-flow), a new sample is started. + Supports multi-prefix matching to handle interleaved agents. For example, + [agent1-step1, agent1-step2, agent2-step1, agent1-step3] produces two samples: + agent1 steps merged together, agent2 step separate. - return [interleaved_rollout] + Returns a list of samples - could be 1 (extension always held) or up to T + (extension never held). + For VLM models, pass vlm_cache to attach cumulative pixel_values per sample. + Each sample gets the images accumulated up to its last merged step. -def branch_rollout(state: vf.State) -> list[TrainingSample] | None: - """Convert vf.State to *multiple* trainable rollouts using branching trajectories strategy.""" + Args: + output: vf.RolloutOutput containing trajectory data + vlm_cache: Pre-computed VLM image cache for multimodal training + cache_key: Cache key to use when retrieving images from the VLM cache + """ logger = get_logger() - rollouts = [] - trajectory = state["trajectory"] + trajectory = output["trajectory"] if len(trajectory) == 0: - logger.warning(f"No trajectory steps for example {state['example_id']}. Skipping rollout.") + error = output.get("error") + stop = output.get("stop_condition") + logger.warning( + f"No trajectory steps for example {output['example_id']} (error={error}, stop={stop}). Skipping rollout." + ) return None - has_error = state["error"] is not None - for step in state["trajectory"]: - assert "tokens" in step + has_error = output["error"] is not None + # this field should be guaranteed because we set temperature in get_sampling_args + temperature = output["sampling_args"]["temperature"] + + def prepare_step_tokens(step: vf.TrajectoryStep, step_idx: int) -> dict[str, Any] | None: tokens = step["tokens"] + if tokens is not None: + return { + "prompt_ids": list(tokens["prompt_ids"]), + "prompt_mask": [bool(i) for i in tokens["prompt_mask"]], + "completion_ids": list(tokens["completion_ids"]), + "completion_mask": [bool(i) for i in tokens["completion_mask"]], + "completion_logprobs": list(tokens["completion_logprobs"]), + "routed_experts": tokens.get("routed_experts"), + } + + logger.warning(f"Missing rollout tokens for example {output['example_id']} step {step_idx}.") + return None + + prepared_steps: list[dict[str, Any]] = [] + for step_idx, step in enumerate(trajectory): + prepared = prepare_step_tokens(step, step_idx) + if prepared is None: + return None + prepared_steps.append(prepared) + + def make_sample(tokens: dict[str, Any]) -> TrainingSample: + """Create a new TrainingSample from a trajectory step.""" if has_error: completion_mask = [False] * len(tokens["completion_mask"]) else: completion_mask = [bool(i) for i in tokens["completion_mask"]] - rollout = TrainingSample( - prompt_ids=deepcopy(tokens["prompt_ids"]), + completion_ids = list(tokens["completion_ids"]) + + routed_experts = _align_routed_experts( + tokens.get("routed_experts"), + len(tokens["prompt_ids"]) + len(tokens["completion_ids"]), + ) + prompt_ids = list(tokens["prompt_ids"]) + return TrainingSample( + prompt_ids=prompt_ids, prompt_mask=[bool(i) for i in tokens["prompt_mask"]], - completion_ids=deepcopy(tokens["completion_ids"]), + completion_ids=completion_ids, completion_mask=completion_mask, - completion_logprobs=deepcopy(tokens["completion_logprobs"]), - advantage=None, + completion_logprobs=list(tokens["completion_logprobs"]), + completion_temperatures=[temperature] * len(completion_ids), teacher_logprobs=None, + advantage=None, + routed_experts=routed_experts, + mm_token_type_ids=None, ) - rollouts.append(rollout) - return rollouts + + def extend_sample(sample: TrainingSample, prefix_len: int, step_idx: int) -> None: + """Extend an existing sample with a new trajectory step (extension property holds).""" + tokens = prepared_steps[step_idx] + + # Extend with new prompt tokens (mask=False, no gradient) + new_prompt_ids = tokens["prompt_ids"][prefix_len:] + sample.completion_ids.extend(new_prompt_ids) + sample.completion_mask.extend([False] * len(new_prompt_ids)) + sample.completion_logprobs.extend([0.0] * len(new_prompt_ids)) + sample.completion_temperatures.extend([temperature] * len(new_prompt_ids)) + + # Extend with new completion tokens + completion_ids = tokens["completion_ids"] + sample.completion_ids.extend(completion_ids) + if has_error: + sample.completion_mask.extend([False] * len(tokens["completion_mask"])) + else: + sample.completion_mask.extend(bool(i) for i in tokens["completion_mask"]) + sample.completion_logprobs.extend(tokens["completion_logprobs"]) + sample.completion_temperatures.extend([temperature] * len(completion_ids)) + + if tokens.get("routed_experts") is not None and sample.routed_experts is not None: + step_routed = tokens["routed_experts"] + # The previous step's last routing entry was zero-padded by _align_routed_experts + # (vLLM only captures num_tokens-1 routings per request). This step actually + # processed that boundary token as part of its prompt, so replace the zero-fill + # with the real routing decision before appending new entries. + if prefix_len > 0 and prefix_len <= len(step_routed): + sample.routed_experts[prefix_len - 1] = step_routed[prefix_len - 1] + sample.routed_experts.extend(step_routed[prefix_len:]) + expected_len = len(sample.prompt_ids) + len(sample.completion_ids) + sample.routed_experts = _align_routed_experts(sample.routed_experts, expected_len) + + # Track [prefix_tokens, sample, last_step_idx] per active sample + active_samples: list[tuple[list[int], TrainingSample, int]] = [] + + first_tokens = prepared_steps[0] + first_prefix = first_tokens["prompt_ids"] + first_tokens["completion_ids"] + active_samples.append((first_prefix, make_sample(first_tokens), 0)) + + for step_idx, _step in enumerate(trajectory[1:], start=1): + tokens = prepared_steps[step_idx] + step_prompt_ids = tokens["prompt_ids"] + + # Check if this step extends ANY active prefix + matched_idx = None + for idx, (prefix_tokens, _, _) in enumerate(active_samples): + if step_prompt_ids[: len(prefix_tokens)] == prefix_tokens: + matched_idx = idx + break + + if matched_idx is not None: + # Extension holds - merge into matched sample + prefix_tokens, sample, _ = active_samples[matched_idx] + extend_sample(sample, len(prefix_tokens), step_idx=step_idx) + active_samples[matched_idx] = (tokens["prompt_ids"] + tokens["completion_ids"], sample, step_idx) + else: + # No prefix matches - start a new sample + logger.debug( + f"Extension property broke at step {step_idx + 1} for example {output['example_id']}. " + f"Starting new sample (active_prefixes={len(active_samples)}, step_prompt_len={len(step_prompt_ids)})." + ) + new_prefix = tokens["prompt_ids"] + tokens["completion_ids"] + active_samples.append((new_prefix, make_sample(tokens), step_idx)) + + # Attach images once per sample using only the last merged step. Prompt + # tokens already contain fully expanded <|image_pad|> placeholders because + # VLMs go through MITO (chat completions), which runs apply_chat_template + # server-side; the orchestrator re-tokenizes via the same processor in the + # fallback path so features and tokens stay 1:1. + if vlm_cache is not None: + key = output["example_id"] if cache_key is None else cache_key + for _, sample, last_step_idx in active_samples: + pv, shape, grids = vlm_cache.get_for_step(key, last_step_idx) + sample.pixel_values = pv + sample.pixel_values_shape = shape + sample.image_grid_thw = grids + if mm_token_type_ids_mapping is not None: + sample.mm_token_type_ids = [ + mm_token_type_ids_mapping.get(token_id, 0) for token_id in sample.prompt_ids + sample.completion_ids + ] + + return [sample for _, sample, _ in active_samples] + + +# ============================================================================= +# VLM-specific functions +# ============================================================================= + + +_FILE_URL_PREFIX = "file://" + + +def offload_images_to_disk(rollouts: list[vf.RolloutOutput], output_dir: Path) -> int: + """Replace base64 image data in rollout trajectories with file paths on disk. + + Scans all trajectory step prompts for data:image URLs, writes the decoded + image bytes to ``{output_dir}/assets/images/{hash}.png``, and replaces the + URL in-place with ``file://{path}``. Deduplicates by content hash so each + unique image is written only once. + + Returns the number of unique images written to disk. + """ + images_dir = output_dir / "assets" / "images" + images_dir.mkdir(parents=True, exist_ok=True) + + written: set[str] = set() + + for output in rollouts: + for step in output.get("trajectory", []): + prompt = step.get("prompt") + if not prompt or not isinstance(prompt, list): + continue + for msg in prompt: + content = msg.get("content", []) + if not isinstance(content, list): + continue + for item in content: + if item.get("type") != "image_url": + continue + url = item.get("image_url", {}).get("url", "") + if not url.startswith("data:image"): + continue + b64_data = url.split(",", 1)[1] + content_hash = hashlib.sha256(b64_data.encode()).hexdigest()[:16] + path = images_dir / f"{content_hash}.png" + if content_hash not in written: + if not path.exists(): + path.write_bytes(base64.b64decode(b64_data)) + written.add(content_hash) + item["image_url"]["url"] = f"{_FILE_URL_PREFIX}{path}" + + return len(written) + + +def _load_file_image(path_str: str) -> Image.Image: + """Load an image from a file:// path.""" + return Image.open(path_str.removeprefix(_FILE_URL_PREFIX)) + + +def _extract_images_from_messages(messages: list) -> list[tuple[Image.Image, str]]: + """Extract (image, key) pairs from OpenAI-style chat messages. + + Handles both base64 data URLs and file:// paths from disk offloading. + """ + images = [] + if not messages or not isinstance(messages, list): + return images + + for msg in messages: + content = msg.get("content", []) + if isinstance(content, list): + for item in content: + if item.get("type") == "image_url": + url = item.get("image_url", {}).get("url", "") + if url.startswith(_FILE_URL_PREFIX): + img = _load_file_image(url) + images.append((img, url)) + elif url.startswith("data:image"): + b64_data = url.split(",", 1)[1] + img_bytes = base64.b64decode(b64_data) + img = Image.open(BytesIO(img_bytes)) + images.append((img, b64_data)) + return images + + +def _collect_image_keys_from_messages(messages: list) -> list[str]: + """Extract image keys from OpenAI-style chat messages without decoding. + + Handles both base64 data URLs and file:// paths from disk offloading. + """ + keys = [] + if not messages or not isinstance(messages, list): + return keys + for msg in messages: + content = msg.get("content", []) + if isinstance(content, list): + for item in content: + if item.get("type") == "image_url": + url = item.get("image_url", {}).get("url", "") + if url.startswith("data:image"): + keys.append(url.split(",", 1)[1]) + elif url.startswith(_FILE_URL_PREFIX): + keys.append(url) + return keys + + +def _decode_image(key: str) -> Image.Image: + """Decode an image from a base64 string or load from a file:// path.""" + if key.startswith(_FILE_URL_PREFIX): + return _load_file_image(key) + return Image.open(BytesIO(base64.b64decode(key))) + + +_PARALLEL_DECODE_THRESHOLD = 4 + + +_IMAGE_STRIPPED_PLACEHOLDER = "[preprocessed image]" + + +def strip_base64_images(examples: list[tuple[int, vf.RolloutOutput]]) -> None: + """Strip image data from rollout prompts to free memory. + + Handles both base64 data URLs and file:// paths from disk offloading. + The images have been decoded and indexed; the original data is no longer needed. + """ + for _, output in examples: + for step in output.get("trajectory", []): + prompt = step.get("prompt") + if not prompt or not isinstance(prompt, list): + continue + for msg in prompt: + content = msg.get("content", []) + if isinstance(content, list): + for item in content: + if item.get("type") == "image_url": + url = item.get("image_url", {}).get("url", "") + if url.startswith("data:image") or url.startswith(_FILE_URL_PREFIX): + item["image_url"]["url"] = _IMAGE_STRIPPED_PLACEHOLDER + + +def _extract_images_from_examples( + examples: list[tuple[int, vf.RolloutOutput]], +) -> tuple[list[Image.Image], dict[int, list[list[int]]]]: + """ + Extract images from all trajectory steps of each example. + + Two-pass approach: first collects unique base64 keys (fast, string-only), + then decodes unique images in parallel via ThreadPoolExecutor. + + Args: + examples: List of (cache_key, output) tuples where output contains a "trajectory" + list with steps that have "prompt" messages in OpenAI chat format. + + Returns: + Tuple of (all_images, step_image_indices_per_example) + - all_images: deduplicated flat list of decoded PIL images + - step_image_indices_per_example: dict mapping cache_key to per-step lists of + indices into all_images (e.g., [[0], [0, 1], [1]] for the decreasing-images case) + """ + # Pass 1: collect unique b64 keys and build step indices + unique_keys: list[str] = [] + key_to_index: dict[str, int] = {} + step_image_indices_per_example: dict[int, list[list[int]]] = {} + + for eid, output in examples: + trajectory = output.get("trajectory", []) + if not trajectory: + step_image_indices_per_example[eid] = [] + continue + + step_image_indices = [] + for step in trajectory: + prompt = step.get("prompt") + image_keys = _collect_image_keys_from_messages(prompt) + indices = [] + for key in image_keys: + if key not in key_to_index: + key_to_index[key] = len(unique_keys) + unique_keys.append(key) + indices.append(key_to_index[key]) + step_image_indices.append(indices) + + step_image_indices_per_example[eid] = step_image_indices + + # Pass 2: decode unique images (parallel when worthwhile) + if len(unique_keys) > _PARALLEL_DECODE_THRESHOLD: + with ThreadPoolExecutor(max_workers=min(len(unique_keys), 16)) as pool: + all_images = list(pool.map(_decode_image, unique_keys)) + else: + all_images = [_decode_image(k) for k in unique_keys] + del unique_keys, key_to_index + + strip_base64_images(examples) + + return all_images, step_image_indices_per_example + + +_DEFAULT_IMAGE_CHUNK_SIZE = 32 + + +class _ImageStore: + """Holds per-unique-image data, assembled lazily on demand. + + Instead of duplicating pixel bytes for every step that references an image, + we store each image's bytes once and assemble the concatenation at retrieval time. + """ + + def __init__( + self, + image_bytes: list[bytes], + image_num_patches: list[int], + patch_dim: int, + image_grids: list[list[int]], + ): + self.image_bytes = image_bytes + self.image_num_patches = image_num_patches + self.patch_dim = patch_dim + self.image_grids = image_grids + self._cache: dict[tuple[int, ...], tuple[bytes, list[int], list[list[int]]]] = {} + + def assemble(self, indices: list[int]) -> tuple[bytes, list[int], list[list[int]]]: + """Assemble pixel bytes, shape, and grids for a set of image indices. + + Results are cached by index tuple — multi-turn rollouts with the same + cumulative image set (common across rollouts of the same example) hit + the cache and skip the join. + """ + cache_key = tuple(indices) + cached = self._cache.get(cache_key) + if cached is not None: + return cached + + total_patches = sum(self.image_num_patches[i] for i in indices) + pixel_bytes = b"".join(self.image_bytes[i] for i in indices) + shape = [total_patches, self.patch_dim] + grids = [self.image_grids[i] for i in indices] + result = (pixel_bytes, shape, grids) + self._cache[cache_key] = result + return result + + +def _preprocess_images_batched( + images: list[Image.Image], + step_image_indices_per_example: dict[int, list[list[int]]], + processor, + chunk_size: int = _DEFAULT_IMAGE_CHUNK_SIZE, +) -> tuple["_ImageStore | None", dict[int, list[list[int]]]]: + """ + Preprocess all images in chunked batches, returning an _ImageStore and step indices. + + Images are processed in chunks to avoid OOM on large batches. Per-image bytes are + stored once in the _ImageStore and assembled lazily at retrieval time. + + Returns: + Tuple of (_ImageStore or None, step_image_indices_per_example). + The store is None when there are no images or no processor. + """ + if not images or processor is None: + return None, step_image_indices_per_example + + logger = get_logger() + image_sizes = [(img.width, img.height) for img in images] + + # Process images in chunks to avoid OOM, parallelized across threads + # (PIL/numpy release the GIL so threads give real concurrency here) + chunks = [images[i : i + chunk_size] for i in range(0, len(images), chunk_size)] + + def _process_chunk(chunk: list[Image.Image]) -> tuple[torch.Tensor, torch.Tensor]: + processed = processor.image_processor(images=chunk, return_tensors="pt") + return processed["pixel_values"], processed["image_grid_thw"] + + if len(chunks) > 1: + with ThreadPoolExecutor(max_workers=min(len(chunks), 8)) as pool: + results = list(pool.map(_process_chunk, chunks)) + else: + results = [_process_chunk(chunks[0])] + + # Free PIL images now that preprocessing is done + del chunks + images.clear() + + all_pixel_values_list = [r[0] for r in results] + all_grid_thw_list = [r[1] for r in results] + + all_pixel_values = torch.cat(all_pixel_values_list, dim=0) + all_grid_thw = torch.cat(all_grid_thw_list, dim=0) + del all_pixel_values_list, all_grid_thw_list, results + + logger.debug( + f"VLM image processing: {len(image_sizes)} images, sizes={image_sizes}, " + f"pixel_values={all_pixel_values.shape}, grid_thw={all_grid_thw.tolist()}" + ) + + # Pre-compute patch start offset for each image + patch_starts = [0] + for g in all_grid_thw: + patch_starts.append(patch_starts[-1] + int(g[0] * g[1] * g[2])) + + patch_dim = all_pixel_values.shape[1] + + # Convert to bytes per-image and free the tensor immediately after + image_bytes_list: list[bytes] = [] + image_num_patches_list: list[int] = [] + image_grids_list: list[list[int]] = [] + for i in range(len(image_sizes)): + img_slice = all_pixel_values[patch_starts[i] : patch_starts[i + 1]] + image_bytes_list.append(img_slice.numpy().tobytes()) + image_num_patches_list.append(img_slice.shape[0]) + image_grids_list.append(all_grid_thw[i].tolist()) + del all_pixel_values, all_grid_thw + + store = _ImageStore( + image_bytes=image_bytes_list, + image_num_patches=image_num_patches_list, + patch_dim=patch_dim, + image_grids=image_grids_list, + ) + + return store, step_image_indices_per_example + + +class VLMImageCache: + """Result of building VLM image cache with per-step image data.""" + + def __init__( + self, + cache: dict[int, list[tuple[bytes | None, list[int] | None, list[list[int]] | None]]], + num_unique_examples: int, + extract_time: float, + preprocess_time: float, + ): + self._store: _ImageStore | None = None + self._step_indices: dict[int, list[list[int]]] | None = None + self.cache = cache + self.num_unique_examples = num_unique_examples + self.num_unique_images = 0 + self.extract_time = extract_time + self.preprocess_time = preprocess_time + + @classmethod + def from_store( + cls, + store: _ImageStore | None, + step_indices: dict[int, list[list[int]]], + num_unique_examples: int, + num_unique_images: int, + extract_time: float, + preprocess_time: float, + ) -> "VLMImageCache": + """Create a store-backed cache that assembles bytes lazily.""" + obj = cls.__new__(cls) + obj._store = store + obj._step_indices = step_indices + obj.cache = {} + obj.num_unique_examples = num_unique_examples + obj.num_unique_images = num_unique_images + obj.extract_time = extract_time + obj.preprocess_time = preprocess_time + return obj + + def _assemble(self, indices: list[int]) -> tuple[bytes | None, list[int] | None, list[list[int]] | None]: + if not indices: + return (None, None, None) + return self._store.assemble(indices) + + def get_for_step( + self, cache_key: int, step_idx: int + ) -> tuple[bytes | None, list[int] | None, list[list[int]] | None]: + """Get cumulative images up to and including the given step.""" + if self._store is not None: + steps = self._step_indices.get(cache_key, []) + if not steps or step_idx >= len(steps): + return (None, None, None) + return self._assemble(steps[step_idx]) + + steps = self.cache.get(cache_key, []) + if not steps or step_idx >= len(steps): + return (None, None, None) + return steps[step_idx] + + def get_all(self, cache_key: int) -> tuple[bytes | None, list[int] | None, list[list[int]] | None]: + """Get all images for the cache key (last step's cumulative images).""" + if self._store is not None: + steps = self._step_indices.get(cache_key, []) + if not steps: + return (None, None, None) + return self._assemble(steps[-1]) + + steps = self.cache.get(cache_key, []) + if not steps: + return (None, None, None) + return steps[-1] + + +def build_vlm_image_cache(rollouts: list[vf.RolloutOutput], processor) -> VLMImageCache: + """ + Build image cache for VLM training by extracting and preprocessing images. + + Caches per rollout to keep images aligned with divergent multi-turn trajectories. + """ + examples = [(idx, rollout) for idx, rollout in enumerate(rollouts)] + unique_example_ids = {(rollout.get("env_name"), rollout["example_id"]) for rollout in rollouts} + + # Extract images (also strips base64 data from rollout prompts to free memory) + extract_start = time.perf_counter() + all_images, images_per_example = _extract_images_from_examples(examples) + num_unique_images = len(all_images) + extract_time = time.perf_counter() - extract_start + + # Preprocess images (clears PIL image list when done) + preprocess_start = time.perf_counter() + store, step_indices = _preprocess_images_batched(all_images, images_per_example, processor) + preprocess_time = time.perf_counter() - preprocess_start + + return VLMImageCache.from_store( + store=store, + step_indices=step_indices, + num_unique_examples=len(unique_example_ids), + num_unique_images=num_unique_images, + extract_time=extract_time, + preprocess_time=preprocess_time, + ) diff --git a/src/prime_rl/orchestrator/utils.py b/src/prime_rl/orchestrator/utils.py index b0f80438a91..d4d7ec923bd 100644 --- a/src/prime_rl/orchestrator/utils.py +++ b/src/prime_rl/orchestrator/utils.py @@ -1,96 +1,38 @@ import asyncio +import time +from concurrent.futures import ThreadPoolExecutor from itertools import cycle from pathlib import Path -from typing import Any, AsyncContextManager +from typing import Any import pandas as pd -from openai import AsyncOpenAI -from openai.types.chat.chat_completion import ChatCompletion, Choice -from openai.types.completion_usage import CompletionUsage +import verifiers as vf from rich.console import Console from rich.table import Table -from verifiers.utils.async_utils import maybe_semaphore +from verifiers.utils.client_utils import setup_openai_client -from prime_rl.orchestrator.config import SamplingConfig +from prime_rl.configs.orchestrator import OrchestratorConfig from prime_rl.transport import TrainingSample +from prime_rl.utils.logger import get_logger from prime_rl.utils.utils import ( - format_num, format_time, get_broadcast_dir, get_ckpt_dir, get_step_path, ) -SEMAPHORE: AsyncContextManager | None = None - -async def set_semaphore(limit: int): - global SEMAPHORE - SEMAPHORE = await maybe_semaphore(limit) - - -async def get_semaphore() -> AsyncContextManager: - global SEMAPHORE - assert SEMAPHORE is not None, "Semaphore not set" - return SEMAPHORE - - -def get_sampling_args(sampling_config: SamplingConfig) -> dict: - # Convert SamplingConfig to vLLM OAI sampling args - # https://docs.vllm.ai/en/latest/serving/openai_compatible_server.html#extra-parameters_2 - sampling_args = dict(sampling_config) - sampling_args["top_p"] = 1.0 - sampling_args["logprobs"] = True - sampling_args["extra_body"] = { - **sampling_config.extra_body, - "return_token_ids": True, # Always return token IDs - "top_k": -1, - "min_p": 0.0, - } - sampling_args["extra_body"]["min_tokens"] = sampling_args.pop("min_tokens") - sampling_args["extra_body"]["repetition_penalty"] = sampling_args.pop("repetition_penalty") - return sampling_args - - -def parse_num_completion_tokens(responses: list[list[ChatCompletion]]) -> list[int]: - """Parses the number of tokens from a list of chat completions returned by OAI API.""" - all_num_completion_tokens = [] - for response in responses: - num_completion_tokens = 0 - for chat_completion in response: - assert isinstance(chat_completion, ChatCompletion) - assert chat_completion.usage is not None, "Usage should be present in the response" - usage = chat_completion.usage - assert isinstance(usage, CompletionUsage) - num_completion_tokens += usage.completion_tokens - all_num_completion_tokens.append(num_completion_tokens) - assert len(all_num_completion_tokens) == len(responses), ( - "Number of completion tokens should be the same as the number of responses" - ) - return all_num_completion_tokens - - -def parse_is_truncated_completions(responses: list[list[ChatCompletion]]) -> list[bool]: - """Parses whether the completions were truncated from a list of (multi-turn) OAI chat completions""" - all_is_truncated = [] - for response in responses: - is_truncated = False - for chat_completion in response: - assert isinstance(chat_completion, ChatCompletion) - assert len(chat_completion.choices) == 1, "Response should always have one choice" - choice = chat_completion.choices[0] - assert isinstance(choice, Choice) - if choice.finish_reason == "length": - is_truncated = True - all_is_truncated.append(is_truncated) - return all_is_truncated +def set_default_executor(max_workers: int = 64) -> None: + """Scale the default asyncio thread pool so asyncio.to_thread has enough capacity.""" + get_logger().info(f"Setting default executor to ThreadPoolExecutor(max_workers={max_workers})") + asyncio.get_event_loop().set_default_executor(ThreadPoolExecutor(max_workers=max_workers)) def print_benchmark(history: dict[str, list[Any]]) -> None: """ - Print benchmark results as rich table. Shows formatted values for the - inference throughput and overall step time. First first N rows show the - per-step values, and the last row shows the mean, std, min, and max values. + Print benchmark results as rich table. Shows formatted step time values. + First N rows show the per-step values, and the last row shows the mean, + std, min, and max values. """ history.pop("step") assert all(len(v) for v in history.values()), "All metrics must have logged the same number of steps" @@ -98,7 +40,6 @@ def print_benchmark(history: dict[str, list[Any]]) -> None: # Turn metric history into pd.DataFrame df = pd.DataFrame(dict(history.items())) columns = { - "perf/throughput": "Throughput", "time/step": "Step Time", } df = df.rename(columns=columns) @@ -117,7 +58,6 @@ def print_benchmark(history: dict[str, list[Any]]) -> None: # Add formatted rows formatted_df = pd.DataFrame(columns=df.columns) formatted_df["Step Time"] = df["Step Time"].apply(format_time) - formatted_df["Throughput"] = df["Throughput"].apply(format_num, precision=2) for step, row in formatted_df.iterrows(): table.add_row(*([str(step)] + [str(x) for x in row])) @@ -129,7 +69,6 @@ def print_benchmark(history: dict[str, list[Any]]) -> None: mean_df = df.describe().loc[["mean", "std", "min", "max"], :] formatted_mean_df = pd.DataFrame(columns=mean_df.columns) formatted_mean_df["Step Time"] = mean_df["Step Time"].apply(format_time) - formatted_mean_df["Throughput"] = mean_df["Throughput"].apply(format_num, precision=2) mean_row = ["Overall"] + formatted_mean_df.T.apply( lambda row: f"{row['mean']} ± {row['std']} [{row['min']}, {row['max']}]", axis=1 ).tolist() @@ -140,41 +79,119 @@ def print_benchmark(history: dict[str, list[Any]]) -> None: async def compute_teacher_logprobs( - clients: list[AsyncOpenAI], + clients: list[vf.ClientConfig], model_name: str, samples: list[TrainingSample], ) -> list[list[float]]: """Compute teacher model logprobs for a batch of training samples via prefill.""" - - async def _compute_single(client: AsyncOpenAI, sample: TrainingSample) -> list[float]: - async with await get_semaphore(): - response = await client.post( - "/chat/completions/tokens", - body={ - "model": model_name, - "messages": [{"role": "user", "content": ""}], - "tokens": sample.prompt_ids + sample.completion_ids, + import httpx + from vllm.entrypoints.serve.disagg.protocol import GenerateResponse + + async def _compute_single(client_config: vf.ClientConfig, sample: TrainingSample) -> list[float]: + client = setup_openai_client(client_config) + + # Two escape hatches from ``AsyncOpenAI.post``: + # 1. URL — ``/inference/v1/generate`` is mounted at server root, not + # under ``/v1``. Pass an absolute URL so the SDK's + # ``_prepare_url`` skips the base-url merge (it short-circuits + # when the path passes ``httpx.URL.is_relative_url`` as False). + # 2. Parse — vLLM's ``GenerateResponse`` is a plain + # ``pydantic.BaseModel`` and the SDK's parse layer rejects any + # ``cast_to`` that doesn't subclass ``openai.BaseModel``. Use + # ``cast_to=httpx.Response`` so the SDK still builds the request + # (preserving ``auth_headers``, retries, timeouts, idempotency + # keys) and just hands us the raw response to validate ourselves. + base = str(client.base_url).rstrip("/").removesuffix("/v1") + http_response = await client.post( + f"{base}/inference/v1/generate", + cast_to=httpx.Response, + body={ + "model": model_name, + "token_ids": list(sample.prompt_ids) + list(sample.completion_ids), + "sampling_params": { "max_tokens": 1, "temperature": 1.0, "top_p": 1.0, - "skip_special_tokens": False, - "prompt_logprobs": True, + "prompt_logprobs": 1, }, - cast_to=ChatCompletion, - ) - return [0.0 if lp is None else float(next(iter(lp.values()))["logprob"]) for lp in response.prompt_logprobs] + }, + ) + response = GenerateResponse.model_validate_json(http_response.content) + # ``prompt_logprobs[i]`` is a ``{token_id: Logprob}`` dict for tokens + # the engine could score, or ``None`` for the leading token which has + # no preceding context. Flatten to ``list[float]`` with 0.0 in the + # unscored slot. + flat: list[float] = [] + for entry in response.prompt_logprobs or []: + if not entry: + flat.append(0.0) + continue + first = next(iter(entry.values())) + lp = first.logprob if hasattr(first, "logprob") else first.get("logprob") + flat.append(float(lp) if lp is not None else 0.0) + return flat return await asyncio.gather(*[_compute_single(client, sample) for client, sample in zip(cycle(clients), samples)]) -def get_weight_dir(output_dir: Path, step: int) -> Path: - """Get the weight directory for a given checkpoint step.""" +def get_weight_dir(output_dir: Path, step: int, check_exists: bool = True, wait_timeout: int | None = None) -> Path: + """Get the weight directory for a given checkpoint step. + + Args: + output_dir: The output directory for the run. + step: The checkpoint step. + check_exists: If True, raises FileNotFoundError if no weight directory exists. + If False, returns the broadcast directory path without checking existence + (useful for NCCL mode where weights are broadcasted, not stored on disk). + wait_timeout: Maximum time in seconds to wait for a stable directory to appear. + If None, no waiting is performed. + """ ckpt_weight_dir = get_step_path(get_ckpt_dir(output_dir), step) / "weight" broadcast_weight_dir = get_step_path(get_broadcast_dir(output_dir), step) - if ckpt_weight_dir.exists(): - return ckpt_weight_dir - if broadcast_weight_dir.exists(): + + def find_stable_dir() -> Path | None: + # For checkpoint weights, check STABLE file in parent directory (checkpoints/step_{step}/STABLE) + ckpt_step_dir = get_step_path(get_ckpt_dir(output_dir), step) + if (ckpt_step_dir / "STABLE").exists() and ckpt_weight_dir.exists(): + return ckpt_weight_dir + + # For broadcast weights, check STABLE file in the broadcast directory itself + if (broadcast_weight_dir / "STABLE").exists() and broadcast_weight_dir.exists(): + return broadcast_weight_dir + + return None + + # Check immediately, then wait if needed + result = find_stable_dir() + if result is None and wait_timeout: + start_time = time.time() + while time.time() - start_time < wait_timeout: + time.sleep(1) + result = find_stable_dir() + if result: + break + + if result: + return result + if not check_exists: return broadcast_weight_dir - raise FileNotFoundError( - f"No weight directory found for checkpoint step {step}. Expected to find it in {ckpt_weight_dir} or {broadcast_weight_dir}." - ) + + raise FileNotFoundError(f"No weight directory found for checkpoint step {step}") + + +def setup_external_rollout_model(config: OrchestratorConfig, logger) -> tuple[Any, str, bool]: + """Resolve rollout client/model and whether policy updates should be enabled.""" + rollout_client_config = config.client + rollout_model_name = config.model.name + enable_policy_updates = True + + if config.teacher_rollout_model is not None: + rollout_client_config = config.teacher_rollout_model.client + rollout_model_name = config.teacher_rollout_model.model.name + enable_policy_updates = False + logger.info( + f"Using external teacher rollout model (base_url={', '.join(rollout_client_config.base_url)}, " + f"model={rollout_model_name})" + ) + + return rollout_client_config, rollout_model_name, enable_policy_updates diff --git a/src/prime_rl/orchestrator/vf_utils.py b/src/prime_rl/orchestrator/vf_utils.py new file mode 100644 index 00000000000..17b62f53cc3 --- /dev/null +++ b/src/prime_rl/orchestrator/vf_utils.py @@ -0,0 +1,98 @@ +import json +import logging +from pathlib import Path + +import verifiers as vf +from verifiers.utils.save_utils import make_serializable + +from prime_rl.utils.logger import InterceptHandler + + +# TODO: remove once usage is tracked by verifiers +def get_prompt_len(output: vf.RolloutOutput) -> int: + """ + Computes the number of prompt tokens from vf.RolloutOutput. Defined as the + number of prompt ids from the first trajectory step. If raw tokens are not + available, falls back to checking the usage of the first response. + """ + if not output["trajectory"]: + return 0 + first_step = output["trajectory"][0] + if first_step["tokens"] is not None: + return len(first_step["tokens"]["prompt_ids"]) + first_step_response = first_step["response"] + return (first_step_response.get("usage") or {}).get("prompt_tokens", 0) + + +# TODO: remove once usage is tracked by verifiers +def get_seq_len(output: vf.RolloutOutput) -> int: + """ + Computes the number of tokens from vf.RolloutOutput. Defined as the sum of prompt + and completion tokens from the last trajectory step. If raw tokens are not + available, falls back to checking the usage of the last response. + """ + if not output["trajectory"]: + return 0 + last_step = output["trajectory"][-1] + if last_step["tokens"] is not None: + return len(last_step["tokens"]["prompt_ids"]) + len(last_step["tokens"]["completion_ids"]) + last_step_response = last_step["response"] + return (last_step_response.get("usage") or {}).get("total_tokens", 0) + + +# TODO: remove once usage is tracked by verifiers +def get_completion_len(output: vf.RolloutOutput) -> int: + """ + Computes the number of completion tokens from vf.RolloutOutput. Defined as + the difference between the total number of tokens and the number of prompt + tokens. + """ + return get_seq_len(output) - get_prompt_len(output) + + +def get_model_completion_len(output: vf.RolloutOutput) -> int: + """ + Computes the number of model-generated completion tokens across all turns. + Unlike get_completion_len, this excludes environment responses injected + between turns in multi-turn rollouts. + """ + return sum(len(step["tokens"]["completion_ids"]) for step in output["trajectory"] if step.get("tokens")) + + +def get_num_turns(output: vf.RolloutOutput) -> int: + """Number of turns (trajectory steps) in a rollout.""" + return len(output["trajectory"]) + + +def get_tool_response_len(output: vf.RolloutOutput) -> int: + """ + Total tool-response tokens consumed across the whole rollout. + + Read from a harness-emitted metric (e.g. RLM's `rlm_total_tool_response_tokens`, + deduped across turns/branches/sub-RLMs). Returns 0 if no harness metric is + present, which makes this a no-op for envs without tool-response accounting. + """ + metrics = output.get("metrics") or {} + for key, value in metrics.items(): + if key.endswith("total_tool_response_tokens") and isinstance(value, (int, float)): + return int(value) + return 0 + + +def save_rollouts(rollouts: list[vf.RolloutOutput], path: Path, exclude_keys: set[str] | None = None) -> None: + """Save rollouts to a JSONL file using verifiers serialization.""" + path.parent.mkdir(parents=True, exist_ok=True) + with open(path, "w") as f: + for rollout in rollouts: + row = {k: v for k, v in rollout.items() if k not in exclude_keys} if exclude_keys else rollout + json.dump(row, f, default=make_serializable) + f.write("\n") + + +def intercept_vf_logging(logger: str = "verifiers", level: str = "DEBUG", prefix: str | None = None): + """Intercepts verifiers logging and routes through prime-rl logger with optional prefix.""" + vf_logger = logging.getLogger(logger) + vf_logger.handlers.clear() + vf_logger.addHandler(InterceptHandler(prefix=prefix)) + vf_logger.setLevel(level.upper()) + vf_logger.propagate = False diff --git a/src/prime_rl/rl.py b/src/prime_rl/rl.py deleted file mode 100644 index b0ec87967e0..00000000000 --- a/src/prime_rl/rl.py +++ /dev/null @@ -1,819 +0,0 @@ -import json -import os -import shutil -import subprocess -import sys -import time -import uuid -import warnings -from pathlib import Path -from subprocess import Popen -from threading import Event, Thread -from typing import Annotated, Literal - -import tomli_w -from pydantic import Field, model_validator - -from prime_rl.inference.config import InferenceConfig -from prime_rl.inference.config import WeightBroadcastConfig as InferenceWeightBroadcastConfig -from prime_rl.orchestrator.config import CheckpointConfig as OrchestratorCheckpointConfig -from prime_rl.orchestrator.config import FileSystemWeightBroadcastConfig as OrchestratorFileSystemWeightBroadcastConfig -from prime_rl.orchestrator.config import NCCLWeightBroadcastConfig as OrchestratorNCCLWeightBroadcastConfig -from prime_rl.orchestrator.config import OrchestratorConfig -from prime_rl.trainer.config import CheckpointConfig as TrainerCheckpointConfig -from prime_rl.trainer.rl.config import FakeDataLoaderConfig -from prime_rl.trainer.rl.config import FileSystemWeightBroadcastConfig as TrainerFileSystemWeightBroadcastConfig -from prime_rl.trainer.rl.config import NCCLWeightBroadcastConfig as TrainerNCCLWeightBroadcastConfig -from prime_rl.trainer.rl.config import RLTrainerConfig as TrainerConfig -from prime_rl.utils.config import WandbConfig, WandbWithExtrasConfig -from prime_rl.utils.logger import setup_logger -from prime_rl.utils.pydantic_config import BaseSettings, get_temp_toml_file, parse_argv -from prime_rl.utils.utils import ( - get_broadcast_dir, - get_free_port, - get_log_dir, - get_rollout_dir, -) -from prime_rl.utils.validation import ( - validate_shared_ckpt_config, - validate_shared_max_async_level, - validate_shared_max_steps, - validate_shared_model_name, - validate_shared_output_dir, - validate_shared_wandb_config, - validate_shared_weight_broadcast, -) - - -class SharedLogConfig(BaseSettings): - """Configures shared logging.""" - - level: Annotated[str | None, Field(description="The log level to use.")] = "info" - - file: Annotated[bool | None, Field(description="Whether to log to a file.")] = True - - -class SharedWandbConfig(BaseSettings): - """Configures shared W&B configs.""" - - project: Annotated[str | None, Field(description="The W&B project to use.")] = "prime-rl" - - name: Annotated[str | None, Field(description="The W&B run name to use.")] = None - - offline: Annotated[bool | None, Field(description="Whether to run W&B in offline mode.")] = False - - -class SharedCheckpointConfig(BaseSettings): - """Configures shared checkpoint configs.""" - - interval: Annotated[int | None, Field(description="The interval at which to save checkpoints.")] = None - - resume_step: Annotated[ - int | None, Field(description="The step to resume from. If None, will not resume from a checkpoint.") - ] = None - - keep_last: Annotated[ - int | None, - Field( - ge=1, - description="Keep at most this many recent step checkpoints on disk. If None, never clean old checkpoints based on recency.", - ), - ] = None - - keep_interval: Annotated[ - int | None, - Field( - ge=1, - description="Keep checkpoints at every N steps permanently (e.g., keep_interval=100 keeps step 100, 200, ...). If None, no interval-based keeping.", - ), - ] = None - - -class SharedModelConfig(BaseSettings): - """Configures shared model settings.""" - - name: Annotated[ - str, - Field(description="The name of the model to use."), - ] = "Qwen/Qwen3-0.6B" - - -class SharedWeightBroadcastConfig(BaseSettings): - """Configures shared weight broadcast settings.""" - - type: Annotated[Literal["nccl", "filesystem"], Field(description="The type of weight broadcast to use.")] = ( - "filesystem" - ) - - -class RLConfig(BaseSettings): - """Configures an RL training run.""" - - ### Submodule configurations - - trainer: TrainerConfig - orchestrator: OrchestratorConfig - inference: Annotated[ - InferenceConfig | None, - Field( - description="The inference config. If None, will not start an inference process. Only viable, if an inference server was started manually." - ), - ] = None - - ### Top-level configurations - - log: Annotated[ - SharedLogConfig, - Field( - description="Shared log configs. If None, will fallback to the log configs specified on submodule configs." - ), - ] = SharedLogConfig() - - clean: Annotated[ - bool, - Field( - description="Whether to clean the rollouts, checkpoint, checkpoint weights and logs directories at the beginning of the run. If True, will forceably, and irreversibly, delete all directories.", - ), - ] = True - - inference_gpu_ids: Annotated[list[int], Field(description="The GPU IDs to use for inference.")] = [0] - trainer_gpu_ids: Annotated[list[int], Field(description="The GPU IDs to use for trainer.")] = [1] - teacher_gpu_ids: Annotated[ - list[int] | None, - Field( - description="The GPU IDs to use for teacher inference. If None, teacher inference server will not be started." - ), - ] = None - - teacher_inference: Annotated[ - InferenceConfig | None, - Field( - description="The teacher inference config. If None, will use the same config as inference (if available) or a default config. Only used when teacher_gpu_ids is set." - ), - ] = None - - ### Shared configurations - - output_dir: Annotated[ - Path, - Field(description="The directory to store the outputs. Should typically be set to an experiment identifier."), - ] = Path("outputs") # NOTE: Must match `OUTPUT_DIR` in `tmux.sh` to see logs - - ckpt: Annotated[ - SharedCheckpointConfig | None, - Field( - description="Shared checkpoint configs. If None, will fallback to the checkpoint configs specified on submodule configs." - ), - ] = None - - wandb: Annotated[ - SharedWandbConfig | None, - Field( - description="Shared W&B configs. If None, will fallback to the W&B configs specified on submodule configs." - ), - ] = None - - model: Annotated[ - SharedModelConfig | None, - Field( - description="Shared model configs. If None, will fallback to the model configs specified on submodule configs." - ), - ] = None - - max_steps: Annotated[ - int | None, - Field( - description="The maximum number of steps to train for. If None, will fallback to the max steps specified on submodule configs." - ), - ] = None - - max_model_len: Annotated[ - int | None, - Field( - description="The maximum model length to use. If None, will fallback to the max model length specified on submodule configs." - ), - ] = None - - seq_len: Annotated[ - int | None, - Field( - description="The sequence length to use. If set, will configure both trainer.model.seq_len and orchestrator.seq_len to this value. If None, each can be set independently." - ), - ] = None - - max_async_level: Annotated[ - int | None, - Field( - description="The async level to use. If None, will fallback to the async level specified on submodule configs." - ), - ] = None - - bench: Annotated[ - bool, - Field( - description="Whether to run in benchmark mode. It will automatically set the trainer and orchestrator to benchmark mode and, if present, configure the W&B project by suffixing the project with `-bench`.", - ), - ] = False - - weight_broadcast: Annotated[ - SharedWeightBroadcastConfig | None, Field(description="The weight broadcast config.") - ] = None - - @model_validator(mode="after") - def auto_setup_dp(self): - if self.inference and len(self.inference_gpu_ids) != self.inference.parallel.dp * self.inference.parallel.tp: - assert len(self.inference_gpu_ids) % self.inference.parallel.tp == 0, ( - "Number of inference GPUs must be divisible by the tensor parallel size" - ) - self.inference.parallel.dp = len(self.inference_gpu_ids) // self.inference.parallel.tp - return self - - @model_validator(mode="after") - def auto_setup_num_train_workers(self): - non_data_parallel_size = self.trainer.model.cp * self.trainer.model.tp - - if len(self.trainer_gpu_ids) > 1: - self.orchestrator.num_train_workers = len(self.trainer_gpu_ids) // non_data_parallel_size - - return self - - @model_validator(mode="after") - def auto_setup_logs(self): - # Copy log level - if self.log is not None: - if self.log.level is not None: - self.trainer.log.level = self.log.level - self.orchestrator.log.level = self.log.level - if self.log.file is not None: - self.trainer.log.file = self.log.file - self.orchestrator.log.file = self.log.file - - return self - - ### Setup and validate shared configs - - @model_validator(mode="after") - def auto_setup_ckpt(self): - # If specified, automatically setup checkpoint configs for trainer and orchestrator - if self.ckpt is not None: - # Create checkpoint configs if not specified - if self.trainer.ckpt is None: - self.trainer.ckpt = TrainerCheckpointConfig() - if self.orchestrator.ckpt is None: - self.orchestrator.ckpt = OrchestratorCheckpointConfig() - - # If specified, use the same ckpt interval - if self.ckpt.interval is not None: - self.trainer.ckpt.interval = self.ckpt.interval - self.orchestrator.ckpt.interval = self.ckpt.interval - - # If resuming training, ensure orchestrator resume from the same step - if self.ckpt.resume_step is not None: - self.trainer.ckpt.resume_step = self.ckpt.resume_step - self.orchestrator.ckpt.resume_step = self.ckpt.resume_step - - # If specified, propagate keep policy - if self.ckpt.keep_last is not None: - self.trainer.ckpt.keep_last = self.ckpt.keep_last - self.orchestrator.ckpt.keep_last = self.ckpt.keep_last - - if self.ckpt.keep_interval is not None: - self.trainer.ckpt.keep_interval = self.ckpt.keep_interval - self.orchestrator.ckpt.keep_interval = self.ckpt.keep_interval - - validate_shared_ckpt_config(self.trainer, self.orchestrator) - - return self - - @model_validator(mode="after") - def auto_setup_wandb(self): - # If specified, automatically use shared W&B project for orchestrator and trainer - if self.wandb is not None: - if not self.trainer.wandb: - self.trainer.wandb = WandbConfig() - if not self.orchestrator.wandb: - self.orchestrator.wandb = WandbWithExtrasConfig() - - if self.wandb.project: - self.trainer.wandb.project = self.wandb.project - self.orchestrator.wandb.project = self.wandb.project - - # If specified, automatically use shared W&B name for orchestrator and trainer with suffixes - if self.wandb.name: - self.trainer.wandb.name = f"{self.wandb.name}-trainer" - self.orchestrator.wandb.name = f"{self.wandb.name}-orchestrator" - - # If specified, automatically use shared W&B offline mode for orchestrator and trainer - if self.wandb.offline: - self.trainer.wandb.offline = self.wandb.offline - self.orchestrator.wandb.offline = self.wandb.offline - - validate_shared_wandb_config(self.trainer, self.orchestrator) - - return self - - @model_validator(mode="after") - def auto_setup_bench(self): - if self.bench: - # Set trainer and orchestrator to benchmark mode - self.trainer.bench = True - self.orchestrator.bench = True - - # Configure the trainer fake data to match the orchestrator config - self.trainer.data.fake = FakeDataLoaderConfig( - batch_size=self.orchestrator.batch_size, - ) - - - trainer_bench_enabled = self.trainer.bench is not None - if trainer_bench_enabled != self.orchestrator.bench: - raise ValueError( - f"Trainer benchmark mode ({self.trainer.bench}) and orchestrator benchmark mode ({self.orchestrator.bench}) are not the same. Please specify the same benchmark mode for both." - ) - - return self - - @model_validator(mode="after") - def auto_setup_model(self): - # Use the same model for trainer, orchestrator and inference - if self.model is not None: - self.trainer.model.name = self.model.name - self.orchestrator.model.name = self.model.name - if self.inference is not None: - self.inference.model.name = self.model.name - - validate_shared_model_name(self.trainer, self.orchestrator, self.inference) - - return self - - @model_validator(mode="after") - def auto_setup_max_steps(self): - # If specified, use the same max steps for trainer and orchestrator - if self.max_steps is not None: - self.trainer.max_steps = self.max_steps - self.orchestrator.max_steps = self.max_steps - - validate_shared_max_steps(self.trainer, self.orchestrator) - - return self - - @model_validator(mode="after") - def auto_setup_async_level(self): - # If specified, use the same async level for trainer and orchestrator - if self.max_async_level is not None: - self.trainer.max_async_level = self.max_async_level - self.orchestrator.max_async_level = self.max_async_level - - validate_shared_max_async_level(self.trainer, self.orchestrator) - - return self - - @model_validator(mode="after") - def auto_setup_output_dir(self): - # If specified, use the same outputs directory for trainer and orchestrator - if self.output_dir is not None: - self.trainer.output_dir = self.output_dir - self.orchestrator.output_dir = self.output_dir / "run_default" - - validate_shared_output_dir(self.trainer, self.orchestrator) - - return self - - @model_validator(mode="after") - def auto_setup_seq_len(self): - # If specified, use the same seq_len for trainer and orchestrator - if self.seq_len is not None: - self.trainer.model.seq_len = self.seq_len - self.orchestrator.seq_len = self.seq_len - - if self.trainer.model.seq_len < self.orchestrator.seq_len: - raise ValueError( - f"Trainer model seq_len ({self.trainer.model.seq_len}) must be >= orchestrator seq_len ({self.orchestrator.seq_len}). " - f"The trainer needs to be able to handle sequences at least as long as those produced by the orchestrator." - ) - - return self - - @model_validator(mode="after") - def auto_setup_weight_broadcast(self): - if self.weight_broadcast is not None: - if self.weight_broadcast.type == "nccl": - inference_world_size = self.inference.parallel.dp * self.inference.parallel.tp if self.inference else 1 - self.trainer.weight_broadcast = TrainerNCCLWeightBroadcastConfig( - type=self.weight_broadcast.type, inference_world_size=inference_world_size - ) - self.orchestrator.weight_broadcast = OrchestratorNCCLWeightBroadcastConfig( - type=self.weight_broadcast.type - ) - elif self.weight_broadcast.type == "filesystem": - self.trainer.weight_broadcast = TrainerFileSystemWeightBroadcastConfig() - self.orchestrator.weight_broadcast = OrchestratorFileSystemWeightBroadcastConfig() - if self.inference is not None: - self.inference.weight_broadcast = InferenceWeightBroadcastConfig(type=self.weight_broadcast.type) - - validate_shared_weight_broadcast(self.trainer, self.orchestrator, self.inference) - - return self - - @model_validator(mode="after") - def auto_setup_lora(self): - if self.trainer.model.lora is not None: - if self.trainer.weight_broadcast.type == "nccl": - raise ValueError("NCCL weight broadcast does not support LoRA yet.") - - # Ensure orchestrator has LoRA config - if self.orchestrator.model.lora is None: - from prime_rl.orchestrator.config import LoRAConfig - - self.orchestrator.model.lora = LoRAConfig() - - # Auto-generate name if not provided - if self.orchestrator.model.lora.name is None: - self.orchestrator.model.lora.name = f"r{self.trainer.model.lora.rank}-a{self.trainer.model.lora.alpha}" - - if self.inference is not None: - self.inference.enable_lora = True - self.inference.max_lora_rank = self.trainer.model.lora.rank - else: - warnings.warn( - "LoRA is enabled, but inference is not configured. When manually starting the inference server, make sure to set `--enable_lora` and `--max-lora-rank`." - ) - - return self - - @model_validator(mode="after") - def warn_wandb_resume_id_missing(self): - if self.trainer.ckpt is not None and self.trainer.ckpt.resume_step is not None: - if self.trainer.wandb and not self.trainer.wandb.id: - warnings.warn( - "W&B run ID is not set for trainer even though resuming training. The current run will be created as a new run." - ) - if self.orchestrator.ckpt is not None and self.orchestrator.ckpt.resume_step is not None: - if self.orchestrator.wandb and not self.orchestrator.wandb.id: - warnings.warn( - "W&B run ID is not set for orchestrator even though resuming training. The current run will be created as a new run." - ) - return self - - @model_validator(mode="after") - def validate_enough_devices_for_nccl(self): - if self.trainer.weight_broadcast.type == "nccl": - num_gpus = len(set(self.trainer_gpu_ids + self.inference_gpu_ids)) - if num_gpus < 2: - raise ValueError("NCCL weight broadcast requires at least 2 GPUs to build the broadcast process group.") - return self - - @model_validator(mode="after") - def auto_setup_teacher_inference(self): - """Auto-configure teacher inference server and orchestrator teacher_model client.""" - if self.teacher_gpu_ids is None or len(self.teacher_gpu_ids) == 0: - return self - - import copy - - from prime_rl.orchestrator.config import TeacherModelConfig - - # Create or complete teacher_inference config - if self.teacher_inference is None: - self.teacher_inference = copy.deepcopy(self.inference) if self.inference else InferenceConfig() - # Avoid port conflict with main inference by using next port - if self.inference is not None: - self.teacher_inference.server.port = self.inference.server.port + 1 - elif self.inference is not None and self.teacher_inference.server.port == self.inference.server.port: - raise ValueError( - f"teacher_inference.server.port ({self.teacher_inference.server.port}) conflicts with " - f"inference.server.port ({self.inference.server.port}). " - "Either use different ports or let teacher_inference be auto-configured." - ) - - # Auto-configure DP based on GPU count - tp = self.teacher_inference.parallel.tp - if len(self.teacher_gpu_ids) != self.teacher_inference.parallel.dp * tp: - assert len(self.teacher_gpu_ids) % tp == 0, ( - "Number of teacher GPUs must be divisible by tensor parallel size" - ) - assert len(self.teacher_gpu_ids) > 0, "teacher_gpu_ids cannot be empty" - self.teacher_inference.parallel.dp = len(self.teacher_gpu_ids) // tp - - # Auto-configure orchestrator's teacher_model client - if self.orchestrator.teacher_model is None: - self.orchestrator.teacher_model = TeacherModelConfig() - host = self.teacher_inference.server.host or "localhost" - port = self.teacher_inference.server.port - self.orchestrator.teacher_model.client.base_url = [f"http://{host}:{port}/v1"] - self.orchestrator.teacher_model.model.name = self.teacher_inference.model.name - - return self - - @model_validator(mode="after") - def validate_teacher_model(self): - if self.trainer.loss.teacher_tau > 0 and not self.orchestrator.teacher_model: - raise ValueError( - "teacher_model must be configured when teacher_tau > 0. " - "Either set teacher_tau = 0, set teacher_gpu_ids, or configure teacher_model manually." - ) - return self - - -def cleanup_threads(threads: list[Thread]): - for thread in threads: - thread.join(timeout=5) - - -def cleanup_processes(processes: list[Popen]): - for process in processes: - if process.poll() is None: # Process is still running - process.terminate() - try: - process.wait(timeout=60) # 60 seconds to terminate gracefully - except subprocess.TimeoutExpired: - process.kill() - - -def monitor_process(process: Popen, stop_event: Event, error_queue: list, process_name: str): - """Monitor a subprocess and signal errors via shared queue""" - try: - # Wait for process to complete - process.wait() - - if process.returncode != 0: - err_msg = f"{process_name.capitalize()} failed with exit code {process.returncode}" - if process.stderr: - err_msg += f"\n{process.stderr.read().decode('utf-8')}" - error_queue.append(RuntimeError(err_msg)) - stop_event.set() - except Exception as e: - error_queue.append(RuntimeError(f"Error monitoring {process_name}: {e}")) - stop_event.set() - - -def rl(config: RLConfig): - # Setup logger - logger = setup_logger( - config.log.level or "info", log_file=config.output_dir / "logs" / "rl.log" if config.log.file else None - ) - start_command = sys.argv - logger.info("Starting RL run") - logger.debug(f"RL start command: {' '.join(start_command)}") - - # Prepare paths to communicate with the trainer - log_dir = get_log_dir(config.output_dir) - orch_log_dir = get_log_dir(config.orchestrator.output_dir) - rollout_dir = get_rollout_dir(config.orchestrator.output_dir) - broadcast_dir = get_broadcast_dir(config.orchestrator.output_dir) - - # Clean up directories if specified - if config.clean: - logger.info("Cleaning checkpoint, logs, weights, broadcast and rollout directories") - - # Cleaning logs (so that streaming logs to terminal works) - logger.info(f"Cleaning log dir ({log_dir})") - shutil.rmtree(log_dir, ignore_errors=True) - log_dir.mkdir(parents=True, exist_ok=True) - - logger.info(f"Cleaning orchestrator log dir ({orch_log_dir})") - shutil.rmtree(orch_log_dir, ignore_errors=True) - orch_log_dir.mkdir(parents=True, exist_ok=True) - - # Cleaning broadcast dir (so that orchestrator does not pre-maturely update weights) - if not ( - config.ckpt - and config.ckpt.resume_step - and config.trainer.weight_broadcast - and config.trainer.weight_broadcast.type == "filesystem" - ): - logger.info(f"Cleaning broadcast directory ({broadcast_dir})") - shutil.rmtree(broadcast_dir, ignore_errors=True) - - # Cleaning rollouts (so that trainer does not train on old rollouts) - logger.info(f"Cleaning rollout dir ({rollout_dir})") - shutil.rmtree(rollout_dir, ignore_errors=True) - - # Start processes - processes: list[Popen] = [] - monitor_threads: list[Thread] = [] - error_queue: list[Exception] = [] - stop_events: dict[str, Event] = {} - - try: - # Optionally, start inference process - if config.inference: - inference_file = get_temp_toml_file() - with open(inference_file, "wb") as f: - tomli_w.dump(config.inference.model_dump(exclude_none=True, mode="json"), f) - - inference_cmd = ["uv", "run", "inference", "@", inference_file.as_posix()] - logger.info(f"Starting inference process on GPU(s) {' '.join(map(str, config.inference_gpu_ids))}") - logger.debug(f"Inference start command: {' '.join(inference_cmd)}") - # If we don't log stdout, the server hangs - with open(log_dir / "inference.stdout", "w") as log_file: - inference_process = Popen( - inference_cmd, - env={**os.environ, "CUDA_VISIBLE_DEVICES": ",".join(map(str, config.inference_gpu_ids))}, - stdout=log_file, - stderr=log_file, - ) - processes.append(inference_process) - - # Start monitoring thread - stop_event = Event() - stop_events["inference"] = stop_event - monitor_thread = Thread( - target=monitor_process, - args=(inference_process, stop_event, error_queue, "inference"), - daemon=True, - ) - monitor_thread.start() - monitor_threads.append(monitor_thread) - else: - logger.warning( - "No inference config specified, skipping starting inference server. Is your inference server running?" - ) - - # Optionally, start teacher inference process - if config.teacher_inference: - if config.teacher_gpu_ids is None or len(config.teacher_gpu_ids) == 0: - raise ValueError( - "teacher_inference is configured but teacher_gpu_ids is not set or is empty. " - "Either set teacher_gpu_ids to start a teacher inference server, " - "or omit teacher_inference and configure orchestrator.teacher_model to use an existing server." - ) - teacher_inference_file = get_temp_toml_file() - with open(teacher_inference_file, "wb") as f: - tomli_w.dump(config.teacher_inference.model_dump(exclude_none=True, mode="json"), f) - - teacher_inference_cmd = ["uv", "run", "inference", "@", teacher_inference_file.as_posix()] - logger.info(f"Starting teacher inference process on GPU(s) {' '.join(map(str, config.teacher_gpu_ids))}") - logger.debug(f"Teacher inference start command: {' '.join(teacher_inference_cmd)}") - with open(log_dir / "teacher_inference.stdout", "w") as log_file: - teacher_inference_process = Popen( - teacher_inference_cmd, - env={**os.environ, "CUDA_VISIBLE_DEVICES": ",".join(map(str, config.teacher_gpu_ids))}, - stdout=log_file, - stderr=log_file, - ) - processes.append(teacher_inference_process) - - # Start monitoring thread - stop_event = Event() - stop_events["teacher_inference"] = stop_event - monitor_thread = Thread( - target=monitor_process, - args=(teacher_inference_process, stop_event, error_queue, "teacher_inference"), - daemon=True, - ) - monitor_thread.start() - monitor_threads.append(monitor_thread) - elif config.trainer.loss.teacher_tau > 0 or config.orchestrator.teacher_model: - logger.warning( - "No teacher_inference config specified, skipping starting teacher inference server. " - "Is your teacher inference server running? Make sure orchestrator.teacher_model is configured." - ) - - # Start orchestrator process - orchestrator_file = get_temp_toml_file() - with open(orchestrator_file, "wb") as f: - tomli_w.dump(config.orchestrator.model_dump(exclude_none=True, mode="json"), f) - - orchestrator_cmd = [ - "uv", - "run", - "orchestrator", - "@", - orchestrator_file.as_posix(), - ] - logger.info("Starting orchestrator process") - logger.debug(f"Orchestrator start command: {' '.join(orchestrator_cmd)}") - with open(log_dir / "orchestrator.stdout", "w") as log_file: - orchestrator_process = Popen( - orchestrator_cmd, - stdout=log_file, - stderr=log_file, - env={ - **os.environ, - "LOGURU_FORCE_COLORS": "1", - "WANDB_PROGRAM": "uv run rl", - "WANDB_ARGS": json.dumps(start_command), - }, - ) - processes.append(orchestrator_process) - - # Start monitoring thread - stop_event = Event() - stop_events["orchestrator"] = stop_event - monitor_thread = Thread( - target=monitor_process, - args=(orchestrator_process, stop_event, error_queue, "orchestrator"), - daemon=True, - ) - monitor_thread.start() - monitor_threads.append(monitor_thread) - - # Start training process - trainer_file = get_temp_toml_file() - with open(trainer_file, "wb") as f: - tomli_w.dump(config.trainer.model_dump(exclude_none=True, mode="json"), f) - - trainer_cmd = [ - "uv", - "run", - "env", - "PYTHONUNBUFFERED=1", - "PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True", - "torchrun", - f"--rdzv-endpoint=localhost:{get_free_port()}", - f"--rdzv-id={uuid.uuid4().hex}", - # Pipe all logs to file, and only master rank logs to stdout - f"--log-dir={config.output_dir / 'torchrun'}", - "--local-ranks-filter=0", - "--redirect=3", - "--tee=3", - f"--nproc-per-node={len(config.trainer_gpu_ids)}", - "-m", - "prime_rl.trainer.rl.train", - "@", - trainer_file.as_posix(), - ] - logger.info(f"Starting trainer process on GPU(s) {' '.join(map(str, config.trainer_gpu_ids))}") - logger.debug(f"Training start command: {' '.join(trainer_cmd)}") - with open(log_dir / "trainer.stdout", "w") as log_file: - trainer_process = Popen( - trainer_cmd, - env={ - **os.environ, - "CUDA_VISIBLE_DEVICES": ",".join(map(str, config.trainer_gpu_ids)), - "PYTORCH_CUDA_ALLOC_CONF": "expandable_segments:True", - "LOGURU_FORCE_COLORS": "1", - "WANDB_PROGRAM": "uv run rl", - "WANDB_ARGS": json.dumps(start_command), - }, - stdout=log_file, - stderr=log_file, - ) - processes.append(trainer_process) - - # Start monitoring thread - stop_event = Event() - stop_events["trainer"] = stop_event - monitor_thread = Thread( - target=monitor_process, args=(trainer_process, stop_event, error_queue, "trainer"), daemon=True - ) - monitor_thread.start() - monitor_threads.append(monitor_thread) - - # Monitor all processes for failures - logger.success("Startup complete. Showing trainer logs...") - - tail_process = Popen(["tail", "-F", log_dir / "trainer.stdout"]) - processes.append(tail_process) - - # Check for errors from monitor threads - while not (stop_events["orchestrator"].is_set() and stop_events["trainer"].is_set()): - if error_queue: - error = error_queue[0] - logger.error(f"Error: {error}") - logger.error("Terminating all processes...") - cleanup_threads(monitor_threads) - cleanup_processes(processes) - sys.exit(1) - - # Small delay to avoid busy waiting - time.sleep(1) - - # Check if any critical process failed - if orchestrator_process.returncode != 0: - logger.error(f"Orchestrator failed with exit code {orchestrator_process.returncode}") - cleanup_threads(monitor_threads) - cleanup_processes(processes) - sys.exit(1) - - if trainer_process.returncode != 0: - logger.error(f"Trainer failed with exit code {trainer_process.returncode}") - cleanup_threads(monitor_threads) - cleanup_processes(processes) - sys.exit(1) - - logger.success("RL training finished!") - - # Cleanup threads and processes - cleanup_threads(monitor_threads) - cleanup_processes(processes) - - except KeyboardInterrupt: - logger.warning("Received interrupt signal, terminating all processes...") - cleanup_threads(monitor_threads) - cleanup_processes(processes) - sys.exit(1) - except Exception as e: - logger.error(f"Error occurred: {e}") - cleanup_threads(monitor_threads) - cleanup_processes(processes) - raise - - -def main(): - rl(parse_argv(RLConfig)) - - -if __name__ == "__main__": - main() diff --git a/src/prime_rl/synthesize/__init__.py b/src/prime_rl/synthesize/__init__.py deleted file mode 100644 index e69de29bb2d..00000000000 diff --git a/src/prime_rl/synthesize/config.py b/src/prime_rl/synthesize/config.py deleted file mode 100644 index 396ff185bda..00000000000 --- a/src/prime_rl/synthesize/config.py +++ /dev/null @@ -1,44 +0,0 @@ -from pathlib import Path -from typing import Annotated - -from pydantic import Field - -from prime_rl.orchestrator.config import EvalConfig -from prime_rl.utils.config import ClientConfig, LogConfig, ModelConfig -from prime_rl.utils.pydantic_config import BaseSettings - - -class SynthesizeConfig(EvalConfig, BaseSettings): - """Configures synthetic data generation.""" - - # The client configuration - client: ClientConfig = ClientConfig( - timeout=36000, base_url=["https://api.openai.com/v1"], api_key_var="OPENAI_API_KEY" - ) - - # The model configuration - model: ModelConfig = ModelConfig(name="gpt-4.1-mini") - - # The logging configuration - log: LogConfig = LogConfig() - - reasoning_field: Annotated[ - str, - Field( - description="The field in the raw model response that contains the reasoning content. Defaults to 'reasoning_content', which is the default for vLLM when serving a model with a reasoning parser. Other APIs (e.g. DeepSeek, GLM, etc.) may use different field names.", - ), - ] = "reasoning_content" - - output_dir: Annotated[ - Path, - Field( - description="Directory to write outputs to. Will be populated with artifacts such as reports and HF datasets as subdirectories. Should be set to a persistent directory with enough disk space." - ), - ] = Path("outputs") - - max_concurrent: Annotated[ - int | None, - Field( - description="Maximum number of concurrent rollouts to generate and score. Will create a global semaphore and pass to verifiers Environment. If None, will not limit concurrency.", - ), - ] = None diff --git a/src/prime_rl/synthesize/synthesize.py b/src/prime_rl/synthesize/synthesize.py deleted file mode 100644 index cb2c0e0427c..00000000000 --- a/src/prime_rl/synthesize/synthesize.py +++ /dev/null @@ -1,82 +0,0 @@ -import asyncio - -from prime_rl.orchestrator.utils import ( - set_semaphore, -) -from prime_rl.synthesize.config import SynthesizeConfig -from prime_rl.synthesize.utils import generate_synthetic_data -from prime_rl.utils.client import ( - check_has_model, - check_health, - setup_admin_clients, - setup_clients, -) -from prime_rl.utils.logger import intercept_verifiers_logging, setup_logger -from prime_rl.utils.pydantic_config import parse_argv -from prime_rl.utils.utils import clean_exit, get_env_ids_to_install, install_env - - -@clean_exit -async def synthesize(config: SynthesizeConfig): - # Initialize the logger - logger = setup_logger( - config.log.level, log_file=config.output_dir / "logs" / "synthesize.log" if config.log.file else None - ) - intercept_verifiers_logging(level=config.log.vf_level) - - env_names = [env.name or env.id for env in config.env] - logger.info(f"Starting synthetic data generation for {config.model.name} in environments {', '.join(env_names)}") - logger.info(f"Using sampling config {config.sampling}") - - # Install environments - env_ids_to_install = get_env_ids_to_install(config.env) - for env_id in env_ids_to_install: - install_env(env_id) - - # Setup clients - logger.info( - f"Initializing OpenAI client (base_url={', '.join(config.client.base_url)}, api_key_var={config.client.api_key_var}, headers={config.client.headers})" - ) - clients = setup_clients(config.client) - admin_clients = setup_admin_clients(config.client) - - # Check health of the client - logger.info("Waiting for inference pool to be ready") - await check_health(admin_clients) - await check_has_model(clients, config.model.name) - logger.success(f"Inference pool is healthy and serves {config.model.name}") - - # Set global semaphore - await set_semaphore(config.max_concurrent or -1) - - # Generate synthetic data - logger.info("Starting synthetic data generation") - await asyncio.gather( - *[ - generate_synthetic_data( - clients=clients, - env_id=env.id, - env_name=env.name, - env_args=env.args, - num_examples=env.num_examples or config.num_examples, - rollouts_per_example=env.rollouts_per_example or config.rollouts_per_example, - reasoning_field=config.reasoning_field, - skip_first=env.skip_first, - output_dir=config.output_dir, - model_config=config.model, - sampling_config=config.sampling, - client_config=config.client, - ) - for env in config.env - ] - ) - - logger.success("Synthetic data generation finished!") - - -def main(): - asyncio.run(synthesize(parse_argv(SynthesizeConfig))) - - -if __name__ == "__main__": - main() diff --git a/src/prime_rl/synthesize/utils.py b/src/prime_rl/synthesize/utils.py deleted file mode 100644 index 62c41dd2d5e..00000000000 --- a/src/prime_rl/synthesize/utils.py +++ /dev/null @@ -1,189 +0,0 @@ -import asyncio -import json -from itertools import cycle -from pathlib import Path -from typing import Any - -import aiofiles -import verifiers as vf -from openai import AsyncOpenAI -from tqdm.asyncio import tqdm -from verifiers import load_environment -from verifiers.envs.environment import get_results_path - -from prime_rl.orchestrator.config import ClientConfig, EvalSamplingConfig, ModelConfig -from prime_rl.utils.logger import get_logger -from prime_rl.utils.vf import generate_group - -WRITE_LOCK = asyncio.Lock() - - -def prepare_sampling_args(sampling_config: EvalSamplingConfig, client_config: ClientConfig) -> dict[str, Any]: - """Prepare sampling args for synthetic data generation.""" - # Initialize sampling args - sampling_args: dict[str, Any] = {} - - # Apply sampling arguments, if specified - if sampling_config.temperature is not None: - sampling_args["temperature"] = sampling_config.temperature - if sampling_config.max_tokens is not None: - sampling_args["max_tokens"] = sampling_config.max_tokens - if sampling_config.top_p is not None: - sampling_args["top_p"] = sampling_config.top_p - if sampling_config.reasoning_effort is not None: - sampling_args["reasoning_effort"] = sampling_config.reasoning_effort - - extra_body: dict[str, Any] = sampling_config.extra_body.copy() - - # Apply vLLM-specific sampling arguments, if specified - if sampling_config.top_k is not None: - extra_body["top_k"] = sampling_config.top_k - if sampling_config.min_p is not None: - extra_body["min_p"] = sampling_config.min_p - if sampling_config.min_tokens is not None: - extra_body["min_tokens"] = sampling_config.min_tokens - if sampling_config.repetition_penalty is not None: - extra_body["repetition_penalty"] = sampling_config.repetition_penalty - - sampling_args["extra_body"] = extra_body - - return sampling_args - - -# TODO: This is a hotfix for as long as verifiers doesn't support reasoning content parsing -def merge_reasoning_content( - completion: list[vf.ChatMessage], - trajectory: list[vf.TrajectoryStep], - reasoning_field: str = "reasoning_content", -) -> list[vf.ChatMessage]: - """Parse reasoning content from the raw model response and add it to the completion.""" - # Parse responses from trajectory - responses: list[vf.ModelResponse] = [trajectory_step["response"] for trajectory_step in trajectory] - assistant_messages: list[vf.ChatMessage] = [c for c in completion if c.get("role") == "assistant"] - assert len(assistant_messages) == len(responses), "Number of assistant messages and responses must match" - - for assistant_message, response in zip(assistant_messages, responses): - assert isinstance(response, vf.ChatCompletion) - response_message = response.choices[0].message - if getattr(response_message, reasoning_field, None) is not None: - assistant_message[reasoning_field] = getattr(response_message, reasoning_field) - - return completion - - -# TODO: Move to verifiers to avoid code drift -def make_result(state: vf.State, reasoning_field: str) -> dict: - """Translates a finished rollout state to a synthetic dataset row.""" - completion = merge_reasoning_content(state["completion"], state["trajectory"], reasoning_field) - result_dict = { - "example_id": state["example_id"], - "prompt": state["prompt"], - "completion": completion, - "task": state["task"], - "reward": state["reward"], - "generation_ms": state["timing"]["generation_ms"], - "scoring_ms": state["timing"]["scoring_ms"], - "total_ms": state["timing"]["total_ms"], - "info": state.get("info", {}), - "answer": state.get("answer", ""), - } - for metric_name, metric_value in state["metrics"].items(): - result_dict[metric_name] = metric_value - - result_dict["oai_tools"] = json.dumps(state["oai_tools"]) - - return result_dict - - -async def save_result(result_dict: dict, save_file: Path): - """Saves a finished rollout to a file.""" - async with WRITE_LOCK: - async with aiofiles.open(save_file, "a") as f: - await f.write(json.dumps(result_dict) + "\n") - - -async def make_and_save_result(state: vf.State, save_file: Path, reasoning_field: str): - """Translates and saves a finished rollout state to a synthetic dataset row.""" - result_dict = await asyncio.to_thread(make_result, state, reasoning_field) - await save_result(result_dict, save_file) - - -async def generate_and_save_group( - client: AsyncOpenAI, - env: vf.Environment, - model_name: str, - example: dict, - index: int, - rollouts_per_example: int, - sampling_args: dict, - save_file: Path, - reasoning_field: str, - pbar: tqdm, -) -> None: - logger = get_logger() - try: - states = await generate_group(client, env, model_name, example, rollouts_per_example, sampling_args) - await asyncio.gather(*[make_and_save_result(state, save_file, reasoning_field) for state in states]) - pbar.update(rollouts_per_example) - except Exception as e: - logger.error(f"Error generating synthetic data for group {index}: {repr(e)}") - - -async def generate_synthetic_data( - clients: list[AsyncOpenAI], - env_id: str, - env_name: str | None, - env_args: dict, - num_examples: int, - rollouts_per_example: int, - skip_first: int, - reasoning_field: str, - output_dir: Path, - model_config: ModelConfig, - sampling_config: EvalSamplingConfig, - client_config: ClientConfig, -) -> None: - """Generates synthetic data for an environment.""" - # Get the logger - logger = get_logger() - - # Load the environment - env_name_or_id = env_name or env_id - env = load_environment(env_id, **env_args) - try: - dataset = env.get_dataset(n=num_examples + skip_first) - except ValueError: - logger.warning(f"Could not find a training dataset for {env_name_or_id}. Falling back to eval dataset.") - dataset = env.get_eval_dataset(n=num_examples + skip_first) - if skip_first > 0: - dataset = dataset.skip(skip_first) - sampling_args = prepare_sampling_args(sampling_config, client_config) - path_to_save = get_results_path(env_name_or_id, model_config.name, base_path=output_dir) / "results.jsonl" - path_to_save.parent.mkdir(parents=True, exist_ok=True) - - logger.info( - f"Generating synthetic data for {env_name_or_id} (num_examples={len(dataset)}, {rollouts_per_example=}) {'with default args' if env_args == {} else f'with args {env_args}'}" - ) - - total_rollouts = len(dataset) * rollouts_per_example - pbar = tqdm(total=total_rollouts, desc="Generating synthetic data") - - await asyncio.gather( - *[ - generate_and_save_group( - client, - env, - model_config.name, - example, - index, - rollouts_per_example, - sampling_args, - path_to_save, - reasoning_field, - pbar, - ) - for index, (client, example) in enumerate(zip(cycle(clients), dataset.to_list())) - ] - ) - - logger.info(f"Synthetic data generated for {env_name_or_id} and saved to {path_to_save}") diff --git a/src/prime_rl/templates/inference.sbatch.j2 b/src/prime_rl/templates/inference.sbatch.j2 new file mode 100755 index 00000000000..85a1dd69fe2 --- /dev/null +++ b/src/prime_rl/templates/inference.sbatch.j2 @@ -0,0 +1,302 @@ +#!/bin/bash + +#SBATCH --job-name={{ job_name }} +#SBATCH --nodes={{ num_nodes }} +#SBATCH --ntasks-per-node=1 +#SBATCH --gres=gpu:{{ gpus_per_node }} +#SBATCH --partition={{ partition }} +{%- if account %} +#SBATCH --account={{ account }} +{%- endif %} +{%- if time %} +#SBATCH --time={{ time }} +{%- endif %} +{%- if nodelist %} +#SBATCH --nodelist={{ nodelist }} +{%- endif %} +{%- if exclude %} +#SBATCH --exclude={{ exclude }} +{%- endif %} +#SBATCH --exclusive +#SBATCH --output={{ output_dir }}/job_%j.log +#SBATCH --error={{ output_dir }}/job_%j.log + +# Configs +export NUM_NODES={{ num_nodes }} +export GPUS_PER_NODE={{ gpus_per_node }} +{% if disaggregated -%} +export NUM_PREFILL_NODES={{ num_prefill_nodes }} +export NUM_DECODE_NODES={{ num_decode_nodes }} +export NUM_PREFILL_REPLICAS={{ num_prefill_replicas }} +export NUM_DECODE_REPLICAS={{ num_decode_replicas }} +export NODES_PER_REPLICA=$((NUM_PREFILL_NODES + NUM_DECODE_NODES)) +export NODES_PER_PREFILL_REPLICA=$((NUM_PREFILL_NODES / NUM_PREFILL_REPLICAS)) +export NODES_PER_DECODE_REPLICA=$((NUM_DECODE_NODES / NUM_DECODE_REPLICAS)) +export PREFILL_PORT={{ prefill_port }} +export DECODE_PORT={{ decode_port }} +export ROUTER_PORT={{ router_port }} +export RPC_PORT={{ data_parallel_rpc_port }} +{%- elif num_nodes > 1 %} +export ROUTER_PORT={{ router_port }} +export BACKEND_PORT={{ backend_port }} +{%- endif %} + +# Paths +export PROJECT_DIR={{ project_dir }} +export CONFIG_PATH={{ config_path }} +export OUTPUT_DIR={{ output_dir }} +mkdir -p $OUTPUT_DIR/logs/inference +ln -sfn inference/node_0.log $OUTPUT_DIR/logs/inference.log + +# General +export CUDA_DEVICE_ORDER=PCI_BUS_ID +export PYTHONUNBUFFERED=1 +export OMP_NUM_THREADS=1 +export TRITON_CACHE_DIR=/tmp/triton-cache-$USER-$SLURM_JOB_ID +export VLLM_CACHE_ROOT=/tmp/.vllm_cache_$(whoami) + +# Setup environment +cd $PROJECT_DIR +[ -f .env ] && source .env +source .venv/bin/activate +uv sync --all-extras + +# Print inference URLs +export HOSTNAMES=( $( scontrol show hostnames $SLURM_JOB_NODELIST ) ) +export HOSTNAMES_STR="${HOSTNAMES[*]}" +{% if disaggregated -%} +# Build per-replica router URLs +INFER_URLS="" +ALL_ROUTER_ARGS="" +for ((r=0; r 1 %} +# Single router on node 0, all nodes as backends +INFER_URLS="http://${HOSTNAMES[0]}:${ROUTER_PORT}/v1" +ROUTER_ARGS="" +for host in ${HOSTNAMES[@]}; do + ROUTER_ARGS="$ROUTER_ARGS http://${host}:${BACKEND_PORT}" +done +export INFER_URLS +export ROUTER_ARGS +{%- else %} +INFER_URLS="http://${HOSTNAMES[0]}:{{ port }}/v1" +{%- endif %} +echo "HOSTNAMES=${HOSTNAMES[@]}" +echo "INFER_URLS=${INFER_URLS}" + +{% if pre_run_command %} +# Pre-run command +{{ pre_run_command }} +{% endif %} + +# Cleanup stale node-local state from prior jobs. Orphan python/torchrun/vllm +# processes and vLLM/torch IPC files under /dev/shm and /tmp can survive SLURM +# termination and cause the next launch to hang (e.g. decode engines stuck at +# "Waiting for READY message from DP Coordinator"). Harmless on clean nodes. +srun bash -c ' + pkill -9 -f "python.*prime_rl" 2>/dev/null + pkill -9 -f "torchrun" 2>/dev/null + pkill -9 -f "vllm-router" 2>/dev/null + pkill -9 -f "vllm" 2>/dev/null + pkill -9 -f "prime_rl" 2>/dev/null + # Also match prctl-set process names (kernel comm) like "vllm::router" + # which pkill -f does not see because -f only matches the cmdline. + pkill -9 "vllm" 2>/dev/null + pkill -9 "vllm::.*" 2>/dev/null + sleep 2 + rm -rf /dev/shm/vllm-* /dev/shm/vllm_* /tmp/vllm-* /tmp/vllm_* /tmp/torch-* /tmp/torchelastic_* 2>/dev/null + procs=$(ps -eo comm,args | grep -E "python|torchrun|vllm|vllm::" | grep -v grep | wc -l) + gpu=$(nvidia-smi --query-gpu=memory.used --format=csv,noheader,nounits | awk "{s+=\$1} END {print s}") + echo "[node-cleanup] $(hostname) procs=$procs gpu_mem=${gpu}MiB" +' + +# Run inference +srun bash -c ' + cd $PROJECT_DIR + [ -f .env ] && source .env + source .venv/bin/activate + + # Infiniband setup + IB_HCA=$(ibv_devinfo 2>/dev/null | sed -n -e "/hca_id/p" -e "/link_layer:/p" | grep -B1 InfiniBand | grep hca_id | sed -e "s/^hca_id://g" | tr -d "[[:blank:]]" | paste -sd,) + [ -n "$IB_HCA" ] && export NCCL_IB_HCA=$IB_HCA + + INFER_NODE_RANK=$SLURM_PROCID + LOCAL_IP=$(hostname -I | awk "{print \$1}") + +{%- if kv_offload %} + ulimit -l unlimited +{%- endif %} + +{% if disaggregated -%} + # NVSHMEM libs for DeepEP — ensure the rpath target exists on all nodes + if [ ! -d /tmp/deepep_build/nvshmem/lib ] && [ -d "$OUTPUT_DIR/nvshmem/lib" ]; then + mkdir -p /tmp/deepep_build/nvshmem + ln -sfn "$OUTPUT_DIR/nvshmem/lib" /tmp/deepep_build/nvshmem/lib + fi + + # Reconstruct hostnames array + read -ra HOSTNAMES <<< "$HOSTNAMES_STR" + + export VLLM_WORKER_MULTIPROC_METHOD=spawn +{%- if use_deep_gemm %} + export VLLM_USE_DEEP_GEMM=1 +{%- endif %} + export PYTORCH_CUDA_ALLOC_CONF="expandable_segments:False" + export GLOO_SOCKET_IFNAME=bond0 + export VLLM_ENGINE_READY_TIMEOUT_S=4200 + export UCX_TLS=all + export UCX_NET_DEVICES=all + export LD_LIBRARY_PATH="$PROJECT_DIR/third_party/ucx/lib:$PROJECT_DIR/third_party/ucx/lib/ucx:${LD_LIBRARY_PATH:-}" + + export VLLM_NIXL_SIDE_CHANNEL_HOST=$LOCAL_IP + export VLLM_NIXL_SIDE_CHANNEL_PORT=5600 + +{%- if kv_offload %} + PREFILL_KV_CFG='"'"'{"kv_connector":"MultiConnector","kv_role":"kv_both","kv_connector_extra_config":{"connectors":[{"kv_connector":"NixlConnector","kv_role":"kv_both","kv_connector_extra_config":{"num_threads":1}},{"kv_connector":"OffloadingConnector","kv_role":"kv_both","kv_connector_extra_config":{"cpu_bytes_to_use":{{ kv_offload_cpu_bytes }}}}]}}'"'"' +{%- else %} + PREFILL_KV_CFG='"'"'{"kv_connector":"NixlConnector","kv_role":"kv_both","kv_connector_extra_config":{"num_threads":1}}'"'"' +{%- endif %} + +{%- if kv_offload %} + DECODE_KV_CFG='"'"'{"kv_connector":"MultiConnector","kv_role":"kv_both","kv_connector_extra_config":{"connectors":[{"kv_connector":"NixlConnector","kv_role":"kv_both","kv_connector_extra_config":{"num_threads":1}},{"kv_connector":"OffloadingConnector","kv_role":"kv_both","kv_connector_extra_config":{"cpu_bytes_to_use":{{ kv_offload_cpu_bytes }}}}]}}'"'"' +{%- else %} + DECODE_KV_CFG='"'"'{"kv_connector":"NixlConnector","kv_role":"kv_both","kv_connector_extra_config":{"num_threads":1}}'"'"' +{%- endif %} + + DECODE_COMPILE_CFG='"'"'{"cudagraph_mode":"FULL_DECODE_ONLY"}'"'"' + + # Determine outer replica (P/D pair) and role + REPLICA_IDX=$((INFER_NODE_RANK / NODES_PER_REPLICA)) + RANK_IN_REPLICA=$((INFER_NODE_RANK % NODES_PER_REPLICA)) + REPLICA_BASE=$((REPLICA_IDX * NODES_PER_REPLICA)) + + if [ "$RANK_IN_REPLICA" -lt "$NUM_PREFILL_NODES" ]; then + # ── Prefill node ── + ROLE="prefill" + SUB_REPLICA=$((RANK_IN_REPLICA / NODES_PER_PREFILL_REPLICA)) + ROLE_RANK=$((RANK_IN_REPLICA % NODES_PER_PREFILL_REPLICA)) + EP=$((NODES_PER_PREFILL_REPLICA * GPUS_PER_NODE)) + PORT=$PREFILL_PORT + START_RANK=$((ROLE_RANK * GPUS_PER_NODE)) + HEAD_HOST="${HOSTNAMES[$((REPLICA_BASE + SUB_REPLICA * NODES_PER_PREFILL_REPLICA))]}" + +{% for key, value in prefill_env_overrides.items() %} + export {{ key }}="{{ value }}" +{% endfor %} + ROLE_EXTRA="\"all2all_backend\": \"deepep_high_throughput\"" + + if [ "$ROLE_RANK" -eq 0 ]; then + VLLM_EXTRA="{\"data_parallel_address\": \"$LOCAL_IP\", \"data_parallel_hybrid_lb\": true, \"kv_transfer_config\": $PREFILL_KV_CFG, $ROLE_EXTRA}" + else + VLLM_EXTRA="{\"data_parallel_address\": \"$HEAD_HOST\", \"data_parallel_start_rank\": $START_RANK, \"data_parallel_hybrid_lb\": true, \"kv_transfer_config\": $PREFILL_KV_CFG, $ROLE_EXTRA}" + fi + else + # ── Decode node ── + export UCX_NET_DEVICES=mlx5_0:1 + ROLE="decode" + RANK_IN_DECODE=$((RANK_IN_REPLICA - NUM_PREFILL_NODES)) + SUB_REPLICA=$((RANK_IN_DECODE / NODES_PER_DECODE_REPLICA)) + ROLE_RANK=$((RANK_IN_DECODE % NODES_PER_DECODE_REPLICA)) + EP=$((NODES_PER_DECODE_REPLICA * GPUS_PER_NODE)) + PORT=$DECODE_PORT + START_RANK=$((ROLE_RANK * GPUS_PER_NODE)) + HEAD_HOST="${HOSTNAMES[$((REPLICA_BASE + NUM_PREFILL_NODES + SUB_REPLICA * NODES_PER_DECODE_REPLICA))]}" + +{% for key, value in decode_env_overrides.items() %} + export {{ key }}="{{ value }}" +{% endfor %} + ROLE_EXTRA="\"all2all_backend\": \"deepep_low_latency\", \"compilation_config\": $DECODE_COMPILE_CFG" + + if [ "$ROLE_RANK" -eq 0 ]; then + VLLM_EXTRA="{\"data_parallel_address\": \"$LOCAL_IP\", \"data_parallel_hybrid_lb\": true, \"kv_transfer_config\": $DECODE_KV_CFG, $ROLE_EXTRA}" + else + VLLM_EXTRA="{\"data_parallel_address\": \"$HEAD_HOST\", \"data_parallel_start_rank\": $START_RANK, \"data_parallel_hybrid_lb\": true, \"kv_transfer_config\": $DECODE_KV_CFG, $ROLE_EXTRA}" + fi + fi + + echo "NODE_RANK=$INFER_NODE_RANK ROLE=$ROLE REPLICA=$REPLICA_IDX SUB_REPLICA=$SUB_REPLICA ROLE_RANK=$ROLE_RANK EP=$EP PORT=$PORT LOCAL_IP=$LOCAL_IP" \ + | tee $OUTPUT_DIR/logs/inference/node_${INFER_NODE_RANK}.log + + # Start vllm-router on the first prefill node of each replica (background) + if [ "$RANK_IN_REPLICA" -eq 0 ]; then + ROUTER_LOG="$OUTPUT_DIR/logs/inference/router_${REPLICA_IDX}.log" + echo "Starting vllm-router on $LOCAL_IP:$ROUTER_PORT for replica $REPLICA_IDX" | tee $ROUTER_LOG + + REPLICA_ROUTER_ARGS=$(echo "$ALL_ROUTER_ARGS" | cut -d"|" -f$((REPLICA_IDX + 1))) + + vllm-router \ + --policy {{ router_policy }} \ + --vllm-pd-disaggregation \ + $REPLICA_ROUTER_ARGS \ + --host 0.0.0.0 \ + --port $ROUTER_PORT \ + --intra-node-data-parallel-size {{ dp_per_node }} \ + --worker-startup-timeout-secs 4200 \ + --log-level debug \ + >> $ROUTER_LOG 2>&1 & + fi + + uv run inference \ + @ $CONFIG_PATH \ + --server.host 0.0.0.0 \ + --server.port $PORT \ + --parallel.dp $EP \ + --data-parallel-rpc-port $RPC_PORT \ + --vllm-extra "$VLLM_EXTRA" \ + 2>&1 | tee -a $OUTPUT_DIR/logs/inference/node_${INFER_NODE_RANK}.log +{%- elif num_nodes > 1 %} + echo "INFER_NODE_RANK=${INFER_NODE_RANK}, LOCAL_IP=${LOCAL_IP}" \ + | tee $OUTPUT_DIR/logs/inference/node_${INFER_NODE_RANK}.log + + export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:False + + # Start router on node 0 + if [ "$INFER_NODE_RANK" -eq 0 ]; then + ROUTER_LOG="$OUTPUT_DIR/logs/inference/router.log" + echo "Starting vllm-router on $LOCAL_IP:$ROUTER_PORT" | tee $ROUTER_LOG + + vllm-router \ + --policy {{ router_policy }} \ + --worker-urls $ROUTER_ARGS \ + --host 0.0.0.0 \ + --port $ROUTER_PORT \ + --intra-node-data-parallel-size {{ dp_per_node }} \ + --worker-startup-timeout-secs 4200 \ + --log-level debug \ + >> $ROUTER_LOG 2>&1 & + fi + + uv run inference \ + @ $CONFIG_PATH \ + --server.host 0.0.0.0 \ + --server.port $BACKEND_PORT \ + 2>&1 | tee -a $OUTPUT_DIR/logs/inference/node_${INFER_NODE_RANK}.log +{%- else %} + echo "INFER_NODE_RANK=${INFER_NODE_RANK}, LOCAL_IP=${LOCAL_IP}" | tee $OUTPUT_DIR/logs/inference/node_${INFER_NODE_RANK}.log + # Required for graph compilation to work + export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:False + + uv run inference \ + @ $CONFIG_PATH \ + 2>&1 | tee -a $OUTPUT_DIR/logs/inference/node_${INFER_NODE_RANK}.log +{%- endif %} +' diff --git a/src/prime_rl/templates/multi_node_rl.sbatch.j2 b/src/prime_rl/templates/multi_node_rl.sbatch.j2 new file mode 100755 index 00000000000..8016234e355 --- /dev/null +++ b/src/prime_rl/templates/multi_node_rl.sbatch.j2 @@ -0,0 +1,399 @@ +#!/bin/bash + +#SBATCH --job-name={{ job_name }} +#SBATCH --nodes={{ num_train_nodes + num_infer_nodes }} +#SBATCH --ntasks-per-node=1 +#SBATCH --gres=gpu:{{ gpus_per_node }} +#SBATCH --partition={{ partition }} +{%- if account %} +#SBATCH --account={{ account }} +{%- endif %} +{%- if time %} +#SBATCH --time={{ time }} +{%- endif %} +{%- if nodelist %} +#SBATCH --nodelist={{ nodelist }} +{%- endif %} +{%- if exclude %} +#SBATCH --exclude={{ exclude }} +{%- endif %} +#SBATCH --exclusive +#SBATCH --output={{ output_dir }}/job_%j.log +#SBATCH --error={{ output_dir }}/job_%j.log + +# Configs +export NUM_TRAIN_NODES={{ num_train_nodes }} +export NUM_INFER_NODES={{ num_infer_nodes }} +export NODES_PER_INFER_REPLICA={{ nodes_per_infer_replica }} +export NUM_INFER_REPLICAS={{ num_infer_replicas }} +export GPUS_PER_NODE={{ gpus_per_node }} +{% if num_infer_nodes > 0 -%} +export ROUTER_PORT={{ router_port }} +export INFERENCE_TP={{ inference_tp }} +export INFERENCE_DP_LOCAL=$((GPUS_PER_NODE / INFERENCE_TP)) +export INFERENCE_DATA_PARALLEL_RPC_PORT={{ inference_data_parallel_rpc_port }} +{% if is_disaggregated -%} +export NUM_PREFILL_NODES={{ num_prefill_nodes }} +export NUM_DECODE_NODES={{ num_decode_nodes }} +export NUM_PREFILL_REPLICAS={{ num_prefill_replicas }} +export NUM_DECODE_REPLICAS={{ num_decode_replicas }} +export NODES_PER_PREFILL_REPLICA=$((NUM_PREFILL_NODES / NUM_PREFILL_REPLICAS)) +export NODES_PER_DECODE_REPLICA=$((NUM_DECODE_NODES / NUM_DECODE_REPLICAS)) +export PREFILL_PORT={{ prefill_port }} +export DECODE_PORT={{ decode_port }} +{%- else -%} +export BACKEND_PORT={{ backend_port }} +export INFERENCE_ENABLE_EXPERT_PARALLEL={{ "1" if inference_enable_expert_parallel else "0" }} + +# Enable distributed EP mode only when expert parallelism spans multiple inference nodes per replica. +export USE_DISTRIBUTED_EP=0 +if [ "$INFERENCE_ENABLE_EXPERT_PARALLEL" -eq 1 ] && [ "$NODES_PER_INFER_REPLICA" -gt 1 ]; then + USE_DISTRIBUTED_EP=1 +fi +{%- endif %} +{%- endif %} + +# Paths +export PROJECT_DIR={{ project_dir }} +export CONFIG_DIR={{ config_dir }} +export OUTPUT_DIR={{ output_dir }} +export ORCHESTRATOR_OUTPUT_DIR={{ orchestrator_output_dir }} +mkdir -p $OUTPUT_DIR/logs/trainer $OUTPUT_DIR/logs/inference +ln -sfn trainer/node_0.log $OUTPUT_DIR/logs/trainer.log +ln -sfn inference/node_0.log $OUTPUT_DIR/logs/inference.log + +# General +export CUDA_DEVICE_ORDER=PCI_BUS_ID +export PYTHONUNBUFFERED=1 +export OMP_NUM_THREADS=1 +export TRITON_CACHE_DIR=/tmp/triton-cache-$USER-$SLURM_JOB_ID + +{% if wandb_shared -%} +export WANDB_SHARED_MODE=1 +export WANDB_SHARED_RUN_ID=$(python3 -c "import uuid; print(uuid.uuid4().hex)") +{%- endif %} + +# Networking +HOSTNAMES=( $( scontrol show hostnames $SLURM_JOB_NODELIST ) ) +TRAIN_HOSTS=( ${HOSTNAMES[@]:$NUM_INFER_NODES:$SLURM_JOB_NUM_NODES} ) +export HOSTNAMES_STR="${HOSTNAMES[*]}" + +{% if num_infer_nodes > 0 -%} +INFER_HOSTS=( ${HOSTNAMES[@]:0:$NUM_INFER_NODES} ) + +# Build per-replica router URLs, admin URLs, and backend args for the router +INFER_URLS="" +ADMIN_URLS="" +ALL_ROUTER_ARGS="" +{% if is_disaggregated -%} +for ((r=0; r/dev/null + pkill -9 -f "torchrun" 2>/dev/null + pkill -9 -f "vllm-router" 2>/dev/null + pkill -9 -f "vllm" 2>/dev/null + pkill -9 -f "prime_rl" 2>/dev/null + # Also match prctl-set process names (kernel comm) like "vllm::router" + # which pkill -f does not see because -f only matches the cmdline. + pkill -9 "vllm" 2>/dev/null + pkill -9 "vllm::.*" 2>/dev/null + sleep 2 + rm -rf /dev/shm/vllm-* /dev/shm/vllm_* /tmp/vllm-* /tmp/vllm_* /tmp/torch-* /tmp/torchelastic_* 2>/dev/null + procs=$(ps -eo comm,args | grep -E "python|torchrun|vllm|vllm::" | grep -v grep | wc -l) + gpu=$(nvidia-smi --query-gpu=memory.used --format=csv,noheader,nounits | awk "{s+=\$1} END {print s}") + echo "[node-cleanup] $(hostname) procs=$procs gpu_mem=${gpu}MiB" +' + +# Run RL +srun bash -c ' + # Source environment + cd $PROJECT_DIR + [ -f .env ] && source .env + source .venv/bin/activate + # Do not sync to avoid conflicts with lockfile + + # Higher ulimit + ulimit -n 65536 + export GIT_LFS_SKIP_SMUDGE=1 + + # Infiniband setup + IB_HCA=$(ibv_devinfo | sed -n -e "/hca_id/p" -e "/link_layer:/p" | grep -B1 InfiniBand | grep hca_id | sed -e "s/^hca_id://g" | tr -d "[[:blank:]]" |paste -sd,) + export NCCL_IB_HCA=$IB_HCA + +{% if num_infer_nodes > 0 -%} + +if [ "$SLURM_PROCID" -lt "$NUM_INFER_NODES" ]; then + export PYTORCH_CUDA_ALLOC_CONF="expandable_segments:False" + INFER_NODE_RANK=$SLURM_PROCID + export VLLM_WORKER_MULTIPROC_METHOD=spawn + LOCAL_IP=$(hostname -I | awk "{print \$1}") + +{%- if kv_offload %} + ulimit -l unlimited +{%- endif %} + + REPLICA_IDX=$((INFER_NODE_RANK / NODES_PER_INFER_REPLICA)) + RANK_IN_REPLICA=$((INFER_NODE_RANK % NODES_PER_INFER_REPLICA)) + + echo "INFER_NODE_RANK=$INFER_NODE_RANK REPLICA=$REPLICA_IDX RANK_IN_REPLICA=$RANK_IN_REPLICA LOCAL_IP=$LOCAL_IP" \ + | tee $OUTPUT_DIR/logs/inference/node_${INFER_NODE_RANK}.log + +{% if is_disaggregated -%} + # PD disaggregated: NIXL KV transfer + role-based inference + read -ra HOSTNAMES <<< "$HOSTNAMES_STR" +{%- if use_deep_gemm %} + export VLLM_USE_DEEP_GEMM=1 +{%- endif %} + + export GLOO_SOCKET_IFNAME=bond0 + export VLLM_ENGINE_READY_TIMEOUT_S=4200 + export UCX_TLS=all + export UCX_NET_DEVICES=all + export LD_LIBRARY_PATH="$PROJECT_DIR/third_party/ucx/lib:$PROJECT_DIR/third_party/ucx/lib/ucx:${LD_LIBRARY_PATH:-}" + export VLLM_NIXL_SIDE_CHANNEL_HOST=$LOCAL_IP + export VLLM_NIXL_SIDE_CHANNEL_PORT=5600 + +{%- if kv_offload %} + PREFILL_KV_CFG='"'"'{"kv_connector":"MultiConnector","kv_role":"kv_both","kv_connector_extra_config":{"connectors":[{"kv_connector":"NixlConnector","kv_role":"kv_both","kv_connector_extra_config":{"num_threads":1}},{"kv_connector":"OffloadingConnector","kv_role":"kv_both","kv_connector_extra_config":{"cpu_bytes_to_use":{{ kv_offload_cpu_bytes }}}}]}}'"'"' +{%- else %} + PREFILL_KV_CFG='"'"'{"kv_connector":"NixlConnector","kv_role":"kv_both","kv_connector_extra_config":{"num_threads":1}}'"'"' +{%- endif %} + +{%- if kv_offload %} + DECODE_KV_CFG='"'"'{"kv_connector":"MultiConnector","kv_role":"kv_both","kv_connector_extra_config":{"connectors":[{"kv_connector":"NixlConnector","kv_role":"kv_both","kv_connector_extra_config":{"num_threads":1}},{"kv_connector":"OffloadingConnector","kv_role":"kv_both","kv_connector_extra_config":{"cpu_bytes_to_use":{{ kv_offload_cpu_bytes }}}}]}}'"'"' +{%- else %} + DECODE_KV_CFG='"'"'{"kv_connector":"NixlConnector","kv_role":"kv_both","kv_connector_extra_config":{"num_threads":1}}'"'"' +{%- endif %} + DECODE_COMPILE_CFG='"'"'{"cudagraph_mode":"FULL_DECODE_ONLY"}'"'"' + + REPLICA_BASE=$((REPLICA_IDX * NODES_PER_INFER_REPLICA)) + + if [ "$RANK_IN_REPLICA" -lt "$NUM_PREFILL_NODES" ]; then + ROLE="prefill" + SUB_REPLICA=$((RANK_IN_REPLICA / NODES_PER_PREFILL_REPLICA)) + ROLE_RANK=$((RANK_IN_REPLICA % NODES_PER_PREFILL_REPLICA)) + DP=$((NODES_PER_PREFILL_REPLICA * INFERENCE_DP_LOCAL)) + PORT=$PREFILL_PORT + START_RANK=$((ROLE_RANK * INFERENCE_DP_LOCAL)) + HEAD_HOST="${HOSTNAMES[$((REPLICA_BASE + SUB_REPLICA * NODES_PER_PREFILL_REPLICA))]}" +{% for key, value in prefill_env_overrides.items() %} + export {{ key }}="{{ value }}" +{% endfor %} + ROLE_EXTRA="\"all2all_backend\": \"deepep_high_throughput\"" + else + export UCX_NET_DEVICES=mlx5_0:1 + ROLE="decode" + RANK_IN_DECODE=$((RANK_IN_REPLICA - NUM_PREFILL_NODES)) + SUB_REPLICA=$((RANK_IN_DECODE / NODES_PER_DECODE_REPLICA)) + ROLE_RANK=$((RANK_IN_DECODE % NODES_PER_DECODE_REPLICA)) + DP=$((NODES_PER_DECODE_REPLICA * INFERENCE_DP_LOCAL)) + PORT=$DECODE_PORT + START_RANK=$((ROLE_RANK * INFERENCE_DP_LOCAL)) + HEAD_HOST="${HOSTNAMES[$((REPLICA_BASE + NUM_PREFILL_NODES + SUB_REPLICA * NODES_PER_DECODE_REPLICA))]}" +{% for key, value in decode_env_overrides.items() %} + export {{ key }}="{{ value }}" +{% endfor %} + ROLE_EXTRA="\"all2all_backend\": \"deepep_low_latency\", \"compilation_config\": $DECODE_COMPILE_CFG" + fi + + echo "ROLE=$ROLE REPLICA=$REPLICA_IDX SUB_REPLICA=$SUB_REPLICA ROLE_RANK=$ROLE_RANK DP=$DP PORT=$PORT" \ + | tee -a $OUTPUT_DIR/logs/inference/node_${INFER_NODE_RANK}.log + + if [ "$DP" -gt 1 ]; then + if [ "$ROLE" = "prefill" ]; then + KV_CFG_TO_USE="$PREFILL_KV_CFG" + else + KV_CFG_TO_USE="$DECODE_KV_CFG" + fi + + if [ "$ROLE_RANK" -eq 0 ]; then + VLLM_EXTRA="{\"data_parallel_address\": \"$LOCAL_IP\", \"data_parallel_hybrid_lb\": true, \"kv_transfer_config\": $KV_CFG_TO_USE, $ROLE_EXTRA}" + else + VLLM_EXTRA="{\"data_parallel_address\": \"$HEAD_HOST\", \"data_parallel_start_rank\": $START_RANK, \"data_parallel_hybrid_lb\": true, \"kv_transfer_config\": $KV_CFG_TO_USE, $ROLE_EXTRA}" + fi + else + if [ "$ROLE" = "prefill" ]; then + VLLM_EXTRA="{\"kv_transfer_config\": $PREFILL_KV_CFG, $ROLE_EXTRA}" + else + VLLM_EXTRA="{\"kv_transfer_config\": $DECODE_KV_CFG, $ROLE_EXTRA}" + fi + fi + + # Start vllm-router on the first node of each replica (PD mode) + if [ "$RANK_IN_REPLICA" -eq 0 ]; then + ROUTER_LOG="$OUTPUT_DIR/logs/inference/router_${REPLICA_IDX}.log" + echo "Starting vllm-router (PD) on $LOCAL_IP:$ROUTER_PORT for replica $REPLICA_IDX" | tee $ROUTER_LOG + + REPLICA_ROUTER_ARGS=$(echo "$ALL_ROUTER_ARGS" | cut -d"|" -f$((REPLICA_IDX + 1))) + + vllm-router \ + --policy consistent_hash \ + --vllm-pd-disaggregation \ + $REPLICA_ROUTER_ARGS \ + --host 0.0.0.0 \ + --port $ROUTER_PORT \ + --intra-node-data-parallel-size {{ dp_per_node }} \ + --worker-startup-timeout-secs 4200 \ + >> $ROUTER_LOG 2>&1 & + fi + + uv run inference \ + @ $CONFIG_DIR/inference.toml \ + --server.host 0.0.0.0 \ + --server.port $PORT \ + --parallel.dp $DP \ + --data-parallel-rpc-port $INFERENCE_DATA_PARALLEL_RPC_PORT \ + --vllm-extra "$VLLM_EXTRA" \ + 2>&1 | tee -a $OUTPUT_DIR/logs/inference/node_${INFER_NODE_RANK}.log +{%- else -%} + # Start vllm-router on the first node of each replica + if [ "$RANK_IN_REPLICA" -eq 0 ]; then + ROUTER_LOG="$OUTPUT_DIR/logs/inference/router_${REPLICA_IDX}.log" + echo "Starting vllm-router on $LOCAL_IP:$ROUTER_PORT for replica $REPLICA_IDX" | tee $ROUTER_LOG + + REPLICA_ROUTER_ARGS=$(echo "$ALL_ROUTER_ARGS" | cut -d"|" -f$((REPLICA_IDX + 1))) + + vllm-router \ + --policy consistent_hash \ + --worker-urls $REPLICA_ROUTER_ARGS \ + --host 0.0.0.0 \ + --port $ROUTER_PORT \ + --intra-node-data-parallel-size {{ dp_per_node }} \ + --worker-startup-timeout-secs 4200 \ + >> $ROUTER_LOG 2>&1 & + fi + + if [ "$USE_DISTRIBUTED_EP" -eq 1 ]; then + read -ra HOSTNAMES <<< "$HOSTNAMES_STR" + REPLICA_HEAD_HOST="${HOSTNAMES[$((REPLICA_IDX * NODES_PER_INFER_REPLICA))]}" + INFER_DP_START_RANK=$((RANK_IN_REPLICA * INFERENCE_DP_LOCAL)) + + if [ "$RANK_IN_REPLICA" -eq 0 ]; then + VLLM_EXTRA="{\"data_parallel_address\": \"$LOCAL_IP\", \"data_parallel_hybrid_lb\": true}" + else + VLLM_EXTRA="{\"data_parallel_address\": \"$REPLICA_HEAD_HOST\", \"data_parallel_start_rank\": $INFER_DP_START_RANK, \"data_parallel_hybrid_lb\": true}" + fi + + uv run inference \ + @ $CONFIG_DIR/inference.toml \ + --server.host 0.0.0.0 \ + --server.port $BACKEND_PORT \ + --data-parallel-rpc-port $INFERENCE_DATA_PARALLEL_RPC_PORT \ + --vllm-extra "$VLLM_EXTRA" \ + 2>&1 | tee -a $OUTPUT_DIR/logs/inference/node_${INFER_NODE_RANK}.log + else + uv run inference \ + @ $CONFIG_DIR/inference.toml \ + --server.host 0.0.0.0 \ + --server.port $BACKEND_PORT \ + 2>&1 | tee -a $OUTPUT_DIR/logs/inference/node_${INFER_NODE_RANK}.log + fi +{%- endif %} + +else + TRAIN_NODE_RANK=$((SLURM_PROCID - NUM_INFER_NODES)) + + if [ "$TRAIN_NODE_RANK" -eq 0 ]; then + ORCHESTRATOR_ARGS="@ $CONFIG_DIR/orchestrator.toml" + ORCHESTRATOR_ARGS="$ORCHESTRATOR_ARGS --client.base-url $INFER_URLS" + ORCHESTRATOR_ARGS="$ORCHESTRATOR_ARGS --client.admin-base-url $ADMIN_URLS" + {% if use_nccl_broadcast %}ORCHESTRATOR_ARGS="$ORCHESTRATOR_ARGS --weight_broadcast.host $MASTER_ADDR" + {% endif %}{% if wandb_shared %}WANDB_SHARED_LABEL=orchestrator {% endif %}uv run orchestrator $ORCHESTRATOR_ARGS \ + 2>&1 | tee $OUTPUT_DIR/logs/orchestrator.log & + fi +{%- else %} + TRAIN_NODE_RANK=$SLURM_PROCID +{%- endif %} + + export HF_HUB_OFFLINE=1 + + # This is required for compilation to work correctly + export PYTORCH_CUDA_ALLOC_CONF="expandable_segments:True" + + echo $HOSTNAMES_STR | tee $OUTPUT_DIR/logs/trainer/node_${TRAIN_NODE_RANK}.log + + {% if wandb_shared %}WANDB_SHARED_LABEL=trainer {% endif %}uv run torchrun \ + --role=trainer \ + --nnodes=$NUM_TRAIN_NODES \ + --nproc-per-node=$GPUS_PER_NODE \ + --node-rank=$TRAIN_NODE_RANK \ + --rdzv-endpoint=$MASTER_ADDR:$MASTER_PORT \ + --rdzv-id=job_$SLURM_JOB_ID \ + --log-dir=$OUTPUT_DIR/logs/trainer/torchrun \ + --tee=3 \ + --redirects=3 \ + --local-ranks-filter={{ ranks_filter }} \ + -m prime_rl.trainer.rl.train \ + @ $CONFIG_DIR/trainer.toml \ + 2>&1 | sed -u 's/^\[[a-zA-Z]*[0-9]*\]://' | tee -a $OUTPUT_DIR/logs/trainer/node_${TRAIN_NODE_RANK}.log +{% if num_infer_nodes > 0 %} fi{% endif %} +' diff --git a/src/prime_rl/templates/multi_node_sft.sbatch.j2 b/src/prime_rl/templates/multi_node_sft.sbatch.j2 new file mode 100644 index 00000000000..1d5ce7e9615 --- /dev/null +++ b/src/prime_rl/templates/multi_node_sft.sbatch.j2 @@ -0,0 +1,119 @@ +#!/bin/bash + +#SBATCH --job-name={{ job_name }} +#SBATCH --nodes={{ num_nodes }} +#SBATCH --ntasks-per-node=1 +#SBATCH --gres=gpu:{{ gpus_per_node }} +#SBATCH --partition={{ partition }} +{%- if account %} +#SBATCH --account={{ account }} +{%- endif %} +{%- if time %} +#SBATCH --time={{ time }} +{%- endif %} +{%- if nodelist %} +#SBATCH --nodelist={{ nodelist }} +{%- endif %} +{%- if exclude %} +#SBATCH --exclude={{ exclude }} +{%- endif %} +#SBATCH --exclusive +#SBATCH --output={{ output_dir }}/job_%j.log +#SBATCH --error={{ output_dir }}/job_%j.log + +# Configs +export NUM_NODES={{ num_nodes }} +export GPUS_PER_NODE={{ gpus_per_node }} + +# Paths +export PROJECT_DIR={{ project_dir }} +export CONFIG_PATH={{ config_path }} +export OUTPUT_DIR={{ output_dir }} +mkdir -p $OUTPUT_DIR/logs/trainer +ln -sfn trainer/node_0.log $OUTPUT_DIR/logs/trainer.log + +# General +export CUDA_DEVICE_ORDER=PCI_BUS_ID +export PYTHONUNBUFFERED=1 +export OMP_NUM_THREADS=1 +export TRITON_CACHE_DIR=/tmp/triton-cache-$USER-$SLURM_JOB_ID + +# Networking +export HOSTNAMES=( $( scontrol show hostnames $SLURM_JOB_NODELIST ) ) +echo "HOSTNAMES=${HOSTNAMES[@]}" + +export MASTER_ADDR="${HOSTNAMES[0]}" +export MASTER_PORT=29500 +echo "MASTER_ADDR=${MASTER_ADDR}" +echo "MASTER_PORT=${MASTER_PORT}" + +# Install environment +[ -f $PROJECT_DIR/.env ] && source $PROJECT_DIR/.env +source $PROJECT_DIR/.venv/bin/activate + +# Install environment as local package +cd $PROJECT_DIR && uv sync --all-extras + + +{% if pre_run_command %} +# Pre-run command +{{ pre_run_command }} +{% endif %} + +# Cleanup stale node-local state from prior jobs. Orphan python/torchrun/vllm +# processes and vLLM/torch IPC files under /dev/shm and /tmp can survive SLURM +# termination and cause the next launch to hang. Harmless on clean nodes. +srun bash -c ' + pkill -9 -f "python.*prime_rl" 2>/dev/null + pkill -9 -f "torchrun" 2>/dev/null + pkill -9 -f "vllm-router" 2>/dev/null + pkill -9 -f "vllm" 2>/dev/null + pkill -9 -f "prime_rl" 2>/dev/null + # Also match prctl-set process names (kernel comm) like "vllm::router" + # which pkill -f does not see because -f only matches the cmdline. + pkill -9 "vllm" 2>/dev/null + pkill -9 "vllm::.*" 2>/dev/null + sleep 2 + rm -rf /dev/shm/vllm-* /dev/shm/vllm_* /tmp/vllm-* /tmp/vllm_* /tmp/torch-* /tmp/torchelastic_* 2>/dev/null + procs=$(ps -eo comm,args | grep -E "python|torchrun|vllm|vllm::" | grep -v grep | wc -l) + gpu=$(nvidia-smi --query-gpu=memory.used --format=csv,noheader,nounits | awk "{s+=\$1} END {print s}") + echo "[node-cleanup] $(hostname) procs=$procs gpu_mem=${gpu}MiB" +' + +# Run SFT +srun bash -c ' + # Setup environment + [ -f $PROJECT_DIR/.env ] && source $PROJECT_DIR/.env + source $PROJECT_DIR/.venv/bin/activate + + # Higher ulimit + ulimit -n 65536 + export GIT_LFS_SKIP_SMUDGE=1 + + # Infiniband setup + IB_HCA=$(ibv_devinfo | sed -n -e "/hca_id/p" -e "/link_layer:/p" | grep -B1 InfiniBand | grep hca_id | sed -e "s/^hca_id://g" | tr -d "[[:blank:]]" |paste -sd,) + export NCCL_IB_HCA=$IB_HCA + + TRAIN_NODE_RANK=$SLURM_PROCID + + export HF_HUB_OFFLINE=1 + + # This is required for compilation to work correctly + export PYTORCH_CUDA_ALLOC_CONF="expandable_segments:True" + + echo $HOSTNAMES | tee $OUTPUT_DIR/logs/trainer/node_${TRAIN_NODE_RANK}.log + uv run torchrun \ + --role=trainer \ + --nnodes=$NUM_NODES \ + --nproc-per-node=$GPUS_PER_NODE \ + --node-rank=$TRAIN_NODE_RANK \ + --rdzv-endpoint=$MASTER_ADDR:$MASTER_PORT \ + --rdzv-id=job_$SLURM_JOB_ID \ + --log-dir=$OUTPUT_DIR/logs/trainer/torchrun \ + --tee=3 \ + --redirects=3 \ + --local-ranks-filter={{ ranks_filter }} \ + -m prime_rl.trainer.sft.train \ + @ $CONFIG_PATH \ + 2>&1 | sed -u 's/^\[[a-zA-Z]*[0-9]*\]://' | tee -a $OUTPUT_DIR/logs/trainer/node_${TRAIN_NODE_RANK}.log +' diff --git a/src/prime_rl/templates/single_node_rl.sbatch.j2 b/src/prime_rl/templates/single_node_rl.sbatch.j2 new file mode 100644 index 00000000000..9bec929d387 --- /dev/null +++ b/src/prime_rl/templates/single_node_rl.sbatch.j2 @@ -0,0 +1,36 @@ +#!/bin/bash + +#SBATCH --job-name={{ job_name }} +#SBATCH --nodes=1 +#SBATCH --ntasks-per-node=1 +#SBATCH --gres=gpu:{{ gpus_per_node }} +#SBATCH --exclusive +#SBATCH --partition={{ partition }} +{%- if account %} +#SBATCH --account={{ account }} +{%- endif %} +{%- if time %} +#SBATCH --time={{ time }} +{%- endif %} +{%- if nodelist %} +#SBATCH --nodelist={{ nodelist }} +{%- endif %} +{%- if exclude %} +#SBATCH --exclude={{ exclude }} +{%- endif %} +#SBATCH --output={{ output_dir }}/job_%j.log +#SBATCH --error={{ output_dir }}/job_%j.log + +export PROJECT_DIR={{ project_dir }} + +# Setup environment +cd $PROJECT_DIR +[ -f .env ] && source .env +source .venv/bin/activate +uv sync --all-extras + +{% if pre_run_command %} +# Pre-run command +{{ pre_run_command }} +{% endif %} +uv run rl @ {{ config_path }} diff --git a/src/prime_rl/templates/single_node_sft.sbatch.j2 b/src/prime_rl/templates/single_node_sft.sbatch.j2 new file mode 100644 index 00000000000..68694ed13e6 --- /dev/null +++ b/src/prime_rl/templates/single_node_sft.sbatch.j2 @@ -0,0 +1,36 @@ +#!/bin/bash + +#SBATCH --job-name={{ job_name }} +#SBATCH --nodes=1 +#SBATCH --ntasks-per-node=1 +#SBATCH --gres=gpu:{{ gpus_per_node }} +#SBATCH --exclusive +#SBATCH --partition={{ partition }} +{%- if account %} +#SBATCH --account={{ account }} +{%- endif %} +{%- if time %} +#SBATCH --time={{ time }} +{%- endif %} +{%- if nodelist %} +#SBATCH --nodelist={{ nodelist }} +{%- endif %} +{%- if exclude %} +#SBATCH --exclude={{ exclude }} +{%- endif %} +#SBATCH --output={{ output_dir }}/job_%j.log +#SBATCH --error={{ output_dir }}/job_%j.log + +export PROJECT_DIR={{ project_dir }} + +# Setup environment +cd $PROJECT_DIR +[ -f .env ] && source .env +source .venv/bin/activate +uv sync --all-extras + +{% if pre_run_command %} +# Pre-run command +{{ pre_run_command }} +{% endif %} +uv run sft @ {{ config_path }} diff --git a/src/prime_rl/trainer/batch.py b/src/prime_rl/trainer/batch.py index 9044eb03788..662df36a805 100644 --- a/src/prime_rl/trainer/batch.py +++ b/src/prime_rl/trainer/batch.py @@ -3,26 +3,27 @@ from prime_rl.transport.types import MicroBatch, TrainingSample -def prepare_sample( - training_example: TrainingSample, - seq_len: int, - temperature: float, -) -> MicroBatch: +def prepare_sample(training_example: TrainingSample, seq_len: int) -> MicroBatch: """ Prepare a problem for sequence packing training. Tokenize and prepare tensors. """ - # Prepare input_ids, loss_mask, position_ids, inference_logprobs, and advantages input_ids = training_example.prompt_ids + training_example.completion_ids loss_mask = training_example.prompt_mask + training_example.completion_mask - # Inference logprobs only cover completion tokens, so prepend zeros for prompt tokens inference_logprobs = [0.0] * len(training_example.prompt_ids) + training_example.completion_logprobs advantages = [training_example.advantage] * len(input_ids) position_ids = list(range(len(input_ids))) + mm_token_type_ids = training_example.mm_token_type_ids + + # Per-token temperatures: prompt tokens use first completion temp (masked out anyway) + # Default to 1.0 if completion is empty (e.g., model generated only tool calls with no text) + prompt_temp = training_example.completion_temperatures[0] if training_example.completion_temperatures else 1.0 + temperatures = [prompt_temp] * len(training_example.prompt_ids) + training_example.completion_temperatures # Teacher logprobs already cover the full sequence (prompt + completion), # computed via prefill in the orchestrator when a teacher model is configured teacher_logprobs = training_example.teacher_logprobs + routed_experts = training_example.routed_experts if len(input_ids) > seq_len: input_ids = input_ids[:seq_len] @@ -30,8 +31,13 @@ def prepare_sample( inference_logprobs = inference_logprobs[:seq_len] position_ids = position_ids[:seq_len] advantages = advantages[:seq_len] + temperatures = temperatures[:seq_len] if teacher_logprobs is not None: teacher_logprobs = teacher_logprobs[:seq_len] + if routed_experts is not None: + routed_experts = routed_experts[:seq_len] + if mm_token_type_ids is not None: + mm_token_type_ids = mm_token_type_ids[:seq_len] assert ( len(input_ids) @@ -39,11 +45,23 @@ def prepare_sample( == len(loss_mask) == len(position_ids) == len(inference_logprobs) + == len(temperatures) ), ( - f"input_ids: {len(input_ids)}, advantages: {len(advantages)}, loss_mask: {len(loss_mask)}, position_ids: {len(position_ids)}, inference_logprobs: {len(inference_logprobs)}" + f"input_ids: {len(input_ids)}, advantages: {len(advantages)}, loss_mask: {len(loss_mask)}, position_ids: {len(position_ids)}, inference_logprobs: {len(inference_logprobs)}, temperatures: {len(temperatures)}" ) if teacher_logprobs is not None: assert len(teacher_logprobs) == len(input_ids), f"teacher_logprobs: {len(teacher_logprobs)}" + + if routed_experts is not None: + assert len(routed_experts) == len(input_ids), ( + f"routed_experts: {len(routed_experts)}, input_ids: {len(input_ids)}" + ) + + if mm_token_type_ids is not None: + assert len(mm_token_type_ids) == len(input_ids), ( + f"mm_token_type_ids: {len(mm_token_type_ids)}, input_ids: {len(input_ids)}" + ) + return MicroBatch( input_ids=input_ids, advantages=advantages, @@ -51,35 +69,74 @@ def prepare_sample( position_ids=position_ids, inference_logprobs=inference_logprobs, teacher_logprobs=teacher_logprobs, - temperature=temperature, + temperatures=temperatures, + routed_experts=routed_experts, + mm_token_type_ids=mm_token_type_ids, + # Multimodal fields (Qwen3-VL) - passed through without modification + pixel_values=training_example.pixel_values, + pixel_values_shape=training_example.pixel_values_shape, + image_grid_thw=training_example.image_grid_thw, + sft_loss=training_example.sft_loss, ) +def _is_multimodal_sample(sample: MicroBatch) -> bool: + """Check if a sample contains multimodal data (images).""" + return sample.pixel_values is not None + + def packed_samples_into_micro_bs( samples: list[tuple[int, MicroBatch]], max_seq_len: int, num_loras: int ) -> list[MicroBatch]: """ Pack samples into micro_batch efficiently. We follow the First Fit Decreasing algorithm to pack the samples into bins and minimize potential padding while never truncating. + With per-token temperatures, samples can be packed together regardless of their temperature values. + + NOTE: Multimodal samples (with pixel_values) are NOT packed together as they have variable-sized + vision data that doesn't pack well. Each multimodal sample becomes its own micro batch. """ + # Sort by (lora_idx, -length) for packing efficiency samples.sort(key=lambda x: (x[0], -len(x[1].input_ids))) ## we create bins micro_batches: list[MicroBatch] = [] for idx, sample in samples: - # Try to find a bin that can fit this sequence + # Multimodal samples cannot be packed - each becomes its own micro batch + if _is_multimodal_sample(sample): + sample.lora_num_tokens = [0] * num_loras + sample.lora_num_tokens[idx] = len(sample.input_ids) + micro_batches.append(sample) + continue + + # Try to find a bin that can fit this sequence (only pack text-only samples) for bin_content in micro_batches: + # Don't pack into multimodal micro batches + if _is_multimodal_sample(bin_content): + continue # Check if sequence fits in this bin - if len(bin_content.input_ids) + len(sample.input_ids) <= max_seq_len: + if ( + len(bin_content.input_ids) + len(sample.input_ids) <= max_seq_len + and bin_content.sft_loss == sample.sft_loss + ): bin_content.input_ids.extend(sample.input_ids) bin_content.loss_mask.extend(sample.loss_mask) bin_content.advantages.extend(sample.advantages) bin_content.inference_logprobs.extend(sample.inference_logprobs) + bin_content.temperatures.extend(sample.temperatures) if sample.teacher_logprobs is not None: if bin_content.teacher_logprobs is None: bin_content.teacher_logprobs = [] bin_content.teacher_logprobs.extend(sample.teacher_logprobs) + if sample.routed_experts is not None: + if bin_content.routed_experts is None: + bin_content.routed_experts = [] + bin_content.routed_experts.extend(sample.routed_experts) + if sample.mm_token_type_ids is not None: + if bin_content.mm_token_type_ids is None: + bin_content.mm_token_type_ids = [] + bin_content.mm_token_type_ids.extend(sample.mm_token_type_ids) bin_content.position_ids.extend(sample.position_ids) bin_content.lora_num_tokens[idx] += len(sample.input_ids) break @@ -112,18 +169,38 @@ def pad_micro_batch(micro_batch: MicroBatch, pad_to_multiple_of: int) -> MicroBa micro_batch.loss_mask.extend([False] * padding_size) micro_batch.position_ids.extend(list(range(padding_size))) micro_batch.inference_logprobs.extend([0.0] * padding_size) + # Use temperature 1.0 for padding tokens (doesn't matter since loss_mask is False) + micro_batch.temperatures.extend([1.0] * padding_size) if micro_batch.teacher_logprobs is not None: micro_batch.teacher_logprobs.extend([0.0] * padding_size) micro_batch.lora_num_tokens[-1] += ( padding_size # We send padding to the last lora so that tokens have ascending lora idx ) + if micro_batch.mm_token_type_ids is not None: + micro_batch.mm_token_type_ids.extend([0] * padding_size) return micro_batch +def _make_dummy_batch(source: MicroBatch) -> MicroBatch: + """Create a zero-loss dummy batch from an existing batch, preserving its modality.""" + dummy = copy.deepcopy(source) + dummy.advantages = [0.0] * len(dummy.input_ids) + dummy.loss_mask = [False] * len(dummy.input_ids) + return dummy + + +def _pad_group_for_distribution(group: list[MicroBatch], num_train_workers: int) -> list[MicroBatch]: + """Pad a group of micro batches so its length is divisible by num_train_workers.""" + num_padding = -len(group) % num_train_workers + if num_padding > 0 and len(group) > 0: + dummy = _make_dummy_batch(group[0]) + group.extend([dummy] * num_padding) + return group + + def prepare_batch( rollouts: list[TrainingSample], - temperature: float, seq_len: int, num_train_workers: int, idxs: list[int], @@ -133,34 +210,34 @@ def prepare_batch( """ Prepare a batch of problems for each GPU. Each batch is a list of micro batches. Each micro batch is shape [1, seq_len], the number of samples is not fixed per micro batch. - """ - max_seq_len = seq_len - all_samples = [(idx, prepare_sample(rollout, max_seq_len, temperature)) for idx, rollout in zip(idxs, rollouts)] + FSDP requires all ranks to execute the same operations at each step. If one rank + processes a multimodal batch (triggering the vision encoder) while another processes + a text-only batch, the all-gather will hang. We separate micro batches by modality + and distribute them so that at each step index, all ranks see the same modality. + """ + all_samples = [(idx, prepare_sample(rollout, seq_len)) for idx, rollout in zip(idxs, rollouts)] - micro_batches = packed_samples_into_micro_bs(all_samples, max_seq_len, num_loras) + micro_batches = packed_samples_into_micro_bs(all_samples, seq_len, num_loras) micro_batches = [pad_micro_batch(micro_batch, pad_to_multiple_of) for micro_batch in micro_batches] - num_padding_batch = -len(micro_batches) % num_train_workers + # Separate by modality so each step index has uniform modality across all ranks + mm_batches = [b for b in micro_batches if _is_multimodal_sample(b)] + text_batches = [b for b in micro_batches if not _is_multimodal_sample(b)] - # because of fsdp we need to make sure that each data ran has the same number of micro batches otherwise training will hang. - # We create fake micro batches to fill the gap with real data but zero advantages, they would not contribute to the loss. - if num_train_workers > 1 and num_padding_batch > 0: - padded_batch = copy.deepcopy(micro_batches[0]) - padded_batch.advantages = [0.0] * len(padded_batch.input_ids) - padded_batch.loss_mask = [False] * len(padded_batch.input_ids) - micro_batches.extend([padded_batch for _ in range(num_padding_batch)]) + # Pad each group independently so its count is divisible by num_train_workers + mm_batches = _pad_group_for_distribution(mm_batches, num_train_workers) + text_batches = _pad_group_for_distribution(text_batches, num_train_workers) - assert len(micro_batches) % num_train_workers == 0, ( - "Number of micro batches is not divisible by number of data ranks" - ) + # Combine: all multimodal first, then all text-only. Since each group's length is + # divisible by num_train_workers, the modality boundary aligns with distribution rows. + ordered = mm_batches + text_batches + + assert len(ordered) % num_train_workers == 0, "Number of micro batches is not divisible by number of data ranks" - per_gpu_micro_batches = len(micro_batches) // num_train_workers - batches_per_gpu = [] - for _ in range(num_train_workers): - batches = [] - for _ in range(per_gpu_micro_batches): - batches.append(micro_batches.pop(0)) - batches_per_gpu.append(batches) + # Distribute in strided order so each step index has the same modality across ranks + batches_per_gpu: list[list[MicroBatch]] = [[] for _ in range(num_train_workers)] + for i, batch in enumerate(ordered): + batches_per_gpu[i % num_train_workers].append(batch) return batches_per_gpu diff --git a/src/prime_rl/trainer/ckpt.py b/src/prime_rl/trainer/ckpt.py index 6ced9c721a8..296554e0e89 100644 --- a/src/prime_rl/trainer/ckpt.py +++ b/src/prime_rl/trainer/ckpt.py @@ -1,3 +1,5 @@ +import bisect +import gc import shutil import time import warnings @@ -11,27 +13,35 @@ from torch.distributed.checkpoint.state_dict_loader import load as dcp_load from torch.distributed.checkpoint.state_dict_saver import save as dcp_save from torch.distributed.checkpoint.stateful import Stateful +from torch.distributed.tensor import DTensor from torch.nn import Module from torch.optim.lr_scheduler import LRScheduler from torch.optim.optimizer import Optimizer from torchdata.stateful_dataloader import StatefulDataLoader from transformers.tokenization_utils import PreTrainedTokenizer -from prime_rl.trainer.config import CheckpointConfig, LoRAConfig, WeightCheckpointConfig +from prime_rl.configs.trainer import CheckpointConfig, LoRAConfig, WeightCheckpointConfig from prime_rl.trainer.lora import has_lora_layers, save_lora_config from prime_rl.trainer.models import PreTrainedModelPrimeRL -from prime_rl.trainer.runs import Progress +from prime_rl.trainer.optim import CPUOffloadOptimizer +from prime_rl.trainer.runs import Progress, get_multi_run_manager from prime_rl.trainer.weights import ( gather_weights_on_master, - get_adapter_state_dict, save_state_dict, ) from prime_rl.trainer.world import get_world from prime_rl.utils.logger import get_logger -from prime_rl.utils.tensor_hashing import get_module_signature, get_optimizer_signature from prime_rl.utils.utils import get_all_ckpt_steps, get_ckpt_dir, get_step_path, get_weights_dir +def _try_rmtree(path: Path, logger) -> None: + """Remove a directory tree, logging and skipping on failure.""" + try: + shutil.rmtree(path) + except OSError as e: + logger.warning(f"Failed to remove {path}: {e}, skipping cleanup") + + class AppState(Stateful): """ A wrapper for checkpointing the trainer with sharded weights and optimizer @@ -53,9 +63,20 @@ def __init__( self.scheduler = scheduler self.progress = progress + def _get_base_optimizers(self) -> list[Optimizer]: + """Extract base optimizers from wrappers like CPUOffloadOptimizer.""" + return [opt.base_optimizer if isinstance(opt, CPUOffloadOptimizer) else opt for opt in self.optimizers] + def state_dict(self) -> dict[str, Any]: + # Move CPU-offloaded states to GPU before checkpointing + for opt in self.optimizers: + if isinstance(opt, CPUOffloadOptimizer) and opt._initialized: + opt._move_states("cuda") + torch.cuda.synchronize() + # Automatically manages FSDP FQN's, as well as sets the default state dict type to FSDP.SHARDED_STATE_DICT - model_state_dict, optimizer_state_dict = get_state_dict(self.model, self.optimizers) + base_optimizers = self._get_base_optimizers() + model_state_dict, optimizer_state_dict = get_state_dict(self.model, base_optimizers) state_dict = { "model": model_state_dict, "optimizers": optimizer_state_dict, @@ -66,24 +87,51 @@ def state_dict(self) -> dict[str, Any]: if self.progress is not None: progress_state_dict = asdict(self.progress) state_dict["progress"] = progress_state_dict + + # Move states back to CPU + for opt in self.optimizers: + if isinstance(opt, CPUOffloadOptimizer) and opt._initialized: + opt._move_states("cpu") + return state_dict def load_state_dict(self, state_dict: dict[str, Any]): + base_optimizers = self._get_base_optimizers() set_state_dict( - self.model, self.optimizers, model_state_dict=state_dict["model"], optim_state_dict=state_dict["optimizers"] + self.model, base_optimizers, model_state_dict=state_dict["model"], optim_state_dict=state_dict["optimizers"] ) + + # Re-initialize CPU offload wrappers after loading + has_cpu_offload = False + for opt in self.optimizers: + if isinstance(opt, CPUOffloadOptimizer): + opt._move_states("cpu") + opt._initialized = True + has_cpu_offload = True + if self.scheduler is not None: self.scheduler.load_state_dict(state_dict["scheduler"]) if self.progress is not None: for key, value in state_dict["progress"].items(): setattr(self.progress, key, value) + # Reclaim GPU memory freed by moving optimizer states to CPU. + # After set_state_dict + _move_states("cpu"), the optimizer states live on CPU, + # but the state_dict (owned by dcp_load) still holds references to stale GPU + # optimizer tensors. Clearing them and flushing the CUDA cache prevents OOM on + # the first training step. + if has_cpu_offload: + state_dict.clear() # drop stale GPU tensor references from dcp_load + gc.collect() # break any circular references so tensors are freed + torch.cuda.empty_cache() # return freed GPU memory to CUDA + class CheckpointManager: """Utility class to save and load trainer checkpoints to resume SFT and RL training.""" def __init__(self, output_dir: Path, config: CheckpointConfig): self.config = config + self.skip_optimizer = config.skip_optimizer self.ckpt_dir = get_ckpt_dir(output_dir) self.logger = get_logger() self.world = get_world() @@ -145,7 +193,7 @@ def load_from_path( start_time = time.perf_counter() # Load sharded state - app_state = AppState(model, optimizers, scheduler, progress) + app_state = AppState(model, optimizers if not self.skip_optimizer else [], scheduler, progress) state_dict = {"app": app_state} dcp_load(state_dict=state_dict, checkpoint_id=path) @@ -161,7 +209,7 @@ def load_from_path( raise RuntimeError( f"Couldn't fallback to using the master rank's dataloader checkpoint, because dataloder checkpoint was not found at path {dataloader_path}. Cannot resume training." ) - dataloader.load_state_dict(torch.load(dataloader_path)) + dataloader.load_state_dict(torch.load(dataloader_path, weights_only=False)) self.logger.debug(f"Training checkpoint loaded in {time.perf_counter() - start_time:.2f} seconds") @@ -179,9 +227,6 @@ def load( if not ckpt_path.exists(): raise FileNotFoundError(f"Checkpoint not found at {ckpt_path}") self.load_from_path(ckpt_path, model, optimizers, scheduler, progress, dataloader) - self.logger.debug( - f"Signatures after loading training checkpoint: model={get_module_signature(model, compress=True)}, optimizers={', '.join(get_optimizer_signature(optimizer, compress=True) for optimizer in optimizers)}" - ) def save( self, @@ -195,12 +240,9 @@ def save( """Save the full checkpoint state for a specified step.""" ckpt_path = self.get_ckpt_path(step) ckpt_path.parent.mkdir(parents=True, exist_ok=True) - self.logger.debug( - f"Signatures before saving training checkpoint: model={get_module_signature(model, compress=True)}, optimizers={', '.join(get_optimizer_signature(optimizer, compress=True) for optimizer in optimizers)}" - ) self.save_to_path(ckpt_path, model, optimizers, scheduler, progress, dataloader) - self.ckpt_steps.append(step) + bisect.insort(self.ckpt_steps, step) def maybe_clean(self) -> None: """Deletes past checkpoints based on keep_last and keep_interval policies. No-op if both are None.""" @@ -223,14 +265,15 @@ def maybe_clean(self) -> None: if step % self.config.keep_interval == 0: steps_to_keep.add(step) - # Delete steps not in steps_to_keep + # Delete steps not in steps_to_keep (only master rank deletes to avoid race condition) ckpt_steps_to_delete = [step for step in self.ckpt_steps if step not in steps_to_keep] - for ckpt_step in ckpt_steps_to_delete: - trainer_ckpt_path = self.get_ckpt_path(ckpt_step) - ckpt_path = trainer_ckpt_path.parent - if ckpt_path.exists(): - self.logger.debug(f"Removing past checkpoint for step {ckpt_step} ({ckpt_path})") - shutil.rmtree(ckpt_path) + if self.world.is_master: + for ckpt_step in ckpt_steps_to_delete: + trainer_ckpt_path = self.get_ckpt_path(ckpt_step) + ckpt_path = trainer_ckpt_path.parent + if ckpt_path.exists(): + self.logger.debug(f"Removing past checkpoint for step {ckpt_step} ({ckpt_path})") + _try_rmtree(ckpt_path, self.logger) # Update checkpoint steps self.ckpt_steps = [step for step in self.ckpt_steps if step in steps_to_keep] @@ -269,6 +312,25 @@ def get_step_path(self, step: int) -> Path: """Get the path to write the weight checkpoint for a given step.""" return get_step_path(self.weights_dir, step) + def mark_stable(self, step: int) -> None: + """Write STABLE file to indicate weight checkpoint is complete.""" + if self.world.is_master: + step_path = self.get_step_path(step) + (step_path / "STABLE").touch() + + def get_run_adapter_state_dict(self) -> dict[str, Tensor]: + lora_state_dict = { + f"base_model.model.{key}": (value.full_tensor() if isinstance(value, DTensor) else value).to( + "cpu", non_blocking=False + ) + for key, value in get_multi_run_manager().get_state_dict_for_run(0).items() + } + + if not lora_state_dict: + raise ValueError("The LoRA state dict is empty. Something went wrong.") + + return lora_state_dict + def save_to_path( self, path: Path, @@ -278,31 +340,47 @@ def save_to_path( tokenizer: PreTrainedTokenizer, ): """Save HF-compatible weight checkpoint to a given path.""" - path.mkdir(parents=True, exist_ok=True) - start_time = time.perf_counter() + if self.world.is_master: + path.mkdir(parents=True, exist_ok=True) + start_time = time.perf_counter() - self.logger.debug(f"Saving weight checkpoint to {path}") - # Suppress torch.distributed warnings during checkpoint saving - with warnings.catch_warnings(): - warnings.filterwarnings("ignore", category=FutureWarning, module="torch.distributed") - warnings.filterwarnings("ignore", category=UserWarning, module="torch.distributed.*") - - # Save weights - save_state_dict(state_dict, path, self.config.save_format, self.config.save_sharded) - - # Save model config, generation arguments and tokenizer - model.config.save_pretrained(path) - if model.generation_config: - model.generation_config.save_pretrained(path) - tokenizer.save_pretrained(path) - - if self.config.save_adapter_separately and lora_state_dict is not None: - adapter_path = path / "lora_adapters" - adapter_path.mkdir(parents=True, exist_ok=True) - torch.save(lora_state_dict, adapter_path / "adapter_model.bin") - if self.lora_config: - save_lora_config(self.lora_config, model, adapter_path) # Pass model - self.logger.debug(f"Saved weight checkpoint to {path} in {time.perf_counter() - start_time:.2f} seconds") + self.logger.debug(f"Saving weight checkpoint to {path}") + # Suppress torch.distributed warnings during checkpoint saving + with warnings.catch_warnings(): + warnings.filterwarnings("ignore", category=FutureWarning, module="torch.distributed") + warnings.filterwarnings("ignore", category=UserWarning, module="torch.distributed.*") + + # Save weights + save_state_dict(state_dict, path, self.config.save_format, self.config.save_sharded) + + # Save model config, generation arguments and tokenizer + model.config.save_pretrained(path) + if model.generation_config: + # training sets use_cache=False which can conflict with + # cache_implementation — save with use_cache=True without + # mutating the model's config + from copy import deepcopy + + gen_config = deepcopy(model.generation_config) + gen_config.use_cache = True + gen_config.save_pretrained(path) + tokenizer.save_pretrained(path) + + if lora_state_dict is not None: + adapter_path = path / "lora_adapters" + adapter_path.mkdir(parents=True, exist_ok=True) + save_state_dict( + lora_state_dict, adapter_path, self.config.save_format, save_sharded=False, adapter=True + ) + if self.lora_config: + save_lora_config( + model, + adapter_path, + rank=self.lora_config.rank, + alpha=self.lora_config.alpha, + dropout=self.lora_config.dropout, + ) + self.logger.debug(f"Saved weight checkpoint to {path} in {time.perf_counter() - start_time:.2f} seconds") def save( self, @@ -320,15 +398,20 @@ def save( state_dict = gather_weights_on_master(model, self.world.is_master, dtype=torch.bfloat16) self.logger.debug(f"Gathered weights on master rank in {time.perf_counter() - start_time:.2f} seconds") - if has_lora_layers(model): - self.logger.debug("Getting LoRA state dict on master rank for weight checkpoint") + # Remove tied weight keys to match original model format + if getattr(model.config, "tie_word_embeddings", False): + for key in getattr(model, "_tied_weights_keys", []): + state_dict.pop(key, None) + + if has_lora_layers(model) and self.config.save_adapter_separately: + self.logger.debug("Getting run adapter state dict for weight checkpoint") start_time = time.perf_counter() - lora_state_dict = get_adapter_state_dict(model, self.world.is_master) - self.logger.debug(f"Got LoRA state dict on master rank in {time.perf_counter() - start_time:.2f} seconds") + lora_state_dict = self.get_run_adapter_state_dict() + self.logger.debug(f"Got run adapter state dict in {time.perf_counter() - start_time:.2f} seconds") else: lora_state_dict = None - # Convert PrimeRL format to HF format if needed + # Convert to HF hub format if needed if isinstance(model, PreTrainedModelPrimeRL) and model.is_prime_state_dict(state_dict): self.logger.debug("Converting PrimeRL format to HF format for weight checkpoint") start_time = time.perf_counter() @@ -336,13 +419,19 @@ def save( self.logger.debug( f"Converted PrimeRL format to HF format in {time.perf_counter() - start_time:.2f} seconds" ) + else: + # For regular transformers models, revert internal format to original HF hub format + from transformers.core_model_loading import revert_weight_conversion + + self.logger.debug("Reverting transformers internal format to HF hub format for weight checkpoint") + start_time = time.perf_counter() + state_dict = revert_weight_conversion(model, state_dict) + self.logger.debug(f"Reverted to HF hub format in {time.perf_counter() - start_time:.2f} seconds") # Save weight checkpoint on master rank - if self.world.is_master: - self.save_to_path(step_path, state_dict, lora_state_dict, model, tokenizer) - # Write STABLE file to indicate checkpoint is complete (for eval to safely read) - (step_path / "STABLE").touch() - self.ckpt_steps.append(step) + self.save_to_path(step_path, state_dict, lora_state_dict, model, tokenizer) + self.mark_stable(step) + bisect.insort(self.ckpt_steps, step) def maybe_clean(self) -> None: """Deletes past checkpoints based on keep_last and keep_interval policies. No-op if both are None.""" @@ -365,13 +454,14 @@ def maybe_clean(self) -> None: if step % self.keep_interval == 0: steps_to_keep.add(step) - # Delete steps not in steps_to_keep + # Delete steps not in steps_to_keep (only master rank deletes to avoid race condition) ckpt_steps_to_delete = [step for step in self.ckpt_steps if step not in steps_to_keep] - for ckpt_step in ckpt_steps_to_delete: - ckpt_path = self.get_step_path(ckpt_step) - if ckpt_path.exists(): - self.logger.debug(f"Removing past checkpoint for step {ckpt_step} ({ckpt_path})") - shutil.rmtree(ckpt_path) + if self.world.is_master: + for ckpt_step in ckpt_steps_to_delete: + ckpt_path = self.get_step_path(ckpt_step) + if ckpt_path.exists(): + self.logger.debug(f"Removing past checkpoint for step {ckpt_step} ({ckpt_path})") + _try_rmtree(ckpt_path, self.logger) # Update checkpoint steps self.ckpt_steps = [step for step in self.ckpt_steps if step in steps_to_keep] @@ -382,10 +472,11 @@ def setup_ckpt_managers( ) -> tuple[CheckpointManager | None, WeightCheckpointManager | None]: if ckpt_config is None: return None, None - ckpt_manager = CheckpointManager(output_dir, ckpt_config) - if ckpt_config.weights: + ckpt_output_dir = ckpt_config.output_dir or output_dir + ckpt_manager = CheckpointManager(ckpt_output_dir, ckpt_config) + if ckpt_config.weights and not ckpt_config.skip_gather_master_weights: weight_ckpt_manager = WeightCheckpointManager( - output_dir, + ckpt_output_dir, ckpt_config.weights, lora_config=lora_config, keep_last=ckpt_config.keep_last, diff --git a/src/prime_rl/trainer/config.py b/src/prime_rl/trainer/config.py deleted file mode 100644 index 7a65c4e71b2..00000000000 --- a/src/prime_rl/trainer/config.py +++ /dev/null @@ -1,470 +0,0 @@ -from pathlib import Path -from typing import Annotated, Literal, TypeAlias - -from pydantic import BaseModel, Field, model_validator - -from prime_rl.utils.pydantic_config import BaseConfig - -AttnImplementation: TypeAlias = Literal["sdpa", "flash_attention_2", "flash_attention_3"] - -MOE_MODEL_MAPS = { - "Qwen/Qwen3-30B-A3B": "Jackmin108/Qwen3-30B-A3B-Fast", - "moonshotai/Moonlight-16B-A3B-Instruct": "Jackmin108/Moonlight-16B-A3B-Instruct-Fast", -} - - -class ActivationCheckpointConfig(BaseConfig): - """Configures activation checkpointing.""" - - freq: Annotated[ - int, - Field( - ge=1, - description="Applies activation checkpointing to every `freq` layers. Defaults to 1, which will is full activation checkpointing.", - ), - ] = 1 - - -class ActivationOffloadingConfig(BaseConfig): - """Configures the activation offloading.""" - - pin_memory: Annotated[bool, Field(description="Whether to pin the offloaded activations to CPU memory.")] = True - - max_inflight_activations: Annotated[ - int, - Field( - ge=1, - description="The maximum number of activations to keep in while offloading further. (More activations means smoother overlap, but more gpu memory usage)", - ), - ] = 5 - - -class CompileConfig(BaseConfig): - """Configures model compilation.""" - - fullgraph: Annotated[ - bool, - Field(description="Whether to compile the transformer blocks with fullgraph."), - ] = False - - -class BenchConfig(BaseConfig): - """Configures benchmark mode.""" - - output_json: Annotated[ - Path | None, - Field(description="Path to write benchmark results as JSON. If not set, only prints to console."), - ] = None - - -class DebugModelConfig(BaseConfig): - """Debugging feature around model and distributed training.""" - - num_layers: Annotated[ - int | None, - Field(description="The number of layers in the model."), - ] = None - - random_init: Annotated[ - bool, - Field( - description="Whether to random initialize the model.", - ), - ] = False - - -class LoRAConfig(BaseConfig): - """Configuration for LoRA (Low-Rank Adaptation).""" - - rank: Annotated[ - int, - Field( - ge=1, - description="Rank of the low-rank decomposition matrices.", - ), - ] = 16 - - alpha: Annotated[ - float, - Field( - ge=0, - description="LoRA scaling parameter.", - ), - ] = 32.0 - - dropout: Annotated[ - float, - Field( - ge=0, - le=1, - description="LoRA dropout rate.", - ), - ] = 0.0 - - target_modules: Annotated[ - list[str], - Field( - description="Module names or regex patterns for modules to apply LoRA to. Simple names (e.g., 'q_proj') match any component in the module path. Regex patterns match anywhere in the name.", - ), - ] = [ - "q_proj", - "k_proj", - "v_proj", - "o_proj", - "gate_proj", - "up_proj", - "down_proj", - "experts", - ] - - modules_to_save: Annotated[ - list[str], - Field( - description="Module names or regex patterns for modules to keep fully trainable (not freeze). Simple names match any component in the module path. Regex patterns match anywhere in the name.", - ), - ] = [] - - -class ModelConfig(BaseConfig): - """Configures the model for training.""" - - name: Annotated[ - str, - Field( - description="Name or path of the HF model to use.", - ), - ] = "Qwen/Qwen3-0.6B" - - seq_len: Annotated[int, Field(description="The sequence length to use for the model.")] = 2048 - - attn: Annotated[AttnImplementation, Field(description="The attention implementation to use.")] = "flash_attention_2" - - compile: Annotated[ - CompileConfig | None, - Field( - description="Whether to compile the model using `torch.compile`.", - ), - ] = None - - ac: Annotated[ - ActivationCheckpointConfig | None, - Field( - description="Whether to apply activation checkpointing to the model. If None, will not apply activation checkpointing.", - ), - ] = None - - ac_offloading: Annotated[ - ActivationOffloadingConfig | None, - Field( - description="Whether to apply activation offloading to the model. If None, will not apply activation offloading.", - ), - ] = None - - fsdp_cpu_offload: Annotated[ - bool, - Field( - description="Whether to enable FSDP CPU offloading for parameters, gradients, and optimizer states. When enabled, uses pinned memory for efficient CPU-GPU transfers.", - ), - ] = False - - reshard_after_forward: Annotated[ - bool, Field(description="Whether to reshard the model after each forward pass.") - ] = True - - trust_remote_code: Annotated[ - bool, - Field( - description="Whether to trust remote code for model and tokenizer initialization.", - ), - ] = False - - dp_replicate: Annotated[ - int, - Field( - description="The data parallel dim where model weights are replicated.", - ), - ] = 1 - - ep: Annotated[ - int, - Field( - description="The expert parallelism to use if the model has MoE layers. If 1, then no EP will be used.", - ), - ] = 1 - - tp: Annotated[ - int, - Field( - description="The tensor parallelism size to use. If 1, then no TP will be used.", - ), - ] = 1 - - cp: Annotated[ - int, - Field( - description="The context parallelism size to use. If 1, then no CP will be used.", - ), - ] = 1 - - impl: Annotated[ - Literal["hf", "liger_kernel", "custom", "auto"], - Field( - description="Model implementation to use. 'auto' (default) selects 'custom' if supported by the model, otherwise 'hf'.", - ), - ] = "auto" - - optimization_dtype: Annotated[ - Literal["bfloat16", "float32"], - Field( - description="The dtype to use for the model optimization.", - ), - ] = "float32" - - reduce_dtype: Annotated[ - Literal["bfloat16", "float32"], - Field( - description="The dtype to use for the model reduce.", - ), - ] = "float32" - - moe_use_grouped_mm: Annotated[ - bool, - Field( - description="Whether to use grouped mm for the MoE layers. Require compute capability >= 9.0", - ), - ] = True - - lora: Annotated[ - LoRAConfig | None, - Field( - description="Whether to apply LoRA to the model. If None, will not apply LoRA.", - ), - ] = None - - debug: Annotated[ - DebugModelConfig, - Field( - description="Debugging feature around model and distributed training.", - ), - ] = DebugModelConfig() - - fused_lm_head_chunk_size: Annotated[ - int | None, - Field( - description="The chunk size to use for the fused LM head. If None, will not use chunking. RL training auto-sets this to 2048 if not specified (except when impl='liger_kernel').", - ), - ] = None - - @model_validator(mode="after") - def _map_model_name_for_moe(self): - """Map model name if it exists in MOE_MODEL_MAPS.""" - if self.name in MOE_MODEL_MAPS: - self.name = MOE_MODEL_MAPS[self.name] - return self - - @model_validator(mode="after") - def trust_remote_code_only_with_hf(self): - """Trust remote code only if the model is from HF.""" - if self.trust_remote_code: - if self.impl not in ("hf", "auto"): - raise ValueError("Trust remote code is only supported with the HF implementation or auto mode.") - return self - - @model_validator(mode="after") - def cp_only_with_flash_attn(self): - if self.cp > 1 and self.attn not in ["flash_attention_2", "flash_attention_3"]: - raise ValueError("CP is only supported with flash attention 2 or flash attention 3") - return self - - @model_validator(mode="after") - def ac_offloading_requires_ac(self): - """Automatically enable activation checkpointing when activation offloading is enabled.""" - if self.ac_offloading is not None and self.ac is None: - self.ac = ActivationCheckpointConfig() - return self - - @model_validator(mode="after") - def fused_lm_head_chunk_size_is_valid(self): - if self.fused_lm_head_chunk_size is not None: - low = 512 - if self.fused_lm_head_chunk_size < low: - raise ValueError(f"Fused LM head chunk size must be greater than {low}") - - return self - - -class TokenizerConfig(BaseConfig): - """Configuration for the tokenizer.""" - - name: Annotated[ - str | None, - Field(description="The name or path of the tokenizer to use. If None, will use the model's default tokenizer."), - ] = None - - trust_remote_code: Annotated[ - bool | None, - Field( - description="Whether to trust remote code for tokenizer initialization. If None, will use the model's default trust remote code setting.", - ), - ] = None - - chat_template: Annotated[ - str | None, - Field( - description="The chat template to use for the tokenizer. If None, will use the tokenizer's default chat template." - ), - ] = None - - -class ConstantSchedulerConfig(BaseModel): - """Configuration for constant learning rate scheduler.""" - - type: Literal["constant"] = "constant" - - -class LinearSchedulerConfig(BaseModel): - """Configuration for linear learning rate scheduler.""" - - type: Literal["linear"] = "linear" - - warmup_steps: Annotated[int, Field(ge=0, description="Number of warmup steps for the learning rate scheduler.")] = ( - 10 - ) - - decay_steps: Annotated[ - int, - Field( - ge=0, - description="Number of steps to decay the learning rate during the final portion of training.", - ), - ] = 10 - - min_lr: Annotated[float, Field(ge=0, description="Minimum learning rate to converge to.")] = 0.0 - - -class CosineSchedulerConfig(BaseModel): - """Configuration for cosine learning rate scheduler.""" - - type: Literal["cosine"] = "cosine" - - warmup_steps: Annotated[int, Field(ge=0, description="Number of warmup steps for the learning rate scheduler.")] = ( - 10 - ) - - min_lr: Annotated[float, Field(ge=0, description="Minimum learning rate to converge to.")] = 0.0 - - -SchedulerConfigType: TypeAlias = ConstantSchedulerConfig | LinearSchedulerConfig | CosineSchedulerConfig - - -class BaseOptimizerConfig(BaseModel): - lr: Annotated[float, Field(ge=0)] = 1e-6 - weight_decay: Annotated[float, Field(ge=0)] = 0.01 - max_norm: Annotated[float, Field(ge=0, description="Maximum gradient norm to clip.")] = 1.0 - - -class SGDConfig(BaseOptimizerConfig): - type: Literal["sgd"] = "sgd" - nesterov: bool = True - momentum: float = 0.9 - - -class AdamWConfig(BaseOptimizerConfig): - type: Literal["adamw"] = "adamw" - - betas1: Annotated[float, Field(ge=0)] = 0.9 - betas2: Annotated[float, Field(ge=0)] = 0.999 - - -class MuonConfig(BaseOptimizerConfig): - type: Literal["muon"] = "muon" - - betas1: Annotated[float, Field(ge=0)] = 0.9 - betas2: Annotated[float, Field(ge=0)] = 0.999 - - -OptimizerConfigType: TypeAlias = SGDConfig | AdamWConfig | MuonConfig - - -class WeightCheckpointConfig(BaseConfig): - """Configures saving HF-compatible weight checkpoints.""" - - save_sharded: Annotated[ - bool, - Field( - description="Whether to save the weight checkpoint in sharded format.", - ), - ] = True - - save_format: Annotated[ - Literal["safetensors", "torch"], - Field( - description="The format to save the weight checkpoint in.", - ), - ] = "safetensors" - - save_adapter_separately: Annotated[ - bool, - Field( - description="Whether to save LoRA adapters separately before merging into full model weights.", - ), - ] = False - - -class CheckpointConfig(BaseConfig): - """Configures checkpointing the full model, optimizer and training state for resuming training.""" - - interval: Annotated[ - int | None, - Field( - ge=1, - description="Interval at which to save the training checkpoint. If None, will only checkpoint at the end of training.", - ), - ] = None - - weights: WeightCheckpointConfig | None = WeightCheckpointConfig() - - resume_step: Annotated[ - int | None, - Field( - ge=-1, - description="Step to resume training from. If None, will start from scratch. If -1, will restart from latest checkpoint available.", - ), - ] = None - - keep_last: Annotated[ - int | None, - Field( - ge=1, - description="Keep at most this many recent step checkpoints on disk. If None, never clean old checkpoints based on recency.", - ), - ] = None - - keep_interval: Annotated[ - int | None, - Field( - ge=1, - description="Keep checkpoints at every N steps permanently (e.g., keep_interval=100 keeps step 100, 200, ...). If None, no interval-based keeping.", - ), - ] = None - - skip_progress: Annotated[ - bool, - Field( - description="Whether to skip loading the progress from checkpoint.", - ), - ] = False - - skip_scheduler: Annotated[ - bool, - Field( - description="Whether to skip loading the scheduler from checkpoint.", - ), - ] = False - - skip_dataloader: Annotated[ - bool, - Field( - description="Whether to skip loading the dataloader from checkpoint.", - ), - ] = False diff --git a/src/prime_rl/trainer/distributed/__init__.py b/src/prime_rl/trainer/distributed/__init__.py new file mode 100644 index 00000000000..896372e5819 --- /dev/null +++ b/src/prime_rl/trainer/distributed/__init__.py @@ -0,0 +1,3 @@ +from prime_rl.trainer.distributed.expert_parallel import DeepEPExpertParallel, get_ep_group + +__all__ = ["DeepEPExpertParallel", "get_ep_group"] diff --git a/src/prime_rl/trainer/distributed/deepep.py b/src/prime_rl/trainer/distributed/deepep.py new file mode 100644 index 00000000000..c3072480616 --- /dev/null +++ b/src/prime_rl/trainer/distributed/deepep.py @@ -0,0 +1,344 @@ +from dataclasses import dataclass + +import torch +from deep_ep import Buffer +from deep_ep.utils import EventHandle, EventOverlap +from torch.distributed import ProcessGroup + +_buffer: Buffer | None = None +_handle_cache: dict[int, object] = {} +_pending_dispatch_events: dict[int, EventOverlap] = {} +_handle_counter = 0 +_pending_combine_event: EventOverlap | None = None +_deepep_cuda_ops_registered = False +_deepep_cuda_lib: torch.library.Library | None = None + + +def _get_next_handle_id() -> torch.Tensor: + global _handle_counter + _handle_counter += 1 + return torch.tensor([_handle_counter], dtype=torch.int64, device="cpu") + + +def _new_event_overlap() -> EventOverlap: + return EventOverlap(EventHandle()) + + +def register_deepep_cuda_ops() -> None: + global _deepep_cuda_lib, _deepep_cuda_ops_registered + if _deepep_cuda_ops_registered: + return + + # Keep the Library alive so PyTorch does not deregister the custom ops. + _deepep_cuda_lib = torch.library.Library("deepep", "DEF") + _deepep_cuda_lib.define( + "dispatch(Tensor x, Tensor topk_idx, Tensor topk_weights, " + "Tensor num_tokens_per_rank, Tensor num_tokens_per_rdma_rank, " + "Tensor is_token_in_rank, Tensor num_tokens_per_expert) " + "-> (Tensor, Tensor, Tensor, Tensor, Tensor)" + ) + + torch.library.impl(_deepep_cuda_lib, "dispatch", "CUDA")(_dispatch_op_impl) + + torch.library.register_autograd("deepep::dispatch", _dispatch_backward, setup_context=_dispatch_setup_context) + + _deepep_cuda_ops_registered = True + + +def _dispatch_op_impl( + x: torch.Tensor, + topk_idx: torch.Tensor, + topk_weights: torch.Tensor, + num_tokens_per_rank: torch.Tensor, + num_tokens_per_rdma_rank: torch.Tensor, + is_token_in_rank: torch.Tensor, + num_tokens_per_expert: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + assert _buffer is not None, "DeepEP buffer must be initialized before dispatch." + + previous_event = _new_event_overlap() + recv_x, recv_indices, recv_scores, recv_num_tokens_per_expert_list, handle, after_event = _buffer.dispatch( + x=x, + topk_idx=topk_idx, + topk_weights=topk_weights.to(torch.float32), + num_tokens_per_rank=num_tokens_per_rank, + num_tokens_per_rdma_rank=num_tokens_per_rdma_rank, + is_token_in_rank=is_token_in_rank, + num_tokens_per_expert=num_tokens_per_expert, + previous_event=previous_event, + async_finish=True, + allocate_on_comm_stream=True, + ) + handle_id = _get_next_handle_id() + _handle_cache[handle_id.item()] = handle + _pending_dispatch_events[handle_id.item()] = after_event + recv_num_tokens_per_expert = torch.tensor(recv_num_tokens_per_expert_list, dtype=torch.int32, device="cpu") + return recv_x, recv_indices, recv_scores, recv_num_tokens_per_expert, handle_id + + +def _dispatch_setup_context(ctx, inputs, output) -> None: + x, *_ = inputs + *_, handle_id = output + ctx.input_dtype = x.dtype + ctx.saved_handle = _handle_cache.get(handle_id.item()) + + +def _dispatch_backward( + ctx, + grad_recv_x, + grad_recv_indices, + grad_recv_scores, + grad_recv_num_tokens_per_expert, + grad_handle_id, +): + if grad_recv_x is None: + return None, None, None, None, None, None, None + + handle = ctx.saved_handle + assert handle is not None + + previous_event = _new_event_overlap() + grad_x, grad_scores, after_event = _buffer.combine( + x=grad_recv_x, + handle=handle, + topk_weights=grad_recv_scores.float() if grad_recv_scores is not None else None, + previous_event=previous_event, + async_finish=True, + allocate_on_comm_stream=True, + ) + after_event.current_stream_wait() + + grad_x = grad_x.to(ctx.input_dtype) + grad_topk_weights = grad_scores.to(ctx.input_dtype) if grad_scores is not None else None + return grad_x, None, grad_topk_weights, None, None, None, None + + +class _DeepEPCombine(torch.autograd.Function): + @staticmethod + def forward(ctx, x: torch.Tensor, handle_id: torch.Tensor) -> torch.Tensor: + global _pending_combine_event + + assert _buffer is not None, "DeepEP buffer must be initialized before combine." + handle = _handle_cache.pop(handle_id.item(), None) + assert handle is not None, f"Handle not found for handle_id={handle_id.item()}" + + previous_event = _new_event_overlap() + combined, _, after_event = _buffer.combine( + x=x, + handle=handle, + previous_event=previous_event, + async_finish=True, + allocate_on_comm_stream=True, + ) + _pending_combine_event = after_event + ctx.handle = handle + return combined + + @staticmethod + def backward(ctx, grad_combined: torch.Tensor) -> tuple[torch.Tensor, None]: + handle = ctx.handle + assert handle is not None, "Handle not found in DeepEP combine backward." + + previous_event = _new_event_overlap() + grad_x, _, _, _, _, after_event = _buffer.dispatch( + x=grad_combined, + topk_idx=None, + topk_weights=None, + num_tokens_per_rank=None, + num_tokens_per_rdma_rank=None, + is_token_in_rank=None, + num_tokens_per_expert=None, + handle=handle, + previous_event=previous_event, + async_finish=True, + allocate_on_comm_stream=True, + ) + after_event.current_stream_wait() + return grad_x, None + + +@torch.compiler.disable() +def sync_combine() -> None: + global _pending_combine_event + + if _pending_combine_event is not None: + _pending_combine_event.current_stream_wait() + _pending_combine_event = None + + +@torch.compiler.disable() +def _sync_dispatch(handle_id: torch.Tensor | int) -> None: + handle_key = handle_id if isinstance(handle_id, int) else handle_id.item() + pending_event = _pending_dispatch_events.pop(handle_key, None) + if pending_event is not None: + pending_event.current_stream_wait() + + +def configure_num_sms(num_sms: int) -> None: + """Set the number of SMs for DeepEP intranode dispatch/combine kernels. + + Must be called before the first dispatch/combine. Also determines + internode RDMA channel count (num_channels = num_sms / 2). + """ + Buffer.set_num_sms(num_sms) + + +def get_hidden_bytes(x: torch.Tensor) -> int: + return x.size(1) * max(x.element_size(), 2) + + +def get_buffer(group: ProcessGroup, hidden_bytes: int) -> Buffer: + global _buffer + + num_nvl_bytes, num_rdma_bytes = 0, 0 + for config in (Buffer.get_dispatch_config(group.size()), Buffer.get_combine_config(group.size())): + num_nvl_bytes = max(config.get_nvl_buffer_size_hint(hidden_bytes, group.size()), num_nvl_bytes) + num_rdma_bytes = max(config.get_rdma_buffer_size_hint(hidden_bytes, group.size()), num_rdma_bytes) + + if ( + _buffer is None + or _buffer.group != group + or _buffer.num_nvl_bytes < num_nvl_bytes + or _buffer.num_rdma_bytes < num_rdma_bytes + ): + _buffer = Buffer(group, num_nvl_bytes, num_rdma_bytes) + + return _buffer + + +def _permute_tokens( + hidden_states: torch.Tensor, + dispatched_indices: torch.Tensor, + dispatched_scores: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + mask = dispatched_indices != -1 + valid_expert_ids = dispatched_indices[mask] + valid_scores = dispatched_scores[mask] + + sort_order = torch.argsort(valid_expert_ids, stable=True) + permuted_indices = torch.arange(len(hidden_states), device=hidden_states.device).repeat_interleave(mask.sum(dim=1))[ + sort_order + ] + permuted_hidden_states = hidden_states.index_select(0, permuted_indices) + permuted_scores = valid_scores[sort_order] + return permuted_hidden_states, permuted_scores, permuted_indices + + +def _unpermute_tokens( + permuted_hidden_states: torch.Tensor, + permuted_indices: torch.Tensor, + num_tokens: int, +) -> torch.Tensor: + hidden_dim = permuted_hidden_states.shape[1] + output_hidden_states = permuted_hidden_states.new_zeros((num_tokens, hidden_dim)) + output_hidden_states.scatter_add_(0, permuted_indices.unsqueeze(1).expand(-1, hidden_dim), permuted_hidden_states) + return output_hidden_states + + +@dataclass +class _DispatchState: + handle_id: torch.Tensor + permuted_indices: torch.Tensor + num_recv_tokens: int + permuted_scores: torch.Tensor | None = None + + +@dataclass +class _PendingDispatchState: + hidden_states: torch.Tensor + dispatched_indices: torch.Tensor + dispatched_scores: torch.Tensor + num_tokens_per_expert: torch.Tensor + handle_id: torch.Tensor + score_before_experts: bool + + +def dispatch_tokens_async( + hidden_states: torch.Tensor, + selected_experts_indices: torch.Tensor, + top_scores: torch.Tensor, + num_experts: int, + group: ProcessGroup, + *, + score_before_experts: bool = True, +) -> _PendingDispatchState: + selected_experts_indices = selected_experts_indices.contiguous() + top_scores = top_scores.contiguous() + selected_experts_indices = selected_experts_indices.masked_fill(top_scores == 0, -1) + if top_scores.dtype != torch.float32: + top_scores = top_scores.float() + + buffer = get_buffer(group, get_hidden_bytes(hidden_states)) + num_tokens_per_rank, num_tokens_per_rdma_rank, num_tokens_per_expert_dispatch, is_token_in_rank, _ = ( + buffer.get_dispatch_layout(topk_idx=selected_experts_indices, num_experts=num_experts) + ) + + hidden_states, dispatched_indices, dispatched_expert_scores, num_tokens_per_expert, handle_id = ( + torch.ops.deepep.dispatch( + hidden_states, + selected_experts_indices, + top_scores, + num_tokens_per_rank, + num_tokens_per_rdma_rank, + is_token_in_rank, + num_tokens_per_expert_dispatch, + ) + ) + + return _PendingDispatchState( + hidden_states=hidden_states, + dispatched_indices=dispatched_indices, + dispatched_scores=dispatched_expert_scores, + num_tokens_per_expert=num_tokens_per_expert, + handle_id=handle_id, + score_before_experts=score_before_experts, + ) + + +def finalize_dispatch_tokens(pending_state: _PendingDispatchState) -> tuple[torch.Tensor, torch.Tensor, _DispatchState]: + _sync_dispatch(pending_state.handle_id) + + hidden_states = pending_state.hidden_states + num_recv_tokens = hidden_states.shape[0] + hidden_states, permuted_scores, permuted_indices = _permute_tokens( + hidden_states, + pending_state.dispatched_indices, + pending_state.dispatched_scores, + ) + num_tokens_per_expert = pending_state.num_tokens_per_expert.to(hidden_states.device) + + if pending_state.score_before_experts and permuted_scores is not None: + hidden_states = (hidden_states.to(torch.float32) * permuted_scores.to(torch.float32).reshape(-1, 1)).to( + hidden_states.dtype + ) + permuted_scores_for_state = None + else: + permuted_scores_for_state = permuted_scores + + state = _DispatchState( + handle_id=pending_state.handle_id, + permuted_indices=permuted_indices, + num_recv_tokens=num_recv_tokens, + permuted_scores=permuted_scores_for_state, + ) + return hidden_states, num_tokens_per_expert, state + + +def combine_tokens(hidden_states: torch.Tensor, state: _DispatchState) -> torch.Tensor: + if state.permuted_scores is not None: + hidden_states = (hidden_states.to(torch.float32) * state.permuted_scores.to(torch.float32).reshape(-1, 1)).to( + hidden_states.dtype + ) + hidden_states = _unpermute_tokens(hidden_states, state.permuted_indices, state.num_recv_tokens) + return _DeepEPCombine.apply(hidden_states, state.handle_id) + + +register_deepep_cuda_ops() + +__all__ = [ + "combine_tokens", + "configure_num_sms", + "dispatch_tokens_async", + "finalize_dispatch_tokens", + "sync_combine", +] diff --git a/src/prime_rl/trainer/distributed/expert_parallel.py b/src/prime_rl/trainer/distributed/expert_parallel.py new file mode 100644 index 00000000000..23d0eb22086 --- /dev/null +++ b/src/prime_rl/trainer/distributed/expert_parallel.py @@ -0,0 +1,27 @@ +import torch.nn as nn +from torch.distributed import ProcessGroup +from torch.distributed.tensor import DeviceMesh, Shard, distribute_module, distribute_tensor +from torch.distributed.tensor.parallel import ParallelStyle + + +class DeepEPExpertParallel(ParallelStyle): + """Expert-parallel style backed by DeepEP dispatch/combine. + + Only handles weight sharding (Shard(0) on expert dim) and stores the EP + process group on the module. PrimeRL drives DeepEP dispatch/combine from + `MoE.forward()` so communication stays outside the selective-AC checkpoint + boundary while local expert matmuls remain checkpointable. + """ + + @staticmethod + def _partition_fn(name: str, mod: nn.Module, device_mesh: DeviceMesh) -> None: + for param_name, param in mod.named_parameters(recurse=False): + mod.register_parameter(param_name, nn.Parameter(distribute_tensor(param, device_mesh, [Shard(0)]))) + mod._ep_group = device_mesh.get_group() + + def _apply(self, module: nn.Module, device_mesh: DeviceMesh) -> nn.Module: + return distribute_module(module, device_mesh, partition_fn=self._partition_fn) + + +def get_ep_group(experts: nn.Module) -> ProcessGroup: + return experts._ep_group diff --git a/src/prime_rl/trainer/envs.py b/src/prime_rl/trainer/envs.py deleted file mode 100644 index b3511751780..00000000000 --- a/src/prime_rl/trainer/envs.py +++ /dev/null @@ -1,46 +0,0 @@ -from typing import TYPE_CHECKING, Any - -from prime_rl.utils.envs import _ENV_PARSERS as _BASE_ENV_PARSERS, get_env_value, get_dir, set_defaults - -if TYPE_CHECKING: - # Enable type checking for shared envs - # ruff: noqa - from prime_rl.utils.envs import * - - # PyTorch - RANK: int - WORLD_SIZE: int - LOCAL_RANK: int - LOCAL_WORLD_SIZE: int - MASTER_ADDR: str - MASTER_PORT: int - - -_TRAINING_ENV_PARSERS = { - "RANK": int, - "WORLD_SIZE": int, - "LOCAL_RANK": int, - "LOCAL_WORLD_SIZE": int, - "MASTER_ADDR": str, - "MASTER_PORT": int, - **_BASE_ENV_PARSERS, -} - -_TRAINING_ENV_DEFAULTS = { - "RANK": "0", - "WORLD_SIZE": "1", - "LOCAL_RANK": "0", - "LOCAL_WORLD_SIZE": "1", - "MASTER_ADDR": "localhost", - "MASTER_PORT": "29500", -} - -set_defaults(_TRAINING_ENV_DEFAULTS) - - -def __getattr__(name: str) -> Any: - return get_env_value(_TRAINING_ENV_PARSERS, name) - - -def __dir__() -> list[str]: - return get_dir(_TRAINING_ENV_PARSERS) diff --git a/src/prime_rl/trainer/lora.py b/src/prime_rl/trainer/lora.py index 61dd5540ec5..113903602ae 100644 --- a/src/prime_rl/trainer/lora.py +++ b/src/prime_rl/trainer/lora.py @@ -4,11 +4,15 @@ import torch import torch.nn as nn -from prime_rl.trainer.config import LoRAConfig +from prime_rl.configs.trainer import LoRAConfig from prime_rl.trainer.models.layers.lora import MultiLoRALinear, MultiLoRAModule -from prime_rl.trainer.models.layers.lora.multi_moe import MultiLoRAGroupedExperts -from prime_rl.trainer.models.layers.moe import GroupedExperts -from prime_rl.trainer.runs import get_runs +from prime_rl.trainer.models.layers.lora.multi_moe import ( + MultiLoRAGptOssGroupedExperts, + MultiLoRAGroupedExperts, + MultiLoRANonGatedGroupedExperts, +) +from prime_rl.trainer.models.layers.moe import GptOssGroupedExperts, GroupedExperts, NonGatedGroupedExperts +from prime_rl.trainer.runs import get_multi_run_manager from prime_rl.utils.logger import get_logger @@ -74,8 +78,8 @@ def _find_target_modules(model: nn.Module, target_patterns: List[str]) -> List[s target_modules = [] for name, module in model.named_modules(): - # Check if module is Linear or GroupedExperts - if not (isinstance(module, nn.Linear) or isinstance(module, GroupedExperts)): + # Check if module is Linear or one of the supported expert classes + if not isinstance(module, (nn.Linear, GroupedExperts, NonGatedGroupedExperts, GptOssGroupedExperts)): continue for pattern in target_patterns: @@ -137,7 +141,11 @@ def apply_lora_to_model(model: nn.Module, config: LoRAConfig) -> None: config: LoRA configuration """ logger = get_logger() - n_loras = get_runs().max_runs + from prime_rl.trainer.models import PreTrainedModelPrimeRL + + if isinstance(model, PreTrainedModelPrimeRL): + get_multi_run_manager().register_adapter_state_dict_converter(type(model).convert_adapter_to_hf) + n_loras = get_multi_run_manager().max_runs from torch.distributed.fsdp import FSDPModule @@ -178,14 +186,32 @@ def apply_lora_to_model(model: nn.Module, config: LoRAConfig) -> None: alpha=config.alpha, dropout=config.dropout, ) + # Handle NonGatedGroupedExperts (relu2 experts used by NemotronH's LatentMoE) + elif isinstance(base_module, NonGatedGroupedExperts): + lora_module = MultiLoRANonGatedGroupedExperts( + base_layer=base_module, + rank=config.rank, + n_adapters=n_loras, + alpha=config.alpha, + dropout=config.dropout, + ) + # Handle GptOssGroupedExperts (gpt-oss fused gate_up + biases) + elif isinstance(base_module, GptOssGroupedExperts): + lora_module = MultiLoRAGptOssGroupedExperts( + base_layer=base_module, + rank=config.rank, + n_adapters=n_loras, + alpha=config.alpha, + dropout=config.dropout, + ) else: logger.warning( f"Module {module_name} is type {type(base_module).__name__}, " - f"expected nn.Linear or GroupedExperts. Skipping." + f"expected nn.Linear, GroupedExperts, NonGatedGroupedExperts, or GptOssGroupedExperts. Skipping." ) continue - lora_module.register_with_runs(get_runs(), module_name) + lora_module.register_with_runs(get_multi_run_manager(), module_name) _set_module_by_name(model, module_name, lora_module) freeze_all_except_lora_and_specified(model, config) @@ -234,14 +260,16 @@ def clean_lora_state_dict(state_dict: Dict[str, torch.Tensor]) -> Dict[str, torc return clean_state_dict -def save_lora_config(config: LoRAConfig, model: nn.Module, save_path) -> None: +def save_lora_config(model: nn.Module, save_path, rank: int, alpha: float, dropout: float) -> None: """ Save LoRA configuration as JSON for adapter portability. Args: - config: LoRA configuration to save model: Model with LoRA layers to introspect save_path: Path object or string pointing to directory where adapter_config.json will be saved + rank: LoRA rank + alpha: LoRA alpha scaling parameter + dropout: LoRA dropout rate """ import json from pathlib import Path @@ -266,9 +294,9 @@ def save_lora_config(config: LoRAConfig, model: nn.Module, save_path) -> None: "peft_type": "LORA", "task_type": "CAUSAL_LM", "base_model_name_or_path": model.config._name_or_path, - "r": config.rank, - "lora_alpha": config.alpha, - "lora_dropout": config.dropout, + "r": rank, + "lora_alpha": alpha, + "lora_dropout": dropout, "bias": "none", "target_modules": sorted(list(target_modules)), "modules_to_save": sorted(list(modules_to_save)) if modules_to_save else None, diff --git a/src/prime_rl/trainer/model.py b/src/prime_rl/trainer/model.py index 7479b84796d..910a978a663 100644 --- a/src/prime_rl/trainer/model.py +++ b/src/prime_rl/trainer/model.py @@ -1,14 +1,20 @@ import logging +import os import time from pathlib import Path from typing import cast +# Disable transformers hub kernel interception by default. The `kernels` package, when installed, +# causes transformers to auto-replace modules (e.g. mamba-ssm) with hub kernel versions that may +# have incompatible CUDA requirements. We only enable it explicitly for models that need it (GPT-OSS). +os.environ.setdefault("USE_HUB_KERNELS", "NO") + import torch +import torch._dynamo import torch.nn as nn from beartype import beartype as typechecker from huggingface_hub import snapshot_download -from jaxtyping import Int, jaxtyped -from liger_kernel.transformers import AutoLigerKernelForCausalLM +from jaxtyping import Float, Int, jaxtyped from torch import Tensor from torch.distributed.algorithms._checkpoint.checkpoint_wrapper import checkpoint_wrapper from torch.distributed.checkpoint.hf_storage import HuggingFaceStorageReader @@ -17,29 +23,275 @@ from torch.distributed.fsdp import CPUOffloadPolicy, FSDPModule, MixedPrecisionPolicy, OffloadPolicy, fully_shard from torch.distributed.tensor.parallel import parallelize_module from torchtitan.distributed.expert_parallel import ExpertParallel -from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer, PretrainedConfig +from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer, GenerationConfig, PretrainedConfig from transformers.tokenization_utils import PreTrainedTokenizer from transformers.utils.import_utils import is_flash_attn_3_available -from prime_rl.trainer.config import ActivationCheckpointConfig, CompileConfig, ModelConfig, TokenizerConfig -from prime_rl.trainer.lora import apply_lora_to_model, strip_lora_from_state_dict +from prime_rl.configs.trainer import ActivationCheckpointConfig, CompileConfig, ModelConfig, TokenizerConfig +from prime_rl.trainer.distributed import DeepEPExpertParallel +from prime_rl.trainer.lora import apply_lora_to_model, freeze_all_except_lora_and_specified, strip_lora_from_state_dict from prime_rl.trainer.models import ( AutoModelForCausalLMPrimeRL, PreTrainedModelPrimeRL, PrimeLmOutput, cast_float_and_contiguous, + get_custom_vlm_cls, supports_custom_impl, ) +from prime_rl.trainer.models.layers.checkpointing import ( + get_supported_targets, + set_selective_activation_checkpointing, + supports_selective_activation_checkpointing, +) +from prime_rl.trainer.models.layers.fp8_linear import replace_linear_with_fp8_blockwise_linear from prime_rl.trainer.models.layers.lm_head import inject_prime_lm_head -from prime_rl.trainer.models.layers.moe import MoE +from prime_rl.trainer.models.layers.moe import LatentMoE, MoE from prime_rl.trainer.parallel_dims import ParallelDims from prime_rl.trainer.weights import ( load_state_dict, + load_state_dict_keys, save_state_dict, ) from prime_rl.trainer.world import get_world from prime_rl.utils.logger import get_logger -from prime_rl.utils.tensor_hashing import get_module_signature +from prime_rl.utils.sequence import get_cu_seqlens_from_position_ids +from prime_rl.utils.utils import format_time +from prime_rl.utils.vlm import get_language_model, get_vision_encoder, is_vlm_architecture + + +def pre_download_model(model_name: str) -> None: + """Pre-download model from HuggingFace Hub so all nodes have cached weights before training.""" + if Path(model_name).exists(): + get_logger().info(f"Model {model_name} found at local path, skipping download") + return + get_logger().info(f"Pre-downloading model {model_name}") + t0 = time.perf_counter() + path = snapshot_download(repo_id=model_name, repo_type="model") + get_logger().debug( + f"Finished pre-downloading model {model_name} to {path} in {format_time(time.perf_counter() - t0)}" + ) + + +def _patch_qwen3_5_moe_conversion_mapping(): + """Fix Qwen3.5 MoE conversion mapping incorrectly applying qwen2_moe expert weight splitting. + + Qwen3.5 MoE stores expert weights as fused 3D tensors natively in the checkpoint + (e.g. experts.gate_up_proj [num_experts, 2*intermediate, hidden]). The upstream mapping + incorrectly maps qwen3_5_moe → qwen2_moe, which assumes per-expert 2D checkpoint weights, + causing revert_weight_conversion to produce wrong shapes during weight broadcasting. + + Remove once the pinned transformers commit fixes this. + """ + from transformers.conversion_mapping import ( + get_checkpoint_conversion_mapping, + register_checkpoint_conversion_mapping, + ) + + # qwen3_5_moe_text: keep only the qwen3_5_text renaming, remove qwen2_moe expert conversion + qwen3_5_text_mapping = get_checkpoint_conversion_mapping("qwen3_5_text") + if qwen3_5_text_mapping is not None: + register_checkpoint_conversion_mapping("qwen3_5_moe_text", qwen3_5_text_mapping, overwrite=True) + + # qwen3_5_moe: remove the qwen2_moe fallback entirely + register_checkpoint_conversion_mapping("qwen3_5_moe", [], overwrite=True) + + +def _patch_qwen3_5_text_position_ids(): + """Fix Qwen3.5 passing 3D MRoPE position_ids to decoder layers instead of 2D text_position_ids. + + Upstream fix: https://github.com/huggingface/transformers/pull/44399 + Remove once the pinned transformers commit includes this fix. + """ + import inspect + + from transformers.models.qwen3_5.modeling_qwen3_5 import Qwen3_5DecoderLayer, Qwen3_5TextModel + from transformers.models.qwen3_5_moe.modeling_qwen3_5_moe import Qwen3_5MoeDecoderLayer, Qwen3_5MoeTextModel + + for text_model_cls, decoder_layer_cls in [ + (Qwen3_5TextModel, Qwen3_5DecoderLayer), + (Qwen3_5MoeTextModel, Qwen3_5MoeDecoderLayer), + ]: + source = inspect.getsource(text_model_cls.forward) + if "decoder_layer" in source and "position_ids=text_position_ids" in source.split("decoder_layer")[-1]: + continue # already fixed upstream + + _original_forward = decoder_layer_cls.forward + + def _make_patched_forward(original): + def _patched_forward(self, hidden_states, position_ids=None, **kwargs): + if position_ids is not None and position_ids.ndim == 3: + position_ids = position_ids[0] + return original(self, hidden_states, position_ids=position_ids, **kwargs) + + return _patched_forward + + decoder_layer_cls.forward = _make_patched_forward(_original_forward) + + +def _patch_qwen3_5_linear_attn_varlen(): + """Thread cu_seqlens through Qwen3.5 GatedDeltaNet so packed batches don't + leak conv/SSM state across sequences. + + HF's forward hardcodes seq_idx=None for causal_conv1d and omits cu_seqlens + for chunk_gated_delta_rule, so packed RL training sees ~0.23 Mismatch KL vs + vLLM (target <0.01). Mirrors the NemotronH mamba fix. + """ + import torch.nn.functional as F + from transformers.models.qwen3_5.modeling_qwen3_5 import ( + Qwen3_5DecoderLayer, + Qwen3_5GatedDeltaNet, + Qwen3_5TextModel, + apply_mask_to_padding_states, + ) + + if getattr(Qwen3_5GatedDeltaNet.forward, "_prl_varlen_patched", False): + return + + _gdn_orig = Qwen3_5GatedDeltaNet.forward + + def _gdn_forward(self, hidden_states, cache_params=None, attention_mask=None, cu_seqlens=None): + if cu_seqlens is None or cache_params is not None: + return _gdn_orig(self, hidden_states, cache_params=cache_params, attention_mask=attention_mask) + + hidden_states = apply_mask_to_padding_states(hidden_states, attention_mask) + batch_size, seq_len, _ = hidden_states.shape + + mixed_qkv = self.in_proj_qkv(hidden_states).transpose(1, 2) + z = self.in_proj_z(hidden_states).reshape(batch_size, seq_len, -1, self.head_v_dim) + b = self.in_proj_b(hidden_states) + a = self.in_proj_a(hidden_states) + + if self.causal_conv1d_fn is not None: + seg_lens = cu_seqlens[1:] - cu_seqlens[:-1] + seq_idx = torch.repeat_interleave( + torch.arange(seg_lens.numel(), dtype=torch.int32, device=hidden_states.device), + seg_lens, + ).unsqueeze(0) + mixed_qkv = self.causal_conv1d_fn( + x=mixed_qkv, + weight=self.conv1d.weight.squeeze(1), + bias=self.conv1d.bias, + activation=self.activation, + seq_idx=seq_idx, + ) + else: + # Per-segment conv1d so the kernel-1 left pad only draws from within each sequence. + cu = cu_seqlens.tolist() + conv_outs = [] + for i in range(len(cu) - 1): + s, e = cu[i], cu[i + 1] + if s == e: + continue + conv_outs.append(self.conv1d(mixed_qkv[:, :, s:e])[:, :, : e - s]) + mixed_qkv = F.silu(torch.cat(conv_outs, dim=-1)) + + mixed_qkv = mixed_qkv.transpose(1, 2) + query, key, value = torch.split(mixed_qkv, [self.key_dim, self.key_dim, self.value_dim], dim=-1) + query = query.reshape(batch_size, seq_len, -1, self.head_k_dim) + key = key.reshape(batch_size, seq_len, -1, self.head_k_dim) + value = value.reshape(batch_size, seq_len, -1, self.head_v_dim) + + beta = b.sigmoid() + g = -self.A_log.float().exp() * F.softplus(a.float() + self.dt_bias) + if self.num_v_heads // self.num_k_heads > 1: + query = query.repeat_interleave(self.num_v_heads // self.num_k_heads, dim=2) + key = key.repeat_interleave(self.num_v_heads // self.num_k_heads, dim=2) + + core_attn_out, _ = self.chunk_gated_delta_rule( + query, + key, + value, + g=g, + beta=beta, + initial_state=None, + output_final_state=False, + use_qk_l2norm_in_kernel=True, + cu_seqlens=cu_seqlens, + ) + + core_attn_out = core_attn_out.reshape(-1, self.head_v_dim) + z = z.reshape(-1, self.head_v_dim) + core_attn_out = self.norm(core_attn_out, z) + core_attn_out = core_attn_out.reshape(batch_size, seq_len, -1) + return self.out_proj(core_attn_out) + + _gdn_forward._prl_varlen_patched = True + Qwen3_5GatedDeltaNet.forward = _gdn_forward + + _dec_orig = Qwen3_5DecoderLayer.forward + + def _dec_forward( + self, + hidden_states, + position_embeddings, + attention_mask=None, + position_ids=None, + past_key_values=None, + cu_seqlens=None, + **kwargs, + ): + if position_ids is not None and position_ids.ndim == 3: + position_ids = position_ids[0] + if self.layer_type != "linear_attention": + return _dec_orig( + self, + hidden_states, + position_embeddings=position_embeddings, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=past_key_values, + **kwargs, + ) + + residual = hidden_states + hidden_states = self.input_layernorm(hidden_states) + hidden_states = self.linear_attn( + hidden_states=hidden_states, + cache_params=past_key_values, + attention_mask=attention_mask, + cu_seqlens=cu_seqlens, + ) + hidden_states = residual + hidden_states + residual = hidden_states + hidden_states = self.post_attention_layernorm(hidden_states) + hidden_states = self.mlp(hidden_states) + return residual + hidden_states + + Qwen3_5DecoderLayer.forward = _dec_forward + + _text_orig = Qwen3_5TextModel.forward + + def _text_forward( + self, + input_ids=None, + attention_mask=None, + position_ids=None, + past_key_values=None, + inputs_embeds=None, + use_cache=None, + **kwargs, + ): + attn_impl = getattr(self.config, "_attn_implementation", None) + cu_seqlens = None + if attn_impl in ("flash_attention_2", "flash_attention_3", "fa4") and position_ids is not None: + pids = position_ids + if pids.ndim == 3: + pids = pids[0] + cu_seqlens, _ = get_cu_seqlens_from_position_ids(pids) + kwargs["cu_seqlens"] = cu_seqlens + return _text_orig( + self, + input_ids=input_ids, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=past_key_values, + inputs_embeds=inputs_embeds, + use_cache=use_cache, + **kwargs, + ) + + Qwen3_5TextModel.forward = _text_forward + # Add filter to the standard logging module for transformers.modeling_utils to supress the # flash attention dtype warnings since FSDP is used to handle mixed precision. @@ -53,29 +305,109 @@ "float32": torch.float32, } +# We increase the torch.compile recompile limit and cache size as we found this +# necessary for training INTELLECT-3 with Muon. +torch._dynamo.config.recompile_limit = 16 # default: 8 +torch._dynamo.config.cache_size_limit = 64 # default: 8 + + +def freeze_vision_encoder(model: nn.Module, override_attr: str | None = None) -> None: + logger = get_logger() + vision_encoder = get_vision_encoder(model, override=override_attr) + if vision_encoder is None: + raise ValueError("Could not find vision encoder to freeze") + num_frozen = 0 + for param in vision_encoder.parameters(): + param.requires_grad = False + num_frozen += 1 + logger.info(f"Froze {num_frozen} parameters in vision encoder") + + +def freeze_moe_router(model: nn.Module) -> None: + """Freeze MoE router parameters to maintain stable routing during training.""" + logger = get_logger() + language_model = get_language_model(model) + num_frozen = 0 + + for layer in language_model.layers: + mlp = layer.mlp if hasattr(layer, "mlp") else layer.feed_forward if hasattr(layer, "feed_forward") else None + if mlp is None: + continue + + # Custom implementation: MoE/LatentMoE class with router attribute + if isinstance(mlp, (MoE, LatentMoE)): + for param in mlp.router.parameters(): + param.requires_grad = False + num_frozen += 1 + # HuggingFace implementation: gate attribute (nn.Linear) + elif hasattr(mlp, "gate") and isinstance(mlp.gate, nn.Linear): + for param in mlp.gate.parameters(): + param.requires_grad = False + num_frozen += 1 + + if num_frozen == 0: + raise ValueError("No MoE router parameters found to freeze. Is this a MoE model?") + + logger.info(f"Froze {num_frozen} MoE router parameters") + + +def apply_force_balanced_routing(model: nn.Module) -> None: + """Force MoE token-choice routers into round-robin assignment for fake-data smoke tests.""" + logger = get_logger() + language_model = get_language_model(model) + num_routers = 0 + + for layer in language_model.layers: + mlp = layer.mlp if hasattr(layer, "mlp") else layer.feed_forward if hasattr(layer, "feed_forward") else None + if isinstance(mlp, (MoE, LatentMoE)): + mlp.router.force_balanced = True + num_routers += 1 + + if num_routers == 0: + raise ValueError("No MoE routers found to force-balance. Is this a custom-impl MoE model?") + + logger.warning( + f"Forced balanced routing on {num_routers} MoE layers (debug.force_balanced_routing=True). " + "Expert assignment is round-robin; gradient flow through the router is broken." + ) + def is_tt_moe_model(model: nn.Module) -> bool: return hasattr(model.config, "num_experts") or hasattr(model.config, "n_routed_experts") +def configure_moe_ep_backend(model: nn.Module, config: ModelConfig) -> None: + backend = config.ep_comm_backend + if backend == "deepep": + from prime_rl.trainer.distributed.deepep import configure_num_sms + + configure_num_sms(config.deepep_num_sms) + language_model = get_language_model(model) + for transformer_block in language_model.layers: + if not isinstance(transformer_block.mlp, (MoE, LatentMoE)): + continue + transformer_block.mlp.set_ep_comm_backend(backend) + transformer_block.mlp.set_deepep_token_chunk_size(config.deepep_token_chunk_size) + + def get_load_balance_stats( model: nn.Module, reset_stats: bool = True, try_to_avoid_padding_experts: bool = True ) -> dict[str, Tensor | None]: per_layer_max_vio = [] - for transformer_block in model.model.layers: + language_model = get_language_model(model) + for transformer_block in language_model.layers: # This is necessary for models that have mixed dense layers - if not hasattr(transformer_block.mlp, "tokens_per_expert"): + block_mlp = getattr(transformer_block, "mlp", None) + if block_mlp is None or not hasattr(block_mlp, "tokens_per_expert"): continue - tokens_per_expert: torch.Tensor = transformer_block.mlp.tokens_per_expert + tokens_per_expert: torch.Tensor = block_mlp.tokens_per_expert if try_to_avoid_padding_experts: - tokens_per_expert = tokens_per_expert.sort(dim=0, descending=True).values[ - transformer_block.mlp.router.top_k : - ] + tokens_per_expert = tokens_per_expert.sort(dim=0, descending=True).values[block_mlp.router.top_k :] balanced_load = tokens_per_expert.mean() max_vio = (tokens_per_expert.max() - balanced_load) / balanced_load per_layer_max_vio.append(max_vio.item()) if reset_stats: - transformer_block.mlp.tokens_per_expert.zero_() + block_mlp.tokens_per_expert.zero_() if len(per_layer_max_vio) == 0: return {"max_vio": None} return {"max_vio": torch.tensor(per_layer_max_vio, device=torch.device("cuda"))} @@ -88,6 +420,14 @@ def get_model( logger.info( f"Loading model config (name={config.name}, attn={config.attn}, trust_remote_code={config.trust_remote_code})" ) + + is_vlm_training = config.vlm is not None + + if "Qwen3.5" in config.name or "qwen3_5" in config.name.lower(): + _patch_qwen3_5_text_position_ids() + _patch_qwen3_5_moe_conversion_mapping() + _patch_qwen3_5_linear_attn_varlen() + model_config = cast( PretrainedConfig, AutoConfig.from_pretrained( @@ -95,48 +435,130 @@ def get_model( ), ) model_config.use_cache = False + is_vlm_arch = is_vlm_architecture(model_config) + + if is_vlm_training: + logger.info(f"Detected vision-language model: {config.name}") + if config.optimization_dtype != "bfloat16" or config.reduce_dtype != "bfloat16": + raise ValueError( + "VLM models must use optimization_dtype='bfloat16' and reduce_dtype='bfloat16' to match vLLM inference." + ) + + # GPT-OSS only supports FlashAttention via kernels-community/vllm-flash-attn3, which requires Hopper (SM 90). + # On other architectures (e.g. Blackwell), users must fall back to eager attention. + HOPPER_MAJOR = 9 + if getattr(model_config, "model_type", "") == "gpt_oss": + if config.attn != "eager": + major, minor = torch.cuda.get_device_capability() + if major != HOPPER_MAJOR: + raise ValueError( + f"GPT-OSS requires 'attn = \"eager\"' on non-Hopper GPUs (detected SM {major}{minor}). " + f"The only flash attention kernel supported by GPT-OSS (kernels-community/vllm-flash-attn3) is Hopper-only. " + f'Set [trainer.model] attn = "eager" in your config.' + ) + # Enable hub kernels for GPT-OSS (disabled by default to avoid interfering with other models). + import transformers.integrations.hub_kernels as _hub_kernels + + _hub_kernels._kernels_enabled = True + + # Fallback Qwen3.5 patch detection from loaded config model_type + if getattr(model_config, "model_type", "").startswith("qwen3_5_moe"): + _patch_qwen3_5_text_position_ids() + _patch_qwen3_5_moe_conversion_mapping() + _patch_qwen3_5_linear_attn_varlen() + for subconfig_key in getattr(model_config, "sub_configs", {}): + subconfig = getattr(model_config, subconfig_key, None) + if subconfig is not None and hasattr(subconfig, "use_cache"): + subconfig.use_cache = False model_config.use_grouped_mm = config.moe_use_grouped_mm + model_config.fp8 = config.fp8 + + # Ensure pad_token_id is set (some models like Qwen3MoE don't have it). + # In transformers v5, token IDs moved from PretrainedConfig to GenerationConfig. + if not hasattr(model_config, "pad_token_id") or model_config.pad_token_id is None: + gen_config = GenerationConfig.from_model_config(model_config) + # Use `is not None` instead of truthiness: token ID 0 is valid. + pad_token_id = next( + ( + v + for v in [gen_config.pad_token_id, gen_config.eos_token_id, getattr(model_config, "eos_token_id", None)] + if v is not None + ), + None, + ) + # Some HF configs (e.g. Llama 3.2) set pad_token_id to a list, which + # crashes both huggingface_hub's strict setter and transformers' + # GenerationConfig.validate(). Unwrap before assigning. + if isinstance(pad_token_id, list): + pad_token_id = pad_token_id[0] + model_config.pad_token_id = pad_token_id + + # Handle list pad_token_id that was already set on the config (not from our + # fallback above, but directly in the model's config.json). + if isinstance(getattr(model_config, "pad_token_id", None), list): + model_config.pad_token_id = model_config.pad_token_id[0] + + # NOTE: For VLM models, we do NOT propagate dtype to sub_configs. + # The model should load in its default dtype (bf16) to match vLLM inference. + # The FSDP MixedPrecisionPolicy handles compute dtype separately. + logger.debug(f"Loaded model config ({model_config.to_dict()})") if config.debug.num_layers is not None: - num_hidden_layers = min(config.debug.num_layers, model_config.num_hidden_layers) + # VLM configs nest num_hidden_layers under text_config + target_config = getattr(model_config, "text_config", model_config) + num_hidden_layers = min(config.debug.num_layers, target_config.num_hidden_layers) logger.warning( - f"Setting the number of layers to {config.debug.num_layers} in the model config. This means {model_config.num_hidden_layers - num_hidden_layers} layers will not be loaded." + f"Setting the number of layers to {config.debug.num_layers} in the model config. This means {target_config.num_hidden_layers - num_hidden_layers} layers will not be loaded." ) - model_config.num_hidden_layers = num_hidden_layers + target_config.num_hidden_layers = num_hidden_layers # Determine the implementation to use + custom_vlm_cls = get_custom_vlm_cls(model_config) if is_vlm_arch else None if config.impl == "auto": - impl_to_use = "custom" if supports_custom_impl(model_config) else "hf" - logger.info( - f"Auto-selected implementation: {impl_to_use} (custom implementation {'supported' if supports_custom_impl(model_config) else 'not supported'})" - ) + if is_vlm_arch: + impl_to_use = "custom" if custom_vlm_cls is not None else "hf" + else: + impl_to_use = "custom" if supports_custom_impl(model_config) else "hf" + logger.info(f"Auto-selected implementation: {impl_to_use}") else: impl_to_use = config.impl with device: - match impl_to_use: - case "hf": - model_cls = AutoModelForCausalLM - case "liger_kernel": - model_cls = AutoLigerKernelForCausalLM - case "custom": - model_cls = AutoModelForCausalLMPrimeRL + if impl_to_use == "custom" and custom_vlm_cls is not None: + model_cls = custom_vlm_cls + elif is_vlm_arch: + from transformers import AutoModelForImageTextToText + + model_cls = AutoModelForImageTextToText + else: + match impl_to_use: + case "hf": + model_cls = AutoModelForCausalLM + case "custom": + model_cls = AutoModelForCausalLMPrimeRL load_model_start_time = time.perf_counter() + # HF VLM models require torch_dtype; custom PrimeRL models and text Auto models use dtype + use_torch_dtype = is_vlm_arch and model_cls is not custom_vlm_cls + dtype_kwarg = {"torch_dtype": dtype} if use_torch_dtype else {"dtype": dtype} if device == torch.device("meta"): logger.info(f"Loading model {config.name} using {model_cls.__name__} to meta device") - model = model_cls.from_config(model_config, trust_remote_code=config.trust_remote_code, dtype=dtype) + model = model_cls.from_config(model_config, trust_remote_code=config.trust_remote_code, **dtype_kwarg) else: logger.info(f"Loading model {config.name} using {model_cls.__name__} to CPU") model = model_cls.from_pretrained( pretrained_model_name_or_path=config.name, config=model_config, trust_remote_code=config.trust_remote_code, - dtype=dtype, + **dtype_kwarg, ) logger.debug(f"Loaded model {config.name} in {time.perf_counter() - load_model_start_time:.2f} seconds") + # For VLM models, optionally freeze the vision encoder + if is_vlm_training and config.vlm.freeze_vision_encoder: + freeze_vision_encoder(model, override_attr=config.vlm.vision_encoder_attr) + assert model.lm_head.weight.dtype == dtype, ( f"LM head dtype wasnt loaded correctly {model.lm_head.weight.dtype} != {dtype}" ) @@ -144,9 +566,17 @@ def get_model( def setup_tokenizer(config: TokenizerConfig) -> PreTrainedTokenizer: + logger = get_logger() tokenizer = AutoTokenizer.from_pretrained(config.name, trust_remote_code=config.trust_remote_code) if config.chat_template is not None: - tokenizer.chat_template = config.chat_template + path = Path(config.chat_template) + if path.is_file(): + logger.info(f"Loading custom chat template from file: {path}") + tokenizer.chat_template = path.read_text() + logger.debug(f"Chat template content:\n{tokenizer.chat_template}") + else: + logger.info("Using inline custom chat template") + tokenizer.chat_template = config.chat_template tokenizer.pad_token_id = tokenizer.eos_token_id return tokenizer @@ -162,10 +592,7 @@ def setup_fsdp(model: nn.Module, config: ModelConfig, parallel_dims: ParallelDim "reshard_after_forward": config.reshard_after_forward, } - if config.dp_replicate > 1: - hsdp_mesh = parallel_dims.world_mesh["dp_replicate", "dp_shard_cp"] - else: - hsdp_mesh = parallel_dims.world_mesh["dp_shard_cp"] + hsdp_mesh = parallel_dims.get_mesh("hsdp") dp_mod_ep_mesh: DeviceMesh | None = None if parallel_dims.ep_enabled: @@ -176,11 +603,24 @@ def setup_fsdp(model: nn.Module, config: ModelConfig, parallel_dims: ParallelDim dp_mod_ep_mesh = parallel_dims.world_mesh[tuple(dp_mod_ep_mesh_dim_names)] - for transformer_block in model.model.layers: - if parallel_dims.ep_enabled and isinstance(transformer_block.mlp, MoE): - fully_shard(transformer_block.mlp.experts, mesh=dp_mod_ep_mesh, **fsdp_config) + is_vlm_training = config.vlm is not None + if is_vlm_training: + vision_encoder = get_vision_encoder(model, override=config.vlm.vision_encoder_attr) + if vision_encoder is None: + raise ValueError(f"VLM model {config.name} has no recognized vision encoder") + + fully_shard(vision_encoder, mesh=hsdp_mesh, **fsdp_config) + get_logger().info(f"Applied FSDP to vision encoder (frozen={config.vlm.freeze_vision_encoder})") - transformer_block.mlp.experts.set_gradient_divide_factor(parallel_dims.fsdp_gradient_divide_factor) + language_model = get_language_model(model, override=config.vlm.language_model_attr if is_vlm_training else None) + transformer_layers = language_model.layers + + for transformer_block in transformer_layers: + block_mlp = getattr(transformer_block, "mlp", None) + if parallel_dims.ep_enabled and block_mlp is not None and isinstance(block_mlp, (MoE, LatentMoE)): + fully_shard(block_mlp.experts, mesh=dp_mod_ep_mesh, **fsdp_config) + + block_mlp.experts.set_gradient_divide_factor(parallel_dims.fsdp_gradient_divide_factor) fully_shard( transformer_block, @@ -192,20 +632,22 @@ def setup_fsdp(model: nn.Module, config: ModelConfig, parallel_dims: ParallelDim if shard_norm_and_lm_head: # This optimization breaks weight tying + embed_module = getattr(language_model, "embed_tokens", None) or getattr(language_model, "embeddings", None) fully_shard( - model.model.embed_tokens, + embed_module, mesh=hsdp_mesh, **fsdp_config, ) + norm_module = getattr(language_model, "norm", None) or language_model.norm_f fully_shard( - [model.lm_head, model.model.norm], + [model.lm_head, norm_module], mesh=hsdp_mesh, mp_policy=mp_policy, offload_policy=offload_policy, reshard_after_forward=False, ) else: - get_logger().warning("Model is tied word embeddings, so not doing the last layer not resharding optimization") + get_logger().warning("Model uses tied word embeddings, so skipping the last-layer no-reshard optimization.") fully_shard( model, @@ -221,30 +663,30 @@ def setup_fsdp(model: nn.Module, config: ModelConfig, parallel_dims: ParallelDim # if EP is enabled, d2h syncs in the dispatch/combine can interfere with FSDP prefetch, that's why we set it below manually # the rest of the function handles only that - transformer_blocks = list(model.model.layers) + transformer_blocks = list(language_model.layers) next_transformer_blocks = transformer_blocks[1:] + [None] - if model.model.embed_tokens is not None and len(model.model.layers) > 0: + embed_module = getattr(language_model, "embed_tokens", None) or getattr(language_model, "embeddings", None) + if embed_module is not None and len(language_model.layers) > 0: if shard_norm_and_lm_head: - model.model.embed_tokens.set_modules_to_forward_prefetch([transformer_blocks[0]]) + embed_module.set_modules_to_forward_prefetch([transformer_blocks[0]]) for transformer_block, next_transformer_block in zip(transformer_blocks, next_transformer_blocks): if next_transformer_block is not None: - if isinstance(next_transformer_block.mlp, MoE): - transformer_block.set_modules_to_forward_prefetch( - [next_transformer_block, next_transformer_block.mlp.experts] - ) + next_mlp = getattr(next_transformer_block, "mlp", None) + if next_mlp is not None and isinstance(next_mlp, (MoE, LatentMoE)): + transformer_block.set_modules_to_forward_prefetch([next_transformer_block, next_mlp.experts]) else: transformer_block.set_modules_to_forward_prefetch([next_transformer_block]) - elif model.model.norm is not None and model.lm_head is not None: + elif language_model.norm is not None and model.lm_head is not None: if shard_norm_and_lm_head: - transformer_block.set_modules_to_forward_prefetch([model.model.norm, model.lm_head]) + transformer_block.set_modules_to_forward_prefetch([language_model.norm, model.lm_head]) # backward - reversed_transformer_blocks = list(reversed(model.model.layers)) + reversed_transformer_blocks = list(reversed(language_model.layers)) prev_transformer_blocks = reversed_transformer_blocks[1:] + [None] - if model.model.norm is not None and model.lm_head is not None and len(model.model.layers) > 0: + if language_model.norm is not None and model.lm_head is not None and len(language_model.layers) > 0: if shard_norm_and_lm_head: model.lm_head.set_modules_to_backward_prefetch([reversed_transformer_blocks[0]]) else: @@ -252,24 +694,32 @@ def setup_fsdp(model: nn.Module, config: ModelConfig, parallel_dims: ParallelDim for transformer_block, prev_transformer_block in zip(reversed_transformer_blocks, prev_transformer_blocks): if prev_transformer_block is not None: - if isinstance(prev_transformer_block.mlp, MoE): - transformer_block.set_modules_to_backward_prefetch( - [prev_transformer_block, prev_transformer_block.mlp.experts] - ) + prev_mlp = getattr(prev_transformer_block, "mlp", None) + if prev_mlp is not None and isinstance(prev_mlp, (MoE, LatentMoE)): + transformer_block.set_modules_to_backward_prefetch([prev_transformer_block, prev_mlp.experts]) else: transformer_block.set_modules_to_backward_prefetch([prev_transformer_block]) - elif model.model.embed_tokens is not None: + elif embed_module is not None: if shard_norm_and_lm_head: - transformer_block.set_modules_to_backward_prefetch([model.model.embed_tokens]) + transformer_block.set_modules_to_backward_prefetch([embed_module]) def load_dcp_from_hf(model: nn.Module, config: ModelConfig, parallel_dims: ParallelDims): - model.to_empty(device="cuda") + device = "cpu" if config.fsdp_cpu_offload else "cuda" + model.to_empty(device=device) torch.distributed.barrier() + def _init_buffers_post_meta(): + if isinstance(model, PreTrainedModelPrimeRL): + model.init_buffers_post_meta() + else: + fix_model_post_empty(model) + logger = get_logger() if config.debug.random_init: logger.warning("Randomly initializing model. Skipping loading weights from HF.") + _init_buffers_post_meta() + _move_buffers_to_cuda(model, config) return if not Path(config.name).exists(): @@ -280,41 +730,40 @@ def load_dcp_from_hf(model: nn.Module, config: ModelConfig, parallel_dims: Paral ) snapshot_path = Path(config.name) - # Load the snapshot state - snapshot_state_dict = load_state_dict(snapshot_path) - model_state_dict = model.state_dict() - - # Dynamically convert between different weight formats if needed + # Dynamically convert between different weight formats if needed. + # All ranks read just the key names (cheap) to determine the path independently. + # Only master loads the full state dict when conversion is actually needed. if isinstance(model, PreTrainedModelPrimeRL): - if model.is_hf_state_dict(snapshot_state_dict) and model.is_prime_state_dict(model_state_dict): + snapshot_keys = dict.fromkeys(load_state_dict_keys(snapshot_path)) + model_keys = dict.fromkeys(model.state_dict().keys()) + + if model.is_hf_state_dict(snapshot_keys) and model.is_prime_state_dict(model_keys): logger.warning( "Found HF weight format in snapshot state dict and PrimeRL weight format in model state dict. Trying to auto-convert..." ) snapshot_path = snapshot_path / "prime" - if snapshot_path.exists(): - logger.debug(f"Conversion found at {snapshot_path}.") - else: - if get_world().is_master: - logger.debug( - f"Converting snapshot state dict to PrimeRL format and saving to {snapshot_path} on master rank. This is a one-time operation." - ) - model.convert_to_prime(snapshot_state_dict) - save_state_dict(snapshot_state_dict, snapshot_path) - - elif model.is_prime_state_dict(snapshot_state_dict) and model.is_hf_state_dict(model_state_dict): + if not snapshot_path.exists() and get_world().is_master: + logger.debug( + f"Converting snapshot state dict to PrimeRL format and saving to {snapshot_path} on master rank. This is a one-time operation." + ) + snapshot_state_dict = load_state_dict(snapshot_path.parent) + model.convert_to_prime(snapshot_state_dict) + save_state_dict(snapshot_state_dict, snapshot_path) + del snapshot_state_dict + + elif model.is_prime_state_dict(snapshot_keys) and model.is_hf_state_dict(model_keys): logger.warning( "Found PrimeRL weight format in snapshot state dict and HF weight format in model state dict. Trying to auto-convert..." ) snapshot_path = snapshot_path / "hf" - if snapshot_path.exists(): - logger.debug(f"Conversion found at {snapshot_path}.") - else: - if get_world().is_master: - logger.debug( - f"Converting snapshot state dict to HF format and saving to {snapshot_path} on master rank. This is a one-time operation." - ) - model.convert_to_hf(snapshot_state_dict) - save_state_dict(snapshot_state_dict, snapshot_path) + if not snapshot_path.exists() and get_world().is_master: + logger.debug( + f"Converting snapshot state dict to HF format and saving to {snapshot_path} on master rank. This is a one-time operation." + ) + snapshot_state_dict = load_state_dict(snapshot_path.parent) + model.convert_to_hf(snapshot_state_dict) + save_state_dict(snapshot_state_dict, snapshot_path) + del snapshot_state_dict # All ranks wait for master rank to finish conversion torch.distributed.barrier() @@ -329,10 +778,13 @@ def load_dcp_from_hf(model: nn.Module, config: ModelConfig, parallel_dims: Paral state_dict, storage_reader=HuggingFaceStorageReader(path=snapshot_path.as_posix()), ) - if isinstance(model, PreTrainedModelPrimeRL): - model.init_buffers_post_meta() - else: - fix_model_post_empty(model) + # Restore weight tying broken by to_empty() for HF models + if not isinstance(model, PreTrainedModelPrimeRL) and model.config.tie_word_embeddings: + model.tie_weights() + _init_buffers_post_meta() + + _move_buffers_to_cuda(model, config) + lora_modules = [m for m in model.modules() if hasattr(m, "_init_lora_parameters")] if lora_modules: generator: torch.Generator | None = None @@ -355,6 +807,10 @@ def can_reinit_empty_buffers(model: nn.Module): The main issue is with anything that is not in the checkpoint. This is usually any non-persistent buffers. """ + # Custom PrimeRL models handle buffer reinit via init_buffers_post_meta + if isinstance(model, PreTrainedModelPrimeRL): + return True + buffer_names = [name for name, _ in model.named_buffers()] # TT MoE buffers @@ -370,6 +826,16 @@ def can_reinit_empty_buffers(model: nn.Module): if len(buffer_names) == 1 and buffer_names[0] == "model.rotary_emb.inv_freq": return True + # GPT-OSS (has original_inv_freq alongside inv_freq from dynamic rope scaling) + gpt_oss_buffers = {"model.rotary_emb.inv_freq", "model.rotary_emb.original_inv_freq"} + if set(buffer_names) == gpt_oss_buffers: + return True + + # Gemma3 model (has embed_scale and local rotary emb) + gemma3_buffers = {"model.embed_tokens.embed_scale", "model.rotary_emb.inv_freq", "model.rotary_emb_local.inv_freq"} + if set(buffer_names) == gemma3_buffers: + return True + get_logger().warning(f"Model cannot be loaded using meta device because of buffers: {buffer_names}") return False @@ -379,8 +845,32 @@ def fix_model_post_empty(model: nn.Module): # HF standard transformer model if "model.rotary_emb.inv_freq" in buffer_names: rotary_emb = model.model.rotary_emb - inv_freq, rotary_emb.attention_scaling = rotary_emb.rope_init_fn(rotary_emb.config, rotary_emb.inv_freq.device) + if hasattr(rotary_emb, "rope_init_fn"): + rope_init_fn = rotary_emb.rope_init_fn + else: + # GPT-OSS stores rope_init_fn only as a local in __init__; re-derive it + from transformers.modeling_rope_utils import ROPE_INIT_FUNCTIONS + + rope_init_fn = ( + ROPE_INIT_FUNCTIONS[rotary_emb.rope_type] + if rotary_emb.rope_type != "default" + else rotary_emb.compute_default_rope_parameters + ) + inv_freq, rotary_emb.attention_scaling = rope_init_fn(rotary_emb.config, rotary_emb.inv_freq.device) rotary_emb.inv_freq.copy_(inv_freq) + if "model.rotary_emb.original_inv_freq" in buffer_names: + rotary_emb.original_inv_freq.copy_(inv_freq) + # Gemma3 local rotary emb + if "model.rotary_emb_local.inv_freq" in buffer_names: + rotary_emb_local = model.model.rotary_emb_local + inv_freq_local, rotary_emb_local.attention_scaling = rotary_emb_local.rope_init_fn( + rotary_emb_local.config, rotary_emb_local.inv_freq.device + ) + rotary_emb_local.inv_freq.copy_(inv_freq_local) + # Gemma3 embed_scale (scalar computed from hidden_size) + if "model.embed_tokens.embed_scale" in buffer_names: + embed_scale = model.config.hidden_size**0.5 + model.model.embed_tokens.embed_scale.fill_(embed_scale) def reshard_module(model: nn.Module): @@ -390,42 +880,148 @@ def reshard_module(model: nn.Module): def apply_ac(model: nn.Module, ac_config: ActivationCheckpointConfig): - for layer_id, (layer_name, transformer_block) in enumerate(model.model.layers.named_children()): - if layer_id % ac_config.freq == 0: + logger = get_logger() + language_model = get_language_model(model) + target_list = sorted(frozenset(ac_config.targets)) + selective_layers = 0 + full_layers = 0 + fallback_layer_types: set[str] = set() + model_supported_targets: set[str] = set() + + for layer_id, (layer_name, transformer_block) in enumerate(language_model.layers.named_children()): + if layer_id % ac_config.freq != 0: + continue + + if ac_config.mode == "selective" and supports_selective_activation_checkpointing(transformer_block): + model_supported_targets.update(get_supported_targets(transformer_block)) + set_selective_activation_checkpointing(transformer_block, target_list) + selective_layers += 1 + else: + if ac_config.mode == "selective": + fallback_layer_types.add(type(transformer_block).__name__) transformer_block = checkpoint_wrapper(transformer_block, preserve_rng_state=False) - model.model.layers.register_module(layer_name, transformer_block) - get_logger().info(f"Applied activation checkpointing (freq={ac_config.freq})") + full_layers += 1 + + language_model.layers.register_module(layer_name, transformer_block) + + if ac_config.mode == "selective": + unsupported_targets = frozenset(target_list) - model_supported_targets + if unsupported_targets: + raise ValueError( + f"Selective activation checkpoint targets {sorted(unsupported_targets)} are not supported " + f"by the selected model layers. Supported targets across the model: {sorted(model_supported_targets)}" + ) + if fallback_layer_types: + logger.warning( + "Selective activation checkpointing is not supported for layer types " + f"{sorted(fallback_layer_types)}; falling back to full checkpointing for those layers." + ) + logger.info( + "Applied selective activation checkpointing " + f"(freq={ac_config.freq}, targets={target_list}, selective_layers={selective_layers}, " + f"full_fallback_layers={full_layers})" + ) + return + + logger.info(f"Applied activation checkpointing (freq={ac_config.freq})") def apply_compile(model: nn.Module, compile_config: CompileConfig): torch._dynamo.config.capture_scalar_outputs = True - for layer_id in range(len(model.model.layers)): + language_model = get_language_model(model) + for layer_id in range(len(language_model.layers)): # Doing it in-place avoids mangled fqn which can break checkpoint loading - model.model.layers[layer_id].compile(fullgraph=compile_config.fullgraph) - get_logger().info(f"Compiled {len(model.model.layers)} layers (fullgraph={compile_config.fullgraph})") + language_model.layers[layer_id].compile(fullgraph=compile_config.fullgraph) + get_logger().info(f"Compiled {len(language_model.layers)} layers (fullgraph={compile_config.fullgraph})") -def apply_ep(model: nn.Module, parallel_dims: ParallelDims): - for transformer_block in model.model.layers: - if isinstance(transformer_block.mlp, MoE): +def apply_ep(model: nn.Module, config: ModelConfig, parallel_dims: ParallelDims): + language_model = get_language_model(model) + for transformer_block in language_model.layers: + block_mlp = getattr(transformer_block, "mlp", None) + if block_mlp is not None and isinstance(block_mlp, (MoE, LatentMoE)): + if config.ep_comm_backend == "torch": + parallelize_plan = ExpertParallel() + else: + parallelize_plan = DeepEPExpertParallel() parallelize_module( - transformer_block.mlp.experts, - device_mesh=parallel_dims.world_mesh["ep"], - parallelize_plan=ExpertParallel(), + block_mlp.experts, + device_mesh=parallel_dims.get_mesh("ep"), + parallelize_plan=parallelize_plan, ) +def _move_buffers_to_cuda(model: nn.Module, config: ModelConfig) -> None: + """FSDP CPU offloading only manages parameters, not buffers. Move buffers to CUDA.""" + if not config.fsdp_cpu_offload: + return + for _, buffer in model.named_buffers(): + if buffer.device.type == "cpu": + buffer.data = buffer.data.to("cuda") + + +def _reset_runtime_moe_buffers(model: nn.Module) -> None: + for module in model.modules(): + if isinstance(module, (MoE, LatentMoE)) and module.tokens_per_expert.device.type != "meta": + module.tokens_per_expert.zero_() + + +def _validate_flash_attn_4_installed() -> None: + """Validate that flash-attn-cute is installed and not overwritten by flash-attn. + + Both flash-attn and flash-attn-cute ship a `flash_attn.cute` sub-package. + When both extras are installed, the older stub from flash-attn can shadow the + real implementation. We detect this by checking the line count of the interface + module (the real one is >1000 lines). + """ + import flash_attn.cute.interface as fa4_interface + + with open(fa4_interface.__file__, "r") as f: + num_lines = sum(1 for _ in f) + + if num_lines < 1000: + raise ValueError( + "flash-attn-cute has probably been overwritten by flash-attn, " + "run `scripts/fix-flash-attn-cute.sh` to fix this behaviour." + ) + + +def _register_fa4_attention_interface() -> None: + """Register a dummy `fa4` attention with transformers so AutoConfig accepts it. + + The `flash_attention_*` naming pattern triggers transformers to attempt + installing a kernel from the hub, so we use the short name `fa4` internally. + This dummy is never called because fa4 is only supported with our custom + model implementation. + """ + from transformers import AttentionInterface + + def _noop(*args, **kwargs) -> None: + pass + + AttentionInterface.register("fa4", _noop) + + def setup_model( - config: ModelConfig, parallel_dims: ParallelDims, loading_from_checkpoint_later: bool = False + config: ModelConfig, + parallel_dims: ParallelDims, + loading_from_checkpoint_later: bool = False, + fused_cross_entropy: bool | str = False, ) -> nn.Module: if config.attn == "flash_attention_3" and not is_flash_attn_3_available(): raise ValueError( "Flash attention 3 is only supported if the flash_attn_3 package is installed. Install with `uv pip install 'flash-attn-3 @ git+https://github.com/Dao-AILab/flash-attention.git@main#subdirectory=hopper' --no-build-isolation`" ) + + if config.attn == "fa4": + _validate_flash_attn_4_installed() + _register_fa4_attention_interface() + logger = get_logger() # 1. We load to meta device by default model = get_model(config, device=torch.device("meta"), dtype=DTYPE_MAP[config.optimization_dtype]) + configure_moe_ep_backend(model, config) possible_to_load_to_meta = can_reinit_empty_buffers(model) @@ -438,15 +1034,33 @@ def setup_model( if not possible_to_load_to_meta: logger.warning("Cannot load model to meta device only, loading to CPU instead.") model = get_model(config, device=torch.device("cpu"), dtype=DTYPE_MAP[config.optimization_dtype]) + configure_moe_ep_backend(model, config) + + lm_head_chunk_size: int | None = None + if isinstance(config.fused_lm_head_token_chunk_size, int): + lm_head_chunk_size = config.fused_lm_head_token_chunk_size - inject_prime_lm_head(model, chunk_size=config.fused_lm_head_chunk_size) + inject_prime_lm_head(model, chunk_size=lm_head_chunk_size, fused_cross_entropy=fused_cross_entropy) + + if config.fp8: + replace_linear_with_fp8_blockwise_linear(model) # Apply LoRA before FSDP setup if config.lora is not None: apply_lora_to_model(model, config.lora) + if config.freeze_moe_router: + freeze_moe_router(model) + + if config.debug.force_balanced_routing: + apply_force_balanced_routing(model) + if parallel_dims.ep_enabled: - apply_ep(model, parallel_dims) + apply_ep(model, config, parallel_dims) + # EP replaces params with DTensors that default to requires_grad=True, + # re-freeze base params that LoRA froze earlier. + if config.lora is not None: + freeze_all_except_lora_and_specified(model, config.lora) # the right order is AC -> Compile -> FSDP if config.ac is not None: @@ -456,6 +1070,9 @@ def setup_model( setup_fsdp(model, config, parallel_dims) + if not possible_to_load_to_meta: + _move_buffers_to_cuda(model, config) + # 2. if we can load to meta, we either: if possible_to_load_to_meta: # - load from checkpoint later if needed @@ -463,29 +1080,80 @@ def setup_model( logger.warning( "Skipping loading weights. Initializing an empty model on device, loading from checkpoint later." ) - model.to_empty(device="cuda") + device = "cpu" if config.fsdp_cpu_offload else "cuda" + model.to_empty(device=device) torch.distributed.barrier() if isinstance(model, PreTrainedModelPrimeRL): model.init_buffers_post_meta() else: fix_model_post_empty(model) + # Restore weight tying broken by to_empty() for HF models + if model.config.tie_word_embeddings: + model.tie_weights() + + _move_buffers_to_cuda(model, config) # - or load from HF with dcp else: load_dcp_from_hf(model, config, parallel_dims) - logger.debug(f"Model signature: {get_module_signature(model, compress=True)}") + _reset_runtime_moe_buffers(model) return model +def _get_qwen3_vl_mm_token_type_ids(model: nn.Module, input_ids: Tensor) -> Tensor | None: + config = getattr(model, "config", None) + if getattr(config, "model_type", None) != "qwen3_vl": + return None + + mm_token_type_ids = torch.zeros_like(input_ids) + + image_token_id = getattr(config, "image_token_id", None) + if image_token_id is not None: + mm_token_type_ids = mm_token_type_ids.masked_fill(input_ids == image_token_id, 1) + + video_token_id = getattr(config, "video_token_id", None) + if video_token_id is not None: + mm_token_type_ids = mm_token_type_ids.masked_fill(input_ids == video_token_id, 2) + + return mm_token_type_ids + + @jaxtyped(typechecker=typechecker) def forward( model: nn.Module, input_ids: Int[Tensor, "batch seq"], position_ids: Int[Tensor, "batch seq"], labels: Int[Tensor, "batch seq"] | None = None, - temperature: float | None = None, + temperature: Tensor | None = None, + routed_experts: Int[Tensor, "batch seq layers topk"] | None = None, + # Multimodal fields (Qwen3-VL) + pixel_values: Float[Tensor, "num_patches patch_dim"] | None = None, + image_grid_thw: Int[Tensor, "num_images 3"] | None = None, + mm_token_type_ids: Int[Tensor, "batch seq"] | None = None, ) -> PrimeLmOutput: - out = model(input_ids=input_ids, position_ids=position_ids, labels=labels, temperature=temperature) + # Build kwargs for model forward + kwargs = { + "input_ids": input_ids, + "labels": labels, + "temperature": temperature, + } + + # For multimodal (VLM), don't pass position_ids - let the model compute MRoPE internally + # using image_grid_thw. Qwen3-VL only computes proper MRoPE when position_ids is None. + if pixel_values is not None: + assert image_grid_thw is not None, "pixel_values requires image_grid_thw for MRoPE computation" + kwargs["pixel_values"] = pixel_values + kwargs["image_grid_thw"] = image_grid_thw + mm_token_type_ids = _get_qwen3_vl_mm_token_type_ids(model, input_ids) + if mm_token_type_ids is not None: + kwargs["mm_token_type_ids"] = mm_token_type_ids + else: + kwargs["position_ids"] = position_ids + + if routed_experts is not None: + kwargs["routed_experts"] = routed_experts + + out = model(**kwargs) # PrimeLmOutput is a TypedDict (dict at runtime), HF outputs are dataclass-like objects if isinstance(out, dict): diff --git a/src/prime_rl/trainer/models/__init__.py b/src/prime_rl/trainer/models/__init__.py index e1049a9c1de..7e4d27fc331 100644 --- a/src/prime_rl/trainer/models/__init__.py +++ b/src/prime_rl/trainer/models/__init__.py @@ -11,20 +11,35 @@ from prime_rl.trainer.models.afmoe import AfmoeConfig, AfmoeForCausalLM from prime_rl.trainer.models.base import PreTrainedModelPrimeRL from prime_rl.trainer.models.glm4_moe import Glm4MoeConfig, Glm4MoeForCausalLM +from prime_rl.trainer.models.glm_moe_dsa import GlmMoeDsaConfig, GlmMoeDsaForCausalLM +from prime_rl.trainer.models.gpt_oss import GptOssConfig, GptOssForCausalLM from prime_rl.trainer.models.layers.lm_head import PrimeLmOutput, cast_float_and_contiguous from prime_rl.trainer.models.llama import LlamaForCausalLM +from prime_rl.trainer.models.minimax_m2 import MiniMaxM2Config, MiniMaxM2ForCausalLM +from prime_rl.trainer.models.nemotron_h import NemotronHConfig, NemotronHForCausalLM +from prime_rl.trainer.models.qwen3_5_moe import Qwen3_5MoeConfig, Qwen3_5MoeForCausalLM from prime_rl.trainer.models.qwen3_moe import Qwen3MoeConfig, Qwen3MoeForCausalLM # Make custom config discoverable by AutoConfig AutoConfig.register("afmoe", AfmoeConfig, exist_ok=True) AutoConfig.register("glm4_moe", Glm4MoeConfig, exist_ok=True) +AutoConfig.register("glm_moe_dsa", GlmMoeDsaConfig, exist_ok=True) +AutoConfig.register("minimax_m2", MiniMaxM2Config, exist_ok=True) +AutoConfig.register("nemotron_h", NemotronHConfig, exist_ok=True) AutoConfig.register("qwen3_moe", Qwen3MoeConfig, exist_ok=True) +AutoConfig.register("qwen3_5_moe_text", Qwen3_5MoeConfig, exist_ok=True) +# GptOssConfig is just HF's class - already registered by transformers, no override needed. _CUSTOM_CAUSAL_LM_MAPPING = _LazyAutoMapping(CONFIG_MAPPING_NAMES, OrderedDict()) _CUSTOM_CAUSAL_LM_MAPPING.register(LlamaConfig, LlamaForCausalLM, exist_ok=True) _CUSTOM_CAUSAL_LM_MAPPING.register(AfmoeConfig, AfmoeForCausalLM, exist_ok=True) _CUSTOM_CAUSAL_LM_MAPPING.register(Glm4MoeConfig, Glm4MoeForCausalLM, exist_ok=True) +_CUSTOM_CAUSAL_LM_MAPPING.register(GlmMoeDsaConfig, GlmMoeDsaForCausalLM, exist_ok=True) +_CUSTOM_CAUSAL_LM_MAPPING.register(MiniMaxM2Config, MiniMaxM2ForCausalLM, exist_ok=True) +_CUSTOM_CAUSAL_LM_MAPPING.register(NemotronHConfig, NemotronHForCausalLM, exist_ok=True) _CUSTOM_CAUSAL_LM_MAPPING.register(Qwen3MoeConfig, Qwen3MoeForCausalLM, exist_ok=True) +_CUSTOM_CAUSAL_LM_MAPPING.register(Qwen3_5MoeConfig, Qwen3_5MoeForCausalLM, exist_ok=True) +_CUSTOM_CAUSAL_LM_MAPPING.register(GptOssConfig, GptOssForCausalLM, exist_ok=True) class AutoModelForCausalLMPrimeRL(_BaseAutoModelClass): @@ -46,4 +61,24 @@ def supports_custom_impl(model_config: PretrainedConfig) -> bool: return type(model_config) in _CUSTOM_CAUSAL_LM_MAPPING -__all__ = ["AutoModelForCausalLMPrimeRL", "PreTrainedModelPrimeRL", "supports_custom_impl", "PrimeLmOutput", "cast_float_and_contiguous"] +# Mapping from HF composite VLM model_type to custom PrimeRL class. +# Used by get_model() to dispatch VLMs that have a custom text model implementation. +# Points to the same unified class — the config drives text-only vs VLM behavior. +_CUSTOM_VLM_MAPPING: dict[str, type] = { + "qwen3_5_moe": Qwen3_5MoeForCausalLM, +} + + +def get_custom_vlm_cls(model_config: PretrainedConfig) -> type | None: + """Return the custom PrimeRL VLM class for this config, or None if unsupported.""" + return _CUSTOM_VLM_MAPPING.get(getattr(model_config, "model_type", None)) + + +__all__ = [ + "AutoModelForCausalLMPrimeRL", + "PreTrainedModelPrimeRL", + "supports_custom_impl", + "get_custom_vlm_cls", + "PrimeLmOutput", + "cast_float_and_contiguous", +] diff --git a/src/prime_rl/trainer/models/afmoe/__init__.py b/src/prime_rl/trainer/models/afmoe/__init__.py index 4715f9a0f86..eadc8e48a9c 100644 --- a/src/prime_rl/trainer/models/afmoe/__init__.py +++ b/src/prime_rl/trainer/models/afmoe/__init__.py @@ -11,4 +11,3 @@ "AfmoeModel", "AfmoePreTrainedModel", ] - diff --git a/src/prime_rl/trainer/models/afmoe/configuration_afmoe.py b/src/prime_rl/trainer/models/afmoe/configuration_afmoe.py index 4bbc7db892e..24e8013722b 100644 --- a/src/prime_rl/trainer/models/afmoe/configuration_afmoe.py +++ b/src/prime_rl/trainer/models/afmoe/configuration_afmoe.py @@ -1,12 +1,14 @@ from transformers.configuration_utils import PretrainedConfig, layer_type_validation -from transformers.modeling_rope_utils import rope_config_validation from transformers.utils import logging logger = logging.get_logger(__name__) + class AfmoeConfig(PretrainedConfig): """Configuration for AFMoE.""" + model_type = "afmoe" + attribute_map = {"num_local_experts": "num_experts"} base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), "layers": (["hidden_states", "attention_mask"], ["hidden_states"]), @@ -88,7 +90,8 @@ def __init__( self.layer_types = layer_types if self.layer_types is None: self.layer_types = [ - "sliding_attention" if bool((i + 1) % global_attn_every_n_layers) else "full_attention" for i in range(self.num_hidden_layers) + "sliding_attention" if bool((i + 1) % global_attn_every_n_layers) else "full_attention" + for i in range(self.num_hidden_layers) ] layer_type_validation(self.layer_types) @@ -102,7 +105,7 @@ def __init__( # Validate rope configs if self.rope_scaling is not None and "type" in self.rope_scaling: self.rope_scaling["rope_type"] = self.rope_scaling["type"] - rope_config_validation(self) + self.standardize_rope_params() super().__init__( tie_word_embeddings=tie_word_embeddings, diff --git a/src/prime_rl/trainer/models/afmoe/converting_afmoe.py b/src/prime_rl/trainer/models/afmoe/converting_afmoe.py index aee5e23a54b..2a82b9f4452 100644 --- a/src/prime_rl/trainer/models/afmoe/converting_afmoe.py +++ b/src/prime_rl/trainer/models/afmoe/converting_afmoe.py @@ -27,22 +27,22 @@ def is_moe_layer_tt(state_dict: dict[str, Tensor], layer_idx: int) -> bool: def convert_hf_layer_to_tt(state_dict: dict[str, Tensor], layer_idx: int) -> None: """Convert a single MoE layer from HF format to TT format in-place. - + Args: state_dict: The state dict to modify in-place. layer_idx: The layer index to convert. """ prefix = f"model.layers.{layer_idx}.mlp" - + # Check if this is a MoE layer (has individual experts) is_moe_layer = is_moe_layer_hf(state_dict, layer_idx) if not is_moe_layer: return # Dense layer, no conversion needed - + num_experts = get_num_experts_from_state_dict(state_dict, layer_idx) if num_experts == 0: return - + # Convert shared experts: shared_experts -> shared_expert, gate_proj/up_proj/down_proj -> w1/w3/w2 shared_mappings = [ (f"{prefix}.shared_experts.gate_proj.weight", f"{prefix}.shared_expert.w1"), @@ -61,14 +61,14 @@ def convert_hf_layer_to_tt(state_dict: dict[str, Tensor], layer_idx: int) -> Non ("down_proj", "w2"), ("up_proj", "w3"), ] - + for hf_name, tt_name in expert_mappings: expert_weights = [] for i in range(num_experts): hf_key = f"{prefix}.experts.{i}.{hf_name}.weight" if hf_key in state_dict: expert_weights.append(state_dict.pop(hf_key)) - + if expert_weights: # Stack along first dimension: [num_experts, out_dim, in_dim] stacked = torch.stack(expert_weights, dim=0) @@ -77,17 +77,17 @@ def convert_hf_layer_to_tt(state_dict: dict[str, Tensor], layer_idx: int) -> Non def convert_tt_layer_to_hf(state_dict: dict[str, Tensor], layer_idx: int) -> None: """Convert a single MoE layer from TT format to HF format in-place. - + Args: state_dict: The state dict to modify in-place. layer_idx: The layer index to convert. """ prefix = f"model.layers.{layer_idx}.mlp" - + # Check if this is a MoE layer in TT format if not is_moe_layer_tt(state_dict, layer_idx): return # Dense layer, no conversion needed - + # Convert shared expert: shared_expert -> shared_experts, w1/w2/w3 -> gate_proj/down_proj/up_proj shared_mappings = [ (f"{prefix}.shared_expert.w1", f"{prefix}.shared_experts.gate_proj.weight"), @@ -97,14 +97,14 @@ def convert_tt_layer_to_hf(state_dict: dict[str, Tensor], layer_idx: int) -> Non for tt_key, hf_key in shared_mappings: if tt_key in state_dict: state_dict[hf_key] = state_dict.pop(tt_key) - + # Unstack grouped expert weights into individual experts expert_mappings = [ ("w1", "gate_proj"), ("w2", "down_proj"), ("w3", "up_proj"), ] - + for tt_name, hf_name in expert_mappings: tt_key = f"{prefix}.experts.{tt_name}" if tt_key in state_dict: @@ -113,7 +113,7 @@ def convert_tt_layer_to_hf(state_dict: dict[str, Tensor], layer_idx: int) -> Non for i in range(num_experts): hf_key = f"{prefix}.experts.{i}.{hf_name}.weight" state_dict[hf_key] = stacked[i] - + # Remove TT-specific buffers that don't exist in HF for key in list(state_dict.keys()): if f"{prefix}.tokens_per_expert" in key: diff --git a/src/prime_rl/trainer/models/afmoe/modeling_afmoe.py b/src/prime_rl/trainer/models/afmoe/modeling_afmoe.py index cf75dcee015..4bdef935782 100644 --- a/src/prime_rl/trainer/models/afmoe/modeling_afmoe.py +++ b/src/prime_rl/trainer/models/afmoe/modeling_afmoe.py @@ -1,3 +1,4 @@ +import functools from dataclasses import dataclass from typing import Optional, Union @@ -21,12 +22,28 @@ from prime_rl.trainer.models.layers.lm_head import PrimeLmOutput from prime_rl.trainer.models.layers.mlp import MLP, MLPConfig from prime_rl.trainer.models.layers.moe import MoE, MoEArgs -from prime_rl.trainer.models.layers.rms_norm import RMSNorm, RMSNormConfig +from prime_rl.trainer.models.layers.norms import RMSNorm, RMSNormConfig from prime_rl.trainer.models.layers.rotary_emb import ( RotaryEmbedding, RotaryEmbeddingConfig, apply_rotary_pos_emb, ) +from prime_rl.utils.sequence import get_cu_seqlens_from_position_ids + +try: + from flash_attn import flash_attn_varlen_func +except ImportError: + flash_attn_varlen_func = None # type: ignore + +try: + from flash_attn_interface import flash_attn_varlen_func as flash_attn_3_varlen_func +except ImportError: + flash_attn_3_varlen_func = None # type: ignore + +try: + from flash_attn.cute import flash_attn_varlen_func as flash_attn_4_varlen_func +except ImportError: + flash_attn_4_varlen_func = None # type: ignore from .configuration_afmoe import AfmoeConfig from .converting_afmoe import ( @@ -85,11 +102,27 @@ def __init__(self, config: AfmoeAttentionConfig): self.q_norm = RMSNorm(RMSNormConfig(hidden_size=self.head_dim, eps=config.rms_norm_eps)) self.k_norm = RMSNorm(RMSNormConfig(hidden_size=self.head_dim, eps=config.rms_norm_eps)) - def _project_states( + def output_proj( + self, + attn_output: torch.Tensor, + gate_states: torch.Tensor, + ) -> torch.Tensor: + input_shape = gate_states.shape[:-1] + if attn_output.dim() == 4: + attn_output = attn_output.transpose(1, 2).contiguous() + attn_output = attn_output.contiguous().view(*input_shape, -1) + attn_output = attn_output * torch.sigmoid(gate_states) + return self.o_proj(attn_output) + + +class AfmoeSDPAAttention(AfmoeAttentionBase): + """AFMoE attention using PyTorch's scaled_dot_product_attention.""" + + def attn_projections( self, hidden_states: torch.Tensor, position_embeddings: tuple[torch.Tensor, torch.Tensor], - ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, tuple[int, ...]]: + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: input_shape = hidden_states.shape[:-1] hidden_shape = (*input_shape, -1, self.head_dim) @@ -112,50 +145,138 @@ def _project_states( key_states = _repeat_kv(key_states, self.num_key_value_groups) value_states = _repeat_kv(value_states, self.num_key_value_groups) - return query_states, key_states, value_states, gate_states, input_shape + return query_states, key_states, value_states, gate_states - def _finalize_output( + def _attention_core( self, - attn_output: torch.Tensor, - gate_states: torch.Tensor, - input_shape: tuple[int, ...], - ) -> tuple[torch.Tensor, None]: - attn_output = attn_output.transpose(1, 2).contiguous() - attn_output = attn_output.view(*input_shape, -1) - attn_output = attn_output * torch.sigmoid(gate_states) - attn_output = self.o_proj(attn_output) - return attn_output, None - - -class AfmoeSDPAAttention(AfmoeAttentionBase): - """AFMoE attention using PyTorch's scaled_dot_product_attention.""" + query_states: torch.Tensor, + key_states: torch.Tensor, + value_states: torch.Tensor, + attention_mask: torch.Tensor | None = None, + dropout_p: float = 0.0, + ) -> torch.Tensor: + return F.scaled_dot_product_attention( + query_states, + key_states, + value_states, + attn_mask=attention_mask, + dropout_p=dropout_p, + is_causal=attention_mask is None, + scale=self.scaling, + ) def forward( self, hidden_states: torch.Tensor, position_embeddings: tuple[torch.Tensor, torch.Tensor], attention_mask: torch.Tensor | None = None, + cu_seqlens: torch.LongTensor | None = None, + max_seqlen: int | None = None, ) -> tuple[torch.Tensor, None]: - query_states, key_states, value_states, gate_states, input_shape = self._project_states( - hidden_states, position_embeddings - ) + query_states, key_states, value_states, gate_states = self.attn_projections(hidden_states, position_embeddings) dropout_p = self.attention_dropout if self.training else 0.0 - attn_output = F.scaled_dot_product_attention( + attn_output = self._attention_core( query_states, key_states, value_states, - attn_mask=attention_mask, + attention_mask=attention_mask, dropout_p=dropout_p, - is_causal=attention_mask is None, # Use causal if no explicit mask - scale=self.scaling, ) - return self._finalize_output(attn_output, gate_states, input_shape) + return self.output_proj(attn_output, gate_states), None + + +class AfmoeFlashAttention(AfmoeAttentionBase): + """AFMoE attention using Flash Attention varlen functions.""" + + _funcs = { + 2: flash_attn_varlen_func, + 3: flash_attn_3_varlen_func, + 4: flash_attn_4_varlen_func, + } + + def __init__(self, config: AfmoeAttentionConfig, flash_attn_version: int = 4): + super().__init__(config) + self._flash_attn_version = flash_attn_version + self.func = self._funcs[flash_attn_version] + self._flash_attn_call = self.func + if self._flash_attn_version == 4: + self._flash_attn_call = torch._dynamo.disable(self.func) + + def _compute_attention(self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, cu_seqlens, max_seqlen): + """Run the flash attention kernel. q/k/v are [total_tokens, heads, dim].""" + args = [q, k, v, cu_seqlens, cu_seqlens] + if self._flash_attn_version != 4: + args.extend([max_seqlen, max_seqlen]) + kwargs: dict = {"causal": True} + if self.sliding_window is not None: + kwargs["window_size"] = (self.sliding_window - 1, 0) + out = self._flash_attn_call(*args, **kwargs) + if isinstance(out, tuple): + out = out[0] + return out + + def _attention_core( + self, + query_states: torch.Tensor, + key_states: torch.Tensor, + value_states: torch.Tensor, + cu_seqlens: torch.LongTensor | None = None, + max_seqlen: int | None = None, + ) -> torch.Tensor: + return self._compute_attention(query_states[0], key_states[0], value_states[0], cu_seqlens, max_seqlen) + + def attn_projections( + self, + hidden_states: torch.Tensor, + position_embeddings: tuple[torch.Tensor, torch.Tensor], + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + input_shape = hidden_states.shape[:-1] + hidden_shape = (*input_shape, -1, self.head_dim) + + query_states = self.q_proj(hidden_states).view(hidden_shape) + key_states = self.k_proj(hidden_states).view(hidden_shape) + value_states = self.v_proj(hidden_states).view(hidden_shape) + gate_states = self.gate_proj(hidden_states) + + query_states = self.q_norm(query_states) + key_states = self.k_norm(key_states) + + if self.is_local_attention: + query_states = query_states.transpose(1, 2) + key_states = key_states.transpose(1, 2) + cos, sin = position_embeddings + query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin) + query_states = query_states.transpose(1, 2) + key_states = key_states.transpose(1, 2) + + return query_states, key_states, value_states, gate_states + + def forward( + self, + hidden_states: torch.Tensor, + position_embeddings: tuple[torch.Tensor, torch.Tensor], + attention_mask: torch.Tensor | None = None, + cu_seqlens: torch.LongTensor | None = None, + max_seqlen: int | None = None, + ) -> tuple[torch.Tensor, None]: + query_states, key_states, value_states, gate_states = self.attn_projections(hidden_states, position_embeddings) + attn_output = self._attention_core( + query_states, + key_states, + value_states, + cu_seqlens=cu_seqlens, + max_seqlen=max_seqlen, + ) + return self.output_proj(attn_output, gate_states), None AFMOE_ATTN_IMPL2CLASS = { "sdpa": AfmoeSDPAAttention, + "flash_attention_2": functools.partial(AfmoeFlashAttention, flash_attn_version=2), + "flash_attention_3": functools.partial(AfmoeFlashAttention, flash_attn_version=3), + "fa4": functools.partial(AfmoeFlashAttention, flash_attn_version=4), } @@ -194,9 +315,7 @@ def _get_afmoe_attention(config: AfmoeConfig, layer_idx: int) -> nn.Module: if attn_impl not in AFMOE_ATTN_IMPL2CLASS: supported = list(AFMOE_ATTN_IMPL2CLASS.keys()) raise ValueError( - f"AFMoE attention does not support '{config._attn_implementation}'. " - f"Supported implementations: {supported}. " - f"Note: flash_attention is not supported for AFMoE due to sliding window + RoPE constraints." + f"AFMoE attention does not support '{config._attn_implementation}'. Supported implementations: {supported}." ) return AFMOE_ATTN_IMPL2CLASS[attn_impl](attn_config) @@ -234,6 +353,7 @@ def __init__(self, config: AfmoeConfig, layer_idx: int): top_k=config.num_experts_per_tok, use_grouped_mm=getattr(config, "use_grouped_mm", True), load_balance_coeff=getattr(config, "load_balance_coeff", None), + fp8=getattr(config, "fp8", False), ) if self.moe_enabled: self.mlp = MoE(moe_args, dim=config.hidden_size, hidden_dim=config.moe_intermediate_size) @@ -245,6 +365,9 @@ def forward( hidden_states: torch.Tensor, attention_mask: Optional[torch.Tensor] = None, position_embeddings: Optional[tuple[torch.Tensor, torch.Tensor]] = None, + cu_seqlens: torch.LongTensor | None = None, + max_seqlen: int | None = None, + routed_experts: Optional[torch.LongTensor] = None, ) -> torch.FloatTensor: residual = hidden_states @@ -253,13 +376,15 @@ def forward( hidden_states=hidden_states, position_embeddings=position_embeddings, attention_mask=attention_mask, + cu_seqlens=cu_seqlens, + max_seqlen=max_seqlen, ) hidden_states = self.post_attention_layernorm(hidden_states) hidden_states = residual + hidden_states residual = hidden_states hidden_states = self.pre_mlp_layernorm(hidden_states) - hidden_states = self.mlp(hidden_states) + hidden_states = self.mlp(hidden_states, routed_experts=routed_experts) hidden_states = self.post_mlp_layernorm(hidden_states) hidden_states = residual + hidden_states return hidden_states @@ -280,7 +405,7 @@ class AfmoePreTrainedModel(PreTrainedModelPrimeRL): "norm", ] _supports_sdpa = True - _supports_flash_attn = False + _supports_flash_attn = True _supports_flex_attn = False _supports_attention_backend = True _can_compile_fullgraph = False @@ -348,7 +473,12 @@ def forward( attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, inputs_embeds: Optional[torch.FloatTensor] = None, + routed_experts: Optional[torch.LongTensor] = None, ) -> MoeModelOutputWithPast: + """ + routed_experts (`torch.LongTensor` of shape `(batch_size, sequence_length, num_hidden_layers, num_experts_per_tok)`, *optional*): + Routed experts for each token in the sequence. Only used for router replay. + """ if (input_ids is None) ^ (inputs_embeds is not None): raise ValueError("You must specify exactly one of input_ids or inputs_embeds") @@ -357,21 +487,30 @@ def forward( if position_ids is None: position_ids = torch.arange(inputs_embeds.shape[1], device=inputs_embeds.device).unsqueeze(0) - cache_position = torch.arange(inputs_embeds.shape[1], device=inputs_embeds.device) - - if not isinstance(causal_mask_mapping := attention_mask, dict): - mask_kwargs = { - "config": self.config, - "input_embeds": inputs_embeds, - "attention_mask": attention_mask, - "cache_position": cache_position, - "past_key_values": None, - "position_ids": position_ids, - } - causal_mask_mapping = { - "full_attention": create_causal_mask(**mask_kwargs), - "sliding_attention": create_sliding_window_causal_mask(**mask_kwargs), - } + + use_flash = self.config._attn_implementation in ("flash_attention_2", "flash_attention_3", "fa4") + + if use_flash: + cu_seqlens, max_seqlen = get_cu_seqlens_from_position_ids(position_ids) + torch._dynamo.mark_dynamic(cu_seqlens, 0) + causal_mask_mapping = None + else: + cu_seqlens = None + max_seqlen = None + cache_position = torch.arange(inputs_embeds.shape[1], device=inputs_embeds.device) + if not isinstance(causal_mask_mapping := attention_mask, dict): + mask_kwargs = { + "config": self.config, + "inputs_embeds": inputs_embeds, + "attention_mask": attention_mask, + "cache_position": cache_position, + "past_key_values": None, + "position_ids": position_ids, + } + causal_mask_mapping = { + "full_attention": create_causal_mask(**mask_kwargs), + "sliding_attention": create_sliding_window_causal_mask(**mask_kwargs), + } hidden_states = inputs_embeds @@ -380,13 +519,17 @@ def forward( position_embeddings = self.rotary_emb(hidden_states, position_ids) - for decoder_layer in self.layers: - mask = causal_mask_mapping[decoder_layer.attention_type] + for layer_idx, decoder_layer in enumerate(self.layers): + mask = causal_mask_mapping[decoder_layer.attention_type] if causal_mask_mapping is not None else None + routed_experts_layer = routed_experts[:, :, layer_idx, :] if routed_experts is not None else None hidden_states = decoder_layer( hidden_states, attention_mask=mask, position_embeddings=position_embeddings, + cu_seqlens=cu_seqlens, + max_seqlen=max_seqlen, + routed_experts=routed_experts_layer, ) hidden_states = self.norm(hidden_states) @@ -396,7 +539,7 @@ def forward( class AfmoeForCausalLM(AfmoePreTrainedModel, GenerationMixin): - _tied_weights_keys = ["lm_head.weight"] + _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": "colwise_rep"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} @@ -438,15 +581,18 @@ def forward( use_cache: Optional[bool] = None, logits_to_keep: Union[int, torch.Tensor] = 0, token_type_ids: Optional[torch.Tensor] = None, # will be ignored - temperature: Optional[float] = 1.0, + temperature: Optional[torch.Tensor] = None, + routed_experts: Optional[torch.LongTensor] = None, **kwargs: Unpack[TransformersKwargs], ) -> PrimeLmOutput: r""" labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): Labels used by PrimeRL's wrapped LM head to optionally compute per-token logprobs/entropy. If not provided, the wrapped LM head returns logits only. - temperature (`float`, *optional*, defaults to 1.0): - Temperature used for the logprobs/entropy computation when `labels` are provided. + temperature (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*): + Per-token temperatures for logprobs/entropy computation when `labels` are provided. + routed_experts (`torch.LongTensor` of shape `(batch_size, sequence_length, num_hidden_layers, num_experts_per_tok)`, *optional*): + Routed experts for each token in the sequence. Only used for router replay. """ assert use_cache is None, "use_cache is not supported for custom afmoe for now" assert past_key_values is None, "past_key_values is not supported for custom afmoe for now" @@ -456,6 +602,7 @@ def forward( attention_mask=attention_mask, position_ids=position_ids, inputs_embeds=inputs_embeds, + routed_experts=routed_experts, ) hidden_states = outputs.last_hidden_state diff --git a/src/prime_rl/trainer/models/base.py b/src/prime_rl/trainer/models/base.py index f5159d2385b..df669848477 100644 --- a/src/prime_rl/trainer/models/base.py +++ b/src/prime_rl/trainer/models/base.py @@ -6,11 +6,25 @@ class PreTrainedModelPrimeRL(PreTrainedModel): """ Base class for all PrimeRL models that extends HuggingFace PreTrainedModel. - This class provides a unified interface for state dict conversion between different - formats (e.g., HuggingFace format vs. training-optimized format) and buffer initialization + Provides a unified interface for state dict conversion between different formats + (e.g., HuggingFace format vs. training-optimized format) and buffer initialization after loading with meta device. """ + @classmethod + def from_config(cls, config, **kwargs): + """Public from_config that mirrors the Auto class API.""" + return cls._from_config(config, **kwargs) + + @classmethod + def _can_set_experts_implementation(cls) -> bool: + """PrimeRL models use custom MoE implementations and don't support dynamic experts implementation.""" + return False + + def get_correct_experts_implementation(self, requested_experts: str | None) -> str: + """PrimeRL models always use eager experts implementation.""" + return "eager" + @classmethod def is_hf_state_dict(cls, state_dict: dict[str, Tensor]) -> bool: """ @@ -89,6 +103,38 @@ def convert_layer_to_prime(cls, state_dict: dict[str, Tensor], layer_idx: int) - """ raise NotImplementedError(f"convert_layer_to_prime is not implemented for {cls.__name__}") + @classmethod + def convert_adapter_to_hf(cls, state_dict: dict[str, Tensor]) -> dict[str, Tensor]: + """ + Convert a LoRA adapter state dict from PrimeRL training format to HuggingFace format. + + Unlike convert_to_hf, this operates on a partial state dict containing only LoRA + adapter parameters (e.g. `model.layers.N...lora_A.weight`). Models + whose HF naming differs from PrimeRL naming at the submodule level (e.g. NemotronH's + unified `mixer` attribute) should override this to perform the rename. + + Implementations may mutate state_dict in-place or return a new dict; callers must + use the returned value. Default implementation is a no-op. + """ + return state_dict + + @classmethod + def convert_layer_to_vllm_kernel( + cls, + state_dict: dict[str, Tensor], + layer_idx: int, + quantize_fp8: bool = False, + ) -> dict[str, Tensor]: + """ + Convert a single layer's state dict from PrimeRL format to vLLM kernel format. + + Args: + state_dict: Layer weights in PrimeRL format. + layer_idx: Layer index to convert. + quantize_fp8: Whether to emit FP8 (e4m3) kernel weights with per-block scales. + """ + raise NotImplementedError(f"convert_layer_to_vllm_kernel is not implemented for {cls.__name__}") + def init_buffers_post_meta(self) -> None: """ Initialize buffers that are not in the state dict after loading with meta device. diff --git a/src/prime_rl/trainer/models/fp8.py b/src/prime_rl/trainer/models/fp8.py new file mode 100644 index 00000000000..dc3411cc777 --- /dev/null +++ b/src/prime_rl/trainer/models/fp8.py @@ -0,0 +1,39 @@ +import torch +from torch import Tensor + + +def quantize_to_fp8_blockwise(weight: Tensor, block_size: int = 128) -> tuple[Tensor, Tensor]: + """Quantize a 2D tensor to FP8 e4m3 with per-block scales.""" + if weight.ndim != 2: + raise ValueError(f"FP8 quantization expects a 2D tensor, got shape={tuple(weight.shape)}") + + rows, cols = weight.shape + pad_rows = (block_size - rows % block_size) % block_size + pad_cols = (block_size - cols % block_size) % block_size + + if pad_rows or pad_cols: + padded = torch.zeros( + rows + pad_rows, + cols + pad_cols, + dtype=weight.dtype, + device=weight.device, + ) + padded[:rows, :cols] = weight + else: + padded = weight.contiguous() + + padded_rows, padded_cols = padded.shape + blocks = padded.view( + padded_rows // block_size, + block_size, + padded_cols // block_size, + block_size, + ).permute(0, 2, 1, 3) + + fp8_max = torch.finfo(torch.float8_e4m3fn).max + max_abs = blocks.float().abs().amax(dim=(2, 3)) + scales = (max_abs / fp8_max).clamp(min=1e-12) + blocks_fp8 = (blocks.float() / scales[:, :, None, None]).clamp(-fp8_max, fp8_max).to(torch.float8_e4m3fn) + + quantized = blocks_fp8.permute(0, 2, 1, 3).reshape(padded_rows, padded_cols)[:rows, :cols].contiguous() + return quantized, scales.float().contiguous() diff --git a/src/prime_rl/trainer/models/glm4_moe/configuration_glm4_moe.py b/src/prime_rl/trainer/models/glm4_moe/configuration_glm4_moe.py index 132f66b3fe2..bbc5eafc336 100644 --- a/src/prime_rl/trainer/models/glm4_moe/configuration_glm4_moe.py +++ b/src/prime_rl/trainer/models/glm4_moe/configuration_glm4_moe.py @@ -1,7 +1,6 @@ import warnings from transformers.configuration_utils import PretrainedConfig -from transformers.modeling_rope_utils import rope_config_validation class Glm4MoeConfig(PretrainedConfig): @@ -121,6 +120,7 @@ class Glm4MoeConfig(PretrainedConfig): model_type = "glm4_moe" keys_to_ignore_at_inference = ["past_key_values"] + attribute_map = {"num_local_experts": "n_routed_experts"} # Default tensor parallel plan for base model `Glm4Moe` base_model_tp_plan = { @@ -171,6 +171,7 @@ def __init__( norm_topk_prob=True, use_qk_norm=False, use_grouped_mm=True, + pad_token_id=None, **kwargs, ): self.vocab_size = vocab_size @@ -190,11 +191,7 @@ def __init__( self.rope_scaling = rope_scaling self.attention_bias = attention_bias self.attention_dropout = attention_dropout - # Validate the correctness of rotary position embeddings parameters - # BC: if there is a 'type' field, move it to 'rope_type'. - if self.rope_scaling is not None and "type" in self.rope_scaling: - self.rope_scaling["rope_type"] = self.rope_scaling["type"] - rope_config_validation(self) + self.pad_token_id = pad_token_id # MoE arguments self.moe_intermediate_size = moe_intermediate_size diff --git a/src/prime_rl/trainer/models/glm4_moe/converting_glm4_moe.py b/src/prime_rl/trainer/models/glm4_moe/converting_glm4_moe.py index 56f96a159d1..85025821a11 100644 --- a/src/prime_rl/trainer/models/glm4_moe/converting_glm4_moe.py +++ b/src/prime_rl/trainer/models/glm4_moe/converting_glm4_moe.py @@ -9,31 +9,54 @@ def get_max_layer_num(state_dict: dict[str, Tensor]) -> int: def convert_hf_layer_to_tt(state_dict: dict[str, Tensor], layer_idx: int): i = layer_idx - num_experts = len([j for j in state_dict.keys() if f"model.layers.{i}.mlp.experts" in j]) // 3 - if num_experts == 0: + + # Check if this is a MoE layer by looking for gate weight + if f"model.layers.{i}.mlp.gate.weight" not in state_dict: return state_dict[f"model.layers.{i}.mlp.router.gate.weight"] = state_dict[f"model.layers.{i}.mlp.gate.weight"] del state_dict[f"model.layers.{i}.mlp.gate.weight"] - dim, moe_dim = state_dict[f"model.layers.{i}.mlp.experts.0.down_proj.weight"].shape - w1 = torch.empty( - (num_experts, moe_dim, dim), dtype=state_dict[f"model.layers.{i}.mlp.experts.1.down_proj.weight"].dtype - ) # Gate - w2 = torch.empty( - (num_experts, dim, moe_dim), dtype=state_dict[f"model.layers.{i}.mlp.experts.1.down_proj.weight"].dtype - ) # Down - w3 = torch.empty( - (num_experts, moe_dim, dim), dtype=state_dict[f"model.layers.{i}.mlp.experts.1.down_proj.weight"].dtype - ) # Up - for j in range(num_experts): - w1[j].copy_(state_dict[f"model.layers.{i}.mlp.experts.{j}.gate_proj.weight"]) - w2[j].copy_(state_dict[f"model.layers.{i}.mlp.experts.{j}.down_proj.weight"]) - w3[j].copy_(state_dict[f"model.layers.{i}.mlp.experts.{j}.up_proj.weight"]) - - del state_dict[f"model.layers.{i}.mlp.experts.{j}.gate_proj.weight"] - del state_dict[f"model.layers.{i}.mlp.experts.{j}.down_proj.weight"] - del state_dict[f"model.layers.{i}.mlp.experts.{j}.up_proj.weight"] + # Check if this is the new fused format (transformers 5.0+) or old per-expert format + if f"model.layers.{i}.mlp.experts.gate_up_proj" in state_dict: + # New fused format: gate_up_proj has shape (num_experts, 2*moe_dim, dim) + gate_up_proj = state_dict[f"model.layers.{i}.mlp.experts.gate_up_proj"] + down_proj = state_dict[f"model.layers.{i}.mlp.experts.down_proj"] + + num_experts, fused_dim, dim = gate_up_proj.shape + moe_dim = fused_dim // 2 + + # Split gate_up_proj into w1 (gate) and w3 (up) + w1 = gate_up_proj[:, :moe_dim, :] # Gate: (num_experts, moe_dim, dim) + w3 = gate_up_proj[:, moe_dim:, :] # Up: (num_experts, moe_dim, dim) + w2 = down_proj # Down: (num_experts, dim, moe_dim) + + del state_dict[f"model.layers.{i}.mlp.experts.gate_up_proj"] + del state_dict[f"model.layers.{i}.mlp.experts.down_proj"] + else: + # Old per-expert format + num_experts = len([j for j in state_dict.keys() if f"model.layers.{i}.mlp.experts" in j]) // 3 + if num_experts == 0: + return + + dim, moe_dim = state_dict[f"model.layers.{i}.mlp.experts.0.down_proj.weight"].shape + w1 = torch.empty( + (num_experts, moe_dim, dim), dtype=state_dict[f"model.layers.{i}.mlp.experts.1.down_proj.weight"].dtype + ) # Gate + w2 = torch.empty( + (num_experts, dim, moe_dim), dtype=state_dict[f"model.layers.{i}.mlp.experts.1.down_proj.weight"].dtype + ) # Down + w3 = torch.empty( + (num_experts, moe_dim, dim), dtype=state_dict[f"model.layers.{i}.mlp.experts.1.down_proj.weight"].dtype + ) # Up + for j in range(num_experts): + w1[j].copy_(state_dict[f"model.layers.{i}.mlp.experts.{j}.gate_proj.weight"]) + w2[j].copy_(state_dict[f"model.layers.{i}.mlp.experts.{j}.down_proj.weight"]) + w3[j].copy_(state_dict[f"model.layers.{i}.mlp.experts.{j}.up_proj.weight"]) + + del state_dict[f"model.layers.{i}.mlp.experts.{j}.gate_proj.weight"] + del state_dict[f"model.layers.{i}.mlp.experts.{j}.down_proj.weight"] + del state_dict[f"model.layers.{i}.mlp.experts.{j}.up_proj.weight"] state_dict[f"model.layers.{i}.mlp.experts.w1"] = w1 state_dict[f"model.layers.{i}.mlp.experts.w2"] = w2 @@ -109,21 +132,16 @@ def convert_tt_layer_to_hf(state_dict: dict[str, Tensor], layer_index: int): state_dict[f"model.layers.{i}.mlp.gate.weight"] = state_dict[f"model.layers.{i}.mlp.router.gate.weight"] del state_dict[f"model.layers.{i}.mlp.router.gate.weight"] - # Routed experts - num_experts, moe_dim, dim = state_dict[f"model.layers.{i}.mlp.experts.w1"].shape + # Routed experts - convert to per-expert format (compatible with vLLM and transformers) + w1 = state_dict.pop(f"model.layers.{i}.mlp.experts.w1") # (num_experts, moe_dim, dim) + w2 = state_dict.pop(f"model.layers.{i}.mlp.experts.w2") # (num_experts, dim, moe_dim) + w3 = state_dict.pop(f"model.layers.{i}.mlp.experts.w3") # (num_experts, moe_dim, dim) + + num_experts = w1.shape[0] for j in range(num_experts): - state_dict[f"model.layers.{i}.mlp.experts.{j}.gate_proj.weight"] = state_dict[ - f"model.layers.{i}.mlp.experts.w1" - ][j] - state_dict[f"model.layers.{i}.mlp.experts.{j}.down_proj.weight"] = state_dict[ - f"model.layers.{i}.mlp.experts.w2" - ][j] - state_dict[f"model.layers.{i}.mlp.experts.{j}.up_proj.weight"] = state_dict[ - f"model.layers.{i}.mlp.experts.w3" - ][j] - del state_dict[f"model.layers.{i}.mlp.experts.w1"] - del state_dict[f"model.layers.{i}.mlp.experts.w2"] - del state_dict[f"model.layers.{i}.mlp.experts.w3"] + state_dict[f"model.layers.{i}.mlp.experts.{j}.gate_proj.weight"] = w1[j] + state_dict[f"model.layers.{i}.mlp.experts.{j}.down_proj.weight"] = w2[j] + state_dict[f"model.layers.{i}.mlp.experts.{j}.up_proj.weight"] = w3[j] def convert_tt_to_hf_moe(state_dict: dict[str, Tensor]): diff --git a/src/prime_rl/trainer/models/glm4_moe/modeling_glm4_moe.py b/src/prime_rl/trainer/models/glm4_moe/modeling_glm4_moe.py index 8c2bbe8056c..710659fd5db 100644 --- a/src/prime_rl/trainer/models/glm4_moe/modeling_glm4_moe.py +++ b/src/prime_rl/trainer/models/glm4_moe/modeling_glm4_moe.py @@ -37,8 +37,9 @@ from prime_rl.trainer.models.layers.lm_head import PrimeLmOutput from prime_rl.trainer.models.layers.mlp import MLP, MLPConfig from prime_rl.trainer.models.layers.moe import MoE, MoEArgs -from prime_rl.trainer.models.layers.rms_norm import RMSNorm, RMSNormConfig +from prime_rl.trainer.models.layers.norms import RMSNorm, RMSNormConfig from prime_rl.trainer.models.layers.rotary_emb import RotaryEmbedding, RotaryEmbeddingConfig +from prime_rl.utils.sequence import get_cu_seqlens_from_position_ids class Glm4MoeDecoderLayer(GradientCheckpointingLayer): @@ -67,6 +68,7 @@ def __init__(self, config: Glm4MoeConfig, layer_idx: int): top_k=config.num_experts_per_tok, load_balance_coeff=1e-3, use_grouped_mm=config.use_grouped_mm, + fp8=getattr(config, "fp8", False), ) mlp_config = MLPConfig( hidden_size=config.hidden_size, @@ -90,6 +92,7 @@ def forward( position_embeddings: Optional[tuple[torch.Tensor, torch.Tensor]] = None, # necessary, but kept here for BC cu_seqlens: Optional[torch.LongTensor] = None, max_seqlen: Optional[int] = None, + routed_experts: Optional[torch.LongTensor] = None, ) -> torch.Tensor: residual = hidden_states hidden_states = self.input_layernorm(hidden_states) @@ -105,7 +108,7 @@ def forward( # Fully Connected residual = hidden_states hidden_states = self.post_attention_layernorm(hidden_states) - hidden_states = self.mlp(hidden_states) + hidden_states = self.mlp(hidden_states, routed_experts=routed_experts) hidden_states = residual + hidden_states return hidden_states @@ -132,7 +135,8 @@ def _init_weights(self, module): @classmethod def is_hf_state_dict(cls, state_dict: dict[str, Tensor]) -> bool: """Check if the state dict contains MoE layers in HuggingFace format.""" - return any("mlp.experts.1.up_proj" in module_name for module_name in state_dict.keys()) + # Check for old per-expert format or new fused format (transformers 5.0+) + return any("mlp.experts.1.up_proj" in name or "mlp.experts.gate_up_proj" in name for name in state_dict.keys()) @classmethod def is_prime_state_dict(cls, state_dict: dict[str, Tensor]) -> bool: @@ -200,24 +204,20 @@ def forward( input_ids: Optional[torch.LongTensor] = None, position_ids: Optional[torch.LongTensor] = None, inputs_embeds: Optional[torch.FloatTensor] = None, + routed_experts: Optional[torch.LongTensor] = None, ) -> BaseModelOutputWithPast: + """ + routed_experts (`torch.LongTensor` of shape `(batch_size, sequence_length, num_hidden_layers, num_experts_per_tok)`, *optional*): + Routed experts for each token in the sequence. Only used for router replay. + """ if (input_ids is None) ^ (inputs_embeds is not None): raise ValueError("You must specify exactly one of input_ids or inputs_embeds") if inputs_embeds is None: inputs_embeds: torch.Tensor = self.embed_tokens(input_ids) - if self.config._attn_implementation in ("flash_attention_2", "flash_attention_3"): - flat_position_ids = position_ids.view(-1) - seqlens = torch.cat( - [ - flat_position_ids[0:1], - flat_position_ids[:-1][(flat_position_ids == 0)[1:]] + 1, - flat_position_ids[-1:] + 1, - ] - ) - max_seqlen = seqlens.max().item() - cu_seqlens = seqlens.cumsum(dim=0, dtype=torch.int32) + if self.config._attn_implementation in ("flash_attention_2", "flash_attention_3", "fa4"): + cu_seqlens, max_seqlen = get_cu_seqlens_from_position_ids(position_ids) torch._dynamo.mark_dynamic(cu_seqlens, 0) else: max_seqlen = None @@ -226,12 +226,14 @@ def forward( hidden_states = inputs_embeds position_embeddings = self.rotary_emb(hidden_states, position_ids) - for decoder_layer in self.layers[: self.config.num_hidden_layers]: + for layer_idx, decoder_layer in enumerate(self.layers[: self.config.num_hidden_layers]): + routed_experts_layer = routed_experts[:, :, layer_idx, :] if routed_experts is not None else None hidden_states = decoder_layer( hidden_states, position_embeddings=position_embeddings, cu_seqlens=cu_seqlens, max_seqlen=max_seqlen, + routed_experts=routed_experts_layer, ) hidden_states = self.norm(hidden_states) @@ -240,7 +242,7 @@ def forward( @auto_docstring class Glm4MoeForCausalLM(Glm4MoePreTrainedModel, GenerationMixin): - _tied_weights_keys = ["lm_head.weight"] + _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": "colwise_rep"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} @@ -264,15 +266,22 @@ def forward( use_cache: Optional[bool] = None, cache_position: Optional[torch.LongTensor] = None, logits_to_keep: Union[int, torch.Tensor] = 0, - temperature: Optional[float] = 1.0, + temperature: Optional[torch.Tensor] = None, + routed_experts: Optional[torch.LongTensor] = None, **kwargs: Unpack[TransformersKwargs], ) -> PrimeLmOutput: r""" + cache_position (`torch.LongTensor` of shape `(sequence_length)`, *optional*): + Indices of input tokens in the KV cache. Accepted only for HuggingFace API + compatibility — prime-rl asserts `use_cache is None` since training does not + perform autoregressive decoding, so this argument is unused. labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): Labels used by PrimeRL's wrapped LM head to optionally compute per-token logprobs/entropy. If not provided, the wrapped LM head returns logits only. - temperature (`float`, *optional*, defaults to 1.0): - Temperature used for the logprobs/entropy computation when `labels` are provided. + temperature (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*): + Per-token temperatures for logprobs/entropy computation when `labels` are provided. + routed_experts (`torch.LongTensor` of shape `(batch_size, sequence_length, num_hidden_layers, num_experts_per_tok)`, *optional*): + Routed experts for each token in the sequence. Only used for router replay. Example: @@ -303,6 +312,7 @@ def forward( input_ids=input_ids, position_ids=position_ids, inputs_embeds=inputs_embeds, + routed_experts=routed_experts, ) hidden_states = outputs.last_hidden_state diff --git a/src/prime_rl/trainer/models/glm_moe_dsa/__init__.py b/src/prime_rl/trainer/models/glm_moe_dsa/__init__.py new file mode 100644 index 00000000000..7e898a6203f --- /dev/null +++ b/src/prime_rl/trainer/models/glm_moe_dsa/__init__.py @@ -0,0 +1,13 @@ +from prime_rl.trainer.models.glm_moe_dsa.configuration_glm_moe_dsa import GlmMoeDsaConfig +from prime_rl.trainer.models.glm_moe_dsa.modeling_glm_moe_dsa import ( + GlmMoeDsaForCausalLM, + GlmMoeDsaModel, + GlmMoeDsaPreTrainedModel, +) + +__all__ = [ + "GlmMoeDsaConfig", + "GlmMoeDsaForCausalLM", + "GlmMoeDsaModel", + "GlmMoeDsaPreTrainedModel", +] diff --git a/src/prime_rl/trainer/models/glm_moe_dsa/configuration_glm_moe_dsa.py b/src/prime_rl/trainer/models/glm_moe_dsa/configuration_glm_moe_dsa.py new file mode 100644 index 00000000000..caca4fc5e9e --- /dev/null +++ b/src/prime_rl/trainer/models/glm_moe_dsa/configuration_glm_moe_dsa.py @@ -0,0 +1,209 @@ +import warnings + +from transformers.configuration_utils import PretrainedConfig + + +class GlmMoeDsaConfig(PretrainedConfig): + r""" + Configuration class for the GLM-5 (GlmMoeDsa) model. + + GLM-5 uses DeepSeek V3.2-style Multi-head Latent Attention (MLA) with LoRA-compressed + KV projections, combined with Mixture-of-Experts feed-forward layers. + + Args: + vocab_size (`int`, defaults to 154880): + Vocabulary size. + hidden_size (`int`, defaults to 6144): + Dimension of the hidden representations. + intermediate_size (`int`, defaults to 12288): + Dimension of the dense MLP representations. + moe_intermediate_size (`int`, defaults to 2048): + Dimension of the MoE expert representations. + num_hidden_layers (`int`, defaults to 78): + Number of hidden layers in the Transformer decoder. + num_attention_heads (`int`, defaults to 64): + Number of attention heads. + num_key_value_heads (`int`, defaults to 64): + Number of key-value heads (after kv_b_proj up-projection). + n_shared_experts (`int`, defaults to 1): + Number of shared experts. + n_routed_experts (`int`, defaults to 256): + Number of routed experts. + routed_scaling_factor (`float`, defaults to 2.5): + Scaling factor for routed experts. + kv_lora_rank (`int`, defaults to 512): + Rank of the LoRA compression for key-value projections. + q_lora_rank (`int`, defaults to 2048): + Rank of the LoRA compression for query projections. + qk_rope_head_dim (`int`, defaults to 64): + Dimension of query/key heads that use rotary position embeddings. + v_head_dim (`int`, defaults to 256): + Dimension of value heads. + qk_nope_head_dim (`int`, defaults to 192): + Dimension of query/key heads that don't use rotary position embeddings. + n_group (`int`, defaults to 1): + Number of groups for routed experts. + topk_group (`int`, defaults to 1): + Number of selected groups per token. + num_experts_per_tok (`int`, defaults to 8): + Number of active experts per token. + first_k_dense_replace (`int`, defaults to 3): + Number of initial layers that use dense MLP instead of MoE. + norm_topk_prob (`bool`, defaults to `True`): + Whether to normalize the top-k routing probabilities. + hidden_act (`str`, defaults to `"silu"`): + Activation function in the feed-forward layers. + max_position_embeddings (`int`, defaults to 202752): + Maximum sequence length. + initializer_range (`float`, defaults to 0.02): + Standard deviation for weight initialization. + rms_norm_eps (`float`, defaults to 1e-5): + Epsilon for RMS normalization. + rope_interleave (`bool`, defaults to `True`): + Whether to use interleaved (non-neox) RoPE style. + attention_bias (`bool`, defaults to `False`): + Whether to use bias in attention projections. + attention_dropout (`float`, defaults to 0.0): + Dropout for attention probabilities. + index_n_heads (`int`, defaults to 32): + Number of heads used by the sparse indexer. + index_head_dim (`int`, defaults to 128): + Head dimension used by the sparse indexer. + indexer_rope_interleave (`bool`, defaults to `True`): + Whether to use interleaved RoPE style in the sparse indexer. + index_topk (`int`, defaults to 2048): + Number of top tokens selected by the sparse indexer. + scoring_func (`str`, defaults to `"sigmoid"`): + Scoring function for MoE router. Must match the vLLM inference + server's expectation (vLLM defaults to ``"softmax"`` when this + field is absent from the config). + topk_method (`str`, defaults to `"noaux_tc"`): + MoE routing top-k method used by GLM-5 checkpoints. + use_grouped_mm (`bool`, defaults to `True`): + Whether to use grouped matrix multiplication for MoE. + """ + + model_type = "glm_moe_dsa" + keys_to_ignore_at_inference = ["past_key_values"] + attribute_map = {"num_local_experts": "n_routed_experts"} + + base_model_tp_plan = { + "layers.*.self_attn.o_proj": "rowwise", + "layers.*.mlp.experts.*.gate_proj": "colwise", + "layers.*.mlp.experts.*.up_proj": "colwise", + "layers.*.mlp.experts.*.down_proj": "rowwise", + "layers.*.mlp.gate_proj": "colwise", + "layers.*.mlp.up_proj": "colwise", + "layers.*.mlp.down_proj": "rowwise", + } + base_model_pp_plan = { + "embed_tokens": (["input_ids"], ["inputs_embeds"]), + "layers": (["hidden_states"], ["hidden_states"]), + "norm": (["hidden_states"], ["hidden_states"]), + } + + def __init__( + self, + vocab_size=154880, + hidden_size=6144, + intermediate_size=12288, + moe_intermediate_size=2048, + num_hidden_layers=78, + num_attention_heads=64, + num_key_value_heads=64, + n_shared_experts=1, + n_routed_experts=256, + routed_scaling_factor=2.5, + kv_lora_rank=512, + q_lora_rank=2048, + qk_rope_head_dim=64, + v_head_dim=256, + qk_nope_head_dim=192, + n_group=1, + topk_group=1, + num_experts_per_tok=8, + first_k_dense_replace=3, + norm_topk_prob=True, + hidden_act="silu", + max_position_embeddings=202752, + initializer_range=0.02, + rms_norm_eps=1e-5, + use_cache=True, + tie_word_embeddings=False, + rope_interleave=True, + rope_parameters=None, + rope_theta=10000.0, + mlp_layer_types=None, + attention_bias=False, + attention_dropout=0.0, + index_n_heads=32, + index_head_dim=128, + indexer_rope_interleave=True, + pad_token_id=154820, + index_topk=2048, + scoring_func="sigmoid", + topk_method="noaux_tc", + use_grouped_mm=True, + **kwargs, + ): + self.vocab_size = vocab_size + self.max_position_embeddings = max_position_embeddings + self.hidden_size = hidden_size + self.intermediate_size = intermediate_size + self.moe_intermediate_size = moe_intermediate_size + self.num_hidden_layers = num_hidden_layers + self.num_attention_heads = num_attention_heads + self.num_key_value_heads = num_key_value_heads + + # MLA dimensions + self.kv_lora_rank = kv_lora_rank + self.q_lora_rank = q_lora_rank + self.qk_rope_head_dim = qk_rope_head_dim + self.v_head_dim = v_head_dim + self.qk_nope_head_dim = qk_nope_head_dim + self.qk_head_dim = qk_nope_head_dim + qk_rope_head_dim + self.head_dim = qk_rope_head_dim + + # MoE + self.n_shared_experts = n_shared_experts + self.n_routed_experts = n_routed_experts + self.routed_scaling_factor = routed_scaling_factor + self.n_group = n_group + self.topk_group = topk_group + self.num_experts_per_tok = num_experts_per_tok + self.first_k_dense_replace = first_k_dense_replace + self.norm_topk_prob = norm_topk_prob + + # MLP layer types: per-layer dense/sparse dispatch + if mlp_layer_types is None: + mlp_layer_types = ["dense"] * min(first_k_dense_replace, num_hidden_layers) + ["sparse"] * max( + 0, num_hidden_layers - first_k_dense_replace + ) + self.mlp_layer_types = mlp_layer_types + + self.hidden_act = hidden_act + self.initializer_range = initializer_range + self.rms_norm_eps = rms_norm_eps + self.use_cache = use_cache + self.rope_interleave = rope_interleave + self.rope_parameters = rope_parameters + self.rope_theta = rope_theta + self.attention_bias = attention_bias + self.attention_dropout = attention_dropout + self.index_n_heads = index_n_heads + self.index_head_dim = index_head_dim + self.indexer_rope_interleave = indexer_rope_interleave + self.index_topk = index_topk + self.scoring_func = scoring_func + self.topk_method = topk_method + self.use_grouped_mm = use_grouped_mm + self.pad_token_id = pad_token_id + + if not self.use_grouped_mm: + warnings.warn("not using grouped mm for moe is very slow, should only be used for debugging") + + super().__init__( + pad_token_id=pad_token_id, + tie_word_embeddings=tie_word_embeddings, + **kwargs, + ) diff --git a/src/prime_rl/trainer/models/glm_moe_dsa/converting_glm_moe_dsa.py b/src/prime_rl/trainer/models/glm_moe_dsa/converting_glm_moe_dsa.py new file mode 100644 index 00000000000..e8e5b311b83 --- /dev/null +++ b/src/prime_rl/trainer/models/glm_moe_dsa/converting_glm_moe_dsa.py @@ -0,0 +1,266 @@ +import torch +from torch import Tensor + +from prime_rl.trainer.models.fp8 import quantize_to_fp8_blockwise + + +def get_max_layer_num(state_dict: dict[str, Tensor]) -> int: + return max(int(i.split(".")[2]) for i in state_dict.keys() if "model.layers." in i) + 1 + + +def _is_moe_layer(state_dict: dict[str, Tensor], layer_idx: int) -> bool: + """Check if a layer is an MoE layer by looking for the router gate weight.""" + return f"model.layers.{layer_idx}.mlp.gate.weight" in state_dict + + +def convert_hf_layer_to_tt(state_dict: dict[str, Tensor], layer_idx: int): + i = layer_idx + + if not _is_moe_layer(state_dict, i): + return + + # Router: gate.weight -> router.gate.weight + state_dict[f"model.layers.{i}.mlp.router.gate.weight"] = state_dict[f"model.layers.{i}.mlp.gate.weight"] + del state_dict[f"model.layers.{i}.mlp.gate.weight"] + + # Routed experts: fused or per-expert format -> stacked w1/w2/w3 + if f"model.layers.{i}.mlp.experts.gate_up_proj" in state_dict: + gate_up_proj = state_dict[f"model.layers.{i}.mlp.experts.gate_up_proj"] + down_proj = state_dict[f"model.layers.{i}.mlp.experts.down_proj"] + + num_experts, fused_dim, dim = gate_up_proj.shape + moe_dim = fused_dim // 2 + + w1 = gate_up_proj[:, :moe_dim, :] + w3 = gate_up_proj[:, moe_dim:, :] + w2 = down_proj + + del state_dict[f"model.layers.{i}.mlp.experts.gate_up_proj"] + del state_dict[f"model.layers.{i}.mlp.experts.down_proj"] + else: + num_experts = len([j for j in state_dict.keys() if f"model.layers.{i}.mlp.experts" in j]) // 3 + if num_experts == 0: + return + + dim, moe_dim = state_dict[f"model.layers.{i}.mlp.experts.0.down_proj.weight"].shape + dtype = state_dict[f"model.layers.{i}.mlp.experts.0.down_proj.weight"].dtype + w1 = torch.empty((num_experts, moe_dim, dim), dtype=dtype) + w2 = torch.empty((num_experts, dim, moe_dim), dtype=dtype) + w3 = torch.empty((num_experts, moe_dim, dim), dtype=dtype) + for j in range(num_experts): + w1[j].copy_(state_dict[f"model.layers.{i}.mlp.experts.{j}.gate_proj.weight"]) + w2[j].copy_(state_dict[f"model.layers.{i}.mlp.experts.{j}.down_proj.weight"]) + w3[j].copy_(state_dict[f"model.layers.{i}.mlp.experts.{j}.up_proj.weight"]) + + del state_dict[f"model.layers.{i}.mlp.experts.{j}.gate_proj.weight"] + del state_dict[f"model.layers.{i}.mlp.experts.{j}.down_proj.weight"] + del state_dict[f"model.layers.{i}.mlp.experts.{j}.up_proj.weight"] + + state_dict[f"model.layers.{i}.mlp.experts.w1"] = w1 + state_dict[f"model.layers.{i}.mlp.experts.w2"] = w2 + state_dict[f"model.layers.{i}.mlp.experts.w3"] = w3 + + # Shared experts + state_dict[f"model.layers.{i}.mlp.shared_expert.w1"] = state_dict[ + f"model.layers.{i}.mlp.shared_experts.gate_proj.weight" + ] + state_dict[f"model.layers.{i}.mlp.shared_expert.w2"] = state_dict[ + f"model.layers.{i}.mlp.shared_experts.down_proj.weight" + ] + state_dict[f"model.layers.{i}.mlp.shared_expert.w3"] = state_dict[ + f"model.layers.{i}.mlp.shared_experts.up_proj.weight" + ] + del state_dict[f"model.layers.{i}.mlp.shared_experts.gate_proj.weight"] + del state_dict[f"model.layers.{i}.mlp.shared_experts.down_proj.weight"] + del state_dict[f"model.layers.{i}.mlp.shared_experts.up_proj.weight"] + + # Expert bias for load balancing + state_dict[f"model.layers.{i}.mlp.expert_bias"] = state_dict[f"model.layers.{i}.mlp.gate.e_score_correction_bias"] + del state_dict[f"model.layers.{i}.mlp.gate.e_score_correction_bias"] + + +def convert_hf_to_tt_moe(state_dict: dict[str, Tensor]): + num_layers = get_max_layer_num(state_dict) + for i in range(num_layers): + convert_hf_layer_to_tt(state_dict, i) + + +def convert_tt_layer_to_hf(state_dict: dict[str, Tensor], layer_index: int): + i = layer_index + + # Expert bias + if f"model.layers.{i}.mlp.expert_bias" in state_dict: + state_dict[f"model.layers.{i}.mlp.gate.e_score_correction_bias"] = state_dict[ + f"model.layers.{i}.mlp.expert_bias" + ] + del state_dict[f"model.layers.{i}.mlp.expert_bias"] + if f"model.layers.{i}.mlp.tokens_per_expert" in state_dict: + del state_dict[f"model.layers.{i}.mlp.tokens_per_expert"] + + # Shared experts + if f"model.layers.{i}.mlp.shared_expert.w1" in state_dict: + state_dict[f"model.layers.{i}.mlp.shared_experts.gate_proj.weight"] = state_dict[ + f"model.layers.{i}.mlp.shared_expert.w1" + ] + state_dict[f"model.layers.{i}.mlp.shared_experts.down_proj.weight"] = state_dict[ + f"model.layers.{i}.mlp.shared_expert.w2" + ] + state_dict[f"model.layers.{i}.mlp.shared_experts.up_proj.weight"] = state_dict[ + f"model.layers.{i}.mlp.shared_expert.w3" + ] + + if state_dict[f"model.layers.{i}.mlp.shared_experts.up_proj.weight"].shape[0] == 1: + state_dict[f"model.layers.{i}.mlp.shared_experts.up_proj.weight"] = state_dict[ + f"model.layers.{i}.mlp.shared_experts.up_proj.weight" + ][0] + state_dict[f"model.layers.{i}.mlp.shared_experts.down_proj.weight"] = state_dict[ + f"model.layers.{i}.mlp.shared_experts.down_proj.weight" + ][0] + state_dict[f"model.layers.{i}.mlp.shared_experts.gate_proj.weight"] = state_dict[ + f"model.layers.{i}.mlp.shared_experts.gate_proj.weight" + ][0] + del state_dict[f"model.layers.{i}.mlp.shared_expert.w1"] + del state_dict[f"model.layers.{i}.mlp.shared_expert.w2"] + del state_dict[f"model.layers.{i}.mlp.shared_expert.w3"] + + # Router + if f"model.layers.{i}.mlp.router.gate.weight" in state_dict: + state_dict[f"model.layers.{i}.mlp.gate.weight"] = state_dict[f"model.layers.{i}.mlp.router.gate.weight"] + del state_dict[f"model.layers.{i}.mlp.router.gate.weight"] + + # Routed experts - convert to per-expert format (compatible with vLLM and transformers) + w1 = state_dict.pop(f"model.layers.{i}.mlp.experts.w1") # (num_experts, moe_dim, dim) + w2 = state_dict.pop(f"model.layers.{i}.mlp.experts.w2") # (num_experts, dim, moe_dim) + w3 = state_dict.pop(f"model.layers.{i}.mlp.experts.w3") # (num_experts, moe_dim, dim) + + num_experts = w1.shape[0] + for j in range(num_experts): + state_dict[f"model.layers.{i}.mlp.experts.{j}.gate_proj.weight"] = w1[j] + state_dict[f"model.layers.{i}.mlp.experts.{j}.down_proj.weight"] = w2[j] + state_dict[f"model.layers.{i}.mlp.experts.{j}.up_proj.weight"] = w3[j] + + +def convert_tt_to_hf_moe(state_dict: dict[str, Tensor]): + num_layers = get_max_layer_num(state_dict) + for i in range(num_layers): + convert_tt_layer_to_hf(state_dict, i) + + +def convert_tt_layer_to_vllm_kernel( + state_dict: dict[str, Tensor], + layer_idx: int, + quantize_fp8: bool = False, +) -> dict[str, Tensor]: + """Convert a single GLM layer from PrimeRL format to vLLM kernel format.""" + out: dict[str, Tensor] = {} + prefix = f"model.layers.{layer_idx}" + + def add(name: str, tensor: Tensor) -> None: + out[name] = tensor + + def add_maybe_fp8(name: str, tensor: Tensor) -> None: + if quantize_fp8 and tensor.ndim == 2: + fp8_weight, scale = quantize_to_fp8_blockwise(tensor) + out[name] = fp8_weight + scale_name = name.removesuffix(".weight") + ".weight_scale_inv" + out[scale_name] = scale + return + out[name] = tensor + + for name in [f"{prefix}.input_layernorm.weight", f"{prefix}.post_attention_layernorm.weight"]: + if name in state_dict: + add(name, state_dict[name]) + + q_a_key = f"{prefix}.self_attn.q_a_proj.weight" + kv_a_key = f"{prefix}.self_attn.kv_a_proj_with_mqa.weight" + if q_a_key in state_dict and kv_a_key in state_dict: + add_maybe_fp8( + f"{prefix}.self_attn.fused_qkv_a_proj.weight", torch.cat([state_dict[q_a_key], state_dict[kv_a_key]], dim=0) + ) + + for suffix in ["q_a_layernorm.weight", "kv_a_layernorm.weight"]: + key = f"{prefix}.self_attn.{suffix}" + if key in state_dict: + add(key, state_dict[key]) + + for suffix in ["q_b_proj.weight", "kv_b_proj.weight", "o_proj.weight"]: + key = f"{prefix}.self_attn.{suffix}" + if key in state_dict: + add_maybe_fp8(key, state_dict[key]) + + indexer_wq_b_key = f"{prefix}.self_attn.indexer.wq_b.weight" + if indexer_wq_b_key in state_dict: + add_maybe_fp8(indexer_wq_b_key, state_dict[indexer_wq_b_key]) + + indexer_wk_key = f"{prefix}.self_attn.indexer.wk.weight" + indexer_weights_proj_key = f"{prefix}.self_attn.indexer.weights_proj.weight" + if indexer_wk_key in state_dict and indexer_weights_proj_key in state_dict: + add( + f"{prefix}.self_attn.indexer.wk_weights_proj.weight", + torch.cat([state_dict[indexer_wk_key], state_dict[indexer_weights_proj_key]], dim=0), + ) + for suffix in ["indexer.k_norm.weight", "indexer.k_norm.bias"]: + key = f"{prefix}.self_attn.{suffix}" + if key in state_dict: + add(key, state_dict[key]) + + gate_key = f"{prefix}.mlp.gate_proj.weight" + up_key = f"{prefix}.mlp.up_proj.weight" + down_key = f"{prefix}.mlp.down_proj.weight" + if gate_key in state_dict and up_key in state_dict: + add_maybe_fp8(f"{prefix}.mlp.gate_up_proj.weight", torch.cat([state_dict[gate_key], state_dict[up_key]], dim=0)) + if down_key in state_dict: + add_maybe_fp8(f"{prefix}.mlp.down_proj.weight", state_dict[down_key]) + + router_key = f"{prefix}.mlp.router.gate.weight" + if router_key in state_dict: + add(f"{prefix}.mlp.gate.weight", state_dict[router_key]) + expert_bias_key = f"{prefix}.mlp.expert_bias" + if expert_bias_key in state_dict: + add(f"{prefix}.mlp.gate.e_score_correction_bias", state_dict[expert_bias_key]) + + w1_key = f"{prefix}.mlp.experts.w1" + w2_key = f"{prefix}.mlp.experts.w2" + w3_key = f"{prefix}.mlp.experts.w3" + if w1_key in state_dict and w2_key in state_dict and w3_key in state_dict: + w1 = state_dict[w1_key] + w2 = state_dict[w2_key] + w3 = state_dict[w3_key] + w13 = torch.cat([w1, w3], dim=1) + + if quantize_fp8: + w13_fp8: list[Tensor] = [] + w13_scales: list[Tensor] = [] + w2_fp8: list[Tensor] = [] + w2_scales: list[Tensor] = [] + for expert_idx in range(w1.shape[0]): + expert_w13_fp8, expert_w13_scales = quantize_to_fp8_blockwise(w13[expert_idx]) + expert_w2_fp8, expert_w2_scales = quantize_to_fp8_blockwise(w2[expert_idx]) + w13_fp8.append(expert_w13_fp8) + w13_scales.append(expert_w13_scales) + w2_fp8.append(expert_w2_fp8) + w2_scales.append(expert_w2_scales) + + out[f"{prefix}.mlp.experts.w13_weight"] = torch.stack(w13_fp8) + out[f"{prefix}.mlp.experts.w13_weight_scale_inv"] = torch.stack(w13_scales) + out[f"{prefix}.mlp.experts.w2_weight"] = torch.stack(w2_fp8) + out[f"{prefix}.mlp.experts.w2_weight_scale_inv"] = torch.stack(w2_scales) + else: + out[f"{prefix}.mlp.experts.w13_weight"] = w13 + out[f"{prefix}.mlp.experts.w2_weight"] = w2 + + sw1_key = f"{prefix}.mlp.shared_expert.w1" + sw2_key = f"{prefix}.mlp.shared_expert.w2" + sw3_key = f"{prefix}.mlp.shared_expert.w3" + if sw1_key in state_dict and sw2_key in state_dict and sw3_key in state_dict: + sw1 = state_dict[sw1_key] + sw2 = state_dict[sw2_key] + sw3 = state_dict[sw3_key] + if sw1.ndim == 3: + sw1 = sw1.squeeze(0) + sw2 = sw2.squeeze(0) + sw3 = sw3.squeeze(0) + add_maybe_fp8(f"{prefix}.mlp.shared_experts.gate_up_proj.weight", torch.cat([sw1, sw3], dim=0)) + add_maybe_fp8(f"{prefix}.mlp.shared_experts.down_proj.weight", sw2) + + return out diff --git a/src/prime_rl/trainer/models/glm_moe_dsa/modeling_glm_moe_dsa.py b/src/prime_rl/trainer/models/glm_moe_dsa/modeling_glm_moe_dsa.py new file mode 100644 index 00000000000..2bb90d94ff1 --- /dev/null +++ b/src/prime_rl/trainer/models/glm_moe_dsa/modeling_glm_moe_dsa.py @@ -0,0 +1,351 @@ +import warnings +from typing import Optional, Union + +import torch +import torch.distributed as dist +from torch import Tensor, nn +from transformers.cache_utils import Cache +from transformers.generation import GenerationMixin +from transformers.modeling_layers import GradientCheckpointingLayer +from transformers.modeling_outputs import BaseModelOutputWithPast +from transformers.processing_utils import Unpack +from transformers.utils import TransformersKwargs, auto_docstring +from transformers.utils.deprecation import deprecate_kwarg + +from prime_rl.trainer.models.base import PreTrainedModelPrimeRL +from prime_rl.trainer.models.glm_moe_dsa.configuration_glm_moe_dsa import GlmMoeDsaConfig +from prime_rl.trainer.models.glm_moe_dsa.converting_glm_moe_dsa import ( + convert_hf_layer_to_tt, + convert_hf_to_tt_moe, + convert_tt_layer_to_hf, + convert_tt_layer_to_vllm_kernel, + convert_tt_to_hf_moe, +) +from prime_rl.trainer.models.glm_moe_dsa.sparse_mla_attention import GlmMoeDsaAttention, SparseMlaAttentionArgs +from prime_rl.trainer.models.layers.lm_head import PrimeLmOutput +from prime_rl.trainer.models.layers.mlp import MLP, MLPConfig +from prime_rl.trainer.models.layers.moe import MoE, MoEArgs +from prime_rl.trainer.models.layers.norms import RMSNorm, RMSNormConfig +from prime_rl.trainer.models.layers.rotary_emb import RotaryEmbedding, RotaryEmbeddingConfig + + +def _sparse_mla_attention_args(config: GlmMoeDsaConfig) -> SparseMlaAttentionArgs: + if config.q_lora_rank is None: + raise ValueError("Sparse MLA attention requires q_lora_rank to be set") + return SparseMlaAttentionArgs( + hidden_size=config.hidden_size, + num_attention_heads=config.num_attention_heads, + kv_lora_rank=config.kv_lora_rank, + q_lora_rank=config.q_lora_rank, + qk_rope_head_dim=config.qk_rope_head_dim, + qk_nope_head_dim=config.qk_nope_head_dim, + qk_head_dim=config.qk_head_dim, + v_head_dim=config.v_head_dim, + attention_bias=config.attention_bias, + rms_norm_eps=config.rms_norm_eps, + index_n_heads=config.index_n_heads, + index_head_dim=config.index_head_dim, + index_topk=config.index_topk, + ) + + +class GlmMoeDsaDecoderLayer(GradientCheckpointingLayer): + def __init__(self, config: GlmMoeDsaConfig, layer_idx: int): + super().__init__() + self.hidden_size = config.hidden_size + self.self_attn = GlmMoeDsaAttention(_sparse_mla_attention_args(config)) + + moe_args = MoEArgs( + num_experts=config.n_routed_experts, + num_shared_experts=config.n_shared_experts, + score_func="sigmoid", + route_norm=config.norm_topk_prob, + route_scale=config.routed_scaling_factor, + score_before_experts=False, + top_k=config.num_experts_per_tok, + load_balance_coeff=1e-3, + use_grouped_mm=config.use_grouped_mm, + fp8=getattr(config, "fp8", False), + ) + mlp_config = MLPConfig( + hidden_size=config.hidden_size, + intermediate_size=config.intermediate_size, + gate_act=config.hidden_act, + bias=False, + ) + + if layer_idx >= config.first_k_dense_replace: + self.mlp = MoE(moe_args, dim=config.hidden_size, hidden_dim=config.moe_intermediate_size) + else: + self.mlp = MLP(mlp_config) + + self.input_layernorm = RMSNorm(RMSNormConfig(hidden_size=config.hidden_size, eps=config.rms_norm_eps)) + self.post_attention_layernorm = RMSNorm(RMSNormConfig(hidden_size=config.hidden_size, eps=config.rms_norm_eps)) + + def set_context_parallel_attributes(self, cp_group: dist.ProcessGroup, cp_rank: int, cp_world_size: int) -> None: + self._cp_group = cp_group + self._cp_rank = cp_rank + self._cp_world_size = cp_world_size + self.self_attn.set_context_parallel_attributes(cp_group, cp_rank, cp_world_size) + + @deprecate_kwarg("past_key_value", new_name="past_key_values", version="4.58") + def forward( + self, + hidden_states: torch.Tensor, + position_embeddings: Optional[tuple[torch.Tensor, torch.Tensor]] = None, + ks: Optional[torch.Tensor] = None, + ke: Optional[torch.Tensor] = None, + routed_experts: Optional[torch.LongTensor] = None, + ) -> torch.Tensor: + residual = hidden_states + hidden_states = self.input_layernorm(hidden_states) + hidden_states, _ = self.self_attn( + hidden_states=hidden_states, + position_embeddings=position_embeddings, + ks=ks, + ke=ke, + ) + hidden_states = residual + hidden_states + + residual = hidden_states + hidden_states = self.post_attention_layernorm(hidden_states) + hidden_states = self.mlp(hidden_states, routed_experts=routed_experts) + hidden_states = residual + hidden_states + return hidden_states + + +@auto_docstring +class GlmMoeDsaPreTrainedModel(PreTrainedModelPrimeRL): + config: GlmMoeDsaConfig + base_model_prefix = "model" + supports_gradient_checkpointing = True + _no_split_modules = ["GlmMoeDsaDecoderLayer"] + _skip_keys_device_placement = ["past_key_values"] + _supports_flash_attn = True + _supports_sdpa = True + _supports_flex_attn = True + _can_compile_fullgraph = False + _supports_attention_backend = True + _can_record_outputs = { + "hidden_states": GlmMoeDsaDecoderLayer, + } + + def _init_weights(self, module): + super()._init_weights(module) + + @classmethod + def is_hf_state_dict(cls, state_dict: dict[str, Tensor]) -> bool: + return any("mlp.experts.1.up_proj" in name or "mlp.experts.gate_up_proj" in name for name in state_dict.keys()) + + @classmethod + def is_prime_state_dict(cls, state_dict: dict[str, Tensor]) -> bool: + return any("mlp.experts.w1" in module_name for module_name in state_dict.keys()) + + @classmethod + def convert_to_hf(cls, state_dict: dict[str, Tensor]) -> dict[str, Tensor]: + convert_tt_to_hf_moe(state_dict) + return state_dict + + @classmethod + def convert_to_prime(cls, state_dict: dict[str, Tensor]) -> dict[str, Tensor]: + convert_hf_to_tt_moe(state_dict) + return state_dict + + @classmethod + def convert_layer_to_hf(cls, state_dict: dict[str, Tensor], layer_idx: int) -> dict[str, Tensor]: + convert_tt_layer_to_hf(state_dict, layer_idx) + return state_dict + + @classmethod + def convert_layer_to_prime(cls, state_dict: dict[str, Tensor], layer_idx: int) -> dict[str, Tensor]: + convert_hf_layer_to_tt(state_dict, layer_idx) + return state_dict + + @classmethod + def convert_layer_to_vllm_kernel( + cls, state_dict: dict[str, Tensor], layer_idx: int, quantize_fp8: bool = False + ) -> dict[str, Tensor]: + return convert_tt_layer_to_vllm_kernel(state_dict, layer_idx, quantize_fp8=quantize_fp8) + + +@auto_docstring +class GlmMoeDsaModel(GlmMoeDsaPreTrainedModel): + def __init__(self, config: GlmMoeDsaConfig): + super().__init__(config) + + self.padding_idx = config.pad_token_id + self.vocab_size = config.vocab_size + + self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx) + self.layers = nn.ModuleList( + [GlmMoeDsaDecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)] + ) + self.norm = RMSNorm(RMSNormConfig(hidden_size=config.hidden_size, eps=config.rms_norm_eps)) + + rope_parameters = getattr(config, "rope_parameters", None) or {} + rope_type = rope_parameters.get("rope_type", "default") if isinstance(rope_parameters, dict) else "default" + rotary_config = RotaryEmbeddingConfig( + max_position_embeddings=config.max_position_embeddings, + rope_type=rope_type, + model_config=config, + ) + self.rotary_emb = RotaryEmbedding(rotary_config) + self.gradient_checkpointing = False + + self.post_init() + + def _context_parallel_state(self) -> tuple[dist.ProcessGroup | None, int, int]: + if len(self.layers) == 0: + return None, 0, 1 + + layer = self.layers[0] + return getattr(layer, "_cp_group", None), getattr(layer, "_cp_rank", 0), getattr(layer, "_cp_world_size", 1) + + def _gather_position_ids_for_cp( + self, + position_ids: torch.LongTensor, + cp_group: dist.ProcessGroup, + cp_world_size: int, + ) -> torch.LongTensor: + gathered_position_ids = [torch.empty_like(position_ids) for _ in range(cp_world_size)] + dist.all_gather(gathered_position_ids, position_ids.contiguous(), group=cp_group) + return torch.cat(gathered_position_ids, dim=1) + + @auto_docstring + def forward( + self, + input_ids: Optional[torch.LongTensor] = None, + position_ids: Optional[torch.LongTensor] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + routed_experts: Optional[torch.LongTensor] = None, + ) -> BaseModelOutputWithPast: + """ + routed_experts (`torch.LongTensor` of shape `(batch_size, sequence_length, num_hidden_layers, num_experts_per_tok)`, *optional*): + Routed experts for each token in the sequence. Only used for router replay. + """ + if (input_ids is None) ^ (inputs_embeds is not None): + raise ValueError("You must specify exactly one of input_ids or inputs_embeds") + + if inputs_embeds is None: + inputs_embeds: torch.Tensor = self.embed_tokens(input_ids) + + cp_group, cp_rank, cp_world_size = self._context_parallel_state() + if cp_group is not None and cp_world_size > 1: + position_ids_full = self._gather_position_ids_for_cp(position_ids, cp_group, cp_world_size) + else: + position_ids_full = position_ids + + flat_position_ids = position_ids_full.view(-1) + S_full = flat_position_ids.shape[0] + ks_full = torch.arange(S_full, dtype=torch.int32, device=flat_position_ids.device) - flat_position_ids.to( + torch.int32 + ) + ke_full = torch.arange(1, S_full + 1, dtype=torch.int32, device=flat_position_ids.device) + + # Position embeddings are computed over the full sequence: K uses cos/sin[0:S_full] + # while Q uses the local CP slice. ks/ke are computed in K's global coordinate + # system and then sharded to the local Q range so the indexer's per-token + # causal/varlen mask aligns with the gathered K. + hidden_states = inputs_embeds + position_embeddings = self.rotary_emb(hidden_states, position_ids_full) + + if cp_world_size > 1: + s_local = S_full // cp_world_size + ks = ks_full[cp_rank * s_local : (cp_rank + 1) * s_local].contiguous() + ke = ke_full[cp_rank * s_local : (cp_rank + 1) * s_local].contiguous() + else: + ks, ke = ks_full, ke_full + + for layer_idx, decoder_layer in enumerate(self.layers[: self.config.num_hidden_layers]): + routed_experts_layer = routed_experts[:, :, layer_idx, :] if routed_experts is not None else None + hidden_states = decoder_layer( + hidden_states, + position_embeddings=position_embeddings, + ks=ks, + ke=ke, + routed_experts=routed_experts_layer, + ) + + hidden_states = self.norm(hidden_states) + return BaseModelOutputWithPast(last_hidden_state=hidden_states) + + +@auto_docstring +class GlmMoeDsaForCausalLM(GlmMoeDsaPreTrainedModel, GenerationMixin): + _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} + _tp_plan = {"lm_head": "colwise_rep"} + _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + + def __init__(self, config): + super().__init__(config) + self.model = GlmMoeDsaModel(config) + self.vocab_size = config.vocab_size + self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False) + + warnings.warn("GlmMoeDsaForCausalLM is experimental, higher trainer<->inference KL mismatch may be observed.") + warnings.warn("`model.attn` is ignored, GlmMoeDsa uses only sparse attention.") + + self.post_init() + + @auto_docstring + def forward( + self, + input_ids: Optional[torch.LongTensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_values: Optional[Cache] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + labels: Optional[torch.LongTensor] = None, + use_cache: Optional[bool] = None, + cache_position: Optional[torch.LongTensor] = None, + logits_to_keep: Union[int, torch.Tensor] = 0, + temperature: Optional[torch.Tensor] = None, + routed_experts: Optional[torch.LongTensor] = None, + **kwargs: Unpack[TransformersKwargs], + ) -> PrimeLmOutput: + r""" + cache_position (`torch.LongTensor` of shape `(sequence_length)`, *optional*): + Indices of input tokens in the KV cache. Accepted only for HuggingFace API + compatibility — prime-rl asserts `use_cache is None` since training does not + perform autoregressive decoding, so this argument is unused. + labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): + Labels used by PrimeRL's wrapped LM head to optionally compute per-token logprobs/entropy. + temperature (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*): + Per-token temperatures for logprobs/entropy computation when `labels` are provided. + routed_experts (`torch.LongTensor` of shape `(batch_size, sequence_length, num_hidden_layers, num_experts_per_tok)`, *optional*): + Routed experts for each token in the sequence. Only used for router replay. + """ + assert use_cache is None, "use_cache is not supported for custom glm_moe_dsa for now" + assert past_key_values is None, "past_key_values is not supported for custom glm_moe_dsa for now" + + if position_ids is None: + if inputs_embeds is not None: + position_ids = torch.arange(inputs_embeds.shape[1], device=inputs_embeds.device).unsqueeze(0) + else: + position_ids = torch.arange(input_ids.shape[1], device=input_ids.device).unsqueeze(0) + + outputs: BaseModelOutputWithPast = self.model( + input_ids=input_ids, + position_ids=position_ids, + inputs_embeds=inputs_embeds, + routed_experts=routed_experts, + ) + + hidden_states = outputs.last_hidden_state + slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep + return self.lm_head( + hidden_states[:, slice_indices, :], + labels[:, slice_indices] if labels is not None else None, + temperature=temperature, + ) + + def init_buffers_post_meta(self): + buffer_names = [name for name, _ in self.named_buffers()] + if "model.rotary_emb.inv_freq" in buffer_names: + rotary_emb = self.model.rotary_emb + inv_freq, rotary_emb.attention_scaling = rotary_emb.rope_init_fn( + rotary_emb.config, rotary_emb.inv_freq.device + ) + rotary_emb.inv_freq.copy_(inv_freq) + + +__all__ = ["GlmMoeDsaConfig", "GlmMoeDsaPreTrainedModel", "GlmMoeDsaModel", "GlmMoeDsaForCausalLM"] diff --git a/src/prime_rl/trainer/models/glm_moe_dsa/sparse_mla_attention.py b/src/prime_rl/trainer/models/glm_moe_dsa/sparse_mla_attention.py new file mode 100644 index 00000000000..2488505e9f1 --- /dev/null +++ b/src/prime_rl/trainer/models/glm_moe_dsa/sparse_mla_attention.py @@ -0,0 +1,264 @@ +from dataclasses import dataclass + +import torch +import torch.distributed as dist +from torch import nn + +from prime_rl.trainer.models.kernels.fp8_indexer import fp8_indexer +from prime_rl.trainer.models.layers.norms import LayerNorm, RMSNorm, RMSNormConfig +from prime_rl.trainer.models.layers.rotary_emb import rotate_half +from prime_rl.utils.cp import gather_for_cp + +try: + from prime_rl.trainer.models.kernels.sparse_mla_bwd import sparse_mla_bwd + from prime_rl.trainer.models.kernels.sparse_mla_fwd import sparse_mla_fwd_interface +except ImportError: + sparse_mla_fwd_interface = None # type: ignore + sparse_mla_bwd = None # type: ignore + + +@dataclass(frozen=True) +class SparseMlaAttentionArgs: + hidden_size: int + num_attention_heads: int + kv_lora_rank: int + q_lora_rank: int + qk_rope_head_dim: int + qk_nope_head_dim: int + qk_head_dim: int + v_head_dim: int + attention_bias: bool + rms_norm_eps: float + index_n_heads: int + index_head_dim: int + index_topk: int + + +class _SparseMLA(torch.autograd.Function): + """Autograd wrapper for tilelang sparse MLA forward/backward kernels.""" + + @staticmethod + def forward(ctx, q, kv, indices, sm_scale): + out, lse = sparse_mla_fwd_interface(q, kv, indices, sm_scale=sm_scale) + ctx.save_for_backward(q, kv, out, indices, lse) + ctx.sm_scale = sm_scale + return out + + @staticmethod + def backward(ctx, do): + q, kv, out, indices, lse = ctx.saved_tensors + dq, dkv = sparse_mla_bwd(q, kv, out, do.contiguous(), indices, lse, sm_scale=ctx.sm_scale) + return dq, dkv, None, None + + +def apply_rope_interleave_single( + t: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor, unsqueeze_dim: int = 1 +) -> torch.Tensor: + cos = cos.unsqueeze(unsqueeze_dim) + sin = sin.unsqueeze(unsqueeze_dim) + b, h, s, d = t.shape + t = t.view(b, h, s, d // 2, 2).transpose(4, 3).reshape(b, h, s, d) + return (t * cos) + (rotate_half(t) * sin) + + +class Indexer(nn.Module): + def __init__(self, args: SparseMlaAttentionArgs): + super().__init__() + self.n_head = args.index_n_heads + self.head_dim = args.index_head_dim + self.rope_dim = args.qk_rope_head_dim + self.wq_b = nn.Linear(args.q_lora_rank, self.n_head * self.head_dim, bias=False) + self.wk = nn.Linear(args.hidden_size, self.head_dim, bias=args.attention_bias) + self.k_norm = LayerNorm(dim=self.head_dim, eps=1e-6) + self.weights_proj = nn.Linear(args.hidden_size, self.n_head, bias=False) + self.weight_scale = (self.head_dim**-0.5) * (self.n_head**-0.5) + + @torch.no_grad() + def compute_sparse_indices( + self, + hidden_states_local: torch.Tensor, + q_latent_local: torch.Tensor, + ks: torch.Tensor, + ke: torch.Tensor, + index_topk: int, + position_embeddings_full: tuple[torch.Tensor, torch.Tensor], + cp_group: dist.ProcessGroup | None, + cp_world_size: int, + cp_rank: int, + ) -> torch.Tensor: + assert index_topk % 64 == 0, f"index_topk must be divisible by 64 (block_I), got {index_topk}" + + s_local = hidden_states_local.shape[1] + + q_idx = self.wq_b(q_latent_local[0]).view(s_local, self.n_head, self.head_dim) + k_idx_local = self.k_norm(self.wk(hidden_states_local[0])) + w = self.weights_proj(hidden_states_local[0]) + + if cp_world_size > 1: + k_idx = gather_for_cp(k_idx_local.unsqueeze(0), cp_group).squeeze(0) + else: + k_idx = k_idx_local + + cos_full, sin_full = position_embeddings_full + if cp_world_size > 1: + cos_local = cos_full[:, cp_rank * s_local : (cp_rank + 1) * s_local, :] + sin_local = sin_full[:, cp_rank * s_local : (cp_rank + 1) * s_local, :] + else: + cos_local, sin_local = cos_full, sin_full + + q_pe = q_idx[..., : self.rope_dim] + q_nope = q_idx[..., self.rope_dim :] + k_pe = k_idx[..., : self.rope_dim] + k_nope = k_idx[..., self.rope_dim :] + + q_pe = q_pe.unsqueeze(0).transpose(1, 2) # [1, H, S_local, rope_dim] + q_pe = apply_rope_interleave_single(q_pe, cos_local, sin_local) + q_pe = q_pe.transpose(1, 2).squeeze(0) + + k_pe = k_pe.unsqueeze(0).unsqueeze(1) # [1, 1, S_full, rope_dim] + k_pe = apply_rope_interleave_single(k_pe, cos_full, sin_full) + k_pe = k_pe.squeeze(1).squeeze(0) + + q_idx = torch.cat([q_pe, q_nope], dim=-1) + k_idx = torch.cat([k_pe, k_nope], dim=-1) + + indices = fp8_indexer(q_idx, k_idx, w, ks, ke, index_topk, self.weight_scale) + # indices shape: [S_local, topk] in K's coordinate space (sentinel = s_full) + # KV passed to sparse MLA has length s_full + 1 (sentinel zeros at index s_full). + return indices.view(1, s_local, 1, index_topk) + + +class GlmMoeDsaAttention(nn.Module): + def __init__(self, args: SparseMlaAttentionArgs): + super().__init__() + self.args = args + self.num_heads = args.num_attention_heads + self.kv_lora_rank = args.kv_lora_rank + self.qk_rope_head_dim = args.qk_rope_head_dim + self.qk_nope_head_dim = args.qk_nope_head_dim + self.qk_head_dim = args.qk_head_dim + self.v_head_dim = args.v_head_dim + + self.q_a_proj = nn.Linear(args.hidden_size, args.q_lora_rank, bias=args.attention_bias) + self.q_a_layernorm = RMSNorm(RMSNormConfig(hidden_size=args.q_lora_rank, eps=args.rms_norm_eps)) + self.q_b_proj = nn.Linear(args.q_lora_rank, self.num_heads * self.qk_head_dim, bias=False) + + self.kv_a_proj_with_mqa = nn.Linear( + args.hidden_size, + self.kv_lora_rank + self.qk_rope_head_dim, + bias=args.attention_bias, + ) + self.kv_a_layernorm = RMSNorm(RMSNormConfig(hidden_size=self.kv_lora_rank, eps=args.rms_norm_eps)) + self.kv_b_proj = nn.Linear( + self.kv_lora_rank, + self.num_heads * (self.qk_nope_head_dim + self.v_head_dim), + bias=False, + ) + + self.o_proj = nn.Linear(self.num_heads * self.v_head_dim, args.hidden_size, bias=args.attention_bias) + self.indexer = Indexer(args) + self.scaling = self.qk_head_dim ** (-0.5) + + self._cp_group: dist.ProcessGroup | None = None + self._cp_rank: int = 0 + self._cp_world_size: int = 1 + + def set_context_parallel_attributes(self, cp_group: dist.ProcessGroup, cp_rank: int, cp_world_size: int) -> None: + self._cp_group = cp_group + self._cp_rank = cp_rank + self._cp_world_size = cp_world_size + + @property + def cp_enabled(self) -> bool: + return self._cp_world_size > 1 + + def attn_projections(self, hidden_states: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + q_latent = self.q_a_layernorm(self.q_a_proj(hidden_states)) + compressed_kv = self.kv_a_proj_with_mqa(hidden_states) + k_compressed, k_rope = compressed_kv.split([self.kv_lora_rank, self.qk_rope_head_dim], dim=-1) + return q_latent, self.kv_a_layernorm(k_compressed), k_rope + + def mla_up_proj( + self, + q_latent_local: torch.Tensor, + k_compressed_normed_full: torch.Tensor, + k_rope_full: torch.Tensor, + position_embeddings_full: tuple[torch.Tensor, torch.Tensor], + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + batch_size, s_local, _ = q_latent_local.shape + s_full = k_compressed_normed_full.shape[1] + + q_full = self.q_b_proj(q_latent_local).view(batch_size, s_local, self.num_heads, self.qk_head_dim) + q_nope, q_rope = q_full.split([self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1) + + cos_full, sin_full = position_embeddings_full + if self.cp_enabled: + cos_local = cos_full[:, self._cp_rank * s_local : (self._cp_rank + 1) * s_local, :] + sin_local = sin_full[:, self._cp_rank * s_local : (self._cp_rank + 1) * s_local, :] + else: + cos_local, sin_local = cos_full, sin_full + + q_rope_r = q_rope.transpose(1, 2) # [B, H, S_local, rope_dim] + q_rope_r = apply_rope_interleave_single(q_rope_r, cos_local, sin_local) + q_rope = q_rope_r.transpose(1, 2) + + k_rope_r = k_rope_full.unsqueeze(1) # [B, 1, S_full, rope_dim] + k_rope_r = apply_rope_interleave_single(k_rope_r, cos_full, sin_full) + k_rope_full = k_rope_r.squeeze(1) + + kv_b_w = self.kv_b_proj.weight.view(self.num_heads, self.qk_nope_head_dim + self.v_head_dim, self.kv_lora_rank) + w_k_nope = kv_b_w[:, : self.qk_nope_head_dim, :] + w_v = kv_b_w[:, self.qk_nope_head_dim :, :] + q_absorbed = torch.einsum("bshd,hdk->bshk", q_nope, w_k_nope) + + sparse_q = torch.cat([q_absorbed, q_rope], dim=-1) + sparse_kv = torch.cat([k_compressed_normed_full, k_rope_full], dim=-1).unsqueeze(2) + + sentinel = torch.zeros(batch_size, 1, 1, sparse_kv.shape[-1], dtype=sparse_kv.dtype, device=sparse_kv.device) + sparse_kv = torch.cat([sparse_kv, sentinel], dim=1) + assert sparse_kv.shape[1] == s_full + 1 + return sparse_q, sparse_kv, w_v + + def _mla_unabsorb(self, out: torch.Tensor, w_v: torch.Tensor) -> torch.Tensor: + return torch.einsum("bshk,hdk->bshd", out, w_v) + + def output_proj(self, attn_output: torch.Tensor, w_v: torch.Tensor) -> torch.Tensor: + attn_output = self._mla_unabsorb(attn_output, w_v) + batch_size, total_tokens = attn_output.shape[:2] + attn_output = attn_output.reshape(batch_size, total_tokens, -1) + return self.o_proj(attn_output) + + def forward( + self, + hidden_states: torch.Tensor, + position_embeddings: tuple[torch.Tensor, torch.Tensor], + ks: torch.Tensor | None = None, + ke: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, torch.Tensor | None]: + q_latent, k_compressed_normed, k_rope = self.attn_projections(hidden_states) + + if self.cp_enabled: + k_compressed_normed = gather_for_cp(k_compressed_normed, self._cp_group) + k_rope = gather_for_cp(k_rope, self._cp_group) + + indices = self.indexer.compute_sparse_indices( + hidden_states_local=hidden_states, + q_latent_local=q_latent, + ks=ks, + ke=ke, + index_topk=self.args.index_topk, + position_embeddings_full=position_embeddings, + cp_group=self._cp_group, + cp_world_size=self._cp_world_size, + cp_rank=self._cp_rank, + ) + + sparse_q, sparse_kv, w_v = self.mla_up_proj( + q_latent_local=q_latent, + k_compressed_normed_full=k_compressed_normed, + k_rope_full=k_rope, + position_embeddings_full=position_embeddings, + ) + + out = _SparseMLA.apply(sparse_q, sparse_kv, indices, self.scaling) + return self.output_proj(out, w_v), None diff --git a/src/prime_rl/trainer/models/gpt_oss/__init__.py b/src/prime_rl/trainer/models/gpt_oss/__init__.py new file mode 100644 index 00000000000..6c72b3a1d7f --- /dev/null +++ b/src/prime_rl/trainer/models/gpt_oss/__init__.py @@ -0,0 +1,13 @@ +from prime_rl.trainer.models.gpt_oss.configuration_gpt_oss import GptOssConfig +from prime_rl.trainer.models.gpt_oss.modeling_gpt_oss import ( + GptOssForCausalLM, + GptOssModel, + GptOssPreTrainedModel, +) + +__all__ = [ + "GptOssConfig", + "GptOssForCausalLM", + "GptOssModel", + "GptOssPreTrainedModel", +] diff --git a/src/prime_rl/trainer/models/gpt_oss/configuration_gpt_oss.py b/src/prime_rl/trainer/models/gpt_oss/configuration_gpt_oss.py new file mode 100644 index 00000000000..b22850b9b4f --- /dev/null +++ b/src/prime_rl/trainer/models/gpt_oss/configuration_gpt_oss.py @@ -0,0 +1,6 @@ +# Re-export HF's GptOssConfig so the parameter naming and rope/sliding-window defaults +# stay in sync with upstream. We don't need to subclass it: the only PrimeRL-specific +# behavior lives in modeling_gpt_oss.py (custom MoE for LoRA detection). +from transformers.models.gpt_oss.configuration_gpt_oss import GptOssConfig + +__all__ = ["GptOssConfig"] diff --git a/src/prime_rl/trainer/models/gpt_oss/converting_gpt_oss.py b/src/prime_rl/trainer/models/gpt_oss/converting_gpt_oss.py new file mode 100644 index 00000000000..b91d7326f8f --- /dev/null +++ b/src/prime_rl/trainer/models/gpt_oss/converting_gpt_oss.py @@ -0,0 +1,28 @@ +"""State-dict conversion for GPT-OSS. + +The custom prime-rl GPT-OSS implementation mirrors HuggingFace's parameter naming +exactly (gate_up_proj/gate_up_proj_bias/down_proj/down_proj_bias plus router.weight +and router.bias as nn.Parameters). So loading the unsloth BF16 checkpoint requires +no key conversion - HF and prime formats are identical for this model. +""" + +from torch import Tensor + + +def is_hf_state_dict(state_dict: dict[str, Tensor]) -> bool: + return any("mlp.experts.gate_up_proj" in name for name in state_dict.keys()) + + +def is_prime_state_dict(state_dict: dict[str, Tensor]) -> bool: + # Prime format equals HF format for GPT-OSS, so we never claim to be a separate + # prime format - this disables the auto-conversion path in load_dcp_from_hf and + # lets DCP load HF safetensors directly. + return False + + +def convert_to_hf(state_dict: dict[str, Tensor]) -> dict[str, Tensor]: + return state_dict + + +def convert_to_prime(state_dict: dict[str, Tensor]) -> dict[str, Tensor]: + return state_dict diff --git a/src/prime_rl/trainer/models/gpt_oss/modeling_gpt_oss.py b/src/prime_rl/trainer/models/gpt_oss/modeling_gpt_oss.py new file mode 100644 index 00000000000..75a60a9ec9e --- /dev/null +++ b/src/prime_rl/trainer/models/gpt_oss/modeling_gpt_oss.py @@ -0,0 +1,356 @@ +# GPT-OSS model with PrimeRL custom MoE. +# +# This file mirrors HuggingFace's modeling_gpt_oss.py but swaps the experts module for +# `GptOssGroupedExperts` so the LoRA detection logic in `prime_rl/trainer/lora.py` can +# find and adapt the expert weights. Attention, rotary embedding, and RMSNorm are reused +# from HF as-is - they correctly handle attention sinks and gpt-oss's split-rotate RoPE. + +from typing import Union + +import torch +import torch.nn.functional as F +from torch import Tensor, nn +from transformers.cache_utils import Cache +from transformers.generation import GenerationMixin +from transformers.masking_utils import create_causal_mask, create_sliding_window_causal_mask +from transformers.modeling_layers import GradientCheckpointingLayer +from transformers.modeling_outputs import MoeModelOutputWithPast +from transformers.models.gpt_oss.modeling_gpt_oss import ( + GptOssAttention, + GptOssRMSNorm, + GptOssRotaryEmbedding, +) +from transformers.processing_utils import Unpack +from transformers.utils import TransformersKwargs, auto_docstring, can_return_tuple + +from prime_rl.trainer.models.base import PreTrainedModelPrimeRL +from prime_rl.trainer.models.gpt_oss.configuration_gpt_oss import GptOssConfig +from prime_rl.trainer.models.gpt_oss.converting_gpt_oss import ( + convert_to_hf, + convert_to_prime, + is_hf_state_dict, + is_prime_state_dict, +) +from prime_rl.trainer.models.layers.lm_head import PrimeLmOutput +from prime_rl.trainer.models.layers.moe import GptOssGroupedExperts + + +class GptOssTopKRouter(nn.Module): + """Token-choice top-k router matching HF's GptOssTopKRouter parameter naming. + + Stores `weight` (num_experts, hidden_size) and `bias` (num_experts) as raw nn.Parameters + so the unsloth BF16 checkpoint loads with no key conversion. Returns the same outputs as + `TokenChoiceTopKRouter` so the surrounding MoE plumbing matches the rest of the repo. + """ + + def __init__(self, config: GptOssConfig): + super().__init__() + self.num_experts = config.num_local_experts + self.top_k = config.num_experts_per_tok + self.hidden_size = config.hidden_size + self.weight = nn.Parameter(torch.empty(self.num_experts, self.hidden_size)) + self.bias = nn.Parameter(torch.empty(self.num_experts)) + + def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """x: (T, hidden). Returns (top_scores, top_indices, num_tokens_per_expert). + + Matches HF's softmax-after-topk semantics: we softmax the top-k logits only, + not the full distribution. + """ + logits = F.linear(x, self.weight, self.bias) # (T, num_experts) + top_logits, top_indices = torch.topk(logits, self.top_k, dim=-1) + top_scores = F.softmax(top_logits, dim=-1, dtype=top_logits.dtype) + num_tokens_per_expert = torch.histc( + top_indices.reshape(-1).float(), + bins=self.num_experts, + min=0, + max=self.num_experts, + ) + return top_scores, top_indices, num_tokens_per_expert + + def init_weights(self, init_std: float): + nn.init.trunc_normal_(self.weight, mean=0.0, std=init_std) + nn.init.zeros_(self.bias) + + +class GptOssMoE(nn.Module): + """GPT-OSS sparse MLP: top-k router + grouped experts with biases. + + Wraps `GptOssGroupedExperts` (which expects pre-permuted tokens + num_tokens_per_expert) + with the standard MoE permute/unpermute pipeline. Routing weights are applied AFTER + the expert computation, matching HF's `routing_weights[token_idx, top_k_pos]` index_add. + """ + + def __init__(self, config: GptOssConfig): + super().__init__() + # GptOssGroupedExperts has fused gate_up_proj + per-expert biases; the + # current FP8 grouped-GEMM path doesn't model that yet. + assert not getattr(config, "fp8", False), "FP8 training is not supported for GPT-OSS" + self.num_experts = config.num_local_experts + self.top_k = config.num_experts_per_tok + self.router = GptOssTopKRouter(config) + self.experts = GptOssGroupedExperts( + hidden_size=config.hidden_size, + intermediate_size=config.intermediate_size, + num_experts=self.num_experts, + use_grouped_mm=getattr(config, "use_grouped_mm", True), + ) + + def forward(self, hidden_states: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + bs, slen, dim = hidden_states.shape + x = hidden_states.reshape(-1, dim) # (T, dim) + + top_scores, top_indices, num_tokens_per_expert = self.router(x) + # top_scores, top_indices: (T, top_k) + + # Sort tokens by expert. Each token contributes top_k entries (one per chosen expert). + flat_expert_indices = top_indices.reshape(-1) # (T*top_k,) + sorted_perm = torch.argsort(flat_expert_indices, stable=True) # (T*top_k,) + # token id of each (token, k) entry, in expert-sorted order + token_indices_experts_sorted = sorted_perm // self.top_k + # routing weight for each (token, k) entry, in expert-sorted order + top_scores_experts_sorted = top_scores.reshape(-1)[sorted_perm] + + # Gather inputs for each expert in expert-sorted order + gather_indices = token_indices_experts_sorted.reshape(-1, 1).expand(-1, dim) + routed_input = torch.gather(x, dim=0, index=gather_indices) # (T*top_k, dim) + + routed_output = self.experts(routed_input, num_tokens_per_expert) # (T*top_k, dim) + + # Apply routing weights post-experts (HF: weighted_output = out * routing_weights) + routed_output = (routed_output.float() * top_scores_experts_sorted.reshape(-1, 1)).to(routed_output.dtype) + + # Scatter-add back to original token positions + out = torch.zeros_like(x) + scatter_indices = token_indices_experts_sorted.reshape(-1, 1).expand(-1, dim) + out = out.scatter_add(dim=0, index=scatter_indices, src=routed_output) + + out = out.reshape(bs, slen, dim) + return out, top_scores # router_scores returned only for compatibility; trainer ignores + + +class GptOssDecoderLayer(GradientCheckpointingLayer): + def __init__(self, config: GptOssConfig, layer_idx: int): + super().__init__() + self.hidden_size = config.hidden_size + self.self_attn = GptOssAttention(config=config, layer_idx=layer_idx) + self.mlp = GptOssMoE(config) + self.input_layernorm = GptOssRMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.post_attention_layernorm = GptOssRMSNorm(config.hidden_size, eps=config.rms_norm_eps) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: torch.Tensor | None = None, + position_ids: torch.LongTensor | None = None, + past_key_values: Cache | None = None, + use_cache: bool | None = False, + position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> torch.Tensor: + residual = hidden_states + hidden_states = self.input_layernorm(hidden_states) + hidden_states, _ = self.self_attn( + hidden_states=hidden_states, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=past_key_values, + use_cache=use_cache, + position_embeddings=position_embeddings, + **kwargs, + ) + hidden_states = residual + hidden_states + + residual = hidden_states + hidden_states = self.post_attention_layernorm(hidden_states) + hidden_states, _ = self.mlp(hidden_states) + hidden_states = residual + hidden_states + return hidden_states + + +@auto_docstring +class GptOssPreTrainedModel(PreTrainedModelPrimeRL): + config: GptOssConfig + base_model_prefix = "model" + supports_gradient_checkpointing = True + _no_split_modules = ["GptOssDecoderLayer"] + _skip_keys_device_placement = ["past_key_values"] + _supports_flash_attn = True + _supports_sdpa = False + _supports_flex_attn = True + _can_compile_fullgraph = False + _supports_attention_backend = True + _keep_in_fp32_modules = ["post_attention_layernorm", "input_layernorm", "norm"] + _compatible_flash_implementations = ["kernels-community/vllm-flash-attn3", "flash_attention_4"] + _can_record_outputs = {"hidden_states": GptOssDecoderLayer} + + @classmethod + def is_hf_state_dict(cls, state_dict: dict[str, Tensor]) -> bool: + return is_hf_state_dict(state_dict) + + @classmethod + def is_prime_state_dict(cls, state_dict: dict[str, Tensor]) -> bool: + return is_prime_state_dict(state_dict) + + @classmethod + def convert_to_hf(cls, state_dict: dict[str, Tensor]) -> dict[str, Tensor]: + return convert_to_hf(state_dict) + + @classmethod + def convert_to_prime(cls, state_dict: dict[str, Tensor]) -> dict[str, Tensor]: + return convert_to_prime(state_dict) + + +@auto_docstring +class GptOssModel(GptOssPreTrainedModel): + _no_split_modules = ["GptOssDecoderLayer"] + + def __init__(self, config: GptOssConfig): + super().__init__(config) + self.padding_idx = config.pad_token_id + self.vocab_size = config.vocab_size + + self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx) + self.layers = nn.ModuleList( + [GptOssDecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)] + ) + self.norm = GptOssRMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.rotary_emb = GptOssRotaryEmbedding(config=config) + self.gradient_checkpointing = False + + self.post_init() + + @auto_docstring + def forward( + self, + input_ids: torch.LongTensor | None = None, + attention_mask: torch.Tensor | None = None, + position_ids: torch.LongTensor | None = None, + past_key_values: Cache | None = None, + inputs_embeds: torch.FloatTensor | None = None, + use_cache: bool | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> MoeModelOutputWithPast: + if (input_ids is None) ^ (inputs_embeds is not None): + raise ValueError("You must specify exactly one of input_ids or inputs_embeds") + + if inputs_embeds is None: + inputs_embeds = self.embed_tokens(input_ids) + + if position_ids is None: + position_ids = torch.arange(inputs_embeds.shape[1], device=inputs_embeds.device).unsqueeze(0) + + if not isinstance(causal_mask_mapping := attention_mask, dict): + mask_kwargs = { + "config": self.config, + "inputs_embeds": inputs_embeds, + "attention_mask": attention_mask, + "past_key_values": past_key_values, + "position_ids": position_ids, + } + causal_mask_mapping = { + "full_attention": create_causal_mask(**mask_kwargs), + "sliding_attention": create_sliding_window_causal_mask(**mask_kwargs), + } + + hidden_states = inputs_embeds + position_embeddings = self.rotary_emb(hidden_states, position_ids) + + for i, decoder_layer in enumerate(self.layers): + hidden_states = decoder_layer( + hidden_states, + attention_mask=causal_mask_mapping[self.config.layer_types[i]], + position_embeddings=position_embeddings, + position_ids=position_ids, + past_key_values=past_key_values, + use_cache=use_cache, + **kwargs, + ) + hidden_states = self.norm(hidden_states) + return MoeModelOutputWithPast(last_hidden_state=hidden_states, past_key_values=past_key_values) + + +@auto_docstring +class GptOssForCausalLM(GptOssPreTrainedModel, GenerationMixin): + _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} + _tp_plan = {"lm_head": "colwise_gather_output"} + _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + + def __init__(self, config: GptOssConfig): + super().__init__(config) + self.model = GptOssModel(config) + self.vocab_size = config.vocab_size + self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False) + self.num_experts = config.num_local_experts + self.num_experts_per_tok = config.num_experts_per_tok + + self.post_init() + + def set_decoder(self, decoder): + self.model = decoder + + def get_decoder(self): + return self.model + + @can_return_tuple + @auto_docstring + def forward( + self, + input_ids: torch.LongTensor | None = None, + attention_mask: torch.Tensor | None = None, + position_ids: torch.LongTensor | None = None, + past_key_values: Cache | None = None, + inputs_embeds: torch.FloatTensor | None = None, + labels: torch.LongTensor | None = None, + use_cache: bool | None = None, + logits_to_keep: Union[int, torch.Tensor] = 0, + temperature: torch.Tensor | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> PrimeLmOutput: + r""" + temperature (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*): + Per-token temperatures for logprobs/entropy computation when `labels` are provided. + """ + assert use_cache is None or not use_cache, "use_cache is not supported for custom GPT-OSS" + assert past_key_values is None, "past_key_values is not supported for custom GPT-OSS" + + outputs: MoeModelOutputWithPast = self.model( + input_ids=input_ids, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=past_key_values, + inputs_embeds=inputs_embeds, + use_cache=use_cache, + **kwargs, + ) + + hidden_states = outputs.last_hidden_state + slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep + return self.lm_head( + hidden_states[:, slice_indices, :], + labels[:, slice_indices] if labels is not None else None, + temperature=temperature, + ) + + def init_buffers_post_meta(self): + buffer_names = [name for name, _ in self.named_buffers()] + if "model.rotary_emb.inv_freq" in buffer_names: + rotary_emb = self.model.rotary_emb + from transformers.modeling_rope_utils import ROPE_INIT_FUNCTIONS + + rope_init_fn = ( + ROPE_INIT_FUNCTIONS[rotary_emb.rope_type] + if rotary_emb.rope_type != "default" + else rotary_emb.compute_default_rope_parameters + ) + inv_freq, rotary_emb.attention_scaling = rope_init_fn(rotary_emb.config, rotary_emb.inv_freq.device) + rotary_emb.inv_freq.copy_(inv_freq) + if "model.rotary_emb.original_inv_freq" in buffer_names: + rotary_emb.original_inv_freq.copy_(inv_freq) + + +__all__ = [ + "GptOssForCausalLM", + "GptOssModel", + "GptOssPreTrainedModel", +] diff --git a/src/prime_rl/eval/__init__.py b/src/prime_rl/trainer/models/kernels/__init__.py similarity index 100% rename from src/prime_rl/eval/__init__.py rename to src/prime_rl/trainer/models/kernels/__init__.py diff --git a/src/prime_rl/trainer/models/kernels/fp8_indexer.py b/src/prime_rl/trainer/models/kernels/fp8_indexer.py new file mode 100644 index 00000000000..e2d33de9547 --- /dev/null +++ b/src/prime_rl/trainer/models/kernels/fp8_indexer.py @@ -0,0 +1,239 @@ +"""Triton FP8 indexer kernel for DSA sparse token selection. + +Scoring formula: I_{t,s} = Σ_j w_{t,j} · ReLU(q_{t,j} · k_s) + +FP8 quantization vendored from vLLM (Apache 2.0). +""" + +import torch +import triton +import triton.language as tl + +FP8_MAX = 448.0 +FP8_MIN = -448.0 +FP8_EPS = 1e-10 + + +@triton.jit +def _per_token_group_quant_fp8( + y_ptr, + y_q_ptr, + y_s_ptr, + group_size, + y_num_columns, + y_row_stride, + eps, + fp8_min, + fp8_max, + use_ue8m0: tl.constexpr, + BLOCK: tl.constexpr, +): + groups_per_row = y_num_columns // group_size + + g_id = tl.program_id(0) + row = g_id // groups_per_row + row_g_id = g_id % groups_per_row + + y_ptr_offset = (row.to(tl.int64) * y_row_stride) + (row_g_id.to(tl.int64) * group_size) + y_ptr += y_ptr_offset + + y_q_ptr_offset = g_id.to(tl.int64) * group_size + y_q_ptr += y_q_ptr_offset + y_s_ptr += g_id + + cols = tl.arange(0, BLOCK) + mask = cols < group_size + + y = tl.load(y_ptr + cols, mask=mask, other=0.0).to(tl.float32) + + _absmax = tl.maximum(tl.max(tl.abs(y)), eps) + scale_raw = _absmax / fp8_max + y_s = tl.math.exp2(tl.ceil(tl.log2(scale_raw))) if use_ue8m0 else scale_raw + y_q = tl.clamp(y / y_s, fp8_min, fp8_max).to(y_q_ptr.dtype.element_ty) + + tl.store(y_q_ptr + cols, y_q, mask=mask) + tl.store(y_s_ptr, y_s) + + +def per_token_group_quant_fp8( + x: torch.Tensor, group_size: int, use_ue8m0: bool = True +) -> tuple[torch.Tensor, torch.Tensor]: + assert x.shape[-1] % group_size == 0 + assert x.stride(-1) == 1 + + x_q = torch.empty(x.shape, device=x.device, dtype=torch.float8_e4m3fn) + shape = x.shape[:-1] + (x.shape[-1] // group_size,) + x_s = torch.empty(shape, device=x.device, dtype=torch.float32) + + M = x.numel() // group_size + N = group_size + BLOCK = triton.next_power_of_2(N) + num_warps = min(max(BLOCK // 256, 1), 8) + + _per_token_group_quant_fp8[(M,)]( + x, + x_q, + x_s, + group_size, + x.shape[-1], + x.stride(-2), + FP8_EPS, + FP8_MIN, + FP8_MAX, + use_ue8m0=use_ue8m0, + BLOCK=BLOCK, + num_warps=num_warps, + num_stages=1, + ) + return x_q, x_s + + +@triton.autotune( + configs=[ + triton.Config({"BLOCK_M": 64, "BLOCK_N": 64}, num_warps=4, num_stages=2), + triton.Config({"BLOCK_M": 64, "BLOCK_N": 128}, num_warps=4, num_stages=2), + triton.Config({"BLOCK_M": 128, "BLOCK_N": 64}, num_warps=8, num_stages=2), + triton.Config({"BLOCK_M": 128, "BLOCK_N": 128}, num_warps=8, num_stages=2), + triton.Config({"BLOCK_M": 64, "BLOCK_N": 64}, num_warps=8, num_stages=3), + triton.Config({"BLOCK_M": 64, "BLOCK_N": 128}, num_warps=8, num_stages=3), + ], + key=["S_Q", "S_K"], +) +@triton.jit +def _triton_fp8_indexer_kernel( + Q_fp8, + K_fp8, + K_scales, + W, + Out, + KS, + KE, + S_Q, + S_K, + stride_qh, + stride_qs, + stride_ks, + stride_ws, + H: tl.constexpr, + D: tl.constexpr, + BLOCK_M: tl.constexpr, + BLOCK_N: tl.constexpr, +): + pid_m = tl.program_id(0) + pid_n = tl.program_id(1) + + offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) + offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N) + offs_d = tl.arange(0, D) + + mask_m = offs_m < S_Q + mask_n = offs_n < S_K + + ks_vals = tl.load(KS + offs_m, mask=mask_m, other=S_K) + ke_vals = tl.load(KE + offs_m, mask=mask_m, other=0) + + offs_m_64 = offs_m.to(tl.int64) + offs_n_64 = offs_n.to(tl.int64) + out_ptrs = Out + offs_m_64[:, None] * S_K + offs_n_64[None, :] + out_mask = mask_m[:, None] & mask_n[None, :] + + ks_min = tl.min(ks_vals) + ke_max = tl.max(ke_vals) + n_lo = pid_n * BLOCK_N + + if n_lo >= ke_max or n_lo + BLOCK_N <= ks_min: + tl.store(out_ptrs, tl.full([BLOCK_M, BLOCK_N], float("-inf"), tl.float32), mask=out_mask) + return + + # Clamp indices for FP8 loads (can't use mask+other with fp8 dtype) + offs_n_safe = tl.minimum(offs_n, S_K - 1) + offs_m_safe = tl.minimum(offs_m, S_Q - 1) + + k_block = tl.load(K_fp8 + offs_n_safe[:, None] * stride_ks + offs_d[None, :]) + k_sc = tl.load(K_scales + offs_n_safe, mask=mask_n, other=0.0).to(tl.float32) + + acc = tl.zeros([BLOCK_M, BLOCK_N], dtype=tl.float32) + + for h in range(H): + q_h = tl.load(Q_fp8 + h * stride_qh + offs_m_safe[:, None] * stride_qs + offs_d[None, :]) + w_h = tl.load(W + offs_m * stride_ws + h, mask=mask_m, other=0.0).to(tl.float32) + + scores = tl.dot(q_h, tl.trans(k_block)).to(tl.float32) + scores = tl.maximum(scores, 0.0) * k_sc[None, :] + acc += scores * w_h[:, None] + + valid = (offs_n[None, :] >= ks_vals[:, None]) & (offs_n[None, :] < ke_vals[:, None]) + acc = tl.where(valid, acc, float("-inf")) + + tl.store(out_ptrs, acc, mask=out_mask) + + +def fp8_indexer(q, k, w, ks, ke, topk, weight_scale=1.0): + """Triton FP8 indexer: UE8M0 quantization + fused scoring kernel + topk. + + Args: + q: [S_q, H, D] bf16 query vectors per head + k: [S_k, D] bf16 key vectors (shared across heads) + w: [S_q, H] bf16 per-head weights + ks: [S_q] int32 sequence start per token (in K's coordinate system) + ke: [S_q] int32 causal end per token (in K's coordinate system, = global position + 1) + topk: number of top indices to return + weight_scale: constant scaling factor for weights + + Returns: + [S_q, topk] int32 selected token indices per query (sentinel = S_k) + """ + # NOTE: We don't use weight scale in this kernel as it produces higher KL mismatch for some reason + # This is not a problem result-wise, as it is a constant multiplier + _weight_scale = weight_scale + S_q, H, D = q.shape + S_k = k.shape[0] + device = q.device + + q_flat = q.reshape(S_q * H, D).contiguous() + q_fp8, q_scales = per_token_group_quant_fp8(q_flat, group_size=D, use_ue8m0=True) + k_fp8, k_scales = per_token_group_quant_fp8(k.contiguous(), group_size=D, use_ue8m0=True) + + q_fp8 = q_fp8.view(S_q, H, D).permute(1, 0, 2).contiguous() + + # Fold q_scale into weights + q_scales = q_scales.view(S_q, H) + w = w * q_scales + + logits = torch.empty(S_q, S_k, dtype=torch.float32, device=device) + + grid = lambda meta: ( + triton.cdiv(S_q, meta["BLOCK_M"]), + triton.cdiv(S_k, meta["BLOCK_N"]), + ) + _triton_fp8_indexer_kernel[grid]( + q_fp8, + k_fp8, + k_scales, + w, + logits, + ks, + ke, + S_q, + S_k, + q_fp8.stride(0), + q_fp8.stride(1), + k_fp8.stride(0), + w.stride(0), + H=H, + D=D, + ) + + actual_topk = min(topk, S_k) + _, indices = torch.topk(logits, actual_topk, dim=-1) + if actual_topk < topk: + padding = torch.full((S_q, topk - actual_topk), S_k, dtype=indices.dtype, device=device) + indices = torch.cat([indices, padding], dim=-1) + + # Replace cross-sequence indices with sentinel S_k (maps to zero-valued KV) + ks_exp = ks.unsqueeze(1).expand_as(indices) + ke_exp = ke.unsqueeze(1).expand_as(indices) + out_of_range = (indices < ks_exp) | (indices >= ke_exp) + indices = indices.masked_fill(out_of_range, S_k) + + return indices.to(torch.int32) diff --git a/src/prime_rl/trainer/models/kernels/fp8_utils.py b/src/prime_rl/trainer/models/kernels/fp8_utils.py new file mode 100644 index 00000000000..82c26fb7e4a --- /dev/null +++ b/src/prime_rl/trainer/models/kernels/fp8_utils.py @@ -0,0 +1,552 @@ +from __future__ import annotations + +from typing import Tuple + +import torch +import triton +import triton.language as tl + +FP8_MAX = 448.0 +MIN_SCALE = 1e-4 +GROUP_ALIGNMENT = 128 + + +def ceil_div(x: int, y: int) -> int: + return (x + y - 1) // y + + +def align(x: int, y: int) -> int: + return ceil_div(x, y) * y + + +# --------------------------------------------------------------------------- +# Layout building +# --------------------------------------------------------------------------- + + +def build_grouped_layout(offs: torch.Tensor, *, total_m: int | None = None): + assert offs.dim() == 1 + assert offs.dtype == torch.int32 + device = offs.device + total_m = (total_m if total_m is not None else int(offs[-1].item())) if offs.numel() else 0 + starts_tensor = torch.empty_like(offs) + if offs.numel() > 0: + starts_tensor[0] = 0 + if offs.numel() > 1: + starts_tensor[1:] = offs[:-1] + actual_ms_tensor = offs - starts_tensor + aligned_ms_tensor = ((actual_ms_tensor + GROUP_ALIGNMENT - 1) // GROUP_ALIGNMENT) * GROUP_ALIGNMENT + padded_ends = aligned_ms_tensor.cumsum(0) + block_starts_tensor = (padded_ends - aligned_ms_tensor) // GROUP_ALIGNMENT + ks_tensor = aligned_ms_tensor.contiguous() + padded_total_m = int(padded_ends[-1].item()) if offs.numel() else 0 + total_blocks = padded_total_m // GROUP_ALIGNMENT + grouped_layout = torch.empty((padded_total_m,), dtype=torch.int32, device=device) + block_to_group = torch.empty((total_blocks,), dtype=torch.int32, device=device) + if offs.numel(): + _build_grouped_layout_triton( + grouped_layout, + block_to_group, + starts_tensor, + actual_ms_tensor, + aligned_ms_tensor, + block_starts_tensor, + ) + return ( + total_m, + padded_total_m, + grouped_layout, + block_to_group, + ks_tensor, + starts_tensor, + actual_ms_tensor, + block_starts_tensor, + ) + + +def _build_grouped_layout_triton( + grouped_layout: torch.Tensor, + block_to_group: torch.Tensor, + starts_tensor: torch.Tensor, + actual_ms_tensor: torch.Tensor, + aligned_ms_tensor: torch.Tensor, + block_starts_tensor: torch.Tensor, +) -> None: + _build_grouped_layout_kernel[(actual_ms_tensor.numel(),)]( + grouped_layout, + block_to_group, + starts_tensor, + actual_ms_tensor, + aligned_ms_tensor, + block_starts_tensor, + BLOCK_M=128, + num_warps=4, + ) + + +# --------------------------------------------------------------------------- +# Triton kernels +# --------------------------------------------------------------------------- + + +@triton.jit +def _build_grouped_layout_kernel( + grouped_layout_ptr, + block_to_group_ptr, + starts_ptr, + actual_ms_ptr, + aligned_ms_ptr, + block_starts_ptr, + BLOCK_M: tl.constexpr, +): + pid_g = tl.program_id(axis=0) + actual_m = tl.load(actual_ms_ptr + pid_g) + aligned_m = tl.load(aligned_ms_ptr + pid_g) + block_start = tl.load(block_starts_ptr + pid_g) + dst_start = block_start * BLOCK_M + block_idx = 0 + while block_idx < (aligned_m // BLOCK_M): + row_offsets = block_idx * BLOCK_M + tl.arange(0, BLOCK_M) + values = tl.where(row_offsets < actual_m, pid_g, -1) + tl.store(grouped_layout_ptr + dst_start + row_offsets, values) + tl.store(block_to_group_ptr + block_start + block_idx, pid_g) + block_idx += 1 + + +@triton.jit +def _unpack_grouped_rows_kernel( + x_ptr, + block_to_group_ptr, + starts_ptr, + actual_ms_ptr, + block_starts_ptr, + out_ptr, + cols, + stride_xm, + stride_xn, + stride_ym, + stride_yn, + BLOCK_M: tl.constexpr, + BLOCK_N: tl.constexpr, + GROUP_BLOCK_M: tl.constexpr, +): + pid_blk = tl.program_id(axis=0) + pid_sub = tl.program_id(axis=1) + pid_n = tl.program_id(axis=2) + pid_g = tl.load(block_to_group_ptr + pid_blk) + block_start = tl.load(block_starts_ptr + pid_g) + dst_start = tl.load(starts_ptr + pid_g) + actual_m = tl.load(actual_ms_ptr + pid_g) + row_offsets = (pid_blk - block_start) * GROUP_BLOCK_M + pid_sub * BLOCK_M + tl.arange(0, BLOCK_M) + col_offsets = pid_n * BLOCK_N + tl.arange(0, BLOCK_N) + src_rows_i64 = (pid_blk * GROUP_BLOCK_M + pid_sub * BLOCK_M + tl.arange(0, BLOCK_M)).to(tl.int64) + dst_rows_i64 = (dst_start + row_offsets).to(tl.int64) + col_offsets_i64 = col_offsets.to(tl.int64) + valid_rows = row_offsets < actual_m + valid_cols = col_offsets < cols + x = tl.load( + x_ptr + src_rows_i64[:, None] * stride_xm + col_offsets_i64[None, :] * stride_xn, + mask=valid_rows[:, None] & valid_cols[None, :], + other=0.0, + ) + tl.store( + out_ptr + dst_rows_i64[:, None] * stride_ym + col_offsets_i64[None, :] * stride_yn, + x, + mask=valid_rows[:, None] & valid_cols[None, :], + ) + + +@triton.jit +def _per_token_fp8_kernel( + x_ptr, + out_ptr, + sf_ptr, + rows, + cols, + stride_xm, + stride_xn, + stride_ym, + stride_yn, + stride_sm, + stride_sk, + USE_UE8M0: tl.constexpr, + BLOCK_M: tl.constexpr, + BLOCK_K: tl.constexpr, +): + pid_m = tl.program_id(axis=0) + pid_k = tl.program_id(axis=1) + row_offsets = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) + col_offsets = pid_k * BLOCK_K + tl.arange(0, BLOCK_K) + row_offsets_i64 = row_offsets.to(tl.int64) + col_offsets_i64 = col_offsets.to(tl.int64) + mask = (row_offsets[:, None] < rows) & (col_offsets[None, :] < cols) + x = tl.load( + x_ptr + row_offsets_i64[:, None] * stride_xm + col_offsets_i64[None, :] * stride_xn, + mask=mask, + other=0.0, + ).to(tl.float32) + amax = tl.max(tl.abs(x), axis=1) + scale = tl.maximum(amax / 448.0, 1e-4) + if USE_UE8M0: + scale = tl.exp2(tl.ceil(tl.log2(scale))) + y = x / scale[:, None] + tl.store( + out_ptr + row_offsets_i64[:, None] * stride_ym + col_offsets_i64[None, :] * stride_yn, + y.to(tl.float8e4nv), + mask=mask, + ) + tl.store(sf_ptr + row_offsets_i64 * stride_sm + pid_k * stride_sk, scale, mask=row_offsets < rows) + + +@triton.jit +def _grouped_per_token_fp8_kernel( + x_ptr, + block_to_group_ptr, + starts_ptr, + actual_ms_ptr, + block_starts_ptr, + out_ptr, + sf_ptr, + cols, + stride_xm, + stride_xn, + stride_ym, + stride_yn, + stride_sm, + stride_sk, + USE_UE8M0: tl.constexpr, + BLOCK_M: tl.constexpr, + BLOCK_K: tl.constexpr, + GROUP_BLOCK_M: tl.constexpr, +): + pid_blk = tl.program_id(axis=0) + pid_sub = tl.program_id(axis=1) + pid_k = tl.program_id(axis=2) + pid_g = tl.load(block_to_group_ptr + pid_blk) + src_start = tl.load(starts_ptr + pid_g) + actual_m = tl.load(actual_ms_ptr + pid_g) + block_start = tl.load(block_starts_ptr + pid_g) + local_block = pid_blk - block_start + row_offsets = local_block * GROUP_BLOCK_M + pid_sub * BLOCK_M + tl.arange(0, BLOCK_M) + col_offsets = pid_k * BLOCK_K + tl.arange(0, BLOCK_K) + src_rows_i64 = (src_start + row_offsets).to(tl.int64) + dst_rows_i64 = (pid_blk * GROUP_BLOCK_M + pid_sub * BLOCK_M + tl.arange(0, BLOCK_M)).to(tl.int64) + col_offsets_i64 = col_offsets.to(tl.int64) + valid_rows = row_offsets < actual_m + valid_cols = col_offsets < cols + x = tl.load( + x_ptr + src_rows_i64[:, None] * stride_xm + col_offsets_i64[None, :] * stride_xn, + mask=valid_rows[:, None] & valid_cols[None, :], + other=0.0, + ).to(tl.float32) + amax = tl.max(tl.abs(x), axis=1) + scale = tl.maximum(amax / 448.0, 1e-4) + if USE_UE8M0: + scale = tl.exp2(tl.ceil(tl.log2(scale))) + y = x / scale[:, None] + tl.store( + out_ptr + dst_rows_i64[:, None] * stride_ym + col_offsets_i64[None, :] * stride_yn, + y.to(tl.float8e4nv), + mask=valid_rows[:, None] & valid_cols[None, :], + ) + tl.store(sf_ptr + dst_rows_i64 * stride_sm + pid_k * stride_sk, scale, mask=valid_rows) + + +@triton.jit +def _grouped_per_channel_fp8_sm90_kmajor_kernel( + x_ptr, + block_to_group_ptr, + starts_ptr, + actual_ms_ptr, + aligned_ms_ptr, + block_starts_ptr, + out_ptr, + sf_ptr, + cols, + stride_xm, + stride_xn, + stride_sf0, + stride_sf1, + USE_UE8M0: tl.constexpr, + BLOCK_K: tl.constexpr, + BLOCK_N: tl.constexpr, +): + pid_blk = tl.program_id(axis=0) + pid_n = tl.program_id(axis=1) + pid_g = tl.load(block_to_group_ptr + pid_blk) + src_start = tl.load(starts_ptr + pid_g) + actual_m = tl.load(actual_ms_ptr + pid_g) + aligned_m = tl.load(aligned_ms_ptr + pid_g) + block_start = tl.load(block_starts_ptr + pid_g) + local_block = pid_blk - block_start + row_offsets = local_block * BLOCK_K + tl.arange(0, BLOCK_K) + col_offsets = pid_n * BLOCK_N + tl.arange(0, BLOCK_N) + src_rows_i64 = (src_start + row_offsets).to(tl.int64) + row_offsets_i64 = row_offsets.to(tl.int64) + col_offsets_i64 = col_offsets.to(tl.int64) + valid_rows = row_offsets < actual_m + valid_cols = col_offsets < cols + x = tl.load( + x_ptr + src_rows_i64[:, None] * stride_xm + col_offsets_i64[None, :] * stride_xn, + mask=valid_rows[:, None] & valid_cols[None, :], + other=0.0, + ).to(tl.float32) + amax = tl.max(tl.abs(x), axis=0) + scale = tl.maximum(amax / 448.0, 1e-4) + if USE_UE8M0: + scale = tl.exp2(tl.ceil(tl.log2(scale))) + y = x / scale[None, :] + flat_base = block_start.to(tl.int64) * BLOCK_K * cols + out_ptrs = out_ptr + flat_base + col_offsets_i64[:, None] * aligned_m + row_offsets_i64[None, :] + tl.store(out_ptrs, tl.trans(y).to(tl.float8e4nv), mask=valid_cols[:, None] & (row_offsets[None, :] < aligned_m)) + tl.store( + sf_ptr + pid_blk * stride_sf0 + col_offsets_i64 * stride_sf1, + scale, + mask=valid_cols, + ) + + +@triton.jit +def _grouped_per_block_fp8_kernel( + x_ptr, + out_ptr, + sf_ptr, + groups, + rows, + cols, + stride_xg, + stride_xm, + stride_xn, + stride_yg, + stride_ym, + stride_yn, + stride_sg, + stride_sm, + stride_sn, + USE_UE8M0: tl.constexpr, + BLOCK_M: tl.constexpr, + BLOCK_N: tl.constexpr, +): + pid_g = tl.program_id(axis=0) + pid_m = tl.program_id(axis=1) + pid_n = tl.program_id(axis=2) + row_offsets = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) + col_offsets = pid_n * BLOCK_N + tl.arange(0, BLOCK_N) + mask = (pid_g < groups) & (row_offsets[:, None] < rows) & (col_offsets[None, :] < cols) + x = tl.load( + x_ptr + pid_g * stride_xg + row_offsets[:, None] * stride_xm + col_offsets[None, :] * stride_xn, + mask=mask, + other=0.0, + ).to(tl.float32) + amax = tl.max(tl.abs(x)) + scale = tl.maximum(amax / 448.0, 1e-4) + if USE_UE8M0: + scale = tl.exp2(tl.ceil(tl.log2(scale))) + y = x / scale + tl.store( + out_ptr + pid_g * stride_yg + row_offsets[:, None] * stride_ym + col_offsets[None, :] * stride_yn, + y.to(tl.float8e4nv), + mask=mask, + ) + tl.store(sf_ptr + pid_g * stride_sg + pid_m * stride_sm + pid_n * stride_sn, scale, mask=pid_g < groups) + + +# --------------------------------------------------------------------------- +# Public quantization functions +# --------------------------------------------------------------------------- + + +def unpack_rows_triton( + x: torch.Tensor, + total_m: int, + block_to_group: torch.Tensor, + starts_tensor: torch.Tensor, + actual_ms_tensor: torch.Tensor, + block_starts_tensor: torch.Tensor, +) -> torch.Tensor: + out = torch.empty((total_m, x.size(1)), device=x.device, dtype=x.dtype) + if total_m == 0: + return out + grid = (block_to_group.numel(), GROUP_ALIGNMENT // 32, ceil_div(x.size(1), 128)) + _unpack_grouped_rows_kernel[grid]( + x, + block_to_group, + starts_tensor, + actual_ms_tensor, + block_starts_tensor, + out, + x.size(1), + x.stride(0), + x.stride(1), + out.stride(0), + out.stride(1), + BLOCK_M=32, + BLOCK_N=128, + GROUP_BLOCK_M=GROUP_ALIGNMENT, + num_warps=4, + ) + return out + + +def per_token_cast_to_fp8_triton( + x: torch.Tensor, use_ue8m0: bool, gran_k: int = GROUP_ALIGNMENT +) -> Tuple[torch.Tensor, torch.Tensor]: + assert x.dim() == 2 + assert gran_k == GROUP_ALIGNMENT + rows, cols = x.shape + out = torch.empty_like(x, dtype=torch.float8_e4m3fn) + sf = torch.empty((rows, ceil_div(cols, gran_k)), device=x.device, dtype=torch.float32) + grid = lambda meta: (ceil_div(rows, meta["BLOCK_M"]), ceil_div(cols, meta["BLOCK_K"])) + _per_token_fp8_kernel[grid]( + x, + out, + sf, + rows, + cols, + x.stride(0), + x.stride(1), + out.stride(0), + out.stride(1), + sf.stride(0), + sf.stride(1), + USE_UE8M0=use_ue8m0, + BLOCK_M=8, + BLOCK_K=gran_k, + num_warps=4, + ) + return out, sf + + +def grouped_per_token_cast_to_fp8_triton( + x: torch.Tensor, + padded_total_m: int, + block_to_group: torch.Tensor, + starts_tensor: torch.Tensor, + actual_ms_tensor: torch.Tensor, + block_starts_tensor: torch.Tensor, + use_ue8m0: bool, + gran_k: int = GROUP_ALIGNMENT, +) -> Tuple[torch.Tensor, torch.Tensor]: + assert x.dim() == 2 + assert gran_k == GROUP_ALIGNMENT + out = torch.empty((padded_total_m, x.size(1)), device=x.device, dtype=torch.float8_e4m3fn) + sf = torch.empty( + (padded_total_m, ceil_div(x.size(1), gran_k)), + device=x.device, + dtype=torch.float32, + ) + if block_to_group.numel() == 0: + return out, sf + grid = (block_to_group.numel(), GROUP_ALIGNMENT // 8, ceil_div(x.size(1), gran_k)) + _grouped_per_token_fp8_kernel[grid]( + x, + block_to_group, + starts_tensor, + actual_ms_tensor, + block_starts_tensor, + out, + sf, + x.size(1), + x.stride(0), + x.stride(1), + out.stride(0), + out.stride(1), + sf.stride(0), + sf.stride(1), + USE_UE8M0=use_ue8m0, + BLOCK_M=8, + BLOCK_K=gran_k, + GROUP_BLOCK_M=GROUP_ALIGNMENT, + num_warps=4, + ) + return out, sf + + +def grouped_per_channel_cast_to_fp8_sm90_kmajor_triton( + x: torch.Tensor, + padded_total_m: int, + block_to_group: torch.Tensor, + starts_tensor: torch.Tensor, + actual_ms_tensor: torch.Tensor, + ks_tensor: torch.Tensor, + block_starts_tensor: torch.Tensor, + use_ue8m0: bool, + gran_k: int = GROUP_ALIGNMENT, +) -> Tuple[torch.Tensor, torch.Tensor]: + assert x.dim() == 2 + assert gran_k == GROUP_ALIGNMENT + out = torch.empty((padded_total_m * x.size(1),), device=x.device, dtype=torch.float8_e4m3fn) + total_blocks = padded_total_m // gran_k + sf = torch.empty((total_blocks, x.size(1)), device=x.device, dtype=torch.float32) + if block_to_group.numel() == 0: + return out, sf.T + grid = (block_to_group.numel(), ceil_div(x.size(1), 128)) + _grouped_per_channel_fp8_sm90_kmajor_kernel[grid]( + x, + block_to_group, + starts_tensor, + actual_ms_tensor, + ks_tensor, + block_starts_tensor, + out, + sf, + x.size(1), + x.stride(0), + x.stride(1), + sf.stride(0), + sf.stride(1), + USE_UE8M0=use_ue8m0, + BLOCK_K=gran_k, + BLOCK_N=128, + num_warps=4, + ) + return out, sf.T + + +def grouped_per_block_cast_to_fp8_triton( + x: torch.Tensor, use_ue8m0: bool, gran_k: int = GROUP_ALIGNMENT +) -> Tuple[torch.Tensor, torch.Tensor]: + assert x.dim() == 3 + assert gran_k == GROUP_ALIGNMENT + groups, rows, cols = x.shape + out = torch.empty((groups, rows, cols), device=x.device, dtype=torch.float8_e4m3fn) + sf = torch.empty( + (groups, ceil_div(rows, gran_k), ceil_div(cols, gran_k)), + device=x.device, + dtype=torch.float32, + ) + grid = (groups, ceil_div(rows, gran_k), ceil_div(cols, gran_k)) + _grouped_per_block_fp8_kernel[grid]( + x, + out, + sf, + groups, + rows, + cols, + x.stride(0), + x.stride(1), + x.stride(2), + out.stride(0), + out.stride(1), + out.stride(2), + sf.stride(0), + sf.stride(1), + sf.stride(2), + USE_UE8M0=use_ue8m0, + BLOCK_M=gran_k, + BLOCK_N=gran_k, + num_warps=8, + ) + return out, sf + + +def per_block_cast_to_fp8_triton( + x: torch.Tensor, use_ue8m0: bool, gran_k: int = GROUP_ALIGNMENT +) -> Tuple[torch.Tensor, torch.Tensor]: + assert x.dim() == 2 + out, sf = grouped_per_block_cast_to_fp8_triton( + x.unsqueeze(0), + use_ue8m0, + gran_k, + ) + return out[0], sf[0] diff --git a/src/prime_rl/trainer/models/kernels/sparse_mla_bwd.py b/src/prime_rl/trainer/models/kernels/sparse_mla_bwd.py new file mode 100644 index 00000000000..36501688f2f --- /dev/null +++ b/src/prime_rl/trainer/models/kernels/sparse_mla_bwd.py @@ -0,0 +1,302 @@ +# Vendored from tile-ai/tilelang (Apache 2.0), modified for dynamic shapes. + +# TileLang ships a libcudart stub that proxies to the real CUDA runtime via +# dlsym(RTLD_DEFAULT, ...). If the stub's own symbols are the first ones found +# (because nothing loaded the real libcudart globally yet), the self-check fails +# and the stub calls abort(). Pre-loading the real library with RTLD_GLOBAL +# ensures dlsym finds it before the stub's own exports. +import ctypes as _ctypes + +try: + _ctypes.CDLL("libcudart.so", mode=_ctypes.RTLD_GLOBAL) +except Exception: + # This is expected on CPU-only machines + pass + +import tilelang +import torch +from tilelang import language as T + + +@tilelang.jit(out_idx=[-1]) +def preprocess( + H, + D, + block_ND=32, + num_stages=5, + dtype=T.bfloat16, + accum_dtype=T.float32, +): + assert dtype == T.bfloat16 + assert accum_dtype == T.float32 + B = T.dynamic("B") + S = T.dynamic("S") + shape = [B, S, H, D] + + @T.prim_func + def preprocess_kernel( + O: T.Tensor(shape, dtype), + dO: T.Tensor(shape, dtype), + Delta: T.Tensor([B, S, H], accum_dtype), + ): + with T.Kernel(H, T.ceildiv(S, block_ND), B) as (bx, by, bz): + o = T.alloc_fragment([block_ND, block_ND], accum_dtype) + do = T.alloc_fragment([block_ND, block_ND], accum_dtype) + delta = T.alloc_fragment([block_ND], accum_dtype) + acc = T.alloc_fragment([block_ND, block_ND], accum_dtype) + T.clear(acc) + for k in T.Pipelined(T.ceildiv(D, block_ND), num_stages=num_stages): + T.copy(O[bz, by * block_ND : (by + 1) * block_ND, bx, k * block_ND : (k + 1) * block_ND], o) + T.copy(dO[bz, by * block_ND : (by + 1) * block_ND, bx, k * block_ND : (k + 1) * block_ND], do) + for i, j in T.Parallel(block_ND, block_ND): + acc[i, j] += o[i, j] * do[i, j] + T.reduce_sum(acc, delta, 1) + T.copy(delta, Delta[bz, by * block_ND : (by + 1) * block_ND, bx]) + + return preprocess_kernel + + +@tilelang.jit(out_idx=[-1]) +def postprocess( + D, + D_tail, + kv_group=1, + block_N=64, + threads=128, + dtype=T.bfloat16, + accum_dtype=T.float32, +): + assert dtype == T.bfloat16 + assert accum_dtype == T.float32 + B = T.dynamic("B") + S_kv = T.dynamic("S_kv") + dkv_shape = [B, S_kv, kv_group, D + D_tail] + + @T.prim_func + def postprocess_kernel( + dKV: T.Tensor(dkv_shape, accum_dtype), + dKV_out: T.Tensor(dkv_shape, dtype), + ): + with T.Kernel(T.ceildiv(S_kv, block_N), kv_group, B, threads=threads) as (bx, by, bz): + T.copy( + dKV[bz, bx * block_N : (bx + 1) * block_N, by, :], + dKV_out[bz, bx * block_N : (bx + 1) * block_N, by, :], + ) + + return postprocess_kernel + + +@tilelang.jit( + out_idx=[-2], + pass_configs={ + tilelang.PassConfigKey.TL_DISABLE_TMA_LOWER: True, + tilelang.PassConfigKey.TL_DISABLE_WARP_SPECIALIZED: True, + tilelang.PassConfigKey.TL_ENABLE_AGGRESSIVE_SHARED_MEMORY_MERGE: True, + }, +) +def bwd( + H, + D, + D_tail, + topk, + kv_group=1, + sm_scale=None, + is_causal=True, + block_size=32, + num_stages=0, + threads=256, + indices_dtype=T.int32, + dtype=T.bfloat16, + accum_dtype=T.float32, +): + assert is_causal is True, "non-casual is not supported now" + assert topk % block_size == 0, "otherwise will load some index=0 thus causing wrong kv to be loaded" + assert dtype == T.bfloat16 + assert accum_dtype == T.float32 + assert indices_dtype == T.int32 + + if sm_scale is None: + sm_scale = (D + D_tail) ** (-0.5) + sm_scale_mul_reciprocal_log2 = sm_scale * 1.44269504 # log2(e) + + B = T.dynamic("B") + S = T.dynamic("S") + S_kv = T.dynamic("S_kv") + + H_kv = H // kv_group + q_shape = [B, S, H, D + D_tail] + k_shape = [B, S_kv, kv_group, D + D_tail] + o_shape = [B, S, H, D] + indices_shape = [B, S, kv_group, topk] + delta_shape = [B, S, H] + lse_shape = [B, S, H] + + H = H_kv + padded_H = max(tilelang.math.next_power_of_2(H_kv), 16) + block_H = min(64, padded_H) + assert padded_H % block_H == 0 + NH = padded_H // block_H + BS = block_size + NS = tilelang.cdiv(topk, block_size) + + split_store = 2 + + @T.prim_func + def sparse_mla_bwd_kernel( + Q: T.Tensor(q_shape, dtype), + KV: T.Tensor(k_shape, dtype), + dO: T.Tensor(o_shape, dtype), + Indices: T.Tensor(indices_shape, indices_dtype), + Lse: T.Tensor(lse_shape, accum_dtype), + Delta: T.Tensor(delta_shape, accum_dtype), + dQ: T.Tensor(q_shape, dtype), + dKV: T.Tensor(k_shape, accum_dtype), + ): + with T.Kernel(S, B, kv_group * NH, threads=threads) as (s_i, by, bz): + Q_shared = T.alloc_shared([block_H, D], dtype) + Q_tail_shared = T.alloc_shared([block_H, D_tail], dtype) + KV_shared = T.alloc_shared([BS, D], dtype) + KV_tail_shared = T.alloc_shared([BS, D_tail], dtype) + dO_shared = T.alloc_shared([block_H, D], dtype) + mask = T.alloc_fragment([BS], "bool") + + P_shared_cast = T.alloc_shared([block_H, BS], dtype) + dP_shared_cast = T.alloc_shared([block_H, BS], dtype) + dQ_shared = T.alloc_shared([block_H, D], dtype) + dQ_tail_shared = T.alloc_shared([block_H, D_tail], dtype) + + acc_p = T.alloc_fragment([block_H, BS], accum_dtype) + acc_dp = T.alloc_fragment([block_H, BS], accum_dtype) + acc_dq = T.alloc_fragment([block_H, D], accum_dtype) + acc_dq_tail = T.alloc_fragment([block_H, D_tail], accum_dtype) + acc_dkv = T.alloc_fragment([BS, D], accum_dtype) + acc_dkv_tail = T.alloc_fragment([BS, D_tail], accum_dtype) + acc_dkv_shared = T.alloc_shared([BS // split_store, D], accum_dtype) + acc_dkv_tail_shared = T.alloc_shared([BS // split_store, D_tail], accum_dtype) + + # See sparse_mla_fwd: sentinel is at index S_kv - 1 (zero KV), valid indices + # live in [0, S_kv - 1). Using this single bound makes the kernel work for + # both full and CP-sharded Q without needing a global Q offset. + max_kv_i = S_kv - 2 + + T.copy(Q[by, s_i, bz * block_H : (bz + 1) * block_H, :D], Q_shared) + T.copy(Q[by, s_i, bz * block_H : (bz + 1) * block_H, D:], Q_tail_shared) + T.copy(dO[by, s_i, bz * block_H : (bz + 1) * block_H, :D], dO_shared) + + T.clear(acc_dq) + T.clear(acc_dq_tail) + + for i_i in T.Pipelined(NS, num_stages=num_stages): + for bi_i in T.Parallel(BS): + mask[bi_i] = Indices[by, s_i, bz // NH, i_i * BS + bi_i] <= max_kv_i + + for h_i, bi_i in T.Parallel(block_H, BS): + acc_p[h_i, bi_i] = T.if_then_else(mask[bi_i], 0, -T.infinity(acc_p.dtype)) + + for bi_i, d_i in T.Parallel(BS, D): + KV_shared[bi_i, d_i] = KV[by, Indices[by, s_i, bz // NH, i_i * BS + bi_i], bz // NH, d_i] + + T.gemm(Q_shared, KV_shared, acc_p, transpose_B=True, policy=T.GemmWarpPolicy.FullCol) + + for bi_i, d_i in T.Parallel(BS, D_tail): + KV_tail_shared[bi_i, d_i] = KV[by, Indices[by, s_i, bz // NH, i_i * BS + bi_i], bz // NH, D + d_i] + T.gemm(Q_tail_shared, KV_tail_shared, acc_p, transpose_B=True, policy=T.GemmWarpPolicy.FullCol) + + for h_i, bi_i in T.Parallel(block_H, BS): + acc_p[h_i, bi_i] = T.exp2( + acc_p[h_i, bi_i] * sm_scale_mul_reciprocal_log2 - Lse[by, s_i, bz * block_H + h_i] + ) + + T.copy(acc_p, P_shared_cast) + + T.gemm( + dO_shared, KV_shared, acc_dp, transpose_B=True, policy=T.GemmWarpPolicy.FullCol, clear_accum=True + ) + + for h_i, bi_i in T.Parallel(block_H, BS): + acc_dp[h_i, bi_i] = ( + acc_p[h_i, bi_i] * (acc_dp[h_i, bi_i] - Delta[by, s_i, bz * block_H + h_i]) * sm_scale + ) + + T.copy(acc_dp, dP_shared_cast) + T.gemm(dP_shared_cast, KV_shared, acc_dq, policy=T.GemmWarpPolicy.FullCol) + T.gemm(dP_shared_cast, KV_tail_shared, acc_dq_tail, policy=T.GemmWarpPolicy.FullCol) + + T.gemm( + dP_shared_cast, + Q_shared, + acc_dkv, + transpose_A=True, + policy=T.GemmWarpPolicy.FullCol, + clear_accum=True, + ) + T.gemm(P_shared_cast, dO_shared, acc_dkv, transpose_A=True, policy=T.GemmWarpPolicy.FullCol) + + T.clear(acc_dkv_tail) + T.gemm(dP_shared_cast, Q_tail_shared, acc_dkv_tail, transpose_A=True, policy=T.GemmWarpPolicy.FullCol) + + for s in range(split_store): + for bi_i, d_i in T.Parallel(BS, D): + if bi_i < BS // split_store: + acc_dkv_shared[bi_i, d_i] = acc_dkv[bi_i + s * (BS // split_store), d_i] + + for bi_i, d_i in T.Parallel(BS, D_tail): + if bi_i < BS // split_store: + acc_dkv_tail_shared[bi_i, d_i] = acc_dkv_tail[bi_i + s * (BS // split_store), d_i] + + for bi_i, d_i in T.Parallel(BS // split_store, D // 4): + T.atomic_addx4( + dKV[ + by, + Indices[by, s_i, bz // NH, i_i * BS + bi_i + s * (BS // split_store)], + bz // NH, + d_i * 4, + ], + acc_dkv_shared[bi_i, d_i * 4], + ) + + for bi_i, d_i in T.Parallel(BS // split_store, D_tail // 4): + T.atomic_addx4( + dKV[ + by, + Indices[by, s_i, bz // NH, i_i * BS + bi_i + s * (BS // split_store)], + bz // NH, + D + d_i * 4, + ], + acc_dkv_tail_shared[bi_i, d_i * 4], + ) + + T.copy(acc_dq, dQ_shared) + T.copy(acc_dq_tail, dQ_tail_shared) + + T.copy(dQ_shared, dQ[by, s_i, bz * block_H : (bz + 1) * block_H, :D]) + T.copy(dQ_tail_shared, dQ[by, s_i, bz * block_H : (bz + 1) * block_H, D:]) + + return sparse_mla_bwd_kernel + + +def sparse_mla_bwd(q, kv, o, do, indices, lse, sm_scale=None): + assert q.is_contiguous() + assert kv.is_contiguous() + assert indices.is_contiguous() + assert lse.is_contiguous() + B, S, H, dim_plus_tail_dim = q.shape + _, S_kv, kv_group, _ = kv.shape + assert kv.shape[-1] == dim_plus_tail_dim + assert kv.shape[0] == B + D = 512 + D_tail = dim_plus_tail_dim - D + topk = indices.shape[-1] + assert indices.shape == (B, S, kv_group, topk) + assert lse.shape == (B, S, H) + + preprocess_kernel = preprocess(H, D) + bwd_kernel = bwd(H, D, D_tail, topk, kv_group, sm_scale, True) + postprocess_kernel = postprocess(D, D_tail, kv_group) + + delta = preprocess_kernel(o, do) + dkv = torch.zeros_like(kv, dtype=torch.float32) + dq = bwd_kernel(q, kv, do, indices, lse, delta, dkv) + dkv = postprocess_kernel(dkv) + + return dq, dkv diff --git a/src/prime_rl/trainer/models/kernels/sparse_mla_fwd.py b/src/prime_rl/trainer/models/kernels/sparse_mla_fwd.py new file mode 100644 index 00000000000..ac360a02427 --- /dev/null +++ b/src/prime_rl/trainer/models/kernels/sparse_mla_fwd.py @@ -0,0 +1,212 @@ +# Vendored from tile-ai/tilelang (Apache 2.0), modified for dynamic shapes. + +# TileLang ships a libcudart stub that proxies to the real CUDA runtime via +# dlsym(RTLD_DEFAULT, ...). If the stub's own symbols are the first ones found +# (because nothing loaded the real libcudart globally yet), the self-check fails +# and the stub calls abort(). Pre-loading the real library with RTLD_GLOBAL +# ensures dlsym finds it before the stub's own exports. +import ctypes as _ctypes + +try: + _ctypes.CDLL("libcudart.so", mode=_ctypes.RTLD_GLOBAL) +except Exception: + # This is expected on CPU-only machines + pass + +import tilelang +from tilelang import language as T + + +@tilelang.jit( + out_idx=[-2, -1], + pass_configs={ + tilelang.PassConfigKey.TL_DISABLE_TMA_LOWER: True, + tilelang.PassConfigKey.TL_DISABLE_WARP_SPECIALIZED: True, + }, +) +def sparse_mla_fwd( + heads, + dim, + tail_dim, + topk, + kv_group=1, + sm_scale=None, + is_causal=True, + CP0=True, + block_I=64, + num_stages=2, + threads=256, +): + assert dim == tilelang.math.next_power_of_2(dim), f"haven't check padding correctness yet, dim={dim}" + assert tail_dim == tilelang.math.next_power_of_2(tail_dim), f"haven't check padding correctness yet, dim={tail_dim}" + assert is_causal is True, "non-casual is not supported" + assert topk % block_I == 0, "otherwise will load some index=0 thus causing wrong kv to be loaded" + if sm_scale is None: + sm_scale = (1.0 / (dim + tail_dim)) ** 0.5 * 1.44269504 # log2(e) + else: + sm_scale = sm_scale * 1.44269504 # log2(e) + + batch = T.dynamic("batch") + seq_len = T.dynamic("seq_len") + seq_len_kv = T.dynamic("seq_len_kv") + + head_kv = heads // kv_group + q_shape = [batch, seq_len, heads, dim + tail_dim] + kv_shape = [batch, seq_len_kv, kv_group, dim + tail_dim] + o_shape = [batch, seq_len, heads, dim] + indices_shape = [batch, seq_len, kv_group, topk] + lse_shape = [batch, seq_len, heads] + indices_dtype = T.int32 + dtype = T.bfloat16 + accum_dtype = T.float32 + + H = head_kv + padded_H = max(tilelang.math.next_power_of_2(head_kv), 16) + if padded_H != H: + assert kv_group == 1, ( + "here we solve the H padding automatically, other wise you should handle Q copy and Output copy" + " with your mask (when kv_group == 1, use g_i * padded_H:(g_i+1) * padded_H would be handled" + " automatically)" + ) + BI = block_I + NI = tilelang.cdiv(topk, block_I) + D = dim + D_tail = tail_dim + + if head_kv > 64: + assert head_kv % 64 == 0, "head_kv should be a multiple of 64" + REPLICATE_H = head_kv // 64 + else: + REPLICATE_H = 1 + + H_per_block = padded_H if REPLICATE_H == 1 else 64 + + @T.prim_func + def main( + Q: T.Tensor(q_shape, dtype), # type: ignore + KV: T.Tensor(kv_shape, dtype), # type: ignore + Indices: T.Tensor(indices_shape, indices_dtype), # type: ignore + Output: T.Tensor(o_shape, dtype), # type: ignore + Lse: T.Tensor(lse_shape, accum_dtype), # type: ignore + ): + with T.Kernel(seq_len * REPLICATE_H, batch, kv_group, threads=threads) as ( + bx, + by, + bz, + ): + Q_shared = T.alloc_shared([H_per_block, D], dtype) + Q_tail_shared = T.alloc_shared([H_per_block, D_tail], dtype) + KV_shared = T.alloc_shared([BI, D], dtype) + K_tail_shared = T.alloc_shared([BI, D_tail], dtype) + O_shared = T.alloc_shared([H_per_block, D], dtype) # noqa: F841 + Lse_shared = T.alloc_shared([H_per_block], accum_dtype) # noqa: F841 + mask = T.alloc_fragment([BI], "bool") + + acc_o = T.alloc_fragment([H_per_block, D], accum_dtype) + acc_s = T.alloc_fragment([H_per_block, BI], accum_dtype) + S_shared = T.alloc_shared([H_per_block, BI], dtype) + sumexp = T.alloc_fragment([H_per_block], accum_dtype) + sumexp_i = T.alloc_fragment([H_per_block], accum_dtype) + alpha = T.alloc_fragment([H_per_block], accum_dtype) + m_i = T.alloc_fragment([H_per_block], accum_dtype) + m_i_prev = T.alloc_fragment([H_per_block], accum_dtype) + + T.fill(acc_o, 0) + T.fill(sumexp, 0) + T.fill(m_i, -(2**30)) + + b_i, g_i = by, bz + s_i = bx if REPLICATE_H == 1 else (bx // REPLICATE_H) + # The indexer pre-filters indices using per-token `ke` and replaces out-of-range + # entries with the sentinel value `seq_len_kv - 1` (the last KV slot is a + # zero sentinel; valid K indices live in [0, seq_len_kv - 1)). This single + # bound preserves causality + varlen masking for both full and CP-sharded Q + # (where local q_i no longer matches the global K position). + max_kv_i = seq_len_kv - 2 + + H0 = g_i * padded_H + (0 if REPLICATE_H == 1 else (bx % REPLICATE_H) * 64) + H1 = H0 + H_per_block + + T.copy(Q[b_i, s_i, H0:H1, :D], Q_shared) + T.copy(Q[b_i, s_i, H0:H1, D:], Q_tail_shared) + + for i_i in T.Pipelined(NI, num_stages=num_stages): + for bi_i in T.Parallel(BI): + mask[bi_i] = Indices[b_i, s_i, g_i, i_i * BI + bi_i] <= max_kv_i + + for bi_i, d_i in T.Parallel(BI, D): + KV_shared[bi_i, d_i] = KV[b_i, Indices[b_i, s_i, g_i, i_i * BI + bi_i], g_i, d_i] + for bi_i, d_i in T.Parallel(BI, D_tail): + K_tail_shared[bi_i, d_i] = KV[b_i, Indices[b_i, s_i, g_i, i_i * BI + bi_i], g_i, D + d_i] + + for h_i, bi_i in T.Parallel(H_per_block, BI): + acc_s[h_i, bi_i] = T.if_then_else(mask[bi_i], 0, -T.infinity(acc_s.dtype)) + T.gemm( + Q_shared, + KV_shared, + acc_s, + transpose_B=True, + policy=T.GemmWarpPolicy.FullRow, + ) + T.gemm( + Q_tail_shared, + K_tail_shared, + acc_s, + transpose_B=True, + policy=T.GemmWarpPolicy.FullRow, + ) + T.copy(m_i, m_i_prev) + T.reduce_max(acc_s, m_i, dim=1, clear=False) + for h_i in T.Parallel(H_per_block): + m_i[h_i] = T.max(m_i[h_i], m_i_prev[h_i]) + for h_i in T.Parallel(H_per_block): + alpha[h_i] = T.exp2((m_i_prev[h_i] - m_i[h_i]) * sm_scale) + for h_i, bi_i in T.Parallel(H_per_block, BI): + acc_s[h_i, bi_i] = T.exp2(acc_s[h_i, bi_i] * sm_scale - m_i[h_i] * sm_scale) + T.reduce_sum(acc_s, sumexp_i, dim=1) + for h_i in T.Parallel(H_per_block): + sumexp[h_i] = sumexp[h_i] * alpha[h_i] + sumexp_i[h_i] + for h_i, d_i in T.Parallel(H_per_block, D): + acc_o[h_i, d_i] = acc_o[h_i, d_i] * alpha[h_i] + + T.copy(acc_s, S_shared) + T.gemm(S_shared, KV_shared, acc_o, policy=T.GemmWarpPolicy.FullRow) + + for h_i, d_i in T.Parallel(H_per_block, D): + acc_o[h_i, d_i] /= sumexp[h_i] + for h_i in T.Parallel(H_per_block): + sumexp[h_i] = T.log2(sumexp[h_i]) + m_i[h_i] * sm_scale + + T.copy(acc_o, Output[b_i, s_i, H0:H1, :]) + T.copy(sumexp, Lse[b_i, s_i, H0:H1]) + + return main + + +def sparse_mla_fwd_interface(q, kv, indices, sm_scale=None, d_v=512, block_I=64, num_stages=2, threads=256): + assert q.is_contiguous() and kv.is_contiguous() and indices.is_contiguous() + batch, seq_len, heads, dim_plus_tail_dim = q.shape + _, seq_len_kv, kv_group, _ = kv.shape + + assert dim_plus_tail_dim == 576, "you should assign dim otherwise" + dim = d_v + assert kv.shape[-1] == dim_plus_tail_dim + assert kv.shape[0] == batch + tail_dim = dim_plus_tail_dim - dim + _, _, _, topk = indices.shape + assert indices.shape == (batch, seq_len, kv_group, topk) + + kernel = sparse_mla_fwd( + heads, + dim, + tail_dim, + topk, + kv_group, + sm_scale, + True, + block_I=block_I, + num_stages=num_stages, + threads=threads, + ) + out, lse = kernel(q, kv, indices) + return out, lse diff --git a/src/prime_rl/trainer/models/layers/attn.py b/src/prime_rl/trainer/models/layers/attn.py index 13d2828fa38..d59a2099260 100644 --- a/src/prime_rl/trainer/models/layers/attn.py +++ b/src/prime_rl/trainer/models/layers/attn.py @@ -1,22 +1,30 @@ import functools from dataclasses import dataclass +from typing import Literal import torch import torch.nn.functional as F from torch import nn -from .rms_norm import RMSNorm, RMSNormConfig +from .norms import RMSNorm, RMSNormConfig from .rotary_emb import apply_rotary_pos_emb +# flash-attention-2 try: from flash_attn import flash_attn_varlen_func except ImportError: - flash_attn_varlen_func = None + flash_attn_varlen_func = None # type: ignore +# flash-attention-3 try: from flash_attn_interface import flash_attn_varlen_func as flash_attn_3_varlen_func except ImportError: - flash_attn_3_varlen_func = None + flash_attn_3_varlen_func = None # type: ignore + +try: + from flash_attn.cute import flash_attn_varlen_func as flash_attn_4_varlen_func +except ImportError: + flash_attn_4_varlen_func = None # type: ignore @dataclass @@ -29,6 +37,7 @@ class AttentionConfig: attention_bias: bool use_qk_norm: bool rms_norm_eps: float + qk_norm_type: Literal["per_head", "per_layer"] = "per_head" # TODO: Does torch compile support config._attn_implementation forking? @@ -39,6 +48,12 @@ class AttentionConfig: class FlashAttention(nn.Module): """Flash Attention""" + _funcs = { + 2: flash_attn_varlen_func, + 3: flash_attn_3_varlen_func, + 4: flash_attn_4_varlen_func, + } + def __init__(self, config: AttentionConfig, flash_attn_version: int = 2): super().__init__() self.head_dim = config.head_dim @@ -57,27 +72,75 @@ def __init__(self, config: AttentionConfig, flash_attn_version: int = 2): ) self.o_proj = nn.Linear(config.num_attention_heads * self.head_dim, config.hidden_size, bias=False) self.use_qk_norm = config.use_qk_norm + self.qk_norm_type = config.qk_norm_type if self.use_qk_norm: - self.q_norm = RMSNorm(RMSNormConfig(hidden_size=self.head_dim, eps=config.rms_norm_eps)) - self.k_norm = RMSNorm(RMSNormConfig(hidden_size=self.head_dim, eps=config.rms_norm_eps)) - - self.func = flash_attn_3_varlen_func if flash_attn_version == 3 else flash_attn_varlen_func + if self.qk_norm_type == "per_layer": + self.q_norm = RMSNorm( + RMSNormConfig(hidden_size=config.num_attention_heads * self.head_dim, eps=config.rms_norm_eps) + ) + self.k_norm = RMSNorm( + RMSNormConfig(hidden_size=config.num_key_value_heads * self.head_dim, eps=config.rms_norm_eps) + ) + else: + self.q_norm = RMSNorm(RMSNormConfig(hidden_size=self.head_dim, eps=config.rms_norm_eps)) + self.k_norm = RMSNorm(RMSNormConfig(hidden_size=self.head_dim, eps=config.rms_norm_eps)) + + self._flash_attn_version = flash_attn_version + self.func = self._funcs[flash_attn_version] + self._flash_attn_call = self.func + if self._flash_attn_version == 4: + self._flash_attn_call = torch._dynamo.disable(self.func) + + def _compute_attention(self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, cu_seqlens, max_seqlen): + """Run the flash attention kernel. q/k/v are [total_tokens, heads, dim].""" + kwargs: dict = {"causal": True} + sliding_window = getattr(self, "sliding_window", None) + if sliding_window is not None: + kwargs["window_size"] = (sliding_window - 1, 0) + if self._flash_attn_version == 4: + # FA4's flash_attn_varlen_func has qv as the 4th positional arg, + # so cu_seqlens must be passed as keyword args to avoid misalignment. + kwargs["cu_seqlens_q"] = cu_seqlens + kwargs["cu_seqlens_k"] = cu_seqlens + out = self._flash_attn_call(q, k, v, **kwargs) + else: + out = self._flash_attn_call(q, k, v, cu_seqlens, cu_seqlens, max_seqlen, max_seqlen, **kwargs) + if isinstance(out, tuple): + out = out[0] + return out - def forward( + def _attention_core( self, - hidden_states: torch.Tensor, - position_embeddings: tuple[torch.Tensor, torch.Tensor], + query_states: torch.Tensor, + key_states: torch.Tensor, + value_states: torch.Tensor, cu_seqlens: torch.LongTensor | None = None, max_seqlen: int | None = None, - ) -> tuple[torch.Tensor, torch.Tensor | None]: + ) -> torch.Tensor: + out = self._compute_attention(query_states[0], key_states[0], value_states[0], cu_seqlens, max_seqlen) + return out.contiguous().view(1, out.shape[0], -1) + + def attn_projections( + self, + hidden_states: torch.Tensor, + position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: input_shape = hidden_states.shape[:-1] hidden_shape = (*input_shape, -1, self.head_dim) - query_states = self.q_proj(hidden_states).view(hidden_shape) - key_states = self.k_proj(hidden_states).view(hidden_shape) - value_states = self.v_proj(hidden_states).view(hidden_shape) + query_states = self.q_proj(hidden_states) + key_states = self.k_proj(hidden_states) + value_states = self.v_proj(hidden_states) - if self.use_qk_norm: # main diff from Llama + if self.use_qk_norm and self.qk_norm_type == "per_layer": + query_states = self.q_norm(query_states) + key_states = self.k_norm(key_states) + + query_states = query_states.view(hidden_shape) + key_states = key_states.view(hidden_shape) + value_states = value_states.view(hidden_shape) + + if self.use_qk_norm and self.qk_norm_type == "per_head": query_states = self.q_norm(query_states) key_states = self.k_norm(key_states) @@ -85,32 +148,38 @@ def forward( key_states = key_states.transpose(1, 2) value_states = value_states.transpose(1, 2) - cos, sin = position_embeddings - query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin) + if position_embeddings is not None: + cos, sin = position_embeddings + query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin) # TODO: Can we optimize the rotary application instead of double transpose? query_states = query_states.transpose(1, 2) key_states = key_states.transpose(1, 2) value_states = value_states.transpose(1, 2) - out = self.func( - query_states[0], - key_states[0], - value_states[0], - cu_seqlens, - cu_seqlens, - max_seqlen, - max_seqlen, - causal=True, - ) - if isinstance(out, tuple): - out = out[0] - out = out.contiguous() - attn_output = out.view(1, out.shape[0], -1) - attn_weights = None + return query_states, key_states, value_states + + def output_proj(self, attn_output: torch.Tensor) -> torch.Tensor: + return self.o_proj(attn_output) - attn_output = self.o_proj(attn_output) - return attn_output, attn_weights + def forward( + self, + hidden_states: torch.Tensor, + position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None, + cu_seqlens: torch.LongTensor | None = None, + max_seqlen: int | None = None, + ) -> tuple[torch.Tensor, torch.Tensor | None]: + query_states, key_states, value_states = self.attn_projections(hidden_states, position_embeddings) + + attn_output = self._attention_core( + query_states, + key_states, + value_states, + cu_seqlens=cu_seqlens, + max_seqlen=max_seqlen, + ) + attn_output = self.output_proj(attn_output) + return attn_output, None class SDPAAttention(nn.Module): @@ -134,25 +203,40 @@ def __init__(self, config: AttentionConfig): ) self.o_proj = nn.Linear(config.num_attention_heads * self.head_dim, config.hidden_size, bias=False) self.use_qk_norm = config.use_qk_norm + self.qk_norm_type = config.qk_norm_type if self.use_qk_norm: - self.q_norm = RMSNorm(RMSNormConfig(hidden_size=self.head_dim, eps=config.rms_norm_eps)) - self.k_norm = RMSNorm(RMSNormConfig(hidden_size=self.head_dim, eps=config.rms_norm_eps)) - - def forward( + if self.qk_norm_type == "per_layer": + self.q_norm = RMSNorm( + RMSNormConfig(hidden_size=config.num_attention_heads * self.head_dim, eps=config.rms_norm_eps) + ) + self.k_norm = RMSNorm( + RMSNormConfig(hidden_size=config.num_key_value_heads * self.head_dim, eps=config.rms_norm_eps) + ) + else: + self.q_norm = RMSNorm(RMSNormConfig(hidden_size=self.head_dim, eps=config.rms_norm_eps)) + self.k_norm = RMSNorm(RMSNormConfig(hidden_size=self.head_dim, eps=config.rms_norm_eps)) + + def attn_projections( self, hidden_states: torch.Tensor, - position_embeddings: tuple[torch.Tensor, torch.Tensor], - cu_seqlens: torch.LongTensor | None = None, - max_seqlen: int | None = None, - ) -> tuple[torch.Tensor, torch.Tensor | None]: + position_embeddings: tuple[torch.Tensor, torch.Tensor] | None, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: input_shape = hidden_states.shape[:-1] hidden_shape = (*input_shape, -1, self.head_dim) - query_states: torch.Tensor = self.q_proj(hidden_states).view(hidden_shape) - key_states: torch.Tensor = self.k_proj(hidden_states).view(hidden_shape) - value_states: torch.Tensor = self.v_proj(hidden_states).view(hidden_shape) + query_states = self.q_proj(hidden_states) + key_states = self.k_proj(hidden_states) + value_states = self.v_proj(hidden_states) + + if self.use_qk_norm and self.qk_norm_type == "per_layer": + query_states = self.q_norm(query_states) + key_states = self.k_norm(key_states) + + query_states = query_states.view(hidden_shape) + key_states = key_states.view(hidden_shape) + value_states = value_states.view(hidden_shape) - if self.use_qk_norm: # main diff from Llama + if self.use_qk_norm and self.qk_norm_type == "per_head": query_states = self.q_norm(query_states) key_states = self.k_norm(key_states) @@ -160,88 +244,98 @@ def forward( key_states = key_states.transpose(1, 2) value_states = value_states.transpose(1, 2) - cos, sin = position_embeddings - query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin) + if position_embeddings is not None: + cos, sin = position_embeddings + query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin) - # TODO: Can we optimize the rotary application instead of double transpose? + return query_states, key_states, value_states + + def _attention_core( + self, + query_states: torch.Tensor, + key_states: torch.Tensor, + value_states: torch.Tensor, + ) -> torch.Tensor: key_states = key_states.repeat_interleave(self.num_key_value_groups, dim=1) value_states = value_states.repeat_interleave(self.num_key_value_groups, dim=1) out = F.scaled_dot_product_attention(query_states, key_states, value_states, is_causal=True) - out = out.transpose(1, 2).contiguous() # .view(out.shape[0], out.shape[1], -1) - attn_output = out.view(out.shape[0], out.shape[1], -1) - attn_weights = None + out = out.transpose(1, 2).contiguous() + return out.view(out.shape[0], out.shape[1], -1) - # attn_output = attn_output.reshape(*input_shape, -1).contiguous() - attn_output = self.o_proj(attn_output) - return attn_output, attn_weights + def output_proj(self, attn_output: torch.Tensor) -> torch.Tensor: + return self.o_proj(attn_output) + + def forward( + self, + hidden_states: torch.Tensor, + position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None, + cu_seqlens: torch.LongTensor | None = None, + max_seqlen: int | None = None, + ) -> tuple[torch.Tensor, torch.Tensor | None]: + query_states, key_states, value_states = self.attn_projections(hidden_states, position_embeddings) + + attn_output = self._attention_core(query_states, key_states, value_states) + attn_output = self.output_proj(attn_output) + return attn_output, None ATTN_IMPL2CLASS = { "flash_attention_2": functools.partial(FlashAttention, flash_attn_version=2), "sdpa": SDPAAttention, "flash_attention_3": functools.partial(FlashAttention, flash_attn_version=3), + "fa4": functools.partial(FlashAttention, flash_attn_version=4), } -def substitute_prime_rl_flash_attn(process_group: torch.distributed.ProcessGroup, heads_k_stride: int) -> None: +def substitute_ring_attn( + process_group: torch.distributed.ProcessGroup, + heads_k_stride: int, + attn_impl: str = "flash_attention_2", +) -> None: + """Patch _compute_attention on FlashAttention variants to use ring attention.""" from ring_flash_attn import llama3_flash_attn_varlen_func - class RingFlashAttention(FlashAttention): - def forward( - self, - hidden_states: torch.Tensor, - position_embeddings: tuple[torch.Tensor, torch.Tensor], - cu_seqlens: torch.LongTensor | None = None, - max_seqlen: int | None = None, - ) -> tuple[torch.Tensor, torch.Tensor | None]: - input_shape = hidden_states.shape[:-1] - hidden_shape = (*input_shape, -1, self.head_dim) - - query_states = self.q_proj(hidden_states).view(hidden_shape) - key_states = self.k_proj(hidden_states).view(hidden_shape) - value_states = self.v_proj(hidden_states).view(hidden_shape) - - if self.use_qk_norm: # main diff from Llama - query_states = self.q_norm(query_states) - key_states = self.k_norm(key_states) - - query_states = query_states.transpose(1, 2) - key_states = key_states.transpose(1, 2) - value_states = value_states.transpose(1, 2) + from .ring_attn import ring_fa3_varlen_func, ring_fa4_varlen_func + + if attn_impl == "fa4": + ring_func = ring_fa4_varlen_func + elif attn_impl == "flash_attention_3": + ring_func = ring_fa3_varlen_func + else: + ring_func = llama3_flash_attn_varlen_func + + def _ring_compute_attention(self, q, k, v, cu_seqlens, max_seqlen): + from ring_flash_attn.adapters.hf_adapter import DATA_PARAMS + + window_size = (-1, -1) + sliding_window = getattr(self, "sliding_window", None) + if sliding_window is not None: + window_size = (sliding_window - 1, 0) + + out = ring_func( + q, + k, + v, + cu_seqlens_q=DATA_PARAMS["cu_seqlens_q"], + cu_seqlens_k=DATA_PARAMS["cu_seqlens_k"], + max_seqlen_q=DATA_PARAMS["max_seqlen_q"], + max_seqlen_k=DATA_PARAMS["max_seqlen_k"], + local_k_slice=DATA_PARAMS["local_k_slice"], + causal=True, + window_size=window_size, + group=process_group, + heads_k_stride=heads_k_stride, + ) + if isinstance(out, tuple): + out = out[0] + return out - cos, sin = position_embeddings - query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin) + FlashAttention._compute_attention = _ring_compute_attention + + from prime_rl.trainer.models.afmoe.modeling_afmoe import AfmoeFlashAttention + + AfmoeFlashAttention._compute_attention = _ring_compute_attention + + from prime_rl.trainer.models.qwen3_5_moe.modeling_qwen3_5_moe import Qwen3_5MoeGatedFlashAttention - from ring_flash_attn.adapters.hf_adapter import DATA_PARAMS - - cu_seqlens_q = DATA_PARAMS["cu_seqlens_q"] - cu_seqlens_k = DATA_PARAMS["cu_seqlens_k"] - max_seqlen_q = DATA_PARAMS["max_seqlen_q"] - max_seqlen_k = DATA_PARAMS["max_seqlen_k"] - local_k_slice = DATA_PARAMS["local_k_slice"] - - # TODO: Can we optimize the rotary application instead of double transpose? - query_states = query_states.transpose(1, 2) - key_states = key_states.transpose(1, 2) - value_states = value_states.transpose(1, 2) - out = llama3_flash_attn_varlen_func( - query_states[0], - key_states[0], - value_states[0], - cu_seqlens_q=cu_seqlens_q, - cu_seqlens_k=cu_seqlens_k, - max_seqlen_q=max_seqlen_q, - max_seqlen_k=max_seqlen_k, - local_k_slice=local_k_slice, - causal=True, - group=process_group, - heads_k_stride=heads_k_stride, - ) - out = out.contiguous() - attn_output = out.view(1, out.shape[0], -1) - attn_weights = None - - attn_output = self.o_proj(attn_output) - return attn_output, attn_weights - - FlashAttention.forward = RingFlashAttention.forward + Qwen3_5MoeGatedFlashAttention._compute_attention = _ring_compute_attention diff --git a/src/prime_rl/trainer/models/layers/checkpointing.py b/src/prime_rl/trainer/models/layers/checkpointing.py new file mode 100644 index 00000000000..ef4c41cd7d2 --- /dev/null +++ b/src/prime_rl/trainer/models/layers/checkpointing.py @@ -0,0 +1,151 @@ +"""Selective activation checkpointing hooks for custom decoder layers. + +New model integration contract: + +- `norm`: no explicit hook is required. Every child module whose type name + contains `norm` is checkpointed through `forward()`. +- `attn_proj`: expose `layer.self_attn.attn_projections(...)` for + projection-side attention work before the core kernel, plus + `layer.self_attn.output_proj(...)` for the post-attention output projection. + This target is meant for attention-local work outside the main attention + kernel, such as q/k/v projections, attention-local norms, RoPE, gating, and + similar model-specific helpers. +- `mla_up_proj`: expose `layer.self_attn.mla_up_proj(...)` when MLA-style + up-projection work should be checkpointed separately from + `attn_projections(...)`. +- `mlp`: expose a dense `layer.mlp.forward(...)`. A module is treated as dense + when it does not define `_run_routed_experts` or `tokens_per_expert`. +- `routed_experts`: expose `layer.mlp._run_routed_experts(...)` for the MoE + expert path, and optionally `layer.mlp._run_local_routed_experts(...)` when + local expert compute is separated from dispatch/combine. +- `linear_attn`: expose a token-mixer module on `layer.linear_attn` or + `layer.mamba`, or reuse `layer.self_attn` when + `layer.attention_type == "sliding_attention"`. + +These hook names are an intentional public interface for custom model support, +so new models should follow them directly rather than adding private aliases. +""" + +from collections.abc import Iterable +from functools import wraps + +import torch +import torch.nn as nn +from torch.utils.checkpoint import checkpoint + +SELECTIVE_AC_TARGETS = frozenset({"norm", "attn_proj", "mlp", "mla_up_proj", "routed_experts", "linear_attn"}) +_PATCHED_METHODS_ATTR = "_prime_rl_selective_ac_patched_methods" + + +def _is_norm_module(module: nn.Module) -> bool: + return "norm" in type(module).__name__.lower() + + +def _should_checkpoint(module: nn.Module) -> bool: + return torch.is_grad_enabled() and module.training + + +def checkpoint_method(module: nn.Module, method_name: str) -> None: + patched_methods = frozenset(getattr(module, _PATCHED_METHODS_ATTR, ())) + if method_name in patched_methods: + return + + original = getattr(module, method_name) + + @wraps(original) + def checkpointed(*args, **kwargs): + if not _should_checkpoint(module): + return original(*args, **kwargs) + + def fn(*checkpoint_args): + return original(*checkpoint_args, **kwargs) + + return checkpoint(fn, *args, use_reentrant=False, preserve_rng_state=False) + + setattr(module, method_name, checkpointed) + setattr(module, _PATCHED_METHODS_ATTR, patched_methods.union({method_name})) + + +def _configure_norm_checkpointing(layer: nn.Module) -> None: + for module in layer.modules(): + if module is layer or not _is_norm_module(module): + continue + checkpoint_method(module, "forward") + + +def _is_dense_mlp(mlp: nn.Module) -> bool: + return not hasattr(mlp, "_run_routed_experts") and not hasattr(mlp, "tokens_per_expert") + + +def _supports_attn_proj(self_attn: nn.Module | None) -> bool: + return self_attn is not None and hasattr(self_attn, "attn_projections") and hasattr(self_attn, "output_proj") + + +def _get_linear_attn_module(layer: nn.Module) -> nn.Module | None: + linear_attn = getattr(layer, "linear_attn", None) + if linear_attn is not None: + return linear_attn + + mamba = getattr(layer, "mamba", None) + if mamba is not None: + return mamba + + if getattr(layer, "attention_type", None) == "sliding_attention": + return getattr(layer, "self_attn", None) + + return None + + +def get_supported_targets(layer: nn.Module) -> frozenset[str]: + """Infer which selective activation checkpoint targets a decoder layer supports.""" + supported_targets = {"norm"} + self_attn = getattr(layer, "self_attn", None) + mlp = getattr(layer, "mlp", None) + linear_attn = _get_linear_attn_module(layer) + + if _supports_attn_proj(self_attn): + supported_targets.add("attn_proj") + if self_attn is not None and hasattr(self_attn, "mla_up_proj"): + supported_targets.add("mla_up_proj") + if mlp is not None and _is_dense_mlp(mlp): + supported_targets.add("mlp") + if mlp is not None and (hasattr(mlp, "_run_routed_experts") or hasattr(mlp, "_run_local_routed_experts")): + supported_targets.add("routed_experts") + if linear_attn is not None: + supported_targets.add("linear_attn") + + return frozenset(supported_targets) + + +def set_selective_activation_checkpointing(layer: nn.Module, targets: Iterable[str]) -> None: + target_set = frozenset(targets) + invalid_targets = target_set - SELECTIVE_AC_TARGETS + if invalid_targets: + raise ValueError(f"Unsupported selective activation checkpoint targets: {sorted(invalid_targets)}") + + enabled_targets = target_set & get_supported_targets(layer) + self_attn = getattr(layer, "self_attn", None) + mlp = getattr(layer, "mlp", None) + linear_attn = _get_linear_attn_module(layer) + attn_proj_is_subsumed = "linear_attn" in enabled_targets and linear_attn is self_attn + + if self_attn is not None and "attn_proj" in enabled_targets and not attn_proj_is_subsumed: + checkpoint_method(self_attn, "attn_projections") + checkpoint_method(self_attn, "output_proj") + if self_attn is not None and "mla_up_proj" in enabled_targets: + checkpoint_method(self_attn, "mla_up_proj") + if mlp is not None and "mlp" in enabled_targets: + checkpoint_method(mlp, "forward") + if mlp is not None and "routed_experts" in enabled_targets: + if hasattr(mlp, "_run_routed_experts"): + checkpoint_method(mlp, "_run_routed_experts") + if hasattr(mlp, "_run_local_routed_experts"): + checkpoint_method(mlp, "_run_local_routed_experts") + if linear_attn is not None and "linear_attn" in enabled_targets: + checkpoint_method(linear_attn, "forward") + if "norm" in enabled_targets: + _configure_norm_checkpointing(layer) + + +def supports_selective_activation_checkpointing(layer: nn.Module) -> bool: + return type(layer).__module__.startswith("prime_rl.trainer.models.") diff --git a/src/prime_rl/trainer/models/layers/cp_mamba.py b/src/prime_rl/trainer/models/layers/cp_mamba.py new file mode 100644 index 00000000000..4175ada241b --- /dev/null +++ b/src/prime_rl/trainer/models/layers/cp_mamba.py @@ -0,0 +1,240 @@ +"""All-to-all context parallelism for Mamba-2 SSM layers. + +Transposes between sequence-parallel and head-parallel layouts so that each +CP rank processes the full sequence on a subset of heads, then transposes +back. This avoids materializing the full [S, full_heads] tensor on any GPU. + +Communication pattern per Mamba layer: + 1. in_proj (token-parallel, no comm) + 2. all-to-all: [B, S/cp, D] -> [B, S, D/cp] (seq-to-head) + 3. conv1d + SSM on full sequence, local heads + 4. all-to-all: [B, S, D/cp] -> [B, S/cp, D] (head-to-seq) + 5. out_proj (token-parallel, no comm) +""" + +from __future__ import annotations + +import torch +import torch.distributed as dist + + +class _SeqToHeadParallel(torch.autograd.Function): + """[B, S/cp, D] -> [B, S, D/cp] via all-to-all.""" + + @staticmethod + def forward(ctx, x: torch.Tensor, cp_group: dist.ProcessGroup, cp_size: int) -> torch.Tensor: + ctx.cp_group = cp_group + ctx.cp_size = cp_size + return _all_to_all_seq_to_head(x, cp_group, cp_size) + + @staticmethod + def backward(ctx, grad_output: torch.Tensor) -> tuple[torch.Tensor, None, None]: + return _all_to_all_head_to_seq(grad_output, ctx.cp_group, ctx.cp_size), None, None + + +class _HeadToSeqParallel(torch.autograd.Function): + """[B, S, D/cp] -> [B, S/cp, D] via all-to-all.""" + + @staticmethod + def forward(ctx, x: torch.Tensor, cp_group: dist.ProcessGroup, cp_size: int) -> torch.Tensor: + ctx.cp_group = cp_group + ctx.cp_size = cp_size + return _all_to_all_head_to_seq(x, cp_group, cp_size) + + @staticmethod + def backward(ctx, grad_output: torch.Tensor) -> tuple[torch.Tensor, None, None]: + return _all_to_all_seq_to_head(grad_output, ctx.cp_group, ctx.cp_size), None, None + + +def _all_to_all_seq_to_head(x: torch.Tensor, cp_group: dist.ProcessGroup, cp_size: int) -> torch.Tensor: + """Transpose [B, S_local, D] -> [B, S_global, D_local] via all-to-all. + + Each rank sends its S_local tokens split into cp_size chunks along D, + and receives one D_local chunk from each of the cp_size ranks. + """ + B, S_local, D = x.shape + D_local = D // cp_size + + # Split D into cp_size chunks, put cp dim first for all_to_all scatter + # [B, S_local, D] -> [B, S_local, cp, D_local] -> [cp, B, S_local, D_local] + x = x.reshape(B, S_local, cp_size, D_local).permute(2, 0, 1, 3).contiguous() + output = torch.empty_like(x) + dist.all_to_all_single(output, x, group=cp_group) + # output[i] = data from rank i (its S_local tokens, our D_local slice) + # [cp, B, S_local, D_local] -> [B, cp, S_local, D_local] -> [B, cp*S_local, D_local] + return output.permute(1, 0, 2, 3).reshape(B, cp_size * S_local, D_local).contiguous() + + +def _all_to_all_head_to_seq(x: torch.Tensor, cp_group: dist.ProcessGroup, cp_size: int) -> torch.Tensor: + """Transpose [B, S_global, D_local] -> [B, S_local, D] via all-to-all. + + Inverse of _all_to_all_seq_to_head. + """ + B, S_global, D_local = x.shape + S_local = S_global // cp_size + + # Split S_global into cp_size chunks (one per rank), put cp dim first for scatter + # [B, S_global, D_local] -> [B, cp, S_local, D_local] -> [cp, B, S_local, D_local] + x = x.reshape(B, cp_size, S_local, D_local).permute(1, 0, 2, 3).contiguous() + output = torch.empty_like(x) + dist.all_to_all_single(output, x, group=cp_group) + # Each rank now has cp chunks of [B, S_local, D_local] — one D_local slice from each rank + # [cp, B, S_local, D_local] -> [B, S_local, cp, D_local] -> [B, S_local, D] + return output.permute(1, 2, 0, 3).reshape(B, S_local, cp_size * D_local).contiguous() + + +def seq_to_head_parallel(x: torch.Tensor, cp_group: dist.ProcessGroup, cp_size: int) -> torch.Tensor: + """[B, S/cp, D] -> [B, S, D/cp] with correct backward.""" + return _SeqToHeadParallel.apply(x, cp_group, cp_size) + + +def head_to_seq_parallel(x: torch.Tensor, cp_group: dist.ProcessGroup, cp_size: int) -> torch.Tensor: + """[B, S, D/cp] -> [B, S/cp, D] with correct backward.""" + return _HeadToSeqParallel.apply(x, cp_group, cp_size) + + +def mamba_cp_forward(mixer, hidden_states: torch.Tensor, cp_group: dist.ProcessGroup, cp_rank: int, cp_size: int): + """CP-aware Mamba-2 forward: in_proj -> all-to-all -> conv+SSM(local heads) -> all-to-all -> out_proj. + + Replaces the mixer's cuda_kernels_forward when CP > 1. Parameters are + sliced at runtime so FSDP / checkpoint shapes are unchanged. + """ + from mamba_ssm.ops.triton.ssd_combined import mamba_chunk_scan_combined + + batch_size, seq_len, _ = hidden_states.shape # seq_len = S/cp here + + # ── 1. in_proj (token-parallel, no communication needed) ── + projected_states = mixer.in_proj(hidden_states) + A = -torch.exp(mixer.A_log.float()) + + gate, hidden_states_B_C, time_step = torch.split( + projected_states, + [mixer.intermediate_size, mixer.conv_dim, mixer.num_heads], + dim=-1, + ) + + # ── 2. All-to-all: seq-parallel -> head-parallel ── + # gate: [B, S/cp, nheads * hdim] -> [B, S, (nheads/cp) * hdim] + # hidden_states_B_C: [B, S/cp, nheads*hdim + 2*ngroups*dstate] -> [B, S, ...] + # time_step: [B, S/cp, nheads] -> [B, S, nheads/cp] + gate = seq_to_head_parallel(gate, cp_group, cp_size) + time_step = seq_to_head_parallel(time_step, cp_group, cp_size) + + # hidden_states_B_C needs special handling: x part shards by heads, B/C parts shard by groups + # conv_dim = intermediate_size + 2 * ngroups * dstate + groups_time_state_size = mixer.n_groups * mixer.ssm_state_size + x_part, B_part, C_part = torch.split( + hidden_states_B_C, + [mixer.intermediate_size, groups_time_state_size, groups_time_state_size], + dim=-1, + ) + x_part = seq_to_head_parallel(x_part, cp_group, cp_size) + B_part = seq_to_head_parallel(B_part, cp_group, cp_size) + C_part = seq_to_head_parallel(C_part, cp_group, cp_size) + hidden_states_B_C = torch.cat([x_part, B_part, C_part], dim=-1) + + full_seq_len = hidden_states_B_C.shape[1] # S (full sequence) + + # ── 3. Local head dimensions ── + local_num_heads = mixer.num_heads // cp_size + local_intermediate_size = local_num_heads * mixer.head_dim + local_n_groups = mixer.n_groups // cp_size + + # Slice parameters for local heads + head_start = cp_rank * local_num_heads + head_end = head_start + local_num_heads + + local_A = A[head_start:head_end] + local_D = mixer.D[head_start:head_end] + local_dt_bias = mixer.dt_bias[head_start:head_end] + + # Slice conv1d weight/bias for local channels + # conv1d channels = intermediate_size + 2 * ngroups * dstate (same as conv_dim) + # After all-to-all, we have local channels: local_intermediate + 2 * local_ngroups * dstate + local_conv_dim = local_intermediate_size + 2 * local_n_groups * mixer.ssm_state_size + conv_x_start = cp_rank * local_intermediate_size + conv_x_end = conv_x_start + local_intermediate_size + conv_b_start = mixer.intermediate_size + cp_rank * local_n_groups * mixer.ssm_state_size + conv_b_end = conv_b_start + local_n_groups * mixer.ssm_state_size + conv_c_start = mixer.intermediate_size + groups_time_state_size + cp_rank * local_n_groups * mixer.ssm_state_size + conv_c_end = conv_c_start + local_n_groups * mixer.ssm_state_size + + conv_indices = torch.cat( + [ + torch.arange(conv_x_start, conv_x_end, device=mixer.conv1d.weight.device), + torch.arange(conv_b_start, conv_b_end, device=mixer.conv1d.weight.device), + torch.arange(conv_c_start, conv_c_end, device=mixer.conv1d.weight.device), + ] + ) + local_conv_weight = mixer.conv1d.weight[conv_indices] + local_conv_bias = mixer.conv1d.bias[conv_indices] if mixer.conv1d.bias is not None else None + + # ── 4. Conv1d on full sequence, local heads ── + hidden_states_B_C = torch.nn.functional.silu( + torch.nn.functional.conv1d( + hidden_states_B_C.transpose(1, 2), + local_conv_weight, + local_conv_bias, + groups=local_conv_dim, + padding=mixer.conv1d.weight.shape[2] - 1, + ).transpose(1, 2)[:, :full_seq_len] + ) + + local_groups_time_state_size = local_n_groups * mixer.ssm_state_size + hidden_states_local, B_local, C_local = torch.split( + hidden_states_B_C, + [local_intermediate_size, local_groups_time_state_size, local_groups_time_state_size], + dim=-1, + ) + + # ── 5. SSM scan on full sequence, local heads ── + dt_limit_kwargs = {} if mixer.time_step_limit is None else {"dt_limit": mixer.time_step_limit} + + scan_output, _ = mamba_chunk_scan_combined( + hidden_states_local.view(batch_size, full_seq_len, local_num_heads, mixer.head_dim), + time_step, + local_A, + B_local.view(batch_size, full_seq_len, local_n_groups, -1), + C_local.view(batch_size, full_seq_len, local_n_groups, -1), + chunk_size=mixer.chunk_size, + D=local_D, + z=None, + seq_idx=None, + return_final_states=True, + dt_bias=local_dt_bias, + dt_softplus=True, + **dt_limit_kwargs, + ) + scan_output = scan_output.view(batch_size, full_seq_len, local_intermediate_size) + + # ── 6. Gated RMSNorm with local weights ── + # Zamba2RMSNormGated normalizes within groups of size `group_size`. + # In head-parallel mode, we have local_intermediate_size features. + # The group_size stays the same, but group_count is reduced. + local_norm_weight = mixer.norm.weight[cp_rank * local_intermediate_size : (cp_rank + 1) * local_intermediate_size] + scan_output = _gated_rms_norm( + scan_output, gate, local_norm_weight, mixer.norm.variance_epsilon, mixer.norm.group_size + ) + + # ── 7. All-to-all: head-parallel -> seq-parallel ── + scan_output = head_to_seq_parallel(scan_output, cp_group, cp_size) + + # ── 8. out_proj (token-parallel, no communication needed) ── + return mixer.out_proj(scan_output) + + +def _gated_rms_norm( + x: torch.Tensor, gate: torch.Tensor, weight: torch.Tensor, eps: float, group_size: int +) -> torch.Tensor: + """Gated RMSNorm matching Zamba2RMSNormGated: gate -> group-wise RMSNorm -> weight.""" + input_dtype = x.dtype + x = x.to(torch.float32) + if gate is not None: + x = x * torch.nn.functional.silu(gate.to(torch.float32)) + *prefix_dims, last_dim = x.shape + group_count = last_dim // group_size + x_grouped = x.view(*prefix_dims, group_count, group_size) + variance = x_grouped.pow(2).mean(-1, keepdim=True) + x_grouped = x_grouped * torch.rsqrt(variance + eps) + x = x_grouped.view(*prefix_dims, last_dim) + return (weight * x).to(input_dtype) diff --git a/src/prime_rl/trainer/models/layers/fp8_grouped_gemm.py b/src/prime_rl/trainer/models/layers/fp8_grouped_gemm.py new file mode 100644 index 00000000000..b1a40c4543c --- /dev/null +++ b/src/prime_rl/trainer/models/layers/fp8_grouped_gemm.py @@ -0,0 +1,224 @@ +from __future__ import annotations + +try: + import deep_gemm +except ImportError: + deep_gemm = None # CPU-only environments don't ship deep_gemm; FP8 paths + # are GPU-only at runtime, so leaving the symbol None is safe — only the + # autograd Function bodies below actually call into it. +import torch + +from prime_rl.trainer.models.kernels.fp8_utils import ( + GROUP_ALIGNMENT, + build_grouped_layout, + grouped_per_block_cast_to_fp8_triton, + grouped_per_channel_cast_to_fp8_sm90_kmajor_triton, + grouped_per_token_cast_to_fp8_triton, + unpack_rows_triton, +) + + +def _compute_grad_weight( + x: torch.Tensor, + grad_output: torch.Tensor, + weight_shape: torch.Size, + padded_total_m: int, + block_to_group: torch.Tensor, + ks_tensor: torch.Tensor, + starts_tensor: torch.Tensor, + actual_ms_tensor: torch.Tensor, + block_starts_tensor: torch.Tensor, + aligned_ms: list[int], +) -> torch.Tensor: + x_k_major = grouped_per_channel_cast_to_fp8_sm90_kmajor_triton( + x, + padded_total_m, + block_to_group, + starts_tensor, + actual_ms_tensor, + ks_tensor, + block_starts_tensor, + False, + GROUP_ALIGNMENT, + ) + dy_k_major = grouped_per_channel_cast_to_fp8_sm90_kmajor_triton( + grad_output, + padded_total_m, + block_to_group, + starts_tensor, + actual_ms_tensor, + ks_tensor, + block_starts_tensor, + False, + GROUP_ALIGNMENT, + ) + grad_weight = torch.zeros(weight_shape, device=x.device, dtype=torch.float32) + deep_gemm.k_grouped_fp8_gemm_nt_contiguous( + x_k_major, + dy_k_major, + grad_weight, + aligned_ms, + ks_tensor, + grad_weight, + ) + return grad_weight.to(torch.bfloat16) + + +class _GroupedFP8Gemm(torch.autograd.Function): + @staticmethod + def forward( + ctx, + x: torch.Tensor, + weight: torch.Tensor, + offs: torch.Tensor, + ) -> torch.Tensor: + ( + total_m, + padded_total_m, + grouped_layout, + block_to_group, + ks_tensor, + starts_tensor, + actual_ms_tensor, + block_starts_tensor, + ) = build_grouped_layout(offs, total_m=x.size(0)) + + x_fp8 = grouped_per_token_cast_to_fp8_triton( + x, + padded_total_m, + block_to_group, + starts_tensor, + actual_ms_tensor, + block_starts_tensor, + False, + GROUP_ALIGNMENT, + ) + weight_fp8 = grouped_per_block_cast_to_fp8_triton( + weight.transpose(1, 2), + False, + GROUP_ALIGNMENT, + ) + + out_padded = torch.empty( + (padded_total_m, weight.size(2)), + device=x.device, + dtype=x.dtype, + ) + deep_gemm.m_grouped_fp8_gemm_nt_contiguous( + x_fp8, + weight_fp8, + out_padded, + grouped_layout, + use_psum_layout=False, + ) + out = unpack_rows_triton( + out_padded, + total_m, + block_to_group, + starts_tensor, + actual_ms_tensor, + block_starts_tensor, + ) + + ctx.padded_total_m = padded_total_m + ctx.aligned_ms = ks_tensor.tolist() + ctx.save_for_backward( + x, + weight, + grouped_layout, + block_to_group, + ks_tensor, + starts_tensor, + actual_ms_tensor, + block_starts_tensor, + ) + return out + + @staticmethod + def backward(ctx, grad_output: torch.Tensor): + ( + x, + weight, + grouped_layout, + block_to_group, + ks_tensor, + starts_tensor, + actual_ms_tensor, + block_starts_tensor, + ) = ctx.saved_tensors + padded_total_m = ctx.padded_total_m + aligned_ms = ctx.aligned_ms + grad_output = grad_output.contiguous() + + grad_x = grad_weight = None + + if ctx.needs_input_grad[1]: + grad_weight = _compute_grad_weight( + x, + grad_output, + weight.shape, + padded_total_m, + block_to_group, + ks_tensor, + starts_tensor, + actual_ms_tensor, + block_starts_tensor, + aligned_ms, + ) + + if ctx.needs_input_grad[0]: + dy_fp8 = grouped_per_token_cast_to_fp8_triton( + grad_output, + padded_total_m, + block_to_group, + starts_tensor, + actual_ms_tensor, + block_starts_tensor, + False, + GROUP_ALIGNMENT, + ) + weight_dx_fp8 = grouped_per_block_cast_to_fp8_triton( + weight, + False, + GROUP_ALIGNMENT, + ) + grad_x_padded = torch.empty( + (padded_total_m, weight.size(1)), + device=grad_output.device, + dtype=grad_output.dtype, + ) + deep_gemm.m_grouped_fp8_gemm_nt_contiguous( + dy_fp8, + weight_dx_fp8, + grad_x_padded, + grouped_layout, + use_psum_layout=False, + ) + grad_x = unpack_rows_triton( + grad_x_padded, + x.size(0), + block_to_group, + starts_tensor, + actual_ms_tensor, + block_starts_tensor, + ) + + return grad_x, grad_weight, None + + +def grouped_fp8_gemm( + x: torch.Tensor, + weight: torch.Tensor, + offs: torch.Tensor, +) -> torch.Tensor: + """FP8 grouped GEMM, drop-in replacement for torch._grouped_mm. + + Args: + x: (M, K) concatenated token activations in bfloat16. + weight: (G, K, N) expert weights in bfloat16. + offs: (G,) int32 cumulative token counts per expert. + + Returns: + (M, N) output tensor in bfloat16. + """ + return _GroupedFP8Gemm.apply(x, weight, offs) diff --git a/src/prime_rl/trainer/models/layers/fp8_linear.py b/src/prime_rl/trainer/models/layers/fp8_linear.py new file mode 100644 index 00000000000..5f7f675dc34 --- /dev/null +++ b/src/prime_rl/trainer/models/layers/fp8_linear.py @@ -0,0 +1,179 @@ +from __future__ import annotations + +import re + +try: + import deep_gemm +except ImportError: + deep_gemm = None # CPU-only environments don't ship deep_gemm; FP8 paths + # are GPU-only at runtime, so leaving the symbol None is safe — only the + # autograd Function bodies below actually call into it. +import torch +from torch import nn + +from prime_rl.trainer.models.kernels.fp8_utils import ( + per_block_cast_to_fp8_triton, + per_token_cast_to_fp8_triton, +) +from prime_rl.utils.logger import get_logger + + +class _FP8BlockwiseMM(torch.autograd.Function): + @staticmethod + def forward(ctx, x, weight, block_size, out_dtype=torch.bfloat16): + x_shape = x.shape + x_2d = x.reshape(-1, x_shape[-1]).contiguous() + x_fp8 = per_token_cast_to_fp8_triton(x_2d, False, block_size) + weight_fp8 = per_block_cast_to_fp8_triton(weight, False, block_size) + + out = torch.empty((x_2d.size(0), weight.size(0)), device=x.device, dtype=out_dtype) + deep_gemm.fp8_gemm_nt(x_fp8, weight_fp8, out) + + ctx.save_for_backward(x_2d, weight) + ctx.x_shape = x_shape + ctx.block_size = block_size + return out.reshape(*x_shape[:-1], out.size(-1)) + + @staticmethod + def backward(ctx, grad_output): + x_2d, weight = ctx.saved_tensors + block_size = ctx.block_size + grad_output_2d = grad_output.reshape(-1, grad_output.shape[-1]).contiguous() + + grad_x = grad_weight = None + if ctx.needs_input_grad[0]: + grad_output_fp8 = per_token_cast_to_fp8_triton(grad_output_2d, False, block_size) + weight_t = weight.transpose(0, 1).contiguous() + weight_dx_fp8 = per_block_cast_to_fp8_triton(weight_t, False, block_size) + grad_x_2d = torch.empty_like(x_2d) + deep_gemm.fp8_gemm_nt(grad_output_fp8, weight_dx_fp8, grad_x_2d) + grad_x = grad_x_2d.reshape(ctx.x_shape) + + if ctx.needs_input_grad[1]: + # deep_gemm.fp8_gemm_nt with recipe=(1, 1, 128) requires the K (token) + # dim to be a multiple of 128. Zero-pad along the token axis so non- + # aligned per-rank batches (from sequence packing) don't trip the kernel. + M_tok = grad_output_2d.size(0) + M_pad = (M_tok + block_size - 1) // block_size * block_size + if M_pad != M_tok: + pad_rows = M_pad - M_tok + grad_output_2d_padded = torch.nn.functional.pad(grad_output_2d, (0, 0, 0, pad_rows)) + x_2d_padded = torch.nn.functional.pad(x_2d, (0, 0, 0, pad_rows)) + else: + grad_output_2d_padded = grad_output_2d + x_2d_padded = x_2d + grad_output_t = grad_output_2d_padded.transpose(0, 1).contiguous() + x_t = x_2d_padded.transpose(0, 1).contiguous() + grad_output_t_fp8 = per_token_cast_to_fp8_triton(grad_output_t, False, block_size) + x_t_fp8 = per_token_cast_to_fp8_triton(x_t, False, block_size) + grad_weight_fp32 = torch.zeros_like(weight, dtype=torch.float32) + deep_gemm.fp8_gemm_nt( + grad_output_t_fp8, + x_t_fp8, + grad_weight_fp32, + c=grad_weight_fp32, + recipe=(1, 1, 128), + ) + grad_weight = grad_weight_fp32.to(weight.dtype) + + return grad_x, grad_weight, None, None + + +class Float8BlockwiseLinear(nn.Linear): + """nn.Linear replacement that uses FP8 blockwise matmul via DeepGEMM. + + Requires: + - SM90 (Hopper) GPU + - bfloat16 inputs/weights + - No bias + - in_features and out_features divisible by 128 + """ + + def __init__(self, *args, block_size: int = 128, dtype=torch.bfloat16, **kwargs): + super().__init__(*args, **kwargs) + self.block_size = block_size + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return _FP8BlockwiseMM.apply(x, self.weight, self.block_size, torch.bfloat16) + + @classmethod + def from_linear(cls, mod: nn.Linear) -> "Float8BlockwiseLinear": + """Convert an existing nn.Linear to Float8BlockwiseLinear.""" + with torch.device("meta"): + new_mod = cls( + mod.in_features, + mod.out_features, + bias=mod.bias is not None, + ) + new_mod.weight = mod.weight + new_mod.bias = mod.bias + return new_mod + + +DEFAULT_FP8_IGNORE_PATTERNS: list[str] = [ + "lm_head", + "router", + # Use escaped dots — re.search treats `.` as any-char, so the previous + # "mlp.gate." pattern was also matching dense MLP `mlp.gate_proj` (the + # trailing `.` was matching `_`). That left the dense MLP gate projection + # in BF16 on the trainer while inference quantized it to FP8, causing + # hidden-state drift before the MoE router. + r"mlp\.gate\.", + "shared_expert_gate", # Qwen3.5 MoE: nn.Linear(hidden, 1, bias=False) + "eh_proj", + "weights_proj", + "in_proj_a", + "in_proj_b", +] + + +def replace_linear_with_fp8_blockwise_linear(model: nn.Module, ignore_modules: list[str] | None = None) -> None: + """Replace nn.Linear in `model` with Float8BlockwiseLinear, skipping any + module whose qualified name matches an ignore pattern (substring or regex). + + The default ignore list covers layers that should never be quantized: + - lm_head + - MoE routers and gates (router, mlp.gate., shared_expert_gate) + - sparse-MLA scalar projection (weights_proj) + - GLM-5.1 MTP head (eh_proj) + - hybrid-Mamba projections (in_proj_a, in_proj_b) + + Independently of the name-based ignore list, we also skip any nn.Linear + whose in_features or out_features is not a multiple of 128. Float8BlockwiseLinear + documents that requirement and DeepGEMM's fp8_gemm_nt crashes at runtime + on unaligned dims — better to keep them in BF16 with a clear log line than + silently break in the kernel. + + Conv1d, layer norms, and embedding tables are not nn.Linear and are + skipped automatically by the type check; we don't need to list them. + """ + if ignore_modules is None: + ignore_modules = list(DEFAULT_FP8_IGNORE_PATTERNS) + logger = get_logger() + logger.info(f"Replacing linear layers with FP8 blockwise linear layers (ignore={ignore_modules})") + replaced_modules = [] + skipped_modules = [] + skipped_unaligned: list[str] = [] + named_modules = dict(model.named_modules()) + for name, module in named_modules.items(): + if not isinstance(module, nn.Linear): + continue + if any(re.search(pattern, name) for pattern in ignore_modules): + skipped_modules.append(name) + continue + if module.in_features % 128 != 0 or module.out_features % 128 != 0: + skipped_unaligned.append(f"{name}({module.in_features}->{module.out_features})") + continue + parent_name, attr_name = name.rsplit(".", 1) if "." in name else ("", name) + parent = model.get_submodule(parent_name) if parent_name else model + setattr(parent, attr_name, Float8BlockwiseLinear.from_linear(module)) + replaced_modules.append(name) + + logger.info( + f"Replaced {len(replaced_modules)} linear layers with FP8 blockwise linear " + f"(skipped {len(skipped_modules)} by name, " + f"{len(skipped_unaligned)} by 128-divisibility); " + f"first replaced={replaced_modules[:3]}, " + f"first skipped(name)={skipped_modules[:3]}, " + f"first skipped(unaligned)={skipped_unaligned[:3]}" + ) diff --git a/src/prime_rl/trainer/models/layers/lm_head.py b/src/prime_rl/trainer/models/layers/lm_head.py index b14b64099a5..84556425a3c 100644 --- a/src/prime_rl/trainer/models/layers/lm_head.py +++ b/src/prime_rl/trainer/models/layers/lm_head.py @@ -9,6 +9,8 @@ from prime_rl.utils.logger import get_logger +FUSED_CE_IGNORE_INDEX = -100 + class PrimeLmOutput(TypedDict, total=False): """Output from LM head - a TypedDict so pytree can find tensors for FSDP2 hooks.""" @@ -16,6 +18,7 @@ class PrimeLmOutput(TypedDict, total=False): logits: Tensor | None logprobs: Tensor | None entropy: Tensor | None + loss: Tensor | None def cast_float_and_contiguous(output: PrimeLmOutput) -> PrimeLmOutput: @@ -28,6 +31,7 @@ def _float_and_contiguous(tensor: Tensor | None) -> Tensor | None: logits=_float_and_contiguous(output.get("logits")), logprobs=_float_and_contiguous(output.get("logprobs")), entropy=_float_and_contiguous(output.get("entropy")), + loss=output.get("loss"), ) @@ -40,16 +44,19 @@ def forward( self, hidden_states: torch.Tensor, labels: torch.Tensor | None = None, - temperature: float = 1.0, + temperature: Tensor | None = None, ) -> PrimeLmOutput: assert labels is not None, "FusedOutputLinear requires labels for chunked logprob computation" + assert temperature is not None, "FusedOutputLinear requires per-token temperatures" - inv_t = 1.0 / float(temperature) b, s, h = hidden_states.shape hidden_states = hidden_states.reshape(b * s, h).contiguous() labels = labels.reshape(b * s).contiguous() + inv_t = 1.0 / temperature.reshape(b * s).contiguous() # [N] - logprobs, entropy = _ChunkedLogProbEntropyFn.apply(hidden_states, self.weight, labels, inv_t, self.chunk_size) + logprobs, entropy = _SequenceChunkedLogProbEntropyFn.apply( + hidden_states, self.weight, labels, inv_t, self.chunk_size + ) logprobs = logprobs.reshape(b, s) entropy = entropy.reshape(b, s) @@ -61,69 +68,160 @@ def __init__(self, in_features: int, out_features: int): super().__init__(in_features, out_features, bias=False) def forward( - self, hidden_states: torch.Tensor, labels: torch.Tensor | None = None, temperature: float = 1.0 + self, hidden_states: torch.Tensor, labels: torch.Tensor | None = None, temperature: Tensor | None = None ) -> PrimeLmOutput: + # VanillaOutputLinear just returns logits - temperature scaling is done externally in train.py return PrimeLmOutput(logits=super().forward(hidden_states)) -class _ChunkedLogProbEntropyFn(torch.autograd.Function): +class FusedCrossEntropyOutputLinear(torch.nn.Linear): + """Fused lm_head + cross-entropy loss using Liger kernel. + + Avoids materializing the full [N, V] logits tensor by fusing the linear + projection with the cross-entropy loss computation. + """ + + IGNORE_INDEX = FUSED_CE_IGNORE_INDEX + + def __init__(self, in_features: int, out_features: int, softcap: float | None = None): + super().__init__(in_features, out_features, bias=False) + from liger_kernel.transformers.fused_linear_cross_entropy import LigerFusedLinearCrossEntropyLoss + + self.fused_ce = LigerFusedLinearCrossEntropyLoss( + ignore_index=self.IGNORE_INDEX, reduction="mean", softcap=softcap + ) + + def forward( + self, + hidden_states: torch.Tensor, + labels: torch.Tensor | None = None, + temperature: Tensor | None = None, + ) -> PrimeLmOutput: + if labels is None: + return PrimeLmOutput(logits=super().forward(hidden_states)) + + b, s, h = hidden_states.shape + hidden_flat = hidden_states.reshape(b * s, h).contiguous() + labels_flat = labels.reshape(b * s).contiguous() + loss = self.fused_ce(self.weight, hidden_flat, labels_flat) + return PrimeLmOutput(loss=loss) + + +class QuackFusedCrossEntropyOutputLinear(torch.nn.Linear): + """Fused lm_head + cross-entropy loss using quack-kernels. + + Chunks the linear projection and cross-entropy computation to avoid + materializing the full [N, V] logits tensor, using quack's optimized + CuTe DSL kernels for CE and GEMM. + """ + + IGNORE_INDEX = FUSED_CE_IGNORE_INDEX + + def __init__(self, in_features: int, out_features: int, chunk_size: int = 4096): + super().__init__(in_features, out_features, bias=False) + self.chunk_size = chunk_size + + def forward( + self, + hidden_states: torch.Tensor, + labels: torch.Tensor | None = None, + temperature: Tensor | None = None, + ) -> PrimeLmOutput: + if labels is None: + return PrimeLmOutput(logits=super().forward(hidden_states)) + + from quack.linear_cross_entropy import chunked_linear_cross_entropy + + b, s, h = hidden_states.shape + hidden_flat = hidden_states.reshape(b * s, h).contiguous() + labels_flat = labels.reshape(b * s).contiguous() + loss = chunked_linear_cross_entropy( + hidden_flat, + self.weight, + labels_flat, + chunk_size=self.chunk_size, + ignore_index=self.IGNORE_INDEX, + reduction="mean", + ) + return PrimeLmOutput(loss=loss) + + +def _online_logsumexp_and_weighted_update( + m: torch.Tensor, s: torch.Tensor, t: torch.Tensor, chunk_logits: torch.Tensor +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + chunk_m = torch.amax(chunk_logits, dim=-1) + m_new = torch.maximum(m, chunk_m) + exp_old = torch.exp(m - m_new) + + chunk_exp = torch.exp(chunk_logits - m_new.unsqueeze(-1)) + s_new = s * exp_old + chunk_exp.sum(dim=-1) + t_new = t * exp_old + (chunk_exp * chunk_logits).sum(dim=-1) + return m_new, s_new, t_new + + +class _SequenceChunkedLogProbEntropyFn(torch.autograd.Function): @staticmethod def forward( # type: ignore[override] ctx, hidden: torch.Tensor, # [N, H] weight: torch.Tensor, # [V, H] labels: torch.Tensor, # [N] - inv_temperature: float, + inv_temperature: torch.Tensor, # [N] chunk_size: int, ) -> tuple[torch.Tensor, torch.Tensor]: """ - Returns (per-token logprobs, per-token entropy) without materializing [N, V]. - - Important: entropy is computed from the *same* per-chunk logits used for the softmax - normalization (no extra W @ hidden matmul). + Returns per-token logprobs and entropy by chunking over flattened sequence tokens. """ assert hidden.dim() == 2, f"expected hidden [N,H], got {tuple(hidden.shape)}" assert weight.dim() == 2, f"expected weight [V,H], got {tuple(weight.shape)}" assert labels.dim() == 1, f"expected labels [N], got {tuple(labels.shape)}" + assert inv_temperature.dim() == 1, f"expected inv_temperature [N], got {tuple(inv_temperature.shape)}" assert hidden.shape[0] == labels.shape[0], "hidden/labels N mismatch" assert hidden.shape[1] == weight.shape[1], "hidden/weight H mismatch" + assert hidden.shape[0] == inv_temperature.shape[0], "hidden/inv_temperature N mismatch" assert chunk_size > 0 device = hidden.device n = hidden.shape[0] vocab = weight.shape[0] - - # Running stats in fp32. - m = torch.full((n,), float("-inf"), device=device, dtype=torch.float32) - s = torch.zeros((n,), device=device, dtype=torch.float32) - t = torch.zeros((n,), device=device, dtype=torch.float32) - target_logits = torch.zeros((n,), device=device, dtype=torch.float32) - - for start in range(0, vocab, chunk_size): - end = min(start + chunk_size, vocab) - w_chunk = weight[start:end] # [C, H] - logits = hidden @ w_chunk.t() # [N, C] (model dtype) - logits_f = logits.to(torch.float32).mul_(inv_temperature) # [N, C] fp32 - - # Shared intermediates for logZ and entropy stats. - m, s, t = _online_logsumexp_and_weighted_update(m, s, t, logits_f) - - # Fill target logits for labels that fall in this chunk. - mask = (labels >= start) & (labels < end) - if torch.any(mask): - idx = (labels[mask] - start).to(torch.long) - target_logits[mask] = logits_f[mask, idx] - - logz = m + torch.log(s) - logprobs = target_logits - logz - entropy = logz - (t / s) - - # Save for backward (recompute logits per chunk for grad) - ctx.save_for_backward(hidden, weight, labels, logz) - ctx.inv_temperature = inv_temperature + vocab_chunk_size = min(vocab, 8192) + logprobs = torch.empty((n,), device=device, dtype=torch.float32) + entropy = torch.empty((n,), device=device, dtype=torch.float32) + logz = torch.empty((n,), device=device, dtype=torch.float32) + + for start in range(0, n, chunk_size): + end = min(start + chunk_size, n) + hidden_chunk = hidden[start:end] + labels_chunk = labels[start:end] + inv_t_chunk = inv_temperature[start:end].unsqueeze(-1) + token_count = end - start + + m = torch.full((token_count,), float("-inf"), device=device, dtype=torch.float32) + s = torch.zeros((token_count,), device=device, dtype=torch.float32) + t = torch.zeros((token_count,), device=device, dtype=torch.float32) + target_logits = torch.zeros((token_count,), device=device, dtype=torch.float32) + + for vocab_start in range(0, vocab, vocab_chunk_size): + vocab_end = min(vocab_start + vocab_chunk_size, vocab) + weight_chunk = weight[vocab_start:vocab_end] + logits_chunk = hidden_chunk @ weight_chunk.t() + scaled_logits = logits_chunk.to(torch.float32) * inv_t_chunk + + m, s, t = _online_logsumexp_and_weighted_update(m, s, t, scaled_logits) + + mask = (labels_chunk >= vocab_start) & (labels_chunk < vocab_end) + if torch.any(mask): + idx = (labels_chunk[mask] - vocab_start).to(torch.long) + target_logits[mask] = scaled_logits[mask, idx] + + logz_chunk = m + torch.log(s) + logz[start:end] = logz_chunk + logprobs[start:end] = target_logits - logz_chunk + entropy[start:end] = logz_chunk - (t / s) + + ctx.save_for_backward(hidden, weight, labels, inv_temperature, logz) ctx.chunk_size = chunk_size - # Return fp32 for numerical stability (matching baseline behavior). return logprobs, entropy @staticmethod @@ -132,77 +230,63 @@ def backward(ctx, grad_logprobs: torch.Tensor, grad_entropy: torch.Tensor | None "Backward through entropy is not implemented in FusedOutputLinear" ) - hidden, weight, labels, logz = ctx.saved_tensors - inv_temperature: float = ctx.inv_temperature + hidden, weight, labels, inv_temperature, logz = ctx.saved_tensors chunk_size: int = ctx.chunk_size - n, h = hidden.shape + n, _ = hidden.shape vocab = weight.shape[0] + vocab_chunk_size = min(vocab, 8192) grad_hidden = torch.zeros_like(hidden) grad_weight = torch.zeros_like(weight) - g = grad_logprobs.to(torch.float32) # [N] fp32 for stable scaling - - for start in range(0, vocab, chunk_size): - end = min(start + chunk_size, vocab) - w_chunk = weight[start:end] # [C, H] - - logits = hidden @ w_chunk.t() # [N, C] (model dtype) - logits_f = logits.to(torch.float32).mul_(inv_temperature) # [N, C] fp32 - - # p = softmax(logits_f) chunk = exp(logits_f - logz) - p = torch.exp(logits_f - logz.unsqueeze(-1)) # [N, C] fp32 - - # dL/dlogits = g * (1_{label} - p) - grad_logits = (-g).unsqueeze(-1) * p # [N, C] fp32 - mask = (labels >= start) & (labels < end) - if torch.any(mask): - idx = (labels[mask] - start).to(torch.long) - grad_logits[mask, idx] += g[mask] - - # Chain through temperature scaling: logits_f = logits * inv_temperature - grad_logits.mul_(inv_temperature) - - grad_hidden.add_(grad_logits.to(hidden.dtype) @ w_chunk) - grad_w_chunk = grad_logits.to(weight.dtype).t() @ hidden # [C, H] - grad_weight[start:end].add_(grad_w_chunk) + for start in range(0, n, chunk_size): + end = min(start + chunk_size, n) + hidden_chunk = hidden[start:end] + labels_chunk = labels[start:end] + grad_chunk = grad_logprobs[start:end].to(torch.float32) + inv_t_chunk = inv_temperature[start:end].unsqueeze(-1) + logz_chunk = logz[start:end] + + for vocab_start in range(0, vocab, vocab_chunk_size): + vocab_end = min(vocab_start + vocab_chunk_size, vocab) + weight_chunk = weight[vocab_start:vocab_end] + logits_chunk = hidden_chunk @ weight_chunk.t() + scaled_logits = logits_chunk.to(torch.float32) * inv_t_chunk + probs = torch.exp(scaled_logits - logz_chunk.unsqueeze(-1)) + + grad_logits = (-grad_chunk).unsqueeze(-1) * probs + mask = (labels_chunk >= vocab_start) & (labels_chunk < vocab_end) + if torch.any(mask): + idx = (labels_chunk[mask] - vocab_start).to(torch.long) + grad_logits[mask, idx] += grad_chunk[mask] + grad_logits = grad_logits * inv_t_chunk + + grad_hidden[start:end].add_(grad_logits.to(hidden.dtype) @ weight_chunk) + grad_weight[vocab_start:vocab_end].add_(grad_logits.to(weight.dtype).t() @ hidden_chunk) return grad_hidden, grad_weight, None, None, None -def _online_logsumexp_and_weighted_update( - m: torch.Tensor, s: torch.Tensor, t: torch.Tensor, chunk_logits: torch.Tensor -) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - """ - Online logsumexp + weighted-sum accumulator for entropy. - - Maintains: - m: running max - s: running sum(exp(x - m)) - t: running sum(exp(x - m) * x) - """ - chunk_m = torch.amax(chunk_logits, dim=-1) # [N] - m_new = torch.maximum(m, chunk_m) # [N] - exp_old = torch.exp(m - m_new) - - chunk_exp = torch.exp(chunk_logits - m_new.unsqueeze(-1)) # [N, C] - s_new = s * exp_old + chunk_exp.sum(dim=-1) - t_new = t * exp_old + (chunk_exp * chunk_logits).sum(dim=-1) - return m_new, s_new, t_new - - -def inject_prime_lm_head(model: nn.Module, chunk_size: int | None = None) -> None: +def inject_prime_lm_head( + model: nn.Module, + chunk_size: int | None = None, + fused_cross_entropy: bool | str = False, +) -> None: """ - Inject a PrimeRL LM head (FusedOutputLinear or VanillaOutputLinear) into a model. + Inject a PrimeRL LM head into a model. This replaces the model's lm_head and overrides the forward method to use labels and temperature for chunked loss computation. Args: model: The model to wrap. - chunk_size: If provided, use FusedOutputLinear with chunked logprob/entropy computation. - If None, use VanillaOutputLinear which just returns logits. + chunk_size: When set to an int, uses FusedOutputLinear with sequence-token chunked + logprob/entropy computation (for RL). + fused_cross_entropy: Controls fused lm_head + CE loss. Accepts: + - False: no fusion + - True or "liger": Liger kernel fusion + - "quack": quack-kernels fusion (chunked linear + CE with CuTe DSL kernels) """ # Guards so we have nicer error messages when a non-standard model is used assert hasattr(model, "model"), f"model doesnt have backbone in model.model:\n{model}" @@ -214,19 +298,51 @@ def inject_prime_lm_head(model: nn.Module, chunk_size: int | None = None) -> Non ) logger = get_logger() - logger.info(f"Injecting Prime LM head with chunk size {chunk_size}") + + # Check for Gemma-style softcapping - dispatch to specialized implementation + final_logit_softcapping = getattr(model.config, "final_logit_softcapping", None) + if final_logit_softcapping: + if fused_cross_entropy == "quack": + raise ValueError( + "quack_fused does not support Gemma logit softcapping. " + "Use loss_impl='liger_fused' or loss_impl='torch' instead." + ) + if not fused_cross_entropy: + from prime_rl.trainer.models.layers.lm_head_gemma import inject_gemma_lm_head + + inject_gemma_lm_head(model, chunk_size, final_logit_softcapping) + return # Replace the lm_head with the appropriate wrapper old_lm_head = model.lm_head - if chunk_size is not None: + if fused_cross_entropy == "quack": + logger.info("Injecting fused cross-entropy LM head (quack-kernels)") + model.lm_head = QuackFusedCrossEntropyOutputLinear( + in_features=old_lm_head.in_features, + out_features=old_lm_head.out_features, + ) + elif fused_cross_entropy: + logger.info("Injecting fused cross-entropy LM head (Liger kernel)") + model.lm_head = FusedCrossEntropyOutputLinear( + in_features=old_lm_head.in_features, + out_features=old_lm_head.out_features, + softcap=final_logit_softcapping, + ) + elif isinstance(chunk_size, int): + logger.info(f"Injecting chunked LM head with chunk size {chunk_size}") model.lm_head = FusedOutputLinear( in_features=old_lm_head.in_features, out_features=old_lm_head.out_features, chunk_size=chunk_size ) else: + logger.info("Injecting vanilla LM head") model.lm_head = VanillaOutputLinear(in_features=old_lm_head.in_features, out_features=old_lm_head.out_features) model.lm_head.weight = old_lm_head.weight del old_lm_head + _patch_model_forward(model) + + +def _patch_model_forward(model: nn.Module) -> None: # Patch the forward method to use the new lm_head with labels and temperature def new_forward( self: nn.Module, @@ -235,10 +351,12 @@ def new_forward( inputs_embeds: torch.Tensor | None = None, labels: torch.Tensor | None = None, logits_to_keep: int = 0, - temperature: float = 1.0, + temperature: torch.Tensor | None = None, **kwargs: object, ) -> PrimeLmOutput: - if position_ids is None: + # For VLM with images, don't create position_ids - let model compute MRoPE internally + is_multimodal = kwargs.get("pixel_values") is not None + if position_ids is None and not is_multimodal: reference_tensor = input_ids if input_ids is not None else inputs_embeds position_ids = torch.arange(1, reference_tensor.shape[1] + 1, device=reference_tensor.device).unsqueeze(0) outputs = self.model( @@ -258,7 +376,7 @@ def new_forward( return self.lm_head( hidden_states[:, slice_indices, :], labels[:, slice_indices] if labels is not None else None, - temperature=temperature, + temperature=temperature[:, slice_indices] if temperature is not None else None, ) # Bind the new forward to the model diff --git a/src/prime_rl/trainer/models/layers/lm_head_gemma.py b/src/prime_rl/trainer/models/layers/lm_head_gemma.py new file mode 100644 index 00000000000..7fee1a8d87f --- /dev/null +++ b/src/prime_rl/trainer/models/layers/lm_head_gemma.py @@ -0,0 +1,196 @@ +from __future__ import annotations + +import torch +import torch.nn as nn +from torch import Tensor + +from prime_rl.trainer.models.layers.lm_head import ( + PrimeLmOutput, + _online_logsumexp_and_weighted_update, + _patch_model_forward, +) +from prime_rl.utils.logger import get_logger + + +class GemmaFusedOutputLinear(torch.nn.Linear): + def __init__(self, in_features: int, out_features: int, chunk_size: int, softcap: float): + super().__init__(in_features, out_features, bias=False) + self.chunk_size = chunk_size + self.softcap = softcap + + def forward( + self, + hidden_states: torch.Tensor, + labels: torch.Tensor | None = None, + temperature: Tensor | None = None, + ) -> PrimeLmOutput: + assert labels is not None, "GemmaFusedOutputLinear requires labels for chunked logprob computation" + assert temperature is not None, "GemmaFusedOutputLinear requires per-token temperatures" + + b, s, h = hidden_states.shape + hidden_states = hidden_states.reshape(b * s, h).contiguous() + labels = labels.reshape(b * s).contiguous() + inv_t = 1.0 / temperature.reshape(b * s).contiguous() # [N] + + logprobs, entropy = _GemmaChunkedLogProbEntropyFn.apply( + hidden_states, self.weight, labels, inv_t, self.chunk_size, self.softcap + ) + + logprobs = logprobs.reshape(b, s) + entropy = entropy.reshape(b, s) + return PrimeLmOutput(logprobs=logprobs, entropy=entropy) + + +class GemmaVanillaOutputLinear(torch.nn.Linear): + def __init__(self, in_features: int, out_features: int, softcap: float): + super().__init__(in_features, out_features, bias=False) + self.softcap = softcap + + def forward( + self, hidden_states: torch.Tensor, labels: torch.Tensor | None = None, temperature: Tensor | None = None + ) -> PrimeLmOutput: + logits = super().forward(hidden_states) + logits = self.softcap * torch.tanh(logits / self.softcap) + return PrimeLmOutput(logits=logits) + + +class _GemmaChunkedLogProbEntropyFn(torch.autograd.Function): + @staticmethod + def forward( # type: ignore[override] + ctx, + hidden: torch.Tensor, # [N, H] + weight: torch.Tensor, # [V, H] + labels: torch.Tensor, # [N] + inv_temperature: torch.Tensor, # [N] + chunk_size: int, + softcap: float, + ) -> tuple[torch.Tensor, torch.Tensor]: + """ + Returns per-token logprobs and entropy by chunking over flattened sequence tokens. + """ + assert hidden.dim() == 2, f"expected hidden [N,H], got {tuple(hidden.shape)}" + assert weight.dim() == 2, f"expected weight [V,H], got {tuple(weight.shape)}" + assert labels.dim() == 1, f"expected labels [N], got {tuple(labels.shape)}" + assert inv_temperature.dim() == 1, f"expected inv_temperature [N], got {tuple(inv_temperature.shape)}" + assert hidden.shape[0] == labels.shape[0], "hidden/labels N mismatch" + assert hidden.shape[1] == weight.shape[1], "hidden/weight H mismatch" + assert hidden.shape[0] == inv_temperature.shape[0], "hidden/inv_temperature N mismatch" + assert chunk_size > 0 + + device = hidden.device + n = hidden.shape[0] + vocab = weight.shape[0] + vocab_chunk_size = min(vocab, 8192) + logprobs = torch.empty((n,), device=device, dtype=torch.float32) + entropy = torch.empty((n,), device=device, dtype=torch.float32) + logz = torch.empty((n,), device=device, dtype=torch.float32) + + for start in range(0, n, chunk_size): + end = min(start + chunk_size, n) + hidden_chunk = hidden[start:end] + labels_chunk = labels[start:end] + inv_t_chunk = inv_temperature[start:end].unsqueeze(-1) + token_count = end - start + + m = torch.full((token_count,), float("-inf"), device=device, dtype=torch.float32) + s = torch.zeros((token_count,), device=device, dtype=torch.float32) + t = torch.zeros((token_count,), device=device, dtype=torch.float32) + target_logits = torch.zeros((token_count,), device=device, dtype=torch.float32) + + for vocab_start in range(0, vocab, vocab_chunk_size): + vocab_end = min(vocab_start + vocab_chunk_size, vocab) + weight_chunk = weight[vocab_start:vocab_end] + logits_chunk = hidden_chunk @ weight_chunk.t() + scaled_logits = logits_chunk.to(torch.float32) + scaled_logits = softcap * torch.tanh(scaled_logits / softcap) + scaled_logits = scaled_logits * inv_t_chunk + + m, s, t = _online_logsumexp_and_weighted_update(m, s, t, scaled_logits) + + mask = (labels_chunk >= vocab_start) & (labels_chunk < vocab_end) + if torch.any(mask): + idx = (labels_chunk[mask] - vocab_start).to(torch.long) + target_logits[mask] = scaled_logits[mask, idx] + + logz_chunk = m + torch.log(s) + logz[start:end] = logz_chunk + logprobs[start:end] = target_logits - logz_chunk + entropy[start:end] = logz_chunk - (t / s) + + ctx.save_for_backward(hidden, weight, labels, inv_temperature, logz) + ctx.chunk_size = chunk_size + ctx.softcap = softcap + + return logprobs, entropy + + @staticmethod + def backward(ctx, grad_logprobs: torch.Tensor, grad_entropy: torch.Tensor | None): + assert grad_entropy is None or torch.all(grad_entropy == 0.0), ( + "Backward through entropy is not implemented in GemmaFusedOutputLinear" + ) + + hidden, weight, labels, inv_temperature, logz = ctx.saved_tensors + chunk_size: int = ctx.chunk_size + softcap: float = ctx.softcap + + n, _ = hidden.shape + vocab = weight.shape[0] + vocab_chunk_size = min(vocab, 8192) + + grad_hidden = torch.zeros_like(hidden) + grad_weight = torch.zeros_like(weight) + + for start in range(0, n, chunk_size): + end = min(start + chunk_size, n) + hidden_chunk = hidden[start:end] + labels_chunk = labels[start:end] + grad_chunk = grad_logprobs[start:end].to(torch.float32) + inv_t_chunk = inv_temperature[start:end].unsqueeze(-1) + logz_chunk = logz[start:end] + + for vocab_start in range(0, vocab, vocab_chunk_size): + vocab_end = min(vocab_start + vocab_chunk_size, vocab) + weight_chunk = weight[vocab_start:vocab_end] + logits_chunk = hidden_chunk @ weight_chunk.t() + logits_f = logits_chunk.to(torch.float32) + tanh_val = torch.tanh(logits_f / softcap) + scaled_logits = softcap * tanh_val + scaled_logits = scaled_logits * inv_t_chunk + probs = torch.exp(scaled_logits - logz_chunk.unsqueeze(-1)) + + grad_logits = (-grad_chunk).unsqueeze(-1) * probs + mask = (labels_chunk >= vocab_start) & (labels_chunk < vocab_end) + if torch.any(mask): + idx = (labels_chunk[mask] - vocab_start).to(torch.long) + grad_logits[mask, idx] += grad_chunk[mask] + grad_logits = grad_logits * inv_t_chunk + grad_logits = grad_logits * (1 - tanh_val**2) + + grad_hidden[start:end].add_(grad_logits.to(hidden.dtype) @ weight_chunk) + grad_weight[vocab_start:vocab_end].add_(grad_logits.to(weight.dtype).t() @ hidden_chunk) + + return grad_hidden, grad_weight, None, None, None, None + + +def inject_gemma_lm_head(model: nn.Module, chunk_size: int | None, softcap: float) -> None: + logger = get_logger() + logger.info(f"Injecting Gemma LM head with chunk size {chunk_size}, softcap={softcap}") + + old_lm_head = model.lm_head + if chunk_size is not None: + model.lm_head = GemmaFusedOutputLinear( + in_features=old_lm_head.in_features, + out_features=old_lm_head.out_features, + chunk_size=chunk_size, + softcap=softcap, + ) + else: + model.lm_head = GemmaVanillaOutputLinear( + in_features=old_lm_head.in_features, + out_features=old_lm_head.out_features, + softcap=softcap, + ) + model.lm_head.weight = old_lm_head.weight + del old_lm_head + + _patch_model_forward(model) diff --git a/src/prime_rl/trainer/models/layers/lora/__init__.py b/src/prime_rl/trainer/models/layers/lora/__init__.py index ecea8bfe3d3..9ceefa44538 100644 --- a/src/prime_rl/trainer/models/layers/lora/__init__.py +++ b/src/prime_rl/trainer/models/layers/lora/__init__.py @@ -6,12 +6,13 @@ set_multilora_scaling, ) from prime_rl.trainer.models.layers.lora.multi_linear import MultiLoRALinear -from prime_rl.trainer.models.layers.lora.multi_moe import MultiLoRAGroupedExperts +from prime_rl.trainer.models.layers.lora.multi_moe import MultiLoRAGroupedExperts, MultiLoRANonGatedGroupedExperts __all__ = [ "MultiLoRAModule", "MultiLoRALinear", "MultiLoRAGroupedExperts", + "MultiLoRANonGatedGroupedExperts", "set_lora_num_tokens", "get_lora_num_tokens", "set_multilora_scaling", diff --git a/src/prime_rl/trainer/models/layers/lora/base.py b/src/prime_rl/trainer/models/layers/lora/base.py index f65230237ee..f4f29cb0a65 100644 --- a/src/prime_rl/trainer/models/layers/lora/base.py +++ b/src/prime_rl/trainer/models/layers/lora/base.py @@ -5,7 +5,7 @@ from torch import nn if TYPE_CHECKING: - from prime_rl.trainer.runs import Runs + from prime_rl.trainer.runs import MultiRunManager _LORA_PREFIX = "base_layer." @@ -123,20 +123,20 @@ def get_lora_param_counts(self) -> tuple[int, int]: """ ... - def register_with_runs(self, runs: "Runs", prefix: str) -> None: - """Register this module with the Runs system. + def register_with_runs(self, runs: "MultiRunManager", prefix: str) -> None: + """Register this module with the MultiRunManager system. This method should be called after FSDP/compile/AC setup as these transformations may change the underlying parameters while preserving the module. - The Runs class will use this registration to: + The MultiRunManager class will use this registration to: - Get named parameters for specific run indices (for optimizer setup) - Reset parameters when new runs are created - Construct sliced state dicts for weight broadcast Args: - runs: The Runs instance to register with + runs: The MultiRunManager instance to register with prefix: The module's name/prefix in the model (e.g., "model.layers.0.self_attn.q_proj") """ runs.register_module(prefix, self) diff --git a/src/prime_rl/trainer/models/layers/lora/multi_linear.py b/src/prime_rl/trainer/models/layers/lora/multi_linear.py index 1656ac0f7ec..fd9b59e789d 100644 --- a/src/prime_rl/trainer/models/layers/lora/multi_linear.py +++ b/src/prime_rl/trainer/models/layers/lora/multi_linear.py @@ -53,13 +53,6 @@ def __init__( if rank <= 0 or n_adapters <= 0: raise ValueError("rank and n_adapters must be > 0") - # Set use_grouped_mm to False if CUDA compute capability < 9.0 - if torch.cuda.is_available(): - cc_major, _ = torch.cuda.get_device_capability() - if cc_major != 9: - use_grouped_mm = False - else: - use_grouped_mm = False if rank % 8 != 0 or base_layer.in_features % 8 != 0 or base_layer.out_features % 8 != 0: use_grouped_mm = False diff --git a/src/prime_rl/trainer/models/layers/lora/multi_moe.py b/src/prime_rl/trainer/models/layers/lora/multi_moe.py index 6d7ca705227..0d55dc24407 100644 --- a/src/prime_rl/trainer/models/layers/lora/multi_moe.py +++ b/src/prime_rl/trainer/models/layers/lora/multi_moe.py @@ -3,9 +3,17 @@ import torch import torch.nn.functional as F from torch import nn +from torch.distributed.tensor import DTensor from prime_rl.trainer.models.layers.lora.base import MultiLoRAModule, get_lora_num_tokens, get_multilora_scaling -from prime_rl.trainer.models.layers.moe import GroupedExperts +from prime_rl.trainer.models.layers.moe import ( + GptOssGroupedExperts, + GroupedExperts, + NonGatedGroupedExperts, + _broadcast_expert_bias, + _gpt_oss_apply_gate, + relu2, +) def _run_lora_grouped_mm( @@ -80,6 +88,7 @@ def __init__( n_adapters: int, alpha: float = 32.0, dropout: float = 0.0, + use_grouped_mm: bool = True, ): super().__init__(base_layer) if rank <= 0 or n_adapters <= 0: @@ -89,12 +98,6 @@ def __init__( self.dim = base_layer.w1.shape[2] self.hidden_dim = base_layer.w1.shape[1] - use_grouped_mm = False - if torch.cuda.is_available(): - cc_major, _ = torch.cuda.get_device_capability() - if cc_major == 9: # Hopper GPUs (H100, H200) - use_grouped_mm = True - # grouped_mm requires 8-byte alignment if rank % 8 != 0 or self.dim % 8 != 0 or self.hidden_dim % 8 != 0: use_grouped_mm = False @@ -290,6 +293,16 @@ def state_dict_for_adapter(self, idx: int) -> dict[str, torch.Tensor]: detached_w3_lora_a = self.w3_lora_A[idx].detach() detached_w3_lora_b = self.w3_lora_B[idx].detach() + # With EP, LoRA weights are DTensors sharded across expert-parallel ranks. + # Gather them before per-expert indexing. + if isinstance(detached_w1_lora_a, DTensor): + detached_w1_lora_a = detached_w1_lora_a.full_tensor() + detached_w1_lora_b = detached_w1_lora_b.full_tensor() + detached_w2_lora_a = detached_w2_lora_a.full_tensor() + detached_w2_lora_b = detached_w2_lora_b.full_tensor() + detached_w3_lora_a = detached_w3_lora_a.full_tensor() + detached_w3_lora_b = detached_w3_lora_b.full_tensor() + # The clone is necessary to avoid views that cause giant memory spikes # TODO: There's probably a better way to do this for expert_id in range(self.num_experts): @@ -325,6 +338,40 @@ def forward(self, x: torch.Tensor, num_tokens_per_expert: torch.Tensor) -> torch base_w2 = self.base_layer.w2 # [num_experts, dim, hidden_dim] base_w3 = self.base_layer.w3 # [num_experts, hidden_dim, dim] + # EP handling: convert DTensors to local shards. + # Standard EP also needs token permutation; DeepEP tokens are already dispatched. + permuted_indices = None + if isinstance(base_w1, DTensor): + base_w1 = base_w1.to_local() + base_w2 = base_w2.to_local() + base_w3 = base_w3.to_local() + w1_lora_a = w1_lora_a.to_local() + w1_lora_b = w1_lora_b.to_local() + w2_lora_a = w2_lora_a.to_local() + w2_lora_b = w2_lora_b.to_local() + w3_lora_a = w3_lora_a.to_local() + w3_lora_b = w3_lora_b.to_local() + + if getattr(self.base_layer, "ep_comm_backend", "torch") != "deepep": + from torchtitan.distributed.expert_parallel import TOKEN_GROUP_ALIGN_SIZE_M + from torchtitan.experiments.kernels.moe.indices import generate_permute_indices + + experts_per_ep_rank = base_w1.shape[0] + num_ep_ranks = num_tokens_per_expert.shape[0] // experts_per_ep_rank + + with torch.no_grad(): + permuted_indices, num_tokens_per_expert, _ = generate_permute_indices( + num_tokens_per_expert, + experts_per_ep_rank, + num_ep_ranks, + x.shape[0] + experts_per_ep_rank * TOKEN_GROUP_ALIGN_SIZE_M, + TOKEN_GROUP_ALIGN_SIZE_M, + ) + + x = torch.vstack((x, x.new_zeros((x.shape[-1])))) + input_shape = x.shape + x = x[permuted_indices, :] + # Compute offsets for grouped_mm offsets = torch.cumsum(num_tokens_per_expert, dim=0, dtype=torch.int32) @@ -350,10 +397,9 @@ def forward(self, x: torch.Tensor, num_tokens_per_expert: torch.Tensor) -> torch w2_lora_out = _run_lora_grouped_mm(lora_h, w2_lora_a, w2_lora_b, offsets) out = h2_base + scaling * w2_lora_out.bfloat16() - return out.type_as(x) + out = out.type_as(x) else: - # For-loop fallback for older GPUs - return self._forward_for_loop( + out = self._forward_for_loop( x, num_tokens_per_expert, w1_lora_a, @@ -368,6 +414,19 @@ def forward(self, x: torch.Tensor, num_tokens_per_expert: torch.Tensor) -> torch scaling, ) + # EP handling: unpermute output back to dispatched token order + if permuted_indices is not None: + # For-loop may produce fewer rows than permuted input (no padding rows), + # pad to match before unpermuting + if out.shape[0] < len(permuted_indices): + num_padding = len(permuted_indices) - out.shape[0] + out = torch.vstack((out, out.new_zeros((num_padding, out.shape[-1])))) + out_unpermuted = out.new_zeros(input_shape) + out_unpermuted[permuted_indices, :] = out + out = out_unpermuted[:-1] + + return out + def _forward_for_loop( self, x: torch.Tensor, @@ -431,3 +490,585 @@ def __repr__(self) -> str: f"alpha={self.alpha}, dropout={self.lora_dropout}, " f"use_grouped_mm={self.use_grouped_mm})" ) + + +class MultiLoRANonGatedGroupedExperts(MultiLoRAModule): + """ + NonGatedGroupedExperts + multi-LoRA with grouped GEMM. + Adapts the two projections (w1: up, w2: down) of NemotronH-style relu2 experts. + """ + + def __init__( + self, + base_layer: NonGatedGroupedExperts, + rank: int, + n_adapters: int, + alpha: float = 32.0, + dropout: float = 0.0, + use_grouped_mm: bool = True, + ): + super().__init__(base_layer) + if rank <= 0 or n_adapters <= 0: + raise ValueError("rank and n_adapters must be > 0") + + self.num_experts = base_layer.num_experts + # w1 shape: [num_experts, intermediate_dim, input_dim] + self.hidden_dim = base_layer.w1.shape[1] + self.dim = base_layer.w1.shape[2] + + if rank % 8 != 0 or self.dim % 8 != 0 or self.hidden_dim % 8 != 0: + use_grouped_mm = False + + self.rank = rank + self.n_adapters = n_adapters + self.alpha = alpha + self.lora_dropout = nn.Dropout(dropout) if dropout > 0.0 else nn.Identity() + self.use_grouped_mm = use_grouped_mm + + self._lora_num_tokens = get_lora_num_tokens() + self._scaling_factors = get_multilora_scaling() + + # w1 (up: dim -> hidden_dim) + self.w1_lora_A = nn.ParameterList( + [ + nn.Parameter( + torch.empty( + self.num_experts, + rank, + self.dim, + device=base_layer.w1.device, + dtype=base_layer.w1.dtype, + ) + ) + for _ in range(n_adapters) + ] + ) + self.w1_lora_B = nn.ParameterList( + [ + nn.Parameter( + torch.empty( + self.num_experts, + self.hidden_dim, + rank, + device=base_layer.w1.device, + dtype=base_layer.w1.dtype, + ) + ) + for _ in range(n_adapters) + ] + ) + + # w2 (down: hidden_dim -> dim) + self.w2_lora_A = nn.ParameterList( + [ + nn.Parameter( + torch.empty( + self.num_experts, + rank, + self.hidden_dim, + device=base_layer.w2.device, + dtype=base_layer.w2.dtype, + ) + ) + for _ in range(n_adapters) + ] + ) + self.w2_lora_B = nn.ParameterList( + [ + nn.Parameter( + torch.empty( + self.num_experts, + self.dim, + rank, + device=base_layer.w2.device, + dtype=base_layer.w2.dtype, + ) + ) + for _ in range(n_adapters) + ] + ) + + self.reset_parameters() + + def reset_parameters(self, index: int | None = None) -> None: + if index is None: + for i in range(self.n_adapters): + self.reset_parameters(i) + else: + nn.init.kaiming_uniform_(self.w1_lora_A[index], a=math.sqrt(5)) + nn.init.zeros_(self.w1_lora_B[index]) + nn.init.kaiming_uniform_(self.w2_lora_A[index], a=math.sqrt(5)) + nn.init.zeros_(self.w2_lora_B[index]) + + def named_parameters_for_adapter(self, idx: int) -> list[tuple[str, nn.Parameter]]: + return [ + ("w1_lora_A", self.w1_lora_A[idx]), + ("w1_lora_B", self.w1_lora_B[idx]), + ("w2_lora_A", self.w2_lora_A[idx]), + ("w2_lora_B", self.w2_lora_B[idx]), + ] + + def get_lora_param_counts(self) -> tuple[int, int]: + adapter_params = ( + self.w1_lora_A[0].numel() + + self.w1_lora_B[0].numel() + + self.w2_lora_A[0].numel() + + self.w2_lora_B[0].numel() + ) + adapted_params = self.base_layer.w1.numel() + self.base_layer.w2.numel() + return adapter_params, adapted_params + + def state_dict_for_adapter(self, idx: int) -> dict[str, torch.Tensor]: + state_dict = {} + + detached_w1_lora_a = self.w1_lora_A[idx].detach() + detached_w1_lora_b = self.w1_lora_B[idx].detach() + detached_w2_lora_a = self.w2_lora_A[idx].detach() + detached_w2_lora_b = self.w2_lora_B[idx].detach() + + if isinstance(detached_w1_lora_a, DTensor): + detached_w1_lora_a = detached_w1_lora_a.full_tensor() + detached_w1_lora_b = detached_w1_lora_b.full_tensor() + detached_w2_lora_a = detached_w2_lora_a.full_tensor() + detached_w2_lora_b = detached_w2_lora_b.full_tensor() + + for expert_id in range(self.num_experts): + state_dict[f"{expert_id}.up_proj.lora_A.weight"] = detached_w1_lora_a[expert_id].clone() + state_dict[f"{expert_id}.up_proj.lora_B.weight"] = detached_w1_lora_b[expert_id].clone() + state_dict[f"{expert_id}.down_proj.lora_A.weight"] = detached_w2_lora_a[expert_id].clone() + state_dict[f"{expert_id}.down_proj.lora_B.weight"] = detached_w2_lora_b[expert_id].clone() + + return state_dict + + def forward(self, x: torch.Tensor, num_tokens_per_expert: torch.Tensor) -> torch.Tensor: + adapter_idx = self._lora_num_tokens.argmax().item() + w1_lora_a = self.w1_lora_A[adapter_idx] + w1_lora_b = self.w1_lora_B[adapter_idx] + w2_lora_a = self.w2_lora_A[adapter_idx] + w2_lora_b = self.w2_lora_B[adapter_idx] + + scaling = self._scaling_factors[adapter_idx].item() + + base_w1 = self.base_layer.w1 + base_w2 = self.base_layer.w2 + + permuted_indices = None + if isinstance(base_w1, DTensor): + base_w1 = base_w1.to_local() + base_w2 = base_w2.to_local() + w1_lora_a = w1_lora_a.to_local() + w1_lora_b = w1_lora_b.to_local() + w2_lora_a = w2_lora_a.to_local() + w2_lora_b = w2_lora_b.to_local() + + if getattr(self.base_layer, "ep_comm_backend", "torch") != "deepep": + from torchtitan.distributed.expert_parallel import TOKEN_GROUP_ALIGN_SIZE_M + from torchtitan.experiments.kernels.moe.indices import generate_permute_indices + + experts_per_ep_rank = base_w1.shape[0] + num_ep_ranks = num_tokens_per_expert.shape[0] // experts_per_ep_rank + + with torch.no_grad(): + permuted_indices, num_tokens_per_expert, _ = generate_permute_indices( + num_tokens_per_expert, + experts_per_ep_rank, + num_ep_ranks, + x.shape[0] + experts_per_ep_rank * TOKEN_GROUP_ALIGN_SIZE_M, + TOKEN_GROUP_ALIGN_SIZE_M, + ) + + x = torch.vstack((x, x.new_zeros((x.shape[-1])))) + input_shape = x.shape + x = x[permuted_indices, :] + + offsets = torch.cumsum(num_tokens_per_expert, dim=0, dtype=torch.int32) + lora_x = self.lora_dropout(x) + + if self.use_grouped_mm: + # Up (w1) + h_base = torch._grouped_mm(x.bfloat16(), base_w1.bfloat16().transpose(-2, -1), offs=offsets) + w1_lora_out = _run_lora_grouped_mm(lora_x, w1_lora_a, w1_lora_b, offsets) + h = relu2(h_base + scaling * w1_lora_out.bfloat16()) + + # Down (w2) + lora_h = self.lora_dropout(h) + out_base = torch._grouped_mm(h, base_w2.bfloat16().transpose(-2, -1), offs=offsets) + w2_lora_out = _run_lora_grouped_mm(lora_h, w2_lora_a, w2_lora_b, offsets) + out = out_base + scaling * w2_lora_out.bfloat16() + out = out.type_as(x) + else: + out = self._forward_for_loop( + x, + num_tokens_per_expert, + w1_lora_a, + w1_lora_b, + w2_lora_a, + w2_lora_b, + base_w1, + base_w2, + scaling, + ) + + if permuted_indices is not None: + if out.shape[0] < len(permuted_indices): + num_padding = len(permuted_indices) - out.shape[0] + out = torch.vstack((out, out.new_zeros((num_padding, out.shape[-1])))) + out_unpermuted = out.new_zeros(input_shape) + out_unpermuted[permuted_indices, :] = out + out = out_unpermuted[:-1] + + return out + + def _forward_for_loop( + self, + x: torch.Tensor, + num_tokens_per_expert: torch.Tensor, + w1_lora_a: torch.Tensor, + w1_lora_b: torch.Tensor, + w2_lora_a: torch.Tensor, + w2_lora_b: torch.Tensor, + base_w1: torch.Tensor, + base_w2: torch.Tensor, + scaling: float, + ) -> torch.Tensor: + num_tokens_per_expert_list = num_tokens_per_expert.tolist() + out_splits = [] + + start = 0 + for expert_idx, num_tokens in enumerate(num_tokens_per_expert_list): + if num_tokens == 0: + continue + end = start + num_tokens + x_expert = x[start:end] + x_expert_lora = self.lora_dropout(x_expert) + + h_base = torch.matmul(x_expert, base_w1[expert_idx].transpose(-2, -1)) + w1_lora_tmp = torch.matmul(x_expert_lora, w1_lora_a[expert_idx].transpose(-2, -1)) + w1_lora_out = torch.matmul(w1_lora_tmp, w1_lora_b[expert_idx].transpose(-2, -1)) + h = relu2(h_base + scaling * w1_lora_out) + + h_lora = self.lora_dropout(h) + out_base = torch.matmul(h, base_w2[expert_idx].transpose(-2, -1)) + w2_lora_tmp = torch.matmul(h_lora, w2_lora_a[expert_idx].transpose(-2, -1)) + w2_lora_out = torch.matmul(w2_lora_tmp, w2_lora_b[expert_idx].transpose(-2, -1)) + out = out_base + scaling * w2_lora_out + + out_splits.append(out) + start = end + + return torch.cat(out_splits, dim=0) + + def __repr__(self) -> str: + return ( + f"{self.__class__.__name__}(base={self.base_layer}, rank={self.rank}, " + f"n_adapters={self.n_adapters}, num_experts={self.num_experts}, " + f"alpha={self.alpha}, dropout={self.lora_dropout}, " + f"use_grouped_mm={self.use_grouped_mm})" + ) + + +class MultiLoRAGptOssGroupedExperts(MultiLoRAModule): + """ + GptOssGroupedExperts + multi-LoRA. + + Adapts the two projections of GPT-OSS experts: the fused gate_up projection + (hidden_size -> 2*intermediate_size) and the down projection (intermediate_size -> hidden_size). + """ + + def __init__( + self, + base_layer: GptOssGroupedExperts, + rank: int, + n_adapters: int, + alpha: float = 32.0, + dropout: float = 0.0, + use_grouped_mm: bool = True, + ): + super().__init__(base_layer) + if rank <= 0 or n_adapters <= 0: + raise ValueError("rank and n_adapters must be > 0") + + self.num_experts = base_layer.num_experts + self.hidden_size = base_layer.hidden_size + self.intermediate_size = base_layer.intermediate_size + self.gate_up_out = 2 * self.intermediate_size + + if rank % 8 != 0 or self.hidden_size % 8 != 0 or self.intermediate_size % 8 != 0: + use_grouped_mm = False + + self.rank = rank + self.n_adapters = n_adapters + self.alpha = alpha + self.lora_dropout = nn.Dropout(dropout) if dropout > 0.0 else nn.Identity() + self.use_grouped_mm = use_grouped_mm + + self._lora_num_tokens = get_lora_num_tokens() + self._scaling_factors = get_multilora_scaling() + + self.gate_up_lora_A = nn.ParameterList( + [ + nn.Parameter( + torch.empty( + self.num_experts, + rank, + self.hidden_size, + device=base_layer.gate_up_proj.device, + dtype=base_layer.gate_up_proj.dtype, + ) + ) + for _ in range(n_adapters) + ] + ) + self.gate_up_lora_B = nn.ParameterList( + [ + nn.Parameter( + torch.empty( + self.num_experts, + self.gate_up_out, + rank, + device=base_layer.gate_up_proj.device, + dtype=base_layer.gate_up_proj.dtype, + ) + ) + for _ in range(n_adapters) + ] + ) + + self.down_lora_A = nn.ParameterList( + [ + nn.Parameter( + torch.empty( + self.num_experts, + rank, + self.intermediate_size, + device=base_layer.down_proj.device, + dtype=base_layer.down_proj.dtype, + ) + ) + for _ in range(n_adapters) + ] + ) + self.down_lora_B = nn.ParameterList( + [ + nn.Parameter( + torch.empty( + self.num_experts, + self.hidden_size, + rank, + device=base_layer.down_proj.device, + dtype=base_layer.down_proj.dtype, + ) + ) + for _ in range(n_adapters) + ] + ) + + self.reset_parameters() + + def reset_parameters(self, index: int | None = None) -> None: + if index is None: + for i in range(self.n_adapters): + self.reset_parameters(i) + else: + nn.init.kaiming_uniform_(self.gate_up_lora_A[index], a=math.sqrt(5)) + nn.init.zeros_(self.gate_up_lora_B[index]) + nn.init.kaiming_uniform_(self.down_lora_A[index], a=math.sqrt(5)) + nn.init.zeros_(self.down_lora_B[index]) + + def named_parameters_for_adapter(self, idx: int) -> list[tuple[str, nn.Parameter]]: + return [ + ("gate_up_lora_A", self.gate_up_lora_A[idx]), + ("gate_up_lora_B", self.gate_up_lora_B[idx]), + ("down_lora_A", self.down_lora_A[idx]), + ("down_lora_B", self.down_lora_B[idx]), + ] + + def get_lora_param_counts(self) -> tuple[int, int]: + adapter_params = ( + self.gate_up_lora_A[0].numel() + + self.gate_up_lora_B[0].numel() + + self.down_lora_A[0].numel() + + self.down_lora_B[0].numel() + ) + adapted_params = self.base_layer.gate_up_proj.numel() + self.base_layer.down_proj.numel() + return adapter_params, adapted_params + + def state_dict_for_adapter(self, idx: int) -> dict[str, torch.Tensor]: + """vLLM-compatible 3D MoE adapter format. + + For 3D MoE models (gpt-oss in vLLM has `is_3d_moe_weight = True`), vLLM expects: + - `experts.base_layer.lora_{A,B}.weight` for the gate_up projection + - `experts.lora_{A,B}.weight` for the down projection + with experts stacked into the rank dim. See + vllm/lora/model_manager.py::_stack_moe_lora_weights, which reshapes + lora_A: (num_experts*rank, in) -> (num_experts, rank, in) + lora_B: (out, rank*num_experts) -> (out, rank, num_experts) -> (num_experts, out, rank) + """ + detached_gu_a = self.gate_up_lora_A[idx].detach() + detached_gu_b = self.gate_up_lora_B[idx].detach() + detached_d_a = self.down_lora_A[idx].detach() + detached_d_b = self.down_lora_B[idx].detach() + + if isinstance(detached_gu_a, DTensor): + detached_gu_a = detached_gu_a.full_tensor() + detached_gu_b = detached_gu_b.full_tensor() + detached_d_a = detached_d_a.full_tensor() + detached_d_b = detached_d_b.full_tensor() + + # lora_A: (num_experts, rank, in) -> (num_experts*rank, in) + gu_a_flat = detached_gu_a.reshape(self.num_experts * self.rank, self.hidden_size).clone() + d_a_flat = detached_d_a.reshape(self.num_experts * self.rank, self.intermediate_size).clone() + # lora_B: (num_experts, out, rank) -> (out, rank, num_experts) -> (out, rank*num_experts) + # vLLM's reshape treats the last dim of lora_B as (rank, num_experts) with experts fast-varying. + gu_b_flat = detached_gu_b.permute(1, 2, 0).contiguous().reshape(self.gate_up_out, self.rank * self.num_experts) + d_b_flat = detached_d_b.permute(1, 2, 0).contiguous().reshape(self.hidden_size, self.rank * self.num_experts) + + return { + "base_layer.lora_A.weight": gu_a_flat, + "base_layer.lora_B.weight": gu_b_flat, + "lora_A.weight": d_a_flat, + "lora_B.weight": d_b_flat, + } + + def forward(self, x: torch.Tensor, num_tokens_per_expert: torch.Tensor) -> torch.Tensor: + adapter_idx = self._lora_num_tokens.argmax().item() + gu_a = self.gate_up_lora_A[adapter_idx] + gu_b = self.gate_up_lora_B[adapter_idx] + d_a = self.down_lora_A[adapter_idx] + d_b = self.down_lora_B[adapter_idx] + + scaling = self._scaling_factors[adapter_idx].item() + + base_gu = self.base_layer.gate_up_proj + base_gu_bias = self.base_layer.gate_up_proj_bias + base_d = self.base_layer.down_proj + base_d_bias = self.base_layer.down_proj_bias + + permuted_indices = None + if isinstance(base_gu, DTensor): + base_gu = base_gu.to_local() + base_gu_bias = base_gu_bias.to_local() + base_d = base_d.to_local() + base_d_bias = base_d_bias.to_local() + gu_a = gu_a.to_local() + gu_b = gu_b.to_local() + d_a = d_a.to_local() + d_b = d_b.to_local() + + if getattr(self.base_layer, "ep_comm_backend", "torch") != "deepep": + from torchtitan.distributed.expert_parallel import TOKEN_GROUP_ALIGN_SIZE_M + from torchtitan.experiments.kernels.moe.indices import generate_permute_indices + + experts_per_ep_rank = base_gu.shape[0] + num_ep_ranks = num_tokens_per_expert.shape[0] // experts_per_ep_rank + + with torch.no_grad(): + permuted_indices, num_tokens_per_expert, _ = generate_permute_indices( + num_tokens_per_expert, + experts_per_ep_rank, + num_ep_ranks, + x.shape[0] + experts_per_ep_rank * TOKEN_GROUP_ALIGN_SIZE_M, + TOKEN_GROUP_ALIGN_SIZE_M, + ) + + x = torch.vstack((x, x.new_zeros((x.shape[-1])))) + input_shape = x.shape + x = x[permuted_indices, :] + + offsets = torch.cumsum(num_tokens_per_expert, dim=0, dtype=torch.int32) + lora_x = self.lora_dropout(x) + + if self.use_grouped_mm: + gate_up_base = torch._grouped_mm(x.bfloat16(), base_gu.bfloat16(), offs=offsets) + gate_up_base = ( + gate_up_base + + _broadcast_expert_bias(base_gu_bias, num_tokens_per_expert, gate_up_base.shape[0]).bfloat16() + ) + gate_up_lora = _run_lora_grouped_mm(lora_x, gu_a, gu_b, offsets) + gate_up = gate_up_base + scaling * gate_up_lora.bfloat16() + + h = _gpt_oss_apply_gate(gate_up) + lora_h = self.lora_dropout(h) + + out_base = torch._grouped_mm(h, base_d.bfloat16(), offs=offsets) + out_base = ( + out_base + _broadcast_expert_bias(base_d_bias, num_tokens_per_expert, out_base.shape[0]).bfloat16() + ) + out_lora = _run_lora_grouped_mm(lora_h, d_a, d_b, offsets) + out = out_base + scaling * out_lora.bfloat16() + out = out.type_as(x) + else: + out = self._forward_for_loop( + x, + num_tokens_per_expert, + gu_a, + gu_b, + d_a, + d_b, + base_gu, + base_gu_bias, + base_d, + base_d_bias, + scaling, + ) + + if permuted_indices is not None: + if out.shape[0] < len(permuted_indices): + num_padding = len(permuted_indices) - out.shape[0] + out = torch.vstack((out, out.new_zeros((num_padding, out.shape[-1])))) + out_unpermuted = out.new_zeros(input_shape) + out_unpermuted[permuted_indices, :] = out + out = out_unpermuted[:-1] + + return out + + def _forward_for_loop( + self, + x: torch.Tensor, + num_tokens_per_expert: torch.Tensor, + gu_a: torch.Tensor, + gu_b: torch.Tensor, + d_a: torch.Tensor, + d_b: torch.Tensor, + base_gu: torch.Tensor, + base_gu_bias: torch.Tensor, + base_d: torch.Tensor, + base_d_bias: torch.Tensor, + scaling: float, + ) -> torch.Tensor: + n = num_tokens_per_expert.tolist() + out_splits = [] + + start = 0 + for e, num_tokens in enumerate(n): + if num_tokens == 0: + continue + end = start + num_tokens + x_e = x[start:end] + x_e_lora = self.lora_dropout(x_e) + + gate_up_base = x_e @ base_gu[e] + base_gu_bias[e] + gate_up_lora_tmp = x_e_lora @ gu_a[e].transpose(-2, -1) + gate_up_lora = gate_up_lora_tmp @ gu_b[e].transpose(-2, -1) + gate_up = gate_up_base + scaling * gate_up_lora + + h = _gpt_oss_apply_gate(gate_up) + h_lora = self.lora_dropout(h) + + out_base = h @ base_d[e] + base_d_bias[e] + out_lora_tmp = h_lora @ d_a[e].transpose(-2, -1) + out_lora = out_lora_tmp @ d_b[e].transpose(-2, -1) + out = out_base + scaling * out_lora + + out_splits.append(out) + start = end + + return torch.cat(out_splits, dim=0) + + def __repr__(self) -> str: + return ( + f"{self.__class__.__name__}(base={self.base_layer}, rank={self.rank}, " + f"n_adapters={self.n_adapters}, num_experts={self.num_experts}, " + f"alpha={self.alpha}, dropout={self.lora_dropout}, " + f"use_grouped_mm={self.use_grouped_mm})" + ) diff --git a/src/prime_rl/trainer/models/layers/mlp.py b/src/prime_rl/trainer/models/layers/mlp.py index 3fe9f9e8606..1c94e9fd511 100644 --- a/src/prime_rl/trainer/models/layers/mlp.py +++ b/src/prime_rl/trainer/models/layers/mlp.py @@ -1,5 +1,6 @@ from dataclasses import dataclass +import torch from torch import nn from transformers.activations import ACT2FN @@ -23,6 +24,6 @@ def __init__(self, config: MLPConfig): self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False) self.gate_act_fn = ACT2FN[config.gate_act] - def forward(self, x): + def forward(self, x, routed_experts: torch.Tensor | None = None): down_proj = self.down_proj(self.gate_act_fn(self.gate_proj(x)) * self.up_proj(x)) return down_proj diff --git a/src/prime_rl/trainer/models/layers/moe.py b/src/prime_rl/trainer/models/layers/moe.py index 73ae509fdbd..c59a720539c 100644 --- a/src/prime_rl/trainer/models/layers/moe.py +++ b/src/prime_rl/trainer/models/layers/moe.py @@ -12,6 +12,8 @@ from torch import nn from torchtitan.distributed.expert_parallel import expert_parallel +from prime_rl.configs.trainer import EPCommBackend + @dataclass class MoEArgs: @@ -28,6 +30,7 @@ class MoEArgs: top_k: int = 1 use_grouped_mm: bool = True # grouped mm or for-loop for the experts computation load_balance_coeff: float | None = 1e-3 + fp8: bool = False # use FP8 grouped GEMM via DeepGEMM (requires SM90) # can be used as dense FFN layer or shared experts in MoE layers @@ -85,8 +88,7 @@ def init_weights(self, init_std: float): # TODO: keeping this for-loop implementation for comparison # and readability, may remove later -@expert_parallel -def _run_experts_for_loop( +def _run_experts_for_loop_impl( w1: torch.Tensor, w2: torch.Tensor, w3: torch.Tensor, @@ -121,6 +123,17 @@ def _run_experts_for_loop( return out +@expert_parallel +def _run_experts_for_loop( + w1: torch.Tensor, + w2: torch.Tensor, + w3: torch.Tensor, + x: torch.Tensor, + num_tokens_per_expert: torch.Tensor, +) -> torch.Tensor: + return _run_experts_for_loop_impl(w1, w2, w3, x, num_tokens_per_expert) + + @expert_parallel def _run_experts_grouped_mm( w1: torch.Tensor, @@ -128,14 +141,43 @@ def _run_experts_grouped_mm( w3: torch.Tensor, x: torch.Tensor, num_tokens_per_expert: torch.Tensor, +) -> torch.Tensor: + return _run_experts_grouped_mm_impl(w1, w2, w3, x, num_tokens_per_expert) + + +@expert_parallel +def _run_experts_fp8_grouped_mm( + w1: torch.Tensor, + w2: torch.Tensor, + w3: torch.Tensor, + x: torch.Tensor, + num_tokens_per_expert: torch.Tensor, +) -> torch.Tensor: + return _run_experts_grouped_mm_impl(w1, w2, w3, x, num_tokens_per_expert, fp8=True) + + +def _run_experts_grouped_mm_impl( + w1: torch.Tensor, + w2: torch.Tensor, + w3: torch.Tensor, + x: torch.Tensor, + num_tokens_per_expert: torch.Tensor, + fp8: bool = False, ) -> torch.Tensor: offsets = torch.cumsum(num_tokens_per_expert, dim=0, dtype=torch.int32) # grouped mm between a 2D tensor and a 3D tensor assert x.dim() == 2 - h = F.silu(torch._grouped_mm(x.bfloat16(), w1.bfloat16().transpose(-2, -1), offs=offsets)) - h = h * torch._grouped_mm(x.bfloat16(), w3.bfloat16().transpose(-2, -1), offs=offsets) - out = torch._grouped_mm(h, w2.bfloat16().transpose(-2, -1), offs=offsets).type_as(x) + if fp8: + from prime_rl.trainer.models.layers.fp8_grouped_gemm import grouped_fp8_gemm + + h = F.silu(grouped_fp8_gemm(x.bfloat16(), w1.bfloat16().transpose(-2, -1), offsets)) + h = h * grouped_fp8_gemm(x.bfloat16(), w3.bfloat16().transpose(-2, -1), offsets) + out = grouped_fp8_gemm(h, w2.bfloat16().transpose(-2, -1), offsets).type_as(x) + else: + h = F.silu(torch._grouped_mm(x.bfloat16(), w1.bfloat16().transpose(-2, -1), offs=offsets)) + h = h * torch._grouped_mm(x.bfloat16(), w3.bfloat16().transpose(-2, -1), offs=offsets) + out = torch._grouped_mm(h, w2.bfloat16().transpose(-2, -1), offs=offsets).type_as(x) return out @@ -147,6 +189,7 @@ def __init__( hidden_dim: int, num_experts: int, use_grouped_mm: bool, + fp8: bool = False, ): super().__init__() self.num_experts = num_experts @@ -154,13 +197,31 @@ def __init__( self.w2 = nn.Parameter(torch.empty(num_experts, dim, hidden_dim)) self.w3 = nn.Parameter(torch.empty(num_experts, hidden_dim, dim)) self.use_grouped_mm = use_grouped_mm + self.fp8 = fp8 + self.ep_comm_backend: EPCommBackend = "torch" + + def set_ep_comm_backend(self, backend: EPCommBackend) -> None: + self.ep_comm_backend = backend + + def _forward_deepep(self, x: torch.Tensor, num_tokens_per_expert: torch.Tensor) -> torch.Tensor: + w1 = self.w1.to_local() + w2 = self.w2.to_local() + w3 = self.w3.to_local() + if self.use_grouped_mm: + return _run_experts_grouped_mm_impl(w1, w2, w3, x, num_tokens_per_expert, fp8=self.fp8) + return _run_experts_for_loop_impl(w1, w2, w3, x, num_tokens_per_expert) def forward( self, x: torch.Tensor, num_tokens_per_expert: torch.Tensor, ) -> torch.Tensor: + if self.ep_comm_backend == "deepep": + return self._forward_deepep(x, num_tokens_per_expert) + if self.use_grouped_mm: + if self.fp8: + return _run_experts_fp8_grouped_mm(self.w1, self.w2, self.w3, x, num_tokens_per_expert) return _run_experts_grouped_mm(self.w1, self.w2, self.w3, x, num_tokens_per_expert) else: return _run_experts_for_loop(self.w1, self.w2, self.w3, x, num_tokens_per_expert) @@ -171,6 +232,161 @@ def init_weights(self, init_std: float): nn.init.trunc_normal_(self.w3, mean=0.0, std=init_std) +# GPT-OSS activation constants. Both clamping limit and the sigmoid alpha live here +# rather than as instance attrs so the function is JIT/compile-friendly. +GPT_OSS_LIMIT = 7.0 +GPT_OSS_ALPHA = 1.702 + + +def _gpt_oss_apply_gate(gate_up: torch.Tensor) -> torch.Tensor: + """GPT-OSS expert activation: clamped sigmoid-glu over interleaved gate/up channels.""" + gate, up = gate_up[..., ::2], gate_up[..., 1::2] + gate = gate.clamp(min=None, max=GPT_OSS_LIMIT) + up = up.clamp(min=-GPT_OSS_LIMIT, max=GPT_OSS_LIMIT) + glu = gate * torch.sigmoid(gate * GPT_OSS_ALPHA) + return (up + 1) * glu + + +def _broadcast_expert_bias(bias: torch.Tensor, num_tokens_per_expert: torch.Tensor, target_rows: int) -> torch.Tensor: + """Repeat per-expert bias to per-token, padding to target_rows if EP added padding rows.""" + # repeat_interleave on CUDA requires int counts; histc/router output is float. + bias_per_token = torch.repeat_interleave(bias, num_tokens_per_expert.to(torch.int64), dim=0) + if bias_per_token.shape[0] < target_rows: + pad_rows = target_rows - bias_per_token.shape[0] + bias_per_token = F.pad(bias_per_token, (0, 0, 0, pad_rows)) + return bias_per_token + + +def _run_gpt_oss_experts_for_loop_impl( + gate_up_proj: torch.Tensor, + gate_up_proj_bias: torch.Tensor, + down_proj: torch.Tensor, + down_proj_bias: torch.Tensor, + x: torch.Tensor, + num_tokens_per_expert: torch.Tensor, +) -> torch.Tensor: + n = num_tokens_per_expert.tolist() + num_padding = x.shape[0] - sum(n) + x_split = torch.split(x[: sum(n)], split_size_or_sections=n, dim=0) + out_splits = [] + for e, x_e in enumerate(x_split): + gate_up = x_e @ gate_up_proj[e] + gate_up_proj_bias[e] + h = _gpt_oss_apply_gate(gate_up) + out = h @ down_proj[e] + down_proj_bias[e] + out_splits.append(out) + out = torch.cat(out_splits, dim=0) + return torch.vstack((out, out.new_zeros((num_padding, out.shape[-1])))) + + +@expert_parallel +def _run_gpt_oss_experts_for_loop( + gate_up_proj: torch.Tensor, + gate_up_proj_bias: torch.Tensor, + down_proj: torch.Tensor, + down_proj_bias: torch.Tensor, + x: torch.Tensor, + num_tokens_per_expert: torch.Tensor, +) -> torch.Tensor: + return _run_gpt_oss_experts_for_loop_impl( + gate_up_proj, gate_up_proj_bias, down_proj, down_proj_bias, x, num_tokens_per_expert + ) + + +def _run_gpt_oss_experts_grouped_mm_impl( + gate_up_proj: torch.Tensor, + gate_up_proj_bias: torch.Tensor, + down_proj: torch.Tensor, + down_proj_bias: torch.Tensor, + x: torch.Tensor, + num_tokens_per_expert: torch.Tensor, +) -> torch.Tensor: + offsets = torch.cumsum(num_tokens_per_expert, dim=0, dtype=torch.int32) + assert x.dim() == 2 + + gate_up = torch._grouped_mm(x.bfloat16(), gate_up_proj.bfloat16(), offs=offsets) + gate_up = gate_up + _broadcast_expert_bias(gate_up_proj_bias, num_tokens_per_expert, gate_up.shape[0]).bfloat16() + h = _gpt_oss_apply_gate(gate_up) + out = torch._grouped_mm(h, down_proj.bfloat16(), offs=offsets) + out = out + _broadcast_expert_bias(down_proj_bias, num_tokens_per_expert, out.shape[0]).bfloat16() + return out.type_as(x) + + +@expert_parallel +def _run_gpt_oss_experts_grouped_mm( + gate_up_proj: torch.Tensor, + gate_up_proj_bias: torch.Tensor, + down_proj: torch.Tensor, + down_proj_bias: torch.Tensor, + x: torch.Tensor, + num_tokens_per_expert: torch.Tensor, +) -> torch.Tensor: + return _run_gpt_oss_experts_grouped_mm_impl( + gate_up_proj, gate_up_proj_bias, down_proj, down_proj_bias, x, num_tokens_per_expert + ) + + +class GptOssGroupedExperts(nn.Module): + """GPT-OSS-style grouped experts. + + Mirrors HF's `GptOssExperts` parameter naming (gate_up_proj/down_proj plus per-expert + biases, fused interleaved gate/up channels) so the unsloth BF16 checkpoint loads with + no key conversion. Forward signature matches `GroupedExperts` (`x`, `num_tokens_per_expert`) + so the surrounding MoE plumbing and LoRA wrapper follow the same convention. + """ + + def __init__( + self, + hidden_size: int, + intermediate_size: int, + num_experts: int, + use_grouped_mm: bool, + ): + super().__init__() + self.num_experts = num_experts + self.hidden_size = hidden_size + self.intermediate_size = intermediate_size + self.gate_up_proj = nn.Parameter(torch.empty(num_experts, hidden_size, 2 * intermediate_size)) + self.gate_up_proj_bias = nn.Parameter(torch.empty(num_experts, 2 * intermediate_size)) + self.down_proj = nn.Parameter(torch.empty(num_experts, intermediate_size, hidden_size)) + self.down_proj_bias = nn.Parameter(torch.empty(num_experts, hidden_size)) + self.use_grouped_mm = use_grouped_mm + self.ep_comm_backend: EPCommBackend = "torch" + + def set_ep_comm_backend(self, backend: EPCommBackend) -> None: + self.ep_comm_backend = backend + + def _forward_deepep(self, x: torch.Tensor, num_tokens_per_expert: torch.Tensor) -> torch.Tensor: + gate_up_proj = self.gate_up_proj.to_local() + gate_up_proj_bias = self.gate_up_proj_bias.to_local() + down_proj = self.down_proj.to_local() + down_proj_bias = self.down_proj_bias.to_local() + if self.use_grouped_mm: + return _run_gpt_oss_experts_grouped_mm_impl( + gate_up_proj, gate_up_proj_bias, down_proj, down_proj_bias, x, num_tokens_per_expert + ) + return _run_gpt_oss_experts_for_loop_impl( + gate_up_proj, gate_up_proj_bias, down_proj, down_proj_bias, x, num_tokens_per_expert + ) + + def forward(self, x: torch.Tensor, num_tokens_per_expert: torch.Tensor) -> torch.Tensor: + if self.ep_comm_backend == "deepep": + return self._forward_deepep(x, num_tokens_per_expert) + + if self.use_grouped_mm: + return _run_gpt_oss_experts_grouped_mm( + self.gate_up_proj, self.gate_up_proj_bias, self.down_proj, self.down_proj_bias, x, num_tokens_per_expert + ) + return _run_gpt_oss_experts_for_loop( + self.gate_up_proj, self.gate_up_proj_bias, self.down_proj, self.down_proj_bias, x, num_tokens_per_expert + ) + + def init_weights(self, init_std: float): + nn.init.trunc_normal_(self.gate_up_proj, mean=0.0, std=0.02) + nn.init.zeros_(self.gate_up_proj_bias) + nn.init.trunc_normal_(self.down_proj, mean=0.0, std=init_std) + nn.init.zeros_(self.down_proj_bias) + + class TokenChoiceTopKRouter(nn.Module): """This class implements token-choice routing. In token-choice top-K routing, each token is routed to top K experts based on the router scores. @@ -200,15 +416,17 @@ def __init__( self.score_func = score_func self.route_norm = route_norm self.route_scale = route_scale + self.force_balanced = False def forward( - self, x: torch.Tensor, expert_bias: torch.Tensor | None = None + self, x: torch.Tensor, expert_bias: torch.Tensor | None = None, routed_experts: torch.Tensor | None = None ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: """ Args: x (torch.Tensor): Input tensor with shape ``(bs*slen, dim)``. expert_bias (torch.Tensor | None, optional): Optional bias tensor for experts with shape ``(num_experts,)``. Used for load balancing. Defaults to None. + routed_experts (torch.Tensor | None, optional): Optional tensor with shape ``(bs * slen, top_k)``. Returns: tuple[torch.Tensor, torch.Tensor, torch.Tensor]: @@ -220,6 +438,9 @@ def forward( Number of tokens assigned to each expert with shape ``(num_experts,)``. """ # scores shape (bs*slen, num_experts) + assert routed_experts is None or routed_experts.shape[-1] == self.top_k, ( + f"routed_experts shape: {routed_experts.shape}, top_k: {self.top_k}" + ) scores = self.gate(x) # By default, sigmoid or softmax is performed in float32 to avoid loss explosion @@ -233,7 +454,16 @@ def forward( # top scores shape (bs*slen, top_k) # NOTE: The expert_bias is only used for routing. The gating value # top_scores is still derived from the original scores. - if expert_bias is not None: + + if routed_experts is not None: + top_scores = scores.gather(dim=1, index=routed_experts) + selected_experts_indices = routed_experts + elif self.force_balanced: + num_tokens = scores.shape[0] + arange = torch.arange(num_tokens * self.top_k, device=scores.device) + selected_experts_indices = (arange % self.num_experts).view(num_tokens, self.top_k) + top_scores = scores.gather(dim=1, index=selected_experts_indices) + elif expert_bias is not None: _, selected_experts_indices = torch.topk(scores + expert_bias, k=self.top_k, dim=1) top_scores = scores.gather(dim=1, index=selected_experts_indices) else: @@ -246,7 +476,7 @@ def forward( # group tokens together by expert indices from 0 to num_experts and pass that to experts forward num_tokens_per_expert = torch.histc( - selected_experts_indices.view(-1), + selected_experts_indices.reshape(-1), bins=self.num_experts, min=0, max=self.num_experts, @@ -296,8 +526,9 @@ def forward( - num_tokens_per_expert: Number of tokens assigned to each expert """ # group tokens together by expert indices from 0 to num_experts and pass that to experts forward + selected_experts_indices = selected_experts_indices.reshape(-1) num_tokens_per_expert = torch.histc( - selected_experts_indices.view(-1), + selected_experts_indices, bins=self.num_experts, min=0, max=self.num_experts, @@ -305,7 +536,7 @@ def forward( # Reorder the token indices to match the order of the experts # token_indices_experts_sorted shape (bs*slen*top_k,) - token_indices_experts_sorted = torch.argsort(selected_experts_indices.view(-1), stable=True) + token_indices_experts_sorted = torch.argsort(selected_experts_indices, stable=True) top_scores_experts_sorted = top_scores.view(-1)[token_indices_experts_sorted] token_indices_experts_sorted = token_indices_experts_sorted // self.top_k @@ -322,17 +553,15 @@ def __init__(self, moe_args: MoEArgs, dim: int, hidden_dim: int): super().__init__() num_experts = moe_args.num_experts - if torch.cuda.is_available(): - cc_major, _ = torch.cuda.get_device_capability() - # As of PyTorch 2.9, grouped_mm is supported on Hopper and Blackwell. - if cc_major < 9: - moe_args.use_grouped_mm = False self.experts = GroupedExperts( dim=dim, hidden_dim=hidden_dim, num_experts=num_experts, use_grouped_mm=moe_args.use_grouped_mm, + fp8=moe_args.fp8, ) + self.ep_comm_backend: EPCommBackend = "torch" + self.experts.set_ep_comm_backend(self.ep_comm_backend) self.router = TokenChoiceTopKRouter( dim=dim, num_experts=num_experts, @@ -349,6 +578,7 @@ def __init__(self, moe_args: MoEArgs, dim: int, hidden_dim: int): else None ) self.score_before_experts = moe_args.score_before_experts + self.deepep_token_chunk_size: int | None = None # define fields for auxiliary-loss-free load balancing (https://arxiv.org/abs/2408.15664) # NOTE: tokens_per_expert is accumulated in the model forward pass. @@ -371,10 +601,104 @@ def __init__(self, moe_args: MoEArgs, dim: int, hidden_dim: int): persistent=False, ) - def forward(self, x: torch.Tensor) -> torch.Tensor: + def set_ep_comm_backend(self, backend: EPCommBackend) -> None: + self.ep_comm_backend = backend + self.experts.set_ep_comm_backend(backend) + + def set_deepep_token_chunk_size(self, chunk_size: int | None) -> None: + self.deepep_token_chunk_size = chunk_size + + def _run_local_routed_experts( + self, + x: torch.Tensor, + num_tokens_per_expert: torch.Tensor, + ) -> torch.Tensor: + return self.experts(x, num_tokens_per_expert) + + def _run_routed_experts( + self, + x: torch.Tensor, + token_indices_experts_sorted: torch.Tensor, + num_tokens_per_expert: torch.Tensor, + top_scores_experts_sorted: torch.Tensor, + ) -> torch.Tensor: + dim = x.shape[-1] + routed_indices = token_indices_experts_sorted.reshape(-1, 1).expand(-1, dim) + routed_input = torch.gather(x, dim=0, index=routed_indices) + + if self.score_before_experts: + routed_input = (routed_input.to(torch.float32) * top_scores_experts_sorted.reshape(-1, 1)).to(x.dtype) + + routed_output = self.experts(routed_input, num_tokens_per_expert) + + if not self.score_before_experts: + routed_output = (routed_output.to(torch.float32) * top_scores_experts_sorted.reshape(-1, 1)).to(x.dtype) + + return routed_output + + def _run_deepep_routed_experts( + self, + x: torch.Tensor, + selected_experts_indices: torch.Tensor, + top_scores: torch.Tensor, + ) -> torch.Tensor: + from prime_rl.trainer.distributed.deepep import ( + combine_tokens, + dispatch_tokens_async, + finalize_dispatch_tokens, + sync_combine, + ) + from prime_rl.trainer.distributed.expert_parallel import get_ep_group + + if x.shape[0] == 0: + shared_output = self.shared_expert(x) if self.shared_expert is not None else None + return x.new_zeros(x.shape) if shared_output is None else shared_output + + group = get_ep_group(self.experts) + chunk_size = min(self.deepep_token_chunk_size or x.shape[0], x.shape[0]) + + def dispatch_chunk(start: int, end: int): + return dispatch_tokens_async( + x[start:end], + selected_experts_indices[start:end], + top_scores[start:end], + num_experts=self.experts.num_experts, + group=group, + score_before_experts=self.score_before_experts, + ) + + def run_pending_chunk(pending_state): + hidden_states, num_tokens_per_expert, dispatch_state = finalize_dispatch_tokens(pending_state) + routed_output = self._run_local_routed_experts(hidden_states, num_tokens_per_expert) + # Keep combine outside the checkpointed routed-expert region so + # selective AC only recomputes local expert matmuls. + return combine_tokens(routed_output, dispatch_state) + + pending_state = dispatch_chunk(0, chunk_size) + routed_outputs: list[torch.Tensor] = [] + + for chunk_start in range(chunk_size, x.shape[0], chunk_size): + chunk_end = min(chunk_start + chunk_size, x.shape[0]) + next_pending_state = dispatch_chunk(chunk_start, chunk_end) + routed_outputs.append(run_pending_chunk(pending_state)) + pending_state = next_pending_state + + routed_outputs.append(run_pending_chunk(pending_state)) + + shared_output = self.shared_expert(x) if self.shared_expert is not None else None + sync_combine() + routed_output = routed_outputs[0] if len(routed_outputs) == 1 else torch.cat(routed_outputs, dim=0) + return routed_output if shared_output is None else shared_output + routed_output + + def forward( + self, + x: torch.Tensor, + routed_experts: torch.Tensor | None = None, + ) -> torch.Tensor: """ Args: x (torch.Tensor): Input tensor with shape ``(bs, slen, dim)``. + routed_experts (torch.Tensor | None, optional): Optional tensor with shape ``(bs, slen, top_k)``. Returns: out (torch.Tensor): Output tensor with shape ``(bs, slen, dim)``. @@ -382,22 +706,32 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: bs, slen, dim = x.shape x = x.view(-1, dim) + if routed_experts is not None: + _, _, top_k = routed_experts.shape + routed_experts = routed_experts.reshape( + -1, top_k + ) # we have to reshape here because the original is non-contiguous + # top_scores and selected_experts_indices shape (bs*slen*top_k,) # num_tokens_per_expert shape (num_experts,) ( top_scores, selected_experts_indices, num_tokens_per_expert, - ) = self.router(x, self.expert_bias) + ) = self.router(x, self.expert_bias, routed_experts=routed_experts) # tokens_per_expert will be used to update the expert bias for load balancing. # and also to count the expert usage - # TODO: Activation Checkpointing has the side effect of double counting tokens_per_expert -- - # first in the forward pass, and then in the backward pass. However, this has no - # effect on the expert bias update thanks to the torch.sign() operator. + # Full block checkpointing can double count tokens_per_expert because it reruns the router + # in backward. The selective MoE path avoids that by checkpointing only the + # routed expert compute below. with torch.no_grad(): self.tokens_per_expert.add_(num_tokens_per_expert) + if self.ep_comm_backend == "deepep": + routed_output = self._run_deepep_routed_experts(x, selected_experts_indices, top_scores) + return routed_output.reshape(bs, slen, dim) + # top_scores and token_indices_experts_sorted shape (bs*slen*top_k,) # num_tokens_per_expert shape (num_experts,) # NOTE: the reason we need to compute num_tokens_per_expert again is: @@ -412,28 +746,19 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: num_tokens_per_expert, ) = self.reorderer(top_scores, selected_experts_indices) - # shape (bs*slen*top_k, dim) - token_indices_experts_sorted = token_indices_experts_sorted.reshape(-1, 1).expand(-1, dim) - - # shape (bs*slen*top_k, dim) - routed_input = torch.gather(x, dim=0, index=token_indices_experts_sorted) - - if self.score_before_experts: - routed_input = (routed_input.to(torch.float32) * top_scores_experts_sorted.reshape(-1, 1)).to(x.dtype) - - # shape (bs*slen*top_k, dim) - routed_output = self.experts(routed_input, num_tokens_per_expert) - - if not self.score_before_experts: - routed_output = (routed_output.to(torch.float32) * top_scores_experts_sorted.reshape(-1, 1)).to(x.dtype) - - # shared expert + routed_output = self._run_routed_experts( + x, + token_indices_experts_sorted, + num_tokens_per_expert, + top_scores_experts_sorted, + ) if self.shared_expert is not None: out = self.shared_expert(x) else: out = torch.zeros_like(x) - out = out.scatter_add(dim=0, index=token_indices_experts_sorted, src=routed_output) + routed_indices = token_indices_experts_sorted.reshape(-1, 1).expand(-1, dim) + out = out.scatter_add(dim=0, index=routed_indices, src=routed_output) out = out.reshape(bs, slen, dim) return out @@ -451,3 +776,425 @@ def init_weights( self.tokens_per_expert = torch.zeros(self.experts.num_experts, dtype=torch.float32) if self.load_balance_coeff is not None: self.expert_bias = torch.zeros(self.experts.num_experts, dtype=torch.float32) + + +def relu2(x: torch.Tensor) -> torch.Tensor: + return F.relu(x).square() + + +def _run_nongated_experts_for_loop_impl( + w1: torch.Tensor, + w2: torch.Tensor, + _w3: torch.Tensor, + x: torch.Tensor, + num_tokens_per_expert: torch.Tensor, +) -> torch.Tensor: + num_tokens_per_expert = num_tokens_per_expert.tolist() + num_padding = x.shape[0] - sum(num_tokens_per_expert) + + x = torch.split( + x[: sum(num_tokens_per_expert)], + split_size_or_sections=num_tokens_per_expert, + dim=0, + ) + out_experts_splits = [] + for expert_idx, x_expert in enumerate(x): + h = relu2(torch.matmul(x_expert, w1[expert_idx].transpose(-2, -1))) + h = torch.matmul(h, w2[expert_idx].transpose(-2, -1)) + out_experts_splits.append(h) + out = torch.cat(out_experts_splits, dim=0) + out = torch.vstack((out, out.new_zeros((num_padding, out.shape[-1])))) + return out + + +@expert_parallel +def _run_nongated_experts_for_loop( + w1: torch.Tensor, + w2: torch.Tensor, + _w3: torch.Tensor, + x: torch.Tensor, + num_tokens_per_expert: torch.Tensor, +) -> torch.Tensor: + return _run_nongated_experts_for_loop_impl(w1, w2, _w3, x, num_tokens_per_expert) + + +def _run_nongated_experts_grouped_mm_impl( + w1: torch.Tensor, + w2: torch.Tensor, + _w3: torch.Tensor, + x: torch.Tensor, + num_tokens_per_expert: torch.Tensor, + fp8: bool = False, +) -> torch.Tensor: + offsets = torch.cumsum(num_tokens_per_expert, dim=0, dtype=torch.int32) + assert x.dim() == 2 + + if fp8: + from prime_rl.trainer.models.layers.fp8_grouped_gemm import grouped_fp8_gemm + + h = relu2(grouped_fp8_gemm(x.bfloat16(), w1.bfloat16().transpose(-2, -1), offsets)) + out = grouped_fp8_gemm(h, w2.bfloat16().transpose(-2, -1), offsets).type_as(x) + else: + h = relu2(torch._grouped_mm(x.bfloat16(), w1.bfloat16().transpose(-2, -1), offs=offsets)) + out = torch._grouped_mm(h, w2.bfloat16().transpose(-2, -1), offs=offsets).type_as(x) + return out + + +@expert_parallel +def _run_nongated_experts_grouped_mm( + w1: torch.Tensor, + w2: torch.Tensor, + _w3: torch.Tensor, + x: torch.Tensor, + num_tokens_per_expert: torch.Tensor, +) -> torch.Tensor: + return _run_nongated_experts_grouped_mm_impl(w1, w2, _w3, x, num_tokens_per_expert) + + +@expert_parallel +def _run_nongated_experts_fp8_grouped_mm( + w1: torch.Tensor, + w2: torch.Tensor, + _w3: torch.Tensor, + x: torch.Tensor, + num_tokens_per_expert: torch.Tensor, +) -> torch.Tensor: + return _run_nongated_experts_grouped_mm_impl(w1, w2, _w3, x, num_tokens_per_expert, fp8=True) + + +class NonGatedGroupedExperts(nn.Module): + def __init__( + self, + input_dim: int, + intermediate_dim: int, + num_experts: int, + use_grouped_mm: bool, + fp8: bool = False, + ): + super().__init__() + self.num_experts = num_experts + self.w1 = nn.Parameter(torch.empty(num_experts, intermediate_dim, input_dim)) + self.w2 = nn.Parameter(torch.empty(num_experts, input_dim, intermediate_dim)) + # Dummy w3 for @expert_parallel decorator compatibility (expects w1, w2, w3 signature) + self.w3 = nn.Parameter(torch.empty(0)) + self.use_grouped_mm = use_grouped_mm + self.fp8 = fp8 + self.ep_comm_backend: EPCommBackend = "torch" + + def set_ep_comm_backend(self, backend: EPCommBackend) -> None: + self.ep_comm_backend = backend + + def _forward_deepep(self, x: torch.Tensor, num_tokens_per_expert: torch.Tensor) -> torch.Tensor: + w1 = self.w1.to_local() + w2 = self.w2.to_local() + w3 = self.w3.to_local() + if self.use_grouped_mm: + return _run_nongated_experts_grouped_mm_impl(w1, w2, w3, x, num_tokens_per_expert, fp8=self.fp8) + return _run_nongated_experts_for_loop_impl(w1, w2, w3, x, num_tokens_per_expert) + + def forward( + self, + x: torch.Tensor, + num_tokens_per_expert: torch.Tensor, + ) -> torch.Tensor: + if self.ep_comm_backend == "deepep": + return self._forward_deepep(x, num_tokens_per_expert) + if self.use_grouped_mm: + if self.fp8: + return _run_nongated_experts_fp8_grouped_mm(self.w1, self.w2, self.w3, x, num_tokens_per_expert) + return _run_nongated_experts_grouped_mm(self.w1, self.w2, self.w3, x, num_tokens_per_expert) + else: + return _run_nongated_experts_for_loop(self.w1, self.w2, self.w3, x, num_tokens_per_expert) + + def init_weights(self, init_std: float): + nn.init.trunc_normal_(self.w1, mean=0.0, std=0.02) + nn.init.trunc_normal_(self.w2, mean=0.0, std=init_std) + + +class NemotronHRouter(nn.Module): + """Sigmoid router with group-based expert selection and e_score_correction_bias. + + Follows the DeepseekV3 routing pattern: sigmoid scoring, group-based top-k selection, + and bias correction for load balancing. + """ + + def __init__( + self, + dim: int, + num_experts: int, + top_k: int, + n_group: int, + topk_group: int, + norm_topk_prob: bool, + ): + super().__init__() + self.gate = nn.Parameter(torch.empty(num_experts, dim)) + self.register_buffer("e_score_correction_bias", torch.zeros(num_experts)) + self.num_experts = num_experts + self.top_k = top_k + self.n_group = n_group + self.topk_group = topk_group + self.norm_topk_prob = norm_topk_prob + + def forward( + self, x: torch.Tensor, expert_bias: torch.Tensor | None = None + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + scores = F.linear(x.float(), self.gate.float()).sigmoid() + scores_for_choice = scores + self.e_score_correction_bias + + if expert_bias is not None: + scores_for_choice = scores_for_choice + expert_bias + + # Group-based routing + if self.n_group > 1: + group_scores = ( + scores_for_choice.view(-1, self.n_group, self.num_experts // self.n_group) + .topk(2, dim=-1)[0] + .sum(dim=-1) + ) + group_idx = torch.topk(group_scores, k=self.topk_group, dim=-1, sorted=False)[1] + group_mask = torch.zeros_like(group_scores) + group_mask.scatter_(1, group_idx, 1) + score_mask = ( + group_mask.unsqueeze(-1) + .expand(-1, self.n_group, self.num_experts // self.n_group) + .reshape(-1, self.num_experts) + ) + scores_for_choice = scores_for_choice.masked_fill(~score_mask.bool(), 0.0) + + selected_experts_indices = torch.topk(scores_for_choice, k=self.top_k, dim=-1, sorted=False)[1] + top_scores = scores.gather(1, selected_experts_indices) + + if self.norm_topk_prob: + denominator = top_scores.sum(dim=-1, keepdim=True) + 1e-20 + top_scores = top_scores / denominator + + num_tokens_per_expert = torch.histc( + selected_experts_indices.reshape(-1).float(), + bins=self.num_experts, + min=0, + max=self.num_experts, + ) + + return top_scores, selected_experts_indices, num_tokens_per_expert + + def init_weights(self, init_std: float): + nn.init.trunc_normal_(self.gate, mean=0.0, std=init_std) + + +class BCNonGatedFeedForward(nn.Module): + """Non-gated feed-forward network used as the shared expert in NemotronH. + + Uses relu2 activation: down_proj(relu2(up_proj(x))). + """ + + def __init__(self, dim: int, hidden_dim: int): + super().__init__() + self.up_proj = nn.Linear(dim, hidden_dim, bias=False) + self.down_proj = nn.Linear(hidden_dim, dim, bias=False) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.down_proj(relu2(self.up_proj(x))) + + +class LatentMoE(nn.Module): + """NemotronH-style Mixture of Experts with latent projections. + + The input is projected to a latent space before expert computation, + and the output is projected back. Experts use relu2 activation without gating. + """ + + def __init__( + self, + dim: int, + latent_dim: int | None, + moe_intermediate_size: int, + shared_expert_intermediate_size: int, + num_experts: int, + top_k: int, + n_group: int, + topk_group: int, + norm_topk_prob: bool, + routed_scaling_factor: float, + use_grouped_mm: bool, + load_balance_coeff: float | None, + fp8: bool = False, + ): + super().__init__() + effective_latent_dim = latent_dim if latent_dim is not None else dim + + self.router = NemotronHRouter( + dim=dim, + num_experts=num_experts, + top_k=top_k, + n_group=n_group, + topk_group=topk_group, + norm_topk_prob=norm_topk_prob, + ) + self.experts = NonGatedGroupedExperts( + input_dim=effective_latent_dim, + intermediate_dim=moe_intermediate_size, + num_experts=num_experts, + use_grouped_mm=use_grouped_mm, + fp8=fp8, + ) + self.ep_comm_backend: EPCommBackend = "torch" + self.experts.set_ep_comm_backend(self.ep_comm_backend) + self.reorderer = TokenReorderer(num_experts=num_experts, top_k=top_k) + self.shared_expert = BCNonGatedFeedForward(dim=dim, hidden_dim=shared_expert_intermediate_size) + self.deepep_token_chunk_size: int | None = None + + if latent_dim is not None: + self.fc1_latent_proj = nn.Linear(dim, latent_dim, bias=False) + self.fc2_latent_proj = nn.Linear(latent_dim, dim, bias=False) + else: + self.fc1_latent_proj = nn.Identity() + self.fc2_latent_proj = nn.Identity() + + self.routed_scaling_factor = routed_scaling_factor + self.load_balance_coeff = load_balance_coeff + if self.load_balance_coeff is not None: + assert self.load_balance_coeff > 0.0 + self.register_buffer( + "expert_bias", + torch.zeros(num_experts, dtype=torch.float32), + persistent=True, + ) + else: + self.expert_bias = None + self.register_buffer( + "tokens_per_expert", + torch.zeros(num_experts, dtype=torch.float32), + persistent=False, + ) + + def set_ep_comm_backend(self, backend: EPCommBackend) -> None: + self.ep_comm_backend = backend + self.experts.set_ep_comm_backend(backend) + + def set_deepep_token_chunk_size(self, chunk_size: int | None) -> None: + self.deepep_token_chunk_size = chunk_size + + def _run_local_routed_experts( + self, + x: torch.Tensor, + num_tokens_per_expert: torch.Tensor, + ) -> torch.Tensor: + return self.experts(x, num_tokens_per_expert) + + def _run_routed_experts( + self, + x: torch.Tensor, + token_indices_experts_sorted: torch.Tensor, + num_tokens_per_expert: torch.Tensor, + top_scores_experts_sorted: torch.Tensor, + ) -> torch.Tensor: + dim = x.shape[-1] + token_indices_expanded = token_indices_experts_sorted.reshape(-1, 1).expand(-1, dim) + routed_input = torch.gather(x, dim=0, index=token_indices_expanded) + + routed_input = self.fc1_latent_proj(routed_input) + routed_output = self.experts(routed_input, num_tokens_per_expert) + + routed_output = (routed_output.float() * top_scores_experts_sorted.reshape(-1, 1)).to(routed_output.dtype) + routed_output = routed_output * self.routed_scaling_factor + + routed_output = self.fc2_latent_proj(routed_output) + return routed_output + + def _run_deepep_routed_experts( + self, + x: torch.Tensor, + selected_experts_indices: torch.Tensor, + top_scores: torch.Tensor, + ) -> torch.Tensor: + from prime_rl.trainer.distributed.deepep import ( + combine_tokens, + dispatch_tokens_async, + finalize_dispatch_tokens, + sync_combine, + ) + from prime_rl.trainer.distributed.expert_parallel import get_ep_group + + if x.shape[0] == 0: + return self.shared_expert(x) + + group = get_ep_group(self.experts) + # Project before dispatch so DeepEP communicates the smaller latent activations. + latent_x = self.fc1_latent_proj(x) + chunk_size = min(self.deepep_token_chunk_size or latent_x.shape[0], latent_x.shape[0]) + + def dispatch_chunk(start: int, end: int): + return dispatch_tokens_async( + latent_x[start:end], + selected_experts_indices[start:end], + top_scores[start:end], + num_experts=self.experts.num_experts, + group=group, + score_before_experts=False, + ) + + def run_pending_chunk(pending_state): + hidden_states, num_tokens_per_expert, dispatch_state = finalize_dispatch_tokens(pending_state) + routed_output = self._run_local_routed_experts(hidden_states, num_tokens_per_expert) + return combine_tokens(routed_output, dispatch_state) + + pending_state = dispatch_chunk(0, chunk_size) + routed_outputs: list[torch.Tensor] = [] + + for chunk_start in range(chunk_size, latent_x.shape[0], chunk_size): + chunk_end = min(chunk_start + chunk_size, latent_x.shape[0]) + next_pending_state = dispatch_chunk(chunk_start, chunk_end) + routed_outputs.append(run_pending_chunk(pending_state)) + pending_state = next_pending_state + + routed_outputs.append(run_pending_chunk(pending_state)) + + shared_output = self.shared_expert(x) + sync_combine() + routed_output = routed_outputs[0] if len(routed_outputs) == 1 else torch.cat(routed_outputs, dim=0) + routed_output = routed_output * self.routed_scaling_factor + routed_output = self.fc2_latent_proj(routed_output) + return shared_output + routed_output + + def forward(self, x: torch.Tensor, routed_experts: torch.Tensor | None = None) -> torch.Tensor: + bs, slen, dim = x.shape + x_flat = x.view(-1, dim) + + top_scores, selected_experts_indices, num_tokens_per_expert = self.router(x_flat, self.expert_bias) + + with torch.no_grad(): + self.tokens_per_expert.add_(num_tokens_per_expert) + + if self.ep_comm_backend == "deepep": + routed_output = self._run_deepep_routed_experts(x_flat, selected_experts_indices, top_scores) + return routed_output.reshape(bs, slen, dim) + + ( + top_scores_experts_sorted, + token_indices_experts_sorted, + num_tokens_per_expert, + ) = self.reorderer(top_scores, selected_experts_indices) + + routed_output = self._run_routed_experts( + x_flat, + token_indices_experts_sorted, + num_tokens_per_expert, + top_scores_experts_sorted, + ) + + out = self.shared_expert(x_flat) + + token_indices_full = token_indices_experts_sorted.reshape(-1, 1).expand(-1, dim) + out = out.scatter_add(dim=0, index=token_indices_full, src=routed_output) + out = out.reshape(bs, slen, dim) + return out + + def init_weights(self, init_std: float, buffer_device: torch.device): + self.experts.init_weights(init_std) + self.router.init_weights(init_std) + + with torch.device(buffer_device): + self.tokens_per_expert = torch.zeros(self.experts.num_experts, dtype=torch.float32) + if self.load_balance_coeff is not None: + self.expert_bias = torch.zeros(self.experts.num_experts, dtype=torch.float32) diff --git a/src/prime_rl/trainer/models/layers/norms.py b/src/prime_rl/trainer/models/layers/norms.py new file mode 100644 index 00000000000..6bd63c6bdf7 --- /dev/null +++ b/src/prime_rl/trainer/models/layers/norms.py @@ -0,0 +1,79 @@ +from dataclasses import dataclass +from functools import lru_cache + +import torch +import torch.nn.functional as F +from torch import nn +from transformers.integrations import use_kernel_forward_from_hub + + +@lru_cache(maxsize=1) +def _get_quack_rmsnorm(): + """Lazy-load quack rmsnorm. Returns None if unavailable or GPU is pre-Hopper.""" + if not torch.cuda.is_available() or torch.cuda.get_device_capability()[0] < 9: + return None + try: + from quack import rmsnorm + + return rmsnorm + except ImportError: + return None + + +class _ContiguousGrad(torch.autograd.Function): + @staticmethod + def forward(ctx, x): + return x + + @staticmethod + def backward(ctx, grad): + return grad.contiguous() + + +def _contiguous_grad(x: torch.Tensor) -> torch.Tensor: + """Identity in forward, makes gradient contiguous in backward. + + Quack's RMSNorm backward kernel requires contiguous gradients (stride[-1]==1) + but upstream ops like attention permute can produce non-contiguous ones. + """ + return _ContiguousGrad.apply(x) if x.requires_grad else x + + +@dataclass +class RMSNormConfig: + hidden_size: int + eps: float = 1e-6 + + +@use_kernel_forward_from_hub("RMSNorm") +class RMSNorm(nn.Module): + def __init__(self, config: RMSNormConfig) -> None: + super().__init__() + self.weight = nn.Parameter(torch.ones(config.hidden_size)) + self.variance_epsilon = config.eps + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + quack_fn = _get_quack_rmsnorm() if hidden_states.is_cuda else None + if quack_fn is not None: + out = quack_fn(hidden_states, self.weight, eps=self.variance_epsilon) + return _contiguous_grad(out) + input_dtype = hidden_states.dtype + hidden_states = hidden_states.to(torch.float32) + variance = hidden_states.pow(2).mean(-1, keepdim=True) + hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon) + return self.weight * hidden_states.to(input_dtype) + + def extra_repr(self): + return f"{tuple(self.weight.shape)}, eps={self.variance_epsilon}" + + +class LayerNorm(nn.Module): + def __init__(self, dim: int, eps: float = 1e-6): + super().__init__() + self.dim = dim + self.eps = eps + self.weight = nn.Parameter(torch.ones(dim, dtype=torch.float32)) + self.bias = nn.Parameter(torch.zeros(dim, dtype=torch.float32)) + + def forward(self, x: torch.Tensor): + return F.layer_norm(x.float(), (self.dim,), self.weight.float(), self.bias.float(), self.eps).type_as(x) diff --git a/src/prime_rl/trainer/models/layers/ring_attn.py b/src/prime_rl/trainer/models/layers/ring_attn.py new file mode 100644 index 00000000000..cdb1774200b --- /dev/null +++ b/src/prime_rl/trainer/models/layers/ring_attn.py @@ -0,0 +1,606 @@ +from __future__ import annotations + +# ruff: noqa: I001 — `prime_rl._compat` must run before `ring_flash_attn` imports below. +import prime_rl._compat # noqa: F401 + +import torch +import torch.distributed as dist +from ring_flash_attn.utils import AllGatherComm, get_default_args + + +def _set_fa3_signature_params(params: dict, causal: bool, window_size: tuple[int, int]) -> None: + if "is_causal" in params: + params["is_causal"] = causal + else: + params["causal"] = causal + + if "window_size" in params: + params["window_size"] = window_size + else: + params["window_size_left"] = window_size[0] + params["window_size_right"] = window_size[1] + + +def _fa3_varlen_forward( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + cu_seqlens_q: torch.Tensor, + cu_seqlens_k: torch.Tensor, + max_seqlen_q: int, + max_seqlen_k: int, + softmax_scale: float, + causal: bool, + window_size: tuple[int, int] = (-1, -1), +) -> tuple[torch.Tensor, torch.Tensor]: + from flash_attn_interface import _flash_attn_forward + + params = get_default_args(_flash_attn_forward).copy() + params.update( + { + "q": q, + "k": k, + "v": v, + "cu_seqlens_q": cu_seqlens_q, + "cu_seqlens_k": cu_seqlens_k, + "max_seqlen_q": max_seqlen_q, + "max_seqlen_k": max_seqlen_k, + "softmax_scale": softmax_scale, + } + ) + _set_fa3_signature_params(params, causal, window_size) + out, lse, _, _ = _flash_attn_forward(**params) + return out, lse + + +def _fa3_varlen_backward( + dout: torch.Tensor, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + out: torch.Tensor, + softmax_lse: torch.Tensor, + cu_seqlens_q: torch.Tensor, + cu_seqlens_k: torch.Tensor, + max_seqlen_q: int, + max_seqlen_k: int, + dq: torch.Tensor, + dk: torch.Tensor, + dv: torch.Tensor, + softmax_scale: float, + causal: bool, + window_size: tuple[int, int] = (-1, -1), +) -> None: + from flash_attn_interface import _flash_attn_backward + + params = get_default_args(_flash_attn_backward).copy() + params.update( + { + "dout": dout, + "q": q, + "k": k, + "v": v, + "out": out, + "softmax_lse": softmax_lse, + "cu_seqlens_q": cu_seqlens_q, + "cu_seqlens_k": cu_seqlens_k, + "max_seqlen_q": max_seqlen_q, + "max_seqlen_k": max_seqlen_k, + "dq": dq, + "dk": dk, + "dv": dv, + "softmax_scale": softmax_scale, + } + ) + _set_fa3_signature_params(params, causal, window_size) + _flash_attn_backward(**params) + + +class _RingFA3Varlen(torch.autograd.Function): + """Ring attention using FA3 kernels with all-gather communication. + + Mirrors the llama3_flash_attn_varlen pattern from ring-flash-attn but + calls flash-attention-3 low-level forward/backward instead of FA2. + """ + + @staticmethod + def forward( + ctx, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + cu_seqlens_q: torch.Tensor, + cu_seqlens_k: torch.Tensor, + max_seqlen_q: int, + max_seqlen_k: int, + local_k_slice_start: int, + local_k_slice_stop: int, + heads_k_stride: int, + causal: bool, + group_name: str, + window_size_left: int = -1, + window_size_right: int = -1, + ) -> torch.Tensor: + group = dist.group.WORLD + for pg in dist.distributed_c10d._world.pg_map: + if pg.group_name == group_name: + group = pg + break + + local_k_slice = slice(local_k_slice_start, local_k_slice_stop) + window_size = (window_size_left, window_size_right) + softmax_scale = q.shape[-1] ** (-0.5) + out_list = [] + lse_list = [] + + nheads = q.shape[1] + total_k, nheads_k, head_dim = k.shape + world_size = dist.get_world_size(group) + + kv_buffer = torch.empty((2, total_k * world_size, heads_k_stride, head_dim), dtype=k.dtype, device=k.device) + kv_buffer_copy = torch.empty_like(kv_buffer) + comm = AllGatherComm(group) + + comm.all_gather(kv_buffer_copy[0], k[:, :heads_k_stride].contiguous()) + comm.all_gather(kv_buffer_copy[1], v[:, :heads_k_stride].contiguous()) + + for i in range(0, nheads_k, heads_k_stride): + comm.wait() + kv_buffer, kv_buffer_copy = kv_buffer_copy, kv_buffer + + if i < nheads_k - heads_k_stride: + left = i + heads_k_stride + right = left + heads_k_stride + comm.all_gather(kv_buffer_copy[0], k[:, left:right].contiguous()) + comm.all_gather(kv_buffer_copy[1], v[:, left:right].contiguous()) + + q_i = q[:, i * nheads // nheads_k : (i + heads_k_stride) * nheads // nheads_k] + k_i = kv_buffer[0][local_k_slice] + v_i = kv_buffer[1][local_k_slice] + out_i, lse_i = _fa3_varlen_forward( + q=q_i, + k=k_i, + v=v_i, + cu_seqlens_q=cu_seqlens_q, + cu_seqlens_k=cu_seqlens_k, + max_seqlen_q=max_seqlen_q, + max_seqlen_k=max_seqlen_k, + softmax_scale=softmax_scale, + causal=causal, + window_size=window_size, + ) + out_list.append(out_i) + lse_list.append(lse_i) + + out = torch.cat(out_list, dim=1) + lse = torch.cat(lse_list, dim=-2) + + ctx.save_for_backward(q, k, v, out, lse, cu_seqlens_q, cu_seqlens_k) + ctx.softmax_scale = softmax_scale + ctx.max_seqlen_q = max_seqlen_q + ctx.max_seqlen_k = max_seqlen_k + ctx.local_k_slice = local_k_slice + ctx.heads_k_stride = heads_k_stride + ctx.causal = causal + ctx.group_name = group_name + ctx.window_size = window_size + return out + + @staticmethod + def backward(ctx, dout: torch.Tensor): + q, k, v, out, softmax_lse, cu_seqlens_q, cu_seqlens_k = ctx.saved_tensors + heads_k_stride = ctx.heads_k_stride + local_k_slice = ctx.local_k_slice + causal = ctx.causal + + group = dist.group.WORLD + for pg in dist.distributed_c10d._world.pg_map: + if pg.group_name == ctx.group_name: + group = pg + break + + nheads = q.shape[1] + total_k, nheads_k, head_dim = k.shape + world_size = dist.get_world_size(group) + + kv_buffer = torch.empty((2, total_k * world_size, heads_k_stride, head_dim), dtype=k.dtype, device=k.device) + kv_buffer_copy = torch.empty_like(kv_buffer) + dkv_buffer = torch.empty((2, total_k * world_size, heads_k_stride, head_dim), dtype=k.dtype, device=k.device) + + kv_contiguous_buffer = None + if heads_k_stride != nheads_k: + kv_contiguous_buffer = torch.empty((2, total_k, heads_k_stride, head_dim), dtype=k.dtype, device=k.device) + + dq = torch.empty_like(q) + dk = torch.empty_like(k) + dv = torch.empty_like(v) + + comm = AllGatherComm(group) + comm.all_gather(kv_buffer_copy[0], k[:, :heads_k_stride].contiguous()) + comm.all_gather(kv_buffer_copy[1], v[:, :heads_k_stride].contiguous()) + + for i in range(0, nheads_k, heads_k_stride): + dkv_buffer.zero_() + q_slice = slice(i * nheads // nheads_k, (i + heads_k_stride) * nheads // nheads_k) + q_i = q[:, q_slice] + dout_i = dout[:, q_slice] + out_i = out[:, q_slice] + dq_i = dq[:, q_slice] + lse_i = softmax_lse[q_slice].contiguous() + + comm.wait() + kv_buffer, kv_buffer_copy = kv_buffer_copy, kv_buffer + if i < nheads_k - heads_k_stride: + left = i + heads_k_stride + right = left + heads_k_stride + comm.all_gather(kv_buffer_copy[0], k[:, left:right].contiguous()) + comm.all_gather(kv_buffer_copy[1], v[:, left:right].contiguous()) + + k_i = kv_buffer[0][local_k_slice] + v_i = kv_buffer[1][local_k_slice] + dk_i = dkv_buffer[0][local_k_slice] + dv_i = dkv_buffer[1][local_k_slice] + + _fa3_varlen_backward( + dout=dout_i, + q=q_i, + k=k_i, + v=v_i, + out=out_i, + softmax_lse=lse_i, + cu_seqlens_q=cu_seqlens_q, + cu_seqlens_k=cu_seqlens_k, + max_seqlen_q=ctx.max_seqlen_q, + max_seqlen_k=ctx.max_seqlen_k, + dq=dq_i, + dk=dk_i, + dv=dv_i, + softmax_scale=ctx.softmax_scale, + causal=causal, + window_size=ctx.window_size, + ) + + if heads_k_stride != nheads_k: + dk_i = kv_contiguous_buffer[0] + dv_i = kv_contiguous_buffer[1] + else: + dk_i = dk + dv_i = dv + + dist.reduce_scatter_tensor(dk_i, dkv_buffer[0], group=group) + dist.reduce_scatter_tensor(dv_i, dkv_buffer[1], group=group) + if heads_k_stride != nheads_k: + dk[:, i : i + heads_k_stride] = dk_i + dv[:, i : i + heads_k_stride] = dv_i + + # Grads for: q, k, v, cu_seqlens_q, cu_seqlens_k, max_seqlen_q, max_seqlen_k, + # local_k_slice_start, local_k_slice_stop, heads_k_stride, causal, group_name, + # window_size_left, window_size_right + return dq, dk, dv, None, None, None, None, None, None, None, None, None, None, None + + +def ring_fa3_varlen_func( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + cu_seqlens_q: torch.Tensor, + cu_seqlens_k: torch.Tensor, + max_seqlen_q: int, + max_seqlen_k: int, + local_k_slice: slice, + causal: bool, + heads_k_stride: int, + group: dist.ProcessGroup, + window_size: tuple[int, int] = (-1, -1), +) -> torch.Tensor: + return _RingFA3Varlen.apply( + q, + k, + v, + cu_seqlens_q, + cu_seqlens_k, + max_seqlen_q, + max_seqlen_k, + local_k_slice.start, + local_k_slice.stop, + heads_k_stride, + causal, + group.group_name, + window_size[0], + window_size[1], + ) + + +# --------------------------------------------------------------------------- +# FA4 (flash_attn.cute) ring attention +# --------------------------------------------------------------------------- + + +def _fa4_varlen_forward( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + cu_seqlens_q: torch.Tensor, + cu_seqlens_k: torch.Tensor, + max_seqlen_q: int, + max_seqlen_k: int, + softmax_scale: float, + causal: bool, + window_size: tuple[int, int] = (-1, -1), +) -> tuple[torch.Tensor, torch.Tensor]: + from flash_attn.cute.interface import _flash_attn_fwd + + wl = window_size[0] if window_size[0] != -1 else None + wr = window_size[1] if window_size[1] != -1 else None + out, lse = _flash_attn_fwd( + q, + k, + v, + cu_seqlens_q=cu_seqlens_q, + cu_seqlens_k=cu_seqlens_k, + max_seqlen_q=max_seqlen_q, + max_seqlen_k=max_seqlen_k, + softmax_scale=softmax_scale, + causal=causal, + window_size_left=wl, + window_size_right=wr, + return_lse=True, + ) + return out, lse + + +def _fa4_varlen_backward( + dout: torch.Tensor, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + out: torch.Tensor, + softmax_lse: torch.Tensor, + cu_seqlens_q: torch.Tensor, + cu_seqlens_k: torch.Tensor, + max_seqlen_q: int, + max_seqlen_k: int, + dq: torch.Tensor, + dk: torch.Tensor, + dv: torch.Tensor, + softmax_scale: float, + causal: bool, + window_size: tuple[int, int] = (-1, -1), +) -> None: + from flash_attn.cute.interface import _flash_attn_bwd + + wl = window_size[0] if window_size[0] != -1 else None + wr = window_size[1] if window_size[1] != -1 else None + _flash_attn_bwd( + q, + k, + v, + out, + dout, + softmax_lse, + softmax_scale=softmax_scale, + causal=causal, + window_size_left=wl, + window_size_right=wr, + cu_seqlens_q=cu_seqlens_q, + cu_seqlens_k=cu_seqlens_k, + max_seqlen_q=max_seqlen_q, + max_seqlen_k=max_seqlen_k, + dq=dq, + dk=dk, + dv=dv, + ) + + +class _RingFA4Varlen(torch.autograd.Function): + """Ring attention using FA4 (flash_attn.cute) kernels with all-gather communication. + + Mirrors _RingFA3Varlen but calls the FA4 low-level forward/backward. + """ + + @staticmethod + def forward( + ctx, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + cu_seqlens_q: torch.Tensor, + cu_seqlens_k: torch.Tensor, + max_seqlen_q: int, + max_seqlen_k: int, + local_k_slice_start: int, + local_k_slice_stop: int, + heads_k_stride: int, + causal: bool, + group_name: str, + window_size_left: int = -1, + window_size_right: int = -1, + ) -> torch.Tensor: + group = dist.group.WORLD + for pg in dist.distributed_c10d._world.pg_map: + if pg.group_name == group_name: + group = pg + break + + local_k_slice = slice(local_k_slice_start, local_k_slice_stop) + window_size = (window_size_left, window_size_right) + softmax_scale = q.shape[-1] ** (-0.5) + out_list = [] + lse_list = [] + + nheads = q.shape[1] + total_k, nheads_k, head_dim = k.shape + world_size = dist.get_world_size(group) + + kv_buffer = torch.empty((2, total_k * world_size, heads_k_stride, head_dim), dtype=k.dtype, device=k.device) + kv_buffer_copy = torch.empty_like(kv_buffer) + comm = AllGatherComm(group) + + comm.all_gather(kv_buffer_copy[0], k[:, :heads_k_stride].contiguous()) + comm.all_gather(kv_buffer_copy[1], v[:, :heads_k_stride].contiguous()) + + for i in range(0, nheads_k, heads_k_stride): + comm.wait() + kv_buffer, kv_buffer_copy = kv_buffer_copy, kv_buffer + + if i < nheads_k - heads_k_stride: + left = i + heads_k_stride + right = left + heads_k_stride + comm.all_gather(kv_buffer_copy[0], k[:, left:right].contiguous()) + comm.all_gather(kv_buffer_copy[1], v[:, left:right].contiguous()) + + q_i = q[:, i * nheads // nheads_k : (i + heads_k_stride) * nheads // nheads_k] + k_i = kv_buffer[0][local_k_slice] + v_i = kv_buffer[1][local_k_slice] + out_i, lse_i = _fa4_varlen_forward( + q=q_i, + k=k_i, + v=v_i, + cu_seqlens_q=cu_seqlens_q, + cu_seqlens_k=cu_seqlens_k, + max_seqlen_q=max_seqlen_q, + max_seqlen_k=max_seqlen_k, + softmax_scale=softmax_scale, + causal=causal, + window_size=window_size, + ) + out_list.append(out_i) + lse_list.append(lse_i) + + out = torch.cat(out_list, dim=1) + lse = torch.cat(lse_list, dim=-2) + + ctx.save_for_backward(q, k, v, out, lse, cu_seqlens_q, cu_seqlens_k) + ctx.softmax_scale = softmax_scale + ctx.max_seqlen_q = max_seqlen_q + ctx.max_seqlen_k = max_seqlen_k + ctx.local_k_slice = local_k_slice + ctx.heads_k_stride = heads_k_stride + ctx.causal = causal + ctx.group_name = group_name + ctx.window_size = window_size + return out + + @staticmethod + def backward(ctx, dout: torch.Tensor): + q, k, v, out, softmax_lse, cu_seqlens_q, cu_seqlens_k = ctx.saved_tensors + heads_k_stride = ctx.heads_k_stride + local_k_slice = ctx.local_k_slice + causal = ctx.causal + + group = dist.group.WORLD + for pg in dist.distributed_c10d._world.pg_map: + if pg.group_name == ctx.group_name: + group = pg + break + + nheads = q.shape[1] + total_k, nheads_k, head_dim = k.shape + world_size = dist.get_world_size(group) + + kv_buffer = torch.empty((2, total_k * world_size, heads_k_stride, head_dim), dtype=k.dtype, device=k.device) + kv_buffer_copy = torch.empty_like(kv_buffer) + dkv_buffer = torch.empty((2, total_k * world_size, heads_k_stride, head_dim), dtype=k.dtype, device=k.device) + + kv_contiguous_buffer = None + if heads_k_stride != nheads_k: + kv_contiguous_buffer = torch.empty((2, total_k, heads_k_stride, head_dim), dtype=k.dtype, device=k.device) + + dq = torch.empty_like(q) + dk = torch.empty_like(k) + dv = torch.empty_like(v) + + comm = AllGatherComm(group) + comm.all_gather(kv_buffer_copy[0], k[:, :heads_k_stride].contiguous()) + comm.all_gather(kv_buffer_copy[1], v[:, :heads_k_stride].contiguous()) + + for i in range(0, nheads_k, heads_k_stride): + dkv_buffer.zero_() + q_slice = slice(i * nheads // nheads_k, (i + heads_k_stride) * nheads // nheads_k) + q_i = q[:, q_slice] + dout_i = dout[:, q_slice] + out_i = out[:, q_slice] + dq_i = dq[:, q_slice] + lse_i = softmax_lse[q_slice].contiguous() + + comm.wait() + kv_buffer, kv_buffer_copy = kv_buffer_copy, kv_buffer + if i < nheads_k - heads_k_stride: + left = i + heads_k_stride + right = left + heads_k_stride + comm.all_gather(kv_buffer_copy[0], k[:, left:right].contiguous()) + comm.all_gather(kv_buffer_copy[1], v[:, left:right].contiguous()) + + k_i = kv_buffer[0][local_k_slice] + v_i = kv_buffer[1][local_k_slice] + dk_i = dkv_buffer[0][local_k_slice] + dv_i = dkv_buffer[1][local_k_slice] + + _fa4_varlen_backward( + dout=dout_i, + q=q_i, + k=k_i, + v=v_i, + out=out_i, + softmax_lse=lse_i, + cu_seqlens_q=cu_seqlens_q, + cu_seqlens_k=cu_seqlens_k, + max_seqlen_q=ctx.max_seqlen_q, + max_seqlen_k=ctx.max_seqlen_k, + dq=dq_i, + dk=dk_i, + dv=dv_i, + softmax_scale=ctx.softmax_scale, + causal=causal, + window_size=ctx.window_size, + ) + + if heads_k_stride != nheads_k: + dk_i = kv_contiguous_buffer[0] + dv_i = kv_contiguous_buffer[1] + else: + dk_i = dk + dv_i = dv + + dist.reduce_scatter_tensor(dk_i, dkv_buffer[0], group=group) + dist.reduce_scatter_tensor(dv_i, dkv_buffer[1], group=group) + if heads_k_stride != nheads_k: + dk[:, i : i + heads_k_stride] = dk_i + dv[:, i : i + heads_k_stride] = dv_i + + # Grads for: q, k, v, cu_seqlens_q, cu_seqlens_k, max_seqlen_q, max_seqlen_k, + # local_k_slice_start, local_k_slice_stop, heads_k_stride, causal, group_name, + # window_size_left, window_size_right + return dq, dk, dv, None, None, None, None, None, None, None, None, None, None, None + + +def ring_fa4_varlen_func( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + cu_seqlens_q: torch.Tensor, + cu_seqlens_k: torch.Tensor, + max_seqlen_q: int, + max_seqlen_k: int, + local_k_slice: slice, + causal: bool, + heads_k_stride: int, + group: dist.ProcessGroup, + window_size: tuple[int, int] = (-1, -1), +) -> torch.Tensor: + return _RingFA4Varlen.apply( + q, + k, + v, + cu_seqlens_q, + cu_seqlens_k, + max_seqlen_q, + max_seqlen_k, + local_k_slice.start, + local_k_slice.stop, + heads_k_stride, + causal, + group.group_name, + window_size[0], + window_size[1], + ) diff --git a/src/prime_rl/trainer/models/layers/rms_norm.py b/src/prime_rl/trainer/models/layers/rms_norm.py index 76d190309f9..a150c9302ac 100644 --- a/src/prime_rl/trainer/models/layers/rms_norm.py +++ b/src/prime_rl/trainer/models/layers/rms_norm.py @@ -1,32 +1,3 @@ -from dataclasses import dataclass +from prime_rl.trainer.models.layers.norms import RMSNorm, RMSNormConfig -import torch -from torch import nn -from transformers.integrations import use_kernel_forward_from_hub - - -@dataclass -class RMSNormConfig: - hidden_size: int - eps: float = 1e-6 - - -@use_kernel_forward_from_hub("RMSNorm") -class RMSNorm(nn.Module): - def __init__(self, config: RMSNormConfig) -> None: - """ - Glm4MoeRMSNorm is equivalent to T5LayerNorm - """ - super().__init__() - self.weight = nn.Parameter(torch.ones(config.hidden_size)) - self.variance_epsilon = config.eps - - def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: - input_dtype = hidden_states.dtype - hidden_states = hidden_states.to(torch.float32) - variance = hidden_states.pow(2).mean(-1, keepdim=True) - hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon) - return self.weight * hidden_states.to(input_dtype) - - def extra_repr(self): - return f"{tuple(self.weight.shape)}, eps={self.variance_epsilon}" +__all__ = ["RMSNorm", "RMSNormConfig"] diff --git a/src/prime_rl/trainer/models/layers/rotary_emb.py b/src/prime_rl/trainer/models/layers/rotary_emb.py index d6ed8d18e19..f7b77de9515 100644 --- a/src/prime_rl/trainer/models/layers/rotary_emb.py +++ b/src/prime_rl/trainer/models/layers/rotary_emb.py @@ -1,10 +1,34 @@ from dataclasses import dataclass +from typing import Optional import torch from torch import nn from transformers.modeling_rope_utils import ROPE_INIT_FUNCTIONS, dynamic_rope_update +def _compute_default_rope_parameters( + config: Optional["PreTrainedConfig"] = None, + device: Optional["torch.device"] = None, + seq_len: int | None = None, + layer_type: str | None = None, +) -> tuple["torch.Tensor", float]: + """Computes standard RoPE inverse frequencies without any scaling (factor=1.0).""" + config.standardize_rope_params() + rope_parameters_dict = config.rope_parameters[layer_type] if layer_type is not None else config.rope_parameters + + base = rope_parameters_dict["rope_theta"] + partial_rotary_factor = rope_parameters_dict.get("partial_rotary_factor", 1.0) + head_dim = getattr(config, "head_dim", None) or config.hidden_size // config.num_attention_heads + dim = int(head_dim * partial_rotary_factor) + + inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2, dtype=torch.int64).to(device=device, dtype=torch.float) / dim)) + return inv_freq, 1.0 + + +if "default" not in ROPE_INIT_FUNCTIONS: + ROPE_INIT_FUNCTIONS["default"] = _compute_default_rope_parameters + + @dataclass class RotaryEmbeddingConfig: max_position_embeddings: int @@ -28,6 +52,10 @@ def __init__(self, config: RotaryEmbeddingConfig, device=None): self.register_buffer("inv_freq", inv_freq, persistent=False) self.original_inv_freq = self.inv_freq + def compute_default_rope_parameters(self, config=None, device=None, seq_len=None, layer_type=None): + """Required by transformers 5.0.0 for weight initialization when rope_type is 'default'.""" + return _compute_default_rope_parameters(config or self.config, device, seq_len, layer_type) + @torch.no_grad() @dynamic_rope_update # power user: used with advanced RoPE types (e.g. dynamic rope) def forward(self, x, position_ids): @@ -87,3 +115,34 @@ def apply_rotary_pos_emb(q, k, cos, sin, position_ids=None, unsqueeze_dim=1): q_embed = torch.cat([q_embed, q_pass], dim=-1) k_embed = torch.cat([k_embed, k_pass], dim=-1) return q_embed, k_embed + + +def apply_rotary_pos_emb_interleave(q, k, cos, sin, unsqueeze_dim=1): + """Interleaved RoPE: converts from [r0, i0, r1, i1, ...] layout before applying rotation. + + Args: + q (`torch.Tensor`): The query tensor. + k (`torch.Tensor`): The key tensor. + cos (`torch.Tensor`): The cosine part of the rotary embedding. + sin (`torch.Tensor`): The sine part of the rotary embedding. + unsqueeze_dim (`int`, *optional*, defaults to 1): + The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and + sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note + that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and + k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes + cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have + the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2. + Returns: + `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding. + """ + cos = cos.unsqueeze(unsqueeze_dim) + sin = sin.unsqueeze(unsqueeze_dim) + + b, h, s, d = q.shape + q = q.view(b, h, s, d // 2, 2).transpose(4, 3).reshape(b, h, s, d) + b, h, s, d = k.shape + k = k.view(b, h, s, d // 2, 2).transpose(4, 3).reshape(b, h, s, d) + + q_embed = (q * cos) + (rotate_half(q) * sin) + k_embed = (k * cos) + (rotate_half(k) * sin) + return q_embed, k_embed diff --git a/src/prime_rl/trainer/models/layers/ulysses_attn.py b/src/prime_rl/trainer/models/layers/ulysses_attn.py new file mode 100644 index 00000000000..39139e40982 --- /dev/null +++ b/src/prime_rl/trainer/models/layers/ulysses_attn.py @@ -0,0 +1,296 @@ +"""Ulysses-style context parallelism via all-to-all on the head/seq dimensions. + +Ulysses is the simpler alternative to ring attention for context parallelism: + + sequence-sharded Q/K/V head-sharded Q/K/V + [S/cp, H, D] ── all-to-all ──▶ [S, H/cp, D] + (heads ↔ seq) + │ + ▼ + run vanilla local attention + on the *full* sequence with + *fewer* heads (any kernel works) + │ + ▼ + [S/cp, H, D] ◀── all-to-all ── [S, H/cp, D] + (seq ↔ heads) + +Key benefit: the attention kernel itself does not need to be CP-aware. The +all-to-all is purely on Q/K/V tensors, so this works out of the box with +softmax flash-attn, linear attention, mamba, etc., without rewriting kernels. + +Constraints: +- cp_size must divide both num_attention_heads and num_key_value_heads. +- Sequence length must be divisible by cp_size. +""" + +from __future__ import annotations + +import torch +import torch.distributed as dist +import torch.distributed.nn.functional as dist_nn + +# Populated by `update_ulysses_params` before each forward pass. Mirrors +# ring_flash_attn's DATA_PARAMS pattern so the patched attention path can +# reach the *full* (un-sharded) cu_seqlens / max_seqlen at call time. +ULYSSES_PARAMS: dict = {} + + +def update_ulysses_params(cu_seqlens: torch.Tensor, max_seqlen: int) -> None: + ULYSSES_PARAMS["cu_seqlens"] = cu_seqlens + ULYSSES_PARAMS["max_seqlen"] = int(max_seqlen) + + +def _all_to_all_seq_to_head(t: torch.Tensor, cp_size: int, cp_group: dist.ProcessGroup) -> torch.Tensor: + """Redistribute [S_local, H, D] -> [S_global, H_local, D]. + + Splits the head dim into cp_size groups and exchanges them so each rank + ends up with the full sequence but only H/cp_size heads. Differentiable. + """ + s_local, h, d = t.shape + assert h % cp_size == 0, f"num_heads ({h}) must be divisible by cp_size ({cp_size})" + h_local = h // cp_size + + # [S_local, cp_size, H_local, D] -> [cp_size, S_local, H_local, D] + t = t.reshape(s_local, cp_size, h_local, d).transpose(0, 1).contiguous() + output = torch.empty_like(t) + out = dist_nn.all_to_all_single(output, t, group=cp_group) + # out[i] is the chunk that source-rank i had at position my_rank, i.e. + # source-rank i's local sequence shard for *my* head slice. + return out.reshape(cp_size * s_local, h_local, d) + + +def _all_to_all_head_to_seq(t: torch.Tensor, cp_size: int, cp_group: dist.ProcessGroup) -> torch.Tensor: + """Inverse of `_all_to_all_seq_to_head`: [S_global, H_local, D] -> [S_local, H, D].""" + s_global, h_local, d = t.shape + assert s_global % cp_size == 0 + s_local = s_global // cp_size + h = h_local * cp_size + + t = t.reshape(cp_size, s_local, h_local, d).contiguous() + output = torch.empty_like(t) + out = dist_nn.all_to_all_single(output, t, group=cp_group) + # out[s']: original chunk for sequence-rank s' (which now becomes head-rank s'). + return out.transpose(0, 1).contiguous().reshape(s_local, h, d) + + +def ulysses_flash_attn_varlen_func( + flash_fn, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + cu_seqlens_q: torch.Tensor, + cu_seqlens_k: torch.Tensor, + max_seqlen_q: int, + max_seqlen_k: int, + causal: bool, + cp_group: dist.ProcessGroup, + cp_size: int, + flash_attn_version: int = 2, + window_size: tuple[int, int] = (-1, -1), + softmax_scale: float | None = None, + dropout_p: float = 0.0, + deterministic: bool | None = None, +) -> torch.Tensor: + """Run varlen flash attention under Ulysses CP. + + `cu_seqlens_*` and `max_seqlen_*` describe the *full* (un-sharded) sequence, + because after the seq->head all-to-all each rank holds the full sequence. + """ + q = _all_to_all_seq_to_head(q, cp_size, cp_group) + k = _all_to_all_seq_to_head(k, cp_size, cp_group) + v = _all_to_all_seq_to_head(v, cp_size, cp_group) + + kwargs: dict = {"causal": causal} + if window_size != (-1, -1): + kwargs["window_size"] = window_size + if softmax_scale is not None: + kwargs["softmax_scale"] = softmax_scale + if dropout_p: + kwargs["dropout_p"] = dropout_p + if deterministic is not None: + kwargs["deterministic"] = deterministic + + if flash_attn_version == 4: + # FA4 takes cu_seqlens as keyword args (qv positional collides otherwise). + kwargs["cu_seqlens_q"] = cu_seqlens_q + kwargs["cu_seqlens_k"] = cu_seqlens_k + out = flash_fn(q, k, v, **kwargs) + else: + out = flash_fn(q, k, v, cu_seqlens_q, cu_seqlens_k, max_seqlen_q, max_seqlen_k, **kwargs) + if isinstance(out, tuple): + out = out[0] + + return _all_to_all_head_to_seq(out, cp_size, cp_group) + + +def substitute_ulysses_attn( + process_group: dist.ProcessGroup, + attn_impl: str = "flash_attention_2", +) -> None: + """Patch FlashAttention._compute_attention to do Ulysses all-to-all + local FA. + + Mirrors `substitute_ring_attn` so swapping CP styles is a one-line change + in the trainer. + """ + cp_size = dist.get_world_size(group=process_group) + + # Resolve flash kernel + version from attn_impl. + if attn_impl == "fa4": + from flash_attn.cute import flash_attn_varlen_func as flash_fn + + flash_attn_version = 4 + flash_fn = torch._dynamo.disable(flash_fn) + elif attn_impl == "flash_attention_3": + from flash_attn_interface import flash_attn_varlen_func as flash_fn + + flash_attn_version = 3 + else: + from flash_attn import flash_attn_varlen_func as flash_fn + + flash_attn_version = 2 + + def _ulysses_compute_attention(self, q, k, v, cu_seqlens, max_seqlen): + # cu_seqlens / max_seqlen passed in are for the *local* sharded sequence; + # ulysses needs the *full* ones (each rank holds the full seq after a2a). + cu_seqlens_full = ULYSSES_PARAMS["cu_seqlens"] + max_seqlen_full = ULYSSES_PARAMS["max_seqlen"] + + window_size = (-1, -1) + sliding_window = getattr(self, "sliding_window", None) + if sliding_window is not None: + window_size = (sliding_window - 1, 0) + + return ulysses_flash_attn_varlen_func( + flash_fn, + q, + k, + v, + cu_seqlens_q=cu_seqlens_full, + cu_seqlens_k=cu_seqlens_full, + max_seqlen_q=max_seqlen_full, + max_seqlen_k=max_seqlen_full, + causal=True, + cp_group=process_group, + cp_size=cp_size, + flash_attn_version=flash_attn_version, + window_size=window_size, + ) + + from prime_rl.trainer.models.layers.attn import FlashAttention + + FlashAttention._compute_attention = _ulysses_compute_attention + + from prime_rl.trainer.models.afmoe.modeling_afmoe import AfmoeFlashAttention + + AfmoeFlashAttention._compute_attention = _ulysses_compute_attention + + from prime_rl.trainer.models.qwen3_5_moe.modeling_qwen3_5_moe import Qwen3_5MoeGatedFlashAttention + + Qwen3_5MoeGatedFlashAttention._compute_attention = _ulysses_compute_attention + + +def substitute_hf_ulysses_attn(process_group: dist.ProcessGroup) -> None: + """Patch HF's `_flash_attention_forward` to use Ulysses all-to-all + local FA2. + + Used for HF (non-custom) model paths, e.g. qwen2/qwen3 dense models. Mirrors + ring_flash_attn's `substitute_hf_flash_attn` but with all-to-all instead of + ring all-gather. + """ + import transformers + import transformers.modeling_flash_attention_utils + from flash_attn import flash_attn_varlen_func + + cp_size = dist.get_world_size(group=process_group) + + def _ulysses_flash_attention_forward( + query_states: torch.Tensor, + key_states: torch.Tensor, + value_states: torch.Tensor, + attention_mask, + query_length: int, + is_causal: bool, + dropout: float = 0.0, + position_ids=None, + softmax_scale=None, + sliding_window=None, + use_top_left_mask: bool = False, + softcap=None, + deterministic=None, + **kwargs, + ): + assert is_causal, "ulysses CP only supports causal attention" + assert softcap is None, "ulysses CP path does not support softcap" + assert query_states.size(0) == 1, "varlen data should be processed with batch=1" + + cu_seqlens = ULYSSES_PARAMS["cu_seqlens"] + max_seqlen = ULYSSES_PARAMS["max_seqlen"] + + window_size = (-1, -1) + if sliding_window is not None and key_states.shape[1] > sliding_window: + window_size = (sliding_window, sliding_window) + + out = ulysses_flash_attn_varlen_func( + flash_attn_varlen_func, + query_states.squeeze(0), + key_states.squeeze(0), + value_states.squeeze(0), + cu_seqlens_q=cu_seqlens, + cu_seqlens_k=cu_seqlens, + max_seqlen_q=max_seqlen, + max_seqlen_k=max_seqlen, + causal=True, + cp_group=process_group, + cp_size=cp_size, + flash_attn_version=2, + window_size=window_size, + softmax_scale=softmax_scale, + dropout_p=dropout, + deterministic=deterministic, + ) + return out.unsqueeze(0) + + transformers.modeling_flash_attention_utils._flash_attention_forward = _ulysses_flash_attention_forward + + # Newer transformers route attention through ALL_ATTENTION_FUNCTIONS["flash_attention_2"]. + try: + from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS + except ImportError: + ALL_ATTENTION_FUNCTIONS = None + + if ALL_ATTENTION_FUNCTIONS is not None: + + def _ulysses_flash_attention_forward_v2( + module, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + attention_mask=None, + dropout: float = 0.0, + scaling=None, + sliding_window=None, + softcap=None, + **kw, + ): + # Match HF v2 entrypoint: query/key/value arrive as [B, H, S, D]. + seq_len = query.shape[2] + query = query.transpose(1, 2) + key = key.transpose(1, 2) + value = value.transpose(1, 2) + + kw.pop("is_causal", None) + attn_out = _ulysses_flash_attention_forward( + query, + key, + value, + attention_mask, + query_length=seq_len, + is_causal=module.is_causal, + dropout=dropout, + softmax_scale=scaling, + sliding_window=sliding_window, + softcap=softcap, + ) + return attn_out, None + + ALL_ATTENTION_FUNCTIONS["flash_attention_2"] = _ulysses_flash_attention_forward_v2 diff --git a/src/prime_rl/trainer/models/llama/modeling_llama.py b/src/prime_rl/trainer/models/llama/modeling_llama.py index 03112736980..4e3c2f784fe 100644 --- a/src/prime_rl/trainer/models/llama/modeling_llama.py +++ b/src/prime_rl/trainer/models/llama/modeling_llama.py @@ -38,8 +38,9 @@ from prime_rl.trainer.models.layers.attn import ATTN_IMPL2CLASS, AttentionConfig from prime_rl.trainer.models.layers.lm_head import PrimeLmOutput from prime_rl.trainer.models.layers.mlp import MLP, MLPConfig -from prime_rl.trainer.models.layers.rms_norm import RMSNorm, RMSNormConfig +from prime_rl.trainer.models.layers.norms import RMSNorm, RMSNormConfig from prime_rl.trainer.models.layers.rotary_emb import RotaryEmbedding, RotaryEmbeddingConfig +from prime_rl.utils.sequence import get_cu_seqlens_from_position_ids logger = logging.get_logger(__name__) @@ -187,17 +188,8 @@ def forward( if inputs_embeds is None: inputs_embeds: torch.Tensor = self.embed_tokens(input_ids) - if self.config._attn_implementation in ("flash_attention_2", "flash_attention_3"): - flat_position_ids = position_ids.view(-1) - seqlens = torch.cat( - [ - flat_position_ids[0:1], - flat_position_ids[:-1][(flat_position_ids == 0)[1:]] + 1, - flat_position_ids[-1:] + 1, - ] - ) - max_seqlen = seqlens.max().item() - cu_seqlens = seqlens.cumsum(dim=0, dtype=torch.int32) + if self.config._attn_implementation in ("flash_attention_2", "flash_attention_3", "fa4"): + cu_seqlens, max_seqlen = get_cu_seqlens_from_position_ids(position_ids) torch._dynamo.mark_dynamic(cu_seqlens, 0) else: max_seqlen = None @@ -222,7 +214,7 @@ def forward( @auto_docstring class LlamaForCausalLM(LlamaPreTrainedModel, GenerationMixin): - _tied_weights_keys = ["lm_head.weight"] + _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": "colwise_rep"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} @@ -248,15 +240,19 @@ def forward( use_cache: Optional[bool] = None, cache_position: Optional[torch.LongTensor] = None, logits_to_keep: Union[int, torch.Tensor] = 0, - temperature: Optional[float] = 1.0, + temperature: Optional[torch.Tensor] = None, **kwargs: Unpack[TransformersKwargs], ) -> PrimeLmOutput: r""" + cache_position (`torch.LongTensor` of shape `(sequence_length)`, *optional*): + Indices of input tokens in the KV cache. Accepted only for HuggingFace API + compatibility — prime-rl asserts `use_cache is None` since training does not + perform autoregressive decoding, so this argument is unused. labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): Labels used by PrimeRL's wrapped LM head to optionally compute per-token logprobs/entropy. If not provided, the wrapped LM head returns logits only. - temperature (`float`, *optional*, defaults to 1.0): - Temperature used for the logprobs/entropy computation when `labels` are provided. + temperature (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*): + Per-token temperatures for logprobs/entropy computation when `labels` are provided. Example: diff --git a/src/prime_rl/trainer/models/minimax_m2/__init__.py b/src/prime_rl/trainer/models/minimax_m2/__init__.py new file mode 100644 index 00000000000..075c7be2a46 --- /dev/null +++ b/src/prime_rl/trainer/models/minimax_m2/__init__.py @@ -0,0 +1,13 @@ +from prime_rl.trainer.models.minimax_m2.configuration_minimax_m2 import MiniMaxM2Config +from prime_rl.trainer.models.minimax_m2.modeling_minimax_m2 import ( + MiniMaxM2ForCausalLM, + MiniMaxM2Model, + MiniMaxM2PreTrainedModel, +) + +__all__ = [ + "MiniMaxM2Config", + "MiniMaxM2ForCausalLM", + "MiniMaxM2Model", + "MiniMaxM2PreTrainedModel", +] diff --git a/src/prime_rl/trainer/models/minimax_m2/configuration_minimax_m2.py b/src/prime_rl/trainer/models/minimax_m2/configuration_minimax_m2.py new file mode 100644 index 00000000000..6fcceec309f --- /dev/null +++ b/src/prime_rl/trainer/models/minimax_m2/configuration_minimax_m2.py @@ -0,0 +1,150 @@ +from transformers.configuration_utils import PretrainedConfig + + +class MiniMaxM2Config(PretrainedConfig): + r""" + Configuration class for MiniMax M2.1 MoE model. + + Args: + vocab_size (`int`, *optional*, defaults to 200064): + Vocabulary size of the model. + hidden_size (`int`, *optional*, defaults to 6144): + Dimension of the hidden representations. + intermediate_size (`int`, *optional*, defaults to 24576): + Dimension of the MoE expert FFN. + num_hidden_layers (`int`, *optional*, defaults to 92): + Number of hidden layers in the Transformer. + num_attention_heads (`int`, *optional*, defaults to 48): + Number of attention heads. + num_key_value_heads (`int`, *optional*, defaults to 8): + Number of key-value heads for GQA. + head_dim (`int`, *optional*, defaults to 128): + Dimension of each attention head. + hidden_act (`str`, *optional*, defaults to `"silu"`): + Activation function for the MLP. + max_position_embeddings (`int`, *optional*, defaults to 131072): + Maximum sequence length. + rms_norm_eps (`float`, *optional*, defaults to 1e-6): + Epsilon for RMS normalization. + rope_theta (`float`, *optional*, defaults to 5000000): + Base period for RoPE. + rotary_dim (`int`, *optional*, defaults to 64): + Dimension of the rotary embeddings. Converted to `partial_rotary_factor` for HF rope init. + num_local_experts (`int`, *optional*, defaults to 256): + Number of routed experts per layer. + num_experts_per_tok (`int`, *optional*, defaults to 8): + Number of experts selected per token. + scoring_func (`str`, *optional*, defaults to `"sigmoid"`): + Scoring function for the router. + use_routing_bias (`bool`, *optional*, defaults to `True`): + Whether to use e_score_correction_bias for load balancing. + use_qk_norm (`bool`, *optional*, defaults to `True`): + Whether to use QK normalization. + qk_norm_type (`str`, *optional*, defaults to `"per_layer"`): + Type of QK normalization. "per_layer" means each head has its own scale. + attention_bias (`bool`, *optional*, defaults to `False`): + Whether to use bias in attention projections. + use_grouped_mm (`bool`, *optional*, defaults to `True`): + Whether to use grouped matmul for experts. + """ + + model_type = "minimax_m2" + keys_to_ignore_at_inference = ["past_key_values"] + + base_model_tp_plan = { + "layers.*.self_attn.q_proj": "colwise", + "layers.*.self_attn.k_proj": "colwise", + "layers.*.self_attn.v_proj": "colwise", + "layers.*.self_attn.o_proj": "rowwise", + } + base_model_pp_plan = { + "embed_tokens": (["input_ids"], ["inputs_embeds"]), + "layers": (["hidden_states", "attention_mask"], ["hidden_states"]), + "norm": (["hidden_states"], ["hidden_states"]), + } + + def __init__( + self, + vocab_size=200064, + hidden_size=6144, + intermediate_size=24576, + num_hidden_layers=92, + num_attention_heads=48, + num_key_value_heads=8, + head_dim=128, + hidden_act="silu", + max_position_embeddings=131072, + initializer_range=0.02, + rms_norm_eps=1e-6, + use_cache=True, + tie_word_embeddings=False, + rope_theta=5000000, + rope_scaling=None, + rotary_dim=64, + pad_token_id=None, + bos_token_id=1, + eos_token_id=2, + num_local_experts=256, + num_experts_per_tok=8, + scoring_func="sigmoid", + use_routing_bias=True, + use_qk_norm=True, + qk_norm_type="per_layer", + attention_bias=False, + sliding_window=None, + attention_dropout=0.0, + output_router_logits=False, + router_aux_loss_coef=0.001, + router_jitter_noise=0.0, + use_grouped_mm=True, + **kwargs, + ): + self.vocab_size = vocab_size + self.hidden_size = hidden_size + self.intermediate_size = intermediate_size + self.num_hidden_layers = num_hidden_layers + self.num_attention_heads = num_attention_heads + self.num_key_value_heads = num_key_value_heads + self.head_dim = head_dim + self.hidden_act = hidden_act + self.max_position_embeddings = max_position_embeddings + self.initializer_range = initializer_range + self.rms_norm_eps = rms_norm_eps + self.use_cache = use_cache + self.rope_theta = rope_theta + self.rope_scaling = rope_scaling + self.attention_bias = attention_bias + self.sliding_window = sliding_window + self.attention_dropout = attention_dropout + + # Compute partial_rotary_factor from rotary_dim for HF rope init compatibility + self.rotary_dim = rotary_dim + self.partial_rotary_factor = rotary_dim / head_dim + + if self.rope_scaling is not None and "type" in self.rope_scaling: + self.rope_scaling["rope_type"] = self.rope_scaling["type"] + self.standardize_rope_params() + + # MoE arguments + self.num_local_experts = num_local_experts + self.num_experts_per_tok = num_experts_per_tok + self.scoring_func = scoring_func + self.use_routing_bias = use_routing_bias + self.output_router_logits = output_router_logits + self.router_aux_loss_coef = router_aux_loss_coef + self.router_jitter_noise = router_jitter_noise + + # Attention + self.use_qk_norm = use_qk_norm + self.qk_norm_type = qk_norm_type + + # Training + self.use_grouped_mm = use_grouped_mm + + super().__init__( + pad_token_id=pad_token_id, + bos_token_id=bos_token_id, + eos_token_id=eos_token_id, + tie_word_embeddings=tie_word_embeddings, + **kwargs, + ) diff --git a/src/prime_rl/trainer/models/minimax_m2/converting_minimax_m2.py b/src/prime_rl/trainer/models/minimax_m2/converting_minimax_m2.py new file mode 100644 index 00000000000..706d2f035d0 --- /dev/null +++ b/src/prime_rl/trainer/models/minimax_m2/converting_minimax_m2.py @@ -0,0 +1,101 @@ +import torch +from torch import Tensor + + +def get_max_layer_num(state_dict: dict[str, Tensor]) -> int: + """Get the maximum number of layers in the model.""" + return max(int(i.split(".")[2]) for i in state_dict.keys() if "model.layers." in i) + 1 + + +def convert_hf_layer_to_tt(state_dict: dict[str, Tensor], layer_idx: int): + """Convert a layer from HF to PrimeRL format in-place. + + HF MiniMax M2.1 uses `block_sparse_moe` as the MoE block name, with + expert weights stored as `experts.{j}.w1.weight` (nn.Linear format). + """ + i = layer_idx + prefix = f"model.layers.{i}" + + # Check if this layer has MoE experts + num_experts = len([k for k in state_dict.keys() if f"{prefix}.block_sparse_moe.experts" in k]) // 3 + if num_experts == 0: + return + + # Router: block_sparse_moe.gate.weight -> mlp.router.gate.weight + state_dict[f"{prefix}.mlp.router.gate.weight"] = state_dict[f"{prefix}.block_sparse_moe.gate.weight"] + del state_dict[f"{prefix}.block_sparse_moe.gate.weight"] + + # e_score_correction_bias: direct tensor under block_sparse_moe -> mlp.expert_bias + bias_key = f"{prefix}.block_sparse_moe.e_score_correction_bias" + if bias_key in state_dict: + state_dict[f"{prefix}.mlp.expert_bias"] = state_dict[bias_key] + del state_dict[bias_key] + + # Expert weights: stack individual experts into grouped tensors + # HF uses w1.weight, w2.weight, w3.weight (nn.Linear with .weight suffix) + dim, moe_dim = state_dict[f"{prefix}.block_sparse_moe.experts.0.w2.weight"].shape + dtype = state_dict[f"{prefix}.block_sparse_moe.experts.0.w2.weight"].dtype + + w1 = torch.empty((num_experts, moe_dim, dim), dtype=dtype) + w2 = torch.empty((num_experts, dim, moe_dim), dtype=dtype) + w3 = torch.empty((num_experts, moe_dim, dim), dtype=dtype) + + for j in range(num_experts): + w1[j].copy_(state_dict[f"{prefix}.block_sparse_moe.experts.{j}.w1.weight"]) + w2[j].copy_(state_dict[f"{prefix}.block_sparse_moe.experts.{j}.w2.weight"]) + w3[j].copy_(state_dict[f"{prefix}.block_sparse_moe.experts.{j}.w3.weight"]) + + del state_dict[f"{prefix}.block_sparse_moe.experts.{j}.w1.weight"] + del state_dict[f"{prefix}.block_sparse_moe.experts.{j}.w2.weight"] + del state_dict[f"{prefix}.block_sparse_moe.experts.{j}.w3.weight"] + + state_dict[f"{prefix}.mlp.experts.w1"] = w1 + state_dict[f"{prefix}.mlp.experts.w2"] = w2 + state_dict[f"{prefix}.mlp.experts.w3"] = w3 + + +def convert_tt_layer_to_hf(state_dict: dict[str, Tensor], layer_idx: int): + """Convert a layer from PrimeRL to HF format in-place.""" + i = layer_idx + prefix = f"model.layers.{i}" + + if f"{prefix}.mlp.router.gate.weight" not in state_dict: + return + + # Router: mlp.router.gate.weight -> block_sparse_moe.gate.weight + state_dict[f"{prefix}.block_sparse_moe.gate.weight"] = state_dict[f"{prefix}.mlp.router.gate.weight"] + del state_dict[f"{prefix}.mlp.router.gate.weight"] + + # expert_bias -> block_sparse_moe.e_score_correction_bias + if f"{prefix}.mlp.expert_bias" in state_dict: + state_dict[f"{prefix}.block_sparse_moe.e_score_correction_bias"] = state_dict[f"{prefix}.mlp.expert_bias"] + del state_dict[f"{prefix}.mlp.expert_bias"] + + # tokens_per_expert is a runtime buffer, remove if present + if f"{prefix}.mlp.tokens_per_expert" in state_dict: + del state_dict[f"{prefix}.mlp.tokens_per_expert"] + + # Expert weights: unstack grouped tensors into individual experts + num_experts, moe_dim, dim = state_dict[f"{prefix}.mlp.experts.w1"].shape + for j in range(num_experts): + state_dict[f"{prefix}.block_sparse_moe.experts.{j}.w1.weight"] = state_dict[f"{prefix}.mlp.experts.w1"][j] + state_dict[f"{prefix}.block_sparse_moe.experts.{j}.w2.weight"] = state_dict[f"{prefix}.mlp.experts.w2"][j] + state_dict[f"{prefix}.block_sparse_moe.experts.{j}.w3.weight"] = state_dict[f"{prefix}.mlp.experts.w3"][j] + + del state_dict[f"{prefix}.mlp.experts.w1"] + del state_dict[f"{prefix}.mlp.experts.w2"] + del state_dict[f"{prefix}.mlp.experts.w3"] + + +def convert_hf_to_tt_moe(state_dict: dict[str, Tensor]): + """Convert MoE weights from HF to PrimeRL format in-place.""" + num_layers = get_max_layer_num(state_dict) + for i in range(num_layers): + convert_hf_layer_to_tt(state_dict, i) + + +def convert_tt_to_hf_moe(state_dict: dict[str, Tensor]): + """Convert MoE weights from PrimeRL to HF format in-place.""" + num_layers = get_max_layer_num(state_dict) + for i in range(num_layers): + convert_tt_layer_to_hf(state_dict, i) diff --git a/src/prime_rl/trainer/models/minimax_m2/modeling_minimax_m2.py b/src/prime_rl/trainer/models/minimax_m2/modeling_minimax_m2.py new file mode 100644 index 00000000000..34b25666ad8 --- /dev/null +++ b/src/prime_rl/trainer/models/minimax_m2/modeling_minimax_m2.py @@ -0,0 +1,285 @@ +from typing import Optional, Union + +import torch +from torch import Tensor, nn +from transformers.cache_utils import Cache +from transformers.generation import GenerationMixin +from transformers.modeling_layers import GradientCheckpointingLayer +from transformers.modeling_outputs import BaseModelOutputWithPast +from transformers.processing_utils import Unpack +from transformers.utils import TransformersKwargs, auto_docstring, can_return_tuple, logging + +from prime_rl.trainer.models.base import PreTrainedModelPrimeRL +from prime_rl.trainer.models.layers.attn import ATTN_IMPL2CLASS, AttentionConfig +from prime_rl.trainer.models.layers.lm_head import PrimeLmOutput +from prime_rl.trainer.models.layers.moe import MoE, MoEArgs +from prime_rl.trainer.models.layers.norms import RMSNorm, RMSNormConfig +from prime_rl.trainer.models.layers.rotary_emb import RotaryEmbedding, RotaryEmbeddingConfig +from prime_rl.trainer.models.minimax_m2.configuration_minimax_m2 import MiniMaxM2Config +from prime_rl.trainer.models.minimax_m2.converting_minimax_m2 import ( + convert_hf_layer_to_tt, + convert_hf_to_tt_moe, + convert_tt_layer_to_hf, + convert_tt_to_hf_moe, +) +from prime_rl.utils.sequence import get_cu_seqlens_from_position_ids + +logger = logging.get_logger(__name__) + + +class MiniMaxM2DecoderLayer(GradientCheckpointingLayer): + def __init__(self, config: MiniMaxM2Config, layer_idx: int): + super().__init__() + self.hidden_size = config.hidden_size + + attn_config = AttentionConfig( + hidden_size=config.hidden_size, + head_dim=config.head_dim, + num_attention_heads=config.num_attention_heads, + num_key_value_heads=config.num_key_value_heads, + is_causal=True, + attention_bias=config.attention_bias, + use_qk_norm=config.use_qk_norm, + rms_norm_eps=config.rms_norm_eps, + qk_norm_type=config.qk_norm_type, + ) + self.self_attn = ATTN_IMPL2CLASS[config._attn_implementation](attn_config) + + moe_args = MoEArgs( + num_experts=config.num_local_experts, + num_shared_experts=0, + score_func=config.scoring_func, + route_norm=True, + route_scale=1.0, + score_before_experts=False, + top_k=config.num_experts_per_tok, + use_grouped_mm=config.use_grouped_mm, + load_balance_coeff=1e-3 if config.use_routing_bias else None, + fp8=getattr(config, "fp8", False), + ) + self.mlp = MoE(moe_args, dim=config.hidden_size, hidden_dim=config.intermediate_size) + + self.input_layernorm = RMSNorm(RMSNormConfig(hidden_size=config.hidden_size, eps=config.rms_norm_eps)) + self.post_attention_layernorm = RMSNorm(RMSNormConfig(hidden_size=config.hidden_size, eps=config.rms_norm_eps)) + + def forward( + self, + hidden_states: torch.Tensor, + position_embeddings: tuple[torch.Tensor, torch.Tensor], + cu_seqlens: torch.LongTensor | None = None, + max_seqlen: int | None = None, + routed_experts: Optional[torch.LongTensor] = None, + ) -> torch.FloatTensor: + residual = hidden_states + hidden_states = self.input_layernorm(hidden_states) + hidden_states, _ = self.self_attn( + hidden_states=hidden_states, + position_embeddings=position_embeddings, + cu_seqlens=cu_seqlens, + max_seqlen=max_seqlen, + ) + hidden_states = residual + hidden_states + + residual = hidden_states + hidden_states = self.post_attention_layernorm(hidden_states) + hidden_states = self.mlp(hidden_states, routed_experts=routed_experts) + hidden_states = residual + hidden_states + return hidden_states + + +@auto_docstring +class MiniMaxM2PreTrainedModel(PreTrainedModelPrimeRL): + config: MiniMaxM2Config + base_model_prefix = "model" + supports_gradient_checkpointing = True + _no_split_modules = ["MiniMaxM2DecoderLayer"] + _skip_keys_device_placement = ["past_key_values"] + _supports_flash_attn = True + _supports_sdpa = True + _supports_flex_attn = True + _can_compile_fullgraph = False + _supports_attention_backend = True + _can_record_outputs = { + "hidden_states": MiniMaxM2DecoderLayer, + } + + @classmethod + def is_hf_state_dict(cls, state_dict: dict[str, Tensor]) -> bool: + return any("block_sparse_moe.experts.1.w1" in name for name in state_dict.keys()) + + @classmethod + def is_prime_state_dict(cls, state_dict: dict[str, Tensor]) -> bool: + return any("mlp.experts.w1" in name for name in state_dict.keys()) + + @classmethod + def convert_to_hf(cls, state_dict: dict[str, Tensor]) -> dict[str, Tensor]: + convert_tt_to_hf_moe(state_dict) + return state_dict + + @classmethod + def convert_to_prime(cls, state_dict: dict[str, Tensor]) -> dict[str, Tensor]: + convert_hf_to_tt_moe(state_dict) + return state_dict + + @classmethod + def convert_layer_to_hf(cls, state_dict: dict[str, Tensor], layer_idx: int) -> dict[str, Tensor]: + convert_tt_layer_to_hf(state_dict, layer_idx) + return state_dict + + @classmethod + def convert_layer_to_prime(cls, state_dict: dict[str, Tensor], layer_idx: int) -> dict[str, Tensor]: + convert_hf_layer_to_tt(state_dict, layer_idx) + return state_dict + + +@auto_docstring +class MiniMaxM2Model(MiniMaxM2PreTrainedModel): + def __init__(self, config: MiniMaxM2Config): + super().__init__(config) + self.padding_idx = config.pad_token_id + self.vocab_size = config.vocab_size + + self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx) + self.layers = nn.ModuleList( + [MiniMaxM2DecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)] + ) + self.norm = RMSNorm(RMSNormConfig(hidden_size=config.hidden_size, eps=config.rms_norm_eps)) + + if hasattr(config, "rope_scaling") and isinstance(config.rope_scaling, dict): + rope_type = config.rope_scaling.get("rope_type", config.rope_scaling.get("type")) + else: + rope_type = "default" + rotary_config = RotaryEmbeddingConfig( + max_position_embeddings=config.max_position_embeddings, + rope_type=rope_type, + model_config=config, + ) + self.rotary_emb = RotaryEmbedding(rotary_config) + self.gradient_checkpointing = False + + self.post_init() + + @auto_docstring + def forward( + self, + input_ids: Optional[torch.LongTensor] = None, + position_ids: Optional[torch.LongTensor] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + routed_experts: Optional[torch.LongTensor] = None, + ) -> BaseModelOutputWithPast: + """ + routed_experts (`torch.LongTensor` of shape `(batch_size, sequence_length, num_hidden_layers, num_experts_per_tok)`, *optional*): + Routed experts for each token in the sequence. Only used for router replay. + """ + if (input_ids is None) ^ (inputs_embeds is not None): + raise ValueError("You must specify exactly one of input_ids or inputs_embeds") + + if inputs_embeds is None: + inputs_embeds = self.embed_tokens(input_ids) + + if self.config._attn_implementation in ("flash_attention_2", "flash_attention_3", "fa4"): + cu_seqlens, max_seqlen = get_cu_seqlens_from_position_ids(position_ids) + torch._dynamo.mark_dynamic(cu_seqlens, 0) + else: + max_seqlen = None + cu_seqlens = None + + hidden_states = inputs_embeds + position_embeddings = self.rotary_emb(hidden_states, position_ids) + + for layer_idx, decoder_layer in enumerate(self.layers[: self.config.num_hidden_layers]): + routed_experts_layer = routed_experts[:, :, layer_idx, :] if routed_experts is not None else None + hidden_states = decoder_layer( + hidden_states, + position_embeddings=position_embeddings, + cu_seqlens=cu_seqlens, + max_seqlen=max_seqlen, + routed_experts=routed_experts_layer, + ) + + hidden_states = self.norm(hidden_states) + return BaseModelOutputWithPast(last_hidden_state=hidden_states) + + +@auto_docstring +class MiniMaxM2ForCausalLM(MiniMaxM2PreTrainedModel, GenerationMixin): + _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} + _tp_plan = {"lm_head": "colwise_rep"} + _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + + def __init__(self, config): + super().__init__(config) + self.model = MiniMaxM2Model(config) + self.vocab_size = config.vocab_size + self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False) + + self.post_init() + + @can_return_tuple + @auto_docstring + def forward( + self, + input_ids: Optional[torch.LongTensor] = None, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_values: Optional[Cache] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + labels: Optional[torch.LongTensor] = None, + use_cache: Optional[bool] = None, + cache_position: Optional[torch.LongTensor] = None, + logits_to_keep: Union[int, torch.Tensor] = 0, + temperature: Union[torch.Tensor, None] = None, + routed_experts: Optional[torch.LongTensor] = None, + **kwargs: Unpack[TransformersKwargs], + ) -> PrimeLmOutput: + r""" + cache_position (`torch.LongTensor` of shape `(sequence_length)`, *optional*): + Indices of input tokens in the KV cache. Accepted only for HuggingFace API + compatibility — prime-rl asserts `use_cache is None` since training does not + perform autoregressive decoding, so this argument is unused. + labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): + Labels for computing the masked language modeling loss. + temperature (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*): + Per-token temperatures for logprobs/entropy computation. + routed_experts (`torch.LongTensor` of shape `(batch_size, sequence_length, num_hidden_layers, num_experts_per_tok)`, *optional*): + Routed experts for each token in the sequence. Only used for router replay. + """ + assert use_cache is None, "use_cache is not supported for custom minimax_m2 for now" + assert past_key_values is None, "past_key_values is not supported for custom minimax_m2 for now" + + if position_ids is None: + if inputs_embeds is not None: + position_ids = torch.arange(inputs_embeds.shape[1], device=inputs_embeds.device).unsqueeze(0) + else: + position_ids = torch.arange(input_ids.shape[1], device=input_ids.device).unsqueeze(0) + + outputs: BaseModelOutputWithPast = self.model( + input_ids=input_ids, + position_ids=position_ids, + inputs_embeds=inputs_embeds, + routed_experts=routed_experts, + ) + + hidden_states = outputs.last_hidden_state + slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep + return self.lm_head( + hidden_states[:, slice_indices, :], + labels[:, slice_indices] if labels is not None else None, + temperature=temperature, + ) + + def init_buffers_post_meta(self): + buffer_names = [name for name, _ in self.named_buffers()] + if "model.rotary_emb.inv_freq" in buffer_names: + rotary_emb = self.model.rotary_emb + inv_freq, rotary_emb.attention_scaling = rotary_emb.rope_init_fn( + rotary_emb.config, rotary_emb.inv_freq.device + ) + rotary_emb.inv_freq.copy_(inv_freq) + + +__all__ = [ + "MiniMaxM2ForCausalLM", + "MiniMaxM2Model", + "MiniMaxM2PreTrainedModel", +] diff --git a/src/prime_rl/trainer/models/nemotron_h/__init__.py b/src/prime_rl/trainer/models/nemotron_h/__init__.py new file mode 100644 index 00000000000..391533a501c --- /dev/null +++ b/src/prime_rl/trainer/models/nemotron_h/__init__.py @@ -0,0 +1,13 @@ +from prime_rl.trainer.models.nemotron_h.configuration_nemotron_h import NemotronHConfig +from prime_rl.trainer.models.nemotron_h.modeling_nemotron_h import ( + NemotronHForCausalLM, + NemotronHModel, + NemotronHPreTrainedModel, +) + +__all__ = [ + "NemotronHConfig", + "NemotronHForCausalLM", + "NemotronHModel", + "NemotronHPreTrainedModel", +] diff --git a/src/prime_rl/trainer/models/nemotron_h/configuration_nemotron_h.py b/src/prime_rl/trainer/models/nemotron_h/configuration_nemotron_h.py new file mode 100644 index 00000000000..c1962737a78 --- /dev/null +++ b/src/prime_rl/trainer/models/nemotron_h/configuration_nemotron_h.py @@ -0,0 +1,227 @@ +from transformers.configuration_utils import PretrainedConfig + + +class NemotronHConfig(PretrainedConfig): + """Configuration for NemotronH (Nemotron-3-Super-120B-A12B) hybrid Mamba-Transformer-MoE model. + + The model architecture is defined by `layers_block_type`, a list where each element is + one of "mamba", "attention", or "moe". The 120B model uses a pattern like MEMEMEM*EMEMEMEM*... + (M=Mamba-2, E=LatentMoE, *=Attention) repeated across 88 layers. + + Args: + vocab_size: Vocabulary size. + hidden_size: Dimension of the hidden representations. + layers_block_type: Explicit list of layer types. Determines num_hidden_layers. + tie_word_embeddings: Whether to tie input/output embeddings. + pad_token_id: Padding token id. + bos_token_id: Beginning of sequence token id. + eos_token_id: End of sequence token id. + num_attention_heads: Number of attention heads. + num_key_value_heads: Number of key/value heads for GQA. + head_dim: Dimension of each attention head. + max_position_embeddings: Maximum sequence length. + attention_bias: Whether to use bias in attention layers. + intermediate_size: MLP intermediate dimension (for shared expert). + mlp_hidden_act: Activation function for MLP layers. + mlp_bias: Whether to use bias in MLP layers. + use_mamba_kernels: Whether to use fast mamba CUDA kernels. + ssm_state_size: Mamba state space dimension. + mamba_num_heads: Number of mamba heads. + mamba_n_groups: Number of groups in mamba. + mamba_head_dim: Dimension per mamba head. + mamba_d_conv: Convolution kernel size. + mamba_expand: Expansion factor. + mamba_hidden_act: Activation for mamba layers. + mamba_dt_min: Time step minimum. + mamba_dt_max: Time step maximum. + mamba_dt_limit: Time step limits. + mamba_dt_init_floor: Floor for time step initialization. + mamba_conv_bias: Whether to use bias in mamba convolution. + mamba_proj_bias: Whether to use bias in mamba projections. + mamba_chunk_size: Chunk size for mamba processing. + n_routed_experts: Number of routed experts. + n_shared_experts: Number of shared experts. + moe_intermediate_size: Expert intermediate dimension. + moe_shared_expert_intermediate_size: Shared expert intermediate dimension. + moe_latent_size: Latent projection size (None means use hidden_size). + num_experts_per_tok: Top-k routing parameter. + routed_scaling_factor: Scaling factor for routed expert outputs. + n_group: Number of groups for expert routing. + topk_group: Top-k group parameter. + norm_topk_prob: Whether to normalize top-k probabilities. + use_bias: Global bias setting. + initializer_range: Standard deviation for weight initialization. + layer_norm_epsilon: Epsilon for layer normalization. + residual_in_fp32: Whether to keep residuals in fp32. + rescale_prenorm_residual: Whether to rescale pre-norm residuals. + load_balance_coeff: Auxiliary-loss-free load balancing coefficient. + use_grouped_mm: Whether to use grouped MM for experts. + """ + + model_type = "nemotron_h" + keys_to_ignore_at_inference = ["past_key_values"] + + PATTERN_MAP = {"M": "mamba", "E": "moe", "*": "attention"} + REVERSE_PATTERN_MAP = {"mamba": "M", "moe": "E", "attention": "*"} + + def __init__( + self, + # General + vocab_size=131072, + hidden_size=4096, + layers_block_type=None, + num_hidden_layers=None, + tie_word_embeddings=False, + pad_token_id=0, + bos_token_id=1, + eos_token_id=2, + # Attention + num_attention_heads=32, + num_key_value_heads=8, + head_dim=128, + max_position_embeddings=4096, + attention_bias=False, + # MLP + intermediate_size=21504, + mlp_hidden_act="relu2", + mlp_bias=False, + # Mamba + use_mamba_kernels=True, + ssm_state_size=128, + mamba_num_heads=128, + mamba_n_groups=8, + mamba_head_dim=64, + mamba_d_conv=4, + mamba_expand=2, + mamba_hidden_act="silu", + mamba_dt_min=0.001, + mamba_dt_max=0.1, + mamba_dt_limit=(0.0, float("inf")), + mamba_dt_init_floor=1e-4, + mamba_conv_bias=True, + mamba_proj_bias=False, + mamba_chunk_size=128, + # MoE + n_routed_experts=8, + n_shared_experts=1, + moe_intermediate_size=7688, + moe_shared_expert_intermediate_size=7688, + moe_latent_size=None, + num_experts_per_tok=2, + routed_scaling_factor=1.0, + n_group=1, + topk_group=1, + norm_topk_prob=True, + # General training + use_bias=False, + initializer_range=0.02, + layer_norm_epsilon=1e-5, + residual_in_fp32=False, + rescale_prenorm_residual=True, + # PrimeRL training features + load_balance_coeff=None, + use_grouped_mm=True, + # RoPE + rope_theta=10000.0, + rope_scaling=None, + **kwargs, + ): + # Handle hybrid_override_pattern for backward compatibility + if "hybrid_override_pattern" in kwargs: + pattern = kwargs.pop("hybrid_override_pattern") + if layers_block_type is None: + layers_block_type = [self.PATTERN_MAP[c] for c in pattern] + + if layers_block_type is None: + layers_block_type = ["mamba", "moe", "attention", "moe"] + + self.vocab_size = vocab_size + self.hidden_size = hidden_size + self.layers_block_type = layers_block_type + + # Attention config + self.num_attention_heads = num_attention_heads + self.num_key_value_heads = num_key_value_heads if num_key_value_heads is not None else num_attention_heads + self.head_dim = head_dim + self.max_position_embeddings = max_position_embeddings + self.attention_bias = attention_bias + self.rope_theta = rope_theta + self.rope_scaling = rope_scaling + + # MLP config + self.intermediate_size = intermediate_size + self.mlp_hidden_act = mlp_hidden_act + self.mlp_bias = mlp_bias + + # Mamba config - stored with HF-compatible attribute names used by NemotronHMamba2Mixer + self.use_mamba_kernels = use_mamba_kernels + self.ssm_state_size = ssm_state_size + self.mamba_num_heads = mamba_num_heads + self.n_groups = mamba_n_groups + self.mamba_head_dim = mamba_head_dim + self.conv_kernel = mamba_d_conv + self.expand = mamba_expand + self.mamba_hidden_act = mamba_hidden_act + self.time_step_min = mamba_dt_min + self.time_step_max = mamba_dt_max + self.time_step_limit = mamba_dt_limit + self.time_step_floor = mamba_dt_init_floor + self.use_conv_bias = mamba_conv_bias + self.mamba_proj_bias = mamba_proj_bias + self.chunk_size = mamba_chunk_size + + # Zamba2MambaMixer compat aliases (read by parent __init__ before NemotronHMamba2Mixer overrides). + # mamba_expand must give the correct intermediate_size = mamba_num_heads * mamba_head_dim + # when Zamba2 computes int(mamba_expand * hidden_size); the config's raw "expand" field + # does not satisfy this for all model sizes (e.g. Nemotron-3-Nano-30B). + self.mamba_d_state = ssm_state_size + self.mamba_d_conv = mamba_d_conv + self.mamba_expand = (mamba_num_heads * mamba_head_dim) / hidden_size + self.mamba_ngroups = mamba_n_groups + self.mamba_headdim = mamba_head_dim + self.n_mamba_heads = mamba_num_heads + self.use_mem_eff_path = True + self.add_bias_linear = use_bias + + # MoE config + self.n_routed_experts = n_routed_experts + self.n_shared_experts = n_shared_experts + self.moe_intermediate_size = moe_intermediate_size + self.moe_shared_expert_intermediate_size = moe_shared_expert_intermediate_size + self.moe_latent_size = moe_latent_size + self.num_experts_per_tok = num_experts_per_tok + self.routed_scaling_factor = routed_scaling_factor + self.n_group = n_group + self.topk_group = topk_group + self.norm_topk_prob = norm_topk_prob + + # General training config + self.use_bias = use_bias + self.initializer_range = initializer_range + self.layer_norm_epsilon = layer_norm_epsilon + self.residual_in_fp32 = residual_in_fp32 + self.rescale_prenorm_residual = rescale_prenorm_residual + + # PrimeRL training features + self.load_balance_coeff = load_balance_coeff + self.use_grouped_mm = use_grouped_mm + + super().__init__( + pad_token_id=pad_token_id, + bos_token_id=bos_token_id, + eos_token_id=eos_token_id, + tie_word_embeddings=tie_word_embeddings, + **kwargs, + ) + + @property + def num_hidden_layers(self) -> int: + return len(self.layers_block_type) + + @num_hidden_layers.setter + def num_hidden_layers(self, value): + if value is not None and value < len(self.layers_block_type): + self.layers_block_type = self.layers_block_type[:value] + + +__all__ = ["NemotronHConfig"] diff --git a/src/prime_rl/trainer/models/nemotron_h/converting_nemotron_h.py b/src/prime_rl/trainer/models/nemotron_h/converting_nemotron_h.py new file mode 100644 index 00000000000..e14993ba565 --- /dev/null +++ b/src/prime_rl/trainer/models/nemotron_h/converting_nemotron_h.py @@ -0,0 +1,176 @@ +"""Weight conversion between HuggingFace and PrimeRL formats for NemotronH. + +HF NemotronH uses a unified `mixer` attribute for all layer types: + - Mamba layers: backbone.layers.{i}.mixer.{in_proj, conv1d, ...} + - Attention layers: backbone.layers.{i}.mixer.{q_proj, k_proj, v_proj, o_proj} + - MoE layers: backbone.layers.{i}.mixer.{gate, experts, shared_experts, fc1_latent_proj, fc2_latent_proj} + +PrimeRL separates these into distinct namespaces: + - Mamba layers: model.layers.{i}.mamba.* + - Attention layers: model.layers.{i}.self_attn.* + - MoE layers: model.layers.{i}.mlp.{router, experts, shared_expert, fc1_latent_proj, fc2_latent_proj} + +Global renames: + - HF: backbone.embeddings.weight <-> PrimeRL: model.embed_tokens.weight + - HF: backbone.norm_f.weight <-> PrimeRL: model.norm.weight + - HF uses "backbone." prefix, PrimeRL uses "model." prefix + - HF mtp.* keys (multi-token prediction) are dropped during conversion +""" + +import torch +from torch import Tensor + + +def _rename_keys(state_dict: dict[str, Tensor], old_prefix: str, new_prefix: str): + """Rename all keys matching old_prefix to new_prefix in-place.""" + keys_to_rename = [k for k in state_dict if k.startswith(old_prefix)] + for key in keys_to_rename: + new_key = new_prefix + key[len(old_prefix) :] + state_dict[new_key] = state_dict.pop(key) + + +def convert_hf_layer_to_prime(state_dict: dict[str, Tensor], layer_idx: int, layer_type: str): + """Convert a single layer from HF to PrimeRL format in-place.""" + prefix = f"model.layers.{layer_idx}." + + if layer_type == "moe": + _convert_hf_moe_layer_to_prime(state_dict, prefix) + elif layer_type == "attention": + _convert_hf_attention_layer_to_prime(state_dict, prefix) + elif layer_type == "mamba": + _rename_keys(state_dict, f"{prefix}mixer.", f"{prefix}mamba.") + + +def _convert_hf_moe_layer_to_prime(state_dict: dict[str, Tensor], prefix: str): + """Convert MoE layer: mixer.gate -> mlp.router, mixer.experts -> mlp.experts, etc.""" + mixer = f"{prefix}mixer." + mlp = f"{prefix}mlp." + + # Router: gate.weight -> router.gate (nn.Parameter), gate.e_score_correction_bias -> router.e_score_correction_bias + if f"{mixer}gate.weight" in state_dict: + state_dict[f"{mlp}router.gate"] = state_dict.pop(f"{mixer}gate.weight") + if f"{mixer}gate.e_score_correction_bias" in state_dict: + bias = state_dict.pop(f"{mixer}gate.e_score_correction_bias") + # Shift by per-row min so values land in a range bf16 can represent without + # collapsing. NemotronH biases sit ~57 with ~0.04 inter-expert spread, well + # under bf16's ~0.4 step at that magnitude. topk(scores + bias) is invariant + # to a constant shift of bias, and the returned routing weights are gathered + # from raw sigmoid scores (not biased), so this does not change routing math. + bias = bias - bias.min() + state_dict[f"{mlp}router.e_score_correction_bias"] = bias + + # Experts: check if stored as individual weights (experts.{i}.up_proj.weight) + # or fused 3D tensors (experts.up_proj) + individual_keys = [ + k + for k in state_dict + if k.startswith(f"{mixer}experts.") and k[len(f"{mixer}experts.") :].split(".")[0].isdigit() + ] + + if individual_keys: + expert_indices = sorted({int(k[len(f"{mixer}experts.") :].split(".")[0]) for k in individual_keys}) + + up_projs = [state_dict.pop(f"{mixer}experts.{i}.up_proj.weight") for i in expert_indices] + state_dict[f"{mlp}experts.w1"] = torch.stack(up_projs) + + down_projs = [state_dict.pop(f"{mixer}experts.{i}.down_proj.weight") for i in expert_indices] + state_dict[f"{mlp}experts.w2"] = torch.stack(down_projs) + else: + # Fused 3D tensors + if f"{mixer}experts.up_proj" in state_dict: + state_dict[f"{mlp}experts.w1"] = state_dict.pop(f"{mixer}experts.up_proj") + if f"{mixer}experts.down_proj" in state_dict: + state_dict[f"{mlp}experts.w2"] = state_dict.pop(f"{mixer}experts.down_proj") + + # Dummy w3 required by @expert_parallel decorator compatibility + device = state_dict[f"{mlp}experts.w1"].device if f"{mlp}experts.w1" in state_dict else "cpu" + state_dict[f"{mlp}experts.w3"] = torch.empty(0, device=device) + + # Shared expert + _rename_keys(state_dict, f"{mixer}shared_experts.", f"{mlp}shared_expert.") + + # Latent projections + _rename_keys(state_dict, f"{mixer}fc1_latent_proj.", f"{mlp}fc1_latent_proj.") + _rename_keys(state_dict, f"{mixer}fc2_latent_proj.", f"{mlp}fc2_latent_proj.") + + +def _convert_hf_attention_layer_to_prime(state_dict: dict[str, Tensor], prefix: str): + """Convert attention layer: mixer.{q,k,v,o}_proj -> self_attn.{q,k,v,o}_proj.""" + _rename_keys(state_dict, f"{prefix}mixer.", f"{prefix}self_attn.") + + +def convert_prime_layer_to_hf(state_dict: dict[str, Tensor], layer_idx: int, layer_type: str): + """Convert a single layer from PrimeRL to HF format in-place.""" + prefix = f"model.layers.{layer_idx}." + + if layer_type == "moe": + _convert_prime_moe_layer_to_hf(state_dict, prefix) + elif layer_type == "attention": + _rename_keys(state_dict, f"{prefix}self_attn.", f"{prefix}mixer.") + elif layer_type == "mamba": + _rename_keys(state_dict, f"{prefix}mamba.", f"{prefix}mixer.") + + +def _convert_prime_moe_layer_to_hf(state_dict: dict[str, Tensor], prefix: str): + """Convert MoE layer back to HF format.""" + mlp = f"{prefix}mlp." + mixer = f"{prefix}mixer." + + # Router + if f"{mlp}router.gate" in state_dict: + state_dict[f"{mixer}gate.weight"] = state_dict.pop(f"{mlp}router.gate") + if f"{mlp}router.e_score_correction_bias" in state_dict: + state_dict[f"{mixer}gate.e_score_correction_bias"] = state_dict.pop(f"{mlp}router.e_score_correction_bias") + + # Experts: unstack fused 3D tensors into individual expert weights + if f"{mlp}experts.w1" in state_dict: + w1 = state_dict.pop(f"{mlp}experts.w1") + for i in range(w1.shape[0]): + state_dict[f"{mixer}experts.{i}.up_proj.weight"] = w1[i] + if f"{mlp}experts.w2" in state_dict: + w2 = state_dict.pop(f"{mlp}experts.w2") + for i in range(w2.shape[0]): + state_dict[f"{mixer}experts.{i}.down_proj.weight"] = w2[i] + # Remove dummy w3 (not present in HF format) + state_dict.pop(f"{mlp}experts.w3", None) + + # Shared expert + _rename_keys(state_dict, f"{mlp}shared_expert.", f"{mixer}shared_experts.") + + # Latent projections + _rename_keys(state_dict, f"{mlp}fc1_latent_proj.", f"{mixer}fc1_latent_proj.") + _rename_keys(state_dict, f"{mlp}fc2_latent_proj.", f"{mixer}fc2_latent_proj.") + + +def convert_hf_to_prime(state_dict: dict[str, Tensor], layers_block_type: list[str]): + """Convert full model from HF to PrimeRL format in-place.""" + # Handle backbone.* -> model.* prefix (HF NemotronH uses "backbone" instead of "model") + _rename_keys(state_dict, "backbone.", "model.") + + # Drop multi-token prediction keys (not used in training) + for key in [k for k in state_dict if k.startswith("mtp.")]: + del state_dict[key] + + # Global renames + if "model.embeddings.weight" in state_dict: + state_dict["model.embed_tokens.weight"] = state_dict.pop("model.embeddings.weight") + if "model.norm_f.weight" in state_dict: + state_dict["model.norm.weight"] = state_dict.pop("model.norm_f.weight") + + for i, layer_type in enumerate(layers_block_type): + convert_hf_layer_to_prime(state_dict, i, layer_type) + + +def convert_prime_to_hf(state_dict: dict[str, Tensor], layers_block_type: list[str]): + """Convert full model from PrimeRL to HF format in-place.""" + # Global renames + if "model.embed_tokens.weight" in state_dict: + state_dict["model.embeddings.weight"] = state_dict.pop("model.embed_tokens.weight") + if "model.norm.weight" in state_dict: + state_dict["model.norm_f.weight"] = state_dict.pop("model.norm.weight") + + for i, layer_type in enumerate(layers_block_type): + convert_prime_layer_to_hf(state_dict, i, layer_type) + + # Rename model.* -> backbone.* (HF checkpoint uses "backbone" prefix) + _rename_keys(state_dict, "model.", "backbone.") diff --git a/src/prime_rl/trainer/models/nemotron_h/modeling_nemotron_h.py b/src/prime_rl/trainer/models/nemotron_h/modeling_nemotron_h.py new file mode 100644 index 00000000000..d9c5ec1ede0 --- /dev/null +++ b/src/prime_rl/trainer/models/nemotron_h/modeling_nemotron_h.py @@ -0,0 +1,583 @@ +"""PrimeRL implementation of NemotronH (Nemotron-3-Super-120B-A12B). + +Hybrid Mamba-Transformer-MoE architecture with three distinct layer types: +- Mamba-2 layers (using NemotronHMamba2Mixer from HF transformers) +- LatentMoE layers (non-gated experts with latent projections) +- Attention layers (using shared FlashAttention/SDPA from prime-rl) +""" + +from typing import Optional + +import torch +import torch.distributed as dist +from torch import Tensor, nn +from transformers.generation import GenerationMixin +from transformers.modeling_layers import GradientCheckpointingLayer +from transformers.modeling_outputs import BaseModelOutputWithPast +from transformers.models.nemotron_h.modular_nemotron_h import NemotronHMamba2Mixer +from transformers.utils import auto_docstring, logging + +from prime_rl.trainer.models.base import PreTrainedModelPrimeRL +from prime_rl.trainer.models.layers.attn import ATTN_IMPL2CLASS, AttentionConfig +from prime_rl.trainer.models.layers.cp_mamba import mamba_cp_forward +from prime_rl.trainer.models.layers.lm_head import PrimeLmOutput +from prime_rl.trainer.models.layers.moe import LatentMoE, NemotronHRouter, NonGatedGroupedExperts +from prime_rl.trainer.models.layers.rms_norm import RMSNorm, RMSNormConfig +from prime_rl.trainer.models.nemotron_h.configuration_nemotron_h import NemotronHConfig +from prime_rl.trainer.models.nemotron_h.converting_nemotron_h import ( + convert_hf_layer_to_prime, + convert_hf_to_prime, + convert_prime_layer_to_hf, + convert_prime_to_hf, +) +from prime_rl.utils.sequence import get_cu_seqlens_from_position_ids + +logger = logging.get_logger(__name__) + +_patch_applied = False + + +def _patch_mamba2_use_triton_ssd(): + """Patch NemotronHMamba2Mixer to use mamba_ssm Triton SSD kernels. + + The HF torch_forward computes softplus in model dtype (bf16), while vLLM's + Triton kernels and mamba_ssm use fp32. This causes ~0.4 KL divergence. + + This patch makes the mixer use mamba_chunk_scan_combined (Triton, fp32 + softplus) for the SSD computation, with PyTorch nn.Conv1d for convolution. + Requires mamba_ssm installed; causal_conv1d should NOT be installed (it + needs arch-specific CUDA compilation). The HF cuda_kernels_forward + already falls back to PyTorch nn.Conv1d when causal_conv1d is absent. + """ + global _patch_applied + if _patch_applied: + return + + if not torch.cuda.is_available(): + logger.warning_once("CUDA not available; NemotronH Mamba layers will use torch_forward (bf16 softplus)") + return + + try: + from mamba_ssm.ops.triton.ssd_combined import ( + mamba_chunk_scan_combined as _mamba_chunk_scan_combined, + ) + except ImportError: + logger.warning_once("mamba_ssm not installed; NemotronH Mamba layers will use torch_forward (bf16 softplus)") + return + + if _mamba_chunk_scan_combined is None: + logger.warning_once("mamba_ssm not available; NemotronH Mamba layers will use torch_forward (bf16 softplus)") + return + + def _varlen_mamba_forward(self, hidden_states, cu_seqlens): + """Varlen-aware mamba forward. `cu_seqlens` is required: shape [num_seq+1], + e.g. [0, s1, s1+s2, ...]. Handles packed batches by: + - applying conv1d per-sequence (with kernel-1 leading zeros) + - passing seq_idx to mamba_chunk_scan_combined (SSM state resets at boundaries) + """ + batch_size, seq_len, _ = hidden_states.shape + assert batch_size == 1, "varlen mamba expects batch_size=1 (packed)" + groups_time_state_size = self.n_groups * self.ssm_state_size + + # 1. Input projection (pointwise, no cross-seq mixing) + projected_states = self.in_proj(hidden_states) + A = -torch.exp(self.A_log.float()) + dt_limit_kwargs = {} if self.time_step_limit is None else {"dt_limit": self.time_step_limit} + + gate, hidden_states_B_C, time_step = torch.split( + projected_states, + [self.intermediate_size, self.conv_dim, self.num_heads], + dim=-1, + ) + + # 2. Per-sequence causal conv1d. self.conv1d already has padding=kernel-1 so + # applying it to each segment independently gives the correct causal output + # without cross-sequence history. + hbc_t = hidden_states_B_C.transpose(1, 2) # (1, C, L) + cu = cu_seqlens.tolist() + conv_outs = [] + for i in range(len(cu) - 1): + s, e = cu[i], cu[i + 1] + if s == e: + continue + seg = hbc_t[:, :, s:e] + conv_out = self.conv1d(seg)[:, :, : e - s] + conv_outs.append(conv_out) + hidden_states_B_C = self.act(torch.cat(conv_outs, dim=-1).transpose(1, 2)) + + hidden_states, B, C = torch.split( + hidden_states_B_C, + [self.intermediate_size, groups_time_state_size, groups_time_state_size], + dim=-1, + ) + + # 3. SSM with seq_idx so state resets at sequence boundaries + seq_idx = torch.zeros(seq_len, dtype=torch.int32, device=hidden_states.device) + for i in range(len(cu) - 1): + s, e = cu[i], cu[i + 1] + if e > s: + seq_idx[s:e] = i + seq_idx = seq_idx.unsqueeze(0) # (1, L) + + scan_output = _mamba_chunk_scan_combined( + hidden_states.view(batch_size, seq_len, -1, self.head_dim), + time_step, + A, + B.view(batch_size, seq_len, self.n_groups, -1), + C.view(batch_size, seq_len, self.n_groups, -1), + chunk_size=self.chunk_size, + D=self.D, + z=None, + seq_idx=seq_idx, + return_final_states=False, + dt_bias=self.dt_bias, + dt_softplus=True, + **dt_limit_kwargs, + ) + scan_output = scan_output.view(batch_size, seq_len, -1) + scan_output = self.norm(scan_output, gate) + return self.out_proj(scan_output) + + def _patched_forward(self, hidden_states, cache_params=None, attention_mask=None, cu_seqlens=None): + if cu_seqlens is not None and cache_params is None and "cuda" in self.in_proj.weight.device.type: + return _varlen_mamba_forward(self, hidden_states, cu_seqlens) + if "cuda" in self.in_proj.weight.device.type: + # Disable fused training path (needs causal_conv1d CUDA extension) + orig = self.use_mem_eff_path + self.use_mem_eff_path = False + result = self.cuda_kernels_forward(hidden_states, cache_params, attention_mask) + self.use_mem_eff_path = orig + return result + return self.torch_forward(hidden_states, cache_params, attention_mask) + + NemotronHMamba2Mixer.forward = _patched_forward + _patch_applied = True + logger.info("Patched NemotronHMamba2Mixer: Triton SSD kernels + varlen (cu_seqlens) support") + + +def _ensure_zamba2_compat(config: NemotronHConfig): + """Add Zamba2-compatible attribute aliases to NemotronHConfig. + + The HF modular NemotronHMamba2Mixer inherits from Zamba2MambaMixer, whose + __init__ reads Zamba2-style attribute names before the NemotronH child + overrides them. We set the missing aliases so the parent __init__ doesn't + crash. + + Critically, ``mamba_expand`` must give ``mamba_num_heads * mamba_head_dim`` + when multiplied by ``hidden_size``; the config's ``expand`` field does not + satisfy this for all model sizes (e.g. Nemotron-3-Nano-30B). + """ + correct_intermediate = config.mamba_num_heads * config.mamba_head_dim + correct_expand = correct_intermediate / config.hidden_size + + aliases = { + "mamba_d_state": config.ssm_state_size, + "mamba_d_conv": config.conv_kernel, + "mamba_ngroups": config.n_groups, + "mamba_headdim": config.mamba_head_dim, + "n_mamba_heads": config.mamba_num_heads, + "add_bias_linear": getattr(config, "use_bias", False), + "use_mem_eff_path": True, + } + for attr, value in aliases.items(): + if not hasattr(config, attr): + setattr(config, attr, value) + + config.mamba_expand = correct_expand + + +class NemotronHMambaLayer(GradientCheckpointingLayer): + """Mamba-2 SSM layer: norm -> NemotronHMamba2Mixer -> residual. + + With context parallelism, uses all-to-all head partitioning: transposes + from sequence-parallel to head-parallel before the SSM, so each CP rank + processes the full sequence on a subset of heads, then transposes back. + """ + + def __init__(self, config: NemotronHConfig, layer_idx: int): + super().__init__() + _patch_mamba2_use_triton_ssd() + _ensure_zamba2_compat(config) + self.norm = RMSNorm(RMSNormConfig(hidden_size=config.hidden_size, eps=config.layer_norm_epsilon)) + self.mamba = NemotronHMamba2Mixer(config, layer_idx=layer_idx) + self.mlp = None # No MoE in this layer type + + def set_context_parallel_attributes(self, cp_group: dist.ProcessGroup, cp_rank: int, cp_world_size: int) -> None: + assert self.mamba.num_heads % cp_world_size == 0, ( + f"num_heads ({self.mamba.num_heads}) must be divisible by cp_world_size ({cp_world_size})" + ) + assert self.mamba.n_groups % cp_world_size == 0, ( + f"n_groups ({self.mamba.n_groups}) must be divisible by cp_world_size ({cp_world_size})" + ) + self._cp_group = cp_group + self._cp_rank = cp_rank + self._cp_world_size = cp_world_size + + @property + def cp_enabled(self) -> bool: + return hasattr(self, "_cp_group") + + def forward( + self, + hidden_states: torch.Tensor, + position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None, + cu_seqlens: torch.LongTensor | None = None, + max_seqlen: int | None = None, + ) -> torch.Tensor: + residual = hidden_states + hidden_states = self.norm(hidden_states) + + if self.cp_enabled: + # TODO: This path doesnt support cu_seqlens so packing makes it wrong + hidden_states = mamba_cp_forward( + self.mamba, hidden_states, self._cp_group, self._cp_rank, self._cp_world_size + ) + else: + hidden_states = self.mamba(hidden_states, cu_seqlens=cu_seqlens) + + return residual + hidden_states + + +class NemotronHMoELayer(GradientCheckpointingLayer): + """MoE layer: norm -> LatentMoE -> residual.""" + + def __init__(self, config: NemotronHConfig): + super().__init__() + self.norm = RMSNorm(RMSNormConfig(hidden_size=config.hidden_size, eps=config.layer_norm_epsilon)) + self.mlp = LatentMoE( + dim=config.hidden_size, + latent_dim=config.moe_latent_size, + moe_intermediate_size=config.moe_intermediate_size, + shared_expert_intermediate_size=config.moe_shared_expert_intermediate_size, + num_experts=config.n_routed_experts, + top_k=config.num_experts_per_tok, + n_group=config.n_group, + topk_group=config.topk_group, + norm_topk_prob=config.norm_topk_prob, + routed_scaling_factor=config.routed_scaling_factor, + use_grouped_mm=config.use_grouped_mm, + load_balance_coeff=config.load_balance_coeff, + fp8=getattr(config, "fp8", False), + ) + + def forward( + self, + hidden_states: torch.Tensor, + position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None, + cu_seqlens: torch.LongTensor | None = None, + max_seqlen: int | None = None, + ) -> torch.Tensor: + residual = hidden_states + hidden_states = self.norm(hidden_states) + hidden_states = self.mlp(hidden_states) + return residual + hidden_states + + +class NemotronHAttentionLayer(GradientCheckpointingLayer): + """Attention layer: norm -> FlashAttention/SDPA -> residual.""" + + def __init__(self, config: NemotronHConfig): + super().__init__() + self.norm = RMSNorm(RMSNormConfig(hidden_size=config.hidden_size, eps=config.layer_norm_epsilon)) + attn_config = AttentionConfig( + hidden_size=config.hidden_size, + head_dim=config.head_dim, + num_attention_heads=config.num_attention_heads, + num_key_value_heads=config.num_key_value_heads, + is_causal=True, + attention_bias=config.attention_bias, + use_qk_norm=False, + rms_norm_eps=config.layer_norm_epsilon, + ) + self.self_attn = ATTN_IMPL2CLASS[config._attn_implementation](attn_config) + self.mlp = None # No MoE in this layer type + + def forward( + self, + hidden_states: torch.Tensor, + position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None, + cu_seqlens: torch.LongTensor | None = None, + max_seqlen: int | None = None, + ) -> torch.Tensor: + residual = hidden_states + hidden_states = self.norm(hidden_states) + hidden_states, _ = self.self_attn( + hidden_states=hidden_states, + position_embeddings=position_embeddings, + cu_seqlens=cu_seqlens, + max_seqlen=max_seqlen, + ) + return residual + hidden_states + + +BLOCK_TYPE_MAP = { + "mamba": NemotronHMambaLayer, + "moe": NemotronHMoELayer, + "attention": NemotronHAttentionLayer, +} + + +def _build_layer(config: NemotronHConfig, layer_idx: int): + layer_type = config.layers_block_type[layer_idx] + cls = BLOCK_TYPE_MAP[layer_type] + if layer_type == "mamba": + return cls(config, layer_idx=layer_idx) + return cls(config) + + +@auto_docstring +class NemotronHPreTrainedModel(PreTrainedModelPrimeRL): + config: NemotronHConfig + base_model_prefix = "model" + supports_gradient_checkpointing = True + _no_split_modules = ["NemotronHMambaLayer", "NemotronHMoELayer", "NemotronHAttentionLayer"] + _supports_flash_attn = True + _supports_sdpa = True + _can_compile_fullgraph = False + + def _init_weights(self, module): + std = self.config.initializer_range + if isinstance(module, nn.Linear): + module.weight.data.normal_(mean=0.0, std=std) + if module.bias is not None: + module.bias.data.zero_() + elif isinstance(module, nn.Embedding): + module.weight.data.normal_(mean=0.0, std=std) + if module.padding_idx is not None: + module.weight.data[module.padding_idx].zero_() + elif isinstance(module, (NonGatedGroupedExperts, NemotronHRouter)): + module.init_weights(std) + + @classmethod + def is_hf_state_dict(cls, state_dict: dict[str, Tensor]) -> bool: + return any("mixer." in name for name in state_dict) or any("backbone." in name for name in state_dict) + + @classmethod + def is_prime_state_dict(cls, state_dict: dict[str, Tensor]) -> bool: + return any( + "mamba." in name + or "mlp.experts.w1" in name + or "self_attn." in name + or "model.embed_tokens." in name + or "model.norm." in name + for name in state_dict + ) + + @classmethod + def convert_to_hf(cls, state_dict: dict[str, Tensor]) -> dict[str, Tensor]: + # Need config to know layer types; infer from state dict keys + layers_block_type = _infer_layers_block_type(state_dict) + convert_prime_to_hf(state_dict, layers_block_type) + return state_dict + + @classmethod + def convert_to_prime(cls, state_dict: dict[str, Tensor]) -> dict[str, Tensor]: + layers_block_type = _infer_layers_block_type_from_hf(state_dict) + convert_hf_to_prime(state_dict, layers_block_type) + return state_dict + + @classmethod + def convert_adapter_to_hf(cls, state_dict: dict[str, Tensor]) -> dict[str, Tensor]: + # HF NemotronH unifies attention/mlp under a single `mixer` attribute per layer. + import re + + rename = [ + (re.compile(r"(\.layers\.\d+)\.mlp\."), r"\1.mixer."), + (re.compile(r"(\.layers\.\d+)\.self_attn\."), r"\1.mixer."), + ] + for old_key in list(state_dict.keys()): + new_key = old_key + for pattern, repl in rename: + new_key = pattern.sub(repl, new_key) + if new_key != old_key: + state_dict[new_key] = state_dict.pop(old_key) + return state_dict + + @classmethod + def convert_layer_to_hf(cls, state_dict: dict[str, Tensor], layer_idx: int) -> dict[str, Tensor]: + from prime_rl.trainer.models.nemotron_h.converting_nemotron_h import _rename_keys + + if layer_idx == -1: + # Non-layer weights: rename global keys and model.* -> backbone.* + if "model.embed_tokens.weight" in state_dict: + state_dict["model.embeddings.weight"] = state_dict.pop("model.embed_tokens.weight") + if "model.norm.weight" in state_dict: + state_dict["model.norm_f.weight"] = state_dict.pop("model.norm.weight") + _rename_keys(state_dict, "model.", "backbone.") + else: + layer_type = _infer_layer_type_prime(state_dict, layer_idx) + convert_prime_layer_to_hf(state_dict, layer_idx, layer_type) + _rename_keys(state_dict, "model.layers.", "backbone.layers.") + return state_dict + + @classmethod + def convert_layer_to_prime(cls, state_dict: dict[str, Tensor], layer_idx: int) -> dict[str, Tensor]: + from prime_rl.trainer.models.nemotron_h.converting_nemotron_h import _rename_keys + + # Handle backbone.* -> model.* prefix before layer conversion + _rename_keys(state_dict, "backbone.", "model.") + + if layer_idx == -1: + # Non-layer weights: rename global keys + if "model.embeddings.weight" in state_dict: + state_dict["model.embed_tokens.weight"] = state_dict.pop("model.embeddings.weight") + if "model.norm_f.weight" in state_dict: + state_dict["model.norm.weight"] = state_dict.pop("model.norm_f.weight") + else: + layer_type = _infer_layer_type_hf(state_dict, layer_idx) + convert_hf_layer_to_prime(state_dict, layer_idx, layer_type) + return state_dict + + +def _infer_layer_type_hf(state_dict: dict[str, Tensor], layer_idx: int) -> str: + """Infer layer type from HF state dict keys for a given layer.""" + # HF checkpoints may use either "model." or "backbone." prefix + for root in ("model", "backbone"): + prefix = f"{root}.layers.{layer_idx}.mixer." + layer_keys = [k for k in state_dict if k.startswith(prefix)] + if layer_keys: + for k in layer_keys: + suffix = k[len(prefix) :] + if suffix.startswith("gate.") or suffix.startswith("experts."): + return "moe" + if suffix.startswith("q_proj") or suffix.startswith("k_proj"): + return "attention" + return "mamba" + return "mamba" # fallback + + +def _infer_layer_type_prime(state_dict: dict[str, Tensor], layer_idx: int) -> str: + """Infer layer type from PrimeRL state dict keys for a given layer.""" + prefix = f"model.layers.{layer_idx}." + for k in state_dict: + if not k.startswith(prefix): + continue + suffix = k[len(prefix) :] + if suffix.startswith("mlp."): + return "moe" + if suffix.startswith("self_attn."): + return "attention" + if suffix.startswith("mamba."): + return "mamba" + return "mamba" + + +def _infer_layers_block_type_from_hf(state_dict: dict[str, Tensor]) -> list[str]: + """Infer full layers_block_type list from HF state dict.""" + # HF checkpoints may use either "model." or "backbone." prefix + layer_keys = [k for k in state_dict if k.startswith("model.layers.") or k.startswith("backbone.layers.")] + max_layer = max(int(k.split(".")[2]) for k in layer_keys) + 1 + return [_infer_layer_type_hf(state_dict, i) for i in range(max_layer)] + + +def _infer_layers_block_type(state_dict: dict[str, Tensor]) -> list[str]: + """Infer full layers_block_type list from PrimeRL state dict.""" + max_layer = max(int(k.split(".")[2]) for k in state_dict if k.startswith("model.layers.")) + 1 + return [_infer_layer_type_prime(state_dict, i) for i in range(max_layer)] + + +class NemotronHModel(NemotronHPreTrainedModel): + def __init__(self, config: NemotronHConfig): + super().__init__(config) + self.padding_idx = config.pad_token_id + self.vocab_size = config.vocab_size + + self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx) + self.layers = nn.ModuleList([_build_layer(config, i) for i in range(config.num_hidden_layers)]) + self.norm = RMSNorm(RMSNormConfig(hidden_size=config.hidden_size, eps=config.layer_norm_epsilon)) + + # NemotronH does not use RoPE - position information comes from Mamba layers + self.rotary_emb = None + + self.gradient_checkpointing = False + self.post_init() + + @auto_docstring + def forward( + self, + input_ids: Optional[torch.LongTensor] = None, + position_ids: Optional[torch.LongTensor] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + ) -> BaseModelOutputWithPast: + if (input_ids is None) ^ (inputs_embeds is not None): + raise ValueError("You must specify exactly one of input_ids or inputs_embeds") + + if inputs_embeds is None: + inputs_embeds = self.embed_tokens(input_ids) + + # Compute cu_seqlens and max_seqlen for flash attention + if self.config._attn_implementation in ("flash_attention_2", "flash_attention_3", "fa4"): + cu_seqlens, max_seqlen = get_cu_seqlens_from_position_ids(position_ids) + torch._dynamo.mark_dynamic(cu_seqlens, 0) + else: + max_seqlen = None + cu_seqlens = None + + hidden_states = inputs_embeds + position_embeddings = self.rotary_emb(hidden_states, position_ids) if self.rotary_emb is not None else None + + for decoder_layer in self.layers: + hidden_states = decoder_layer( + hidden_states, + position_embeddings=position_embeddings, + cu_seqlens=cu_seqlens, + max_seqlen=max_seqlen, + ) + + hidden_states = self.norm(hidden_states) + return BaseModelOutputWithPast(last_hidden_state=hidden_states) + + +@auto_docstring +class NemotronHForCausalLM(NemotronHPreTrainedModel, GenerationMixin): + _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} + + def __init__(self, config: NemotronHConfig): + super().__init__(config) + self.model = NemotronHModel(config) + self.vocab_size = config.vocab_size + self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False) + self.post_init() + + def get_decoder(self): + return self.model + + def forward( + self, + input_ids: Optional[torch.LongTensor] = None, + position_ids: Optional[torch.LongTensor] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + labels: Optional[torch.LongTensor] = None, + logits_to_keep: int = 0, + temperature: Optional[torch.Tensor] = None, + **kwargs, + ) -> PrimeLmOutput: + if position_ids is None: + if inputs_embeds is not None: + position_ids = torch.arange(inputs_embeds.shape[1], device=inputs_embeds.device).unsqueeze(0) + else: + position_ids = torch.arange(input_ids.shape[1], device=input_ids.device).unsqueeze(0) + + outputs = self.model( + input_ids=input_ids, + position_ids=position_ids, + inputs_embeds=inputs_embeds, + ) + + hidden_states = outputs.last_hidden_state + slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep + return self.lm_head( + hidden_states[:, slice_indices, :], + labels[:, slice_indices] if labels is not None else None, + temperature=temperature, + ) + + def init_buffers_post_meta(self): + pass + + +__all__ = [ + "NemotronHForCausalLM", + "NemotronHModel", + "NemotronHPreTrainedModel", +] diff --git a/src/prime_rl/trainer/models/qwen3_5_moe/__init__.py b/src/prime_rl/trainer/models/qwen3_5_moe/__init__.py new file mode 100644 index 00000000000..17573101dd3 --- /dev/null +++ b/src/prime_rl/trainer/models/qwen3_5_moe/__init__.py @@ -0,0 +1,13 @@ +from prime_rl.trainer.models.qwen3_5_moe.configuration_qwen3_5_moe import Qwen3_5MoeConfig +from prime_rl.trainer.models.qwen3_5_moe.modeling_qwen3_5_moe import ( + Qwen3_5MoeForCausalLM, + Qwen3_5MoeModel, + Qwen3_5MoePreTrainedModel, +) + +__all__ = [ + "Qwen3_5MoeConfig", + "Qwen3_5MoeForCausalLM", + "Qwen3_5MoeModel", + "Qwen3_5MoePreTrainedModel", +] diff --git a/src/prime_rl/trainer/models/qwen3_5_moe/configuration_qwen3_5_moe.py b/src/prime_rl/trainer/models/qwen3_5_moe/configuration_qwen3_5_moe.py new file mode 100644 index 00000000000..a00a9e30a28 --- /dev/null +++ b/src/prime_rl/trainer/models/qwen3_5_moe/configuration_qwen3_5_moe.py @@ -0,0 +1,108 @@ +from transformers.configuration_utils import PretrainedConfig + + +class Qwen3_5MoeConfig(PretrainedConfig): + r""" + Configuration class for the custom PrimeRL Qwen3.5-MoE model. + + This matches the HuggingFace `Qwen3_5MoeTextConfig` architecture: + hybrid attention (GatedDeltaNet linear + gated softmax), MoE with gated shared expert, + and (1+weight) RMSNorm parameterization. + + Defaults match Qwen3.5-35B-A3B. + """ + + model_type = "qwen3_5_moe_text" + keys_to_ignore_at_inference = ["past_key_values"] + + def __init__( + self, + vocab_size=248320, + hidden_size=2048, + num_hidden_layers=40, + num_attention_heads=16, + num_key_value_heads=2, + hidden_act="silu", + max_position_embeddings=32768, + initializer_range=0.02, + rms_norm_eps=1e-6, + use_cache=True, + tie_word_embeddings=False, + rope_theta=10000.0, + rope_scaling=None, + attention_bias=False, + attention_dropout=0.0, + head_dim=256, + # Linear attention (GatedDeltaNet) params + linear_conv_kernel_dim=4, + linear_key_head_dim=128, + linear_value_head_dim=128, + linear_num_key_heads=16, + linear_num_value_heads=32, + # MoE params + moe_intermediate_size=512, + shared_expert_intermediate_size=512, + num_experts_per_tok=8, + num_experts=256, + output_router_logits=False, + router_aux_loss_coef=0.001, + # Layer type configuration + layer_types=None, + # PrimeRL additions + load_balance_coeff=None, + use_grouped_mm=True, + pad_token_id=None, + **kwargs, + ): + self.vocab_size = vocab_size + self.max_position_embeddings = max_position_embeddings + self.hidden_size = hidden_size + self.num_hidden_layers = num_hidden_layers + self.num_attention_heads = num_attention_heads + self.num_key_value_heads = num_key_value_heads + self.hidden_act = hidden_act + self.initializer_range = initializer_range + self.rms_norm_eps = rms_norm_eps + self.use_cache = use_cache + self.rope_theta = rope_theta + self.rope_scaling = rope_scaling + self.attention_bias = attention_bias + self.attention_dropout = attention_dropout + self.head_dim = head_dim + self.pad_token_id = pad_token_id + + # Partial RoPE: only rotate 25% of head_dim + kwargs.setdefault("partial_rotary_factor", 0.25) + + # Layer types: default every 4th layer is full_attention, rest linear_attention + self.layer_types = layer_types + if self.layer_types is None: + interval_pattern = kwargs.pop("full_attention_interval", 4) + self.layer_types = [ + "linear_attention" if bool((i + 1) % interval_pattern) else "full_attention" + for i in range(self.num_hidden_layers) + ] + + # Linear attention (GatedDeltaNet) params + self.linear_conv_kernel_dim = linear_conv_kernel_dim + self.linear_key_head_dim = linear_key_head_dim + self.linear_value_head_dim = linear_value_head_dim + self.linear_num_key_heads = linear_num_key_heads + self.linear_num_value_heads = linear_num_value_heads + + # MoE params + self.moe_intermediate_size = moe_intermediate_size + self.shared_expert_intermediate_size = shared_expert_intermediate_size + self.num_experts_per_tok = num_experts_per_tok + self.num_experts = num_experts + self.output_router_logits = output_router_logits + self.router_aux_loss_coef = router_aux_loss_coef + + # PrimeRL additions + self.load_balance_coeff = load_balance_coeff + self.use_grouped_mm = use_grouped_mm + + super().__init__( + tie_word_embeddings=tie_word_embeddings, + **kwargs, + ) diff --git a/src/prime_rl/trainer/models/qwen3_5_moe/converting_qwen3_5_moe.py b/src/prime_rl/trainer/models/qwen3_5_moe/converting_qwen3_5_moe.py new file mode 100644 index 00000000000..24e399a44a7 --- /dev/null +++ b/src/prime_rl/trainer/models/qwen3_5_moe/converting_qwen3_5_moe.py @@ -0,0 +1,125 @@ +import torch +from torch import Tensor + + +def get_max_layer_num(state_dict: dict[str, Tensor]) -> int: + """Get the maximum number of layers in the model.""" + return max(int(i.split(".")[2]) for i in state_dict.keys() if "model.layers." in i) + 1 + + +def convert_hf_layer_to_tt(state_dict: dict[str, Tensor], layer_idx: int): + """Convert a single layer from HF to PrimeRL format in-place. + + Handles: + - Router: mlp.gate.weight -> mlp.router.gate.weight + - Routed experts: fused gate_up_proj/down_proj or per-expert -> w1/w2/w3 + - Shared expert: mlp.shared_expert.{gate,up,down}_proj.weight -> shared_expert.{w1,w3,w2}.weight + - Shared expert gate: mlp.shared_expert_gate.weight -> shared_expert_gate.weight + """ + i = layer_idx + + # Router: mlp.gate.weight -> mlp.router.gate.weight + gate_key = f"model.layers.{i}.mlp.gate.weight" + if gate_key not in state_dict: + return + + state_dict[f"model.layers.{i}.mlp.router.gate.weight"] = state_dict.pop(gate_key) + + # Routed experts: convert to fused w1/w2/w3 format + if f"model.layers.{i}.mlp.experts.gate_up_proj" in state_dict: + # New fused format (transformers 5.0+): gate_up_proj shape (num_experts, 2*moe_dim, dim) + gate_up_proj = state_dict.pop(f"model.layers.{i}.mlp.experts.gate_up_proj") + down_proj = state_dict.pop(f"model.layers.{i}.mlp.experts.down_proj") + + moe_dim = gate_up_proj.shape[1] // 2 + w1 = gate_up_proj[:, :moe_dim, :] # gate + w3 = gate_up_proj[:, moe_dim:, :] # up + w2 = down_proj # down + else: + # Old per-expert format + num_experts = len([k for k in state_dict.keys() if f"model.layers.{i}.mlp.experts" in k and "gate_proj" in k]) + if num_experts == 0: + return + + dim, moe_dim = state_dict[f"model.layers.{i}.mlp.experts.0.down_proj.weight"].shape + dtype = state_dict[f"model.layers.{i}.mlp.experts.0.down_proj.weight"].dtype + w1 = torch.empty((num_experts, moe_dim, dim), dtype=dtype) + w2 = torch.empty((num_experts, dim, moe_dim), dtype=dtype) + w3 = torch.empty((num_experts, moe_dim, dim), dtype=dtype) + for j in range(num_experts): + w1[j].copy_(state_dict.pop(f"model.layers.{i}.mlp.experts.{j}.gate_proj.weight")) + w2[j].copy_(state_dict.pop(f"model.layers.{i}.mlp.experts.{j}.down_proj.weight")) + w3[j].copy_(state_dict.pop(f"model.layers.{i}.mlp.experts.{j}.up_proj.weight")) + + state_dict[f"model.layers.{i}.mlp.experts.w1"] = w1 + state_dict[f"model.layers.{i}.mlp.experts.w2"] = w2 + state_dict[f"model.layers.{i}.mlp.experts.w3"] = w3 + + # Shared expert: mlp.shared_expert.{gate,up,down}_proj -> shared_expert.{w1,w3,w2} + se_gate_key = f"model.layers.{i}.mlp.shared_expert.gate_proj.weight" + if se_gate_key in state_dict: + state_dict[f"model.layers.{i}.shared_expert.w1.weight"] = state_dict.pop(se_gate_key) + state_dict[f"model.layers.{i}.shared_expert.w2.weight"] = state_dict.pop( + f"model.layers.{i}.mlp.shared_expert.down_proj.weight" + ) + state_dict[f"model.layers.{i}.shared_expert.w3.weight"] = state_dict.pop( + f"model.layers.{i}.mlp.shared_expert.up_proj.weight" + ) + + # Shared expert gate: mlp.shared_expert_gate.weight -> shared_expert_gate.weight + seg_key = f"model.layers.{i}.mlp.shared_expert_gate.weight" + if seg_key in state_dict: + state_dict[f"model.layers.{i}.shared_expert_gate.weight"] = state_dict.pop(seg_key) + + +def convert_tt_layer_to_hf(state_dict: dict[str, Tensor], layer_idx: int): + """Convert a single layer from PrimeRL to HF format in-place.""" + i = layer_idx + + # Router + router_key = f"model.layers.{i}.mlp.router.gate.weight" + if router_key not in state_dict: + return + + state_dict[f"model.layers.{i}.mlp.gate.weight"] = state_dict.pop(router_key) + + # Routed experts: w1/w2/w3 -> per-expert format + w1 = state_dict.pop(f"model.layers.{i}.mlp.experts.w1") + w2 = state_dict.pop(f"model.layers.{i}.mlp.experts.w2") + w3 = state_dict.pop(f"model.layers.{i}.mlp.experts.w3") + + num_experts = w1.shape[0] + for j in range(num_experts): + state_dict[f"model.layers.{i}.mlp.experts.{j}.gate_proj.weight"] = w1[j] + state_dict[f"model.layers.{i}.mlp.experts.{j}.down_proj.weight"] = w2[j] + state_dict[f"model.layers.{i}.mlp.experts.{j}.up_proj.weight"] = w3[j] + + # Shared expert: shared_expert.{w1,w2,w3} -> mlp.shared_expert.{gate,down,up}_proj + se_w1_key = f"model.layers.{i}.shared_expert.w1.weight" + if se_w1_key in state_dict: + state_dict[f"model.layers.{i}.mlp.shared_expert.gate_proj.weight"] = state_dict.pop(se_w1_key) + state_dict[f"model.layers.{i}.mlp.shared_expert.down_proj.weight"] = state_dict.pop( + f"model.layers.{i}.shared_expert.w2.weight" + ) + state_dict[f"model.layers.{i}.mlp.shared_expert.up_proj.weight"] = state_dict.pop( + f"model.layers.{i}.shared_expert.w3.weight" + ) + + # Shared expert gate + seg_key = f"model.layers.{i}.shared_expert_gate.weight" + if seg_key in state_dict: + state_dict[f"model.layers.{i}.mlp.shared_expert_gate.weight"] = state_dict.pop(seg_key) + + +def convert_hf_to_tt_moe(state_dict: dict[str, Tensor]): + """Convert all MoE weights from HF to PrimeRL format in-place.""" + num_layers = get_max_layer_num(state_dict) + for i in range(num_layers): + convert_hf_layer_to_tt(state_dict, i) + + +def convert_tt_to_hf_moe(state_dict: dict[str, Tensor]): + """Convert all MoE weights from PrimeRL to HF format in-place.""" + num_layers = get_max_layer_num(state_dict) + for i in range(num_layers): + convert_tt_layer_to_hf(state_dict, i) diff --git a/src/prime_rl/trainer/models/qwen3_5_moe/modeling_qwen3_5_moe.py b/src/prime_rl/trainer/models/qwen3_5_moe/modeling_qwen3_5_moe.py new file mode 100644 index 00000000000..a03713a3af8 --- /dev/null +++ b/src/prime_rl/trainer/models/qwen3_5_moe/modeling_qwen3_5_moe.py @@ -0,0 +1,1065 @@ +import functools +from dataclasses import dataclass +from typing import Optional, Union + +import torch +import torch.nn.functional as F +from torch import Tensor, nn +from transformers.cache_utils import Cache +from transformers.configuration_utils import PretrainedConfig +from transformers.generation import GenerationMixin +from transformers.modeling_layers import GradientCheckpointingLayer +from transformers.modeling_outputs import MoeModelOutputWithPast +from transformers.models.qwen3_5_moe.modeling_qwen3_5_moe import Qwen3_5MoeVisionModel +from transformers.processing_utils import Unpack +from transformers.utils import TransformersKwargs, logging + +from prime_rl.trainer.models.base import PreTrainedModelPrimeRL +from prime_rl.trainer.models.layers.lm_head import PrimeLmOutput +from prime_rl.trainer.models.layers.moe import FeedForward, MoE, MoEArgs +from prime_rl.trainer.models.layers.rotary_emb import RotaryEmbedding, RotaryEmbeddingConfig, apply_rotary_pos_emb +from prime_rl.utils.sequence import get_cu_seqlens_from_position_ids + +from .configuration_qwen3_5_moe import Qwen3_5MoeConfig +from .converting_qwen3_5_moe import ( + convert_hf_layer_to_tt, + convert_hf_to_tt_moe, + convert_tt_layer_to_hf, + convert_tt_to_hf_moe, +) + +# Flash attention imports +try: + from flash_attn import flash_attn_varlen_func +except ImportError: + flash_attn_varlen_func = None # type: ignore + +try: + from flash_attn_interface import flash_attn_varlen_func as flash_attn_3_varlen_func +except ImportError: + flash_attn_3_varlen_func = None # type: ignore + +try: + from flash_attn.cute import flash_attn_varlen_func as flash_attn_4_varlen_func +except ImportError: + flash_attn_4_varlen_func = None # type: ignore + +# Flash linear attention imports (for GatedDeltaNet fast path) +try: + from causal_conv1d import causal_conv1d_fn +except ImportError: + causal_conv1d_fn = None # type: ignore + +try: + from fla.modules import FusedRMSNormGated + from fla.ops.cp import FLACPContext, build_cp_context + from fla.ops.gated_delta_rule import chunk_gated_delta_rule +except ImportError: + chunk_gated_delta_rule = None # type: ignore + FusedRMSNormGated = None # type: ignore + FLACPContext = None # type: ignore + build_cp_context = None # type: ignore + +logger = logging.get_logger(__name__) + + +# --------------------------------------------------------------------------- +# RMSNorm variants +# --------------------------------------------------------------------------- + + +class Qwen3_5MoeRMSNorm(nn.Module): + """RMSNorm with (1+weight) parameterization. Weight initialized to zeros.""" + + def __init__(self, dim: int, eps: float = 1e-6): + super().__init__() + self.eps = eps + self.weight = nn.Parameter(torch.zeros(dim)) + + def _norm(self, x): + return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps) + + def forward(self, x): + output = self._norm(x.float()) + output = output * (1.0 + self.weight.float()) + return output.type_as(x) + + +class Qwen3_5MoeRMSNormGated(nn.Module): + """RMSNorm with SiLU gating for GatedDeltaNet output.""" + + def __init__(self, hidden_size, eps=1e-6): + super().__init__() + self.weight = nn.Parameter(torch.ones(hidden_size)) + self.variance_epsilon = eps + + def forward(self, hidden_states, gate=None): + input_dtype = hidden_states.dtype + hidden_states = hidden_states.to(torch.float32) + variance = hidden_states.pow(2).mean(-1, keepdim=True) + hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon) + hidden_states = self.weight * hidden_states.to(input_dtype) + hidden_states = hidden_states * F.silu(gate.to(torch.float32)) + return hidden_states.to(input_dtype) + + +# --------------------------------------------------------------------------- +# GatedDeltaNet linear attention +# --------------------------------------------------------------------------- + + +def l2norm(x: torch.FloatTensor, dim: int = -1, eps: float = 1e-6): + inv_norm = torch.rsqrt((x * x).sum(dim=dim, keepdim=True) + eps) + return x * inv_norm + + +def torch_chunk_gated_delta_rule( + query, + key, + value, + g, + beta, + chunk_size=64, + initial_state=None, + output_final_state=False, + use_qk_l2norm_in_kernel=False, +): + """Pure-PyTorch fallback for chunk_gated_delta_rule.""" + initial_dtype = query.dtype + if use_qk_l2norm_in_kernel: + query = l2norm(query, dim=-1, eps=1e-6) + key = l2norm(key, dim=-1, eps=1e-6) + query, key, value, beta, g = [ + x.transpose(1, 2).contiguous().to(torch.float32) for x in (query, key, value, beta, g) + ] + + batch_size, num_heads, sequence_length, k_head_dim = key.shape + v_head_dim = value.shape[-1] + pad_size = (chunk_size - sequence_length % chunk_size) % chunk_size + query = F.pad(query, (0, 0, 0, pad_size)) + key = F.pad(key, (0, 0, 0, pad_size)) + value = F.pad(value, (0, 0, 0, pad_size)) + beta = F.pad(beta, (0, pad_size)) + g = F.pad(g, (0, pad_size)) + total_sequence_length = sequence_length + pad_size + scale = 1 / (query.shape[-1] ** 0.5) + query = query * scale + + v_beta = value * beta.unsqueeze(-1) + k_beta = key * beta.unsqueeze(-1) + query, key, value, k_beta, v_beta = [ + x.reshape(x.shape[0], x.shape[1], -1, chunk_size, x.shape[-1]) for x in (query, key, value, k_beta, v_beta) + ] + g = g.reshape(g.shape[0], g.shape[1], -1, chunk_size) + mask = torch.triu(torch.ones(chunk_size, chunk_size, dtype=torch.bool, device=query.device), diagonal=0) + + g = g.cumsum(dim=-1) + decay_mask = ((g.unsqueeze(-1) - g.unsqueeze(-2)).tril().exp().float()).tril() + attn = -((k_beta @ key.transpose(-1, -2)) * decay_mask).masked_fill(mask, 0) + for i in range(1, chunk_size): + row = attn[..., i, :i].clone() + sub = attn[..., :i, :i].clone() + attn[..., i, :i] = row + (row.unsqueeze(-1) * sub).sum(-2) + attn = attn + torch.eye(chunk_size, dtype=attn.dtype, device=attn.device) + value = attn @ v_beta + k_cumdecay = attn @ (k_beta * g.exp().unsqueeze(-1)) + last_recurrent_state = ( + torch.zeros(batch_size, num_heads, k_head_dim, v_head_dim).to(value) + if initial_state is None + else initial_state.to(value) + ) + core_attn_out = torch.zeros_like(value) + mask = torch.triu(torch.ones(chunk_size, chunk_size, dtype=torch.bool, device=query.device), diagonal=1) + + for i in range(0, total_sequence_length // chunk_size): + q_i, k_i, v_i = query[:, :, i], key[:, :, i], value[:, :, i] + attn = (q_i @ k_i.transpose(-1, -2) * decay_mask[:, :, i]).masked_fill_(mask, 0) + v_prime = (k_cumdecay[:, :, i]) @ last_recurrent_state + v_new = v_i - v_prime + attn_inter = (q_i * g[:, :, i, :, None].exp()) @ last_recurrent_state + core_attn_out[:, :, i] = attn_inter + attn @ v_new + last_recurrent_state = ( + last_recurrent_state * g[:, :, i, -1, None, None].exp() + + (k_i * (g[:, :, i, -1, None] - g[:, :, i]).exp()[..., None]).transpose(-1, -2) @ v_new + ) + + if not output_final_state: + last_recurrent_state = None + core_attn_out = core_attn_out.reshape(core_attn_out.shape[0], core_attn_out.shape[1], -1, core_attn_out.shape[-1]) + core_attn_out = core_attn_out[:, :, :sequence_length] + core_attn_out = core_attn_out.transpose(1, 2).contiguous().to(initial_dtype) + return core_attn_out, last_recurrent_state + + +class Qwen3_5MoeGatedDeltaNet(nn.Module): + """GatedDeltaNet linear attention with Conv1d, beta/gamma gates, and chunk delta rule.""" + + def __init__(self, config: Qwen3_5MoeConfig): + super().__init__() + self.hidden_size = config.hidden_size + self.num_v_heads = config.linear_num_value_heads + self.num_k_heads = config.linear_num_key_heads + self.head_k_dim = config.linear_key_head_dim + self.head_v_dim = config.linear_value_head_dim + self.key_dim = self.head_k_dim * self.num_k_heads + self.value_dim = self.head_v_dim * self.num_v_heads + + self.conv_kernel_size = config.linear_conv_kernel_dim + self.activation = config.hidden_act + self.layer_norm_epsilon = config.rms_norm_eps + + # QKV convolution (depthwise) + self.conv_dim = self.key_dim * 2 + self.value_dim + self.conv1d = nn.Conv1d( + in_channels=self.conv_dim, + out_channels=self.conv_dim, + bias=False, + kernel_size=self.conv_kernel_size, + groups=self.conv_dim, + padding=self.conv_kernel_size - 1, + ) + + # Time step projection + self.dt_bias = nn.Parameter(torch.ones(self.num_v_heads)) + A = torch.empty(self.num_v_heads).uniform_(0, 16) + self.A_log = nn.Parameter(torch.log(A)) + + # Gated RMSNorm on output (per-head) + if FusedRMSNormGated is not None: + self.norm = FusedRMSNormGated(self.head_v_dim, eps=self.layer_norm_epsilon) + else: + self.norm = Qwen3_5MoeRMSNormGated(self.head_v_dim, eps=self.layer_norm_epsilon) + + self.out_proj = nn.Linear(self.value_dim, self.hidden_size, bias=False) + + # Input projections + self.in_proj_qkv = nn.Linear(self.hidden_size, self.key_dim * 2 + self.value_dim, bias=False) + self.in_proj_z = nn.Linear(self.hidden_size, self.value_dim, bias=False) + self.in_proj_b = nn.Linear(self.hidden_size, self.num_v_heads, bias=False) + self.in_proj_a = nn.Linear(self.hidden_size, self.num_v_heads, bias=False) + + # Select best available kernel + self._causal_conv1d_fn = causal_conv1d_fn + self._chunk_gated_delta_rule = chunk_gated_delta_rule or torch_chunk_gated_delta_rule + + def _build_cp_context(self, local_seq_len: int, device: torch.device) -> "FLACPContext | None": + """Build fla CP context from the local (sharded) sequence length.""" + cp_group = getattr(self, "cp_group", None) + if cp_group is None or build_cp_context is None: + return None + # Reconstruct global cu_seqlens: single contiguous sequence across all CP ranks + global_seq_len = local_seq_len * self.cp_world_size + global_cu_seqlens = torch.tensor([0, global_seq_len], dtype=torch.int32, device=device) + return build_cp_context( + cu_seqlens=global_cu_seqlens, + group=cp_group, + conv1d_kernel_size=self.conv_kernel_size, + ) + + def forward( + self, + hidden_states: torch.Tensor, + cu_seqlens: torch.LongTensor | None = None, + ) -> torch.Tensor: + batch_size, seq_len, _ = hidden_states.shape + + mixed_qkv = self.in_proj_qkv(hidden_states).transpose(1, 2) + z = self.in_proj_z(hidden_states).reshape(batch_size, seq_len, -1, self.head_v_dim) + b = self.in_proj_b(hidden_states) + a = self.in_proj_a(hidden_states) + + # Causal conv1d — must reset at sequence boundaries for packed batches, + # otherwise the kernel-1 left pad leaks state across sequences. + if self._causal_conv1d_fn is not None: + seq_idx = None + if cu_seqlens is not None: + seg_lens = cu_seqlens[1:] - cu_seqlens[:-1] + seq_idx = torch.repeat_interleave( + torch.arange(seg_lens.numel(), dtype=torch.int32, device=hidden_states.device), + seg_lens, + ).unsqueeze(0) + mixed_qkv = self._causal_conv1d_fn( + x=mixed_qkv, + weight=self.conv1d.weight.squeeze(1), + bias=self.conv1d.bias, + activation=self.activation, + seq_idx=seq_idx, + ) + elif cu_seqlens is not None: + cu = cu_seqlens.tolist() + conv_outs = [] + for i in range(len(cu) - 1): + s, e = cu[i], cu[i + 1] + if s == e: + continue + conv_outs.append(self.conv1d(mixed_qkv[:, :, s:e])[:, :, : e - s]) + mixed_qkv = F.silu(torch.cat(conv_outs, dim=-1)) + else: + mixed_qkv = F.silu(self.conv1d(mixed_qkv)[:, :, :seq_len]) + + mixed_qkv = mixed_qkv.transpose(1, 2) + query, key, value = torch.split(mixed_qkv, [self.key_dim, self.key_dim, self.value_dim], dim=-1) + + query = query.reshape(batch_size, seq_len, -1, self.head_k_dim) + key = key.reshape(batch_size, seq_len, -1, self.head_k_dim) + value = value.reshape(batch_size, seq_len, -1, self.head_v_dim) + + beta = b.sigmoid() + g = -self.A_log.float().exp() * F.softplus(a.float() + self.dt_bias) + + if self.num_v_heads // self.num_k_heads > 1: + query = query.repeat_interleave(self.num_v_heads // self.num_k_heads, dim=2) + key = key.repeat_interleave(self.num_v_heads // self.num_k_heads, dim=2) + + # Use fla's native CP when available, otherwise fall back to PyTorch kernel + cp_context = self._build_cp_context(seq_len, hidden_states.device) + if cp_context is not None: + cu_seqlens = cp_context.cu_seqlens + core_attn_out, _ = self._chunk_gated_delta_rule( + query, + key, + value, + g=g, + beta=beta, + use_qk_l2norm_in_kernel=True, + cu_seqlens=cu_seqlens, + cp_context=cp_context, + ) + else: + core_attn_out, _ = self._chunk_gated_delta_rule( + query, + key, + value, + g=g, + beta=beta, + initial_state=None, + output_final_state=False, + use_qk_l2norm_in_kernel=True, + cu_seqlens=cu_seqlens, + ) + + # Gated RMSNorm + core_attn_out = core_attn_out.reshape(-1, self.head_v_dim) + z = z.reshape(-1, self.head_v_dim) + core_attn_out = self.norm(core_attn_out, z) + core_attn_out = core_attn_out.reshape(batch_size, seq_len, -1) + + return self.out_proj(core_attn_out) + + +# --------------------------------------------------------------------------- +# Gated softmax attention (for full_attention layers) +# --------------------------------------------------------------------------- + + +@dataclass +class Qwen3_5MoeGatedAttentionConfig: + hidden_size: int + head_dim: int + num_attention_heads: int + num_key_value_heads: int + rms_norm_eps: float + attention_bias: bool = False + attention_dropout: float = 0.0 + + +def _repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor: + if n_rep == 1: + return hidden_states + batch, num_kv_heads, slen, head_dim = hidden_states.shape + hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_kv_heads, n_rep, slen, head_dim) + return hidden_states.reshape(batch, num_kv_heads * n_rep, slen, head_dim) + + +class Qwen3_5MoeGatedAttentionBase(nn.Module): + """Base class for gated softmax attention (Q projects 2x: query + gate).""" + + def __init__(self, config: Qwen3_5MoeGatedAttentionConfig): + super().__init__() + self.head_dim = config.head_dim + self.num_heads = config.num_attention_heads + self.num_key_value_heads = config.num_key_value_heads + self.num_key_value_groups = config.num_attention_heads // config.num_key_value_heads + self.scaling = self.head_dim**-0.5 + self.attention_dropout = config.attention_dropout + + # Q projects 2x: query + gate + self.q_proj = nn.Linear(config.hidden_size, self.num_heads * self.head_dim * 2, bias=config.attention_bias) + self.k_proj = nn.Linear( + config.hidden_size, self.num_key_value_heads * self.head_dim, bias=config.attention_bias + ) + self.v_proj = nn.Linear( + config.hidden_size, self.num_key_value_heads * self.head_dim, bias=config.attention_bias + ) + self.o_proj = nn.Linear(self.num_heads * self.head_dim, config.hidden_size, bias=config.attention_bias) + + # QK normalization with (1+weight) parameterization + self.q_norm = Qwen3_5MoeRMSNorm(self.head_dim, eps=config.rms_norm_eps) + self.k_norm = Qwen3_5MoeRMSNorm(self.head_dim, eps=config.rms_norm_eps) + + def output_proj( + self, + attn_output: torch.Tensor, + gate: torch.Tensor, + ) -> torch.Tensor: + input_shape = gate.shape[:-1] + if attn_output.dim() == 4: + attn_output = attn_output.transpose(1, 2).contiguous() + attn_output = attn_output.contiguous().view(*input_shape, -1) + attn_output = attn_output * torch.sigmoid(gate) + return self.o_proj(attn_output) + + +class Qwen3_5MoeGatedSDPAAttention(Qwen3_5MoeGatedAttentionBase): + """Gated softmax attention using PyTorch's scaled_dot_product_attention.""" + + def attn_projections( + self, + hidden_states: torch.Tensor, + position_embeddings: tuple[torch.Tensor, torch.Tensor], + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + input_shape = hidden_states.shape[:-1] + hidden_shape = (*input_shape, -1, self.head_dim) + + query_states, gate = torch.chunk( + self.q_proj(hidden_states).view(*input_shape, -1, self.head_dim * 2), 2, dim=-1 + ) + gate = gate.reshape(*input_shape, -1) + + query_states = self.q_norm(query_states.view(hidden_shape)).transpose(1, 2) + key_states = self.k_norm(self.k_proj(hidden_states).view(hidden_shape)).transpose(1, 2) + value_states = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2) + + cos, sin = position_embeddings + query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin) + + return query_states, key_states, value_states, gate + + def _attention_core( + self, + query_states: torch.Tensor, + key_states: torch.Tensor, + value_states: torch.Tensor, + ) -> torch.Tensor: + key_states = _repeat_kv(key_states, self.num_key_value_groups) + value_states = _repeat_kv(value_states, self.num_key_value_groups) + return F.scaled_dot_product_attention( + query_states, + key_states, + value_states, + is_causal=True, + scale=self.scaling, + ) + + def forward( + self, + hidden_states: torch.Tensor, + position_embeddings: tuple[torch.Tensor, torch.Tensor], + cu_seqlens: torch.LongTensor | None = None, + max_seqlen: int | None = None, + ) -> tuple[torch.Tensor, None]: + query_states, key_states, value_states, gate = self.attn_projections(hidden_states, position_embeddings) + attn_output = self._attention_core(query_states, key_states, value_states) + return self.output_proj(attn_output, gate), None + + +class Qwen3_5MoeGatedFlashAttention(Qwen3_5MoeGatedAttentionBase): + """Gated softmax attention using Flash Attention varlen functions.""" + + _funcs = { + 2: flash_attn_varlen_func, + 3: flash_attn_3_varlen_func, + 4: flash_attn_4_varlen_func, + } + + def __init__(self, config: Qwen3_5MoeGatedAttentionConfig, flash_attn_version: int = 4): + super().__init__(config) + self._flash_attn_version = flash_attn_version + self.func = self._funcs[flash_attn_version] + self._flash_attn_call = self.func + if self._flash_attn_version == 4: + self._flash_attn_call = torch._dynamo.disable(self.func) + + def _compute_attention(self, q, k, v, cu_seqlens, max_seqlen): + """Run the flash attention kernel. q/k/v are [total_tokens, heads, dim].""" + kwargs: dict = {"causal": True} + sliding_window = getattr(self, "sliding_window", None) + if sliding_window is not None: + kwargs["window_size"] = (sliding_window - 1, 0) + if self._flash_attn_version == 4: + # FA4's flash_attn_varlen_func has qv as the 4th positional arg, + # so cu_seqlens must be passed as keyword args to avoid misalignment. + kwargs["cu_seqlens_q"] = cu_seqlens + kwargs["cu_seqlens_k"] = cu_seqlens + out = self._flash_attn_call(q, k, v, **kwargs) + else: + out = self._flash_attn_call(q, k, v, cu_seqlens, cu_seqlens, max_seqlen, max_seqlen, **kwargs) + if isinstance(out, tuple): + out = out[0] + return out + + def _attention_core( + self, + query_states: torch.Tensor, + key_states: torch.Tensor, + value_states: torch.Tensor, + cu_seqlens: torch.LongTensor | None = None, + max_seqlen: int | None = None, + ) -> torch.Tensor: + return self._compute_attention(query_states[0], key_states[0], value_states[0], cu_seqlens, max_seqlen) + + def attn_projections( + self, + hidden_states: torch.Tensor, + position_embeddings: tuple[torch.Tensor, torch.Tensor], + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + input_shape = hidden_states.shape[:-1] + hidden_shape = (*input_shape, -1, self.head_dim) + + query_states, gate = torch.chunk( + self.q_proj(hidden_states).view(*input_shape, -1, self.head_dim * 2), 2, dim=-1 + ) + gate = gate.reshape(*input_shape, -1) + + query_states = self.q_norm(query_states.view(hidden_shape)) + key_states = self.k_norm(self.k_proj(hidden_states).view(hidden_shape)) + value_states = self.v_proj(hidden_states).view(hidden_shape) + + query_states = query_states.transpose(1, 2) + key_states = key_states.transpose(1, 2) + cos, sin = position_embeddings + query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin) + query_states = query_states.transpose(1, 2) + key_states = key_states.transpose(1, 2) + + return query_states, key_states, value_states, gate + + def forward( + self, + hidden_states: torch.Tensor, + position_embeddings: tuple[torch.Tensor, torch.Tensor], + cu_seqlens: torch.LongTensor | None = None, + max_seqlen: int | None = None, + ) -> tuple[torch.Tensor, None]: + query_states, key_states, value_states, gate = self.attn_projections(hidden_states, position_embeddings) + attn_output = self._attention_core( + query_states, + key_states, + value_states, + cu_seqlens=cu_seqlens, + max_seqlen=max_seqlen, + ) + return self.output_proj(attn_output, gate), None + + +QWEN35MOE_ATTN_IMPL2CLASS = { + "sdpa": Qwen3_5MoeGatedSDPAAttention, + "flash_attention_2": functools.partial(Qwen3_5MoeGatedFlashAttention, flash_attn_version=2), + "flash_attention_3": functools.partial(Qwen3_5MoeGatedFlashAttention, flash_attn_version=3), + "fa4": functools.partial(Qwen3_5MoeGatedFlashAttention, flash_attn_version=4), +} + + +# --------------------------------------------------------------------------- +# Decoder layer +# --------------------------------------------------------------------------- + + +def _get_gated_attention(config: Qwen3_5MoeConfig) -> nn.Module: + attn_config = Qwen3_5MoeGatedAttentionConfig( + hidden_size=config.hidden_size, + head_dim=config.head_dim, + num_attention_heads=config.num_attention_heads, + num_key_value_heads=config.num_key_value_heads, + rms_norm_eps=config.rms_norm_eps, + attention_bias=config.attention_bias, + attention_dropout=config.attention_dropout, + ) + + attn_impl = config._attn_implementation + if attn_impl == "eager": + attn_impl = "sdpa" + + if attn_impl not in QWEN35MOE_ATTN_IMPL2CLASS: + supported = list(QWEN35MOE_ATTN_IMPL2CLASS.keys()) + raise ValueError( + f"Qwen3.5-MoE attention does not support '{config._attn_implementation}'. " + f"Supported implementations: {supported}." + ) + + return QWEN35MOE_ATTN_IMPL2CLASS[attn_impl](attn_config) + + +class Qwen3_5MoeDecoderLayer(GradientCheckpointingLayer): + def __init__(self, config: Qwen3_5MoeConfig, layer_idx: int): + super().__init__() + self.hidden_size = config.hidden_size + self.layer_type = config.layer_types[layer_idx] + + # Token mixer: either GatedDeltaNet or gated softmax attention + if self.layer_type == "linear_attention": + self.linear_attn = Qwen3_5MoeGatedDeltaNet(config) + elif self.layer_type == "full_attention": + self.self_attn = _get_gated_attention(config) + + # MoE: routed experts via shared MoE class (no shared experts in MoE itself) + moe_args = MoEArgs( + num_experts=config.num_experts, + num_shared_experts=0, + score_func="softmax", + route_norm=True, + route_scale=1.0, + score_before_experts=False, + top_k=config.num_experts_per_tok, + use_grouped_mm=config.use_grouped_mm, + load_balance_coeff=config.load_balance_coeff, + fp8=getattr(config, "fp8", False), + ) + self.mlp = MoE(moe_args, dim=config.hidden_size, hidden_dim=config.moe_intermediate_size) + + # Separate gated shared expert + self.shared_expert = FeedForward(dim=config.hidden_size, hidden_dim=config.shared_expert_intermediate_size) + self.shared_expert_gate = nn.Linear(config.hidden_size, 1, bias=False) + + # Layer norms with (1+weight) parameterization + self.input_layernorm = Qwen3_5MoeRMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.post_attention_layernorm = Qwen3_5MoeRMSNorm(config.hidden_size, eps=config.rms_norm_eps) + + def forward( + self, + hidden_states: torch.Tensor, + position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None, + cu_seqlens: torch.LongTensor | None = None, + max_seqlen: int | None = None, + routed_experts: Optional[torch.LongTensor] = None, + ) -> torch.FloatTensor: + residual = hidden_states + hidden_states = self.input_layernorm(hidden_states) + + # Token mixer + if self.layer_type == "linear_attention": + hidden_states = self.linear_attn(hidden_states, cu_seqlens=cu_seqlens) + elif self.layer_type == "full_attention": + hidden_states, _ = self.self_attn( + hidden_states=hidden_states, + position_embeddings=position_embeddings, + cu_seqlens=cu_seqlens, + max_seqlen=max_seqlen, + ) + + hidden_states = residual + hidden_states + + # MLP: routed experts + gated shared expert + residual = hidden_states + hidden_states = self.post_attention_layernorm(hidden_states) + + # Routed experts + routed_output = self.mlp(hidden_states, routed_experts=routed_experts) + + # Gated shared expert + bs, slen, dim = hidden_states.shape + hidden_flat = hidden_states.view(-1, dim) + shared_output = self.shared_expert(hidden_flat) + shared_output = F.sigmoid(self.shared_expert_gate(hidden_flat)) * shared_output + shared_output = shared_output.view(bs, slen, dim) + + hidden_states = residual + routed_output + shared_output + return hidden_states + + +# --------------------------------------------------------------------------- +# Model classes +# --------------------------------------------------------------------------- + + +def _create_rotary_emb(config: Qwen3_5MoeConfig) -> RotaryEmbedding: + if hasattr(config, "rope_scaling") and isinstance(config.rope_scaling, dict): + rope_type = config.rope_scaling.get("rope_type", config.rope_scaling.get("type", "default")) + else: + rope_type = "default" + + rotary_config = RotaryEmbeddingConfig( + max_position_embeddings=config.max_position_embeddings, + rope_type=rope_type, + model_config=config, + ) + return RotaryEmbedding(rotary_config) + + +class Qwen3_5MoePreTrainedModel(PreTrainedModelPrimeRL): + config_class = Qwen3_5MoeConfig + base_model_prefix = "model" + supports_gradient_checkpointing = True + _no_split_modules = ["Qwen3_5MoeDecoderLayer"] + _skip_keys_device_placement = ["past_key_values"] + _supports_flash_attn = True + _supports_sdpa = True + _supports_flex_attn = False + _supports_attention_backend = True + _can_compile_fullgraph = False + _can_record_outputs = { + "hidden_states": Qwen3_5MoeDecoderLayer, + } + + @classmethod + def is_hf_state_dict(cls, state_dict: dict[str, Tensor]) -> bool: + return any( + "mlp.experts.1.up_proj" in name + or "mlp.experts.gate_up_proj" in name + or "mlp.shared_expert.gate_proj" in name + for name in state_dict.keys() + ) + + @classmethod + def is_prime_state_dict(cls, state_dict: dict[str, Tensor]) -> bool: + return any("mlp.experts.w1" in name for name in state_dict.keys()) + + @classmethod + def convert_to_hf(cls, state_dict: dict[str, Tensor]) -> dict[str, Tensor]: + convert_tt_to_hf_moe(state_dict) + return state_dict + + @classmethod + def convert_to_prime(cls, state_dict: dict[str, Tensor]) -> dict[str, Tensor]: + convert_hf_to_tt_moe(state_dict) + return state_dict + + @classmethod + def convert_layer_to_hf(cls, state_dict: dict[str, Tensor], layer_idx: int) -> dict[str, Tensor]: + convert_tt_layer_to_hf(state_dict, layer_idx) + return state_dict + + @classmethod + def convert_layer_to_prime(cls, state_dict: dict[str, Tensor], layer_idx: int) -> dict[str, Tensor]: + convert_hf_layer_to_tt(state_dict, layer_idx) + return state_dict + + +class Qwen3_5MoeModel(Qwen3_5MoePreTrainedModel): + def __init__(self, config: Qwen3_5MoeConfig): + super().__init__(config) + self.padding_idx = config.pad_token_id + self.vocab_size = config.vocab_size + + self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx) + self.layers = nn.ModuleList( + [Qwen3_5MoeDecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)] + ) + self.norm = Qwen3_5MoeRMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.rotary_emb = _create_rotary_emb(config) + self.gradient_checkpointing = False + + self.post_init() + + def forward( + self, + input_ids: Optional[torch.LongTensor] = None, + position_ids: Optional[torch.LongTensor] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + routed_experts: Optional[torch.LongTensor] = None, + ) -> MoeModelOutputWithPast: + if (input_ids is None) ^ (inputs_embeds is not None): + raise ValueError("You must specify exactly one of input_ids or inputs_embeds") + + if inputs_embeds is None: + inputs_embeds = self.embed_tokens(input_ids) + + if self.config._attn_implementation in ("flash_attention_2", "flash_attention_3", "fa4"): + cu_seqlens, max_seqlen = get_cu_seqlens_from_position_ids(position_ids) + torch._dynamo.mark_dynamic(cu_seqlens, 0) + else: + max_seqlen = None + cu_seqlens = None + + hidden_states = inputs_embeds + position_embeddings = self.rotary_emb(hidden_states, position_ids) + + for layer_idx, decoder_layer in enumerate(self.layers): + routed_experts_layer = routed_experts[:, :, layer_idx, :] if routed_experts is not None else None + hidden_states = decoder_layer( + hidden_states, + position_embeddings=position_embeddings, + cu_seqlens=cu_seqlens, + max_seqlen=max_seqlen, + routed_experts=routed_experts_layer, + ) + + hidden_states = self.norm(hidden_states) + return MoeModelOutputWithPast(last_hidden_state=hidden_states) + + +# --------------------------------------------------------------------------- +# VLM composite model body +# --------------------------------------------------------------------------- + + +def _build_text_config(composite_config: PretrainedConfig) -> Qwen3_5MoeConfig: + """Build custom PrimeRL text config from HF's composite VLM config.""" + text_dict = composite_config.text_config.to_dict() + text_config = Qwen3_5MoeConfig(**text_dict) + attn_impl = getattr( + composite_config.text_config, + "_attn_implementation", + getattr(composite_config, "_attn_implementation", None), + ) + if attn_impl is not None: + text_config._attn_implementation = attn_impl + return text_config + + +class Qwen3_5MoeVLMModel(nn.Module): + """Composite VLM body: HF vision encoder + custom PrimeRL text model.""" + + def __init__(self, config: PretrainedConfig): + super().__init__() + self.config = config + self.visual = Qwen3_5MoeVisionModel._from_config(config.vision_config) + self.language_model = Qwen3_5MoeModel(_build_text_config(config)) + + def get_input_embeddings(self): + return self.language_model.embed_tokens + + def set_input_embeddings(self, value): + self.language_model.embed_tokens = value + + def forward( + self, + input_ids: torch.LongTensor | None = None, + position_ids: torch.LongTensor | None = None, + inputs_embeds: torch.FloatTensor | None = None, + pixel_values: torch.Tensor | None = None, + image_grid_thw: torch.LongTensor | None = None, + routed_experts: torch.LongTensor | None = None, + **kwargs, + ) -> MoeModelOutputWithPast: + if inputs_embeds is None: + inputs_embeds = self.language_model.embed_tokens(input_ids) + + if pixel_values is not None: + pixel_values = pixel_values.type(self.visual.dtype) + vision_output = self.visual(pixel_values, grid_thw=image_grid_thw, return_dict=True) + image_embeds = vision_output.pooler_output.to(inputs_embeds.device, inputs_embeds.dtype) + + image_mask = input_ids == self.config.image_token_id + image_mask = image_mask.unsqueeze(-1).expand_as(inputs_embeds).to(inputs_embeds.device) + inputs_embeds = inputs_embeds.masked_scatter(image_mask, image_embeds) + + if position_ids is None: + position_ids = torch.arange(inputs_embeds.shape[1], device=inputs_embeds.device).unsqueeze(0) + + return self.language_model( + inputs_embeds=inputs_embeds, + position_ids=position_ids, + routed_experts=routed_experts, + ) + + +# --------------------------------------------------------------------------- +# Unified CausalLM / VLM class +# --------------------------------------------------------------------------- + + +def _has_vlm_keys(state_dict: dict[str, Tensor]) -> bool: + return any(k.startswith("model.language_model.") for k in state_dict) + + +def _remap_lm_keys(state_dict: dict[str, Tensor], to_flat: bool = True) -> None: + """Remap language model keys between VLM and flat format for weight conversion. + + to_flat=True: model.language_model.* -> model.* + to_flat=False: model.* -> model.language_model.* + + Vision keys (model.visual.*) are never touched. + """ + src = "model.language_model." if to_flat else "model." + dst = "model." if to_flat else "model.language_model." + for k in [k for k in list(state_dict.keys()) if k.startswith(src) and not k.startswith("model.visual.")]: + state_dict[dst + k[len(src) :]] = state_dict.pop(k) + + +class Qwen3_5MoeForCausalLM(Qwen3_5MoePreTrainedModel, GenerationMixin): + """Unified Qwen3.5 MoE model for both text-only and VLM configs. + + When config has a vision_config, creates a composite model with HF's frozen + vision encoder + custom text model. Otherwise creates a text-only model. + """ + + _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} + _checkpoint_conversion_mapping = {} + _tp_plan = {"lm_head": "colwise_rep"} + _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + + def __init__(self, config, **kwargs): + super().__init__(config, **kwargs) + self._is_vlm = hasattr(config, "vision_config") + + if self._is_vlm: + self.model = Qwen3_5MoeVLMModel(config) + text_config = config.text_config + self._tied_weights_keys = {"lm_head.weight": "model.language_model.embed_tokens.weight"} + else: + self.model = Qwen3_5MoeModel(config) + text_config = config + + self.vocab_size = text_config.vocab_size + self.lm_head = nn.Linear(text_config.hidden_size, text_config.vocab_size, bias=False) + self.post_init() + + def get_input_embeddings(self): + if self._is_vlm: + return self.model.get_input_embeddings() + return self.model.embed_tokens + + def set_input_embeddings(self, value): + if self._is_vlm: + self.model.set_input_embeddings(value) + else: + self.model.embed_tokens = value + + def set_decoder(self, decoder): + self.model = decoder + + def get_decoder(self): + return self.model + + # ------------------------------------------------------------------ + # State dict detection & conversion (handles both text-only and VLM) + # ------------------------------------------------------------------ + + @classmethod + def is_hf_state_dict(cls, state_dict: dict[str, Tensor]) -> bool: + return any( + "mlp.experts.gate_up_proj" in name + or "mlp.experts.1.up_proj" in name + or "mlp.shared_expert.gate_proj" in name + for name in state_dict + ) + + @classmethod + def is_prime_state_dict(cls, state_dict: dict[str, Tensor]) -> bool: + return any("mlp.experts.w1" in name for name in state_dict) + + @classmethod + def convert_to_hf(cls, state_dict: dict[str, Tensor]) -> dict[str, Tensor]: + vlm = _has_vlm_keys(state_dict) + if vlm: + _remap_lm_keys(state_dict, to_flat=True) + convert_tt_to_hf_moe(state_dict) + if vlm: + _remap_lm_keys(state_dict, to_flat=False) + return state_dict + + @classmethod + def convert_to_prime(cls, state_dict: dict[str, Tensor]) -> dict[str, Tensor]: + vlm = _has_vlm_keys(state_dict) + if vlm: + _remap_lm_keys(state_dict, to_flat=True) + convert_hf_to_tt_moe(state_dict) + if vlm: + _remap_lm_keys(state_dict, to_flat=False) + return state_dict + + @classmethod + def convert_layer_to_hf(cls, state_dict: dict[str, Tensor], layer_idx: int) -> dict[str, Tensor]: + vlm = _has_vlm_keys(state_dict) + if vlm: + _remap_lm_keys(state_dict, to_flat=True) + convert_tt_layer_to_hf(state_dict, layer_idx) + if vlm: + _remap_lm_keys(state_dict, to_flat=False) + return state_dict + + @classmethod + def convert_layer_to_prime(cls, state_dict: dict[str, Tensor], layer_idx: int) -> dict[str, Tensor]: + vlm = _has_vlm_keys(state_dict) + if vlm: + _remap_lm_keys(state_dict, to_flat=True) + convert_hf_layer_to_tt(state_dict, layer_idx) + if vlm: + _remap_lm_keys(state_dict, to_flat=False) + return state_dict + + # ------------------------------------------------------------------ + # Forward + # ------------------------------------------------------------------ + + def forward( + self, + input_ids: Optional[torch.LongTensor] = None, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_values: Optional[Cache] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + labels: Optional[torch.LongTensor] = None, + use_cache: Optional[bool] = None, + logits_to_keep: Union[int, torch.Tensor] = 0, + temperature: Union[torch.Tensor, None] = None, + routed_experts: Optional[torch.LongTensor] = None, + pixel_values: Optional[torch.Tensor] = None, + image_grid_thw: Optional[torch.LongTensor] = None, + **kwargs: Unpack[TransformersKwargs], + ) -> PrimeLmOutput: + assert use_cache is None, "use_cache is not supported for custom qwen3_5_moe for now" + assert past_key_values is None, "past_key_values is not supported for custom qwen3_5_moe for now" + + if position_ids is None: + if inputs_embeds is not None: + position_ids = torch.arange(inputs_embeds.shape[1], device=inputs_embeds.device).unsqueeze(0) + elif input_ids is not None: + position_ids = torch.arange(input_ids.shape[1], device=input_ids.device).unsqueeze(0) + + if self._is_vlm: + outputs: MoeModelOutputWithPast = self.model( + input_ids=input_ids, + position_ids=position_ids, + inputs_embeds=inputs_embeds, + pixel_values=pixel_values, + image_grid_thw=image_grid_thw, + routed_experts=routed_experts, + ) + else: + outputs = self.model( + input_ids=input_ids, + position_ids=position_ids, + inputs_embeds=inputs_embeds, + routed_experts=routed_experts, + ) + + hidden_states = outputs.last_hidden_state + slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep + return self.lm_head( + hidden_states[:, slice_indices, :], + labels[:, slice_indices] if labels is not None else None, + temperature=temperature, + ) + + # ------------------------------------------------------------------ + # Buffer init after meta-device loading + # ------------------------------------------------------------------ + + def init_buffers_post_meta(self): + if self._is_vlm: + lm_rope = self.model.language_model.rotary_emb + else: + lm_rope = self.model.rotary_emb + + if hasattr(lm_rope, "rope_init_fn"): + inv_freq, lm_rope.attention_scaling = lm_rope.rope_init_fn(lm_rope.config, lm_rope.inv_freq.device) + lm_rope.inv_freq.copy_(inv_freq) + + if self._is_vlm: + vis_rope = self.model.visual.rotary_pos_emb + if hasattr(vis_rope, "inv_freq"): + dim = vis_rope.inv_freq.shape[0] + inv_freq = 1.0 / ( + 10000.0 + ** (torch.arange(0, dim * 2, 2, dtype=torch.float32, device=vis_rope.inv_freq.device) / (dim * 2)) + ) + vis_rope.inv_freq.copy_(inv_freq) + + +__all__ = [ + "Qwen3_5MoeForCausalLM", + "Qwen3_5MoeModel", + "Qwen3_5MoePreTrainedModel", +] diff --git a/src/prime_rl/trainer/models/qwen3_moe/configuration_qwen3_moe.py b/src/prime_rl/trainer/models/qwen3_moe/configuration_qwen3_moe.py index ae1c7b2da29..07c14fdcb14 100644 --- a/src/prime_rl/trainer/models/qwen3_moe/configuration_qwen3_moe.py +++ b/src/prime_rl/trainer/models/qwen3_moe/configuration_qwen3_moe.py @@ -1,5 +1,4 @@ from transformers.configuration_utils import PretrainedConfig -from transformers.modeling_rope_utils import rope_config_validation class Qwen3MoeConfig(PretrainedConfig): @@ -119,6 +118,7 @@ class Qwen3MoeConfig(PretrainedConfig): model_type = "qwen3_moe" keys_to_ignore_at_inference = ["past_key_values"] + attribute_map = {"num_local_experts": "num_experts"} # Default tensor parallel plan for base model `Qwen3Moe` base_model_tp_plan = { @@ -169,6 +169,7 @@ def __init__( mlp_only_layers=None, load_balance_coeff=None, use_grouped_mm=True, + pad_token_id=None, **kwargs, ): self.vocab_size = vocab_size @@ -189,11 +190,7 @@ def __init__( self.rope_scaling = rope_scaling self.attention_bias = attention_bias self.attention_dropout = attention_dropout - # Validate the correctness of rotary position embeddings parameters - # BC: if there is a 'type' field, move it to 'rope_type'. - if self.rope_scaling is not None and "type" in self.rope_scaling: - self.rope_scaling["rope_type"] = self.rope_scaling["type"] - rope_config_validation(self) + self.pad_token_id = pad_token_id # MoE arguments self.decoder_sparse_step = decoder_sparse_step @@ -206,7 +203,7 @@ def __init__( self.mlp_only_layers = [] if mlp_only_layers is None else mlp_only_layers self.load_balance_coeff = load_balance_coeff self.use_grouped_mm = use_grouped_mm - + super().__init__( tie_word_embeddings=tie_word_embeddings, **kwargs, diff --git a/src/prime_rl/trainer/models/qwen3_moe/converting_qwen3_moe.py b/src/prime_rl/trainer/models/qwen3_moe/converting_qwen3_moe.py index 37584dc9501..48ccc5b9416 100644 --- a/src/prime_rl/trainer/models/qwen3_moe/converting_qwen3_moe.py +++ b/src/prime_rl/trainer/models/qwen3_moe/converting_qwen3_moe.py @@ -11,30 +11,53 @@ def convert_hf_layer_to_tt(state_dict: dict[str, Tensor], layer_idx: int): """Convert a layer from HF to TT format in-place.""" i = layer_idx - num_experts = len([j for j in state_dict.keys() if f"model.layers.{i}.mlp.experts" in j]) // 3 - if num_experts == 0: + # Check if this is a MoE layer by looking for gate weight + if f"model.layers.{i}.mlp.gate.weight" not in state_dict: return state_dict[f"model.layers.{i}.mlp.router.gate.weight"] = state_dict[f"model.layers.{i}.mlp.gate.weight"] del state_dict[f"model.layers.{i}.mlp.gate.weight"] - dim, moe_dim = state_dict[f"model.layers.{i}.mlp.experts.0.down_proj.weight"].shape - w1 = torch.empty( - (num_experts, moe_dim, dim), dtype=state_dict[f"model.layers.{i}.mlp.experts.0.down_proj.weight"].dtype - ) # Gate - w2 = torch.empty( - (num_experts, dim, moe_dim), dtype=state_dict[f"model.layers.{i}.mlp.experts.0.down_proj.weight"].dtype - ) # Down - w3 = torch.empty( - (num_experts, moe_dim, dim), dtype=state_dict[f"model.layers.{i}.mlp.experts.0.down_proj.weight"].dtype - ) # Up - for j in range(num_experts): - w1[j].copy_(state_dict[f"model.layers.{i}.mlp.experts.{j}.gate_proj.weight"]) - w2[j].copy_(state_dict[f"model.layers.{i}.mlp.experts.{j}.down_proj.weight"]) - w3[j].copy_(state_dict[f"model.layers.{i}.mlp.experts.{j}.up_proj.weight"]) - del state_dict[f"model.layers.{i}.mlp.experts.{j}.gate_proj.weight"] - del state_dict[f"model.layers.{i}.mlp.experts.{j}.down_proj.weight"] - del state_dict[f"model.layers.{i}.mlp.experts.{j}.up_proj.weight"] + # Check if this is the new fused format (transformers 5.0+) or old per-expert format + if f"model.layers.{i}.mlp.experts.gate_up_proj" in state_dict: + # New fused format: gate_up_proj has shape (num_experts, 2*moe_dim, dim) + gate_up_proj = state_dict[f"model.layers.{i}.mlp.experts.gate_up_proj"] + down_proj = state_dict[f"model.layers.{i}.mlp.experts.down_proj"] + + num_experts, fused_dim, dim = gate_up_proj.shape + moe_dim = fused_dim // 2 + + # Split gate_up_proj into w1 (gate) and w3 (up) + w1 = gate_up_proj[:, :moe_dim, :] # Gate: (num_experts, moe_dim, dim) + w3 = gate_up_proj[:, moe_dim:, :] # Up: (num_experts, moe_dim, dim) + w2 = down_proj # Down: (num_experts, dim, moe_dim) + + del state_dict[f"model.layers.{i}.mlp.experts.gate_up_proj"] + del state_dict[f"model.layers.{i}.mlp.experts.down_proj"] + else: + # Old per-expert format + num_experts = len([j for j in state_dict.keys() if f"model.layers.{i}.mlp.experts" in j]) // 3 + if num_experts == 0: + return + + dim, moe_dim = state_dict[f"model.layers.{i}.mlp.experts.0.down_proj.weight"].shape + w1 = torch.empty( + (num_experts, moe_dim, dim), dtype=state_dict[f"model.layers.{i}.mlp.experts.0.down_proj.weight"].dtype + ) # Gate + w2 = torch.empty( + (num_experts, dim, moe_dim), dtype=state_dict[f"model.layers.{i}.mlp.experts.0.down_proj.weight"].dtype + ) # Down + w3 = torch.empty( + (num_experts, moe_dim, dim), dtype=state_dict[f"model.layers.{i}.mlp.experts.0.down_proj.weight"].dtype + ) # Up + for j in range(num_experts): + w1[j].copy_(state_dict[f"model.layers.{i}.mlp.experts.{j}.gate_proj.weight"]) + w2[j].copy_(state_dict[f"model.layers.{i}.mlp.experts.{j}.down_proj.weight"]) + w3[j].copy_(state_dict[f"model.layers.{i}.mlp.experts.{j}.up_proj.weight"]) + + del state_dict[f"model.layers.{i}.mlp.experts.{j}.gate_proj.weight"] + del state_dict[f"model.layers.{i}.mlp.experts.{j}.down_proj.weight"] + del state_dict[f"model.layers.{i}.mlp.experts.{j}.up_proj.weight"] state_dict[f"model.layers.{i}.mlp.experts.w1"] = w1 state_dict[f"model.layers.{i}.mlp.experts.w2"] = w2 @@ -51,21 +74,16 @@ def convert_tt_layer_to_hf(state_dict: dict[str, Tensor], layer_index: int): state_dict[f"model.layers.{i}.mlp.gate.weight"] = state_dict[f"model.layers.{i}.mlp.router.gate.weight"] del state_dict[f"model.layers.{i}.mlp.router.gate.weight"] - # Routed experts - num_experts, moe_dim, dim = state_dict[f"model.layers.{i}.mlp.experts.w1"].shape + # Routed experts - convert to per-expert format (compatible with vLLM and transformers) + w1 = state_dict.pop(f"model.layers.{i}.mlp.experts.w1") # (num_experts, moe_dim, dim) + w2 = state_dict.pop(f"model.layers.{i}.mlp.experts.w2") # (num_experts, dim, moe_dim) + w3 = state_dict.pop(f"model.layers.{i}.mlp.experts.w3") # (num_experts, moe_dim, dim) + + num_experts = w1.shape[0] for j in range(num_experts): - state_dict[f"model.layers.{i}.mlp.experts.{j}.gate_proj.weight"] = state_dict[ - f"model.layers.{i}.mlp.experts.w1" - ][j] - state_dict[f"model.layers.{i}.mlp.experts.{j}.down_proj.weight"] = state_dict[ - f"model.layers.{i}.mlp.experts.w2" - ][j] - state_dict[f"model.layers.{i}.mlp.experts.{j}.up_proj.weight"] = state_dict[f"model.layers.{i}.mlp.experts.w3"][ - j - ] - del state_dict[f"model.layers.{i}.mlp.experts.w1"] - del state_dict[f"model.layers.{i}.mlp.experts.w2"] - del state_dict[f"model.layers.{i}.mlp.experts.w3"] + state_dict[f"model.layers.{i}.mlp.experts.{j}.gate_proj.weight"] = w1[j] + state_dict[f"model.layers.{i}.mlp.experts.{j}.down_proj.weight"] = w2[j] + state_dict[f"model.layers.{i}.mlp.experts.{j}.up_proj.weight"] = w3[j] def convert_hf_to_tt_moe(state_dict: dict[str, Tensor]): diff --git a/src/prime_rl/trainer/models/qwen3_moe/modeling_qwen3_moe.py b/src/prime_rl/trainer/models/qwen3_moe/modeling_qwen3_moe.py index d027539143c..25d83f94366 100644 --- a/src/prime_rl/trainer/models/qwen3_moe/modeling_qwen3_moe.py +++ b/src/prime_rl/trainer/models/qwen3_moe/modeling_qwen3_moe.py @@ -31,7 +31,7 @@ from prime_rl.trainer.models.layers.lm_head import PrimeLmOutput from prime_rl.trainer.models.layers.mlp import MLP, MLPConfig from prime_rl.trainer.models.layers.moe import MoE, MoEArgs -from prime_rl.trainer.models.layers.rms_norm import RMSNorm, RMSNormConfig +from prime_rl.trainer.models.layers.norms import RMSNorm, RMSNormConfig from prime_rl.trainer.models.layers.rotary_emb import RotaryEmbedding, RotaryEmbeddingConfig from prime_rl.trainer.models.qwen3_moe.configuration_qwen3_moe import Qwen3MoeConfig from prime_rl.trainer.models.qwen3_moe.converting_qwen3_moe import ( @@ -40,6 +40,7 @@ convert_tt_layer_to_hf, convert_tt_to_hf_moe, ) +from prime_rl.utils.sequence import get_cu_seqlens_from_position_ids logger = logging.get_logger(__name__) @@ -72,6 +73,7 @@ def __init__(self, config: Qwen3MoeConfig, layer_idx: int): top_k=config.num_experts_per_tok, use_grouped_mm=config.use_grouped_mm, load_balance_coeff=config.load_balance_coeff, + fp8=getattr(config, "fp8", False), ) mlp_config = MLPConfig( hidden_size=config.hidden_size, @@ -96,6 +98,7 @@ def forward( position_embeddings: tuple[torch.Tensor, torch.Tensor], cu_seqlens: torch.LongTensor | None = None, max_seqlen: int | None = None, + routed_experts: Optional[torch.LongTensor] = None, ) -> torch.FloatTensor: residual = hidden_states hidden_states = self.input_layernorm(hidden_states) @@ -111,7 +114,7 @@ def forward( # Fully Connected residual = hidden_states hidden_states = self.post_attention_layernorm(hidden_states) - hidden_states = self.mlp(hidden_states) + hidden_states = self.mlp(hidden_states, routed_experts=routed_experts) hidden_states = residual + hidden_states return hidden_states @@ -135,7 +138,8 @@ class Qwen3MoePreTrainedModel(PreTrainedModelPrimeRL): @classmethod def is_hf_state_dict(cls, state_dict: dict[str, Tensor]) -> bool: """Check if the state dict contains MoE layers in HuggingFace format.""" - return any("mlp.experts.1.up_proj" in module_name for module_name in state_dict.keys()) + # Check for old per-expert format or new fused format (transformers 5.0+) + return any("mlp.experts.1.up_proj" in name or "mlp.experts.gate_up_proj" in name for name in state_dict.keys()) @classmethod def is_prime_state_dict(cls, state_dict: dict[str, Tensor]) -> bool: @@ -201,24 +205,20 @@ def forward( input_ids: Optional[torch.LongTensor] = None, position_ids: Optional[torch.LongTensor] = None, inputs_embeds: Optional[torch.FloatTensor] = None, + routed_experts: Optional[torch.LongTensor] = None, ) -> MoeModelOutputWithPast: + """ + routed_experts (`torch.LongTensor` of shape `(batch_size, sequence_length, num_hidden_layers, num_experts_per_tok)`, *optional*): + Routed experts for each token in the sequence. Only used for router replay. + """ if (input_ids is None) ^ (inputs_embeds is not None): raise ValueError("You must specify exactly one of input_ids or inputs_embeds") if inputs_embeds is None: inputs_embeds = self.embed_tokens(input_ids) - if self.config._attn_implementation in ("flash_attention_2", "flash_attention_3"): - flat_position_ids = position_ids.view(-1) - seqlens = torch.cat( - [ - flat_position_ids[0:1], - flat_position_ids[:-1][(flat_position_ids == 0)[1:]] + 1, - flat_position_ids[-1:] + 1, - ] - ) - max_seqlen = seqlens.max().item() - cu_seqlens = seqlens.cumsum(dim=0, dtype=torch.int32) + if self.config._attn_implementation in ("flash_attention_2", "flash_attention_3", "fa4"): + cu_seqlens, max_seqlen = get_cu_seqlens_from_position_ids(position_ids) torch._dynamo.mark_dynamic(cu_seqlens, 0) else: max_seqlen = None @@ -227,12 +227,14 @@ def forward( hidden_states = inputs_embeds position_embeddings = self.rotary_emb(hidden_states, position_ids) - for decoder_layer in self.layers[: self.config.num_hidden_layers]: + for layer_idx, decoder_layer in enumerate(self.layers[: self.config.num_hidden_layers]): + routed_experts_layer = routed_experts[:, :, layer_idx, :] if routed_experts is not None else None hidden_states = decoder_layer( hidden_states, position_embeddings=position_embeddings, cu_seqlens=cu_seqlens, max_seqlen=max_seqlen, + routed_experts=routed_experts_layer, ) hidden_states = self.norm(hidden_states) @@ -324,7 +326,7 @@ def load_balancing_loss_func( @auto_docstring class Qwen3MoeForCausalLM(Qwen3MoePreTrainedModel, GenerationMixin): - _tied_weights_keys = ["lm_head.weight"] + _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": "colwise_rep"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} @@ -360,16 +362,23 @@ def forward( output_router_logits: Optional[bool] = None, cache_position: Optional[torch.LongTensor] = None, logits_to_keep: Union[int, torch.Tensor] = 0, - temperature: Optional[float] = 1.0, + temperature: Union[torch.Tensor, None] = None, + routed_experts: Optional[torch.LongTensor] = None, **kwargs: Unpack[TransformersKwargs], ) -> PrimeLmOutput: r""" + cache_position (`torch.LongTensor` of shape `(sequence_length)`, *optional*): + Indices of input tokens in the KV cache. Accepted only for HuggingFace API + compatibility — prime-rl asserts `use_cache is None` since training does not + perform autoregressive decoding, so this argument is unused. labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): Labels for computing the masked language modeling loss. Indices should either be in `[0, ..., config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`. - temperature (`float`, *optional*, defaults to 1.0): - The temperature to use for the logprobs/entropy. + temperature (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*): + Per-token temperatures for logprobs/entropy computation. + routed_experts (`torch.LongTensor` of shape `(batch_size, sequence_length, num_hidden_layers, num_experts_per_tok)`, *optional*): + Routed experts for each token in the sequence. Only used for router replay. Example: @@ -404,6 +413,7 @@ def forward( input_ids=input_ids, position_ids=position_ids, inputs_embeds=inputs_embeds, + routed_experts=routed_experts, ) hidden_states = outputs.last_hidden_state diff --git a/src/prime_rl/trainer/multi_ckpt.py b/src/prime_rl/trainer/multi_ckpt.py index c73791cb9a0..ec55e6d26de 100644 --- a/src/prime_rl/trainer/multi_ckpt.py +++ b/src/prime_rl/trainer/multi_ckpt.py @@ -13,12 +13,12 @@ import torch.distributed as dist from torch.distributed.checkpoint.stateful import Stateful +from prime_rl.configs.trainer import CheckpointConfig from prime_rl.trainer.ckpt import CheckpointManager -from prime_rl.trainer.config import CheckpointConfig -from prime_rl.trainer.runs import Progress, get_runs +from prime_rl.trainer.runs import Progress, get_multi_run_manager from prime_rl.trainer.world import get_world from prime_rl.utils.logger import get_logger -from prime_rl.utils.utils import resolve_latest_ckpt_step +from prime_rl.utils.pathing import get_stable_ckpt_steps if TYPE_CHECKING: from prime_rl.trainer.optim import MultiLoRAOptimizer @@ -72,12 +72,12 @@ class MultiCheckpointManager: def __init__(self, output_dir: Path): self.output_dir = output_dir - self.runs = get_runs() + self.multi_run_manager = get_multi_run_manager() self.world = get_world() self.logger = get_logger() - self.managers: list[CheckpointManager | None] = [None] * self.runs.max_runs - self.runs.register_deletion_hook(self._run_deletion_hook) - self.runs.register_creation_hook(self._run_creation_hook) + self.managers: list[CheckpointManager | None] = [None] * self.multi_run_manager.max_runs + self.multi_run_manager.register_deletion_hook(self._run_deletion_hook) + self.multi_run_manager.register_creation_hook(self._run_creation_hook) def _run_deletion_hook(self, idx: int, run_id: str) -> None: self.managers[idx] = None @@ -86,7 +86,7 @@ def _run_creation_hook(self, idx: int, run_id: str) -> None: self.managers[idx] = self._maybe_create_manager(idx) def _maybe_create_manager(self, idx: int) -> CheckpointManager | None: - ckpt_config = self.runs.config[idx].ckpt + ckpt_config = self.multi_run_manager.config[idx].ckpt if ckpt_config is None: return None @@ -95,14 +95,14 @@ def _maybe_create_manager(self, idx: int) -> CheckpointManager | None: keep_last=ckpt_config.keep_last, keep_interval=ckpt_config.keep_interval, ) - run_dir = self.runs.get_run_dir(idx) + run_dir = self.multi_run_manager.get_run_dir(idx) manager = CheckpointManager(run_dir, config) self.managers[idx] = manager return manager def _should_save(self, idx: int, step: int) -> bool: """Determine if a checkpoint should be saved for a given run and step.""" - ckpt_config = self.runs.config[idx].ckpt + ckpt_config = self.multi_run_manager.config[idx].ckpt if ckpt_config is None or ckpt_config.interval is None: return False if step <= 0 or step % ckpt_config.interval != 0: @@ -115,8 +115,8 @@ def save( optimizer: "MultiLoRAOptimizer", scheduler: "MultiLoRAScheduler", ) -> None: - for idx in self.runs.used_idxs: - step = self.runs.progress[idx].step + for idx in self.multi_run_manager.used_idxs: + step = self.multi_run_manager.progress[idx].step if not self._should_save(idx, step): continue @@ -124,12 +124,14 @@ def save( # We have a very wide try-except because we dont want to crash the trainer over one run having issues try: - model_state_dict = {k: v.data.detach().clone() for k, v in self.runs.get_named_parameters_for_run(idx)} + model_state_dict = { + k: v.data.detach().clone() for k, v in self.multi_run_manager.get_named_parameters_for_run(idx) + } run_state = RunState( model_state_dict, optimizer.optimizers[idx], scheduler.schedulers[idx], - self.runs.progress[idx], + self.multi_run_manager.progress[idx], ) ckpt_path = manager.get_ckpt_path(step) ckpt_path.mkdir(parents=True, exist_ok=True) @@ -141,7 +143,7 @@ def save( # Copy broadcast folder to checkpoint # This way, we only need to save the checkpoint folder if self.world.is_master: - run_dir = self.runs.get_run_dir(idx) + run_dir = self.multi_run_manager.get_run_dir(idx) broadcast_src = run_dir / "broadcasts" / f"step_{step}" weight_dst = run_dir / "checkpoints" / f"step_{step}" / "weight" try: @@ -150,13 +152,26 @@ def save( self.logger.error( f"Broadcast folder not found for run {idx} at step {step}. Looking for it in {broadcast_src}" ) + dist.barrier() + manager.mark_stable(step) + manager.ckpt_steps.append(step) except FileNotFoundError: self.logger.warning(f"Run {idx} deleted during checkpoint, skipping") except Exception as e: self.logger.error(f"Error checkpointing run {idx}: {e}") dist.barrier() - manager.mark_stable(step) - manager.ckpt_steps.append(step) + # If the run is deleted, remove the run directory + # This is avoid the creation of zombie runs when the directory is deleted while we are checkpointing which recreates the directory + # Ideally we move this to discover but lets have here for now + if ( + self.multi_run_manager.get_orchestrator_config(self.multi_run_manager.idx_2_id[idx]) is None + and self.world.is_master + ): + try: + self.logger.warning(f"Run {idx} deleted during checkpoint, removing run directory") + shutil.rmtree(self.multi_run_manager.get_run_dir(idx)) + except Exception as e: + self.logger.error(f"Error removing run directory for run {idx}: {e}") dist.barrier() def load_run( @@ -165,35 +180,39 @@ def load_run( optimizer: "MultiLoRAOptimizer", scheduler: "MultiLoRAScheduler", ) -> bool: - if self.runs.config[idx].ckpt is None or self.runs.config[idx].ckpt.resume_step is None: + if ( + self.multi_run_manager.config[idx].ckpt is None + or self.multi_run_manager.config[idx].ckpt.resume_step is None + ): return False manager = self.managers[idx] if manager is None: return False - step = self.runs.config[idx].ckpt.resume_step + step = self.multi_run_manager.config[idx].ckpt.resume_step if step == -1: - step = resolve_latest_ckpt_step(manager.ckpt_dir) - if step is None: + stable_steps = get_stable_ckpt_steps(manager.ckpt_dir) + if not stable_steps: return False + step = max(stable_steps) try: - model_state_dict = dict(self.runs.get_named_parameters_for_run(idx)) + model_state_dict = dict(self.multi_run_manager.get_named_parameters_for_run(idx)) run_state = RunState( model_state_dict, optimizer.optimizers[idx], scheduler.schedulers[idx], - self.runs.progress[idx], + self.multi_run_manager.progress[idx], ) ckpt_path = manager.get_ckpt_path(step) if not ckpt_path.exists(): raise FileNotFoundError(f"Checkpoint not found at {ckpt_path}") self.logger.info(f"Loading checkpoint from {ckpt_path}") - state_dict = torch.load(ckpt_path / f"rank_{self.world.rank}.pt") + state_dict = torch.load(ckpt_path / f"rank_{self.world.rank}.pt", weights_only=False) run_state.load_state_dict(state_dict) - self.logger.info(f"Resumed run {self.runs.idx_2_id[idx]} from step {step}") + self.logger.info(f"Resumed run {self.multi_run_manager.idx_2_id[idx]} from step {step}") return True except Exception as e: self.logger.error(f"Error loading checkpoint for run {idx}: {e}") @@ -202,7 +221,7 @@ def load_run( def maybe_clean(self) -> None: if not self.world.is_master: return - for idx in self.runs.used_idxs: + for idx in self.multi_run_manager.used_idxs: if self.managers[idx] is None: continue self.managers[idx].maybe_clean() diff --git a/src/prime_rl/trainer/optim.py b/src/prime_rl/trainer/optim.py index 5f7bd3e2170..d8465cc388d 100644 --- a/src/prime_rl/trainer/optim.py +++ b/src/prime_rl/trainer/optim.py @@ -1,42 +1,147 @@ +import copy import time from typing import Callable +import torch from dion import Muon from torch import nn +from torch.distributed.tensor import DTensor from torch.optim import SGD, AdamW, Optimizer -from prime_rl.trainer.config import OptimizerConfigType +from prime_rl.configs.trainer import OptimizerConfig from prime_rl.trainer.parallel_dims import ParallelDims -from prime_rl.trainer.runs import get_runs +from prime_rl.trainer.runs import get_multi_run_manager +from prime_rl.trainer.sign_sgd import SignSGD from prime_rl.trainer.world import get_world from prime_rl.utils.logger import get_logger +class CPUOffloadOptimizer: + """Wraps an optimizer to keep states on CPU, moving to GPU only for step(). + + Unlike FSDP's CPUOffload which offloads weights too, this keeps weights on GPU. + With activation checkpointing, activations and optimizer states are never on GPU + at the same time: peak memory becomes max(activations, opt_states) instead of sum. + """ + + def __init__(self, optimizer: Optimizer, pin_memory: bool = True): + self.optimizer = optimizer + self.pin_memory = pin_memory + self._initialized = False + + def _move_states(self, device: str): + """Move optimizer states to CPU or back to GPU (matching each parameter's device).""" + for p in self.optimizer.state: + state = self.optimizer.state[p] + for k, v in state.items(): + if isinstance(v, DTensor): + local_tensor = v._local_tensor + if device == "cpu": + non_blocking = not self.pin_memory + new_local = local_tensor.to("cpu", non_blocking=non_blocking) + if self.pin_memory and not new_local.is_pinned(): + new_local = new_local.pin_memory() + else: + new_local = local_tensor.to(device, non_blocking=True) + new_dtensor = copy.copy(v) + new_dtensor._local_tensor = new_local + state[k] = new_dtensor + elif isinstance(v, torch.Tensor): + if device == "cpu": + non_blocking = not self.pin_memory + cpu_tensor = v.to("cpu", non_blocking=non_blocking) + if self.pin_memory and not cpu_tensor.is_pinned(): + cpu_tensor = cpu_tensor.pin_memory() + state[k] = cpu_tensor + else: + state[k] = v.to(device, non_blocking=True) + + def step(self, closure=None): + # First step initializes states on GPU - offload after + if not self._initialized: + result = self.optimizer.step(closure) + self._move_states("cpu") + self._initialized = True + return result + + # Move states to GPU + self._move_states("cuda") + + # Run optimizer step + result = self.optimizer.step(closure) + + # Move states back to CPU + self._move_states("cpu") + + return result + + def zero_grad(self, set_to_none: bool = True): + self.optimizer.zero_grad(set_to_none=set_to_none) + + def state_dict(self): + # Move to GPU temporarily for consistent state dict + if self._initialized: + self._move_states("cuda") + torch.cuda.synchronize() + sd = self.optimizer.state_dict() + if self._initialized: + self._move_states("cpu") + return sd + + def load_state_dict(self, state_dict): + self.optimizer.load_state_dict(state_dict) + self._move_states("cpu") + self._initialized = True + + @property + def param_groups(self): + return self.optimizer.param_groups + + @param_groups.setter + def param_groups(self, value): + self.optimizer.param_groups = value + + @property + def state(self): + return self.optimizer.state + + @property + def base_optimizer(self) -> Optimizer: + return self.optimizer + + def setup_optimizer( - config: OptimizerConfigType, + config: OptimizerConfig, named_params: list[tuple[str, nn.Parameter]], parallel_dims: ParallelDims, lora: bool = False, -) -> Optimizer: + cpu_offload: bool = False, +) -> Optimizer | CPUOffloadOptimizer: if lora: - # Wait for run 0 to be created in the runs system + # Wait for run 0 to be created in the multi run manager # Otherwise, the creation will reset the parameters - runs = get_runs() + multi_run_manager = get_multi_run_manager() world = get_world() logger = get_logger() - while 0 not in runs.idx_2_id: + while 0 not in multi_run_manager.idx_2_id: if world.is_master: - runs.check_for_changes() - runs.sync_runs() - logger.info(f"Waiting for run 0 to be created {runs.id_2_idx=}") + multi_run_manager.discover_runs() + multi_run_manager.synchronize_state() + logger.info(f"Waiting for run 0 to be created {multi_run_manager.id_2_idx=}") time.sleep(1) - named_params = runs.get_named_parameters_for_run(0) + named_params = multi_run_manager.get_named_parameters_for_run(0) + + optimizer = _create_optimizer(config, named_params, parallel_dims) - return _create_optimizer(config, named_params, parallel_dims) + if cpu_offload: + get_logger().info("Wrapping optimizer with CPUOffloadOptimizer for optimizer state CPU offloading") + return CPUOffloadOptimizer(optimizer) + + return optimizer def _create_optimizer( - config: OptimizerConfigType, + config: OptimizerConfig, named_params: list[tuple[str, nn.Parameter]], parallel_dims: ParallelDims, lr: float | None = None, @@ -62,10 +167,16 @@ def _create_optimizer( ) case "muon": return _create_muon_optimizer(config, named_params, parallel_dims, lr) + case "sign_sgd": + return SignSGD( + params=[p for _, p in named_params], + lr=lr, + weight_decay=config.weight_decay, + ) def _create_muon_optimizer( - config: OptimizerConfigType, + config: OptimizerConfig, named_params: list[tuple[str, nn.Parameter]], parallel_dims: ParallelDims, lr: float | None = None, @@ -129,13 +240,15 @@ def muon_enabled(n, p): param_groups.append(dict(params=adamw_params, algorithm="adamw", lr=lr, weight_decay=config.weight_decay)) if parallel_dims.dp_shard_enabled or parallel_dims.cp_enabled: - distributed_mesh = parallel_dims.world_mesh["dp_shard_cp"] + distributed_mesh = parallel_dims.get_mesh("dp_shard_cp") else: distributed_mesh = parallel_dims.world_mesh optimizer = Muon( params=param_groups, lr=lr, + mu=config.mu, + betas=(config.betas1, config.betas2), weight_decay=config.weight_decay, adjust_lr="rms_norm", distributed_mesh=distributed_mesh, @@ -146,18 +259,18 @@ def muon_enabled(n, p): class MultiLoRAOptimizer: - def __init__(self, config: OptimizerConfigType, parallel_dims: ParallelDims): + def __init__(self, config: OptimizerConfig, parallel_dims: ParallelDims): self.config = config self.parallel_dims = parallel_dims - self.runs = get_runs() + self.multi_run_manager = get_multi_run_manager() self.logger = get_logger() - self.optimizers: list[Optimizer | None] = [None] * self.runs.max_runs + self.optimizers: list[Optimizer | None] = [None] * self.multi_run_manager.max_runs self._post_creation_callbacks: list[Callable[[Optimizer, int], None]] = [] # Register creation hook for optimizer setup - # The Runs class handles parameter reset internally when new runs are created - self.runs.register_creation_hook(self.optimizer_creation_hook) + # The MultiRunManager class handles parameter reset internally when new runs are created + self.multi_run_manager.register_creation_hook(self.optimizer_creation_hook) def register_post_creation_callback( self, callback: Callable[[Optimizer, int], None], index: int | None = None @@ -174,10 +287,10 @@ def register_post_creation_callback( self._post_creation_callbacks.insert(index, callback) def optimizer_creation_hook(self, idx: int, run_id: str) -> None: - # Get named parameters for this run from the Runs system - named_params = self.runs.get_named_parameters_for_run(idx) + # Get named parameters for this run from the MultiRunManager system + named_params = self.multi_run_manager.get_named_parameters_for_run(idx) - lr = self.runs.config[idx].optim.lr + lr = self.multi_run_manager.config[idx].optim.lr self.optimizers[idx] = _create_optimizer(self.config, named_params, self.parallel_dims, lr) # Call post-creation callbacks in order @@ -185,16 +298,16 @@ def optimizer_creation_hook(self, idx: int, run_id: str) -> None: callback(self.optimizers[idx], idx) def step(self): - for idx in self.runs.ready_to_update_idxs: + for idx in self.multi_run_manager.ready_to_update_idxs: self.optimizers[idx].step() def zero_grad(self): - for idx in self.runs.ready_to_update_idxs: + for idx in self.multi_run_manager.ready_to_update_idxs: self.optimizers[idx].zero_grad() def get_current_lr(self, idx: int | None = None) -> float: if idx is None: - for idx in self.runs.ready_to_update_idxs: + for idx in self.multi_run_manager.ready_to_update_idxs: return self.optimizers[idx].param_groups[0]["lr"] else: self.logger.warning("No runs are ready to update. Returning 0.0 for current learning rate.") @@ -203,5 +316,5 @@ def get_current_lr(self, idx: int | None = None) -> float: return self.optimizers[idx].param_groups[0]["lr"] -def setup_multi_optimizer(config: OptimizerConfigType, parallel_dims: ParallelDims) -> MultiLoRAOptimizer: +def setup_multi_optimizer(config: OptimizerConfig, parallel_dims: ParallelDims) -> MultiLoRAOptimizer: return MultiLoRAOptimizer(config, parallel_dims) diff --git a/src/prime_rl/trainer/parallel_dims.py b/src/prime_rl/trainer/parallel_dims.py index 3685266b37c..dbfe8557238 100644 --- a/src/prime_rl/trainer/parallel_dims.py +++ b/src/prime_rl/trainer/parallel_dims.py @@ -24,7 +24,7 @@ from torch._utils import _get_available_device_type from torch.distributed.device_mesh import DeviceMesh, init_device_mesh -from prime_rl.trainer.config import ModelConfig +from prime_rl.configs.trainer import ModelConfig from prime_rl.utils.logger import get_logger device_type = _get_available_device_type() or "cuda" @@ -37,36 +37,36 @@ class ParallelDims: dp_replicate: int dp_shard: int cp: int - tp: int pp: int ep: int world_size: int _world_mesh: DeviceMesh = None + _submeshes: dict = None def __post_init__(self): + self._submeshes = {} self._validate() def _validate(self): - dp_replicate, dp_shard, cp, tp, pp, ep = ( + dp_replicate, dp_shard, cp, pp, ep = ( self.dp_replicate, self.dp_shard, self.cp, - self.tp, self.pp, self.ep, ) - for d in (dp_replicate, cp, tp, pp, ep): + for d in (dp_replicate, cp, pp, ep): assert d >= 1, "Parallelism degree should be >= 1, except for dp_shard" assert dp_shard == -1 or dp_shard >= 1, " dp_shard must -1 or >=1." if dp_shard < 0: - self.dp_shard = dp_shard = self.world_size // (dp_replicate * cp * tp * pp) + self.dp_shard = dp_shard = self.world_size // (dp_replicate * cp * pp) assert dp_shard >= 1 - assert dp_replicate * dp_shard * cp * tp * pp == self.world_size, ( + assert dp_replicate * dp_shard * cp * pp == self.world_size, ( f"Invalid parallel dims: dp_replicate({dp_replicate}) * dp_shard({dp_shard}) * " - f"cp({cp}) * tp({tp}) * pp({pp}) != WORLD_SIZE({self.world_size})" + f"cp({cp}) * pp({pp}) != WORLD_SIZE({self.world_size})" ) if ep > 1: @@ -95,9 +95,8 @@ def _build_mesh_with_ep(self) -> DeviceMesh: dp_shard_mod_ep, dp_shard_in_ep, self.cp, - self.tp, ], - ["pp", "dp_replicate", "dp_shard_mod_ep", "dp_shard_in_ep", "cp", "tp"], + ["pp", "dp_replicate", "dp_shard_mod_ep", "dp_shard_in_ep", "cp"], ): # dp_shard_mod_ep is needed even if it's 1, whose FSDP wrapping # helps the MoE layers do mixed precision training @@ -136,10 +135,19 @@ def _build_mesh_with_ep(self) -> DeviceMesh: dp_cp_mesh_dim_names.append("cp") ep_mesh_dim_names.append("cp") - mesh[tuple(dp_mesh_dim_names)]._flatten(mesh_dim_name="dp") - mesh[tuple(dp_shard_cp_mesh_dim_names)]._flatten(mesh_dim_name="dp_shard_cp") - mesh[tuple(dp_cp_mesh_dim_names)]._flatten(mesh_dim_name="dp_cp") - mesh[tuple(ep_mesh_dim_names)]._flatten(mesh_dim_name="ep") + self._submeshes["dp"] = mesh[tuple(dp_mesh_dim_names)]._flatten(mesh_dim_name="dp") + self._submeshes["dp_shard_cp"] = mesh[tuple(dp_shard_cp_mesh_dim_names)]._flatten(mesh_dim_name="dp_shard_cp") + self._submeshes["dp_cp"] = mesh[tuple(dp_cp_mesh_dim_names)]._flatten(mesh_dim_name="dp_cp") + self._submeshes["ep"] = mesh[tuple(ep_mesh_dim_names)]._flatten(mesh_dim_name="ep") + + if self.dp_replicate_enabled: + parent = mesh[tuple(["dp_replicate"] + dp_shard_cp_mesh_dim_names)] + hsdp_tensor = parent.mesh.reshape(self.dp_replicate, -1) + self._submeshes["hsdp"] = DeviceMesh( + device_type, hsdp_tensor, mesh_dim_names=("dp_replicate", "dp_shard_cp") + ) + else: + self._submeshes["hsdp"] = self._submeshes["dp_shard_cp"] return mesh @@ -147,8 +155,8 @@ def _build_mesh_without_ep(self) -> DeviceMesh: dims = [] names = [] for d, name in zip( - [self.pp, self.dp_replicate, self.dp_shard, self.cp, self.tp], - ["pp", "dp_replicate", "dp_shard", "cp", "tp"], + [self.pp, self.dp_replicate, self.dp_shard, self.cp], + ["pp", "dp_replicate", "dp_shard", "cp"], ): if d > 1 or name == "dp_shard": dims.append(d) @@ -177,11 +185,22 @@ def _build_mesh_without_ep(self) -> DeviceMesh: dp_cp_mesh_dim_names.append("cp") if dp_mesh_dim_names != []: - mesh[tuple(dp_mesh_dim_names)]._flatten(mesh_dim_name="dp") + self._submeshes["dp"] = mesh[tuple(dp_mesh_dim_names)]._flatten(mesh_dim_name="dp") if dp_shard_cp_mesh_dim_names != []: - mesh[tuple(dp_shard_cp_mesh_dim_names)]._flatten(mesh_dim_name="dp_shard_cp") + self._submeshes["dp_shard_cp"] = mesh[tuple(dp_shard_cp_mesh_dim_names)]._flatten( + mesh_dim_name="dp_shard_cp" + ) if dp_cp_mesh_dim_names != []: - mesh[tuple(dp_cp_mesh_dim_names)]._flatten(mesh_dim_name="dp_cp") + self._submeshes["dp_cp"] = mesh[tuple(dp_cp_mesh_dim_names)]._flatten(mesh_dim_name="dp_cp") + + if self.dp_replicate_enabled: + parent = mesh[tuple(["dp_replicate"] + dp_shard_cp_mesh_dim_names)] + hsdp_tensor = parent.mesh.reshape(self.dp_replicate, -1) + self._submeshes["hsdp"] = DeviceMesh( + device_type, hsdp_tensor, mesh_dim_names=("dp_replicate", "dp_shard_cp") + ) + else: + self._submeshes["hsdp"] = self._submeshes["dp_shard_cp"] return mesh @@ -193,6 +212,12 @@ def world_mesh(self) -> DeviceMesh: self._world_mesh = self.build_mesh() return self._world_mesh + def get_mesh(self, name: str) -> DeviceMesh: + mesh = self.world_mesh # ensure lazy init has run + if name in self._submeshes: + return self._submeshes[name] + return mesh[name] + @property def dp_enabled(self): return self.dp_replicate > 1 or self.dp_shard > 1 @@ -217,10 +242,6 @@ def dp_cp_enabled(self): def fsdp_enabled(self): return self.dp_shard_enabled or self.cp_enabled - @property - def tp_enabled(self): - return self.tp > 1 - @property def pp_enabled(self): return self.pp > 1 @@ -235,17 +256,14 @@ def fsdp_gradient_divide_factor(self) -> int: @cached_property def non_data_parallel_size(self): - return self.cp * self.tp * self.pp + return self.cp * self.pp @cached_property def seq_len_divisor(self): - # Sequence Parallel requires that seq_len be divisible by TP degree. - # https://github.com/pytorch/torchtitan/pull/640#discussion_r1849481001 - # Context Parallel requires that seq_len be divisible by 2 * CP degree, # when load balancing is enabled (by default). # https://github.com/pytorch/pytorch/blob/4f62dcc/torch/distributed/tensor/experimental/_attention.py#L1246 - return self.tp * (self.cp * 2) + return self.cp * 2 @cached_property def logger(self): @@ -258,7 +276,6 @@ def get_parallel_dims(config: ModelConfig, seq_len: int | None = None) -> Parall dp_replicate=config.dp_replicate, dp_shard=-1, cp=config.cp, - tp=config.tp, pp=1, ep=config.ep, world_size=dist.get_world_size(), @@ -269,8 +286,7 @@ def get_parallel_dims(config: ModelConfig, seq_len: int | None = None) -> Parall raise ValueError( f"Sequence length ({seq_len}) must be divisible by " f"seq_len_divisor ({parallel_dims.seq_len_divisor}) for the given parallel dimensions. " - f"This requirement comes from context parallel (CP={config.cp}) and " - f"tensor parallel (TP={config.tp}) configurations." + f"This requirement comes from context parallel (CP={config.cp})." ) return parallel_dims diff --git a/src/prime_rl/trainer/perf.py b/src/prime_rl/trainer/perf.py index 906b81e1684..595a1d49480 100644 --- a/src/prime_rl/trainer/perf.py +++ b/src/prime_rl/trainer/perf.py @@ -12,16 +12,22 @@ class PerfCounter: """ - A class to count throughput (tokens/s) with a rolling window to obtain - precise throughput and MFU estimates. + Computes throughput (tokens/s) and MFU. - Inspired from https://github.com/pytorch/torchtitan/blob/4b3f2e41a084bf79a8540068ed525539d1244edd/torchtitan/utils.py#L119 + Two modes: + - Sliding window over full-step wall time: `count_tokens(tokens)` each step, + read smoothed values via `get_tokens_per_second()` / `get_mfu()`. + Time is measured between successive `count_tokens` calls (full step). + - Single point on a caller-provided duration: `get_step_tokens_per_second(tokens, fwd_bwd_time)` + / `get_step_mfu(tokens, fwd_bwd_time)`. No smoothing. + + Sliding window inspired by https://github.com/pytorch/torchtitan/blob/4b3f2e41a084bf79a8540068ed525539d1244edd/torchtitan/utils.py#L119 """ - def __init__(self, model: nn.Module, seq_len: int, window_size: int): + def __init__(self, model: nn.Module, seq_len: int, window_size: int = 10): self.window_size = window_size - self.tokens = [] - self.times = [] + self.tokens: list[int] = [] + self.times: list[float] = [] self.model = model self._world = get_world() @@ -36,7 +42,8 @@ def __init__(self, model: nn.Module, seq_len: int, window_size: int): self.num_params = self._get_num_params(model, exclude_embedding=not model.config.tie_word_embeddings) self.num_flop_per_token = self._get_num_flop_per_token(model.config, seq_len=seq_len) - def count_tokens(self, tokens: int): + def count_tokens(self, tokens: int) -> None: + """Push a step into the sliding window. Time is recorded internally.""" self.tokens.append(tokens) self.times.append(time.perf_counter()) if len(self.tokens) > self.window_size: @@ -52,6 +59,16 @@ def get_mfu(self) -> float | None: tokens_per_second = self.get_tokens_per_second() if tokens_per_second is None: return None + return self._mfu_from_tps(tokens_per_second) + + def get_step_tokens_per_second(self, tokens: int, fwd_bwd_time: float) -> float: + """Single-step throughput from a caller-provided duration.""" + return tokens / fwd_bwd_time + + def get_step_mfu(self, tokens: int, fwd_bwd_time: float) -> float: + return self._mfu_from_tps(self.get_step_tokens_per_second(tokens, fwd_bwd_time)) + + def _mfu_from_tps(self, tokens_per_second: float) -> float: return 100 * self.num_flop_per_token * tokens_per_second / self.gpu_peak_flops / self._world.world_size def _get_peak_flops(self, device_name: str) -> float: @@ -75,6 +92,15 @@ def _get_peak_flops(self, device_name: str) -> float: if "B200" in device_name: # https://nvdam.widen.net/s/wwnsxrhm2w/blackwell-datasheet-3384703 return 2.25e15 # This is half of the FLOPS reported in torchtitan + # AMD Instinct GPUs + if "MI300X" in device_name: + # https://www.amd.com/en/products/accelerators/instinct/mi300/mi300x.html + # Peak BF16: 1307.4 TFLOPS (matrix) + return 1307.4e12 + if "MI325X" in device_name: + # https://www.amd.com/en/products/accelerators/instinct/mi300/mi325x.html + # Peak BF16: 1307.4 TFLOPS (matrix) - same compute dies as MI300X, more HBM3e + return 1307.4e12 else: self._logger.warning(f"Peak FLOPS undefined for `{device_name}`. Falling back to A100 (312 TFLOPS)") return 312e12 @@ -82,9 +108,13 @@ def _get_peak_flops(self, device_name: str) -> float: @staticmethod def get_active_mm_params(config: PretrainedConfig) -> float: """Get number of active parameters per token involved in matmuls""" + # Handle VLM models with nested text_config (e.g., Qwen3-VL) + if hasattr(config, "text_config"): + config = config.text_config + vocab_size = config.vocab_size hidden_size = config.hidden_size - intermediate_size = config.intermediate_size + intermediate_size = getattr(config, "intermediate_size", getattr(config, "moe_intermediate_size", 0)) head_dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads) num_attention_heads = config.num_attention_heads num_hidden_layers = config.num_hidden_layers @@ -122,16 +152,16 @@ def get_active_mm_params(config: PretrainedConfig) -> float: sparse_mlp_params = 0 # Some MoE models (e.g. DeepSeek) use moe_intermediate_size, others (e.g. Granite) just use intermediate_size - moe_intermediate_size = getattr(config, "moe_intermediate_size", intermediate_size) - if hasattr(config, "num_shared_experts"): # Shared experts + moe_intermediate_size = getattr(config, "moe_intermediate_size", None) or intermediate_size + if hasattr(config, "num_shared_experts") and config.num_shared_experts: # Shared experts sparse_mlp_params += num_sparse_layers * config.num_shared_experts * 3 * moe_intermediate_size * hidden_size - if hasattr(config, "num_experts_per_tok"): # Routed experts + if hasattr(config, "num_experts_per_tok") and config.num_experts_per_tok: # Routed experts sparse_mlp_params += ( num_sparse_layers * config.num_experts_per_tok * 3 * moe_intermediate_size * hidden_size ) if hasattr(config, "n_routed_experts"): # DeepSeek Router sparse_mlp_params += num_sparse_layers * config.n_routed_experts * hidden_size - elif hasattr(config, "num_experts"): # Qwen Router + elif hasattr(config, "num_experts") and config.num_experts is not None: # Qwen Router sparse_mlp_params += num_sparse_layers * config.num_experts * hidden_size else: sparse_mlp_params = 0 @@ -142,6 +172,10 @@ def get_active_mm_params(config: PretrainedConfig) -> float: return q_params + kv_params + o_params + dense_mlp_params + sparse_mlp_params + lm_head_params def _get_num_flop_per_token(self, model_config: PretrainedConfig, seq_len: int) -> int: + # Handle VLM models with nested text_config (e.g., Qwen3-VL) + if hasattr(model_config, "text_config"): + model_config = model_config.text_config + l, h, q, t = ( # noqa: E741 model_config.num_hidden_layers, model_config.num_attention_heads, diff --git a/src/prime_rl/trainer/rl/broadcast/__init__.py b/src/prime_rl/trainer/rl/broadcast/__init__.py index 92b941fa4f2..e419bda2f29 100644 --- a/src/prime_rl/trainer/rl/broadcast/__init__.py +++ b/src/prime_rl/trainer/rl/broadcast/__init__.py @@ -2,15 +2,14 @@ import torch -from prime_rl.trainer.config import LoRAConfig +from prime_rl.configs.trainer import LoRAConfig, WeightBroadcastConfig from prime_rl.trainer.rl.broadcast.base import WeightBroadcast from prime_rl.trainer.rl.broadcast.filesystem import FileSystemWeightBroadcast from prime_rl.trainer.rl.broadcast.nccl import NCCLWeightBroadcast -from prime_rl.trainer.rl.config import WeightBroadcastConfigType def setup_weight_broadcast( - output_dir: Path, config: WeightBroadcastConfigType, lora_config: LoRAConfig | None = None + output_dir: Path, config: WeightBroadcastConfig, lora_config: LoRAConfig | None = None ) -> WeightBroadcast: if config.type == "nccl": return NCCLWeightBroadcast(output_dir, config, torch.cuda.current_device()) diff --git a/src/prime_rl/trainer/rl/broadcast/base.py b/src/prime_rl/trainer/rl/broadcast/base.py index c444266c1f7..68e502a097f 100644 --- a/src/prime_rl/trainer/rl/broadcast/base.py +++ b/src/prime_rl/trainer/rl/broadcast/base.py @@ -3,7 +3,7 @@ import torch.nn as nn -from prime_rl.trainer.config import LoRAConfig +from prime_rl.configs.trainer import LoRAConfig from prime_rl.utils.logger import get_logger diff --git a/src/prime_rl/trainer/rl/broadcast/filesystem.py b/src/prime_rl/trainer/rl/broadcast/filesystem.py index f9b2f3c2ad7..55a92c832d7 100644 --- a/src/prime_rl/trainer/rl/broadcast/filesystem.py +++ b/src/prime_rl/trainer/rl/broadcast/filesystem.py @@ -6,12 +6,11 @@ import torch.nn as nn from torch.distributed.tensor import DTensor -from prime_rl.trainer.config import LoRAConfig +from prime_rl.configs.trainer import FileSystemWeightBroadcastConfig, LoRAConfig from prime_rl.trainer.lora import save_lora_config from prime_rl.trainer.models import PreTrainedModelPrimeRL from prime_rl.trainer.rl.broadcast.base import WeightBroadcast -from prime_rl.trainer.rl.config import FileSystemWeightBroadcastConfig -from prime_rl.trainer.runs import get_runs +from prime_rl.trainer.runs import get_multi_run_manager from prime_rl.trainer.utils import maybe_clean from prime_rl.trainer.weights import ( gather_weights_on_master, @@ -31,7 +30,7 @@ def __init__( self.save_format: Literal["safetensors", "torch"] = config.save_format self.save_sharded = config.save_sharded if lora_config is None else False self.world = get_world() - self.runs = get_runs() + self.multi_run_manager = get_multi_run_manager() self.logger.debug( f"Filesystem broadcast initialized (save_format={config.save_format}, save_sharded={self.save_sharded})" ) @@ -46,14 +45,21 @@ def broadcast_weights(self, model: nn.Module, step: int) -> None: state_dict = gather_weights_on_master(model, is_master=self.world.is_master) if isinstance(model, PreTrainedModelPrimeRL) and model.is_prime_state_dict(state_dict): model.convert_to_hf(state_dict) + else: + # For regular transformers models, revert internal format to original HF hub format + from transformers.core_model_loading import revert_weight_conversion - for idx in self.runs.ready_to_update_idxs: - self.logger.debug(f"Broadcasting weights for run {idx} (ready_to_update={self.runs.ready_to_update[idx]})") + state_dict = revert_weight_conversion(model, state_dict) + + for idx in self.multi_run_manager.ready_to_update_idxs: + self.logger.debug( + f"Broadcasting weights for run {idx} (ready_to_update={self.multi_run_manager.ready_to_update[idx]})" + ) if adapter_only: - # For adapter-only, Runs creates state dict directly for each run + # For adapter-only, MultiRunManager creates state dict directly for each run # All ranks must participate in DTensor gathering, but only master saves - state_dict = self.runs.get_state_dict_for_run(idx) + state_dict = self.multi_run_manager.get_state_dict_for_run(idx) for key, value in state_dict.items(): if isinstance(value, DTensor): value = value.full_tensor() @@ -64,28 +70,36 @@ def broadcast_weights(self, model: nn.Module, step: int) -> None: if self.world.is_master: try: save_dir = get_step_path( - get_broadcast_dir(self.runs.get_run_dir(idx)), self.runs.progress[idx].step + get_broadcast_dir(self.multi_run_manager.get_run_dir(idx)), + self.multi_run_manager.progress[idx].step, ) save_dir.mkdir(parents=True, exist_ok=True) self.logger.debug(f"Saving weights for run {idx} to {save_dir}") save_state_dict(state_dict, save_dir, self.save_format, self.save_sharded, adapter=adapter_only) if adapter_only: - save_lora_config(self.lora_config, model, save_dir) + orch_lora = self.multi_run_manager.config[idx].model.lora + save_lora_config( + model, + save_dir, + rank=orch_lora.rank, + alpha=orch_lora.alpha, + dropout=self.lora_config.dropout, + ) self._notify_orchestrator(save_dir) # If the run is deleted, remove the run directory # This is avoid the creation of zombie runs when the directory is deleted while we are broadcasting which recreates the directory - if self.runs.get_orchestrator_config(self.runs.idx_2_id[idx]) is None: - shutil.rmtree(self.runs.get_run_dir(idx)) + if self.multi_run_manager.get_orchestrator_config(self.multi_run_manager.idx_2_id[idx]) is None: + shutil.rmtree(self.multi_run_manager.get_run_dir(idx)) except FileNotFoundError: self.logger.warning(f"Run {idx} is deleted, skipping") except Exception as e: self.logger.error(f"Error broadcasting weights for run {idx}: {e}") finally: - self.runs.ready_to_update[idx] = False + self.multi_run_manager.ready_to_update[idx] = False if self.world.is_master: self.logger.debug(f"Weights broadcasted in {time.perf_counter() - start_time:.2f}s") @@ -96,10 +110,10 @@ def _notify_orchestrator(self, save_dir: Path): stable_file.touch() def maybe_clean(self, max_async_level: int, interval_to_keep: int | None): - for idx in self.runs.used_idxs: + for idx in self.multi_run_manager.used_idxs: maybe_clean( - get_broadcast_dir(self.runs.get_run_dir(idx)), - self.runs.progress[idx].step, + get_broadcast_dir(self.multi_run_manager.get_run_dir(idx)), + self.multi_run_manager.progress[idx].step, max_async_level, interval_to_keep, ) diff --git a/src/prime_rl/trainer/rl/broadcast/nccl.py b/src/prime_rl/trainer/rl/broadcast/nccl.py index 5f5ff4df9ef..13a887308ab 100644 --- a/src/prime_rl/trainer/rl/broadcast/nccl.py +++ b/src/prime_rl/trainer/rl/broadcast/nccl.py @@ -1,7 +1,7 @@ import pickle import time from pathlib import Path -from typing import Generator, cast +from typing import Callable, Generator, cast import torch import torch.nn as nn @@ -10,15 +10,17 @@ from vllm.distributed.device_communicators.pynccl import PyNcclCommunicator from vllm.distributed.utils import StatelessProcessGroup +from prime_rl.configs.trainer import NCCLWeightBroadcastConfig from prime_rl.trainer.models import PreTrainedModelPrimeRL from prime_rl.trainer.rl.broadcast.base import WeightBroadcast -from prime_rl.trainer.rl.config import NCCLWeightBroadcastConfig -from prime_rl.trainer.runs import get_runs +from prime_rl.trainer.runs import get_multi_run_manager from prime_rl.trainer.utils import get_world from prime_rl.trainer.weights import get_max_layer_num from prime_rl.utils.logger import get_logger +from prime_rl.utils.nccl import disable_nccl_p2p_if_unavailable from prime_rl.utils.pathing import sync_wait_for_path from prime_rl.utils.utils import get_broadcast_dir, get_step_path +from prime_rl.utils.vlm import get_layer_prefix NCCL_READY_MARKER = "NCCL_READY" @@ -67,22 +69,46 @@ def broadcast_state_dict(state_dict: dict[str, Tensor], communicator: PyNcclComm def filter_state_dict_by_layers( - state_dict: dict[str, torch.Tensor], num_layers: int + state_dict: dict[str, torch.Tensor], num_layers: int, layer_prefix: str ) -> Generator[tuple[int, dict[str, torch.Tensor]], None, None]: - """Yield a generator of state dicts for each layer as well as the remaining weights.""" - yield 0, {key: value for key, value in state_dict.items() if "model.layers" not in key} + """Yield non-layer weights first, then each layer's weights. - for i in range(1, num_layers + 1): # +1 because layer indices start from 1 + Yields (layer_idx, layer_state_dict) where layer_idx is -1 for the non-layer + dict and the actual layer index (0, 1, ...) for layer dicts. + """ + yield -1, {key: value for key, value in state_dict.items() if not key.startswith(layer_prefix)} + + for i in range(num_layers): yield ( i, - { - key: value - for key, value in state_dict.items() - if key.startswith(f"model.layers.{i}.") or key == f"model.layers.{i}" - }, + {key: value for key, value in state_dict.items() if key.startswith(f"{layer_prefix}{i}.")}, ) +def preprocess_layer_checkpoint( + model: nn.Module, + layer_state_dict: dict[str, Tensor], + layer_idx: int, +) -> dict[str, Tensor]: + if isinstance(model, PreTrainedModelPrimeRL) and model.is_prime_state_dict(layer_state_dict): + model.convert_layer_to_hf(layer_state_dict, layer_idx) + return layer_state_dict + + from transformers.core_model_loading import revert_weight_conversion + + return revert_weight_conversion(model, layer_state_dict) + + +def preprocess_layer_quantized( + model: nn.Module, + layer_state_dict: dict[str, Tensor], + layer_idx: int, +) -> dict[str, Tensor]: + if layer_idx < 0: + return layer_state_dict + return model.convert_layer_to_vllm_kernel(layer_state_dict, layer_idx, quantize_fp8=True) + + class NCCLWeightBroadcastSender: def __init__( self, @@ -93,12 +119,15 @@ def __init__( device: int | str | torch.device, timeout: int, dtype: torch.dtype = torch.bfloat16, + quantize_in_weight_transfer: bool = False, ): self.logger = get_logger() self.world = get_world() self.dtype = dtype + self.quantize_in_weight_transfer = quantize_in_weight_transfer if self.world.is_master: + disable_nccl_p2p_if_unavailable() # Trainer is on rank 0 in process group with all inference GPUs pg = StatelessProcessGroup.create( host=host, port=port, rank=rank, world_size=world_size, store_timeout=timeout @@ -112,26 +141,31 @@ def __init__( def broadcast_weights(self, model: nn.Module, step: int) -> None: """Broadcast the state dict of a model into the inference pool using NCCL.""" state_dict = model.state_dict() - num_layers = get_max_layer_num(state_dict) + layer_prefix = get_layer_prefix(model.config) + num_layers = get_max_layer_num(state_dict, layer_prefix) num_state_dict_to_send = num_layers + 1 # we send all layer plus the remaining weights if self.world.is_master: broadcast_integer(num_state_dict_to_send, self.communicator) self.logger.debug(f"Broadcasting {num_state_dict_to_send} layer state dicts") + preprocess_fn: Callable[[nn.Module, dict[str, Tensor], int], dict[str, Tensor]] + if self.quantize_in_weight_transfer: + preprocess_fn = preprocess_layer_quantized + else: + preprocess_fn = preprocess_layer_checkpoint - for layer_id, state_dict in filter_state_dict_by_layers(state_dict, num_layers): - for key, value in list(state_dict.items()): - if isinstance(value, DTensor): - value = cast(DTensor, value.to(self.dtype)).full_tensor() - state_dict[key] = value - - # Convert PrimeRL format to HF format for this layer if needed - if isinstance(model, PreTrainedModelPrimeRL) and model.is_prime_state_dict(state_dict): - model.convert_layer_to_hf(state_dict, layer_id) - + for layer_id, layer_state_dict in filter_state_dict_by_layers(state_dict, num_layers, layer_prefix): + layer_state_dict = self._resolve_dtensors(layer_state_dict) + layer_state_dict = preprocess_fn(model, layer_state_dict, layer_id) if self.world.is_master: - broadcast_state_dict(state_dict, self.communicator) + broadcast_state_dict(layer_state_dict, self.communicator) + + def _resolve_dtensors(self, state_dict: dict[str, Tensor]) -> dict[str, Tensor]: + for key, value in list(state_dict.items()): + if isinstance(value, DTensor): + state_dict[key] = cast(DTensor, value.to(self.dtype)).full_tensor() + return state_dict class NCCLWeightBroadcast(WeightBroadcast): @@ -147,9 +181,16 @@ def __init__( super().__init__(output_dir) self.logger = get_logger() self.world = get_world() - self.runs = get_runs() + self.multi_run_manager = get_multi_run_manager() self.nccl_broadcast_sender = NCCLWeightBroadcastSender( - config.host, config.port, 0, config.inference_world_size + 1, device, config.timeout, dtype + config.host, + config.port, + 0, + config.inference_world_size + 1, + device, + config.timeout, + dtype, + quantize_in_weight_transfer=config.quantize_in_weight_transfer, ) @torch.no_grad() @@ -157,43 +198,62 @@ def broadcast_weights(self, model: nn.Module, step: int) -> None: """Broadcast the state dict of a model into the inference pool using NCCL and notifies the orchestrator.""" self.logger.debug("Starting broadcasting weights to inference engine via NCCL") start_time = time.perf_counter() - notified_runs: list[tuple[int, Path]] = [] + # `_compute_notified_runs` is a pure function of SPMD-replicated state on + # multi_run_manager, so every trainer rank derives the same list. Only + # the master touches the filesystem to notify the orchestrator, but all + # ranks must wait on NCCL_READY before entering the broadcast path: + # the broadcast preparation (DTensor resolution, quantization) enqueues + # collectives on non-master ranks, and if those ranks start prep before + # the orchestrator has paused inference, the collectives sit unmatched + # until NCCL's watchdog kills the process after 10 min. + notified_runs = self._compute_notified_runs() if self.world.is_master: - notified_runs = self._notify_orchestrator() - # Wait for inference workers to signal readiness before starting NCCL broadcast - self._wait_for_nccl_ready(notified_runs) + self._notify_orchestrator(notified_runs) + self._wait_for_nccl_ready(notified_runs) self.nccl_broadcast_sender.broadcast_weights(model, step) self.logger.debug(f"Weights broadcasted in {time.perf_counter() - start_time:.2f}s") - def _notify_orchestrator(self) -> list[tuple[int, Path]]: - """Notify the orchestrator to initiate weight broadcast. + def _compute_notified_runs(self) -> list[tuple[int, Path]]: + """Derive the list of (run_idx, save_dir) pairs that need broadcasting. - Returns: - List of (run_idx, save_dir) tuples for runs that were notified. + Pure function of `multi_run_manager` state, which is replicated across + trainer ranks (SPMD). Returns the same list on every rank so master and + non-master ranks agree on which NCCL_READY markers to wait for. """ notified_runs: list[tuple[int, Path]] = [] - if self.world.is_master: - for idx in self.runs.used_idxs: - if not self.runs.ready_to_update[idx]: - continue - - try: - save_dir = get_step_path( - get_broadcast_dir(self.runs.get_run_dir(idx)), self.runs.progress[idx].step - ) - save_dir.mkdir(parents=True, exist_ok=True) - - stable_file = save_dir / "STABLE" - stable_file.touch() - notified_runs.append((idx, save_dir)) - except FileNotFoundError: - self.logger.warning(f"Run {idx} is deleted, skipping") - except Exception as e: - self.logger.error(f"Error broadcasting weights for run {idx}: {e}") - finally: - self.runs.ready_to_update[idx] = False + for idx in self.multi_run_manager.used_idxs: + if not self.multi_run_manager.ready_to_update[idx]: + continue + try: + save_dir = get_step_path( + get_broadcast_dir(self.multi_run_manager.get_run_dir(idx)), + self.multi_run_manager.progress[idx].step, + ) + notified_runs.append((idx, save_dir)) + except FileNotFoundError: + self.logger.warning(f"Run {idx} is deleted, skipping") + except Exception as e: + self.logger.error(f"Error resolving broadcast dir for run {idx}: {e}") return notified_runs + def _notify_orchestrator(self, notified_runs: list[tuple[int, Path]]) -> None: + """Create STABLE markers for each notified run and clear their ready flags. + + Master-only side effects (filesystem writes + state mutation). Called + after `_compute_notified_runs`; non-master ranks skip this entirely. + """ + for idx, save_dir in notified_runs: + try: + save_dir.mkdir(parents=True, exist_ok=True) + stable_file = save_dir / "STABLE" + stable_file.touch() + except FileNotFoundError: + self.logger.warning(f"Run {idx} is deleted, skipping") + except Exception as e: + self.logger.error(f"Error broadcasting weights for run {idx}: {e}") + finally: + self.multi_run_manager.ready_to_update[idx] = False + def _wait_for_nccl_ready(self, notified_runs: list[tuple[int, Path]]): """Wait for inference workers to signal they are ready to receive NCCL broadcast.""" for idx, save_dir in notified_runs: diff --git a/src/prime_rl/trainer/rl/config.py b/src/prime_rl/trainer/rl/config.py deleted file mode 100644 index 526512b5f1e..00000000000 --- a/src/prime_rl/trainer/rl/config.py +++ /dev/null @@ -1,283 +0,0 @@ -from pathlib import Path -from typing import Annotated, Literal, TypeAlias - -from pydantic import BaseModel, Field, model_validator - -from prime_rl.trainer.config import ( - AdamWConfig, - BenchConfig, - CheckpointConfig, - ConstantSchedulerConfig, - ModelConfig, - OptimizerConfigType, - SchedulerConfigType, - TokenizerConfig, -) -from prime_rl.transport.config import FileSystemTransportConfig, TransportConfigType -from prime_rl.utils.config import HeartbeatConfig, LogConfig, MetricsServerConfig, WandbConfig -from prime_rl.utils.pydantic_config import BaseConfig, BaseSettings - - -class LossConfig(BaseConfig): - """Base config for loss.""" - - ratio_type: Annotated[Literal["token", "sequence"], Field(description="Type of importance ratio to use.")] = "token" - - token_mask_high: Annotated[ - float, Field(ge=0, description="The high threshold for token importance ratio to mask.") - ] = 8.0 - token_mask_low: Annotated[ - float, Field(ge=0, description="The low threshold for token importance ratio to mask.") - ] = 0.125 - sequence_clip_high: Annotated[ - float, Field(ge=0, description="The high threshold for sequence importance ratio to clip.") - ] = 10.0 - geo_mask_high: Annotated[float, Field(ge=0, description="The high threshold for geo importance ratio to mask.")] = ( - 10.0 - ) - geo_mask_low: Annotated[float, Field(ge=0, description="The low threshold for geo importance ratio to mask.")] = 0.1 - sequence_mask_low: Annotated[ - float, - Field( - ge=0, - description="If set, masks entire sequences when any generated token has an importance ratio below this value.", - ), - ] = 0.0 - sequence_mask_high: Annotated[ - float, - Field( - ge=0, - description="If set, masks entire sequences when any generated token has an importance ratio above this value.", - ), - ] = 100.0 - - adv_tau: Annotated[float, Field(ge=0, description="The tau for advantages.")] = 1.0 - teacher_tau: Annotated[float, Field(ge=0, description="The tau for teacher logprobs.")] = 0.0 - kl_tau: Annotated[float, Field(ge=0, description="The tau for KL divergence.")] = 0.0 - - @model_validator(mode="after") - def validate_mask_bounds(self): - if self.token_mask_low >= self.token_mask_high: - raise ValueError( - f"token_mask_low ({self.token_mask_low}) must be less than token_mask_high ({self.token_mask_high})" - ) - if self.geo_mask_low >= self.geo_mask_high: - raise ValueError( - f"geo_mask_low ({self.geo_mask_low}) must be less than geo_mask_high ({self.geo_mask_high})" - ) - if self.sequence_mask_low >= self.sequence_mask_high: - raise ValueError( - f"sequence_mask_low ({self.sequence_mask_low}) must be less than sequence_mask_high ({self.sequence_mask_high})" - ) - return self - - -class FakeDataLoaderConfig(BaseConfig): - """Configures a fake data loader sampling random micro batches for debugging.""" - - batch_size: Annotated[int, Field(ge=1)] = 2 - generate_samples: Annotated[ - bool, Field(description="Whether to generate separate samples and pack them into a single micro batch.") - ] = False - - -class DataLoaderConfig(BaseConfig): - """Configures the data loader used for training.""" - - fake: Annotated[FakeDataLoaderConfig | None, Field(description="Whether to use a fake data loader.")] = None - - -class BaseWeightBroadcastConfig(BaseModel): - """Configures the base weight broadcast.""" - - pass - - -class FileSystemWeightBroadcastConfig(BaseWeightBroadcastConfig): - """Configures the weight broadcast.""" - - type: Literal["filesystem"] = "filesystem" - save_sharded: Annotated[bool, Field(description="Whether to save the weight checkpoint in sharded format.")] = True - save_format: Annotated[ - Literal["safetensors", "torch"], Field(description="The format to save the weight checkpoint in.") - ] = "safetensors" - - -class NCCLWeightBroadcastConfig(BaseWeightBroadcastConfig): - """Configures the NCCL broadcast.""" - - type: Literal["nccl"] = "nccl" - host: Annotated[str, Field(description="The host to use for the NCCL broadcast.")] = "localhost" - port: Annotated[int, Field(description="The port to use for the NCCL broadcast.")] = 29501 - timeout: Annotated[int, Field(description="The timeout in seconds to use for the NCCL broadcast.")] = 1200 - # TODO: Should not be configurable, but auto-inferred - inference_world_size: Annotated[int, Field(description="The number of GPUs used for inference.")] = 1 - - -WeightBroadcastConfigType: TypeAlias = FileSystemWeightBroadcastConfig | NCCLWeightBroadcastConfig - - -class RLTrainerConfig(BaseSettings): - """Configures the RL trainer""" - - # The model configuration - model: ModelConfig = ModelConfig() - - # The tokenizer configuration - tokenizer: TokenizerConfig = TokenizerConfig() - - # The data configuration - data: DataLoaderConfig = DataLoaderConfig() - - # The loss configuration - loss: LossConfig = LossConfig() - - # The optimizer configuration - optim: Annotated[OptimizerConfigType, Field(discriminator="type")] = AdamWConfig() - - # The learning rate scheduler configuration - scheduler: Annotated[SchedulerConfigType, Field(discriminator="type")] = ConstantSchedulerConfig() - - # The checkpoint configuration - ckpt: CheckpointConfig | None = None - - weight_broadcast: Annotated[WeightBroadcastConfigType, Field(discriminator="type")] = ( - FileSystemWeightBroadcastConfig() - ) - - rollout_transport: Annotated[TransportConfigType, Field(discriminator="type")] = FileSystemTransportConfig() - - # The logging configuration - log: LogConfig = LogConfig() - - # The wandb configuration - wandb: WandbConfig | None = None - - output_dir: Annotated[ - Path, - Field( - description="Directory to write outputs to. Will be populated with checkpoints, weights, rollouts and logs as subdirectories. Should be set to a persistent directory with enough disk space. This value should be distinct across experiments running on a single node. See the README for more details." - ), - ] = Path("outputs") - - max_steps: Annotated[ - int | None, - Field( - description="Maximum number of steps to run training for. If None, will run indefinitely.", - ), - ] = None - - max_async_level: Annotated[ - int, - Field( - ge=0, - description="Maximum number of steps that inference can be ahead of training. Determines how 'off-policy' the inference engines can be. Higher values yield better throughput through async execution, but may yield lower performance. If 0, will be fully synchronous.", - ), - ] = 1 - - memory_profiler_path: Annotated[Path | None, Field(description="Path to write memory profile to.")] = None - - bench: Annotated[ - BenchConfig | None, - Field( - description="Whether to run in benchmark mode. It will automatically set the maximum number of steps to run to 4 and use fake data.", - ), - ] = None - - trace_path: Annotated[Path | None, Field(description="Path to write pytorch profiler trace to.")] = None - - dist_timeout_seconds: Annotated[ - int, - Field( - description="Timeout in seconds for torch distributed ops. Defaults to 600 seconds.", - ), - ] = 600 - - heartbeat: Annotated[ - HeartbeatConfig | None, Field(description="The heartbeat config for monitoring training progress.") - ] = None - - metrics_server: Annotated[ - MetricsServerConfig | None, - Field(description="Prometheus metrics server config. If set, exposes /metrics endpoint for scraping."), - ] = None - - max_concurrent_runs: Annotated[ - int, - Field( - ge=1, - description="The maximum number of concurrent runs to allow. If 1, then only one run will be allowed at a time.", - ), - ] = 1 - - @model_validator(mode="after") - def auto_setup_bench(self): - if self.bench is not None: - self.max_steps = 4 # 1 Warmup + 3 Benchmark - if not self.data.fake: - self.data.fake = FakeDataLoaderConfig() - if self.ckpt: # Do not checkpoint - self.ckpt = None - return self - - @model_validator(mode="after") - def dont_do_massive_traces(self): - if self.trace_path: - if self.max_steps is None: - raise ValueError("Must specify max_steps when tracing") - if self.max_steps >= 10: - raise ValueError( - "Tracing more than 10 steps is not recommended as your trace will be massive. Remove this line if you really want to trace more steps." - ) - return self - - @model_validator(mode="after") - def validate_lora_adapter_saving(self): - if self.ckpt and self.ckpt.weights and self.ckpt.weights.save_adapter_separately: - lora_enabled = self.model and self.model.lora - if not lora_enabled: - raise ValueError( - "save_adapter_separately=True requires LoRA to be enabled. " - "Set model.lora or disable save_adapter_separately." - ) - return self - - @model_validator(mode="after") - def validate_weight_broadcast_type(self): - if self.weight_broadcast.type == "nccl" and self.max_async_level != 1: - raise ValueError("NCCL weight broadcast only works with async level 1") - return self - - @model_validator(mode="after") - def validate_opt_and_fsdp_offload(self): - if self.optim.type == "muon" and self.model.fsdp_cpu_offload: - raise ValueError("Muon optimizer does not support FSDP CPU offload") - return self - - @model_validator(mode="after") - def validate_lora_broadcast(self): - if self.model.lora is not None and self.weight_broadcast.type == "nccl": - # TODO: Support this - raise ValueError("NCCL weight broadcast does not support LoRA yet.") - return self - - @model_validator(mode="after") - def auto_setup_tokenizer(self): - if self.tokenizer.name is None: - self.tokenizer.name = self.model.name - if self.tokenizer.trust_remote_code is None: - self.tokenizer.trust_remote_code = self.model.trust_remote_code - return self - - @model_validator(mode="after") - def auto_setup_fused_lm_head_chunk_size(self): - if self.model.fused_lm_head_chunk_size is None and self.model.impl != "liger_kernel": - self.model.fused_lm_head_chunk_size = 2048 - return self - - @model_validator(mode="after") - def ep_only_with_custom_impl(self): - if self.model.ep > 1 and self.model.impl not in ("custom", "auto"): - raise ValueError("EP is only supported with the custom implementation or auto mode") - - return self diff --git a/src/prime_rl/trainer/rl/data.py b/src/prime_rl/trainer/rl/data.py index 68c79ebf54d..ffc4bc627f7 100644 --- a/src/prime_rl/trainer/rl/data.py +++ b/src/prime_rl/trainer/rl/data.py @@ -6,11 +6,11 @@ from torch import Tensor from transformers.tokenization_utils import PreTrainedTokenizer -from prime_rl.trainer.rl.config import FakeDataLoaderConfig +from prime_rl.configs.trainer import FakeDataLoaderConfig from prime_rl.trainer.rl.packer import BasePacker, setup_packer -from prime_rl.trainer.runs import get_runs +from prime_rl.trainer.runs import get_multi_run_manager from prime_rl.trainer.world import get_world -from prime_rl.transport import MicroBatch, MicroBatchReceiver, TransportConfigType, setup_micro_batch_receiver +from prime_rl.transport import MicroBatch, MicroBatchReceiver, TransportConfig, setup_micro_batch_receiver class TensorMicroBatch(TypedDict): @@ -23,11 +23,25 @@ class TensorMicroBatch(TypedDict): inference_logprobs: Float[Tensor, "batch seq"] teacher_logprobs: Float[Tensor, "batch seq"] | None loss_mask: Bool[Tensor, "batch seq"] + temperatures: Float[Tensor, "batch seq"] # Per-token temperatures # Batch level - temperature: float lora_num_tokens: Int[Tensor, "n_loras"] + # MoE router replay + routed_experts: Int[Tensor, "batch seq layers topk"] | None + + # Multimodal fields (Qwen3-VL) + # pixel_values: flattened image patches [num_patches, patch_dim] where patch_dim=1176 for Qwen3-VL + pixel_values: Float[Tensor, "num_patches patch_dim"] | None + # image_grid_thw: grid dimensions [num_images, 3] where each entry is [temporal, height, width] + image_grid_thw: Int[Tensor, "num_images 3"] | None + # mm_token_type_ids: token type per token [batch seq], int64 (0=text, 1=image, 2=video) + mm_token_type_ids: Int[Tensor, "batch seq"] | None + + # When True, trainer uses SFT loss instead of RL loss for this batch + sft_loss: bool + class FakeDataLoader: def __init__(self, config: FakeDataLoaderConfig, seq_len: int, dp_world_size: int): @@ -41,7 +55,7 @@ def __init__(self, config: FakeDataLoaderConfig, seq_len: int, dp_world_size: in self.seq_len = seq_len self.generate_samples = config.generate_samples self.batch_counter = 0 - self.runs = get_runs() + self.multi_run_manager = get_multi_run_manager() def wait_for_batch(self) -> None: return @@ -84,7 +98,7 @@ def _get_sample_micro_batch(self, generator: torch.Generator) -> TensorMicroBatc loss_mask = torch.ones(input_ids.shape[0], dtype=torch.bool) advantages = torch.randn(input_ids.shape[0], generator=generator) inference_logprobs = torch.randn(input_ids.shape[0], generator=generator) - lora_num_tokens = torch.zeros(self.runs.max_runs, dtype=torch.int32) + lora_num_tokens = torch.zeros(self.multi_run_manager.max_runs, dtype=torch.int32) lora_num_tokens[0] = input_ids.shape[0] return { @@ -93,13 +107,18 @@ def _get_sample_micro_batch(self, generator: torch.Generator) -> TensorMicroBatc "advantages": advantages.unsqueeze(0), "inference_logprobs": inference_logprobs.unsqueeze(0), "teacher_logprobs": None, - "temperature": 1.0, + "temperatures": torch.ones(input_ids.shape[0]).unsqueeze(0), "loss_mask": loss_mask.unsqueeze(0), "lora_num_tokens": lora_num_tokens, + "routed_experts": None, + "pixel_values": None, + "image_grid_thw": None, + "mm_token_type_ids": None, + "sft_loss": False, } def _get_micro_batch(self, generator: torch.Generator) -> TensorMicroBatch: - lora_num_tokens = torch.zeros(self.runs.max_runs, dtype=torch.int32) + lora_num_tokens = torch.zeros(self.multi_run_manager.max_runs, dtype=torch.int32) lora_num_tokens[0] = self.seq_len return { "input_ids": torch.randint( @@ -115,9 +134,14 @@ def _get_micro_batch(self, generator: torch.Generator) -> TensorMicroBatch: "advantages": torch.randn(self.seq_len, generator=generator).unsqueeze(0), "inference_logprobs": torch.randn(self.seq_len, generator=generator).unsqueeze(0), "teacher_logprobs": None, - "temperature": 1.0, + "temperatures": torch.ones(self.seq_len).unsqueeze(0), "loss_mask": torch.ones(self.seq_len, dtype=torch.bool).unsqueeze(0), "lora_num_tokens": lora_num_tokens, + "routed_experts": None, + "pixel_values": None, + "image_grid_thw": None, + "mm_token_type_ids": None, + "sft_loss": False, } @@ -132,7 +156,7 @@ def __init__( seq_len: int, pad_to_multiple_of: int, tokenizer: PreTrainedTokenizer, - config: TransportConfigType, + config: TransportConfig, ): self.world = get_world() @@ -148,15 +172,19 @@ def __init__( non_dp_world_size = self.world.world_size // dp_world_size dp_rank = self.world.rank // non_dp_world_size - self.runs = get_runs() + self.multi_run_manager = get_multi_run_manager() self.receiver: MicroBatchReceiver = setup_micro_batch_receiver(output_dir, dp_rank, start_step, config) def wait_for_batch(self) -> None: if self.world.is_master: - self.packer.pack() + self.packer._arm_watchdog() + try: + self.packer.pack() + finally: + self.packer._disarm_watchdog() self.receiver.wait() - self.runs.sync_runs() + self.multi_run_manager.synchronize_state() def get_batch(self) -> list[TensorMicroBatch]: micro_batches = self.receiver.receive() @@ -165,7 +193,7 @@ def get_batch(self) -> list[TensorMicroBatch]: def _micro_batch_to_tensor(self, micro_batch: MicroBatch) -> TensorMicroBatch: """Convert a MicroBatch (msgspec struct with lists) to a TensorMicroBatch (dict with tensors).""" if micro_batch.lora_num_tokens is None: - micro_batch.lora_num_tokens = [0] * self.runs.max_runs + micro_batch.lora_num_tokens = [0] * self.multi_run_manager.max_runs micro_batch.lora_num_tokens[0] = len(micro_batch.input_ids) return TensorMicroBatch( input_ids=torch.tensor(micro_batch.input_ids, dtype=torch.long).unsqueeze(0), @@ -176,6 +204,24 @@ def _micro_batch_to_tensor(self, micro_batch: MicroBatch) -> TensorMicroBatch: if micro_batch.teacher_logprobs is not None else None, loss_mask=torch.tensor(micro_batch.loss_mask, dtype=torch.bool).unsqueeze(0), - temperature=micro_batch.temperature, + temperatures=torch.tensor(micro_batch.temperatures, dtype=torch.float).unsqueeze(0), lora_num_tokens=torch.tensor(micro_batch.lora_num_tokens, dtype=torch.int32), + # Multimodal fields - no batch dimension for these as they are variable-sized + pixel_values=torch.frombuffer(bytearray(micro_batch.pixel_values), dtype=torch.float32).reshape( + micro_batch.pixel_values_shape + ) + if micro_batch.pixel_values is not None + else None, + image_grid_thw=torch.tensor(micro_batch.image_grid_thw, dtype=torch.long) + if micro_batch.image_grid_thw is not None + else None, + mm_token_type_ids=torch.tensor(micro_batch.mm_token_type_ids, dtype=torch.long).unsqueeze(0) + if micro_batch.mm_token_type_ids is not None + else None, + routed_experts=torch.tensor(micro_batch.routed_experts, dtype=torch.int32).unsqueeze( + 0 + ) # [1, seq_len, layers, topk] + if micro_batch.routed_experts is not None + else None, + sft_loss=micro_batch.sft_loss, ) diff --git a/src/prime_rl/trainer/rl/loss.py b/src/prime_rl/trainer/rl/loss.py index 019e2c39e30..d67c055a681 100644 --- a/src/prime_rl/trainer/rl/loss.py +++ b/src/prime_rl/trainer/rl/loss.py @@ -1,11 +1,41 @@ -from typing import Any +from dataclasses import dataclass +from typing import Any, Callable import torch from beartype import beartype as typechecker -from jaxtyping import Float, Int, jaxtyped +from jaxtyping import Bool, Float, Int, jaxtyped from torch import Tensor -from prime_rl.trainer.rl.config import LossConfig +from prime_rl.configs.trainer import CustomLossConfig, DefaultLossConfig, LossConfig, SFTLossConfig +from prime_rl.utils.utils import import_object + + +@dataclass +class LossInputs: + """Inputs for computing loss on a single sample.""" + + trainer_logprobs: Float[Tensor, " seq"] + inference_logprobs: Float[Tensor, " seq"] + teacher_logprobs: Float[Tensor, " seq"] | None + advantages: Float[Tensor, " seq"] + loss_mask: Bool[Tensor, " seq"] + + +@dataclass +class LossOutputs: + """Outputs from computing loss on a single sample.""" + + loss: Float[Tensor, ""] + metrics: dict[str, Tensor] + + +LossFn = Callable[..., LossOutputs] +"""Type for a per-sample loss function. + +Expected signature: + def my_loss(inputs: LossInputs, **kwargs) -> LossOutputs: + ... +""" @jaxtyped(typechecker=typechecker) @@ -74,120 +104,158 @@ def _safe_mean(values: Tensor, mask: Tensor) -> Tensor: return values[mask].sum() / denom +def default_loss_fn(inputs: LossInputs, loss_config: DefaultLossConfig) -> LossOutputs: + """ + DPPO+KL loss, combining: + - DPPO-Binary TV Loss (https://arxiv.org/pdf/2602.04879) + - Kimi-K2.5 KL Loss (https://arxiv.org/pdf/2602.02276) + + The mask is conditioned on the advantage sign: for positive advantages, + we mask tokens whose probability increased too much (trust region violation + in the upweight direction); for negative advantages, we mask tokens whose + probability decreased too much (trust region violation in the downweight + direction). + """ + trainer_logprobs = inputs.trainer_logprobs + inference_logprobs = inputs.inference_logprobs + teacher_logprobs = inputs.teacher_logprobs + advantages = inputs.advantages + loss_mask = inputs.loss_mask + + trainer_probs = torch.exp(trainer_logprobs) + inference_probs = torch.exp(inference_logprobs) + probs_diff = trainer_probs - inference_probs + dppo_invalid_mask_high = probs_diff > loss_config.dppo_mask_high + dppo_invalid_mask_low = probs_diff < -loss_config.dppo_mask_low + dppo_invalid_mask = torch.where(advantages > 0, dppo_invalid_mask_high, dppo_invalid_mask_low) + + is_masked = dppo_invalid_mask + is_masked_high = (advantages > 0) & dppo_invalid_mask_high + is_masked_low = (advantages < 0) & dppo_invalid_mask_low + keep_mask = loss_mask & ~is_masked + + log_importance_ratio = trainer_logprobs - inference_logprobs + importance_ratio = torch.exp(log_importance_ratio) + mismatch_kl = importance_ratio - log_importance_ratio - 1 + + advantages = loss_config.adv_tau * advantages + if teacher_logprobs is not None: + teacher_kl = teacher_logprobs - trainer_logprobs + advantages = advantages + loss_config.teacher_tau * teacher_kl.detach() + else: + teacher_kl = None + + pg_loss = keep_mask * advantages * importance_ratio + kl_loss = loss_mask * log_importance_ratio**2 + loss = (-pg_loss + loss_config.kl_tau * kl_loss).sum() + + metrics = { + "mismatch_kl": _safe_mean(mismatch_kl, loss_mask), # all trainable tokens + "masked_mismatch_kl": _safe_mean(mismatch_kl, loss_mask & is_masked), # all trainable, masked tokens + "unmasked_mismatch_kl": _safe_mean(mismatch_kl, keep_mask), # all trainable, unmasked tokens + "is_masked": _safe_mean(is_masked, loss_mask), + "is_masked_low": _safe_mean(is_masked_low, loss_mask), + "is_masked_high": _safe_mean(is_masked_high, loss_mask), + } + if teacher_kl is not None: + metrics["teacher_kl"] = _safe_mean(teacher_kl, loss_mask) + + return LossOutputs(loss=loss, metrics=metrics) + + +def sft_loss_fn(inputs: LossInputs) -> LossOutputs: + """SFT-style masked negative log-likelihood over trainable tokens.""" + trainer_logprobs = inputs.trainer_logprobs + loss_mask = inputs.loss_mask + + loss = -(trainer_logprobs[loss_mask]).sum() + metrics = { + "nll": _safe_mean(-trainer_logprobs, loss_mask), + } + return LossOutputs(loss=loss, metrics=metrics) + + +def setup_loss_fn(loss_config: LossConfig) -> LossFn: + """Setup the loss function based on config.""" + if isinstance(loss_config, CustomLossConfig): + custom_fn = import_object(loss_config.import_path) + kwargs = loss_config.kwargs + + def loss_fn(inputs: LossInputs) -> LossOutputs: + return custom_fn(inputs, **kwargs) + + return loss_fn + + if isinstance(loss_config, SFTLossConfig): + return sft_loss_fn + + def loss_fn(inputs: LossInputs) -> LossOutputs: + return default_loss_fn(inputs, loss_config) + + return loss_fn + + def compute_loss( - trainer_logprobs: Any, # list of Float[Tensor, "seq_i"] with potentially different seq_i lengths - inference_logprobs: Any, # list of Float[Tensor, "seq_i"] with potentially different seq_i lengths - teacher_logprobs: Any | None, # list of Float[Tensor, "seq_i"] with potentially different seq_i lengths, or None - advantages: Any, # list of Float[Tensor, "seq_i"] with potentially different seq_i lengths - loss_mask: Any, # list of Bool[Tensor, "seq_i"] with potentially different seq_i lengths - loss_config: LossConfig, + trainer_logprobs: list[Float[Tensor, " seq_i"]], + inference_logprobs: list[Float[Tensor, " seq_i"]], + teacher_logprobs: list[Float[Tensor, " seq_i"]] | None, + advantages: list[Float[Tensor, " seq_i"]], + loss_mask: list[Bool[Tensor, " seq_i"]], + loss_fn: LossFn, loss_scale: int, + sft_loss: bool = False, ) -> tuple[Float[Tensor, ""], dict[str, Any]]: """ Compute loss for packed sequences (batch size = 1, multiple sequences packed along sequence dimension). Args: - trainer_logprobs: Log probabilities tensor for packed sequences - inference_logprobs: Old log probabilities tensor for packed sequences - teacher_logprobs: Teacher log probabilities tensor for packed sequences, or None if not configured - advantages: Advantages tensor for packed sequences - loss_mask: Loss mask tensor for packed sequences - loss_config: Loss configuration object + trainer_logprobs: Log probabilities for each sequence + inference_logprobs: Reference log probabilities for each sequence + teacher_logprobs: Teacher log probabilities for each sequence, or None + advantages: Advantages for each sequence + loss_mask: Loss mask for each sequence + loss_fn: Per-sequence loss function loss_scale: Scale factor to normalize the loss + sft_loss: If True, use SFT loss instead of the configured loss_fn for this batch Returns: - Tuple of (scaled_loss, aggregated_loss_tensors) + Tuple of (scaled_loss, aggregated_metrics) """ + effective_loss_fn = sft_loss_fn if sft_loss else loss_fn total_loss = 0.0 - total_mismatch_kl = [] - total_masked_mismatch_kl = [] - total_unmasked_mismatch_kl = [] - total_is_masked = [] - total_is_masked_low = [] - total_is_masked_high = [] - total_sequence_masked_low = [] - total_sequence_masked_high = [] - total_geo_masked_low = [] - total_geo_masked_high = [] - total_geo_seq_ratio = [] - total_teacher_kl = [] + all_metrics: dict[str, list[Tensor]] = {} if teacher_logprobs is None: teacher_logprobs = [None] * len(trainer_logprobs) - for trainer_logprobs, inference_logprobs, teacher_logprobs, advantages, loss_mask in zip( + for t_logp, i_logp, teach_logp, adv, mask in zip( trainer_logprobs, inference_logprobs, teacher_logprobs, advantages, loss_mask ): - log_importance_ratio = trainer_logprobs - inference_logprobs - teacher_kl = teacher_logprobs - trainer_logprobs if teacher_logprobs is not None else None - - # Trainer-inference mismatch KL per token - token_importance_ratio = torch.exp(log_importance_ratio) - geo_seq_ratio = torch.exp(_safe_mean(log_importance_ratio, loss_mask)) - token_mismatch_kl = token_importance_ratio - log_importance_ratio - 1 - - seq_log_importance_ratio = torch.clamp(log_importance_ratio[loss_mask].sum().detach(), max=10.0) - seq_importance_ratio = torch.clamp(torch.exp(seq_log_importance_ratio), max=loss_config.sequence_clip_high) - - seq_min_ratio = torch.where(loss_mask, token_importance_ratio, torch.inf).min() - seq_max_ratio = torch.where(loss_mask, token_importance_ratio, -torch.inf).max() - seq_mask_low = seq_min_ratio < loss_config.sequence_mask_low - seq_mask_high = seq_max_ratio > loss_config.sequence_mask_high - - token_mask_low = token_importance_ratio < loss_config.token_mask_low - token_mask_high = token_importance_ratio > loss_config.token_mask_high - - geo_mask_low = geo_seq_ratio < loss_config.geo_mask_low - geo_mask_high = geo_seq_ratio > loss_config.geo_mask_high - - is_masked = token_mask_low | token_mask_high | geo_mask_low | geo_mask_high | seq_mask_low | seq_mask_high - keep_mask = loss_mask & ~is_masked - - importance_ratio = seq_importance_ratio if loss_config.ratio_type == "sequence" else token_importance_ratio - - advantages = loss_config.adv_tau * advantages - if teacher_logprobs is not None: - advantages = advantages + loss_config.teacher_tau * teacher_kl.detach() - coeff = importance_ratio * (advantages - loss_config.kl_tau * log_importance_ratio) - loss = -(coeff.detach() * trainer_logprobs)[keep_mask].sum() - - if loss_config.ratio_type == "sequence": - loss = loss / torch.clamp_min(loss_mask.sum(), 1) - - total_loss = total_loss + loss - - # Aggregate loss tensors - total_mismatch_kl.append(_safe_mean(token_mismatch_kl, loss_mask)) - total_masked_mismatch_kl.append(_safe_mean(token_mismatch_kl, loss_mask & is_masked)) - total_unmasked_mismatch_kl.append(_safe_mean(token_mismatch_kl, keep_mask)) - total_is_masked.append(is_masked[loss_mask].float()) - total_is_masked_low.append(token_mask_low[loss_mask].float()) - total_is_masked_high.append(token_mask_high[loss_mask].float()) - total_sequence_masked_low.append(seq_mask_low.float()) - total_sequence_masked_high.append(seq_mask_high.float()) - total_geo_masked_low.append(geo_mask_low.float()) - total_geo_masked_high.append(geo_mask_high.float()) - total_geo_seq_ratio.append(geo_seq_ratio) - if teacher_logprobs is not None: - total_teacher_kl.append(_safe_mean(teacher_kl, loss_mask)) - - # Apply loss scaling + inputs = LossInputs( + trainer_logprobs=t_logp, + inference_logprobs=i_logp, + teacher_logprobs=teach_logp, + advantages=adv, + loss_mask=mask, + ) + + result = effective_loss_fn(inputs) + + total_loss = total_loss + result.loss + + for k, v in result.metrics.items(): + if k not in all_metrics: + all_metrics[k] = [] + all_metrics[k].append(v) + scaled_loss = total_loss / loss_scale - result = { - "mismatch_kl": torch.stack(total_mismatch_kl), - "masked_mismatch_kl": torch.stack(total_masked_mismatch_kl), - "unmasked_mismatch_kl": torch.stack(total_unmasked_mismatch_kl), - "is_masked": torch.cat(total_is_masked), - "is_masked_low": torch.cat(total_is_masked_low), - "is_masked_high": torch.cat(total_is_masked_high), - "sequence_masked_low": torch.stack(total_sequence_masked_low), - "sequence_masked_high": torch.stack(total_sequence_masked_high), - "geo_masked_low": torch.stack(total_geo_masked_low), - "geo_masked_high": torch.stack(total_geo_masked_high), - "geo_seq_ratio": torch.stack(total_geo_seq_ratio), - } - if total_teacher_kl: - result["teacher_kl"] = torch.stack(total_teacher_kl) - return scaled_loss, result + aggregated: dict[str, Any] = {} + for k, v in all_metrics.items(): + if v[0].dim() == 0: + aggregated[k] = torch.stack(v) + else: + aggregated[k] = torch.cat(v) + + return scaled_loss, aggregated diff --git a/src/prime_rl/trainer/rl/packer.py b/src/prime_rl/trainer/rl/packer.py index ad7d404ac9a..cf9dcfa02ee 100644 --- a/src/prime_rl/trainer/rl/packer.py +++ b/src/prime_rl/trainer/rl/packer.py @@ -1,4 +1,6 @@ +import os import shutil +import threading import time from abc import ABC, abstractmethod from collections import deque @@ -6,11 +8,12 @@ from transformers.tokenization_utils import PreTrainedTokenizer from prime_rl.trainer.batch import prepare_batch -from prime_rl.trainer.runs import get_runs +from prime_rl.trainer.runs import get_multi_run_manager from prime_rl.transport import ( + MicroBatch, MicroBatchSender, TrainingSample, - TransportConfigType, + TransportConfig, setup_micro_batch_sender, setup_training_batch_receiver, ) @@ -18,6 +21,7 @@ from prime_rl.utils.pathing import get_rollout_dir TIMEOUT_SECONDS = 0.1 +WATCHDOG_TIMEOUT_SECONDS = 1800 # 30 minutes class BasePacker(ABC): @@ -27,20 +31,44 @@ def __init__( seq_len: int, pad_to_multiple_of: int, tokenizer: PreTrainedTokenizer, - config: TransportConfigType, + config: TransportConfig, start_step: int = 0, ): self.logger = get_logger() - self.runs = get_runs() + self.multi_run_manager = get_multi_run_manager() self.dp_world_size = dp_world_size self.seq_len = seq_len self.pad_to_multiple_of = pad_to_multiple_of self.tokenizer = tokenizer self.receiver = setup_training_batch_receiver(config) - shutil.rmtree(get_rollout_dir(self.runs.output_dir), ignore_errors=True) + shutil.rmtree(get_rollout_dir(self.multi_run_manager.output_dir), ignore_errors=True) self.sender: MicroBatchSender = setup_micro_batch_sender( - self.runs.output_dir, dp_world_size, start_step, config + self.multi_run_manager.output_dir, dp_world_size, start_step, config ) + self._last_heartbeat = time.monotonic() + self._watchdog_armed = threading.Event() + self._watchdog = threading.Thread(target=self._watchdog_loop, daemon=True) + self._watchdog.start() + + def _heartbeat(self) -> None: + self._last_heartbeat = time.monotonic() + + def _arm_watchdog(self) -> None: + self._last_heartbeat = time.monotonic() + self._watchdog_armed.set() + + def _disarm_watchdog(self) -> None: + self._watchdog_armed.clear() + + def _watchdog_loop(self) -> None: + while True: + time.sleep(60) + if not self._watchdog_armed.is_set(): + continue + stale = time.monotonic() - self._last_heartbeat + if stale > WATCHDOG_TIMEOUT_SECONDS: + self.logger.error(f"Packer heartbeat stale for {stale:.0f}s, killing process to trigger restart") + os._exit(1) @abstractmethod def pack(self) -> None: @@ -55,33 +83,33 @@ def __init__( seq_len: int, pad_to_multiple_of: int, tokenizer: PreTrainedTokenizer, - config: TransportConfigType, + config: TransportConfig, start_step: int = 0, ): super().__init__(dp_world_size, seq_len, pad_to_multiple_of, tokenizer, config, start_step) - assert self.runs.max_runs == 1, "SinglePacker only supports one run" + assert self.multi_run_manager.max_runs == 1, "SinglePacker only supports one run" def pack(self): # Wait for batch to be available batches = [] while len(batches) == 0: - self.runs.check_for_changes() + self._heartbeat() + self.multi_run_manager.discover_runs() batches = self.receiver.receive() time.sleep(0.2) assert len(batches) == 1, "SinglePacker only supports one batch per step" batch = batches[0] - self.runs.ready_to_update[0] = True - self.runs.progress[0].step += 1 + self.multi_run_manager.ready_to_update[0] = True + self.multi_run_manager.progress[0].step += 1 micro_batch_grid = prepare_batch( rollouts=batch.examples, - temperature=batch.temperature, seq_len=self.seq_len, pad_to_multiple_of=self.pad_to_multiple_of, num_train_workers=self.dp_world_size, idxs=[0] * len(batch.examples), - num_loras=self.runs.max_runs, + num_loras=self.multi_run_manager.max_runs, ) self.sender.send(micro_batch_grid) @@ -94,19 +122,21 @@ def __init__( seq_len: int, pad_to_multiple_of: int, tokenizer: PreTrainedTokenizer, - config: TransportConfigType, + config: TransportConfig, start_step: int = 0, ): super().__init__(dp_world_size, seq_len, pad_to_multiple_of, tokenizer, config, start_step) - # Per-run buffer: stores (TrainingSample, temperature, step) tuples - self.buffers: list[deque[tuple[TrainingSample, float, int]]] = [deque() for _ in range(self.runs.max_runs)] + # Per-run buffer: stores (TrainingSample, step) tuples + self.buffers: list[deque[tuple[TrainingSample, int]]] = [ + deque() for _ in range(self.multi_run_manager.max_runs) + ] # Round-robin position (persists across pack() calls) self._round_robin_position: int = 0 - # Register delete_run_data hook for receiver reset (master only, runs during check_for_changes) + # Register forgotten hook for receiver reset (master only, called during discover_runs) # This must happen when a run is deleted to prevent stale data from remaining - self.runs.register_delete_run_data_hook(self._on_run_data_deleted) + self.multi_run_manager.register_forgotten_hook(self._on_run_data_deleted) def _on_run_data_deleted(self, idx: int, run_id: str) -> None: """Reset run state when run data is deleted (master only).""" @@ -116,27 +146,76 @@ def _on_run_data_deleted(self, idx: int, run_id: str) -> None: # Reset run state self.buffers[idx].clear() + def _validate_sample(self, sample: TrainingSample) -> tuple[bool, str | None]: + """Validate a sample to ensure it won't crash the trainer.""" + sample_length = len(sample.prompt_ids) + len(sample.completion_ids) + if len(sample.prompt_mask) != len(sample.prompt_ids): + return ( + False, + f"Run wrote a sample with prompt mask length != prompt ids length ({len(sample.prompt_mask)} != {len(sample.prompt_ids)})", + ) + if len(sample.completion_mask) != len(sample.completion_ids): + return ( + False, + f"Run wrote a sample with completion mask length != completion ids length ({len(sample.completion_mask)} != {len(sample.completion_ids)})", + ) + if len(sample.completion_logprobs) != len(sample.completion_ids): + return ( + False, + f"Run wrote a sample with completion logprobs length != completion ids length ({len(sample.completion_logprobs)} != {len(sample.completion_ids)})", + ) + if len(sample.completion_temperatures) != len(sample.completion_ids): + return ( + False, + f"Run wrote a sample with completion temperatures length != completion ids length ({len(sample.completion_temperatures)} != {len(sample.completion_ids)})", + ) + if sample_length == 0: + return False, "Run wrote a sample with no tokens" + if sample_length > self.seq_len: + return ( + False, + f"Run wrote a sample with length {sample_length} which exceeds max sequence length {self.seq_len}", + ) + if sample.teacher_logprobs is not None and len(sample.teacher_logprobs) != sample_length: + return ( + False, + f"Run wrote a sample with teacher logprobs length != sample length ({len(sample.teacher_logprobs)} != {sample_length})", + ) + return True, None + def _get_batch(self) -> None: """Receive batches from orchestrator and buffer samples per run.""" - self.runs.check_for_changes() + self._heartbeat() + self.multi_run_manager.discover_runs() batches = self.receiver.receive() for batch in batches: if batch.run_idx is None: self.logger.warning("Received batch with no run index") continue + if len(batch.examples) == 0: + self.multi_run_manager.evict_run(batch.run_idx, "Run wrote a batch with no samples") + continue for sample in batch.examples: - self.buffers[batch.run_idx].append((sample, batch.temperature, batch.step)) + valid, reason = self._validate_sample(sample) + if not valid: + self.multi_run_manager.evict_run(batch.run_idx, f"Run wrote a sample with invalid data: {reason}") + break + self.buffers[batch.run_idx].append((sample, batch.step)) + + # This is necessary to forget evicted runs + self.multi_run_manager.discover_runs() def _count_tokens(self, threshold: int | None = None) -> int: tokens = 0 - for run_idx in self.runs.used_idxs: + for run_idx in self.multi_run_manager.used_idxs: buffer = self.buffers[run_idx] - current_step = self.runs.progress[run_idx].step - for sample, _, step in buffer: + current_step = self.multi_run_manager.progress[run_idx].step + + for sample, step in buffer: if step > current_step: - continue + break tokens += len(sample.prompt_ids) + len(sample.completion_ids) if threshold is not None and tokens >= threshold: return tokens @@ -148,17 +227,17 @@ def _has_enough_tokens(self) -> bool: threshold = self.seq_len * self.dp_world_size return self._count_tokens(threshold) >= threshold - def _select_samples_round_robin(self, token_budget: int) -> list[tuple[int, TrainingSample, float]]: + def _select_samples_round_robin(self, token_budget: int) -> list[tuple[int, TrainingSample, int]]: """Select samples using round-robin from runs with buffered work.""" - selected: list[tuple[int, TrainingSample, float]] = [] + selected: list[tuple[int, TrainingSample, int]] = [] tokens_collected = 0 while tokens_collected < token_budget: # Round-robin until we find a run with work for the current step for _ in range(len(self.buffers)): if len(self.buffers[self._round_robin_position]) > 0: - _, _, step = self.buffers[self._round_robin_position][0] - if step <= self.runs.progress[self._round_robin_position].step: + _, step = self.buffers[self._round_robin_position][0] + if step <= self.multi_run_manager.progress[self._round_robin_position].step: break self._round_robin_position = (self._round_robin_position + 1) % len(self.buffers) else: @@ -167,10 +246,10 @@ def _select_samples_round_robin(self, token_budget: int) -> list[tuple[int, Trai break run_idx = self._round_robin_position self._round_robin_position = (self._round_robin_position + 1) % len(self.buffers) - current_step = self.runs.progress[run_idx].step + current_step = self.multi_run_manager.progress[run_idx].step while len(self.buffers[run_idx]) > 0: - sample, temperature, step = self.buffers[run_idx][0] + sample, step = self.buffers[run_idx][0] if step > current_step: # Samples from different steps should be consumed later break @@ -182,7 +261,7 @@ def _select_samples_round_robin(self, token_budget: int) -> list[tuple[int, Trai self.buffers[run_idx].popleft() continue return selected - selected.append((run_idx, sample, temperature)) + selected.append((run_idx, sample, step)) self.buffers[run_idx].popleft() return selected @@ -192,12 +271,15 @@ def _update_run_progress(self, run_idx: int, num_samples: int, num_tokens: int) # HACK: This fixes the issue with branching rollouts having unpredictable batch size # However, it makes us unable to do incremental orchestrator rollouts # Removing the len(self.buffers[run_idx]) == 0 check would allow incremental orchestrator rollouts - if len(self.buffers[run_idx]) == 0 or self.buffers[run_idx][0][2] > self.runs.progress[run_idx].step: - self.runs.progress[run_idx].step += 1 - self.runs.ready_to_update[run_idx] = True + if ( + len(self.buffers[run_idx]) == 0 + or self.buffers[run_idx][0][1] > self.multi_run_manager.progress[run_idx].step + ): + self.multi_run_manager.progress[run_idx].step += 1 + self.multi_run_manager.ready_to_update[run_idx] = True - self.runs.progress[run_idx].total_tokens += num_tokens - self.runs.progress[run_idx].total_samples += num_samples + self.multi_run_manager.progress[run_idx].total_tokens += num_tokens + self.multi_run_manager.progress[run_idx].total_samples += num_samples def pack(self): """Pack samples from buffers using round-robin fair scheduling.""" @@ -215,40 +297,42 @@ def pack(self): selected_samples = self._select_samples_round_robin(token_budget) assert selected_samples, "No samples selected" - # Group by run for prepare_batch (MultiLoRAMoE requires same run_idx in microbatch) - samples_by_run: dict[int, list[tuple[TrainingSample, float]]] = {} - for run_idx, sample, temperature in selected_samples: + # Group samples by run_idx - each microbatch must contain samples from only ONE run + # because MultiLoRAGroupedExperts (MoE) only supports one adapter per microbatch + samples_by_run: dict[int, list[TrainingSample]] = {} + per_run_stats: dict[int, tuple[int, int]] = {} + for run_idx, sample, step in selected_samples: if run_idx not in samples_by_run: samples_by_run[run_idx] = [] - samples_by_run[run_idx].append((sample, temperature)) + samples_by_run[run_idx].append(sample) - micro_batch_grid = [[] for _ in range(self.dp_world_size)] - - for run_idx, sample_temp_pairs in samples_by_run.items(): - samples = [s for s, _ in sample_temp_pairs] - # We don't support dynamic temperatures in orchestrator yet - # So this works for now - temperature = sample_temp_pairs[0][1] - - num_samples = len(samples) - num_tokens = sum(len(s.prompt_ids) + len(s.completion_ids) for s in samples) + num_tokens = len(sample.prompt_ids) + len(sample.completion_ids) + if run_idx in per_run_stats: + cur_samples, cur_tokens = per_run_stats[run_idx] + per_run_stats[run_idx] = (cur_samples + 1, cur_tokens + num_tokens) + else: + per_run_stats[run_idx] = (1, num_tokens) + for run_idx, (num_samples, num_tokens) in per_run_stats.items(): self._update_run_progress(run_idx, num_samples, num_tokens) - _micro_batch_grid = prepare_batch( - rollouts=samples, - temperature=temperature, + # Pack each run separately to ensure no mixing of runs in microbatches + all_micro_batches: list[list[MicroBatch]] = [[] for _ in range(self.dp_world_size)] + for run_idx in sorted(samples_by_run.keys()): + run_samples = samples_by_run[run_idx] + run_micro_batch_grid = prepare_batch( + rollouts=run_samples, seq_len=self.seq_len, pad_to_multiple_of=self.pad_to_multiple_of, num_train_workers=self.dp_world_size, - idxs=[run_idx] * num_samples, - num_loras=self.runs.max_runs, + idxs=[run_idx] * len(run_samples), + num_loras=self.multi_run_manager.max_runs, ) + # Merge into combined grid + for worker_idx, worker_batches in enumerate(run_micro_batch_grid): + all_micro_batches[worker_idx].extend(worker_batches) - for i, micro_batch in enumerate(_micro_batch_grid): - micro_batch_grid[i].extend(micro_batch) - - self.sender.send(micro_batch_grid) + self.sender.send(all_micro_batches) def setup_packer( @@ -256,11 +340,11 @@ def setup_packer( seq_len: int, pad_to_multiple_of: int, tokenizer: PreTrainedTokenizer, - transport_config: TransportConfigType, + transport_config: TransportConfig, start_step: int = 0, ) -> BasePacker: - runs = get_runs() - if runs.max_runs == 1: + multi_run_manager = get_multi_run_manager() + if multi_run_manager.max_runs == 1: return SinglePacker(dp_world_size, seq_len, pad_to_multiple_of, tokenizer, transport_config, start_step) else: return MultiPacker(dp_world_size, seq_len, pad_to_multiple_of, tokenizer, transport_config, start_step) diff --git a/src/prime_rl/trainer/rl/train.py b/src/prime_rl/trainer/rl/train.py index fd5ddafba1b..4b2b9322978 100644 --- a/src/prime_rl/trainer/rl/train.py +++ b/src/prime_rl/trainer/rl/train.py @@ -1,3 +1,5 @@ +import prime_rl._compat # noqa: F401 — patch ring_flash_attn compat before import + from contextlib import nullcontext import time from datetime import timedelta @@ -5,21 +7,21 @@ # Import environment before any other imports # ruff: noqa: I001 -from prime_rl.trainer.models.layers.attn import substitute_prime_rl_flash_attn +from prime_rl.trainer.models.layers.attn import substitute_ring_attn from prime_rl.trainer.rl.broadcast import setup_weight_broadcast from prime_rl.utils.act_offloading import maybe_activation_offloading import torch import torch.distributed as dist -import torch.distributed.nn as dist_nn from torch.profiler import profile, ProfilerActivity, record_function -from loguru import logger from prime_rl.trainer.ckpt import setup_ckpt_managers from prime_rl.trainer.multi_ckpt import setup_multi_checkpoint_manager from prime_rl.trainer.optim import setup_optimizer, setup_multi_optimizer from prime_rl.trainer.scheduler import setup_scheduler, setup_multi_scheduler -from prime_rl.trainer.rl.config import RLTrainerConfig +from prime_rl.configs.trainer import TrainerConfig from prime_rl.trainer.rl.data import DataLoader, FakeDataLoader from prime_rl.utils.cp import ( + gather_for_cp, + gather_for_cp_wo_grad, setup_cp_params, shard_for_cp, ) @@ -28,6 +30,7 @@ compute_entropy, compute_loss, selective_log_softmax, + setup_loss_fn, shift_tensor_left, shift_tensor_right, ) @@ -41,37 +44,38 @@ from prime_rl.trainer.parallel_dims import get_parallel_dims from prime_rl.trainer.perf import get_perf_counter from prime_rl.trainer.utils import ( + GarbageCollection, MemoryProfiler, Tensors, export_benchmark_json, + filter_rl_trainer_tensor_stats_for_wandb, + get_zero_gradient_ratio, get_ckpt_disk_metrics, setup_torch_distributed, print_benchmark, get_response_lengths, ) from prime_rl.trainer.world import get_world -from prime_rl.trainer.runs import setup_runs, Progress, get_runs +from prime_rl.trainer.runs import setup_multi_run_manager, Progress, get_multi_run_manager from prime_rl.trainer.models.layers.lora import set_lora_num_tokens from prime_rl.utils.heartbeat import Heartbeat from prime_rl.utils.metrics_server import HealthServer, MetricsServer, RunStats from prime_rl.utils.monitor import setup_monitor -from prime_rl.utils.pydantic_config import parse_argv +from prime_rl.utils.config import cli +from prime_rl.utils.process import set_proc_title from prime_rl.utils.utils import clean_exit, resolve_latest_ckpt_step, to_col_format from ring_flash_attn import substitute_hf_flash_attn from torchtitan.distributed.utils import clip_grad_norm_ @clean_exit -@logger.catch(reraise=True) -def train(config: RLTrainerConfig): +def train(config: TrainerConfig): # Setup world and logger world = get_world() logger = setup_logger( config.log.level, - log_file=config.output_dir / "logs" / "trainer" / f"rank_{world.rank}.log" if config.log.file else None, + json_logging=config.log.json_logging, ) - logger.info(f"Starting RL trainer in {world}") - logger.info(f"Starting RL trainer in {world} in {config.output_dir}") # Print warning if running in benchmark mode @@ -105,32 +109,15 @@ def train(config: RLTrainerConfig): setup_torch_distributed( timeout=timedelta(seconds=config.dist_timeout_seconds), enable_gloo=config.model.fsdp_cpu_offload ) - torch.set_float32_matmul_precision("high") - - # Setup runs and offsets - setup_runs(config.output_dir, config.max_concurrent_runs, torch.device("cuda", world.local_rank)) - runs = get_runs() - - # Register validation and scaling hooks for LoRA - if config.model.lora: - trainer_lora = config.model.lora - - def validate_lora_rank(orch_config) -> tuple[bool, str]: - # Default to trainer's rank if not specified - if orch_config.model.lora.rank is None: - orch_config.model.lora.rank = trainer_lora.rank - if orch_config.model.lora.rank > trainer_lora.rank: - return ( - False, - f"model.lora.rank ({orch_config.model.lora.rank}) exceeds trainer max rank ({trainer_lora.rank})", - ) - return True, "" - - def compute_scaling(orch_config) -> float: - return orch_config.model.lora.alpha / orch_config.model.lora.rank - - runs.register_config_validation_hook(validate_lora_rank) - runs.register_scaling_hook(compute_scaling) + # Configurable to support ROCm/AMD GPUs where reduced precision + # matmul corrupts softmax over large vocabularies. Override via config + # (e.g. matmul_precision = "highest") on ROCm. + torch.set_float32_matmul_precision(config.matmul_precision) + + # Setup multi run manager and offsets (including LoRA validation/scaling hooks if applicable) + multi_run_manager = setup_multi_run_manager( + config.output_dir, config.max_concurrent_runs, torch.device("cuda", world.local_rank), config.model.lora + ) # Initialize parallel dimensions parallel_dims = get_parallel_dims(config.model) @@ -160,9 +147,12 @@ def compute_scaling(orch_config) -> float: logger.info(f"Initializing tokenizer ({config.tokenizer})") tokenizer = setup_tokenizer(config.tokenizer) + # Set up the loss function + logger.info(f"Setting up loss function ({config.loss})") + loss_fn = setup_loss_fn(config.loss) + # Set up the optimizer logger.info(f"Initializing optimizer ({config.optim})") - logger.info(f"Using `{config.loss.ratio_type}` importance ratio ({config.loss})") if config.max_concurrent_runs == 1: optimizer = setup_optimizer( @@ -170,6 +160,7 @@ def compute_scaling(orch_config) -> float: list(model.named_parameters()), parallel_dims, lora=config.model.lora is not None, + cpu_offload=config.model.optim_cpu_offload, ) scheduler = setup_scheduler(optimizer, config.scheduler, config.max_steps, config.optim.lr) else: @@ -184,13 +175,43 @@ def load_run_checkpoint(_optimizer, idx: int) -> None: logger.info(f"Using `{config.scheduler.type}` scheduler ({config.scheduler})") - # Set up weight broadcast - logger.info(f"Initializing weight broadcast ({config.weight_broadcast})") - weight_broadcast = setup_weight_broadcast(config.output_dir, config.weight_broadcast, config.model.lora) + # Set up weight broadcast (skip when using fake data since there's no inference server) + if config.data.fake: + weight_broadcast = None + logger.info("Skipping weight broadcast setup (fake data mode)") + else: + logger.info(f"Initializing weight broadcast ({config.weight_broadcast})") + weight_broadcast = setup_weight_broadcast(config.output_dir, config.weight_broadcast, config.model.lora) if parallel_dims.cp_enabled: - substitute_hf_flash_attn(parallel_dims.world_mesh["cp"].get_group(), heads_k_stride=1) - substitute_prime_rl_flash_attn(parallel_dims.world_mesh["cp"].get_group(), heads_k_stride=1) + cp_group = parallel_dims.world_mesh["cp"].get_group() + cp_rank = parallel_dims.world_mesh["cp"].get_local_rank() + if config.model.cp_style == "ring": + substitute_hf_flash_attn(cp_group, heads_k_stride=1) + substitute_ring_attn(cp_group, heads_k_stride=1, attn_impl=config.model.attn) + else: + from prime_rl.trainer.models.layers.ulysses_attn import ( + substitute_hf_ulysses_attn, + substitute_ulysses_attn, + ) + + substitute_hf_ulysses_attn(cp_group) + substitute_ulysses_attn(cp_group, attn_impl=config.model.attn) + from prime_rl.utils.cp import ( + assert_cp_style_supports_model, + setup_hybrid_cp, + setup_nemotron_h_cp, + setup_sparse_mla_cp, + ) + + assert_cp_style_supports_model(config.model.cp_style, model) + # sparse MLA is softmax (works with both ring and ulysses). + setup_sparse_mla_cp(model, cp_group, cp_rank, parallel_dims.cp) + # Linear-attn / Mamba layers are only configured under ulysses; with ring + # we'd have already raised above. + if config.model.cp_style == "ulysses": + setup_hybrid_cp(model, cp_group, cp_rank, parallel_dims.cp) + setup_nemotron_h_cp(model, cp_group, cp_rank, parallel_dims.cp) # Optionally, resume training from a checkpoint progress = Progress() @@ -205,18 +226,20 @@ def load_run_checkpoint(_optimizer, idx: int) -> None: # Set up the data loader (Optionally, use a fake data loader for debugging) logger.info(f"Initializing data loader ({config.data})") if config.data.fake: - dataloader = FakeDataLoader(config.data.fake, config.model.seq_len, parallel_dims.world_mesh["dp"].size()) + dataloader = FakeDataLoader(config.data.fake, config.model.seq_len, parallel_dims.get_mesh("dp").size()) else: dataloader = DataLoader( config.output_dir, progress.step, - parallel_dims.world_mesh["dp"].size(), + parallel_dims.get_mesh("dp").size(), config.model.seq_len, config.model.cp, tokenizer, config.rollout_transport, ) + gc_handler = GarbageCollection(config.gc.interval) if config.gc else None + logger.info(f"Starting training loop (max_steps={config.max_steps or 'infinite'})") is_first_step = True maybe_record_function = nullcontext @@ -227,25 +250,30 @@ def load_run_checkpoint(_optimizer, idx: int) -> None: while True: # Reset peak memory stats torch.cuda.reset_peak_memory_stats() + if gc_handler is not None: + gc_handler.run(progress.step) is_last_step = config.max_steps is not None and progress.step == config.max_steps # Broadcast weights at every step, (except step 0, because no need to broadcast the base model) # Also, with NCCL broadcast, we do not broadcast weights the last async level step as the orchestrator is already finished and will not initialize the receive on the inference; for filesystem broadcast, we do "broadcast" until the final step to allow to resume from the broadcast directory - last_async_level_steps = config.max_steps and progress.step >= config.max_steps - config.max_async_level - if progress.step > 0 and (not last_async_level_steps or config.weight_broadcast.type == "filesystem"): - broadcast_weights_start_time = time.perf_counter() - weight_broadcast.broadcast_weights(model, step=progress.step) - broadcast_weights_time = time.perf_counter() - broadcast_weights_start_time - # Clean up old broadcast directories (unless at ckpt interval if using filesystem weight broadcast) - ckpt_interval = config.ckpt and config.ckpt.interval - interval_to_keep = ckpt_interval if config.weight_broadcast.type == "filesystem" else None - if config.weight_broadcast.type == "filesystem": - weight_broadcast.maybe_clean(config.max_async_level, interval_to_keep) - else: + if weight_broadcast is None: broadcast_weights_time = 0 - # Usually the broadcast will set this. If broadcast is skipped, we need to reset this here. - for idx in runs.used_idxs: - runs.ready_to_update[idx] = False + else: + last_async_level_steps = config.max_steps and progress.step >= config.max_steps - config.max_async_level + if progress.step > 0 and (not last_async_level_steps or config.weight_broadcast.type == "filesystem"): + broadcast_weights_start_time = time.perf_counter() + weight_broadcast.broadcast_weights(model, step=progress.step) + broadcast_weights_time = time.perf_counter() - broadcast_weights_start_time + # Clean up old broadcast directories (unless at ckpt interval if using filesystem weight broadcast) + ckpt_interval = config.ckpt and config.ckpt.interval + interval_to_keep = ckpt_interval if config.weight_broadcast.type == "filesystem" else None + if config.weight_broadcast.type == "filesystem": + weight_broadcast.maybe_clean(config.max_async_level, interval_to_keep) + else: + broadcast_weights_time = 0 + # Usually the broadcast will set this. If broadcast is skipped, we need to reset this here. + for idx in multi_run_manager.used_idxs: + multi_run_manager.ready_to_update[idx] = False if ( ckpt_manager is not None @@ -253,21 +281,23 @@ def load_run_checkpoint(_optimizer, idx: int) -> None: and not (is_first_step or is_last_step) and progress.step % config.ckpt.interval == 0 ): - # Single-run: Save full checkpoint - logger.info(f"Saving checkpoint at step {progress.step}") - save_ckpt_start_time = time.perf_counter() - ckpt_manager.save(progress.step, model, [optimizer], scheduler, progress) - save_ckpt_time = time.perf_counter() - save_ckpt_start_time + save_ckpt_time = 0 + + if not config.ckpt.weights_only: + # Single-run: Save full checkpoint + logger.info(f"Saving checkpoint at step {progress.step}") + save_ckpt_start_time = time.perf_counter() + ckpt_manager.save(progress.step, model, [optimizer], scheduler, progress) + save_ckpt_time += time.perf_counter() - save_ckpt_start_time - # Maybe clean up old checkpoints ckpt_manager.maybe_clean() # Save weight checkpoint if weight_ckpt_manager is not None: logger.info(f"Saving weight checkpoint at step {progress.step}") + save_ckpt_start_time = time.perf_counter() weight_ckpt_manager.save(progress.step, model, tokenizer) - - # Maybe clean up old weight checkpoint + save_ckpt_time += time.perf_counter() - save_ckpt_start_time weight_ckpt_manager.maybe_clean() elif config.max_concurrent_runs > 1: # Multi-run: Save per-run checkpoints (each run has its own interval from orchestrator config) @@ -308,10 +338,7 @@ def load_run_checkpoint(_optimizer, idx: int) -> None: seq_len = micro_batches[0]["input_ids"].shape[1] # Normalize by the local number of unmasked tokens in the batch (per-batch length normalization) - if config.loss.ratio_type == "token": - loss_scale = sum(micro_batch["loss_mask"].sum().item() for micro_batch in micro_batches) - elif config.loss.ratio_type == "sequence": - loss_scale = batch_size + loss_scale = sum(micro_batch["loss_mask"].sum().item() for micro_batch in micro_batches) loss_scale = max(loss_scale, 1) logger.debug(f"Starting forward and backward pass ({batch_size=})") @@ -330,20 +357,52 @@ def load_run_checkpoint(_optimizer, idx: int) -> None: teacher_logprobs = ( micro_batch["teacher_logprobs"].to("cuda") if micro_batch["teacher_logprobs"] is not None else None ) + routed_experts = ( + micro_batch["routed_experts"].to("cuda") if micro_batch["routed_experts"] is not None else None + ) + + if routed_experts is None and config.enable_router_replay: + raise ValueError( + "You must set `enable_return_routed_experts=True` in the inference config or pass `--enable-return-routed-experts` to vLLM server to use router replay." + ) + + if routed_experts is not None and not config.enable_router_replay: + # we could've gotten routed experts from the inference server, but we didn't enable router replay + routed_experts = None + + # Multimodal fields (Qwen3-VL) - only present for VLM training + pixel_values = ( + micro_batch["pixel_values"].to("cuda") if micro_batch.get("pixel_values") is not None else None + ) + image_grid_thw = ( + micro_batch["image_grid_thw"].to("cuda") if micro_batch.get("image_grid_thw") is not None else None + ) + mm_token_type_ids = ( + micro_batch["mm_token_type_ids"].to("cuda") + if micro_batch.get("mm_token_type_ids") is not None + else None + ) labels = shift_tensor_left(input_ids) + # VLM + CP is not supported: MRoPE requires global positions but CP shards the sequence + if cp_enabled and pixel_values is not None: + raise NotImplementedError("Context parallelism is not supported with VLM/multimodal training") + if cp_enabled: - input_ids, forward_position_ids = setup_cp_params(input_ids, position_ids, cp_rank, cp_size, cp_group) + input_ids, forward_position_ids = setup_cp_params( + input_ids, position_ids, cp_rank, cp_size, cp_group, cp_style=config.model.cp_style + ) labels = shard_for_cp(labels, cp_rank=cp_rank, cp_world_size=cp_size) + if routed_experts is not None: + routed_experts = shard_for_cp(routed_experts, cp_rank=cp_rank, cp_world_size=cp_size) else: forward_position_ids = position_ids if config.model.lora: lora_num_tokens = micro_batch["lora_num_tokens"].to("cuda") if cp_enabled: - chunk_size = input_ids.shape[1] # We pad to multiple of cp so this should be fine - logger.debug(f"[Rank {world.rank}] {cp_rank=} {cp_size=} {cp_group=} {chunk_size=}") + chunk_size = input_ids.shape[1] # Convert to cumsum, adjust for CP chunk, convert back to num_tokens cu_offsets = lora_num_tokens.cumsum(dim=0, dtype=torch.int32) adjusted_cu = torch.clip(cu_offsets - chunk_size * cp_rank, min=0, max=chunk_size) @@ -352,27 +411,41 @@ def load_run_checkpoint(_optimizer, idx: int) -> None: ) set_lora_num_tokens(lora_num_tokens) - temperature = micro_batch["temperature"] + temperatures = micro_batch["temperatures"].to("cuda") + + # Shard temperatures for context parallelism if enabled + if cp_enabled: + temperatures = shard_for_cp(temperatures, cp_rank=cp_rank, cp_world_size=cp_size) - # Forward pass + # Forward pass with per-token temperatures with maybe_record_function("forward"), maybe_activation_offloading(config.model.ac_offloading): - out = forward(model, input_ids, forward_position_ids, labels=labels, temperature=temperature) + out = forward( + model, + input_ids, + forward_position_ids, + labels=labels, + temperature=temperatures, + pixel_values=pixel_values, + image_grid_thw=image_grid_thw, + mm_token_type_ids=mm_token_type_ids, + routed_experts=routed_experts, + ) if out.get("logprobs") is None: + # VanillaOutputLinear was used - need to compute logprobs externally with per-token temps assert out.get("logits") is not None, "Logits must be provided to compute logprobs" - logits = out["logits"] / float(temperature) - out["logprobs"] = selective_log_softmax(logits, labels) - out["entropy"] = compute_entropy(logits) + logits = out["logits"] + # Per-token temperature scaling: temperatures is [batch, seq], logits is [batch, seq, vocab] + scaled_logits = logits / temperatures.unsqueeze(-1) + out["logprobs"] = selective_log_softmax(scaled_logits, labels) + out["entropy"] = compute_entropy(scaled_logits) + # else: FusedOutputLinear was used - logprobs already computed with per-token temperatures if cp_enabled: - logprobs = dist_nn.all_gather(out["logprobs"], group=cp_group) - out["logprobs"] = torch.cat(logprobs, dim=1) - - entropies = [torch.zeros_like(out["entropy"]) for _ in range(cp_size)] - dist.all_gather(entropies, out["entropy"], group=cp_group) - out["entropy"] = torch.cat(entropies, dim=1) + out["logprobs"] = gather_for_cp(out["logprobs"], cp_group) + out["entropy"] = gather_for_cp_wo_grad(out["entropy"], cp_size, cp_group) - vocab_size = model.config.vocab_size + vocab_size = getattr(model.config, "vocab_size", None) or model.config.text_config.vocab_size # This is not really necessary as the first token should be masked out, but we do it anyway to be sure out["logprobs"] = shift_tensor_right( out["logprobs"], pad_value=torch.log(torch.tensor(1.0 / vocab_size)).item() @@ -391,8 +464,9 @@ def load_run_checkpoint(_optimizer, idx: int) -> None: else None, advantages=advantages.squeeze().split(response_lengths), loss_mask=loss_mask.squeeze().split(response_lengths), - loss_config=config.loss, + loss_fn=loss_fn, loss_scale=loss_scale, + sft_loss=micro_batch["sft_loss"], ) # Backward pass @@ -400,8 +474,6 @@ def load_run_checkpoint(_optimizer, idx: int) -> None: loss.backward() # Add relevant tensors to tensor dict for logging purposes - tensors["trainer_probs"].append(torch.exp(out["logprobs"])[loss_mask].detach().to("cpu")) - tensors["inference_probs"].append(torch.exp(inference_logprobs)[loss_mask].detach().to("cpu")) tensors["entropy"].append(out["entropy"][loss_mask].detach().to("cpu")) tensors["loss"].append(loss.detach().to("cpu").unsqueeze(0)) @@ -417,18 +489,23 @@ def load_run_checkpoint(_optimizer, idx: int) -> None: tensors[key].append(loss_tensor) # Debug log with *local, micro step* stats - micro_step_message = f"Micro Step {micro_step}/{len(micro_batches)} | Loss: {tensors['loss'][-1].mean().item():.4f} | Entropy: {tensors['entropy'][-1].mean().item():.4f} | Mismatch KL: {tensors['mismatch_kl'][-1].mean().item():.4f}" + micro_step_message = f"Micro Step {micro_step}/{len(micro_batches)} | Loss: {tensors['loss'][-1].mean().item():.4f} | Entropy: {tensors['entropy'][-1].mean().item():.4f}" + if "mismatch_kl" in tensors: + micro_step_message += f" | Mismatch KL: {tensors['mismatch_kl'][-1].mean().item():.4f}" if "max_vio" in tensors: micro_step_message += f" | Max Vio: {tensors['max_vio'][-1].mean().item():.4f}" logger.debug(micro_step_message) # Optionally, clip the gradients + grad_norm: torch.Tensor | None = None + if config.optim.max_norm is not None: + grad_norm = clip_grad_norm_( + model.parameters(), max_norm=config.optim.max_norm, ep_enabled=parallel_dims.ep_enabled + ) + if grad_norm.device.type == "cpu": + grad_norm = grad_norm.to(torch.device("cuda")) - grad_norm = clip_grad_norm_( - model.parameters(), max_norm=config.optim.max_norm, ep_enabled=parallel_dims.ep_enabled - ) - if grad_norm.device.type == "cpu": - grad_norm = grad_norm.to(torch.device("cuda")) + zero_grad_ratio = get_zero_gradient_ratio(model.parameters(), parallel_dims.dp_replicate) # Update the model parameters optimizer.step() @@ -452,18 +529,22 @@ def load_run_checkpoint(_optimizer, idx: int) -> None: # Compute step metrics num_local_tokens = seq_len * batch_size - num_tokens = parallel_dims.world_mesh["dp"].size() * num_local_tokens + num_tokens = parallel_dims.get_mesh("dp").size() * num_local_tokens progress.total_tokens += num_tokens progress.total_samples += batch_size perf_counter = get_perf_counter(model, seq_len) - perf_counter.count_tokens(num_tokens) - throughput = perf_counter.get_tokens_per_second() or 0 - mfu = perf_counter.get_mfu() or 0 + throughput = perf_counter.get_step_tokens_per_second(num_tokens, forward_backward_time) + mfu = perf_counter.get_step_mfu(num_tokens, forward_backward_time) peak_memory = torch.cuda.max_memory_reserved() / 1024**3 # GiB # Log step metrics step_time = time.perf_counter() - step_start_time - step_message = f"Step {progress.step} | Time: {step_time:.2f}s | Loss: {tensor_stats['loss/mean']:.4f} | Entropy: {tensor_stats['entropy/mean']:.4f} | Mismatch KL: {tensor_stats['mismatch_kl/mean']:.4f} | Grad. Norm: {grad_norm:.4f} | LR: {current_lr:.2e} | Throughput: {throughput:.0f} tokens/s | MFU: {mfu:.1f}% | Peak Mem.: {peak_memory:.1f} GiB" + step_message = f"Step {progress.step} | Time: {step_time:.2f}s | Loss: {tensor_stats['loss/mean']:.4f} | Entropy: {tensor_stats['entropy/mean']:.4f}" + if "mismatch_kl/mean" in tensor_stats: + step_message += f" | Mismatch KL: {tensor_stats['mismatch_kl/mean']:.4f}" + if grad_norm is not None: + step_message += f" | Grad. Norm: {grad_norm:.4f}" + step_message += f" | LR: {current_lr:.2e} | Throughput: {throughput:.0f} tokens/s | MFU: {mfu:.1f}% | Peak Mem.: {peak_memory:.1f} GiB" if "max_vio/mean" in tensor_stats: step_message += f" | Max Vio: {tensor_stats['max_vio/mean']:.4f}" logger.success(step_message) @@ -481,14 +562,21 @@ def load_run_checkpoint(_optimizer, idx: int) -> None: # Log optimizer metrics optim_metrics = { "optim/lr": current_lr, - "optim/grad_norm": grad_norm.item(), + "optim/zero_grad_ratio": zero_grad_ratio, "step": progress.step, } + if grad_norm is not None: + optim_metrics["optim/grad_norm"] = grad_norm.item() monitor.log(optim_metrics, step=progress.step) - # Log tensor stats + # Compute derived metrics + entropy_mean = tensor_stats.get("entropy/mean", 0.0) + mismatch_kl_mean = tensor_stats.get("mismatch_kl/mean") + if mismatch_kl_mean is not None and entropy_mean > 0: + tensor_stats["kl_ent_ratio/mean"] = mismatch_kl_mean / entropy_mean + tensor_stats["step"] = progress.step - monitor.log(tensor_stats, step=progress.step) + monitor.log(filter_rl_trainer_tensor_stats_for_wandb(tensor_stats), step=progress.step) # Log time metrics time_metrics = { @@ -513,20 +601,21 @@ def load_run_checkpoint(_optimizer, idx: int) -> None: step=progress.step, loss=tensor_stats["loss/mean"], throughput=throughput, - grad_norm=grad_norm.item(), + grad_norm=grad_norm.item() if grad_norm is not None else None, peak_memory_gib=peak_memory, learning_rate=current_lr, mfu=mfu, entropy=tensor_stats.get("entropy/mean", 0.0), mismatch_kl=tensor_stats.get("mismatch_kl/mean", 0.0), + zero_grad_ratio=zero_grad_ratio, ) # Update run/LoRA metrics - runs = get_runs() + multi_run_manager = get_multi_run_manager() runs_discovered = len(list(config.output_dir.glob("run_*"))) run_stats = [] - for idx in runs.used_idxs: - run_id = runs.idx_2_id[idx] - run_progress = runs.progress[idx] + for idx in multi_run_manager.used_idxs: + run_id = multi_run_manager.idx_2_id[idx] + run_progress = multi_run_manager.progress[idx] if config.max_concurrent_runs == 1: lr = optimizer.param_groups[0]["lr"] else: @@ -537,12 +626,12 @@ def load_run_checkpoint(_optimizer, idx: int) -> None: step=run_progress.step, total_tokens=run_progress.total_tokens, learning_rate=lr, - ready=runs.ready_to_update[idx], + ready=multi_run_manager.ready_to_update[idx], ) ) metrics_server.update_runs( runs_discovered=runs_discovered, - runs_max=runs.max_runs, + runs_max=multi_run_manager.max_runs, run_stats=run_stats, ) @@ -561,14 +650,14 @@ def load_run_checkpoint(_optimizer, idx: int) -> None: prof.export_chrome_trace(trace_file) logger.info(f"Saved trace to {trace_file}") - # Write final checkpoint - if ckpt_manager is not None: - logger.info("Writing final checkpoint") - ckpt_manager.save(progress.step, model, [optimizer], scheduler, progress) + # Write final checkpoint (only for single-run mode; multi-run checkpoints are managed by MultiCheckpointManager) + if config.max_concurrent_runs == 1 and ckpt_manager is not None: + if not (config.ckpt and config.ckpt.weights_only): + logger.info("Writing final checkpoint") + ckpt_manager.save(progress.step, model, [optimizer], scheduler, progress) ckpt_manager.maybe_clean() - # Write final checkpoint - if weight_ckpt_manager is not None: + if config.max_concurrent_runs == 1 and weight_ckpt_manager is not None: logger.info("Writing final weight checkpoint") weight_ckpt_manager.save(progress.step, model, tokenizer) weight_ckpt_manager.maybe_clean() @@ -593,8 +682,8 @@ def load_run_checkpoint(_optimizer, idx: int) -> None: def main(): """Main entry-point for RL trainer. Run using `uv run trainer`""" - - train(parse_argv(RLTrainerConfig)) + set_proc_title("Trainer") + train(cli(TrainerConfig)) if __name__ == "__main__": diff --git a/src/prime_rl/trainer/runs.py b/src/prime_rl/trainer/runs.py index 5bccbaa984a..9372386764b 100644 --- a/src/prime_rl/trainer/runs.py +++ b/src/prime_rl/trainer/runs.py @@ -8,11 +8,13 @@ import torch.distributed as dist import torch.distributed.distributed_c10d as c10d +from prime_rl.configs.trainer import LoRAConfig from prime_rl.trainer.world import get_world from prime_rl.utils.logger import get_logger +from prime_rl.utils.pathing import get_all_ckpt_steps, get_stable_ckpt_steps if TYPE_CHECKING: - from prime_rl.orchestrator.config import OrchestratorConfig + from prime_rl.configs.orchestrator import OrchestratorConfig from prime_rl.trainer.models.layers.lora import MultiLoRALinear @@ -23,9 +25,13 @@ class Progress: total_samples: int = 0 -class Runs: +class MultiRunManager: """This class stores information about the runs in the system.""" + # ========================================================================= + # Initialization + # ========================================================================= + def __init__(self, output_dir: Path, max_runs: int, device: torch.device = torch.device("cpu")): self.output_dir = output_dir self.max_runs = max_runs @@ -41,20 +47,22 @@ def __init__(self, output_dir: Path, max_runs: int, device: torch.device = torch self._creation_hooks: list[Callable[[int, str], None]] = [] self._deletion_hooks: list[Callable[[int, str], None]] = [] - self._create_run_data_hooks: list[Callable[[int, str, "OrchestratorConfig"], None]] = [] - self._delete_run_data_hooks: list[Callable[[int, str], None]] = [] + self._discovered_hooks: list[Callable[[int, str, "OrchestratorConfig"], None]] = [] + self._forgotten_hooks: list[Callable[[int, str], None]] = [] self._config_validation_hooks: list[Callable[["OrchestratorConfig"], tuple[bool, str]]] = [] - self._scaling_hook: Callable[["OrchestratorConfig"], float] | None = None # We use the store to keep other ranks in sync with master self.store = c10d._get_default_store() self.world = get_world() - # Track id_2_idx state at last sync_runs to calculate diffs + # Track id_2_idx state at last synchronize_state to calculate diffs self._last_synced_id_2_idx: dict[str, int] = {} # Store modules with their FQN prefixes for parameter management self._modules: list[tuple[str, "MultiLoRALinear"]] = [] + # Optional conversion applied to adapter state dicts (e.g. PrimeRL -> HF key rename) + self._adapter_state_dict_converter: Callable[[dict[str, torch.Tensor]], dict[str, torch.Tensor]] | None = None + # Initialize lora globals on device so runs.* ARE the global tensors from prime_rl.trainer.models.layers.lora import ( get_lora_num_tokens, @@ -76,6 +84,50 @@ def __init__(self, output_dir: Path, max_runs: int, device: torch.device = torch ) self.scaling_factors = get_multilora_scaling() + # ========================================================================= + # Hook Registration + # ========================================================================= + + def register_creation_hook(self, hook: Callable[[int, str], None]) -> None: + """Register a hook to be called when a new run is created. + + Args: + hook: A callable that takes (idx: int, run_id: str) as arguments. + Called on all ranks when a new run is added to the system. + """ + self._creation_hooks.append(hook) + + def register_deletion_hook(self, hook: Callable[[int, str], None]) -> None: + """Register a hook to be called when a run is deleted. + + Args: + hook: A callable that takes (idx: int, run_id: str) as arguments. + Called on all ranks when a run is removed from the system. + """ + self._deletion_hooks.append(hook) + + def register_discovered_hook(self, hook: Callable[[int, str, "OrchestratorConfig"], None]) -> None: + """Register a hook to be called when a new run is discovered (master only). + + Args: + hook: A callable that takes (idx: int, run_id: str, config: OrchestratorConfig). + Called only on master rank in discover_runs() when a new run is found. + """ + if not self.world.is_master: + raise RuntimeError("register_discovered_hook() must only be called on the master rank") + self._discovered_hooks.append(hook) + + def register_forgotten_hook(self, hook: Callable[[int, str], None]) -> None: + """Register a hook to be called when a run is forgotten/removed (master only). + + Args: + hook: A callable that takes (idx: int, run_id: str). + Called only on master rank in discover_runs() when a run is removed. + """ + if not self.world.is_master: + raise RuntimeError("register_forgotten_hook() must only be called on the master rank") + self._forgotten_hooks.append(hook) + def register_config_validation_hook(self, hook: Callable[["OrchestratorConfig"], tuple[bool, str]]) -> None: """Register a hook to validate orchestrator config when a run is created. @@ -83,33 +135,87 @@ def register_config_validation_hook(self, hook: Callable[["OrchestratorConfig"], hook: A callable that takes (config: OrchestratorConfig) and returns (is_valid: bool, error_message: str). Error message is used when invalid. """ + if not self.world.is_master: + raise RuntimeError("register_config_validation_hook() must only be called on the master rank") self._config_validation_hooks.append(hook) - def register_scaling_hook(self, hook: Callable[["OrchestratorConfig"], float]) -> None: - """Register a hook to compute LoRA scaling factor for a run. + # ========================================================================= + # Module Registration and Parameter Management + # ========================================================================= + + def register_adapter_state_dict_converter( + self, converter: Callable[[dict[str, torch.Tensor]], dict[str, torch.Tensor]] + ) -> None: + """Register a converter applied in-place to adapter state dicts (e.g. model.convert_adapter_to_hf).""" + self._adapter_state_dict_converter = converter + + def register_module(self, prefix: str, module: "MultiLoRALinear") -> None: + """Register a MultiLoRALinear module with its FQN prefix. + + This allows MultiRunManager to manage parameter access, reset, and state dict slicing + for multi-adapter LoRA modules. Args: - hook: A callable that takes (config: OrchestratorConfig) and returns the scaling factor. + prefix: The module's fully qualified name in the model + (e.g., "model.layers.0.self_attn.q_proj") + module: The MultiLoRALinear module to register """ - self._scaling_hook = hook + self._modules.append((prefix, module)) - def register_create_run_data_hook(self, hook: Callable[[int, str, "OrchestratorConfig"], None]) -> None: - """Register a hook to be called when run data is created (master only). + def get_named_parameters_for_run(self, idx: int) -> list[tuple[str, torch.nn.Parameter]]: + """Get named parameters for a specific run index. Args: - hook: A callable that takes (idx: int, run_id: str, config: OrchestratorConfig). - Called only on master rank when a new run's data structures are initialized. + idx: The run index to get parameters for + + Returns: + List of (name, parameter) tuples for the specified run index + """ + params = [] + for prefix, module in self._modules: + for name, param in module.named_parameters_for_adapter(idx): + params.append((f"{prefix}.{name}.weight", param)) + return params + + def get_state_dict_for_run(self, idx: int) -> dict[str, torch.Tensor]: + """Get state dict for a specific run index. + + Args: + idx: The run index to get state dict for + + Returns: + State dict for the specified run index """ - self._create_run_data_hooks.append(hook) + state_dict = {} + for prefix, module in self._modules: + # Check if module has a custom state_dict_for_adapter method (e.g., MoE modules) + # which returns vLLM-compatible per-expert format + if hasattr(module, "state_dict_for_adapter"): + for name, tensor in module.state_dict_for_adapter(idx).items(): + state_dict[f"{prefix}.{name}"] = tensor.detach() + else: + # Default: use named_parameters_for_adapter + for name, param in module.named_parameters_for_adapter(idx): + state_dict[f"{prefix}.{name}.weight"] = param.detach() - def register_delete_run_data_hook(self, hook: Callable[[int, str], None]) -> None: - """Register a hook to be called when run data is deleted (master only). + if self._adapter_state_dict_converter is not None: + state_dict = self._adapter_state_dict_converter(state_dict) + return state_dict + + def reset_run_parameters(self, idx: int) -> None: + """Reset parameters for a specific run index. + + Called when a new run is created to initialize fresh adapter weights. Args: - hook: A callable that takes (idx: int, run_id: str). - Called only on master rank when a run's data structures are being deleted. + idx: The run index to reset parameters for """ - self._delete_run_data_hooks.append(hook) + for _, module in self._modules: + module.reset_parameters(idx) + + # ========================================================================= + # Config Loading + # ========================================================================= def get_orchestrator_config(self, run_id: str) -> Optional["OrchestratorConfig"]: """Load and validate orchestrator config for a run. @@ -117,15 +223,11 @@ def get_orchestrator_config(self, run_id: str) -> Optional["OrchestratorConfig"] Returns None if config doesn't exist, fails to parse, or fails validation. Writes error to config dir for orchestrator to consume. """ - config_path = self.output_dir / run_id / "configs" / "orch.toml" + config_path = self.output_dir / run_id / "control" / "orch.toml" config_dir = config_path.parent - error_path = config_dir / "error.txt" + error_path = config_dir / "config_validation_error.txt" if not config_path.exists(): - if not error_path.exists(): - config_dir.mkdir(parents=True, exist_ok=True) - with open(error_path, "w") as f: - f.write(f"No orchestrator config found at {config_path}\n") self.logger.error(f"Run {run_id}: No orchestrator config found at {config_path}") return None @@ -133,13 +235,13 @@ def get_orchestrator_config(self, run_id: str) -> Optional["OrchestratorConfig"] with open(config_path, "rb") as f: config_dict = tomli.load(f) - from prime_rl.orchestrator.config import OrchestratorConfig + from prime_rl.configs.orchestrator import OrchestratorConfig config = OrchestratorConfig(**config_dict) except Exception as e: - config_dir.mkdir(parents=True, exist_ok=True) - with open(error_path, "w") as f: - f.write(f"Error parsing orchestrator config:\n{str(e)}\n") + if error_path.parent.exists(): + with open(error_path, "w") as f: + f.write(f"Error parsing orchestrator config:\n{str(e)}\n") self.logger.error(f"Run {run_id}: Error parsing orchestrator config: {e}") return None @@ -148,9 +250,9 @@ def get_orchestrator_config(self, run_id: str) -> Optional["OrchestratorConfig"] is_valid, error_message = hook(config) if not is_valid: self.logger.error(f"Run {run_id}: {error_message}") - config_dir.mkdir(parents=True, exist_ok=True) - with open(error_path, "w") as f: - f.write(f"{error_message}\n") + if error_path.parent.exists(): + with open(error_path, "w") as f: + f.write(f"{error_message}\n") return None # Config is valid, remove any stale error file @@ -159,12 +261,34 @@ def get_orchestrator_config(self, run_id: str) -> Optional["OrchestratorConfig"] return config - def _delete_run_data(self, deleted_run: str, deleted_idx: int) -> None: - """Update data structures for a deleted run""" - # Call delete hooks before removing data - for hook in self._delete_run_data_hooks: - hook(deleted_idx, deleted_run) + # ========================================================================= + # Internal Run Data Helpers + # ========================================================================= + def _create_run_data(self, new_run: str, new_id: int, config: "OrchestratorConfig") -> None: + """Update data structures for a new run (no hooks or param reset).""" + self.id_2_idx[new_run] = new_id + self.unused_idxs.remove(new_id) + self.idx_2_id[new_id] = new_run + + # Set progress based on resume_step config (match orchestrator behavior) + self.progress[new_id] = Progress() + if config.ckpt is None or config.ckpt.resume_step is None: + self.progress[new_id].step = 0 + elif config.ckpt.resume_step == -1: + ckpt_dir = self.get_run_dir(new_id) / "checkpoints" + # In multi-run, the trainer writes STABLE after saving LoRA weights to the run's checkpoint dir. + # In single-run, only the orchestrator writes checkpoints here (trainer has its own dir), so no STABLE exists. + steps = get_stable_ckpt_steps(ckpt_dir) if self.max_runs > 1 else get_all_ckpt_steps(ckpt_dir) + self.progress[new_id].step = max(steps) if steps else 0 + else: + self.progress[new_id].step = config.ckpt.resume_step + + # Store the parsed config + self.config[new_id] = config + + def _delete_run_data(self, deleted_run: str, deleted_idx: int) -> None: + """Update data structures for a deleted run (internal cleanup only, no hooks).""" del self.progress[deleted_idx] if deleted_idx in self.config: del self.config[deleted_idx] @@ -177,23 +301,6 @@ def _delete_run_data(self, deleted_run: str, deleted_idx: int) -> None: del self.idx_2_id[deleted_idx] del self.id_2_idx[deleted_run] - def _create_run_data(self, new_run: str, new_id: int, config: "OrchestratorConfig") -> None: - """Update data structures for a new run (no hooks or param reset).""" - self.id_2_idx[new_run] = new_id - self.unused_idxs.remove(new_id) - self.idx_2_id[new_id] = new_run - - # Get progress - self.progress[new_id] = Progress() - - prev_ckpt_steps = [ - int(i.stem.split("_")[-1]) for i in (self.get_run_dir(new_id) / "checkpoints").glob("step_*") - ] - self.progress[new_id].step = max(prev_ckpt_steps) if prev_ckpt_steps else 0 - - # Store the parsed config - self.config[new_id] = config - def _create_run_hooks(self, new_id: int, new_run: str) -> None: """Reset parameters and call creation hooks for a run.""" self.reset_run_parameters(new_id) @@ -205,18 +312,63 @@ def _delete_run_hooks(self, deleted_idx: int, deleted_run: str) -> None: for hook in self._deletion_hooks: hook(deleted_idx, deleted_run) - def check_for_changes(self) -> None: - """Detect run changes and update data structures. Must be followed by sync_runs(). + # ========================================================================= + # Run Discovery, Synchronization, and Eviction + # ========================================================================= - Only updates mappings and data structures. Hooks and parameter resets - are deferred to sync_runs() so all ranks execute them together. + def evict_run(self, idx: int, reason: str) -> None: + """Evict a run by writing the reason to a file for the orchestrator to read. + + The orchestrator will error and surface this reason. The run will be + ignored by future discover_runs() calls. + Note that the run is not deleted on master until the next discover_runs() call. + And not deleted on other ranks until the next synchronize_state() call. + + Args: + idx: The run index to evict + reason: The reason for eviction (will be shown to user) """ + if not self.world.is_master: + raise RuntimeError("evict_run() must only be called on the master rank") + + if idx not in self.idx_2_id: + self.logger.warning(f"Run index {idx} not found, cannot evict") + return + + run_id = self.idx_2_id[idx] + run_dir = self.output_dir / run_id + config_dir = run_dir / "control" + config_dir.mkdir(parents=True, exist_ok=True) + + evicted_path = config_dir / "evicted.txt" + with open(evicted_path, "w") as f: + f.write(reason) + + self.logger.warning(f"Evicted run {run_id} (idx={idx}): {reason}") + + def discover_runs(self) -> None: + """Detect run changes and update data structures (master only). Must be followed by synchronize_state(). + + Scans for new/deleted runs, calls forgotten/discovered hooks (master only), + and updates internal data structures. Hooks and parameter resets for all ranks + are deferred to synchronize_state(). + """ + if not self.world.is_master: + raise RuntimeError("discover_runs() must only be called on the master rank") run_ids = {run_path.stem for run_path in self.output_dir.glob("run_*")} + + # Filter out evicted runs + evicted_runs = {run_id for run_id in run_ids if (self.output_dir / run_id / "control" / "evicted.txt").exists()} + run_ids = run_ids - evicted_runs + deleted_runs = self.id_2_idx.keys() - run_ids new_runs = run_ids - self.id_2_idx.keys() for deleted_run in deleted_runs: deleted_idx = self.id_2_idx[deleted_run] + # Call forgotten hooks (master only) before cleanup + for hook in self._forgotten_hooks: + hook(deleted_idx, deleted_run) self._delete_run_data(deleted_run, deleted_idx) for new_run in new_runs: @@ -228,12 +380,13 @@ def check_for_changes(self) -> None: continue self._create_run_data(new_run, new_id, config) - for hook in self._create_run_data_hooks: + # Call discovered hooks (master only) after data setup + for hook in self._discovered_hooks: hook(new_id, new_run, config) except StopIteration: continue - def sync_runs(self) -> None: + def synchronize_state(self) -> None: """Sync run state across ranks and execute hooks. Master calculates what changed since last sync using _last_synced_id_2_idx. @@ -242,14 +395,8 @@ def sync_runs(self) -> None: """ if self.world.is_master: - # Compute scaling factors for new runs before sync - new_runs = self.id_2_idx.keys() - self._last_synced_id_2_idx.keys() - for new_run in new_runs: - new_id = self.id_2_idx[new_run] - if self._scaling_hook is not None: - self.scaling_factors[new_id] = self._scaling_hook(self.config[new_id]) - # Include configs for new runs so non-master ranks have them + new_runs = self.id_2_idx.keys() - self._last_synced_id_2_idx.keys() new_configs = {new_run: self.config[self.id_2_idx[new_run]] for new_run in new_runs} sync_data = { @@ -306,6 +453,10 @@ def sync_runs(self) -> None: if self.world.is_master: self._last_synced_id_2_idx = self.id_2_idx.copy() + # ========================================================================= + # Properties and Accessors + # ========================================================================= + @property def used_idxs(self): return sorted(self.idx_2_id.keys()) @@ -320,101 +471,61 @@ def run_dirs(self) -> list[Path]: def get_run_dir(self, idx: int) -> Path: return self.output_dir / self.idx_2_id[idx] - def register_creation_hook(self, hook: Callable[[int, str], None]) -> None: - """Register a hook to be called when a new run is created. - - Args: - hook: A callable that takes (idx: int, run_id: str) as arguments. - Called on all ranks when a new run is added to the system. - """ - self._creation_hooks.append(hook) - - def register_deletion_hook(self, hook: Callable[[int, str], None]) -> None: - """Register a hook to be called when a run is deleted. - - Args: - hook: A callable that takes (idx: int, run_id: str) as arguments. - Called on all ranks when a run is removed from the system. - """ - self._deletion_hooks.append(hook) - - def register_module(self, prefix: str, module: "MultiLoRALinear") -> None: - """Register a MultiLoRALinear module with its FQN prefix. - - This allows Runs to manage parameter access, reset, and state dict slicing - for multi-adapter LoRA modules. - - Args: - prefix: The module's fully qualified name in the model - (e.g., "model.layers.0.self_attn.q_proj") - module: The MultiLoRALinear module to register - """ - self._modules.append((prefix, module)) - - def get_named_parameters_for_run(self, idx: int) -> list[tuple[str, torch.nn.Parameter]]: - """Get named parameters for a specific run index. - - Args: - idx: The run index to get parameters for - - Returns: - List of (name, parameter) tuples for the specified run index - """ - params = [] - for prefix, module in self._modules: - for name, param in module.named_parameters_for_adapter(idx): - params.append((f"{prefix}.{name}.weight", param)) - return params - - def get_state_dict_for_run(self, idx: int) -> dict[str, torch.Tensor]: - """Get state dict for a specific run index. - - Args: - idx: The run index to get state dict for - - Returns: - State dict for the specified run index - """ - state_dict = {} - for prefix, module in self._modules: - # Check if module has a custom state_dict_for_adapter method (e.g., MoE modules) - # which returns vLLM-compatible per-expert format - if hasattr(module, "state_dict_for_adapter"): - for name, tensor in module.state_dict_for_adapter(idx).items(): - state_dict[f"{prefix}.{name}"] = tensor.detach() - else: - # Default: use named_parameters_for_adapter - for name, param in module.named_parameters_for_adapter(idx): - state_dict[f"{prefix}.{name}.weight"] = param.detach() - return state_dict - - def reset_run_parameters(self, idx: int) -> None: - """Reset parameters for a specific run index. - - Called when a new run is created to initialize fresh adapter weights. - - Args: - idx: The run index to reset parameters for - """ - for _, module in self._modules: - module.reset_parameters(idx) - def __repr__(self): - return f"Runs(max={self.max_runs})[{self.idx_2_id.keys()}]" - - -# Singleton instance of Tenants -_RUNS: Runs | None = None - - -def get_runs() -> Runs: - """Returns the World. If not initialized, it will initialize.""" - global _RUNS - if _RUNS is None: - raise RuntimeError("Runs not initialized. Please call `setup_runs` first.") - return _RUNS - - -def setup_runs(output_dir: Path, max_runs: int, device: torch.device): - global _RUNS - _RUNS = Runs(output_dir, max_runs, device) + return f"MultiRunManager(max={self.max_runs})[{self.idx_2_id.keys()}]" + + +# Singleton instance of MultiRunManager +_MULTI_RUN_MANAGER: MultiRunManager | None = None + + +def get_multi_run_manager() -> MultiRunManager: + """Returns the MultiRunManager singleton. Must be initialized first via setup_multi_run_manager().""" + global _MULTI_RUN_MANAGER + if _MULTI_RUN_MANAGER is None: + raise RuntimeError("MultiRunManager not initialized. Please call `setup_multi_run_manager` first.") + return _MULTI_RUN_MANAGER + + +def setup_multi_run_manager( + output_dir: Path, max_runs: int, device: torch.device, lora_config: LoRAConfig | None = None +) -> MultiRunManager: + """Initialize the MultiRunManager singleton. + + Args: + output_dir: Directory containing run outputs + max_runs: Maximum number of concurrent runs + device: Device for LoRA tensors + lora_config: Optional trainer LoRA config. If provided, registers validation + and scaling hooks for LoRA rank and alpha. + + Returns: + The initialized MultiRunManager instance. + """ + global _MULTI_RUN_MANAGER + _MULTI_RUN_MANAGER = MultiRunManager(output_dir, max_runs, device) + + # Register validation and scaling hooks for LoRA + if lora_config is not None and _MULTI_RUN_MANAGER.world.is_master: + trainer_lora = lora_config + + def validate_lora_rank(orch_config: "OrchestratorConfig") -> tuple[bool, str]: + # Default to trainer's rank/alpha if not specified + if orch_config.model.lora.rank is None: + orch_config.model.lora.rank = trainer_lora.rank + if orch_config.model.lora.alpha is None: + orch_config.model.lora.alpha = trainer_lora.alpha + if orch_config.model.lora.rank > trainer_lora.rank: + return ( + False, + f"model.lora.rank ({orch_config.model.lora.rank}) exceeds trainer max rank ({trainer_lora.rank})", + ) + return True, "" + + def on_run_discovered(idx: int, run_id: str, orch_config: "OrchestratorConfig") -> None: + _MULTI_RUN_MANAGER.scaling_factors[idx] = orch_config.model.lora.alpha / orch_config.model.lora.rank + + _MULTI_RUN_MANAGER.register_config_validation_hook(validate_lora_rank) + _MULTI_RUN_MANAGER.register_discovered_hook(on_run_discovered) + + return _MULTI_RUN_MANAGER diff --git a/src/prime_rl/trainer/scheduler.py b/src/prime_rl/trainer/scheduler.py index c010737b3b7..4b5d7e4a13a 100644 --- a/src/prime_rl/trainer/scheduler.py +++ b/src/prime_rl/trainer/scheduler.py @@ -1,14 +1,19 @@ -from typing import TYPE_CHECKING +from __future__ import annotations from torch.optim import Optimizer from torch.optim.lr_scheduler import ConstantLR, CosineAnnealingLR, LinearLR, LRScheduler, SequentialLR -from prime_rl.trainer.config import SchedulerConfigType -from prime_rl.trainer.runs import get_runs +from prime_rl.configs.trainer import SchedulerConfig +from prime_rl.trainer.optim import CPUOffloadOptimizer, MultiLoRAOptimizer +from prime_rl.trainer.runs import get_multi_run_manager from prime_rl.utils.logger import get_logger -if TYPE_CHECKING: - from prime_rl.trainer.optim import MultiLoRAOptimizer + +def _get_base_optimizer(optimizer: Optimizer | CPUOffloadOptimizer) -> Optimizer: + """Extract the base optimizer from a potentially wrapped optimizer.""" + if isinstance(optimizer, CPUOffloadOptimizer): + return optimizer.optimizer + return optimizer def setup_constant_scheduler(optimizer: Optimizer) -> LRScheduler: @@ -82,18 +87,23 @@ def setup_cosine_scheduler( def setup_scheduler( - optimizer: Optimizer, - scheduler_config: SchedulerConfigType, + optimizer: Optimizer | CPUOffloadOptimizer, + scheduler_config: SchedulerConfig, max_steps: int | None, lr: float, ) -> LRScheduler: - """Create learning rate scheduler based on config.""" + """Create learning rate scheduler based on config. + + Handles CPUOffloadOptimizer by extracting the base optimizer for the scheduler. + """ + base_optimizer = _get_base_optimizer(optimizer) + match scheduler_config.type: case "constant": - return setup_constant_scheduler(optimizer) + return setup_constant_scheduler(base_optimizer) case "linear": return setup_linear_scheduler( - optimizer, + base_optimizer, max_steps=max_steps, warmup_steps=scheduler_config.warmup_steps, decay_steps=scheduler_config.decay_steps, @@ -102,7 +112,7 @@ def setup_scheduler( ) case "cosine": return setup_cosine_scheduler( - optimizer, + base_optimizer, max_steps=max_steps, warmup_steps=scheduler_config.warmup_steps, lr=lr, @@ -121,22 +131,22 @@ class MultiLoRAScheduler: def __init__( self, - scheduler_config: SchedulerConfigType, + scheduler_config: SchedulerConfig, max_steps: int | None, ): self.scheduler_config = scheduler_config self.max_steps = max_steps - self.runs = get_runs() + self.multi_run_manager = get_multi_run_manager() self.logger = get_logger() - self.schedulers: list[LRScheduler | None] = [None] * self.runs.max_runs + self.schedulers: list[LRScheduler | None] = [None] * self.multi_run_manager.max_runs def scheduler_creation_hook(self, optimizer: Optimizer, idx: int) -> None: """Create a scheduler for a newly created optimizer. This should be called after an optimizer is created for a run. """ - lr = self.runs.config[idx].optim.lr + lr = self.multi_run_manager.config[idx].optim.lr self.schedulers[idx] = setup_scheduler( optimizer, self.scheduler_config, @@ -146,7 +156,7 @@ def scheduler_creation_hook(self, optimizer: Optimizer, idx: int) -> None: def step(self) -> None: """Step all active schedulers.""" - for idx in self.runs.ready_to_update_idxs: + for idx in self.multi_run_manager.ready_to_update_idxs: self.schedulers[idx].step() def get_last_lr(self, idx: int) -> list[float]: @@ -167,8 +177,8 @@ def load_state_dict(self, state_dict: dict) -> None: def setup_multi_scheduler( - optimizer: "MultiLoRAOptimizer", - scheduler_config: SchedulerConfigType, + optimizer: MultiLoRAOptimizer, + scheduler_config: SchedulerConfig, max_steps: int | None, ) -> MultiLoRAScheduler: """Create a MultiLoRAScheduler for managing per-run schedulers.""" diff --git a/src/prime_rl/trainer/sft/config.py b/src/prime_rl/trainer/sft/config.py deleted file mode 100644 index 1f1e7e07b8f..00000000000 --- a/src/prime_rl/trainer/sft/config.py +++ /dev/null @@ -1,254 +0,0 @@ -from pathlib import Path -from typing import Annotated, Literal, TypeAlias - -from pydantic import BaseModel, Field, model_validator - -from prime_rl.trainer.config import ( - AdamWConfig, - BenchConfig, - CheckpointConfig, - ConstantSchedulerConfig, - ModelConfig, - OptimizerConfigType, - SchedulerConfigType, - TokenizerConfig, -) -from prime_rl.utils.config import HeartbeatConfig, LogConfig, WandbConfig -from prime_rl.utils.pydantic_config import BaseConfig, BaseSettings - - -class BaseDataConfig(BaseModel): - """Base config for SFT data.""" - - batch_size: Annotated[int, Field(ge=1)] = 128 - seq_len: Annotated[int, Field(ge=1)] = 128 - pack_function: Literal["cat", "stack"] = "cat" - micro_batch_size: Annotated[int, Field(ge=1)] = 1 - - @model_validator(mode="after") - def validate_batch_size(self): - if self.batch_size % self.micro_batch_size != 0: - raise ValueError("Batch size must be divisible by micro batch size") - if self.batch_size < self.micro_batch_size: - raise ValueError("Batch size must be greater than or equal to micro batch size") - return self - - -class FakeDataConfig(BaseDataConfig): - """Configures fake data used for debugging.""" - - type: Literal["fake"] = "fake" - - length: Literal["fixed", "variable"] = "fixed" - input_ids: Literal["increasing", "random"] = "increasing" - - -class LossMaskConfig(BaseConfig): - """Configures which message types contribute to the loss. If True, the loss_mask will be True and the message type will contribute to the loss.""" - - system: Annotated[bool, Field(description="Whether system messages contribute to the loss.")] = False - user: Annotated[bool, Field(description="Whether user messages contribute to the loss.")] = False - assistant: Annotated[bool, Field(description="Whether assistant messages contribute to the loss.")] = True - tool: Annotated[bool, Field(description="Whether tool messages contribute to the loss.")] = False - - -class SFTDataConfig(BaseDataConfig): - """Configures the data used for training.""" - - type: Literal["sft"] = "sft" - - name: Annotated[str, Field(description="Name or path of the HF dataset to use.")] = ( - "PrimeIntellect/Reverse-Text-SFT" - ) - subsets: Annotated[list[str] | None, Field(description="Subsets to use from the HF dataset.")] = None - splits: Annotated[list[str] | None, Field(description="Splits to use from the HF dataset.")] = None - probabilities: Annotated[list[float] | None, Field(description="Probabilities to use for each subset/split.")] = ( - None - ) - stopping_strategy: Annotated[ - Literal["first_exhausted", "all_exhausted"], - Field(description=""), - ] = "all_exhausted" - shuffle: Annotated[bool, Field(description="Whether to shuffle the dataset at the beginning of each epoch.")] = True - seed: Annotated[ - int, - Field( - description="Random seed to use for shuffling the dataset. We also shuffle at the end of each epoch by adding epoch count to the seed." - ), - ] = 0 - - # Configuring - loss_mask: LossMaskConfig = LossMaskConfig() - - @model_validator(mode="after") - def validate_subsets_and_splits(self): - if self.subsets is not None or self.splits is not None: - if self.subsets is not None and self.splits is not None: - if len(self.subsets) != len(self.splits): - raise ValueError( - "Number of subsets must be equal to number of splits. Please specify which split to load for each subset." - ) - if self.subsets is not None and self.probabilities is not None: - if len(self.probabilities) != len(self.subsets): - raise ValueError( - "Number of probabilities must be equal to number of subsets. Please specify a probability for each subset." - ) - if self.splits is not None and self.probabilities is not None: - if len(self.probabilities) != len(self.splits): - raise ValueError( - "Number of probabilities must be equal to number of splits. Please specify a probability for each split." - ) - return self - - -DataConfigType: TypeAlias = FakeDataConfig | SFTDataConfig - - -class SFTTrainerConfig(BaseSettings): - """Configures the SFT trainer""" - - # The model configuration - model: ModelConfig = ModelConfig() - - # The tokenizer configuration - tokenizer: TokenizerConfig = TokenizerConfig() - - # The data configuration - data: Annotated[DataConfigType, Field(discriminator="type")] = SFTDataConfig() - - # The optimizer configuration - optim: Annotated[OptimizerConfigType, Field(discriminator="type")] = AdamWConfig() - - # The learning rate scheduler configuration - scheduler: Annotated[SchedulerConfigType, Field(discriminator="type")] = ConstantSchedulerConfig() - - # The checkpoint configuration - ckpt: CheckpointConfig | None = None - - # The logging configuration - log: LogConfig = LogConfig() - - # The wandb configuration - wandb: WandbConfig | None = None - - output_dir: Annotated[ - Path, - Field( - description="Directory to write outputs to. Will be populated with checkpoints and logs as subdirectories. Should be set to a persistent directory with enough disk space. This value should be distinct across experiments running on a single node. See the README for more details." - ), - ] = Path("outputs") - - max_steps: Annotated[ - int | None, - Field(description="Maximum number of steps to run training for. If None, will run indefinitely."), - ] = None - - memory_profiler_path: Annotated[Path | None, Field(description="Path to write memory profile to.")] = None - - bench: Annotated[ - BenchConfig | None, - Field( - description="Whether to run in benchmark mode. It will automatically set the maximum number of steps to run to 4 and use fake data.", - ), - ] = None - - trace_path: Annotated[Path | None, Field(description="Path to write pytorch profiler trace to.")] = None - - dist_timeout_seconds: Annotated[ - int, - Field( - description="Timeout in seconds for torch distributed ops. Defaults to 600 seconds.", - ), - ] = 600 - - loss_impl: Annotated[ - Literal["liger", "torch"], Field(description="Implementation of the cross entropy loss function to use.") - ] = "torch" - - heartbeat: Annotated[ - HeartbeatConfig | None, Field(description="The heartbeat config for monitoring training progress.") - ] = None - - @model_validator(mode="after") - def auto_setup_bench(self): - if self.bench is not None: - self.max_steps = 4 # 1 Warmup + 3 Benchmark - if self.ckpt: # Do not checkpoint - self.ckpt = None - return self - - @model_validator(mode="after") - def validate_pack_function(self): - if self.model.cp > 1 and self.data.pack_function != "cat": - raise ValueError("Packing function must be 'cat' when CP is enabled") - return self - - @model_validator(mode="after") - def validate_cp_seq_len(self): - if self.model.cp > 1 and self.data.seq_len % self.model.cp != 0: - raise ValueError("Sequence length must be divisible by CP degree") - return self - - @model_validator(mode="after") - def validate_cp_micro_batch_size(self): - if self.model.cp > 1 and self.data.micro_batch_size != 1: - raise ValueError("Micro batch size must be 1 when CP is enabled") - return self - - @model_validator(mode="after") - def validate_seq_len(self): - if self.data.pack_function == "stack": - if self.data.seq_len % 256 != 0: - raise ValueError("The sequence length must be divisible by 256 when using pack function stack") - return self - - @model_validator(mode="after") - def dont_do_massive_traces(self): - if self.trace_path: - if self.max_steps is None: - raise ValueError("Must specify max_steps when tracing") - if self.max_steps >= 10: - raise ValueError( - "Tracing more than 10 steps is not recommended as your trace will be massive. Remove this line if you really want to trace more steps." - ) - return self - - @model_validator(mode="after") - def validate_lora_adapter_saving(self): - if self.ckpt and self.ckpt.weights and self.ckpt.weights.save_adapter_separately: - lora_enabled = self.model and self.model.lora - if not lora_enabled: - raise ValueError( - "save_adapter_separately=True requires LoRA to be enabled. " - "Set model.lora or disable save_adapter_separately." - ) - return self - - @model_validator(mode="after") - def validate_opt_and_fsdp_offload(self): - if self.optim.type == "muon" and self.model.fsdp_cpu_offload: - raise ValueError("Muon optimizer does not support FSDP CPU offload") - return self - - @model_validator(mode="after") - def auto_setup_tokenizer(self): - if self.tokenizer.name is None: - self.tokenizer.name = self.model.name - if self.tokenizer.trust_remote_code is None: - self.tokenizer.trust_remote_code = self.model.trust_remote_code - return self - - @model_validator(mode="after") - def validate_no_chunked_loss(self): - if self.model.fused_lm_head_chunk_size is not None: - raise ValueError( - "Chunked loss is not supported for SFT training yet, please set `model.fused_lm_head_chunk_size` to None" - ) - return self - - @model_validator(mode="after") - def ep_only_with_custom_impl(self): - if self.model.ep > 1 and self.model.impl not in ("custom", "auto"): - raise ValueError("EP is only supported with the custom implementation or auto mode") - - return self diff --git a/src/prime_rl/trainer/sft/data.py b/src/prime_rl/trainer/sft/data.py index 7372db65401..6aa93cc0e3d 100644 --- a/src/prime_rl/trainer/sft/data.py +++ b/src/prime_rl/trainer/sft/data.py @@ -12,8 +12,15 @@ from torchdata.stateful_dataloader import StatefulDataLoader from transformers.tokenization_utils import PreTrainedTokenizer -from prime_rl.trainer.sft.config import DataConfigType, LossMaskConfig +from prime_rl.configs.sft import DataConfig, LossMaskConfig, SFTDataConfig from prime_rl.trainer.world import get_world +from prime_rl.utils.chat_template import ( + IncrementalTokenizationError, + build_incremental_token_mask, + deserialize_tool_calls, + normalize_messages, + strip_message_content, +) from prime_rl.utils.logger import get_logger STACKING_DATASET_BUCKET_TIMEOUT = 10 @@ -107,7 +114,7 @@ def __iter__(self): class SFTDataset(StatefulIterableDataset): - """A dataset wrapping a HF SFT dataset with prompt + completion format.""" + """A dataset wrapping a HF SFT dataset with prompt/completion or raw messages format.""" def __init__( self, @@ -156,116 +163,61 @@ def _process(self, example: dict) -> dict | None: if self.tokenizer is None: return example - # Assert that the example has a 'prompt' and 'completion' column - if "prompt" not in example or "completion" not in example: - raise ValueError("All examples in the dataset must have a 'prompt' and 'completion' column for SFT") - - def deserialize_tool_calls(messages: list[dict]) -> list[dict]: - """ - Deserialize tool calls in messages, if any are present. Iterates - over all messages in a message list and tries to find - "tool_calls" key. If found, assumes it is a OAI format and has - key "function" with "arguments" key which is stringified. It - will then deserialize the argument so that chat tmeplates like - Qwen3's can be used. - """ - - def deserialize_tool_call(tool_call: dict) -> dict: - return { - **tool_call, - "function": { - **tool_call["function"], - "arguments": json.loads(tool_call["function"]["arguments"]), - }, - } - - return [ - { - **message, - "tool_calls": [deserialize_tool_call(tool_call) for tool_call in message.get("tool_calls") or []], - } - for message in messages - ] - - def strip_content(messages: list[dict]) -> list[dict]: - def _strip_content(message: dict) -> dict: - if isinstance(message.get("content"), str): - return {**message, "content": message["content"].strip()} - return message - - return [_strip_content(message) for message in messages] - - # Deserialize tool call arguments from message list, if present - assumes OAI format - # Reference: https://platform.openai.com/docs/guides/function-calling#handling-function-calls - prompt = deserialize_tool_calls(example["prompt"]) - completion = deserialize_tool_calls(example["completion"]) - - # Strip content from all messages so that incremental tokenization works - # NOTE: This has the side effect that we do never train on leading or trailing whitespace - prompt = strip_content(prompt) - completion = strip_content(completion) + def resolve_messages(example: dict) -> list[dict]: + # `messages` takes precedence over explicit split fields and is interpreted + # as a whole-chat training sample with an empty prompt. + if "messages" in example: + messages = normalize_messages(example["messages"], default_role="assistant") + elif "prompt" in example and "completion" in example: + messages = normalize_messages(example["prompt"], default_role="user") + normalize_messages( + example["completion"], default_role="assistant" + ) + else: + raise ValueError( + "All examples in the dataset must have either a 'messages' column " + "or both 'prompt' and 'completion' columns for SFT" + ) + + # Deserialize tool call arguments from message list, if present - assumes OAI format + # Reference: https://platform.openai.com/docs/guides/function-calling#handling-function-calls + messages = deserialize_tool_calls(messages) + + # Strip content from all messages so that incremental tokenization works + # NOTE: This has the side effect that we do never train on leading or trailing whitespace + return strip_message_content(messages) + + messages = resolve_messages(example) # Parse available tools, if present - assumes OAI format # Reference: https://platform.openai.com/docs/guides/function-calling#function-tool-example tools = json.loads(example.get("tools") or "[]") - def should_mask(message: dict, loss_mask_config: LossMaskConfig) -> bool: + def should_mask(message: dict) -> bool: assert "role" in message, "Message must have a role" match message["role"]: case "user": - return True if loss_mask_config.user else False + return True if self.loss_mask_config.user else False case "assistant": - return True if loss_mask_config.assistant else False + return True if self.loss_mask_config.assistant else False case "system": - return True if loss_mask_config.system else False + return True if self.loss_mask_config.system else False case "tool": - return True if loss_mask_config.tool else False + return True if self.loss_mask_config.tool else False case _: raise ValueError(f"Invalid message role: {message['role']}") - def build_loss_mask(prompt, completion, tokenizer, loss_mask_config: LossMaskConfig) -> list[bool]: - messages = prompt + completion - loss_mask: list[bool] = [] - prev_ids, prev_len = [], 0 - for i, message in enumerate(messages): - assert "role" in message, "Message must have a role" - # Support parallel tool call outputs (treat them as one message for loss mask) - if message["role"] == "tool" and i + 1 < len(messages) and messages[i + 1]["role"] == "tool": - continue - cur_ids = tokenizer.apply_chat_template( - messages[: i + 1], - tools=tools, - # This is to mask out the generation prompt after user and tool messages - # It leads to us not training on <|im_start|>assistant - add_generation_prompt=True - if ( - message["role"] in ["user", "tool"] - and i + 1 < len(messages) - and messages[i + 1]["role"] == "assistant" - ) - else False, - **example.get("chat_template_kwargs", {}), - ) - assert prev_ids == cur_ids[:prev_len], ( - f"Got mismatch in incremental tokenization with chat template at message {i}. Previous ids: {prev_ids} != {cur_ids[:prev_len]=}.\nDecoded prev_ids:\n{tokenizer.decode(prev_ids)}\nDecoded cur_ids:\n{tokenizer.decode(cur_ids[:prev_len])}" - ) - loss_mask.extend([should_mask(message, loss_mask_config)] * (len(cur_ids) - prev_len)) - prev_ids, prev_len = cur_ids, len(cur_ids) - - return loss_mask - - # Build input_ids - input_ids = cast( - list[int], - self.tokenizer.apply_chat_template( - prompt + completion, + try: + input_ids, loss_mask = build_incremental_token_mask( + self.tokenizer, + messages, + role_to_mask=should_mask, tools=tools, - **example.get("chat_template_kwargs", {}), - ), - ) - - # Build loss_mask - loss_mask = build_loss_mask(prompt, completion, self.tokenizer, self.loss_mask_config) + chat_template_kwargs=example.get("chat_template_kwargs", {}), + collapse_consecutive_tool_messages=True, + ) + except IncrementalTokenizationError as e: + self.logger.warning(f"Skipping example {example.get('__index', '')}: {e}") + return None # If EOS token is not found, manually append it if not self.tokenizer.eos_token_id in input_ids: @@ -541,62 +493,73 @@ def setup_and_interleave_datasets( return dataset +def load_sft_dataset(config: SFTDataConfig) -> Dataset: + """Load and interleave the raw HF dataset. This is the expensive I/O step.""" + logger = get_logger() + if config.subsets is None and config.splits is None: + return setup_and_interleave_datasets( + dataset_name=config.name, + subsets_and_splits=[(None, "train")], + probabilities=config.probabilities, + stopping_strategy=config.stopping_strategy, + ) + elif config.subsets is not None and config.splits is None: + logger.debug(f"Loading datasets for subsets {config.subsets} with default split 'train'") + return setup_and_interleave_datasets( + dataset_name=config.name, + subsets_and_splits=[(subset, "train") for subset in config.subsets], + probabilities=config.probabilities, + stopping_strategy=config.stopping_strategy, + ) + elif config.subsets is None and config.splits is not None: + logger.debug(f"Loading datasets for splits {config.splits} with default subset 'None'") + return setup_and_interleave_datasets( + dataset_name=config.name, + subsets_and_splits=[(None, split) for split in config.splits], + probabilities=config.probabilities, + stopping_strategy=config.stopping_strategy, + ) + else: + assert config.subsets is not None and config.splits is not None + logger.debug(f"Loading datasets for subsets {config.subsets} with splits {config.splits}") + return setup_and_interleave_datasets( + dataset_name=config.name, + subsets_and_splits=list(zip(config.subsets, config.splits)), + probabilities=config.probabilities, + stopping_strategy=config.stopping_strategy, + ) + + def setup_dataset( - tokenizer: PreTrainedTokenizer, config: DataConfigType, non_dp_size: int = 1 + tokenizer: PreTrainedTokenizer, + config: DataConfig, + non_dp_size: int = 1, + *, + max_epochs: int | None = None, + raw_dataset: Dataset | None = None, ) -> StatefulIterableDataset: if config.type == "fake": - # Shouldnt matter to handle non_dp_size if dataset is random return FakeDataset( vocab_size=tokenizer.vocab_size, seq_len=config.seq_len, length=config.length, input_ids=config.input_ids ) elif config.type == "sft": - logger = get_logger() - if config.subsets is None and config.splits is None: - dataset = setup_and_interleave_datasets( - dataset_name=config.name, - subsets_and_splits=[(None, "train")], - probabilities=config.probabilities, - stopping_strategy=config.stopping_strategy, - ) - elif config.subsets is not None and config.splits is None: - logger.debug(f"Loading datasets for subsets {config.subsets} with default split 'train'") - dataset = setup_and_interleave_datasets( - dataset_name=config.name, - subsets_and_splits=[(subset, "train") for subset in config.subsets], - probabilities=config.probabilities, - stopping_strategy=config.stopping_strategy, - ) - elif config.subsets is None and config.splits is not None: - logger.debug(f"Loading datasets for splits {config.splits} with default subset 'None'") - dataset = setup_and_interleave_datasets( - dataset_name=config.name, - subsets_and_splits=[(None, split) for split in config.splits], - probabilities=config.probabilities, - stopping_strategy=config.stopping_strategy, - ) - else: - assert config.subsets is not None and config.splits is not None - logger.debug(f"Loading datasets for subsets {config.subsets} with splits {config.splits}") - dataset = setup_and_interleave_datasets( - dataset_name=config.name, - subsets_and_splits=list(zip(config.subsets, config.splits)), - probabilities=config.probabilities, - stopping_strategy=config.stopping_strategy, - ) + if raw_dataset is None: + raw_dataset = load_sft_dataset(config) return SFTDataset( - dataset, + raw_dataset, tokenizer, shuffle=config.shuffle, seed=config.seed, seq_len=config.seq_len, loss_mask_config=config.loss_mask, non_dp_size=non_dp_size, + max_epochs=max_epochs, ) else: raise ValueError(f"Invalid dataset type: {config.type}") -def setup_dataloader(dataset: StatefulIterableDataset, config: DataConfigType) -> StatefulDataLoader: +def setup_dataloader(dataset: StatefulIterableDataset, config: DataConfig) -> StatefulDataLoader: if config.pack_function == "stack": stacking_dataset = StackDataset(dataset, config.seq_len * config.micro_batch_size) return StatefulDataLoader(stacking_dataset, batch_size=1, collate_fn=stack_collate) diff --git a/src/prime_rl/trainer/sft/train.py b/src/prime_rl/trainer/sft/train.py index 2016cfbfee8..032f3b136bb 100644 --- a/src/prime_rl/trainer/sft/train.py +++ b/src/prime_rl/trainer/sft/train.py @@ -1,3 +1,5 @@ +import prime_rl._compat # noqa: F401 — patch ring_flash_attn compat before import + import time from contextlib import nullcontext from datetime import timedelta @@ -8,16 +10,16 @@ # Import environment before any other imports # ruff: noqa: I001 -from prime_rl.trainer.models.layers.attn import substitute_prime_rl_flash_attn +from prime_rl.trainer.models.layers.attn import substitute_ring_attn from prime_rl.utils.act_offloading import maybe_activation_offloading import torch from torch.profiler import profile, ProfilerActivity, record_function -from loguru import logger from prime_rl.trainer.ckpt import setup_ckpt_managers from prime_rl.utils.pathing import resolve_latest_ckpt_step -from prime_rl.trainer.sft.config import SFTTrainerConfig +from prime_rl.configs.sft import SFTConfig from prime_rl.utils.cp import setup_cp_params, shard_for_cp -from prime_rl.trainer.runs import Progress +from prime_rl.trainer.runs import Progress, get_multi_run_manager, setup_multi_run_manager +from prime_rl.trainer.models.layers.lora import set_lora_num_tokens from prime_rl.utils.logger import setup_logger from prime_rl.trainer.optim import setup_optimizer from prime_rl.trainer.scheduler import setup_scheduler @@ -30,10 +32,12 @@ ) from prime_rl.trainer.parallel_dims import get_parallel_dims from prime_rl.trainer.perf import get_perf_counter -from prime_rl.trainer.sft.data import setup_dataloader, setup_dataset +from prime_rl.trainer.sft.data import load_sft_dataset, setup_dataloader, setup_dataset from prime_rl.trainer.utils import ( + GarbageCollection, MemoryProfiler, export_benchmark_json, + get_zero_gradient_ratio, get_ckpt_disk_metrics, print_sample, setup_torch_distributed, @@ -42,22 +46,23 @@ from prime_rl.trainer.world import get_world from prime_rl.utils.heartbeat import Heartbeat from prime_rl.utils.monitor import setup_monitor -from prime_rl.utils.pydantic_config import parse_argv +from prime_rl.utils.config import cli +from prime_rl.utils.process import set_proc_title from prime_rl.utils.utils import clean_exit, to_col_format import torch.distributed as dist from liger_kernel.transformers.cross_entropy import LigerCrossEntropyLoss +from prime_rl.trainer.models.layers.lm_head import FUSED_CE_IGNORE_INDEX from torchtitan.distributed.utils import clip_grad_norm_ @clean_exit -@logger.catch(reraise=True) -def train(config: SFTTrainerConfig): +def train(config: SFTConfig): # Setup world and logger world = get_world() logger = setup_logger( config.log.level, - log_file=config.output_dir / "logs" / "trainer" / f"rank_{world.rank}.log" if config.log.file else None, + json_logging=config.log.json_logging, ) logger.info(f"Starting SFT trainer in {world}") @@ -79,23 +84,41 @@ def train(config: SFTTrainerConfig): setup_torch_distributed( timeout=timedelta(seconds=config.dist_timeout_seconds), enable_gloo=config.model.fsdp_cpu_offload ) - torch.set_float32_matmul_precision("high") + # Configurable to support ROCm/AMD GPUs where reduced precision + # matmul corrupts softmax over large vocabularies. Override via config + # (e.g. matmul_precision = "highest") on ROCm. + torch.set_float32_matmul_precision(config.matmul_precision) + + if config.model.lora is not None: + setup_multi_run_manager(config.output_dir, 1, torch.device("cuda", world.local_rank), config.model.lora) # Initialize parallel dimensions parallel_dims = get_parallel_dims(config.model, config.data.seq_len) - total_micro_batches = config.data.batch_size * config.model.cp * config.model.tp + total_micro_batches = config.data.batch_size * config.model.cp micro_batches_per_step = world.world_size * config.data.micro_batch_size assert total_micro_batches % micro_batches_per_step == 0, ( - f"batch_size * cp * tp ({total_micro_batches}) must be divisible by " + f"batch_size * cp ({total_micro_batches}) must be divisible by " f"world_size * micro_batch_size ({micro_batches_per_step})" ) grad_accum_steps = total_micro_batches // micro_batches_per_step if parallel_dims.cp_enabled: assert config.data.seq_len % parallel_dims.cp == 0, "Sequence length must be divisible by CP degree" - substitute_hf_flash_attn(parallel_dims.world_mesh["cp"].get_group(), heads_k_stride=1) - substitute_prime_rl_flash_attn(parallel_dims.world_mesh["cp"].get_group(), heads_k_stride=1) + cp_group = parallel_dims.world_mesh["cp"].get_group() + cp_rank = parallel_dims.world_mesh["cp"].get_local_rank() + if config.model.cp_style == "ring": + substitute_hf_flash_attn(cp_group, heads_k_stride=1) + substitute_ring_attn(cp_group, heads_k_stride=1, attn_impl=config.model.attn) + else: + from prime_rl.trainer.models.layers.ulysses_attn import ( + substitute_hf_ulysses_attn, + substitute_ulysses_attn, + ) + + substitute_hf_ulysses_attn(cp_group) + substitute_ulysses_attn(cp_group, attn_impl=config.model.attn) + from prime_rl.utils.cp import setup_hybrid_cp, setup_nemotron_h_cp, setup_sparse_mla_cp # Set up checkpoint manager logger.info(f"Initializing checkpoint managers ({config.ckpt})") @@ -111,14 +134,34 @@ def train(config: SFTTrainerConfig): # Initialize the model and tokenizer logger.info(f"Initializing model ({config.model})") loading_from_ckpt_later = config.ckpt and checkpoint_step is not None - model = setup_model(config.model, parallel_dims, loading_from_ckpt_later) + fused_cross_entropy: bool | str = {"liger_fused": "liger", "quack_fused": "quack"}.get(config.loss_impl, False) + model = setup_model(config.model, parallel_dims, loading_from_ckpt_later, fused_cross_entropy=fused_cross_entropy) + + if parallel_dims.cp_enabled: + from prime_rl.utils.cp import assert_cp_style_supports_model + + assert_cp_style_supports_model(config.model.cp_style, model) + # sparse MLA is softmax (works with both ring and ulysses). + setup_sparse_mla_cp(model, cp_group, cp_rank, parallel_dims.cp) + # Linear-attn / Mamba layers are only configured under ulysses; with ring + # we'd have already raised above. + if config.model.cp_style == "ulysses": + setup_hybrid_cp(model, cp_group, cp_rank, parallel_dims.cp) + setup_nemotron_h_cp(model, cp_group, cp_rank, parallel_dims.cp) + + if config.model.lora is not None: + multi_run_manager = get_multi_run_manager() + multi_run_manager.reset_run_parameters(0) + multi_run_manager.scaling_factors[0] = config.model.lora.alpha / config.model.lora.rank logger.info(f"Initializing tokenizer ({config.tokenizer})") tokenizer = setup_tokenizer(config.tokenizer) # Set up the optimizer logger.info(f"Initializing optimizer ({config.optim})") - optimizer = setup_optimizer(config.optim, list(model.named_parameters()), parallel_dims) + optimizer = setup_optimizer( + config.optim, list(model.named_parameters()), parallel_dims, cpu_offload=config.model.optim_cpu_offload + ) # Set up the learning rate scheduler scheduler_steps = ( @@ -132,10 +175,15 @@ def train(config: SFTTrainerConfig): # Set up the dataset and dataloader logger.info(f"Initializing data ({config.data})") - dataset = setup_dataset(tokenizer, config.data, config.model.cp * config.model.tp) + dataset = setup_dataset(tokenizer, config.data, config.model.cp) dataloader = setup_dataloader(dataset, config.data) dataiter = iter(dataloader) + val_raw_dataset = None + if config.val is not None: + logger.info(f"Loading validation dataset ({config.val.data.name})") + val_raw_dataset = load_sft_dataset(config.val.data) + # Optionally, resume training from a checkpoint progress = Progress() @@ -159,27 +207,110 @@ def train(config: SFTTrainerConfig): cp_enabled = parallel_dims.cp_enabled cp_rank = parallel_dims.world_mesh["cp"].get_local_rank() if cp_enabled else 0 cp_group = parallel_dims.world_mesh["cp"].get_group() if cp_enabled else None + dp_cp_group = parallel_dims.get_mesh("dp_cp").get_group() cp_size = parallel_dims.cp + ce_loss = None match config.loss_impl: case "liger": ce_loss = LigerCrossEntropyLoss(reduction="none") case "torch": ce_loss = CrossEntropyLoss(reduction="none") + case "liger_fused" | "quack_fused": + pass # loss is computed inside the fused lm_head case _: raise ValueError(f"Invalid loss implementation: {config.loss_impl}") + def compute_loss(micro_batch: dict) -> tuple[torch.Tensor, torch.Tensor]: + """Forward pass returning (loss_sum, token_count) over unmasked tokens.""" + input_ids = micro_batch["input_ids"].to("cuda") + position_ids = micro_batch["position_ids"].to("cuda") + target_ids = micro_batch["target_ids"].to("cuda") + loss_mask = micro_batch["loss_mask"].to("cuda") + + if cp_enabled: + input_ids, position_ids = setup_cp_params( + input_ids, position_ids, cp_rank, cp_size, cp_group, cp_style=config.model.cp_style + ) + target_ids = shard_for_cp(target_ids, cp_rank=cp_rank, cp_world_size=cp_size) + loss_mask = shard_for_cp(loss_mask, cp_rank=cp_rank, cp_world_size=cp_size) + + if config.model.lora is not None: + set_lora_num_tokens(torch.full((1,), input_ids.numel(), dtype=torch.int32, device="cuda")) + + token_count = loss_mask.sum(dtype=torch.int64) + + with maybe_activation_offloading(config.model.ac_offloading): + if config.loss_impl in ("liger_fused", "quack_fused"): + masked_target_ids = target_ids.clone() + masked_target_ids[~loss_mask] = FUSED_CE_IGNORE_INDEX + out = forward(model, input_ids, position_ids, labels=masked_target_ids) + loss_sum = out["loss"] * token_count + else: + out = forward(model, input_ids, position_ids) + logits = out["logits"] + B, L, V = logits.shape + token_loss = ce_loss(logits.view(-1, V), target_ids.view(-1)).view(B, L) + loss_sum = token_loss[loss_mask].sum() + del logits + + del out + return loss_sum, token_count + + maybe_record_function = nullcontext + + def run_eval_loop(data_iter): + """Validation forward loop. Returns token-weighted global mean loss.""" + total_loss_sum = torch.tensor(0.0, device="cuda") + total_token_count = torch.tensor(0, dtype=torch.int64, device="cuda") + nan_count = torch.tensor(0, device="cuda") + + with torch.no_grad(): + for micro_batch in data_iter: + loss_sum, token_count = compute_loss(micro_batch) + if not torch.isnan(loss_sum.detach()): + total_loss_sum += loss_sum.detach() + total_token_count += token_count + else: + nan_count += 1 + + dist.all_reduce(total_loss_sum, op=dist.ReduceOp.SUM, group=dp_cp_group) + dist.all_reduce(total_token_count, op=dist.ReduceOp.SUM, group=dp_cp_group) + dist.all_reduce(nan_count, op=dist.ReduceOp.SUM) + + mean_loss = (total_loss_sum / total_token_count).item() if total_token_count.item() > 0 else float("nan") + return mean_loss, nan_count.item() + + def run_validation(step: int) -> None: + val_dataset = setup_dataset( + tokenizer, config.val.data, config.model.cp, max_epochs=1, raw_dataset=val_raw_dataset + ) + val_dataloader = setup_dataloader(val_dataset, config.val.data) + + # No train/eval switch: no dropout in these models, and toggling would trigger torch.compile recompilation + mean_loss, nan_count = run_eval_loop(val_dataloader) + if nan_count > 0: + logger.warning(f"Validation at step {step}: {nan_count} batches had NaN loss") + if mean_loss != mean_loss: + logger.warning(f"Validation at step {step} had no valid tokens") + else: + logger.success(f"Validation | Step {step} | Loss: {mean_loss:.4f}") + monitor.log({"val/loss": mean_loss, "step": step}, step=step) + + gc_handler = GarbageCollection(config.gc.interval) if config.gc else None + logger.info(f"Starting training loop (max_steps={config.max_steps or 'infinite'})") max_memory = torch.cuda.mem_get_info()[1] / 1024**3 # GiB is_first_step = True - maybe_record_function = nullcontext if config.trace_path: logger.info(f"Tracing to {config.trace_path}") prof = profile(activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA], record_shapes=True).__enter__() - maybe_record_function = record_function + maybe_record_function = record_function # noqa: F841 – captured by run_forward_loop closure while True: # Reset peak memory stats torch.cuda.reset_peak_memory_stats() + if gc_handler is not None: + gc_handler.run(progress.step) is_last_step = config.max_steps is not None and progress.step == config.max_steps if ( @@ -188,20 +319,23 @@ def train(config: SFTTrainerConfig): and not (is_first_step or is_last_step) and progress.step % config.ckpt.interval == 0 ): - # Save full checkpoint - logger.info(f"Saving checkpoint at step {progress.step}") - save_ckpt_start_time = time.perf_counter() - ckpt_manager.save(progress.step, model, [optimizer], scheduler, progress, dataloader=dataloader) - save_ckpt_time = time.perf_counter() - save_ckpt_start_time + save_ckpt_time = 0 + + if not config.ckpt.weights_only: + # Save full checkpoint + logger.info(f"Saving checkpoint at step {progress.step}") + save_ckpt_start_time = time.perf_counter() + ckpt_manager.save(progress.step, model, [optimizer], scheduler, progress, dataloader=dataloader) + save_ckpt_time += time.perf_counter() - save_ckpt_start_time - # Maybe clean up old checkpoints ckpt_manager.maybe_clean() # Save weight checkpoint if weight_ckpt_manager is not None: logger.info(f"Saving weight checkpoint at step {progress.step}") + save_ckpt_start_time = time.perf_counter() weight_ckpt_manager.save(progress.step, model, tokenizer) - # Maybe clean up old weight checkpoint + save_ckpt_time += time.perf_counter() - save_ckpt_start_time weight_ckpt_manager.maybe_clean() else: save_ckpt_time = 0 @@ -217,61 +351,33 @@ def train(config: SFTTrainerConfig): step_start_time = time.perf_counter() forward_backward_start_time = time.perf_counter() - batch_loss = torch.tensor(0.0).to("cuda") - nan_loss_count = torch.tensor(0).to("cuda") - batch_max_vio, max_vio = torch.tensor(0.0).to("cuda"), None + step_loss_sum = torch.tensor(0.0, device="cuda") + step_local_token_count = torch.tensor(0, dtype=torch.int64, device="cuda") + nan_loss_count = torch.tensor(0, device="cuda") + batch_max_vio = torch.tensor(0.0, device="cuda") for micro_step in range(grad_accum_steps): micro_batch = next(dataiter) - input_ids = micro_batch["input_ids"].to("cuda") - position_ids = micro_batch["position_ids"].to("cuda") - target_ids = micro_batch["target_ids"].to("cuda") - loss_mask = micro_batch["loss_mask"].to("cuda") - - if cp_enabled: - input_ids, position_ids = setup_cp_params(input_ids, position_ids, cp_rank, cp_size, cp_group) - target_ids = shard_for_cp(target_ids, cp_rank=cp_rank, cp_world_size=cp_size) - loss_mask = shard_for_cp(loss_mask, cp_rank=cp_rank, cp_world_size=cp_size) - - assert input_ids.shape == position_ids.shape == target_ids.shape == loss_mask.shape, ( - f"input_ids.shape: {input_ids.shape}, position_ids.shape: {position_ids.shape}, target_ids.shape: {target_ids.shape}, loss_mask.shape: {loss_mask.shape}" - ) if config.log.log_data: - logger.debug("Printing samples of the first micro batch") - print_sample(input_ids.flatten().tolist(), loss_mask.flatten().tolist(), tokenizer) - - # Forward pass - logger.debug("Starting forward pass") - with maybe_record_function("forward"), maybe_activation_offloading(config.model.ac_offloading): - out = forward(model, input_ids, position_ids) + print_sample( + micro_batch["input_ids"].flatten().tolist(), micro_batch["loss_mask"].flatten().tolist(), tokenizer + ) - logits = out["logits"] - B, L, V = logits.shape + with maybe_record_function("forward"): + local_loss_sum, local_token_count = compute_loss(micro_batch) - # Compute loss - loss = ce_loss(logits.view(-1, V), target_ids.view(-1)).view(B, L) + step_local_token_count += local_token_count - # Compute average loss over unmasked tokens - loss = loss[loss_mask].mean() - - # Accumulate average loss over gradient accumulation steps - - current_loss = loss.detach() / grad_accum_steps - - # only add if the loss is not nan - if not torch.isnan(current_loss): - batch_loss += current_loss - else: + if torch.isnan(local_loss_sum.detach()): nan_loss_count += 1 - logger.warning("Loss is nan, not taking into account in the batch loss calculation") - - # Delete logits before backward pass to avoid memory spike - del logits + logger.warning("Local loss is nan, excluding this micro step from backward") + scaled_loss = torch.nan_to_num(local_loss_sum, nan=0.0) / grad_accum_steps + else: + step_loss_sum += local_loss_sum.detach() + scaled_loss = local_loss_sum / grad_accum_steps - # Backward pass - logger.debug("Starting backward pass") with maybe_record_function("backward"): - (loss / grad_accum_steps).backward() + scaled_loss.backward() if is_tt_moe_model(model): max_vio = get_load_balance_stats(model)["max_vio"] @@ -280,18 +386,47 @@ def train(config: SFTTrainerConfig): dist.all_reduce(max_vio, op=dist.ReduceOp.MAX) batch_max_vio += max_vio / grad_accum_steps - # Debug log with *local, micro step* stats - micro_step_message = f"Micro Step {micro_step}/{grad_accum_steps} | Loss: {loss.item():.4f} | Dataloader Step: {dataloader.state_dict()['dataset_state']['dataset']['step']}" - if is_tt_moe_model(model) and max_vio is not None: - micro_step_message += f" | Max Vio: {max_vio.item():.4f}" - logger.debug(micro_step_message) + forward_backward_time = time.perf_counter() - forward_backward_start_time - logger.debug(f"Clipping gradients with max norm {config.optim.max_norm}") - grad_norm = clip_grad_norm_( - model.parameters(), max_norm=config.optim.max_norm, ep_enabled=parallel_dims.ep_enabled - ) - if grad_norm.device.type == "cpu": - grad_norm = grad_norm.to(torch.device("cuda")) + # All-reduce token counts and rescale gradients to get a global token-weighted mean. + # FSDP already divided grads by fsdp_gradient_divide_factor, so we undo that and + # divide by the true global token count instead. + global_step_token_count = step_local_token_count.clone() + dist.all_reduce(global_step_token_count, op=dist.ReduceOp.SUM, group=dp_cp_group) + global_token_count_val = global_step_token_count.item() + + if global_token_count_val > 0: + grad_scale = parallel_dims.fsdp_gradient_divide_factor * grad_accum_steps / global_token_count_val + for param in model.parameters(): + if param.grad is not None: + param.grad.mul_(grad_scale) + + # Run validation after forward-backward (so torch.compile sees training graph first) but before + # optimizer step (so eval_on_start evaluates untrained weights) + if config.val is not None and ( + (is_first_step and config.val.eval_on_start) + or (not is_first_step and progress.step % config.val.interval == 0) + ): + run_validation(progress.step) + + # Compute the global mean loss for logging. + dist.all_reduce(step_loss_sum, op=dist.ReduceOp.SUM, group=dp_cp_group) + dist.all_reduce(nan_loss_count, op=dist.ReduceOp.SUM) + if global_token_count_val > 0: + batch_loss = (step_loss_sum / global_token_count_val).item() + else: + batch_loss = 0.0 + nan_loss_count = nan_loss_count.item() + + grad_norm: torch.Tensor | None = None + if config.optim.max_norm is not None: + logger.debug(f"Clipping gradients with max norm {config.optim.max_norm}") + grad_norm = clip_grad_norm_( + model.parameters(), max_norm=config.optim.max_norm, ep_enabled=parallel_dims.ep_enabled + ) + if grad_norm.device.type == "cpu": + grad_norm = grad_norm.to(torch.device("cuda")) + zero_grad_ratio = get_zero_gradient_ratio(model.parameters(), parallel_dims.dp_replicate) logger.debug("Optimizer step") optimizer.step() @@ -301,20 +436,18 @@ def train(config: SFTTrainerConfig): current_lr = optimizer.param_groups[0]["lr"] scheduler.step() - forward_backward_time = time.perf_counter() - forward_backward_start_time - # Optionally, dump memory snapshot if memory_profiler is not None: memory_profiler.step() - # Synchronize the tensor metrics across all steps and ranks - logger.debug("Synchronizing tensor metrics across all steps and ranks") - dist.all_reduce(batch_loss, op=dist.ReduceOp.AVG) - dist.all_reduce(nan_loss_count, op=dist.ReduceOp.SUM) - - # Compute step metrics - # Divide by CP and TP since those ranks process the same data - num_tokens = config.data.batch_size * config.data.seq_len // (config.model.cp * config.model.tp) + # Compute step metrics. CP shards the same sequences across cp ranks + # (sequence-sharded data parallelism on the seq dim), so the unique + # training tokens per step is dp_size * (batch_per_dp_rank * seq). + # The `dp` mesh excludes cp by construction (parallel_dims.py), mirroring + # the RL trainer's accounting (rl/train.py). + dp_size = parallel_dims.get_mesh("dp").size() + num_local_tokens = config.data.seq_len * (config.data.batch_size // dp_size) + num_tokens = dp_size * num_local_tokens progress.total_tokens += num_tokens progress.total_samples = dataset.step perf_counter = get_perf_counter(model, config.data.seq_len) @@ -325,8 +458,11 @@ def train(config: SFTTrainerConfig): # Log step metrics step_time = time.perf_counter() - step_start_time - step_message = f"Step {progress.step} | Time: {step_time:.2f}s | Loss: {batch_loss.item():.4f} | Grad. Norm: {grad_norm:.4f} | LR: {current_lr:.2e} | Throughput: {throughput:.0f} tokens/s | MFU: {mfu:.1f}% | Peak Mem.: {peak_memory:.1f}/{max_memory:.1f} GiB ({peak_memory / max_memory * 100:.1f}%)" - if is_tt_moe_model(model) and max_vio is not None: + step_message = f"Step {progress.step} | Time: {step_time:.2f}s | Loss: {batch_loss:.4f}" + if grad_norm is not None: + step_message += f" | Grad. Norm: {grad_norm:.4f}" + step_message += f" | LR: {current_lr:.2e} | Throughput: {throughput:.0f} tokens/s | MFU: {mfu:.1f}% | Peak Mem.: {peak_memory:.1f}/{max_memory:.1f} GiB ({peak_memory / max_memory * 100:.1f}%)" + if is_tt_moe_model(model) and batch_max_vio.item() > 0: step_message += f" | Max Vio: {batch_max_vio.item():.4f}" logger.success(step_message) @@ -366,14 +502,16 @@ def train(config: SFTTrainerConfig): # Log optimizer metrics optim_metrics = { "optim/lr": current_lr, - "optim/grad_norm": grad_norm.item(), + "optim/zero_grad_ratio": zero_grad_ratio, "step": progress.step, } + if grad_norm is not None: + optim_metrics["optim/grad_norm"] = grad_norm.item() monitor.log(optim_metrics, step=progress.step) loss_log_metrics = { - "loss/mean": batch_loss.item(), - "loss/nan_count": nan_loss_count.item(), + "loss/mean": batch_loss, + "loss/nan_count": nan_loss_count, "step": progress.step, } # Log tensor stats @@ -393,12 +531,8 @@ def train(config: SFTTrainerConfig): disk_metrics["step"] = progress.step monitor.log(disk_metrics, step=progress.step) - if is_tt_moe_model(model): - max_vio_log_metrics = { - "max_vio/mean": batch_max_vio.item(), - "step": progress.step, - } - monitor.log(max_vio_log_metrics, step=progress.step) + if is_tt_moe_model(model) and batch_max_vio.item() > 0: + monitor.log({"max_vio/mean": batch_max_vio.item(), "step": progress.step}, step=progress.step) is_first_step = False progress.step += 1 @@ -417,8 +551,9 @@ def train(config: SFTTrainerConfig): # Write final checkpoint if ckpt_manager is not None: - logger.info("Writing final checkpoint") - ckpt_manager.save(progress.step, model, [optimizer], scheduler, progress, dataloader=dataloader) + if not (config.ckpt and config.ckpt.weights_only): + logger.info("Writing final checkpoint") + ckpt_manager.save(progress.step, model, [optimizer], scheduler, progress, dataloader=dataloader) ckpt_manager.maybe_clean() # Write final weight checkpoint @@ -440,7 +575,8 @@ def train(config: SFTTrainerConfig): def main(): - train(parse_argv(SFTTrainerConfig)) + set_proc_title("SFTTrainer") + train(cli(SFTConfig)) if __name__ == "__main__": diff --git a/src/prime_rl/trainer/sign_sgd.py b/src/prime_rl/trainer/sign_sgd.py new file mode 100644 index 00000000000..2fe1002e903 --- /dev/null +++ b/src/prime_rl/trainer/sign_sgd.py @@ -0,0 +1,53 @@ +from typing import Callable + +import torch +from torch.optim import Optimizer + + +class SignSGD(Optimizer): + """Sign-based SGD optimizer with minimal memory footprint. + + This optimizer uses the sign of gradients instead of storing momentum and variance, + making it equivalent to AdamW with beta1=0 and beta2=0 (resetting optimizer state each step). + + Mathematical equivalence: + AdamW: W = W - lr * m_t / sqrt(v_t + eps) + With beta1=0, beta2=0: m_t = g_t, v_t = g_t^2 + Simplified: W = W - lr * g_t / sqrt(g_t^2 + eps) + Ignoring eps: W = W - lr * sign(g_t) + """ + + def __init__( + self, + params, + lr: float = 1e-3, + weight_decay: float = 0.01, + ): + if lr < 0.0: + raise ValueError(f"Invalid learning rate: {lr}") + if weight_decay < 0.0: + raise ValueError(f"Invalid weight_decay: {weight_decay}") + + defaults = dict(lr=lr, weight_decay=weight_decay) + super().__init__(params, defaults) + + @torch.no_grad() + def step(self, closure: Callable = None): + loss = None + if closure is not None: + with torch.enable_grad(): + loss = closure() + + for group in self.param_groups: + for p in group["params"]: + if p.grad is None: + continue + + sign_grad = torch.sign(p.grad) + + if group["weight_decay"] > 0.0: + p.add_(p, alpha=-group["lr"] * group["weight_decay"]) + + p.add_(sign_grad, alpha=-group["lr"]) + + return loss diff --git a/src/prime_rl/trainer/utils.py b/src/prime_rl/trainer/utils.py index 8152c72212d..72c4b4999e6 100644 --- a/src/prime_rl/trainer/utils.py +++ b/src/prime_rl/trainer/utils.py @@ -1,8 +1,10 @@ +import gc import json import pickle import shutil import time from collections import defaultdict +from collections.abc import Iterable from datetime import timedelta from pathlib import Path from typing import Any @@ -14,7 +16,8 @@ from rich.console import Console from rich.table import Table from rich.text import Text -from torch import Tensor +from torch import Tensor, nn +from torch.distributed.tensor import DTensor from transformers.tokenization_utils import PreTrainedTokenizer from prime_rl.trainer.world import get_world @@ -25,6 +28,89 @@ DEFAULT_TIMEOUT = timedelta(seconds=600) +class GarbageCollection: + """Controls Python garbage collection to avoid stragglers in distributed training. + + In multi-GPU training, Python's automatic GC can trigger unpredictably on one rank + while others wait at a synchronization point, stalling the entire step. This class + disables automatic GC and runs deterministic collections every `interval` steps so + all ranks collect simultaneously. + + Based on the approach from torchtitan (https://arxiv.org/abs/2505.05713). + """ + + def __init__(self, interval: int = 50): + assert interval > 0, "gc interval must be a positive integer" + self.interval = interval + gc.disable() + self._collect() + + def run(self, step: int): + if step > 0 and step % self.interval == 0: + self._collect() + + def _collect(self, generation: int = 1): + begin = time.monotonic() + gc.collect(generation) + get_logger().info(f"[GC] collection took {time.monotonic() - begin:.2f}s") + + +def _to_local_tensor(tensor: Tensor | DTensor) -> Tensor: + if isinstance(tensor, DTensor): + return tensor.to_local() + return tensor + + +def count_zero_gradient_elements(parameters: Iterable[nn.Parameter]) -> tuple[Tensor, Tensor]: + """Count zero-gradient parameter elements on the local distributed shards. + + Parameters that require gradients but did not receive one in the current step + are counted as fully zero. This makes inactive MoE experts visible in the + metric instead of silently dropping them from the count. + """ + + device = torch.device("cuda", torch.cuda.current_device()) if torch.cuda.is_available() else torch.device("cpu") + num_zeros = torch.zeros((), dtype=torch.long, device=device) + num_tracked = torch.zeros((), dtype=torch.long, device=device) + + for param in parameters: + if not param.requires_grad: + continue + + local_param = _to_local_tensor(param.detach()) + if local_param.numel() == 0: + continue + + if local_param.device != num_zeros.device: + num_zeros = num_zeros.to(local_param.device) + num_tracked = num_tracked.to(local_param.device) + + local_numel = torch.tensor(local_param.numel(), dtype=torch.long, device=local_param.device) + num_tracked += local_numel + + if param.grad is None: + num_zeros += local_numel + continue + + local_grad = _to_local_tensor(param.grad.detach()) + if local_grad.numel() != local_param.numel(): + raise ValueError("Local gradient shape does not match the local parameter shape") + + num_zeros += local_numel - torch.count_nonzero(local_grad) + + return num_zeros, num_tracked + + +def get_zero_gradient_ratio(parameters: Iterable[nn.Parameter], dp_replicate: int = 1) -> float: + num_zero_grad, num_grad_elements = count_zero_gradient_elements(parameters) + dist.all_reduce(num_zero_grad, op=dist.ReduceOp.SUM) + dist.all_reduce(num_grad_elements, op=dist.ReduceOp.SUM) + if dp_replicate > 1: + num_zero_grad = torch.div(num_zero_grad, dp_replicate, rounding_mode="floor") + num_grad_elements = torch.div(num_grad_elements, dp_replicate, rounding_mode="floor") + return (num_zero_grad.float() / num_grad_elements.clamp_min(1).float()).item() + + def get_ckpt_disk_metrics(output_dir: Path) -> dict[str, float]: """ Disk usage metrics for the checkpoint directory (/checkpoints). @@ -45,7 +131,8 @@ def get_ckpt_disk_metrics(output_dir: Path) -> dict[str, float]: def setup_torch_distributed(timeout: timedelta = DEFAULT_TIMEOUT, enable_gloo: bool = False): - torch.cuda.set_device(get_world().local_rank) + device_id = get_world().local_rank + torch.cuda.set_device(device_id) # Use Gloo backend for CPU and NCCL for GPU when CPU offloading is enabled # Otherwise use NCCL for better GPU performance backend = None # by default nccl @@ -53,7 +140,7 @@ def setup_torch_distributed(timeout: timedelta = DEFAULT_TIMEOUT, enable_gloo: b get_logger().info("Using Gloo backend for CPU and NCCL backend for GPU") backend = "cpu:gloo,cuda:nccl" - dist.init_process_group(backend=backend, timeout=timeout) + dist.init_process_group(backend=backend, timeout=timeout, device_id=device_id) def get_response_lengths(position_ids: torch.Tensor) -> list[int]: @@ -295,6 +382,33 @@ def compute_stats(self) -> dict[str, float | int]: return metrics +def filter_rl_trainer_tensor_stats_for_wandb(metrics: dict[str, float | int]) -> dict[str, float | int]: + """Drop noisy per-token distribution keys before sending RL trainer stats to W&B.""" + skip_prefixes = ("trainer_probs/", "inference_probs/") + mean_max_only_prefixes = ( + "is_masked/", + "is_masked_low/", + "is_masked_high/", + "mismatch_kl/", + "masked_mismatch_kl/", + "unmasked_mismatch_kl/", + ) + out: dict[str, float | int] = {} + for k, v in metrics.items(): + if k == "step": + out[k] = v + continue + if any(k.startswith(p) for p in skip_prefixes): + continue + if k.startswith("entropy/") and k != "entropy/mean": + continue + if any(k.startswith(p) for p in mean_max_only_prefixes): + if not (k.endswith("/mean") or k.endswith("/max")): + continue + out[k] = v + return out + + MEMORY_SNAPSHOT_MAX_ENTRIES = 100000 diff --git a/src/prime_rl/trainer/weights.py b/src/prime_rl/trainer/weights.py index c5a356b108c..840197aec84 100644 --- a/src/prime_rl/trainer/weights.py +++ b/src/prime_rl/trainer/weights.py @@ -33,16 +33,30 @@ def _strip_pytorch_wrapper_prefix(key: str) -> str: return key -def get_max_layer_num(state_dict: dict[str, Tensor]) -> int: +def get_max_layer_num(state_dict: dict[str, Tensor], layer_prefix: str = "model.layers.") -> int: """Get the maximum number of layers in the model.""" - return max(int(i.split(".")[2]) for i in state_dict.keys() if "model.layers." in i) + 1 + max_num = -1 + for key in state_dict: + if not key.startswith(layer_prefix): + continue + layer_num_str = key[len(layer_prefix) :].split(".")[0] + if layer_num_str.isdigit(): + max_num = max(max_num, int(layer_num_str)) + return max_num + 1 + + +def load_state_dict_keys(save_dir: Path) -> list[str]: + """Load only the key names from safetensor files without reading tensor data.""" + keys: list[str] = [] + for safetensor_path in save_dir.glob("*.safetensors"): + with safe_open(safetensor_path, framework="pt", device="cpu") as f: + keys.extend(f.keys()) + return keys def load_state_dict(save_dir: Path) -> dict[str, Tensor]: """Load a state dict from a local directory with safetensor files.""" safetensors_paths = list(save_dir.glob("*.safetensors")) - if len(safetensors_paths) > 1: - safetensors_paths.sort(key=lambda x: int(x.stem.split("-")[1].split("of")[0])) state_dict = {} for safetensor_path in safetensors_paths: with safe_open(safetensor_path, framework="pt", device="cpu") as f: diff --git a/src/prime_rl/trainer/world.py b/src/prime_rl/trainer/world.py index 75f4f8a8a8a..679c027caf1 100644 --- a/src/prime_rl/trainer/world.py +++ b/src/prime_rl/trainer/world.py @@ -1,14 +1,14 @@ -from prime_rl.trainer import envs +import os class World: """This class stores topology information for distributed training and inference settings by parsing environment variables set by torchrun.""" def __init__(self): - self.rank = envs.RANK - self.world_size = envs.WORLD_SIZE - self.local_rank = envs.LOCAL_RANK - self.local_world_size = envs.LOCAL_WORLD_SIZE + self.rank = int(os.environ.get("RANK", "0")) + self.world_size = int(os.environ.get("WORLD_SIZE", "1")) + self.local_rank = int(os.environ.get("LOCAL_RANK", "0")) + self.local_world_size = int(os.environ.get("LOCAL_WORLD_SIZE", "1")) self._check_world() self.num_nodes = self.world_size // self.local_world_size diff --git a/src/prime_rl/transport/__init__.py b/src/prime_rl/transport/__init__.py index fe37f1c7794..e4c3153dc73 100644 --- a/src/prime_rl/transport/__init__.py +++ b/src/prime_rl/transport/__init__.py @@ -1,7 +1,7 @@ from pathlib import Path +from prime_rl.configs.shared import TransportConfig from prime_rl.transport.base import MicroBatchReceiver, MicroBatchSender, TrainingBatchReceiver, TrainingBatchSender -from prime_rl.transport.config import TransportConfigType from prime_rl.transport.filesystem import ( FileSystemMicroBatchReceiver, FileSystemMicroBatchSender, @@ -17,7 +17,7 @@ ) -def setup_training_batch_sender(output_dir: Path, transport: TransportConfigType) -> TrainingBatchSender: +def setup_training_batch_sender(output_dir: Path, transport: TransportConfig) -> TrainingBatchSender: if transport.type == "filesystem": return FileSystemTrainingBatchSender(output_dir) elif transport.type == "zmq": @@ -26,7 +26,7 @@ def setup_training_batch_sender(output_dir: Path, transport: TransportConfigType raise ValueError(f"Invalid transport type: {transport.type}") -def setup_training_batch_receiver(transport: TransportConfigType) -> TrainingBatchReceiver: +def setup_training_batch_receiver(transport: TransportConfig) -> TrainingBatchReceiver: if transport.type == "filesystem": return FileSystemTrainingBatchReceiver() elif transport.type == "zmq": @@ -36,7 +36,7 @@ def setup_training_batch_receiver(transport: TransportConfigType) -> TrainingBat def setup_micro_batch_sender( - output_dir: Path, data_world_size: int, current_step: int, transport: TransportConfigType + output_dir: Path, data_world_size: int, current_step: int, transport: TransportConfig ) -> MicroBatchSender: if transport.type == "filesystem": return FileSystemMicroBatchSender(output_dir, data_world_size, current_step) @@ -47,7 +47,7 @@ def setup_micro_batch_sender( def setup_micro_batch_receiver( - output_dir: Path, data_rank: int, current_step: int, transport: TransportConfigType + output_dir: Path, data_rank: int, current_step: int, transport: TransportConfig ) -> MicroBatchReceiver: if transport.type == "filesystem": return FileSystemMicroBatchReceiver(output_dir, data_rank, current_step) diff --git a/src/prime_rl/transport/config.py b/src/prime_rl/transport/config.py deleted file mode 100644 index 2ed2d1a9022..00000000000 --- a/src/prime_rl/transport/config.py +++ /dev/null @@ -1,27 +0,0 @@ -from typing import Annotated, Literal, TypeAlias - -from pydantic import BaseModel, Field - - -class BaseTransportConfig(BaseModel): - """Base configuration for transport.""" - - pass - - -class FileSystemTransportConfig(BaseTransportConfig): - """Configures filesystem-based transport for training examples.""" - - type: Literal["filesystem"] = "filesystem" - - -class ZMQTransportConfig(BaseTransportConfig): - """Configures ZMQ-based transport for training examples.""" - - type: Literal["zmq"] = "zmq" - host: Annotated[str, Field(description="The host address for ZMQ transport.")] = "localhost" - port: Annotated[int, Field(description="The base port for ZMQ transport.")] = 5555 - hwm: Annotated[int, Field(description="High water mark (max messages in queue) for ZMQ sockets.")] = 10 - - -TransportConfigType: TypeAlias = FileSystemTransportConfig | ZMQTransportConfig diff --git a/src/prime_rl/transport/filesystem.py b/src/prime_rl/transport/filesystem.py index e151657c71b..fb9ec68999b 100644 --- a/src/prime_rl/transport/filesystem.py +++ b/src/prime_rl/transport/filesystem.py @@ -1,13 +1,13 @@ from pathlib import Path from time import time -from prime_rl.trainer.runs import get_runs +from prime_rl.trainer.runs import get_multi_run_manager from prime_rl.transport.base import MicroBatchReceiver, MicroBatchSender, TrainingBatchReceiver, TrainingBatchSender from prime_rl.transport.types import MicroBatch, TrainingBatch from prime_rl.utils.pathing import get_rollout_dir, get_step_path, sync_wait_for_path -BATCH_FILE_TMP_NAME = "rollouts.bin.tmp" -BATCH_FILE_NAME = "rollouts.bin" +BATCH_FILE_TMP_NAME = "train_rollouts.bin.tmp" +BATCH_FILE_NAME = "train_rollouts.bin" LOG_FREQ_SECONDS = 10 @@ -35,32 +35,32 @@ class FileSystemTrainingBatchReceiver(TrainingBatchReceiver): def __init__(self) -> None: super().__init__() - self.runs = get_runs() + self.multi_run_manager = get_multi_run_manager() self._last_logged_paths: list[Path] | None = None self._last_logged_time = time() self._waiting_since: float | None = None - # Track received steps per run independently of runs.progress[idx].step + # Track received steps per run independently of multi_run_manager.progress[idx].step # This prevents duplicate reads when trainer step != orchestrator step self._received_steps: dict[int, int] = {} def _get_received_step(self, idx: int) -> int: """Get the next step to receive for a run, initializing from progress if needed.""" if idx not in self._received_steps: - # Initialize from runs.progress on first access (for checkpoint resume) - self._received_steps[idx] = self.runs.progress[idx].step + # Initialize from multi_run_manager.progress on first access (for checkpoint resume) + self._received_steps[idx] = self.multi_run_manager.progress[idx].step return self._received_steps[idx] def _get_batch_path(self, idx: int) -> Path: """Get the batch file path for a specific run at its next step to receive.""" - run_dir = self.runs.get_run_dir(idx) + run_dir = self.multi_run_manager.get_run_dir(idx) rollout_dir = get_rollout_dir(run_dir) step = self._get_received_step(idx) return get_step_path(rollout_dir, step) / BATCH_FILE_NAME def can_receive(self) -> bool: """Check if any run has a batch file available.""" - for idx in self.runs.used_idxs: - if not self.runs.ready_to_update[idx] and self._get_batch_path(idx).exists(): + for idx in self.multi_run_manager.used_idxs: + if not self.multi_run_manager.ready_to_update[idx] and self._get_batch_path(idx).exists(): return True return False @@ -75,7 +75,7 @@ def receive(self) -> list[TrainingBatch]: else: self._waiting_since = self._waiting_since or now - current_paths = [self._get_batch_path(idx) for idx in self.runs.used_idxs] + current_paths = [self._get_batch_path(idx) for idx in self.multi_run_manager.used_idxs] if current_paths != self._last_logged_paths or now - self._last_logged_time > LOG_FREQ_SECONDS: if len(current_paths) == 0: self.logger.debug( @@ -87,8 +87,8 @@ def receive(self) -> list[TrainingBatch]: self.logger.debug(f"Looking for batches in {current_paths}{waiting_suffix}") self._last_logged_paths = current_paths self._last_logged_time = now - for idx in list(self.runs.used_idxs): - if self.runs.ready_to_update[idx]: + for idx in self.multi_run_manager.used_idxs: + if self.multi_run_manager.ready_to_update[idx]: continue batch_path = self._get_batch_path(idx) if batch_path.exists(): @@ -107,7 +107,7 @@ def reset_run(self, idx: int) -> None: """Reset received step tracking for a run index. Called when a run is deleted and a new run takes its place. - The next access to _get_received_step will re-initialize from runs.progress. + The next access to _get_received_step will re-initialize from multi_run_manager.progress. """ if idx in self._received_steps: del self._received_steps[idx] diff --git a/src/prime_rl/transport/types.py b/src/prime_rl/transport/types.py index 52ca56b5b03..4bc594f06d3 100644 --- a/src/prime_rl/transport/types.py +++ b/src/prime_rl/transport/types.py @@ -10,16 +10,29 @@ class TrainingSample(msgspec.Struct, array_like=True, gc=False, omit_defaults=Tr completion_ids: list[int] completion_mask: list[bool] completion_logprobs: list[float] + completion_temperatures: list[float] # Per-token temperatures used during generation teacher_logprobs: list[float] | None = None advantage: float | None = None reward: float | None = None + # Multimodal fields (Qwen3-VL) — pixel_values stored as raw float32 bytes for efficient serialization + pixel_values: bytes | None = None + pixel_values_shape: list[int] | None = None # [num_patches, patch_dim] + # image_grid_thw: grid dimensions [num_images, 3] where each entry is [temporal, height, width] + image_grid_thw: list[list[int]] | None = None + + routed_experts: list[list[list[int]]] | None = None # [seq_len, layers, topk] + + # mm_token_type_ids: token type ids per token [batch seq], int64 (0=text, 1=image, 2=video) + mm_token_type_ids: list[int] | None = None + + sft_loss: bool = False # When True, trainer uses SFT loss instead of RL loss for this sample + class TrainingBatch(msgspec.Struct, array_like=True, gc=False, omit_defaults=True): """A batch of training examples with metadata for transport.""" examples: list[TrainingSample] - temperature: float step: int run_idx: int | None = None @@ -33,6 +46,17 @@ class MicroBatch(msgspec.Struct, array_like=True, gc=False, omit_defaults=True): advantages: list[float] inference_logprobs: list[float] position_ids: list[int] - temperature: float + temperatures: list[float] # Per-token temperatures used during generation teacher_logprobs: list[float] | None = None lora_num_tokens: list[int] | None = None + routed_experts: list[list[list[int]]] | None = None + + # Multimodal fields (Qwen3-VL) — pixel_values stored as raw float32 bytes for efficient serialization + pixel_values: bytes | None = None + pixel_values_shape: list[int] | None = None # [num_patches, patch_dim] + # image_grid_thw: grid dimensions [num_images, 3] where each entry is [temporal, height, width] + image_grid_thw: list[list[int]] | None = None + # mm_token_type_ids: token type ids per token [batch seq], int64 (0=text, 1=image, 2=video) + mm_token_type_ids: list[int] | None = None + + sft_loss: bool = False # When True, trainer uses SFT loss instead of RL loss for this batch diff --git a/src/prime_rl/transport/zmq.py b/src/prime_rl/transport/zmq.py index 7388110edbd..cd5fffbdfe8 100644 --- a/src/prime_rl/transport/zmq.py +++ b/src/prime_rl/transport/zmq.py @@ -3,9 +3,9 @@ import zmq -from prime_rl.trainer.runs import get_runs +from prime_rl.configs.shared import ZMQTransportConfig +from prime_rl.trainer.runs import get_multi_run_manager from prime_rl.transport.base import MicroBatchReceiver, MicroBatchSender, TrainingBatchReceiver, TrainingBatchSender -from prime_rl.transport.config import ZMQTransportConfig from prime_rl.transport.types import MicroBatch, TrainingBatch LOG_FREQ_SECONDS = 10 @@ -57,7 +57,7 @@ class ZMQTrainingBatchReceiver(TrainingBatchReceiver): def __init__(self, transport: ZMQTransportConfig): super().__init__() - self.runs = get_runs() + self.multi_run_manager = get_multi_run_manager() self._last_logged_time = time() self._last_logged_ids: list[str] | None = None self._waiting_since: float | None = None @@ -116,10 +116,10 @@ def receive(self) -> list[TrainingBatch]: # Track how long we've been waiting for any runnable batch. runnable_available = False - for idx in self.runs.used_idxs: - if self.runs.ready_to_update[idx]: + for idx in self.multi_run_manager.used_idxs: + if self.multi_run_manager.ready_to_update[idx]: continue - run_id = self.runs.idx_2_id[idx].encode("utf-8") + run_id = self.multi_run_manager.idx_2_id[idx].encode("utf-8") if self._pending.get(run_id): runnable_available = True break @@ -129,7 +129,7 @@ def receive(self) -> list[TrainingBatch]: else: self._waiting_since = self._waiting_since or now - current_ids = [self.runs.idx_2_id[idx] for idx in self.runs.used_idxs] + current_ids = [self.multi_run_manager.idx_2_id[idx] for idx in self.multi_run_manager.used_idxs] if current_ids != self._last_logged_ids or now - self._last_logged_time > LOG_FREQ_SECONDS: if len(current_ids) == 0: self.logger.debug( @@ -142,9 +142,9 @@ def receive(self) -> list[TrainingBatch]: self._last_logged_ids = current_ids self._last_logged_time = now - for idx in list(self.runs.used_idxs): - run_id = self.runs.idx_2_id[idx].encode("utf-8") - if self.runs.ready_to_update[idx]: + for idx in list(self.multi_run_manager.used_idxs): + run_id = self.multi_run_manager.idx_2_id[idx].encode("utf-8") + if self.multi_run_manager.ready_to_update[idx]: continue per_id_batches = self._pending.get(run_id) diff --git a/src/prime_rl/utils/act_offloading.py b/src/prime_rl/utils/act_offloading.py index b4e422e682f..322655b349d 100644 --- a/src/prime_rl/utils/act_offloading.py +++ b/src/prime_rl/utils/act_offloading.py @@ -6,7 +6,7 @@ import torch from torch.autograd.graph import saved_tensors_hooks -from prime_rl.trainer.config import ActivationOffloadingConfig +from prime_rl.configs.trainer import ActivationOffloadingConfig from .logger import get_logger diff --git a/src/prime_rl/utils/async_utils.py b/src/prime_rl/utils/async_utils.py new file mode 100644 index 00000000000..1de097592f4 --- /dev/null +++ b/src/prime_rl/utils/async_utils.py @@ -0,0 +1,15 @@ +import asyncio + + +async def safe_cancel(task: asyncio.Task) -> None: + """Safely cancels and awaits an asyncio.Task.""" + task.cancel() + try: + await task + except BaseException: + pass + + +async def safe_cancel_all(tasks: list[asyncio.Task]) -> None: + """Safely cancels and awaits all asyncio.Tasks.""" + await asyncio.gather(*[safe_cancel(task) for task in tasks]) diff --git a/src/prime_rl/utils/background_async.py b/src/prime_rl/utils/background_async.py new file mode 100644 index 00000000000..8aaa752437f --- /dev/null +++ b/src/prime_rl/utils/background_async.py @@ -0,0 +1,97 @@ +"""Background async runner for fire-and-forget HTTP work. + +Owns an event loop running on a background thread plus an httpx.AsyncClient, +so callers can submit coroutines without blocking the main thread. +Fork-safe via ``os.register_at_fork``. + +Used by UsageReporter; PrimeMonitor still has its own copy of the same +pattern and can be migrated to this helper in a follow-up. +""" + +import asyncio +import os +from threading import Thread +from typing import Awaitable, Callable + +import httpx + +from prime_rl.utils.logger import get_logger + + +class BackgroundAsync: + """Owns a background event loop + httpx client for fire-and-forget work.""" + + def __init__(self, *, client_timeout: float = 30.0): + self.logger = get_logger() + self._client_timeout = client_timeout + self._closed = False + self._init_runtime() + os.register_at_fork(after_in_child=self._reinit_after_fork) + + @property + def client(self) -> httpx.AsyncClient: + return self._client + + def submit(self, coro_factory: Callable[[], Awaitable[None]]) -> None: + """Schedule a coroutine on the background loop. Fire-and-forget. + + Takes a factory rather than a coroutine so the coroutine is created + on the background thread (avoids "coroutine was never awaited" warnings + if the loop is in the middle of restarting). + """ + if self._closed: + return + future = asyncio.run_coroutine_threadsafe(coro_factory(), self._loop) + self._pending_futures.append(future) + # Drop completed futures to avoid unbounded growth + self._pending_futures = [f for f in self._pending_futures if not f.done()] + + def close(self, *, flush_timeout: float = 15.0) -> None: + if self._closed or not hasattr(self, "_loop"): + return + self._closed = True + self._flush(timeout=flush_timeout) + + async def _close_client() -> None: + await self._client.aclose() + + try: + future = asyncio.run_coroutine_threadsafe(_close_client(), self._loop) + future.result(timeout=5.0) + except Exception as e: + self.logger.debug("Error closing background async client: %s", e) + + self._loop.call_soon_threadsafe(self._loop.stop) + self._thread.join(timeout=5.0) + + def __del__(self) -> None: + self.close() + + # -- internals -- + + def _init_runtime(self) -> None: + self._loop: asyncio.AbstractEventLoop = asyncio.new_event_loop() + self._thread = Thread(target=self._run_loop, daemon=True) + self._thread.start() + self._client = httpx.AsyncClient(timeout=self._client_timeout) + self._pending_futures: list[asyncio.Future] = [] + + def _reinit_after_fork(self) -> None: + # Background thread doesn't survive fork — recreate everything. + self._closed = False + self._init_runtime() + + def _run_loop(self) -> None: + asyncio.set_event_loop(self._loop) + self._loop.run_forever() + + def _flush(self, timeout: float) -> None: + if not self._pending_futures: + return + self.logger.debug("Flushing %d pending background request(s)", len(self._pending_futures)) + for future in self._pending_futures: + try: + future.result(timeout=timeout) + except Exception as e: + self.logger.debug("Pending background request completed with error: %s", e) + self._pending_futures.clear() diff --git a/src/prime_rl/utils/chat_template.py b/src/prime_rl/utils/chat_template.py new file mode 100644 index 00000000000..a30c5e726cd --- /dev/null +++ b/src/prime_rl/utils/chat_template.py @@ -0,0 +1,150 @@ +import json +from typing import Any, Callable + +from transformers.tokenization_utils import PreTrainedTokenizer + + +class IncrementalTokenizationError(ValueError): + """Raised when incremental tokenization produces inconsistent token prefixes.""" + + pass + + +def common_prefix_len(a: list[int], b: list[int]) -> int: + max_len = min(len(a), len(b)) + for idx in range(max_len): + if a[idx] != b[idx]: + return idx + return max_len + + +def normalize_messages(messages: Any, default_role: str) -> list[dict[str, Any]]: + if messages is None: + return [] + if isinstance(messages, str): + return [{"role": default_role, "content": messages}] + if isinstance(messages, dict): + return [dict(messages)] + if isinstance(messages, list): + normalized: list[dict[str, Any]] = [] + for message in messages: + if isinstance(message, str): + normalized.append({"role": default_role, "content": message}) + elif isinstance(message, dict): + normalized.append(dict(message)) + else: + raise TypeError(f"Unsupported message type: {type(message)}") + return normalized + raise TypeError(f"Unsupported messages container type: {type(messages)}") + + +def deserialize_tool_calls(messages: list[dict[str, Any]]) -> list[dict[str, Any]]: + def _deserialize_tool_call(tool_call: dict[str, Any]) -> dict[str, Any]: + function = tool_call.get("function", {}) + arguments = function.get("arguments") + if isinstance(arguments, str): + arguments = json.loads(arguments) + return { + **tool_call, + "function": {**function, "arguments": arguments}, + } + + deserialized_messages: list[dict[str, Any]] = [] + for message in messages: + if "tool_calls" not in message: + deserialized_messages.append(dict(message)) + continue + + tool_calls = message.get("tool_calls") or [] + deserialized_messages.append( + { + **message, + "tool_calls": [_deserialize_tool_call(tc) for tc in tool_calls], + } + ) + + return deserialized_messages + + +def strip_message_content(messages: list[dict[str, Any]]) -> list[dict[str, Any]]: + def _strip(message: dict[str, Any]) -> dict[str, Any]: + content = message.get("content") + if isinstance(content, str): + return {**message, "content": content.strip()} + return message + + return [_strip(message) for message in messages] + + +def should_add_generation_prompt(messages: list[dict[str, Any]], idx: int) -> bool: + role = messages[idx].get("role") + if role not in ("user", "tool"): + return False + if idx + 1 >= len(messages): + return False + return messages[idx + 1].get("role") == "assistant" + + +def render_messages( + tokenizer: PreTrainedTokenizer, + messages: list[dict[str, Any]], + *, + tools: list[dict[str, Any]] | None = None, + chat_template_kwargs: dict[str, Any] | None = None, + add_generation_prompt: bool = False, + processor=None, +) -> list[int]: + kwargs = dict(chat_template_kwargs or {}) + kwargs["add_generation_prompt"] = add_generation_prompt + if tools is not None: + kwargs["tools"] = tools + if processor is not None: + kwargs["tokenize"] = True + kwargs["return_dict"] = True + result = processor.apply_chat_template(messages, **kwargs) + return list(result["input_ids"][0]) + kwargs["return_dict"] = False + return list(tokenizer.apply_chat_template(messages, **kwargs)) + + +def build_incremental_token_mask( + tokenizer: PreTrainedTokenizer, + messages: list[dict[str, Any]], + *, + role_to_mask: Callable[[dict[str, Any]], bool], + tools: list[dict[str, Any]] | None = None, + chat_template_kwargs: dict[str, Any] | None = None, + collapse_consecutive_tool_messages: bool = False, + processor=None, +) -> tuple[list[int], list[bool]]: + token_mask: list[bool] = [] + prev_ids: list[int] = [] + prev_len = 0 + + for idx, message in enumerate(messages): + role = message.get("role") + if collapse_consecutive_tool_messages and role == "tool" and idx + 1 < len(messages): + if messages[idx + 1].get("role") == "tool": + continue + + cur_ids = render_messages( + tokenizer, + messages[: idx + 1], + tools=tools, + chat_template_kwargs=chat_template_kwargs, + add_generation_prompt=should_add_generation_prompt(messages, idx), + processor=processor, + ) + + if prev_ids != cur_ids[:prev_len]: + raise IncrementalTokenizationError( + f"Mismatch in incremental tokenization with chat template at message {idx} (role={role}). " + "This usually means the chat template is not stable under incremental application. " + "The sample will be skipped." + ) + + token_mask.extend([role_to_mask(message)] * (len(cur_ids) - prev_len)) + prev_ids = cur_ids + prev_len = len(cur_ids) + + return prev_ids, token_mask diff --git a/src/prime_rl/utils/client.py b/src/prime_rl/utils/client.py index 754854e0150..21659dfc46b 100644 --- a/src/prime_rl/utils/client.py +++ b/src/prime_rl/utils/client.py @@ -1,46 +1,224 @@ +from __future__ import annotations + import asyncio import os +from itertools import cycle from pathlib import Path +from typing import Protocol, runtime_checkable import httpx +import verifiers as vf from httpx import AsyncClient -from openai import AsyncOpenAI, NotFoundError -from prime_evals import AsyncEvalsClient -from tenacity import retry, retry_if_exception, stop_after_attempt, wait_exponential +from openai import NotFoundError +from tenacity import retry, retry_if_exception, stop_after_attempt, stop_after_delay, wait_exponential -from prime_rl.utils.config import ClientConfig +from prime_rl.configs.shared import ClientConfig from prime_rl.utils.logger import get_logger -def setup_clients(client_config: ClientConfig) -> list[AsyncOpenAI]: - def _setup_client(base_url: str) -> AsyncOpenAI: - # We use a longer request timeout than default, but if more than 20min, we probably need faster inference deployment - timeout = httpx.Timeout(client_config.timeout) - # We use as many concurrent connections as possible, but lower than available ports - limits = httpx.Limits( - max_connections=8192, # OAI default: 1000 - max_keepalive_connections=8192, # OAI default: 100 +@runtime_checkable +class InferencePool(Protocol): + """Protocol for inference pools (static or elastic).""" + + @property + def train_clients(self) -> list[vf.ClientConfig]: + """Get inference clients.""" + ... + + @property + def admin_clients(self) -> list[AsyncClient]: + """Get admin clients.""" + ... + + def update_model_name(self, model_name: str) -> None: + """Update the model name.""" + ... + + async def get_eval_client(self) -> vf.ClientConfig: + """Get next eval client in round-robin fashion.""" + ... + + async def wait_for_ready(self, model_name: str, timeout: int | None = None) -> None: + """Wait for inference pool to be ready.""" + ... + + async def update_weights(self, weight_dir: Path | None, lora_name: str | None = None, step: int = 0) -> None: + """Update weights on all inference servers.""" + ... + + def get_metrics(self) -> dict[str, float]: + """Get pool metrics.""" + ... + + async def stop(self) -> None: + """Stop the inference pool.""" + ... + + +class StaticInferencePool: + """Static inference pool with fixed client list.""" + + def __init__( + self, + client_config: ClientConfig, + model_name: str, + train_client_type: str = "openai_chat_completions_token", + eval_client_type: str = "openai_chat_completions", + renderer_name: str = "auto", + tool_parser: str | None = None, + reasoning_parser: str | None = None, + renderer_pool_size: int | None = None, + ): + renderer_model_name = model_name if train_client_type == "renderer" else None + self._train_clients = setup_clients( + client_config, + client_type=train_client_type, + renderer_name=renderer_name, + renderer_model_name=renderer_model_name, + tool_parser=tool_parser, + reasoning_parser=reasoning_parser, + renderer_pool_size=renderer_pool_size, ) - http_client = httpx.AsyncClient(limits=limits, timeout=timeout, headers=client_config.headers) - return AsyncOpenAI( - base_url=base_url, - api_key=os.getenv(client_config.api_key_var, "EMPTY"), - max_retries=10, # OAI default: 2 (does exponential backoff and reasonable timeout in between retries) - http_client=http_client, + self._eval_clients = setup_clients(client_config, client_type=eval_client_type) + self._admin_clients = setup_admin_clients(client_config) + self._skip_model_check = client_config.skip_model_check + self._wait_for_ready_timeout = client_config.wait_for_ready_timeout + self._eval_cycle = cycle(self._eval_clients) + self.model_name = model_name + + @property + def train_clients(self) -> list[vf.ClientConfig]: + return self._train_clients + + @property + def admin_clients(self) -> list[AsyncClient]: + return self._admin_clients + + def update_model_name(self, model_name: str) -> None: + self.model_name = model_name + + @property + def eval_clients(self) -> list[vf.ClientConfig]: + return self._eval_clients + + async def get_eval_client(self) -> vf.ClientConfig: + return next(self._eval_cycle) + + async def wait_for_ready(self, model_name: str, timeout: int | None = None) -> None: + await check_health( + self._admin_clients, timeout=timeout if timeout is not None else self._wait_for_ready_timeout ) + await maybe_check_has_model(self._admin_clients, model_name, skip_model_check=self._skip_model_check) - return [_setup_client(base_url) for base_url in client_config.base_url] + async def update_weights(self, weight_dir: Path | None, lora_name: str | None = None, step: int = 0) -> None: + await update_weights(self._admin_clients, weight_dir, lora_name=lora_name, step=step) + def get_metrics(self) -> dict[str, float]: + return {} -def setup_evals_client() -> AsyncEvalsClient: - return AsyncEvalsClient() + async def stop(self) -> None: + pass + + +async def setup_inference_pool( + client_config: ClientConfig, + model_name: str, + train_client_type: str = "openai_chat_completions_token", + eval_client_type: str = "openai_chat_completions", + renderer_name: str = "auto", + tool_parser: str | None = None, + reasoning_parser: str | None = None, + renderer_pool_size: int | None = None, +) -> InferencePool: + """Create an inference pool from config (static or elastic).""" + logger = get_logger() + + if train_client_type == "openai_chat_completions_token": + logger.warning( + "Token-in-token-out (TITO) client is enabled for training. Only use " + "this if your environment has a linear history and the chat " + "template has the extension property." + ) + + if client_config.is_elastic: + from prime_rl.utils.elastic import ElasticInferencePool + + return await ElasticInferencePool.from_config( + client_config, + model_name=model_name, + train_client_type=train_client_type, + eval_client_type=eval_client_type, + renderer_name=renderer_name, + tool_parser=tool_parser, + reasoning_parser=reasoning_parser, + renderer_pool_size=renderer_pool_size, + ) + + logger.info( + f"Initializing static inference pool (base_url={', '.join(client_config.base_url)}, " + f"dp_rank_count={client_config.dp_rank_count}, " + f"api_key_var={client_config.api_key_var}, headers={client_config.headers})" + ) + return StaticInferencePool( + client_config, + model_name=model_name, + train_client_type=train_client_type, + eval_client_type=eval_client_type, + renderer_name=renderer_name, + tool_parser=tool_parser, + reasoning_parser=reasoning_parser, + renderer_pool_size=renderer_pool_size, + ) + + +def setup_clients( + client_config: ClientConfig, + client_type: str = "openai_chat_completions", + renderer_name: str = "auto", + renderer_model_name: str | None = None, + tool_parser: str | None = None, + reasoning_parser: str | None = None, + renderer_pool_size: int | None = None, +) -> list[vf.ClientConfig]: + clients = [] + client_idx = 0 + for base_url in client_config.base_url: + for dp_rank in range(client_config.dp_rank_count): + headers = client_config.headers.copy() + if client_config.dp_rank_count > 1: + headers["X-data-parallel-rank"] = str(dp_rank) + clients.append( + vf.ClientConfig( + client_idx=client_idx, + client_type=client_type, + renderer=renderer_name, + renderer_model_name=renderer_model_name, + renderer_pool_size=renderer_pool_size, + tool_parser=tool_parser, + reasoning_parser=reasoning_parser, + api_base_url=base_url, + api_key_var=client_config.api_key_var, + timeout=client_config.timeout, + connect_timeout=client_config.connect_timeout, + max_connections=8192, + max_keepalive_connections=8192, + max_retries=10, + extra_headers=headers, + extra_headers_from_state=client_config.extra_headers_from_state, + ) + ) + client_idx += 1 + return clients def setup_admin_clients(client_config: ClientConfig) -> list[AsyncClient]: - """Create a dedicated admin client for weight update operations. + """Create dedicated admin clients for weight update operations. Uses a separate connection pool to avoid queueing behind streaming requests. + When admin_base_url is set, uses those URLs instead of base_url, allowing + weight updates to bypass routers in disaggregated P/D deployments. """ + urls = client_config.admin_base_url if client_config.admin_base_url else client_config.base_url def _setup_admin_client(base_url: str) -> httpx.AsyncClient: headers = client_config.headers.copy() # avoid mutating config @@ -54,21 +232,25 @@ def _setup_admin_client(base_url: str) -> httpx.AsyncClient: return AsyncClient( base_url=base_url, headers=headers, - limits=httpx.Limits(max_connections=1, max_keepalive_connections=0), - timeout=httpx.Timeout(client_config.timeout), + limits=httpx.Limits(max_connections=4, max_keepalive_connections=1), + timeout=httpx.Timeout(None), ) - return [_setup_admin_client(base_url) for base_url in client_config.base_url] + return [_setup_admin_client(base_url) for base_url in urls] -async def check_has_model(clients: list[AsyncOpenAI], model_name: str) -> None: +async def maybe_check_has_model( + admin_clients: list[AsyncClient], model_name: str, skip_model_check: bool = False +) -> None: + if skip_model_check: + return logger = get_logger() logger.debug(f"Checking if model {model_name} is in the inference pool") - results = await asyncio.gather(*[client.models.list() for client in clients]) - for client, result in zip(clients, results): - models = result.data - if not any(model.id == model_name for model in models): - raise ValueError(f"Model {model_name} was not found in the inference pool on {client.base_url}") + results = await asyncio.gather(*[admin_client.get("/v1/models") for admin_client in admin_clients]) + for admin_client, result in zip(admin_clients, results): + models = result.json()["data"] + if not any(model["id"] == model_name for model in models): + raise ValueError(f"Model {model_name} was not found in the inference pool on {admin_client.base_url}") logger.debug(f"Model {model_name} was found in the inference pool") @@ -105,16 +287,45 @@ async def _check_health(admin_client: AsyncClient) -> None: NCCL_READY_MARKER = "NCCL_READY" +async def _pause_engines(admin_clients: list[AsyncClient]) -> None: + """Pause all inference engines, waiting for in-flight requests to drain.""" + logger = get_logger() + logger.info("Pausing inference engines for weight update") + + async def _pause(client: AsyncClient) -> None: + response = await client.post("/pause", params={"mode": "keep", "clear_cache": "false"}) + response.raise_for_status() + + await asyncio.gather(*[_pause(client) for client in admin_clients]) + logger.info("All inference engines paused") + + +async def _resume_engines(admin_clients: list[AsyncClient]) -> None: + """Resume all inference engines after weight update.""" + logger = get_logger() + + async def _resume(client: AsyncClient) -> None: + response = await client.post("/resume") + response.raise_for_status() + + await asyncio.gather(*[_resume(client) for client in admin_clients]) + logger.info("All inference engines resumed") + + async def update_weights( admin_clients: list[AsyncClient], weight_dir: Path | None, lora_name: str | None = None, + step: int = 0, ) -> None: - """Make a HTTP post request to the vLLM server to update the weights. + """Update weights on static inference servers. - Creates a NCCL_READY marker file before calling the update endpoint to signal - to the trainer that inference workers are about to enter the receive path. - This marker is only used in NCCL broadcast mode but is harmless in filesystem mode. + Pauses all engines first to drain in-flight requests, then performs the + weight update, then resumes. This ensures all DP workers are idle and can + participate in the collective weight transfer. + + Note: The server-side /update_weights endpoint automatically resets the prefix cache + to invalidate any cached KV states computed with the old weights. """ logger = get_logger() @@ -125,41 +336,23 @@ async def update_weights( else: async def _update_weights(admin_client: AsyncClient, weight_dir: str | None) -> None: - try: - response = await admin_client.post("/update_weights", json={"weight_dir": weight_dir}) - response.raise_for_status() - except httpx.HTTPStatusError as e: - if e.response.status_code == 404: - logger.warning("The route /update_weights does not exist. Skipping weight update.") - return - raise - - # Create ready marker before servers enter receive path (used by NCCL broadcast) - if weight_dir is not None: - nccl_ready_file = weight_dir / NCCL_READY_MARKER - nccl_ready_file.parent.mkdir(parents=True, exist_ok=True) - nccl_ready_file.touch() - logger.debug(f"Created NCCL_READY marker at {nccl_ready_file}") - - await asyncio.gather(*[_update_weights(admin_client, weight_dir_posix) for admin_client in admin_clients]) - - -async def reload_weights(admin_clients: list[AsyncClient]) -> None: - """Make a HTTP post request to the vLLM server to reload weights (reset to base model).""" - logger = get_logger() + response = await admin_client.post("/update_weights", json={"weight_dir": weight_dir}) + response.raise_for_status() + + # Pause engines so all DP workers drain in-flight work and can join the NCCL broadcast + await _pause_engines(admin_clients) - async def _reload_weights(admin_client: AsyncClient) -> None: - logger.debug("Sending request to reload weights (reset to base model)") try: - response = await admin_client.post("/reload_weights", json={}) - response.raise_for_status() - except httpx.HTTPStatusError as e: - if e.response.status_code == 404: - logger.warning("The route /reload_weights does not exist. Skipping weight reload.") - return - raise + # Create ready marker before servers enter receive path (used by NCCL broadcast) + if weight_dir is not None: + nccl_ready_file = weight_dir / NCCL_READY_MARKER + nccl_ready_file.parent.mkdir(parents=True, exist_ok=True) + nccl_ready_file.touch() + logger.debug(f"Created NCCL_READY marker at {nccl_ready_file}") - await asyncio.gather(*[_reload_weights(admin_client) for admin_client in admin_clients]) + await asyncio.gather(*[_update_weights(admin_client, weight_dir_posix) for admin_client in admin_clients]) + finally: + await _resume_engines(admin_clients) def _is_retryable_lora_error(exception: BaseException) -> bool: @@ -167,12 +360,31 @@ def _is_retryable_lora_error(exception: BaseException) -> bool: if isinstance(exception, httpx.HTTPStatusError): # Retry on 404 (adapter not found) or 500 (server error during loading) return exception.response.status_code in (404, 500) + # Retry on transport-level failures (timeouts, connection resets, etc.) so + # the per-call read timeout below turns a stuck server into a bounded retry + # loop instead of propagating as a hard failure on the first hiccup. + if isinstance(exception, (httpx.TimeoutException, httpx.TransportError)): + return True return False +# Per-attempt and total bounds for `/load_lora_adapter`. A LoRA load is fast +# (small adapter file + KV cache reset, single-digit seconds in practice) but +# the global admin AsyncClient uses `timeout=None`, so a stuck server hangs +# the orchestrator forever inside `ElasticInferencePool._sync_server_adapter`. +# `_PER_ATTEMPT` converts a hang into a TimeoutException so tenacity retries; +# `_TOTAL` is the wall-clock budget across all retries — pick whichever +# stop condition fires first. +LORA_LOAD_READ_TIMEOUT_S = 30.0 +LORA_LOAD_TOTAL_TIMEOUT_S = 120.0 + + async def load_lora_adapter(admin_clients: list[AsyncClient], lora_name: str, lora_path: Path) -> None: """Make a HTTP post request to the vLLM server to load a LoRA adapter. + Uses our wrapper endpoint that also resets the prefix cache to invalidate + KV states computed with old weights. + Retries with exponential backoff if the adapter files are not found, which can happen due to NFS propagation delays. """ @@ -181,15 +393,16 @@ async def load_lora_adapter(admin_clients: list[AsyncClient], lora_name: str, lo @retry( retry=retry_if_exception(_is_retryable_lora_error), - stop=stop_after_attempt(10), - wait=wait_exponential(multiplier=1, min=1, max=30), + stop=stop_after_delay(LORA_LOAD_TOTAL_TIMEOUT_S) | stop_after_attempt(10), + wait=wait_exponential(multiplier=1, min=1, max=10), reraise=True, ) async def _load_lora_adapter(admin_client: AsyncClient) -> None: logger.debug(f"Sending request to load LoRA adapter {lora_name} from {lora_path}") response = await admin_client.post( - "/v1/load_lora_adapter", + "/load_lora_adapter", json={"lora_name": lora_name, "lora_path": lora_path_posix}, + timeout=httpx.Timeout(connect=10.0, read=LORA_LOAD_READ_TIMEOUT_S, write=60.0, pool=10.0), ) response.raise_for_status() @@ -209,22 +422,46 @@ async def _unload_lora_adapter(admin_client: AsyncClient) -> None: await asyncio.gather(*[_unload_lora_adapter(admin_client) for admin_client in admin_clients]) -async def init_nccl_broadcast(admin_clients: list[AsyncClient], host: str, port: int, timeout: int) -> None: - """Make a HTTP post request to the vLLM server to initialize the NCCL broadcast.""" +async def init_nccl_broadcast( + admin_clients: list[AsyncClient], + host: str, + port: int, + timeout: int, + inference_world_size: int | None = None, + quantize_in_weight_transfer: bool = False, +) -> None: + """Initialize NCCL broadcast on all inference servers. + + Each admin client represents one vLLM server. The function computes + per-server rank_offset and gpus_per_server so that every inference GPU + gets a unique rank in the NCCL broadcast group. + """ logger = get_logger() - async def _init_nccl_broadcast( - admin_client: AsyncClient, host: str, port: int, client_num: int, timeout: int - ) -> None: + if inference_world_size is None: + inference_world_size = len(admin_clients) + logger.warning( + f"inference_world_size not provided, defaulting to {inference_world_size} (one GPU per admin client)" + ) + + gpus_per_server = inference_world_size // len(admin_clients) + + logger.info( + f"Initializing NCCL broadcast: {len(admin_clients)} servers, " + f"inference_world_size={inference_world_size}, gpus_per_server={gpus_per_server}" + ) + + async def _init_nccl_broadcast(admin_client: AsyncClient, rank_offset: int) -> None: try: response = await admin_client.post( "/init_broadcaster", json={ "host": host, "port": port, - "server_rank": client_num, - "num_inference_server": len(admin_clients), + "rank_offset": rank_offset, + "inference_world_size": inference_world_size, "timeout": timeout, + "quantize_in_weight_transfer": quantize_in_weight_transfer, }, ) response.raise_for_status() @@ -235,7 +472,7 @@ async def _init_nccl_broadcast( await asyncio.gather( *[ - _init_nccl_broadcast(admin_client, host, port, client_num, timeout) + _init_nccl_broadcast(admin_client, client_num * gpus_per_server) for client_num, admin_client in enumerate(admin_clients) ] ) diff --git a/src/prime_rl/utils/config.py b/src/prime_rl/utils/config.py deleted file mode 100644 index ed3c4665197..00000000000 --- a/src/prime_rl/utils/config.py +++ /dev/null @@ -1,200 +0,0 @@ -from typing import Annotated, Literal - -from pydantic import Field - -from prime_rl.utils.pydantic_config import BaseConfig - - -class ModelConfig(BaseConfig): - """Configures the model.""" - - name: Annotated[str, Field(description="Name or path of the HF model to use.")] = "Qwen/Qwen3-0.6B" - - trust_remote_code: Annotated[ - bool, - Field( - description="Whether to trust remote code for tokenizer initialization.", - ), - ] = False - - -ServerType = Literal["vllm", "openai"] - - -class ClientConfig(BaseConfig): - """Configures the OAI client.""" - - timeout: Annotated[ - int, - Field( - description="Timeout in seconds. By default, it is set to 1200 seconds.", - ), - ] = 1200 - - base_url: Annotated[ - list[str], - Field( - description="Base URLs to use for the OpenAI API. By default, it is set to a single server on localhost at port 8000 which matches the default local vLLM server configuration. If you specify more than one URL, the client will round-robin (chat) completion requests across all servers.", - ), - ] = ["http://localhost:8000/v1"] - - api_key_var: Annotated[ - str, - Field( - description="Name of environment variable containing the API key to use for the OpenAI API. Will parse using `os.getenv(client_config.api_key_var)`. Can be set to an arbitrary string if the inference server is not protected by an API key. If multiple URLs are specified, the same API key will be used for all servers.", - ), - ] = "OPENAI_API_KEY" - - headers: Annotated[ - dict[str, str], - Field( - description="Headers to use for the OpenAI API. By default, it is set to an empty dictionary.", - ), - ] = {} - - -class LogConfig(BaseConfig): - """Configures the logger.""" - - level: Annotated[ - str, - Field(description="Logging level for the process. Will determine the logging verbosity and format."), - ] = "info" - - vf_level: Annotated[ - str, - Field(description="Logging level for the verifiers package. Will determine the logging verbosity and format."), - ] = "warn" - - file: Annotated[ - bool, - Field( - description="Whether to log to a file. If True, will log to a file in the output directory.", - ), - ] = True - - env_worker_logs: Annotated[ - bool, - Field( - description="Whether env workers log to files. If True, workers write to logs/env_workers/{env_name}.log.", - ), - ] = False - - log_data: Annotated[ - bool, - Field( - description="Whether to log the first data sample to the logger.", - ), - ] = False - - -class LogExtrasConfig(BaseConfig): - """Configures extra logging for monitoring platforms.""" - - samples: Annotated[ - bool, - Field( - description="Whether to log prompt/response samples.", - ), - ] = True - - distributions: Annotated[ - bool, - Field( - description="Whether to log distributions (like rewards, advantages, etc.).", - ), - ] = True - - interval: Annotated[ - int, - Field( - ge=1, - description="Step interval at which to log extras.", - ), - ] = 10 - - -class WandbConfig(BaseConfig): - """Configures logging to Weights and Biases.""" - - # Shared configs (May be overwritten by WandbConfig from `rl.py`) - project: Annotated[str, Field(description="The W&B project to log to.")] = "prime-rl" - - name: Annotated[ - str | None, - Field( - description="The W&B name to to use for logging.", - ), - ] = None - - offline: Annotated[bool, Field(description="Whether to run W&B in offline mode.")] = False - - # Individual configs (can only be specified on trainer or orchestrator) - id: Annotated[ - str | None, - Field( - description="The W&B run ID to log to. If None, a random ID will be generated. If you want to resume a run, you can set the ID to the run ID you want to resume.", - ), - ] = None - - -class WandbWithExtrasConfig(WandbConfig): - """Configures logging to Weights and Biases with extras.""" - - log_extras: Annotated[ - LogExtrasConfig | None, - Field( - description="Configuration for logging extras. If None, no extras are logged.", - ), - ] = LogExtrasConfig() - - -class PrimeMonitorConfig(BaseConfig): - """Configures logging to Prime Intellect API.""" - - base_url: Annotated[ - str, - Field( - description="The base URL for Prime Intellect monitoring API.", - ), - ] = "https://api.primeintellect.ai/api/internal/rft" - - api_key_var: Annotated[ - str, - Field( - description="Name of environment variable containing the API key for Prime Intellect API. Will parse using `os.getenv(config.api_key_var)`.", - ), - ] = "PRIME_API_KEY" - - log_extras: Annotated[ - LogExtrasConfig | None, - Field( - description="Configuration for logging extras. If None, no extras are logged.", - ), - ] = LogExtrasConfig() - - -class HeartbeatConfig(BaseConfig): - """Configures the heartbeat for BetterStack.""" - - url: Annotated[str, Field(description="The URL to send the heartbeat to.")] - - -class MetricsServerConfig(BaseConfig): - """Configures the Prometheus metrics server for trainer observability.""" - - port: Annotated[ - int, - Field( - ge=1, - le=65535, - description="Port to expose metrics and health endpoints. Defaults to 8000.", - ), - ] = 8000 - - host: Annotated[ - str, - Field( - description="Host to bind the server to. Defaults to 0.0.0.0.", - ), - ] = "0.0.0.0" diff --git a/src/prime_rl/utils/cp.py b/src/prime_rl/utils/cp.py index 601b8865f2f..2ac9d45ad73 100644 --- a/src/prime_rl/utils/cp.py +++ b/src/prime_rl/utils/cp.py @@ -1,7 +1,129 @@ +from __future__ import annotations + +# ruff: noqa: I001 — `prime_rl._compat` must run before `ring_flash_attn` imports below. +import prime_rl._compat # noqa: F401 + +from typing import Literal + import torch import torch.distributed as dist +import torch.distributed.nn as dist_nn +import torch.nn as nn from ring_flash_attn import update_ring_flash_attn_params +from prime_rl.utils.sequence import get_cu_seqlens_from_position_ids + +CPStyle = Literal["ring", "ulysses"] + + +def _has_linear_attn_layer(model: nn.Module) -> bool: + """True if the model contains any non-softmax (linear/SSM) attention layer.""" + inner = getattr(model, "model", model) + if hasattr(inner, "language_model"): + inner = inner.language_model + layers = getattr(inner, "layers", None) + if layers is None: + return False + for layer in layers: + # Qwen3.5 hybrid DeltaNet + if getattr(layer, "layer_type", None) == "linear_attention": + return True + # NemotronH Mamba + if hasattr(layer, "mamba"): + return True + return False + + +def assert_cp_style_supports_model(cp_style: CPStyle, model: nn.Module) -> None: + """Refuse `cp_style='ring'` on models that have linear/SSM attention layers. + + Ring CP is a softmax-attention algorithm (sequence ring all-gather of K/V). + For non-softmax layers (DeltaNet, Mamba) we'd need a fundamentally different + CP scheme, which is not implemented. Use `cp_style='ulysses'` for those: + ulysses' all-to-all is purely on Q/K/V tensors, so the linear/SSM kernel + runs unchanged on a sequence shard. + """ + if cp_style == "ring" and _has_linear_attn_layer(model): + raise ValueError( + "cp_style='ring' is not supported for models with linear-attention " + "or Mamba/SSM layers (e.g. Qwen3.5 hybrid, NemotronH). Use " + "cp_style='ulysses' instead — its all-to-all on Q/K/V works " + "out-of-the-box with non-softmax kernels." + ) + + +def setup_hybrid_cp(model: nn.Module, cp_group: dist.ProcessGroup, cp_rank: int, cp_world_size: int) -> None: + """Configure DeltaNet modules in Qwen3.5 hybrid models for ulysses-style CP.""" + layers = None + if hasattr(model, "model"): + inner = model.model + if hasattr(inner, "language_model"): + inner = inner.language_model + if hasattr(inner, "layers"): + layers = inner.layers + + if layers is None: + return + + count = 0 + for layer in layers: + if getattr(layer, "layer_type", None) == "linear_attention": + attn = getattr(layer, "linear_attn", None) + if attn is not None: + attn.cp_group = cp_group + attn.cp_rank = cp_rank + attn.cp_world_size = cp_world_size + count += 1 + + if count > 0: + from prime_rl.utils.logger import get_logger + + get_logger().info(f"Configured hybrid CP on {count} DeltaNet modules (fla native state passing)") + + +def setup_nemotron_h_cp(model: nn.Module, cp_group: dist.ProcessGroup, cp_rank: int, cp_world_size: int) -> None: + """Configure NemotronH Mamba layers for ulysses-style all-to-all head partitioning.""" + layers = None + if hasattr(model, "model") and hasattr(model.model, "layers"): + layers = model.model.layers + + if layers is None: + return + + count = 0 + for layer in layers: + if hasattr(layer, "mamba") and hasattr(layer, "set_context_parallel_attributes"): + layer.set_context_parallel_attributes(cp_group, cp_rank, cp_world_size) + count += 1 + + if count > 0: + from prime_rl.utils.logger import get_logger + + get_logger().info(f"Configured NemotronH CP on {count} Mamba layers (all-to-all head partitioning)") + + +def setup_sparse_mla_cp(model: nn.Module, cp_group: dist.ProcessGroup, cp_rank: int, cp_world_size: int) -> None: + """Configure GLM-5 sparse MLA modules for context-parallel gather/scatter.""" + + count = 0 + if not hasattr(model, "model"): + return + + if not hasattr(model.model, "layers"): + return + + for layer in model.model.layers: + if not hasattr(layer, "set_context_parallel_attributes"): + continue + + layer.set_context_parallel_attributes(cp_group, cp_rank, cp_world_size) + count += 1 + + if count > 0: + from prime_rl.utils.logger import get_logger + + get_logger().info(f"Configured sparse MLA CP on {count} DSA layers") + def shard_for_cp(t: torch.Tensor, cp_rank: int, cp_world_size: int) -> torch.Tensor: """ @@ -21,6 +143,18 @@ def shard_for_cp(t: torch.Tensor, cp_rank: int, cp_world_size: int) -> torch.Ten return chunked_t[cp_rank] +def gather_for_cp(t: torch.Tensor, cp_group: dist.ProcessGroup) -> torch.Tensor: + gathered_t = dist_nn.all_gather(t, group=cp_group) + + return torch.cat(gathered_t, dim=1) + + +def gather_for_cp_wo_grad(t: torch.Tensor, cp_world_size: int, cp_group: dist.ProcessGroup) -> torch.Tensor: + empty_like_t = [torch.empty_like(t) for _ in range(cp_world_size)] + dist.all_gather(empty_like_t, t, group=cp_group) + return torch.cat(empty_like_t, dim=1) + + def get_padding_logit_from_prev_cp_rank( logits: torch.Tensor, cp_rank: int, cp_world_size: int, cp_group: dist.ProcessGroup ) -> torch.Tensor | None: @@ -49,43 +183,38 @@ def get_padding_logit_from_prev_cp_rank( return None -def _get_cu_seqlens_for_cp(position_ids: torch.Tensor) -> torch.Tensor: - flat_position_ids = position_ids.view(-1) - seqlens = torch.cat( - [ - flat_position_ids[0:1], - flat_position_ids[:-1][(flat_position_ids == 0)[1:]] + 1, - flat_position_ids[-1:] + 1, - ] - ) - cu_seqlens = seqlens.cumsum(dim=0, dtype=torch.int32) - return cu_seqlens - - def setup_cp_params( input_ids: torch.Tensor, position_ids: torch.Tensor, cp_rank: int, cp_world_size: int, cp_group: dist.ProcessGroup, + cp_style: CPStyle = "ring", ) -> tuple[torch.Tensor, torch.Tensor]: """ - Prepare the input for context parallelism and sets required parameters for ring flash attention. - Args: - input_ids: The input ids tensor. - position_ids: The position ids tensor. - cp_rank: The rank of the current process. - cp_world_size: The number of processes in the context parallel group. - cp_group: The context parallel group. - Returns: - The sharded input_ids and position_ids for context parallelism. + Prepare the input for context parallelism and set required attention params. + + Both ring and ulysses styles need cu_seqlens computed from the *full* + (un-sharded) position_ids, then publish them to the patched attention layer: + - ring: via ring_flash_attn's DATA_PARAMS (with local_k_slice). + - ulysses: via ULYSSES_PARAMS (just the full cu_seqlens / max_seqlen). + + Returns the sequence-sharded input_ids and position_ids — the rest of the + model still runs sequence-sharded; only attention sees the full sequence. """ - input_ids = shard_for_cp(input_ids, cp_rank=cp_rank, cp_world_size=cp_world_size) + cu_seqlens, max_seqlen = get_cu_seqlens_from_position_ids(position_ids) + + if cp_style == "ring": + update_ring_flash_attn_params(cu_seqlens, cp_group) + elif cp_style == "ulysses": + # Delayed import: ulysses_attn lives under trainer.models, which imports + # back into prime_rl.utils — top-level import would deadlock at startup. + from prime_rl.trainer.models.layers.ulysses_attn import update_ulysses_params - cu_seqlens = _get_cu_seqlens_for_cp(position_ids) - update_ring_flash_attn_params(cu_seqlens, cp_group) + update_ulysses_params(cu_seqlens, max_seqlen) + else: + raise ValueError(f"Unknown cp_style: {cp_style}") + + input_ids = shard_for_cp(input_ids, cp_rank=cp_rank, cp_world_size=cp_world_size) position_ids = shard_for_cp(position_ids, cp_rank=cp_rank, cp_world_size=cp_world_size) - return ( - input_ids, - position_ids, - ) + return input_ids, position_ids diff --git a/src/prime_rl/utils/elastic.py b/src/prime_rl/utils/elastic.py new file mode 100644 index 00000000000..902f873903d --- /dev/null +++ b/src/prime_rl/utils/elastic.py @@ -0,0 +1,512 @@ +""" +Elastic inference pool with DNS-based service discovery. + +Discovers inference servers via DNS (any hostname that resolves to multiple IPs), +tracks which servers have the correct LoRA adapter loaded, and +only exposes ready servers to workers. +""" + +from __future__ import annotations + +import asyncio +import socket +import time +from dataclasses import dataclass +from pathlib import Path +from typing import Literal + +import httpx +import verifiers as vf +from httpx import AsyncClient + +from prime_rl.configs.shared import ClientConfig +from prime_rl.utils.client import load_lora_adapter, setup_admin_clients, setup_clients +from prime_rl.utils.logger import get_logger + +# --- Shared discovery functions --- + + +def discover_server_ips(hostname: str) -> list[str]: + """Discover server IPs via DNS lookup.""" + try: + _, _, ips = socket.gethostbyname_ex(hostname) + return sorted(ips) + except socket.gaierror: + return [] + + +async def check_server_model(url: str, model_name: str, timeout: float = 5.0) -> tuple[bool, bool]: + """Check if a server has a specific model loaded.""" + logger = get_logger() + try: + async with httpx.AsyncClient(timeout=timeout) as client: + response = await client.get(f"{url}/v1/models") + response.raise_for_status() + data = response.json() + models = [m.get("id") for m in data.get("data", [])] + return model_name in models, len(models) > 0 + except Exception as e: + logger.debug(f"Failed to check server {url}: {e}") + return False, False + + +async def discover_ready_servers(hostname: str, port: int, model_name: str) -> list[str]: + """Discover servers via DNS that have the requested model loaded.""" + loop = asyncio.get_event_loop() + ips = await loop.run_in_executor(None, discover_server_ips, hostname) + if not ips: + return [] + + checks = [check_server_model(f"http://{ip}:{port}", model_name) for ip in ips] + results = await asyncio.gather(*checks, return_exceptions=True) + + with_model = set() + for ip, result in zip(ips, results): + if isinstance(result, BaseException): + continue + has_model, _ = result + if has_model: + with_model.add(f"http://{ip}:{port}/v1") + + return sorted(with_model) + + +@dataclass +class AdapterState: + """State of a LoRA adapter (loaded or desired).""" + + name: str | None = None + path: Path | None = None + step: int = 0 + + +ServerStatus = Literal["discovering", "syncing", "ready", "unhealthy"] + + +@dataclass +class ServerState: + """State of an individual inference server.""" + + ip: str + url: str + status: ServerStatus = "discovering" + loaded_adapter: AdapterState | None = None + sync_failures: int = 0 + + +class ElasticInferencePool: + """Manages inference servers with DNS-based discovery and adapter sync. + + Discovers servers via DNS, tracks adapter state, syncs LoRA adapters. + """ + + def __init__( + self, + client_config: ClientConfig, + model_name: str, + train_client_type: str = "openai_chat_completions_token", + eval_client_type: str = "openai_chat_completions", + renderer_name: str = "auto", + tool_parser: str | None = None, + reasoning_parser: str | None = None, + renderer_pool_size: int | None = None, + ): + self.logger = get_logger() + self.client_config = client_config + self.model_name = model_name + self.base_model_name = model_name # Keep original for health checks + self.renderer_model_name = model_name if train_client_type == "renderer" else None + self.hostname = client_config.elastic.hostname + self.port = client_config.elastic.port + self.sync_interval = client_config.elastic.sync_interval + self.train_client_type = train_client_type + self.eval_client_type = eval_client_type + self.renderer_name = renderer_name + self.tool_parser = tool_parser + self.reasoning_parser = reasoning_parser + self.renderer_pool_size = renderer_pool_size + self.router_url = client_config.router_url + + self._servers: dict[str, ServerState] = {} + self._admin_clients: dict[str, AsyncClient] = {} + self._lock = asyncio.Lock() + self._desired: AdapterState = AdapterState() + + self._train_clients: list[vf.ClientConfig] = [] + self._eval_clients: list[vf.ClientConfig] = [] + self._client_urls: list[str] = [] + + self._eval_index = 0 + + self._sync_task: asyncio.Task | None = None + self._started = False + + @classmethod + async def from_config( + cls, + client_config: ClientConfig, + model_name: str, + train_client_type: str = "openai_chat_completions_token", + eval_client_type: str = "openai_chat_completions", + renderer_name: str = "auto", + tool_parser: str | None = None, + reasoning_parser: str | None = None, + renderer_pool_size: int | None = None, + ) -> ElasticInferencePool: + if client_config.elastic is None: + raise ValueError("Elastic inference pool requires elastic config") + pool = cls( + client_config, + model_name=model_name, + train_client_type=train_client_type, + eval_client_type=eval_client_type, + renderer_name=renderer_name, + tool_parser=tool_parser, + reasoning_parser=reasoning_parser, + renderer_pool_size=renderer_pool_size, + ) + await pool.start() + return pool + + def update_model_name(self, model_name: str) -> None: + self.model_name = model_name + + def _build_url(self, ip: str) -> str: + return f"http://{ip}:{self.port}" + + def _build_inference_url(self, ip: str) -> str: + return f"http://{ip}:{self.port}/v1" + + @property + def ready_urls(self) -> list[str]: + return [self._build_inference_url(ip) for ip, s in self._servers.items() if s.status == "ready"] + + def _rebuild_clients(self) -> None: + """Rebuild inference clients when the set of ready URLs changes.""" + # When a router URL is configured, route inference requests through it + # instead of directly to discovered pods. Admin operations still use + # individual pod IPs via admin_clients. + # Only expose the router when backends are actually ready — otherwise + # the orchestrator would send requests into a router with no healthy backends. + if self.router_url and self.ready_urls: + urls = [self.router_url] + else: + urls = self.ready_urls + if set(urls) != set(self._client_urls): + self._client_urls = urls + + self._eval_index = 0 + url_config = ClientConfig( + timeout=self.client_config.timeout, + connect_timeout=self.client_config.connect_timeout, + base_url=urls, + api_key_var=self.client_config.api_key_var, + headers=self.client_config.headers, + dp_rank_count=self.client_config.dp_rank_count, + extra_headers_from_state=self.client_config.extra_headers_from_state, + ) + self._train_clients = ( + setup_clients( + url_config, + client_type=self.train_client_type, + renderer_name=self.renderer_name, + renderer_model_name=self.renderer_model_name, + tool_parser=self.tool_parser, + reasoning_parser=self.reasoning_parser, + renderer_pool_size=self.renderer_pool_size, + ) + if urls + else [] + ) + self._eval_clients = setup_clients(url_config, client_type=self.eval_client_type) if urls else [] + + @property + def train_clients(self) -> list[vf.ClientConfig]: + self._rebuild_clients() + return self._train_clients + + @property + def eval_clients(self) -> list[vf.ClientConfig]: + self._rebuild_clients() + return self._eval_clients + + async def get_eval_client(self) -> vf.ClientConfig: + """Get next eval client in round-robin fashion.""" + while not self.eval_clients: + await asyncio.sleep(self.sync_interval) + client = self._eval_clients[self._eval_index % len(self._eval_clients)] + self._eval_index += 1 + return client + + @property + def admin_clients(self) -> list[AsyncClient]: + return list(self._admin_clients.values()) + + @property + def num_servers(self) -> int: + return len(self._servers) + + @property + def num_ready_servers(self) -> int: + return sum(1 for s in self._servers.values() if s.status == "ready") + + async def _create_admin_client(self, ip: str) -> AsyncClient: + url = self._build_url(ip) + config = ClientConfig( + timeout=self.client_config.timeout, + base_url=[f"{url}/v1"], + api_key_var=self.client_config.api_key_var, + headers=self.client_config.headers, + ) + return setup_admin_clients(config)[0] + + async def _get_loaded_adapter(self, ip: str) -> AdapterState | None: + if ip not in self._admin_clients: + return None + + try: + admin = self._admin_clients[ip] + response = await admin.get("/v1/models") + response.raise_for_status() + data = response.json() + + for model in data.get("data", []): + parent = model.get("parent") + model_id = model.get("id", "") + + if self._desired.name: + if model_id != self._desired.name: + continue + elif not parent: + continue + elif parent != self.model_name: + continue + + root = model.get("root", "") + path = Path(root) + try: + step_part = path.name + if step_part.startswith("step_"): + step = int(step_part.split("_")[1]) + elif step_part.startswith("step-"): + step = int(step_part.split("-")[1]) + else: + step = 0 + except (ValueError, IndexError): + step = 0 + return AdapterState(name=model_id, path=path, step=step) + + self.logger.debug(f"No matching adapter found on {ip} for desired={self._desired.name}") + return None + except Exception as e: + self.logger.warning(f"Failed to query /v1/models on {ip}: {e}") + return None + + def _adapter_matches_desired(self, loaded: AdapterState | None) -> bool: + if self._desired.path is None: + return True + if loaded is None: + return False + if loaded.path == self._desired.path: + return True + if self._desired.step > 0 and loaded.step == self._desired.step: + return True + return False + + async def _sync_server_adapter(self, ip: str) -> bool: + server = self._servers.get(ip) + if not server: + return False + + loaded = await self._get_loaded_adapter(ip) + server.loaded_adapter = loaded + + if self._adapter_matches_desired(loaded): + server.status = "ready" + return True + + # Debug: log why pre-check failed (before attempting load) + self.logger.debug( + f"Pre-check failed on {ip}: loaded={loaded.path if loaded else None} " + f"(step={loaded.step if loaded else None}), desired={self._desired.path} (step={self._desired.step})" + ) + server.status = "syncing" + + if self._desired.name and self._desired.path: + try: + self.logger.debug(f"Loading adapter {self._desired.name} on {ip}") + await load_lora_adapter([self._admin_clients[ip]], self._desired.name, self._desired.path) + except Exception as e: + server.status = "unhealthy" + server.sync_failures += 1 + self.logger.error(f"Failed to sync server {ip}: {e}") + return False + + loaded = await self._get_loaded_adapter(ip) + server.loaded_adapter = loaded + + if self._adapter_matches_desired(loaded): + server.status = "ready" + server.sync_failures = 0 + self.logger.debug(f"Successfully synced server {ip}") + return True + + # Debug: log why adapter didn't match after loading + self.logger.warning( + f"Adapter mismatch on {ip} after loading: " + f"loaded={loaded.path if loaded else None} (step={loaded.step if loaded else None}), " + f"desired={self._desired.path} (step={self._desired.step})" + ) + server.status = "unhealthy" + server.sync_failures += 1 + return False + + async def _check_server_health(self, admin_client: AsyncClient, ip: str) -> bool: + try: + response = await admin_client.get("/health") + response.raise_for_status() + except Exception as e: + self.logger.debug(f"Server {ip} health check failed: {e}") + return False + + try: + response = await admin_client.get("/v1/models") + response.raise_for_status() + data = response.json() + models = [m.get("id") for m in data.get("data", [])] + + if self.base_model_name not in models: + self.logger.debug(f"Server {ip} does not have base model {self.base_model_name}") + return False + except Exception as e: + self.logger.debug(f"Server {ip} model check failed: {e}") + return False + + return True + + async def _add_server(self, ip: str) -> bool: + try: + admin_client = await self._create_admin_client(ip) + except Exception as e: + self.logger.debug(f"Failed to create admin client for {ip}: {e}") + return False + + if not await self._check_server_health(admin_client, ip): + await admin_client.aclose() + return False + + self.logger.debug(f"Discovered new inference server: {ip}") + self._admin_clients[ip] = admin_client + self._servers[ip] = ServerState(ip=ip, url=self._build_url(ip), status="discovering") + await self._sync_server_adapter(ip) + return True + + async def _remove_server(self, ip: str) -> None: + self.logger.debug(f"Inference server removed: {ip}") + self._servers.pop(ip, None) + if ip in self._admin_clients: + await self._admin_clients.pop(ip).aclose() + + async def sync(self) -> tuple[int, int]: + async with self._lock: + discovered_ips = set( + await asyncio.get_event_loop().run_in_executor(None, discover_server_ips, self.hostname) + ) + known_ips = set(self._servers.keys()) + + added = 0 + removed = 0 + + for ip in discovered_ips - known_ips: + if await self._add_server(ip): + added += 1 + + for ip in known_ips - discovered_ips: + await self._remove_server(ip) + removed += 1 + + for ip in list(self._servers.keys()): + if ip not in self._admin_clients: + continue + if not await self._check_server_health(self._admin_clients[ip], ip): + self.logger.debug(f"Server {ip} failed health check, removing") + await self._remove_server(ip) + removed += 1 + elif self._servers[ip].status != "ready": + await self._sync_server_adapter(ip) + + return added, removed + + async def _sync_loop(self) -> None: + while True: + try: + added, removed = await self.sync() + if added > 0 or removed > 0: + self.logger.debug( + f"Elastic pool sync: +{added} -{removed} servers " + f"(total: {self.num_servers}, ready: {self.num_ready_servers})" + ) + except Exception as e: + self.logger.error(f"Error in elastic sync loop: {e}") + await asyncio.sleep(self.sync_interval) + + async def start(self) -> None: + if self._started: + return + + self.logger.debug(f"Starting elastic inference pool (hostname={self.hostname})") + await self.sync() + self.logger.debug(f"Initial discovery: {self.num_servers} server(s), {self.num_ready_servers} ready") + self._sync_task = asyncio.create_task(self._sync_loop()) + self._started = True + + async def stop(self) -> None: + if self._sync_task: + self._sync_task.cancel() + try: + await self._sync_task + except asyncio.CancelledError: + pass + + for ip in list(self._servers.keys()): + await self._remove_server(ip) + + self._train_clients = [] + self._eval_clients = [] + self._client_urls = [] + self._started = False + + async def sync_weights(self, weights_path: Path | None, lora_name: str | None = None, step: int = 0) -> None: + async with self._lock: + self._desired = AdapterState( + name=lora_name, + path=weights_path if lora_name else None, + step=step, + ) + # Sync all servers in parallel for faster weight updates + await asyncio.gather(*[self._sync_server_adapter(ip) for ip in self._servers.keys()]) + + async def wait_for_ready(self, model_name: str = "", timeout: int | None = None, min_servers: int = 1) -> None: + if timeout is None: + timeout = self.client_config.wait_for_ready_timeout + start = time.time() + while time.time() - start < timeout: + await self.sync() + if self.num_ready_servers >= min_servers: + return + self.logger.debug(f"Waiting for servers: {self.num_ready_servers}/{min_servers} ready") + await asyncio.sleep(self.sync_interval) + + raise TimeoutError(f"Timed out waiting for {min_servers} ready servers (got {self.num_ready_servers})") + + async def update_weights(self, weight_dir: Path | None, lora_name: str | None = None, step: int = 0) -> None: + if lora_name is None: + raise ValueError("Elastic inference pool requires LoRA training (lora_name must be set)") + await self.sync_weights(weight_dir, lora_name, step) + + def get_metrics(self) -> dict[str, float]: + return { + "elastic/num_servers": self.num_servers, + "elastic/num_ready_servers": self.num_ready_servers, + "elastic/desired_step": self._desired.step, + } diff --git a/src/prime_rl/utils/envs.py b/src/prime_rl/utils/envs.py deleted file mode 100644 index d2d4969fdfa..00000000000 --- a/src/prime_rl/utils/envs.py +++ /dev/null @@ -1,51 +0,0 @@ -import os -from typing import TYPE_CHECKING, Any, Callable - -if TYPE_CHECKING: - # PyTorch - CUDA_VISIBLE_DEVICES: list[int] - - -# All environment variables with their type parsers -_ENV_PARSERS: dict[str, Callable[[str], Any]] = { - "CUDA_VISIBLE_DEVICES": lambda x: list(map(int, x.split(","))), -} - -# Subset of environment variables with default values -_ENV_DEFAULTS: dict[str, str] = {} - - -# Initialize environment variables that have defaults -def set_defaults(env_defaults: dict[str, str]): - for key, default_value in env_defaults.items(): - if os.getenv(key) is None: - os.environ[key] = default_value - - -set_defaults(_ENV_DEFAULTS) - - -def get_env_value(env_parser: dict[str, Callable[[str], Any]], key: str) -> Any: - """Get an environment variable value with type safety and parsing.""" - if key not in env_parser: - raise AttributeError(f"Invalid environment variable: {key}") - - raw_value = os.getenv(key) - if raw_value is None: - return None - parser = env_parser[key] - - return parser(raw_value) - - -def get_dir(env_parser: dict[str, Callable[[str], Any]]) -> list[str]: - """Get list of available environment variable names.""" - return list(env_parser.keys()) - - -def __getattr__(name: str) -> Any: - return get_env_value(_ENV_PARSERS, name) - - -def __dir__() -> list[str]: - return get_dir(_ENV_PARSERS) diff --git a/src/prime_rl/utils/logger.py b/src/prime_rl/utils/logger.py index 04d4af8240f..953ddbcb745 100644 --- a/src/prime_rl/utils/logger.py +++ b/src/prime_rl/utils/logger.py @@ -1,16 +1,70 @@ +import json as json_module import logging import sys -from pathlib import Path +import traceback +from typing import Any # Global logger instance _LOGGER = None +_JSON_LOGGING = False NO_BOLD = "\033[22m" RESET = "\033[0m" -class _VerifiersInterceptHandler(logging.Handler): - """Intercept stdlib logging from verifiers and route to loguru with [verifiers] tag.""" +def build_log_entry(record) -> dict: + """Build a flat JSON log entry from a loguru record.""" + extra = record["extra"] + + # Handle progress events specially - emit structured progress format + if extra.get("_progress"): + return { + "timestamp": record["time"].isoformat(), + "level": record["level"].name, + "type": "progress", + "desc": extra["desc"], + "current": extra["current"], + "total": extra["total"], + "percent": extra["percent"], + **({"step": extra["step"]} if extra.get("step") is not None else {}), + **({"extra": extra["postfix"]} if extra.get("postfix") else {}), + } + + # Standard log entry + log_entry = { + "timestamp": record["time"].isoformat(), + "level": record["level"].name, + "message": record["message"], + "module": record["module"], + "function": record["function"], + "line": record["line"], + } + if record["exception"] is not None: + exc = record["exception"] + log_entry["exception"] = "".join(traceback.format_exception(exc.type, exc.value, exc.traceback)) + # Extract tag from extra if present (used by workers to identify themselves) + if extra: + if "tag" in extra: + log_entry["tag"] = extra["tag"] + extra = {k: v for k, v in extra.items() if k != "tag"} + if extra: + log_entry["extra"] = extra + return log_entry + + +def json_sink(message) -> None: + """Sink that outputs flat JSON to stdout for log aggregation (Loki, Grafana, etc.).""" + log_entry = build_log_entry(message.record) + sys.stdout.write(json_module.dumps(log_entry) + "\n") + sys.stdout.flush() + + +class InterceptHandler(logging.Handler): + """Intercept standard logging library and routes to our prime-rl logger with specified prefix.""" + + def __init__(self, prefix: str | None): + super().__init__() + self.prefix = prefix def emit(self, record: logging.LogRecord) -> None: logger = get_logger() @@ -24,13 +78,23 @@ def emit(self, record: logging.LogRecord) -> None: frame = frame.f_back depth += 1 - logger.opt(depth=depth, exception=record.exc_info).log(level, f"[verifiers] {record.getMessage()}") + message = record.getMessage() + if self.prefix is not None: + message = f"[{self.prefix}] {message}" + logger.opt(depth=depth, exception=record.exc_info).log(level, message) -def setup_logger(log_level: str, log_file: Path | None = None, append: bool = False, tag: str | None = None): - global _LOGGER +def setup_logger( + log_level: str = "info", + tag: str | None = None, + json_logging: bool = False, +): + global _LOGGER, _JSON_LOGGING + _JSON_LOGGING = json_logging + + # Clean up old logger instance to prevent resource leaks if _LOGGER is not None: - raise RuntimeError("Logger already set. Please call `setup_logger` only once.") + _LOGGER.remove() # Format message with optional tag prefix tag_prefix = f"[{tag}] " if tag else "" @@ -67,14 +131,15 @@ def setup_logger(log_level: str, log_file: Path | None = None, append: bool = Fa extra={}, ) - # Install console handler - logger.add(sys.stdout, format=format, level=log_level.upper(), colorize=True) + # Bind tag to logger context for JSON mode (in non-JSON mode, tag is in the format string) + if json_logging and tag: + logger = logger.bind(tag=tag) - # If specified, install file handler - if log_file is not None: - if not append and log_file.exists(): - log_file.unlink() - logger.add(log_file, format=format, level=log_level.upper(), colorize=True) + # Install console handler (enqueue=True only for JSON mode to avoid blocking in async contexts) + if json_logging: + logger.add(json_sink, level=log_level.upper(), enqueue=True) + else: + logger.add(sys.stdout, format=format, level=log_level.upper(), colorize=True) # Disable critical logging logger.critical = lambda _: None @@ -86,30 +151,83 @@ def setup_logger(log_level: str, log_file: Path | None = None, append: bool = Fa def get_logger(): - """ - Get the global logger. This function is shared across submodules such as - training and inference to accesst the global logger instance. Raises if the - logger has not been set. - - Returns: - The global logger. - """ global _LOGGER if _LOGGER is None: - raise RuntimeError("Logger not set. Please call `set_logger` first.") + _LOGGER = setup_logger() return _LOGGER def reset_logger(): - """Reset the global logger. Useful mainly in test to clear loggers between tests.""" - global _LOGGER + """Reset the logger. Useful mainly in tests.""" + global _LOGGER, _JSON_LOGGING + if _LOGGER is not None: + _LOGGER.remove() _LOGGER = None - - -def intercept_verifiers_logging(level: str = "DEBUG"): - """Intercept all verifiers stdlib logging and route through prime-rl loguru with [verifiers] tag.""" - vf_logger = logging.getLogger("verifiers") - vf_logger.handlers.clear() - vf_logger.addHandler(_VerifiersInterceptHandler()) - vf_logger.setLevel(level.upper()) - vf_logger.propagate = False + _JSON_LOGGING = False + + +class ProgressTracker: + """Progress tracker that uses tqdm or logs progress when JSON logging is enabled.""" + + def __init__( + self, + total: int, + desc: str, + json_logging: bool | None = None, + log_every_percent: int = 10, + step: int | None = None, + ): + self.total = total + self.desc = desc + self.step = step + self.json_logging = json_logging if json_logging is not None else _JSON_LOGGING + self.log_every_percent = log_every_percent + self.current = 0 + self._last_logged_percent = -log_every_percent + self._postfix: dict[str, Any] = {} + + if self.json_logging: + self._pbar = None + # Don't log 0% on init - only log on actual progress + else: + from tqdm import tqdm + + self._pbar = tqdm(total=total, desc=desc) + + def update(self, n: int = 1): + self.current += n + if self._pbar is not None: + self._pbar.update(n) + else: + self._log_progress() + + def set_postfix(self, postfix: dict[str, Any]): + self._postfix = postfix + if self._pbar is not None: + self._pbar.set_postfix(postfix) + + def _log_progress(self): + percent = int(100 * self.current / self.total) if self.total > 0 else 0 + if percent >= self._last_logged_percent + self.log_every_percent or self.current >= self.total: + self._emit_progress(percent) + self._last_logged_percent = percent + + def _emit_progress(self, percent: int): + """Emit progress as structured JSON through loguru (only called in JSON logging mode).""" + get_logger().bind( + _progress=True, + desc=self.desc, + current=self.current, + total=self.total, + percent=percent, + step=self.step, + postfix=self._postfix if self._postfix else None, + ).info("progress") + + def close(self): + if self._pbar is not None: + self._pbar.close() + elif self.current > 0 and self.current < self.total: + percent = int(100 * self.current / self.total) + if percent > self._last_logged_percent: + self._emit_progress(percent) diff --git a/src/prime_rl/utils/messages.py b/src/prime_rl/utils/messages.py new file mode 100644 index 00000000000..13ebf6a6dcf --- /dev/null +++ b/src/prime_rl/utils/messages.py @@ -0,0 +1,62 @@ +"""Message utilities — normalize, deserialize tool calls, strip content.""" + +import json +from typing import Any + + +def normalize_messages(messages: Any, default_role: str) -> list[dict[str, Any]]: + if messages is None: + return [] + if isinstance(messages, str): + return [{"role": default_role, "content": messages}] + if isinstance(messages, dict): + return [dict(messages)] + if isinstance(messages, list): + normalized: list[dict[str, Any]] = [] + for message in messages: + if isinstance(message, str): + normalized.append({"role": default_role, "content": message}) + elif isinstance(message, dict): + normalized.append(dict(message)) + else: + raise TypeError(f"Unsupported message type: {type(message)}") + return normalized + raise TypeError(f"Unsupported messages container type: {type(messages)}") + + +def deserialize_tool_calls(messages: list[dict[str, Any]]) -> list[dict[str, Any]]: + def _deserialize_tool_call(tool_call: dict[str, Any]) -> dict[str, Any]: + function = tool_call.get("function", {}) + arguments = function.get("arguments") + if isinstance(arguments, str): + arguments = json.loads(arguments) + return { + **tool_call, + "function": {**function, "arguments": arguments}, + } + + deserialized_messages: list[dict[str, Any]] = [] + for message in messages: + if "tool_calls" not in message: + deserialized_messages.append(dict(message)) + continue + + tool_calls = message.get("tool_calls") or [] + deserialized_messages.append( + { + **message, + "tool_calls": [_deserialize_tool_call(tc) for tc in tool_calls], + } + ) + + return deserialized_messages + + +def strip_message_content(messages: list[dict[str, Any]]) -> list[dict[str, Any]]: + def _strip(message: dict[str, Any]) -> dict[str, Any]: + content = message.get("content") + if isinstance(content, str): + return {**message, "content": content.strip()} + return message + + return [_strip(message) for message in messages] diff --git a/src/prime_rl/utils/metrics_server.py b/src/prime_rl/utils/metrics_server.py index f4925e15554..5b4b5155f90 100644 --- a/src/prime_rl/utils/metrics_server.py +++ b/src/prime_rl/utils/metrics_server.py @@ -12,16 +12,10 @@ from typing import TYPE_CHECKING from loguru import logger +from prometheus_client import CONTENT_TYPE_LATEST, CollectorRegistry, Gauge, generate_latest if TYPE_CHECKING: - from prime_rl.utils.config import MetricsServerConfig - -try: - from prometheus_client import CONTENT_TYPE_LATEST, CollectorRegistry, Gauge, generate_latest - - PROMETHEUS_AVAILABLE = True -except ImportError: - PROMETHEUS_AVAILABLE = False + from prime_rl.configs.shared import MetricsServerConfig @dataclass @@ -103,53 +97,54 @@ def __init__(self, config: "MetricsServerConfig"): super().__init__(config.port, config.host) self.config = config - if PROMETHEUS_AVAILABLE: - self._registry = CollectorRegistry() - self._step = Gauge("trainer_step", "Current training step", registry=self._registry) - self._loss = Gauge("trainer_loss", "Current training loss", registry=self._registry) - self._throughput = Gauge( - "trainer_throughput_tokens_per_sec", "Training throughput in tokens/sec", registry=self._registry - ) - self._last_step_ts = Gauge( - "trainer_last_step_timestamp_seconds", "Unix timestamp of last step", registry=self._registry - ) - self._grad_norm = Gauge("trainer_grad_norm", "Gradient norm", registry=self._registry) - self._peak_mem = Gauge("trainer_peak_memory_gib", "Peak GPU memory in GiB", registry=self._registry) - self._lr = Gauge("trainer_learning_rate", "Current learning rate", registry=self._registry) - self._mfu = Gauge("trainer_mfu_percent", "Model FLOPS utilization %", registry=self._registry) - self._entropy = Gauge("trainer_entropy", "Mean entropy", registry=self._registry) - self._mismatch_kl = Gauge( - "trainer_mismatch_kl", "KL divergence between trainer and inference model", registry=self._registry - ) - # Aggregate run metrics - self._runs_discovered = Gauge( - "trainer_runs_discovered", "Number of run folders discovered", registry=self._registry - ) - self._runs_active = Gauge( - "trainer_runs_active", "Number of runs with assigned slots", registry=self._registry - ) - self._runs_ready = Gauge( - "trainer_runs_ready", "Number of runs ready for gradient updates", registry=self._registry - ) - self._runs_max = Gauge("trainer_runs_max", "Maximum run capacity", registry=self._registry) - # Per-run metrics with labels - self._run_step = Gauge("trainer_run_step", "Training step for run", ["run"], registry=self._registry) - self._run_tokens = Gauge( - "trainer_run_tokens", "Total tokens processed by run", ["run"], registry=self._registry - ) - self._run_learning_rate = Gauge( - "trainer_run_learning_rate", "Current learning rate for run", ["run"], registry=self._registry - ) - self._run_ready = Gauge( - "trainer_run_ready", - "Whether run is ready for updates (1=ready, 0=not ready)", - ["run"], - registry=self._registry, - ) - # Track known run labels for cleanup - self._known_runs: set[str] = set() - else: - self._registry = None + self._registry = CollectorRegistry() + self._step = Gauge("trainer_step", "Current training step", registry=self._registry) + self._loss = Gauge("trainer_loss", "Current training loss", registry=self._registry) + self._throughput = Gauge( + "trainer_throughput_tokens_per_sec", "Training throughput in tokens/sec", registry=self._registry + ) + self._last_step_ts = Gauge( + "trainer_last_step_timestamp_seconds", "Unix timestamp of last step", registry=self._registry + ) + self._grad_norm = Gauge("trainer_grad_norm", "Gradient norm", registry=self._registry) + self._peak_mem = Gauge("trainer_peak_memory_gib", "Peak GPU memory in GiB", registry=self._registry) + self._lr = Gauge("trainer_learning_rate", "Current learning rate", registry=self._registry) + self._mfu = Gauge("trainer_mfu_percent", "Model FLOPS utilization %", registry=self._registry) + self._entropy = Gauge("trainer_entropy", "Mean entropy", registry=self._registry) + self._mismatch_kl = Gauge( + "trainer_mismatch_kl", "KL divergence between trainer and inference model", registry=self._registry + ) + self._kl_ent_ratio = Gauge("trainer_kl_ent_ratio", "Ratio of mismatch KL to entropy", registry=self._registry) + self._zero_grad_ratio = Gauge( + "trainer_zero_grad_ratio", + "Fraction of tracked parameter elements with zero gradient", + registry=self._registry, + ) + # Aggregate run metrics + self._runs_discovered = Gauge( + "trainer_runs_discovered", "Number of run folders discovered", registry=self._registry + ) + self._runs_active = Gauge("trainer_runs_active", "Number of runs with assigned slots", registry=self._registry) + self._runs_ready = Gauge( + "trainer_runs_ready", "Number of runs ready for gradient updates", registry=self._registry + ) + self._runs_max = Gauge("trainer_runs_max", "Maximum run capacity", registry=self._registry) + # Per-run metrics with labels + self._run_step = Gauge("trainer_run_step", "Training step for run", ["run"], registry=self._registry) + self._run_tokens = Gauge( + "trainer_run_tokens", "Total tokens processed by run", ["run"], registry=self._registry + ) + self._run_learning_rate = Gauge( + "trainer_run_learning_rate", "Current learning rate for run", ["run"], registry=self._registry + ) + self._run_ready = Gauge( + "trainer_run_ready", + "Whether run is ready for updates (1=ready, 0=not ready)", + ["run"], + registry=self._registry, + ) + # Track known run labels for cleanup + self._known_runs: set[str] = set() def _make_handler(self) -> type[BaseHTTPRequestHandler]: """Create handler with /metrics and /health endpoints.""" @@ -167,14 +162,9 @@ def do_GET(self): def _handle_metrics(self): self.send_response(200) - if registry is not None: - self.send_header("Content-Type", CONTENT_TYPE_LATEST) - self.end_headers() - self.wfile.write(generate_latest(registry)) - else: - self.send_header("Content-Type", "text/plain") - self.end_headers() - self.wfile.write(b"# prometheus_client not installed\n") + self.send_header("Content-Type", CONTENT_TYPE_LATEST) + self.end_headers() + self.wfile.write(generate_latest(registry)) def _handle_health(self): self.send_response(200) @@ -193,9 +183,6 @@ def start(self) -> None: logger.warning("Metrics server already started") return - if not PROMETHEUS_AVAILABLE: - logger.warning("prometheus_client not installed. Install with: uv sync --extra metrics") - self._server = HTTPServer((self.host, self.port), self._make_handler()) self._thread = threading.Thread(target=self._server.serve_forever, daemon=True) self._thread.start() @@ -208,26 +195,28 @@ def update( step: int, loss: float, throughput: float, - grad_norm: float, + grad_norm: float | None, peak_memory_gib: float, learning_rate: float, mfu: float = 0.0, entropy: float = 0.0, mismatch_kl: float = 0.0, + zero_grad_ratio: float = 0.0, ) -> None: """Update metrics after a training step.""" - if not PROMETHEUS_AVAILABLE: - return - self._step.set(step) self._loss.set(loss) self._throughput.set(throughput) - self._grad_norm.set(grad_norm) + if grad_norm is not None: + self._grad_norm.set(grad_norm) self._peak_mem.set(peak_memory_gib) self._lr.set(learning_rate) self._mfu.set(mfu) self._entropy.set(entropy) self._mismatch_kl.set(mismatch_kl) + self._zero_grad_ratio.set(zero_grad_ratio) + if entropy > 0: + self._kl_ent_ratio.set(mismatch_kl / entropy) self._last_step_ts.set(time.time()) def update_runs( @@ -243,9 +232,6 @@ def update_runs( runs_max: Maximum run capacity run_stats: List of per-run statistics """ - if not PROMETHEUS_AVAILABLE: - return - # Update aggregate metrics self._runs_discovered.set(runs_discovered) self._runs_active.set(len(run_stats)) diff --git a/src/prime_rl/utils/monitor/__init__.py b/src/prime_rl/utils/monitor/__init__.py index dd3c4644ba7..10d02ddf33b 100644 --- a/src/prime_rl/utils/monitor/__init__.py +++ b/src/prime_rl/utils/monitor/__init__.py @@ -2,12 +2,12 @@ from transformers.tokenization_utils import PreTrainedTokenizer -from prime_rl.utils.config import PrimeMonitorConfig, WandbWithExtrasConfig +from prime_rl.configs.shared import PrimeMonitorConfig, WandbWithExtrasConfig +from prime_rl.utils.config import BaseConfig from prime_rl.utils.monitor.base import Monitor, NoOpMonitor from prime_rl.utils.monitor.multi import MultiMonitor from prime_rl.utils.monitor.prime import PrimeMonitor from prime_rl.utils.monitor.wandb import WandbMonitor -from prime_rl.utils.pydantic_config import BaseSettings __all__ = [ "Monitor", @@ -34,14 +34,19 @@ def setup_monitor( wandb_config: WandbWithExtrasConfig | None = None, output_dir: Path | None = None, tokenizer: PreTrainedTokenizer | None = None, - run_config: BaseSettings | None = None, + run_config: BaseConfig | None = None, *, prime_config: PrimeMonitorConfig | None = None, + keep_full_history: bool = True, # Backward compatibility: support old 'config' keyword argument config: WandbWithExtrasConfig | None = None, ) -> Monitor: """ Sets up monitors to log metrics. + + `keep_full_history`: when False, monitors retain only the most recent + metrics dict. The orchestrator passes False outside `--bench` mode to + avoid an unbounded list growing for the lifetime of the run. """ global _MONITOR if _MONITOR is not None: @@ -59,6 +64,7 @@ def setup_monitor( output_dir=output_dir, tokenizer=tokenizer, run_config=run_config, + keep_full_history=keep_full_history, ) ) @@ -69,15 +75,15 @@ def setup_monitor( output_dir=output_dir, tokenizer=tokenizer, run_config=run_config, + keep_full_history=keep_full_history, ) ) if len(monitors) == 0: - _MONITOR = NoOpMonitor() + _MONITOR = NoOpMonitor(keep_full_history=keep_full_history) elif len(monitors) == 1: _MONITOR = monitors[0] else: _MONITOR = MultiMonitor(monitors) return _MONITOR - diff --git a/src/prime_rl/utils/monitor/base.py b/src/prime_rl/utils/monitor/base.py index b9251484f5e..bc88080451d 100644 --- a/src/prime_rl/utils/monitor/base.py +++ b/src/prime_rl/utils/monitor/base.py @@ -1,26 +1,49 @@ +import random from abc import ABC, abstractmethod from typing import Any import verifiers as vf +def sample_items_for_logging(items: list[Any], sample_ratio: float | None) -> list[Any]: + """Apply monitor sample_ratio semantics to a batch of items. + + - ``None`` keeps the full batch. + - ``<= 0`` logs nothing. + - ``0 < ratio < 1`` logs a random subset with a minimum of 1 item. + - ``>= 1`` keeps the full batch. + """ + if sample_ratio is None: + return items + if sample_ratio <= 0.0: + return [] + if sample_ratio >= 1.0 or len(items) <= 1: + return items + + max_samples = max(1, int(len(items) * sample_ratio)) + if len(items) <= max_samples: + return items + + return random.sample(items, max_samples) + + class Monitor(ABC): """Base class for all monitoring implementations. - + Subclasses should initialize a `history` attribute as a list of dictionaries to store logged metrics. """ @abstractmethod - def log(self, metrics: dict[str, Any], step: int | None = None) -> None: + def log(self, metrics: dict[str, Any], step: int) -> None: pass @abstractmethod - def log_samples(self, rollouts: list[vf.State], step: int) -> None: + def log_samples(self, rollouts: list[vf.RolloutOutput], step: int) -> None: pass @abstractmethod - def log_final_samples(self) -> None: + def log_eval_samples(self, rollouts: list[vf.RolloutOutput], env_name: str, step: int) -> None: pass @abstractmethod @@ -39,16 +62,20 @@ def close(self) -> None: class NoOpMonitor(Monitor): """Monitor that does nothing. Used when no monitors are configured.""" - def __init__(self): + def __init__(self, keep_full_history: bool = True): self.history: list[dict[str, Any]] = [] + self._keep_full_history = keep_full_history - def log(self, metrics: dict[str, Any], step: int | None = None) -> None: - self.history.append(metrics) + def log(self, metrics: dict[str, Any], step: int) -> None: + if self._keep_full_history: + self.history.append(metrics) + else: + self.history = [metrics] - def log_samples(self, rollouts: list[vf.State], step: int) -> None: + def log_samples(self, rollouts: list[vf.RolloutOutput], step: int) -> None: pass - def log_final_samples(self) -> None: + def log_eval_samples(self, rollouts: list[vf.RolloutOutput], env_name: str, step: int) -> None: pass def save_final_summary(self, filename: str = "final_summary.json") -> None: @@ -56,4 +83,3 @@ def save_final_summary(self, filename: str = "final_summary.json") -> None: def log_distributions(self, distributions: dict[str, list[float]], step: int) -> None: pass - diff --git a/src/prime_rl/utils/monitor/multi.py b/src/prime_rl/utils/monitor/multi.py index ecf73eeaca0..0f32ea5906c 100644 --- a/src/prime_rl/utils/monitor/multi.py +++ b/src/prime_rl/utils/monitor/multi.py @@ -19,26 +19,26 @@ def history(self) -> list[dict[str, Any]]: return [] return self.monitors[0].history - def log(self, metrics: dict[str, Any], step: int | None = None) -> None: + def log(self, metrics: dict[str, Any], step: int) -> None: for monitor in self.monitors: try: monitor.log(metrics, step=step) except Exception as e: self.logger.warning(f"Failed to log metrics to {monitor.__class__.__name__}: {e}") - def log_samples(self, rollouts: list[vf.State], step: int) -> None: + def log_samples(self, rollouts: list[vf.RolloutOutput], step: int) -> None: for monitor in self.monitors: try: monitor.log_samples(rollouts=rollouts, step=step) except Exception as e: self.logger.warning(f"Failed to log samples to {monitor.__class__.__name__}: {e}") - def log_final_samples(self) -> None: + def log_eval_samples(self, rollouts: list[vf.RolloutOutput], env_name: str, step: int) -> None: for monitor in self.monitors: try: - monitor.log_final_samples() + monitor.log_eval_samples(rollouts=rollouts, env_name=env_name, step=step) except Exception as e: - self.logger.warning(f"Failed to log final samples to {monitor.__class__.__name__}: {e}") + self.logger.warning(f"Failed to log eval samples to {monitor.__class__.__name__}: {e}") def save_final_summary(self, filename: str = "final_summary.json") -> None: for monitor in self.monitors: diff --git a/src/prime_rl/utils/monitor/prime.py b/src/prime_rl/utils/monitor/prime.py index 8a779f02773..7e63b6c2d73 100644 --- a/src/prime_rl/utils/monitor/prime.py +++ b/src/prime_rl/utils/monitor/prime.py @@ -1,34 +1,129 @@ import asyncio +import io +import json +import math import os import time +from datetime import datetime, timezone from pathlib import Path from threading import Thread from typing import Any import httpx +import pyarrow as pa +import pyarrow.parquet as pq import verifiers as vf +from prime_cli.core.config import Config as PrimeConfig from transformers.tokenization_utils import PreTrainedTokenizer -from prime_rl.utils.config import PrimeMonitorConfig +from prime_rl.configs.shared import PrimeMonitorConfig +from prime_rl.utils.config import BaseConfig from prime_rl.utils.logger import get_logger -from prime_rl.utils.monitor.base import Monitor -from prime_rl.utils.pydantic_config import BaseSettings +from prime_rl.utils.monitor.base import Monitor, sample_items_for_logging + + +def _json(val: Any) -> str: + """JSON-serialize dicts/lists, pass strings through, default to empty string for None.""" + if isinstance(val, str): + return val + if val is None: + return "" + return json.dumps(val) + + +_SAMPLE_SCHEMA = pa.schema( + [ + ("run_id", pa.string()), + ("step", pa.int64()), + ("tag", pa.string()), + ("problem_id", pa.int64()), + ("sample_id", pa.int64()), + ("prompt", pa.string()), + ("completion", pa.string()), + ("trajectory", pa.string()), + ("answer", pa.string()), + ("env_name", pa.string()), + ("task", pa.string()), + ("info", pa.string()), + ("reward", pa.float64()), + ("advantage", pa.float64()), + ("metrics", pa.string()), + ("timing", pa.string()), + ("num_input_tokens", pa.int64()), + ("num_output_tokens", pa.int64()), + ("created_at", pa.timestamp("us", tz="UTC")), + ] +) + + +_DROPPED_JSON_VALUE = object() + + +def _drop_non_finite_json_values(value: Any, dropped_paths: list[str], path: str = "") -> Any: + if isinstance(value, float) and not math.isfinite(value): + dropped_paths.append(path) + return _DROPPED_JSON_VALUE + + if isinstance(value, dict): + return { + key: sanitized_item + for key, item in value.items() + if ( + sanitized_item := _drop_non_finite_json_values( + item, + dropped_paths, + f"{path}.{key}" if path else str(key), + ) + ) + is not _DROPPED_JSON_VALUE + } + + if isinstance(value, list): + return [ + sanitized_item + for idx, item in enumerate(value) + if (sanitized_item := _drop_non_finite_json_values(item, dropped_paths, f"{path}[{idx}]")) + is not _DROPPED_JSON_VALUE + ] + + return value class PrimeMonitor(Monitor): """Logs to Prime Intellect API.""" + def _sanitize_json_payload(self, endpoint: str, payload: dict[str, Any]) -> dict[str, Any]: + """Drop non-finite floats before sending JSON payloads to the public API.""" + dropped_paths: list[str] = [] + sanitized_payload = _drop_non_finite_json_values(payload, dropped_paths) + if not dropped_paths: + return payload + + preview = ", ".join(dropped_paths[:5]) + suffix = " ..." if len(dropped_paths) > 5 else "" + self.logger.warning( + f"Dropping {len(dropped_paths)} non-finite value(s) from Prime monitor {endpoint} payload: " + f"{preview}{suffix}" + ) + return sanitized_payload + def __init__( self, config: PrimeMonitorConfig | None, output_dir: Path | None = None, tokenizer: PreTrainedTokenizer | None = None, - run_config: BaseSettings | None = None, + run_config: BaseConfig | None = None, + keep_full_history: bool = True, ): self.config = config self.logger = get_logger() self.history: list[dict[str, Any]] = [] + self._keep_full_history = keep_full_history self.output_dir = output_dir + self._registered = False + self._finalized = False + self._closed = False + self._owner_pid = os.getpid() rank = int(os.environ.get("RANK", os.environ.get("DP_RANK", "0"))) self.enabled = self.config is not None @@ -41,24 +136,36 @@ def __init__( assert config is not None self.logger.info(f"Initializing {self.__class__.__name__} ({config})") - # Get API key from environment variable api_key = os.getenv(config.api_key_var) + if api_key is None: + prime_config = PrimeConfig() + api_key = prime_config.api_key + if not api_key: self.logger.warning( - f"API key not found. Set {config.api_key_var} environment variable. PrimeMonitor will not be able to upload data." + f"API key not found. Set {config.api_key_var} environment variable or run `prime login`. " + "PrimeMonitor will not be able to upload data." ) self.enabled = False return self.api_key = api_key - self.base_url = config.base_url + self.base_url = config.base_url.rstrip("/") + self._headers = { + "Authorization": f"Bearer {api_key}", + "x-api-key": api_key, + "Content-Type": "application/json", + } - # Get run_id from environment variable (check before allocating resources) run_id = os.getenv("RUN_ID") if not run_id: - self.logger.warning("RUN_ID environment variable not set. PrimeMonitor will not be able to upload data.") - self.enabled = False - return + run_id = self._register_run(config, run_config) + if run_id: + os.environ["RUN_ID"] = run_id + else: + self.enabled = False + return + self.run_id = run_id # Set up async HTTP client with background event loop. @@ -71,15 +178,102 @@ def __init__( os.register_at_fork(after_in_child=self._reinit_after_fork) # Optionally, initialize sample logging attributes - if config is not None and config.log_extras: + if config.log_extras: if config.log_extras.samples: self.last_log_samples_step = -1 + self._pending_sample_steps: set[int] = set() self.tokenizer = tokenizer if config.log_extras.distributions: self.last_log_distributions_step = -1 - def log(self, metrics: dict[str, Any], step: int | None = None) -> None: - self.history.append(metrics) + def _register_run(self, config: PrimeMonitorConfig, run_config: BaseConfig | None) -> str | None: + """Register an external run with the platform. Returns run_id on success, None on failure.""" + prime_config = None + team_id = config.team_id + frontend_url = config.frontend_url + if team_id is None or frontend_url is None: + prime_config = PrimeConfig() + if team_id is None: + team_id = prime_config.team_id + if frontend_url is None: + frontend_url = prime_config.frontend_url + + model = getattr(run_config, "model", None) if run_config else None + environments = getattr(run_config, "env", None) if run_config else None + wandb = getattr(run_config, "wandb", None) if run_config else None + + payload: dict = { + "base_model": model.name if model else "unknown", + "max_steps": getattr(run_config, "max_steps", None) or 0, + } + if config.run_name: + payload["name"] = config.run_name + if team_id: + payload["team_id"] = team_id + if environments: + payload["environments"] = [{"id": env.id} for env in environments if hasattr(env, "id")] + if wandb and getattr(wandb, "project", None): + payload["wandb_project"] = wandb.project + + try: + response = httpx.post( + f"{self.base_url}/external-runs", + headers=self._headers, + json=payload, + timeout=30, + ) + except httpx.HTTPError as e: + self.logger.warning(f"Failed to register platform run: {e}. PrimeMonitor will not be able to upload data.") + return None + + if response.status_code != 201: + self.logger.warning( + f"Failed to create platform run (HTTP {response.status_code}): {response.text}. " + "PrimeMonitor will not be able to upload data." + ) + return None + + run_id = response.json()["run"]["id"] + if frontend_url: + dashboard_url = f"{frontend_url.rstrip('/')}/dashboard/training/{run_id}" + self.logger.success(f"Monitor run at: {dashboard_url}") + else: + self.logger.success(f"Registered platform run {run_id}") + self._registered = True + return run_id + + def _finalize_run(self, success: bool) -> None: + """Mark the run as completed or failed on the platform.""" + if not self._registered: + return + + payload: dict = {"status": "completed" if success else "failed"} + status_label = "completed" if success else "failed" + self.logger.info(f"Finalizing platform run {self.run_id} as {status_label}") + + try: + response = httpx.put( + f"{self.base_url}/external-runs/{self.run_id}/status", + headers=self._headers, + json=payload, + timeout=30, + ) + except httpx.HTTPError as e: + self.logger.warning(f"Failed to finalize platform run {self.run_id}: {e}") + return + + if response.status_code != 200: + self.logger.warning( + f"Failed to finalize platform run {self.run_id} (HTTP {response.status_code}): {response.text}" + ) + return + self.logger.info(f"Platform run {self.run_id} marked as {status_label}") + + def log(self, metrics: dict[str, Any], step: int) -> None: + if self._keep_full_history: + self.history.append(metrics) + else: + self.history = [metrics] if not self.is_master: return if not self.enabled: @@ -92,8 +286,8 @@ def log(self, metrics: dict[str, Any], step: int | None = None) -> None: }, ) - def log_samples(self, rollouts: list[vf.State], step: int) -> None: - """Logs rollouts to Prime Intellect API.""" + def log_samples(self, rollouts: list[vf.RolloutOutput], step: int) -> None: + """Logs rollouts to Prime Intellect API using presigned URLs for direct R2 upload.""" if not self.is_master: return if not self.enabled: @@ -104,82 +298,201 @@ def log_samples(self, rollouts: list[vf.State], step: int) -> None: or not self.config.log_extras.samples or step % self.config.log_extras.interval != 0 ): - # Do not log samples if not enabled or not log interval step + return + + rollouts = sample_items_for_logging( + rollouts, + self.config.log_extras.sample_ratio, + ) + if not rollouts: return assert self.last_log_samples_step <= step, "Step must be greater than last logged step" + assert step not in self._pending_sample_steps, f"Step {step} upload already in progress" assert self.logger is not None, "Logger is required for sample logging" - self.logger.info(f"Logging samples to Prime Intellect API at step {step}") + self.logger.info(f"Logging {len(rollouts)} samples to Prime Intellect API at step {step}") start_time = time.perf_counter() - # Prepare samples for API - samples = [] - for rollout in rollouts: - # Extract prompt and completion separately from the last trajectory step - trajectory = rollout["trajectory"] - if not trajectory: - continue - last_step = trajectory[-1] - prompt_messages = last_step["prompt"] - completion_messages = last_step["completion"] - - # Serialize full trajectory array (excluding large response objects and token arrays) - trajectory_data = [] - for traj_step in rollout["trajectory"]: - trajectory_data.append( - { - "prompt": traj_step["prompt"], - "completion": traj_step["completion"], - "reward": traj_step.get("reward"), - "advantage": traj_step.get("advantage"), - "extras": traj_step.get("extras", {}), - "num_input_tokens": len(traj_step.get("tokens", {}).get("prompt_ids", [])) - if traj_step.get("tokens") - else None, - "num_output_tokens": len(traj_step.get("tokens", {}).get("completion_ids", [])) - if traj_step.get("tokens") - else None, - } - ) + parquet_bytes = self._rollouts_to_parquet_bytes(rollouts, step) - # Get info, timing, and metrics fields - send raw data, backend will serialize - info = rollout.get("info") - timing = rollout.get("timing") - metrics = rollout.get("metrics") + if not parquet_bytes: + self.logger.warning(f"No samples to log at step {step}") + return - sample = { - "step": step, - "example_id": rollout.get("example_id"), - "prompt": prompt_messages, - "completion": completion_messages, - "trajectory": trajectory_data, - "reward": rollout.get("reward"), - "advantage": rollout.get("advantage"), - "answer": rollout.get("answer"), - "task": rollout.get("task"), - "info": info, - "metrics": metrics, - "timing": timing, - } - samples.append(sample) + self._pending_sample_steps.add(step) + + # Use presigned URL flow for uploading samples + self._upload_samples_via_presigned_url(parquet_bytes, step) - # Upload samples - self._make_request( - "samples", - { - "run_id": self.run_id, - "step": step, - "samples": samples, - }, - ) - self.last_log_samples_step = step self.logger.debug( - f"Logged samples at step {step} to Prime Intellect API in {time.perf_counter() - start_time:.2f}s" + f"Initiated samples upload at step {step} to Prime Intellect API in {time.perf_counter() - start_time:.2f}s" + ) + + def _rollouts_to_parquet_bytes(self, rollouts: list[vf.RolloutOutput], step: int) -> bytes | None: + """Convert rollouts directly to Parquet bytes for upload.""" + now = datetime.now(timezone.utc) + rows = [] + + for sample_id, rollout in enumerate(rollouts): + prompt = rollout.get("prompt") + completion = rollout.get("completion") + trajectory = rollout.get("trajectory") or [] + if prompt is None or completion is None or not trajectory: + continue + + example_id = rollout.get("example_id") + try: + problem_id = int(example_id) if example_id is not None else sample_id + except (TypeError, ValueError): + problem_id = sample_id + + trajectory_data = [ + { + "prompt": ts["prompt"], + "completion": ts["completion"], + "reward": ts.get("reward"), + "advantage": ts.get("advantage"), + "extras": ts.get("extras", {}), + "num_input_tokens": len(ts["tokens"]["prompt_ids"]) if ts.get("tokens") else None, + "num_output_tokens": len(ts["tokens"]["completion_ids"]) if ts.get("tokens") else None, + } + for ts in trajectory + ] + + rows.append( + { + "run_id": self.run_id, + "step": step, + "tag": "", + "problem_id": problem_id, + "sample_id": sample_id, + "prompt": json.dumps(prompt), + "completion": json.dumps(completion), + "trajectory": json.dumps(trajectory_data), + "answer": rollout.get("answer") or "", + "env_name": rollout.get("env_name") or "", + "task": rollout.get("task") or "", + "info": _json(rollout.get("info")), + "reward": rollout.get("reward"), + "advantage": rollout.get("advantage"), + "metrics": _json(rollout.get("metrics")), + "timing": _json(rollout.get("timing")), + "num_input_tokens": 0, + "num_output_tokens": 0, + "created_at": now, + } + ) + + if not rows: + return None + + table = pa.Table.from_pylist(rows, schema=_SAMPLE_SCHEMA) + buf = io.BytesIO() + pq.write_table(table, buf, compression="snappy", use_dictionary=True, write_statistics=True) + return buf.getvalue() + + def _upload_samples_via_presigned_url(self, parquet_bytes: bytes, step: int) -> None: + """Upload Parquet samples using presigned URL flow (fire-and-forget).""" + future = asyncio.run_coroutine_threadsafe( + self._upload_samples_via_presigned_url_async(parquet_bytes, step), + self._loop, ) + self._pending_futures.append(future) + # Clean up completed futures to avoid memory growth + self._pending_futures = [f for f in self._pending_futures if not f.done()] + + async def _upload_samples_via_presigned_url_async( + self, parquet_bytes: bytes, step: int, max_retries: int = 3 + ) -> None: + """Upload Parquet bytes via presigned URL flow.""" + try: + presign_data = await self._request_presigned_url(step) + if not presign_data: + self.logger.warning(f"Failed to get presigned URL for samples at step {step}") + return - def log_final_samples(self) -> None: - """Log final samples (no-op - samples are logged per-step only).""" + presigned_url = presign_data["presigned_url"] + s3_key = presign_data["s3_key"] + + upload_success = await self._upload_to_r2( + presigned_url, parquet_bytes, content_type="application/parquet", max_retries=max_retries + ) + if not upload_success: + self.logger.warning(f"Failed to upload samples to R2 at step {step}") + return + + confirm_success = await self._confirm_samples_upload(step, s3_key) + if not confirm_success: + self.logger.warning(f"Failed to confirm samples upload at step {step}") + return + + self.last_log_samples_step = step + self.logger.debug(f"Successfully completed samples upload at step {step}") + + except Exception as e: + self.logger.warning(f"Failed to upload samples via presigned URL at step {step}: {type(e).__name__}: {e}") + finally: + self._pending_sample_steps.discard(step) + + async def _request_presigned_url(self, step: int) -> dict[str, Any] | None: + """Request a presigned URL from the backend.""" + try: + response = await self._client.post( + f"{self.base_url}/samples/presign", + headers=self._headers, + json={"run_id": self.run_id, "step": step}, + ) + response.raise_for_status() + response_data = response.json()["data"] + return { + "presigned_url": response_data["presignedUrl"], + "s3_key": response_data["s3Key"], + } + except Exception as e: + self.logger.warning(f"Failed to request presigned URL: {type(e).__name__}: {e}") + return None + + async def _upload_to_r2( + self, presigned_url: str, data: bytes, content_type: str = "application/json", max_retries: int = 3 + ) -> bool: + """Upload data to R2 using presigned URL.""" + for attempt in range(max_retries): + try: + response = await self._client.put(presigned_url, content=data, headers={"Content-Type": content_type}) + response.raise_for_status() + return True + except Exception as e: + if attempt == max_retries - 1: + self.logger.warning(f"Failed to upload to R2 after {max_retries} attempts: {type(e).__name__}: {e}") + return False + delay = 2**attempt + self.logger.debug(f"Retrying R2 upload in {delay}s (attempt {attempt + 1}/{max_retries})") + await asyncio.sleep(delay) + + async def _confirm_samples_upload(self, step: int, s3_key: str, max_retries: int = 3) -> bool: + """Confirm samples upload with the backend. Returns True on success.""" + for attempt in range(max_retries): + try: + response = await self._client.post( + f"{self.base_url}/samples/confirm", + headers=self._headers, + json={"run_id": self.run_id, "step": step, "s3_key": s3_key}, + ) + response.raise_for_status() + return True + except Exception as e: + if attempt == max_retries - 1: + self.logger.warning( + f"Failed to confirm samples upload after {max_retries} attempts: {type(e).__name__}: {e}" + ) + return False + delay = 2**attempt + self.logger.debug(f"Retrying samples confirm in {delay}s (attempt {attempt + 1}/{max_retries})") + await asyncio.sleep(delay) + return False + + def log_eval_samples(self, rollouts: list[vf.RolloutOutput], env_name: str, step: int) -> None: pass def log_distributions(self, distributions: dict[str, list[float]], step: int) -> None: @@ -217,25 +530,66 @@ def log_distributions(self, distributions: dict[str, list[float]], step: int) -> f"Logged distributions at step {step} to Prime Intellect API in {time.perf_counter() - start_time:.2f}s" ) + def _submit_final_summary(self, summary: dict[str, Any]) -> bool: + """Submit the final summary/finalize request synchronously.""" + payload = self._sanitize_json_payload( + "finalize", + {"run_id": self.run_id, "summary": summary}, + ) + + try: + response = httpx.post( + f"{self.base_url}/finalize", + headers=self._headers, + json=payload, + timeout=30, + ) + except httpx.HTTPError as e: + self.logger.warning(f"Failed to submit final summary for platform run {self.run_id}: {e}") + return False + + if response.status_code != 200: + self.logger.warning( + f"Failed to submit final summary for platform run {self.run_id} " + f"(HTTP {response.status_code}): {response.text}" + ) + return False + + return True + def save_final_summary(self, filename: str = "final_summary.json") -> None: """Save final summary to Prime Intellect API.""" if not self.is_master or not self.enabled: return self.logger.info("Saving final summary to Prime Intellect API") - self._make_request( - "finalize", - { - "run_id": self.run_id, - "summary": self.history[-1] if self.history else {}, - }, - ) + summary = self.history[-1] if self.history else {} + finalized_via_summary = self._submit_final_summary(summary) + + if os.getpid() != self._owner_pid: + return + + if finalized_via_summary: + self._finalized = True + return + + self._finalize_run(success=True) + self._finalized = True def close(self) -> None: """Close the HTTP client and stop the background event loop.""" - if not hasattr(self, "_client"): + if self._closed or not hasattr(self, "_client"): return + self._closed = True + + should_finalize = self.is_master and self.enabled and not self._finalized and os.getpid() == self._owner_pid + if should_finalize: + self._finalize_run(success=False) + self._finalized = True + + self._flush() + # Close the async client within the event loop async def _close_client() -> None: await self._client.aclose() @@ -260,6 +614,9 @@ def _init_async_client(self) -> None: self._thread = Thread(target=self._run_event_loop, daemon=True) self._thread.start() self._client = httpx.AsyncClient(timeout=30) + self._pending_futures: list[asyncio.Future] = [] + if hasattr(self, "_pending_sample_steps") and self._pending_sample_steps: + self._pending_sample_steps.clear() def _reinit_after_fork(self) -> None: """Reinitialize thread and event loop after fork.""" @@ -270,20 +627,34 @@ def _run_event_loop(self) -> None: asyncio.set_event_loop(self._loop) self._loop.run_forever() + def _flush(self, timeout: float = 30.0) -> None: + """Wait for all pending async requests to complete.""" + if not self.enabled or not hasattr(self, "_loop"): + return + + if not self._pending_futures: + return + + self.logger.debug(f"Flushing {len(self._pending_futures)} pending request(s)") + for future in self._pending_futures: + try: + future.result(timeout=timeout) + except Exception as e: + self.logger.debug(f"Pending request completed with error: {e}") + + self._pending_futures.clear() + async def _make_request_async(self, endpoint: str, data: dict[str, Any], max_retries: int = 3) -> None: """Make an async POST request to the Prime Intellect API with retries.""" - headers = { - "x-api-key": self.api_key, - "Content-Type": "application/json", - } full_endpoint = f"{self.base_url}/{endpoint}" + sanitized_data = self._sanitize_json_payload(endpoint, data) for attempt in range(max_retries): try: response = await self._client.post( full_endpoint, - headers=headers, - json=data, + headers=self._headers, + json=sanitized_data, ) response.raise_for_status() return # Success @@ -306,7 +677,10 @@ def _make_request(self, endpoint: str, data: dict[str, Any]) -> None: if not self.enabled: return - asyncio.run_coroutine_threadsafe( + future = asyncio.run_coroutine_threadsafe( self._make_request_async(endpoint, data), self._loop, ) + self._pending_futures.append(future) + # Clean up completed futures to avoid memory growth + self._pending_futures = [f for f in self._pending_futures if not f.done()] diff --git a/src/prime_rl/utils/monitor/wandb.py b/src/prime_rl/utils/monitor/wandb.py index abb8b118374..8b6598cd2a5 100644 --- a/src/prime_rl/utils/monitor/wandb.py +++ b/src/prime_rl/utils/monitor/wandb.py @@ -5,15 +5,16 @@ from pathlib import Path from typing import Any -import pandas as pd import verifiers as vf import wandb from transformers.tokenization_utils import PreTrainedTokenizer +from wandb.errors import CommError -from prime_rl.utils.config import WandbConfig, WandbWithExtrasConfig +from prime_rl.configs.shared import WandbConfig, WandbWithExtrasConfig +from prime_rl.utils.chat_template import deserialize_tool_calls +from prime_rl.utils.config import BaseConfig from prime_rl.utils.logger import get_logger -from prime_rl.utils.monitor.base import Monitor -from prime_rl.utils.pydantic_config import BaseSettings +from prime_rl.utils.monitor.base import Monitor, sample_items_for_logging class WandbMonitor(Monitor): @@ -24,16 +25,19 @@ def __init__( config: WandbConfig | WandbWithExtrasConfig | None, output_dir: Path | None = None, tokenizer: PreTrainedTokenizer | None = None, - run_config: BaseSettings | None = None, + run_config: BaseConfig | None = None, + keep_full_history: bool = True, ): self.config = config self.logger = get_logger() self.history: list[dict[str, Any]] = [] + self._keep_full_history = keep_full_history self.output_dir = output_dir rank = int(os.environ.get("RANK", os.environ.get("DP_RANK", "0"))) self.enabled = self.config is not None self.is_master = rank == 0 + if not self.enabled or not self.is_master: if not self.is_master: self.logger.warning(f"Skipping {self.__class__.__name__} initialization from non-master rank ({rank})") @@ -42,27 +46,77 @@ def __init__( assert config is not None self.logger.info(f"Initializing {self.__class__.__name__} ({config})") self._maybe_overwrite_wandb_command() - self.wandb = wandb.init( - project=config.project, - name=config.name, - id=config.id, - dir=output_dir, - resume="allow", - config=run_config.model_dump() if run_config else None, - mode="offline" if config.offline else None, - ) + + shared_mode = os.environ.get("WANDB_SHARED_MODE") == "1" + if shared_mode: + run_id = os.environ.get("WANDB_SHARED_RUN_ID") + label = os.environ.get("WANDB_SHARED_LABEL") + primary = label == "orchestrator" + settings = wandb.Settings( + mode="shared", + x_label=label, + x_primary=primary, + x_update_finish_state=primary, + ) + self.logger.info( + f"Using shared W&B mode ({label=}, {primary=}). " + "This is an experimental feature. Disable with --wandb.shared False" + ) + else: + run_id = None + primary = False + settings = wandb.Settings( + mode="offline" if config.offline else "online", + ) + + def init_wandb(max_retries: int): + for attempt in range(max_retries): + try: + return wandb.init( + id=run_id, + project=config.project, + entity=config.entity, + name=config.name, + group=config.group, + tags=config.tags, + dir=output_dir, + config=run_config.model_dump() if run_config else None, + settings=settings, + ) + except CommError as e: + if attempt + 1 == max_retries: + raise + if shared_mode and not primary: + msg = ( + f"Shared W&B run not yet created by primary - retrying in 10s ({attempt + 1}/{max_retries})" + ) + else: + msg = f"Transient W&B init error ({e}) - retrying in 10s ({attempt + 1}/{max_retries})" + self.logger.info(msg) + time.sleep(10) + + # Non-primary processes in shared mode wait for the primary to create the run. + # Everyone else still retries to absorb transient W&B server errors (e.g. 404 on upsertBucket). + max_retries = 30 if shared_mode and not primary else 5 + self.wandb = init_wandb(max_retries) + + wandb.define_metric("*", step_metric="step") # Optionally, initialize sample logging attributes if config is not None and isinstance(config, WandbWithExtrasConfig) and config.log_extras: if config.log_extras.samples: self.last_log_samples_step = -1 - self.samples_cols = ["step", "task", "example_id", "messages", "input_ids", "reward"] + self.samples_cols = ["step", "env_name", "task", "example_id", "messages", "input_ids", "reward"] self.samples_table = wandb.Table( columns=self.samples_cols, log_mode="INCREMENTAL", ) self.tokenizer = tokenizer - self.samples = [] + self.eval_samples_cols = ["step", "env", "task", "example_id", "completion", "reward"] + self.eval_samples_table = wandb.Table( + columns=self.eval_samples_cols, + log_mode="INCREMENTAL", + ) def _maybe_overwrite_wandb_command(self) -> None: """Overwrites sys.argv with the start command if it is set in the environment variables.""" @@ -71,15 +125,18 @@ def _maybe_overwrite_wandb_command(self) -> None: self.logger.debug(f"Found WANDB_ARGS in environment variables {wandb_args}") sys.argv = json.loads(wandb_args) - def log(self, metrics: dict[str, Any], step: int | None = None) -> None: - self.history.append(metrics) + def log(self, metrics: dict[str, Any], step: int) -> None: + if self._keep_full_history: + self.history.append(metrics) + else: + self.history = [metrics] if not self.is_master: return if not self.enabled: return - wandb.log(metrics, step=step) + wandb.log({**metrics, "step": step}) - def log_samples(self, rollouts: list[vf.State], step: int) -> None: + def log_samples(self, rollouts: list[vf.RolloutOutput], step: int) -> None: """Logs rollouts to W&B table.""" if not self.is_master: return @@ -93,11 +150,18 @@ def log_samples(self, rollouts: list[vf.State], step: int) -> None: # Do not log samples if not enabled or not log interval step return + rollouts = sample_items_for_logging( + rollouts, + self.config.log_extras.sample_ratio, + ) + if not rollouts: + return + assert self.tokenizer is not None, "Tokenizer is required for sample logging" assert self.last_log_samples_step <= step, "Step must be greater than last logged step" assert self.logger is not None, "Logger is required for sample logging" - self.logger.info(f"Logging samples to W&B table at step {step}") + self.logger.info(f"Logging {len(rollouts)} samples to W&B table at step {step}") start_time = time.perf_counter() for rollout in rollouts: @@ -110,6 +174,7 @@ def log_samples(self, rollouts: list[vf.State], step: int) -> None: messages_text = self.tokenizer.decode(full_ids) sample = { "step": step, + "env_name": rollout.get("env_name"), "task": rollout.get("task"), "example_id": rollout["example_id"], "messages": messages_text, @@ -120,14 +185,13 @@ def log_samples(self, rollouts: list[vf.State], step: int) -> None: "Order of columns in the table must be the same as order of the keys here" ) self.samples_table.add_data(*sample.values()) - self.samples.append(sample) - wandb.log({"samples": self.samples_table}, step=step) + wandb.log({"samples": self.samples_table, "step": step}) self.last_log_samples_step = step self.logger.debug(f"Logged samples at step {step} to W&B table in {time.perf_counter() - start_time:.2f}s") - def log_final_samples(self) -> None: - """Log final samples to W&B table.""" + def log_eval_samples(self, rollouts: list[vf.RolloutOutput], env_name: str, step: int) -> None: + """Logs eval rollouts to a separate W&B table.""" if not self.is_master: return if ( @@ -138,10 +202,26 @@ def log_final_samples(self) -> None: ): return - self.logger.info("Logging final samples to W&B table") - df = pd.DataFrame(self.samples) - table = wandb.Table(dataframe=df) - wandb.log({"final-samples": table}) + for rollout in rollouts: + completion = rollout.get("completion") + if not completion: + continue + if isinstance(completion, list): + try: + completion = self.tokenizer.apply_chat_template(deserialize_tool_calls(completion), tokenize=False) + except Exception: + completion = str(completion) + sample = { + "step": step, + "env": env_name, + "task": rollout.get("task"), + "example_id": rollout["example_id"], + "completion": completion, + "reward": rollout["reward"], + } + self.eval_samples_table.add_data(*sample.values()) + + wandb.log({"eval/samples": self.eval_samples_table, "step": step}) def log_distributions(self, distributions: dict[str, list[float]], step: int) -> None: """Log distributions (no-op for W&B).""" diff --git a/src/prime_rl/utils/nccl.py b/src/prime_rl/utils/nccl.py new file mode 100644 index 00000000000..bb581ceeb04 --- /dev/null +++ b/src/prime_rl/utils/nccl.py @@ -0,0 +1,42 @@ +import os + +import pynvml + +from prime_rl.utils.logger import get_logger + + +def disable_nccl_p2p_if_unavailable() -> None: + """Disable NCCL P2P/SHM transports when GPUs lack NVLink interconnect. + + In environments without NVLink (e.g. PCIe-only or VMs with GPU passthrough), + NCCL's P2P and SHM transports can fail because they rely on CUDA IPC which + requires peer access. Disabling them forces NCCL to use socket-based transport. + + Uses pynvml to check physical GPU topology, which works regardless of + CUDA_VISIBLE_DEVICES restrictions on the current process. + """ + pynvml.nvmlInit() + try: + n = pynvml.nvmlDeviceGetCount() + if n < 2: + return + + for i in range(n): + handle = pynvml.nvmlDeviceGetHandleByIndex(i) + for link in range(18): + try: + if pynvml.nvmlDeviceGetNvLinkState(handle, link): + return # NVLink found, P2P should work + except pynvml.NVMLError: + break + + # No NVLink found between any GPUs — disable P2P/SHM if not already set + if "NCCL_P2P_DISABLE" not in os.environ or "NCCL_SHM_DISABLE" not in os.environ: + os.environ["NCCL_P2P_DISABLE"] = "1" + os.environ["NCCL_SHM_DISABLE"] = "1" + get_logger().warning( + "No NVLink detected, disabling NCCL P2P and SHM transports. " + "Override by setting NCCL_P2P_DISABLE=0 and NCCL_SHM_DISABLE=0 explicitly." + ) + finally: + pynvml.nvmlShutdown() diff --git a/src/prime_rl/utils/pathing.py b/src/prime_rl/utils/pathing.py index ebe6b2a7294..63a230e4181 100644 --- a/src/prime_rl/utils/pathing.py +++ b/src/prime_rl/utils/pathing.py @@ -1,4 +1,5 @@ import asyncio +import shutil import time from pathlib import Path @@ -9,6 +10,57 @@ def get_log_dir(output_dir: Path) -> Path: return output_dir / "logs" +def format_log_message( + log_dir: Path, + trainer: bool = False, + orchestrator: bool = False, + inference: bool = False, + job_log: bool = False, + train_env_names: list[str] | None = None, + eval_env_names: list[str] | None = None, + num_train_nodes: int = 1, + num_infer_nodes: int = 0, +) -> str: + """Format a log message showing where to find all log files.""" + col = 18 + i1 = " " * 2 + i2 = " " * 3 + i3 = " " * 4 + max_name = col - 4 + + log_lines: list[str] = [] + if job_log: + log_lines.append(f"{i1}{'Job:':<{col}}tail -F {log_dir.parent}/job_*.log") + if trainer: + log_lines.append(f"{i1}{'Trainer:':<{col}}tail -F {log_dir}/trainer.log") + if num_train_nodes > 1: + log_lines.append(f"{i2}{'All nodes:':<{col - 1}}tail -F {log_dir}/trainer/node_*.log") + log_lines.append(f"{i2}{'All ranks:':<{col - 1}}tail -F {log_dir}/trainer/torchrun/*/*/*/*.log") + if orchestrator: + log_lines.append(f"{i1}{'Orchestrator:':<{col}}tail -F {log_dir}/orchestrator.log") + if inference: + log_lines.append(f"{i1}{'Inference:':<{col}}tail -F {log_dir}/inference.log") + if num_infer_nodes > 1: + log_lines.append(f"{i2}{'All nodes:':<{col - 1}}tail -F {log_dir}/inference/node_*.log") + if train_env_names: + env_log_dir = log_dir / "envs" + log_lines.append(f"{i1}{'Envs:':<{col}}tail -F {env_log_dir}/*/*/*.log") + log_lines.append(f"{i2}{'Train:':<{col - 1}}tail -F {env_log_dir}/train/*/*.log") + for name in train_env_names: + short = name if len(name) <= max_name else name[: max_name - 3] + "..." + log_lines.append(f"{i3}{f'{short}:':<{col - 2}}tail -F {env_log_dir}/train/{name}/*.log") + if eval_env_names: + log_lines.append(f"{i2}{'Eval:':<{col - 1}}tail -F {env_log_dir}/eval/*/*.log") + for name in eval_env_names: + short = name if len(name) <= max_name else name[: max_name - 3] + "..." + log_lines.append(f"{i3}{f'{short}:':<{col - 2}}tail -F {env_log_dir}/eval/{name}/*.log") + return "Logs:\n" + "\n".join(log_lines) + + +def get_config_dir(output_dir: Path) -> Path: + return output_dir / "configs" + + def get_ckpt_dir(output_dir: Path) -> Path: return output_dir / "checkpoints" @@ -29,10 +81,6 @@ def get_broadcast_dir(output_dir: Path) -> Path: return output_dir / "broadcasts" -def get_env_worker_log_file(output_dir: Path, env_name: str) -> Path: - return output_dir / "logs" / "env_workers" / f"{env_name}.log" - - def get_step_path(path: Path, step: int) -> Path: return path / f"step_{step}" @@ -43,6 +91,12 @@ def get_all_ckpt_steps(ckpt_dir: Path) -> list[int]: return sorted([int(step_dir.name.split("_")[-1]) for step_dir in step_dirs]) +def get_stable_ckpt_steps(ckpt_dir: Path) -> list[int]: + """Gets checkpoint steps that have STABLE file, sorted in ascending order.""" + steps = get_all_ckpt_steps(ckpt_dir) + return [s for s in steps if (ckpt_dir / f"step_{s}" / "STABLE").exists()] + + def resolve_latest_ckpt_step(ckpt_dir: Path) -> int | None: """Gets the latest checkpoint step from the checkpoint directory. Returns None if no checkpoints are found.""" steps = get_all_ckpt_steps(ckpt_dir) @@ -56,6 +110,68 @@ def resolve_latest_ckpt_step(ckpt_dir: Path) -> int | None: return latest_step +def has_checkpoints(output_dir: Path) -> bool: + """Check if the output directory contains any checkpoints.""" + ckpt_dir = get_ckpt_dir(output_dir) + return ckpt_dir.exists() and any(ckpt_dir.iterdir()) + + +def validate_output_dir(output_dir: Path, *, resuming: bool, clean: bool, ckpt_output_dir: Path | None = None) -> None: + """Validate the output directory before training starts. + + Raises if the directory contains checkpoints from a previous run, unless + explicitly resuming or opting into cleaning. Other artifacts (logs, + rollouts, configs) are fine and don't trigger the error. + + When ckpt_output_dir is set, checkpoints live there instead of under + output_dir, so the guard and clean logic check both locations. + """ + dirs_to_check = [output_dir] + if ckpt_output_dir is not None and ckpt_output_dir != output_dir: + dirs_to_check.append(ckpt_output_dir) + + if resuming: + return + if clean: + logger = get_logger() + for d in dirs_to_check: + if d.exists(): + logger.warning(f"Cleaning existing directory: {d}") + shutil.rmtree(d) + return + for d in dirs_to_check: + if has_checkpoints(d): + raise FileExistsError( + f"Directory '{d}' already contains checkpoints from a previous run. " + f"To resume the latest step of the previous run, set ckpt.resume_step=-1 or --ckpt.resume-step -1 via CLI. " + f"To delete the existing directory and start fresh, set clean_output_dir=true or --clean-output-dir via CLI. " + f"Otherwise use a unique output_dir for this experiment." + ) + + +def clean_future_steps(output_dir: Path, resume_step: int) -> None: + """Remove stale rollouts and broadcasts past ``resume_step``. + + Pass ``resume_step=-1`` to wipe every step directory (fresh runs). + """ + run_default = output_dir / "run_default" + dirs = [ + get_rollout_dir(output_dir), + get_rollout_dir(run_default), + get_broadcast_dir(run_default), + ] + + for directory in dirs: + steps_to_delete = [step for step in get_all_ckpt_steps(directory) if step > resume_step] + if not steps_to_delete: + continue + get_logger().info( + f"Deleting {len(steps_to_delete)} step directories in {directory} ({','.join(map(str, steps_to_delete))})" + ) + for step in steps_to_delete: + shutil.rmtree(get_step_path(directory, step)) + + def sync_wait_for_path(path: Path, interval: int = 1, log_interval: int = 10) -> None: logger = get_logger() wait_time = 0 diff --git a/src/prime_rl/utils/process.py b/src/prime_rl/utils/process.py new file mode 100644 index 00000000000..ab14813f62e --- /dev/null +++ b/src/prime_rl/utils/process.py @@ -0,0 +1,74 @@ +import contextlib +import os +import signal +import subprocess +from subprocess import Popen +from threading import Event, Thread + +import psutil +import setproctitle + +from prime_rl.utils.logger import get_logger + +PRIME_RL_PROC_PREFIX = "PRIME-RL" + + +def set_proc_title(name: str) -> None: + """Set the process title for visibility in tools like ``ps`` and ``htop``. + + Args: + name: A short, descriptive label (e.g. ``Trainer``, ``Orchestrator``). + The process title is set to ``{PRIME_RL_PROC_PREFIX}::{name}``. + """ + title = f"{PRIME_RL_PROC_PREFIX}::{name}" + setproctitle.setproctitle(title) + + +def cleanup_threads(threads: list[Thread]): + """Cleanup a list of threads""" + for thread in threads: + thread.join(timeout=5) + + +def cleanup_process(pid: int, sig: int = signal.SIGTERM): + """Kill a process and all its descendants. + + Walks the process tree via ``psutil`` so that grandchildren spawned by + intermediate wrappers (e.g. ``uv``, ``torchrun``) are reliably reached + regardless of process-group boundaries. + """ + try: + parent = psutil.Process(pid) + children = parent.children(recursive=True) + except psutil.NoSuchProcess: + return + for child in children: + with contextlib.suppress(ProcessLookupError): + os.kill(child.pid, sig) + with contextlib.suppress(ProcessLookupError): + os.kill(pid, sig) + + +def cleanup_processes(processes: list[Popen]): + """Cleanup a list of subprocesses by killing their entire process trees.""" + for process in processes: + if process.poll() is not None: + continue + cleanup_process(process.pid, signal.SIGTERM) + try: + process.wait(timeout=60) + except subprocess.TimeoutExpired: + cleanup_process(process.pid, signal.SIGKILL) + get_logger().debug(f"Cleaned up process {process.pid}") + + +def monitor_process(process: Popen, stop_event: Event, error_queue: list, process_name: str): + """Monitor a subprocess and signal errors via shared queue.""" + process.wait() + + if process.returncode != 0: + err_msg = f"{process_name.capitalize()} failed with exit code {process.returncode}" + if process.stderr: + err_msg += f"\n{process.stderr.read().decode('utf-8')}" + error_queue.append(RuntimeError(err_msg)) + stop_event.set() diff --git a/src/prime_rl/utils/pydantic_config.py b/src/prime_rl/utils/pydantic_config.py deleted file mode 100644 index 237b813e3ad..00000000000 --- a/src/prime_rl/utils/pydantic_config.py +++ /dev/null @@ -1,283 +0,0 @@ -import sys -import uuid -import warnings -from pathlib import Path -from typing import Annotated, ClassVar, Type, TypeVar - -import tomli -import tomli_w -from pydantic import BaseModel, ConfigDict, Field, field_validator -from pydantic_settings import BaseSettings as PydanticBaseSettings -from pydantic_settings import ( - PydanticBaseSettingsSource, - SettingsConfigDict, - TomlConfigSettingsSource, -) - - -class BaseConfig(BaseModel): - model_config = ConfigDict(extra="forbid") - - @field_validator("*", mode="before") - @classmethod - def empty_str_to_none(cls, v): - """ - This allow to support setting None via toml files using the string "None" - """ - if v == "None": - return None - return v - - -class BaseSettings(PydanticBaseSettings): - """ - Base settings class for all configs. - """ - - # These are two somewhat hacky workarounds inspired by https://github.com/pydantic/pydantic-settings/issues/259 to ensure backwards compatibility with our old CLI system `pydantic_config` - _TOML_FILES: ClassVar[list[str]] = [] - - toml_files: Annotated[ - list[str] | None, - Field( - description="List of extra TOML files to load (paths are relative to the TOML file containing this field). If provided, will override all other config files. Note: This field is only read from within configuration files - setting --toml-files from CLI has no effect.", - exclude=True, - ), - ] = None - - @classmethod - def set_toml_files(cls, toml_files: list[str]) -> None: - """ - Set the global TOML files to be used for this config. - These are two somewhat hacky workarounds inspired by https://github.com/pydantic/pydantic-settings/issues/259 to ensure backwards compatibility with our old CLI system `pydantic_config` - """ - cls._TOML_FILES = toml_files - - @classmethod - def clear_toml_files(cls) -> None: - """ - Clear the global TOML files to be used for this config. - """ - cls._TOML_FILES = [] - - def set_unknown_args(self, unknown_args: list[str]) -> None: - self._unknown_args = unknown_args - - def get_unknown_args(self) -> list[str]: - return self._unknown_args - - @classmethod - def settings_customise_sources( - cls, - settings_cls: type["BaseSettings"], - init_settings: PydanticBaseSettingsSource, - env_settings: PydanticBaseSettingsSource, - dotenv_settings: PydanticBaseSettingsSource, - file_secret_settings: PydanticBaseSettingsSource, - ) -> tuple[PydanticBaseSettingsSource, ...]: - # This is a hacky way to dynamically load TOML file paths from CLI - # https://github.com/pydantic/pydantic-settings/issues/259 - return ( - TomlConfigSettingsSource(settings_cls, toml_file=cls._TOML_FILES), - init_settings, - env_settings, - dotenv_settings, - file_secret_settings, - ) - - # Pydantic settings configuration - model_config = SettingsConfigDict( - env_prefix="PRIME_", - env_nested_delimiter="__", - cli_parse_args=False, - cli_kebab_case=True, - cli_implicit_flags=True, - cli_use_class_docs_for_groups=True, - nested_model_default_partial_update=True, # Allows partial updates to nested model default objects - ) - - -def check_path_and_handle_inheritance(path: Path, seen_files: list[Path], nested_key: str | None) -> bool | None: - """ - Recursively look for inheritance in a toml file. Return a list of all toml files to load. - - Example: - If config.toml has `toml_files = ["base.toml"]` and base.toml has - `toml_files = ["common.toml"]`, this returns ["config.toml", "base.toml", "common.toml"] - nested_key: smth like "train.optim" - - Returns: - True if some toml inheritance is detected, False otherwise. - """ - if path in seen_files: - return - - if not path.exists(): - raise FileNotFoundError(f"TOML file {path} does not exist") - - with open(path, "rb") as f: - data = tomli.load(f) - - if nested_key is not None: - nested_keys = nested_key.split(".") - for key in nested_keys: - new_data = {} - new_data[key] = data - data = new_data - - path = get_temp_toml_file() - - with open(path, "wb") as f: - tomli_w.dump(data, f) - - seen_files.append(path) - - recurence = False - if "toml_files" in data: - if nested_key is not None: - raise NotImplementedError("--train @ helo.toml where helo.toml point to a tom file is not yet supported") - maybe_new_files = [path.parent / file for file in data["toml_files"]] - - files = [file for file in maybe_new_files if str(file).endswith(".toml")] - # todo which should probably look for infinite inheritance loops here - for file in files: - recurence = True - check_path_and_handle_inheritance(file, seen_files, nested_key=None) - - return recurence - - -# Extract config file paths from CLI to pass to pydantic-settings as toml source -# This enables the use of `@` to pass config file paths to the CLI -def extract_toml_paths(args: list[str]) -> tuple[list[str], list[str]]: - toml_paths = [] - remaining_args = args.copy() - recurence = False - cli_toml_file_count = 0 - for prev_arg, arg, next_arg in zip([""] + args[:-1], args, args[1:] + [""]): - if arg == "@": - toml_path = next_arg - remaining_args.remove(arg) - remaining_args.remove(next_arg) - - if prev_arg.startswith("--"): - remaining_args.remove(prev_arg) - nested_key = prev_arg.replace("--", "") - else: - nested_key = None - - recurence = recurence or check_path_and_handle_inheritance(Path(toml_path), toml_paths, nested_key) - cli_toml_file_count += 1 - - if recurence and cli_toml_file_count > 1: - warnings.warn( - f"{len(toml_paths)} TOML files are added via CLI ({', '.join(toml_paths)}) and at least one of them links to another file. This is not supported yet. Please either compose multiple config files via directly CLI or specify a single file linking to multiple other files" - ) - - return toml_paths, remaining_args - - -def to_kebab_case(args: list[str]) -> list[str]: - """ - Converts CLI argument keys from snake case to kebab case. - - For example, `--max_batch_size 1` will be transformed `--max-batch-size 1`. - """ - for i, arg in enumerate(args): - if arg.startswith("--"): - args[i] = arg.replace("_", "-") - return args - - -def get_all_fields(model: BaseModel | type) -> list[str]: - if isinstance(model, BaseModel): - model_cls = model.__class__ - else: - model_cls = model - - fields = [] - for name, field in model_cls.model_fields.items(): - field_type = field.annotation - fields.append(name) - if field_type is not None and hasattr(field_type, "model_fields"): - sub_fields = get_all_fields(field_type) - fields.extend(f"{name}.{sub}" for sub in sub_fields) - return fields - - -def parse_unknown_args(args: list[str], config_cls: type) -> tuple[list[str], list[str]]: - known_fields = get_all_fields(config_cls) - known_args = [] - unknown_args = [] - i = 0 - n = len(args) - - def get_is_key(arg: str) -> bool: - return arg.startswith("--") or arg.startswith("-") - - while i < n: - is_key = get_is_key(args[i]) - has_value = False if i >= n - 1 or get_is_key(args[i + 1]) else True - if not is_key: - i += 1 - continue - if args[i].startswith("--"): - key = args[i][2:] - else: - key = args[i][1:] - key = key.replace("-", "_") - if key in known_fields: - known_args.append(args[i]) - if has_value: - known_args.append(args[i + 1]) - i += 2 - else: - i += 1 - else: - unknown_args.append(args[i]) - if has_value: - unknown_args.append(args[i + 1]) - i += 2 - else: - i += 1 - - return known_args, unknown_args - - -T = TypeVar("T", bound=BaseSettings) - - -# Class[BaseSettings] -def parse_argv(config_cls: Type[T], allow_extras: bool = False) -> T: - """ - Parse CLI arguments and TOML configuration files into a pydantic settings instance. - - Supports loading TOML files via @ syntax (e.g., @config.toml or @ config.toml). - Automatically converts snake_case CLI args to kebab-case for pydantic compatibility. - TOML files can inherit from other TOML files via the 'toml_files' field. - - Args: - config_cls: A pydantic BaseSettings class to instantiate with parsed configuration. - - Returns: - An instance of config_cls populated with values from TOML files and CLI args. - CLI args take precedence over TOML file values. - """ - toml_paths, cli_args = extract_toml_paths(sys.argv[1:]) - config_cls.set_toml_files(toml_paths) - try: - if allow_extras: - cli_args, unknown_args = parse_unknown_args(cli_args, config_cls) - config = config_cls(_cli_parse_args=to_kebab_case(cli_args)) - if allow_extras: - config.set_unknown_args(unknown_args) - return config - finally: - config_cls.clear_toml_files() # Always clear, even on exception - - -def get_temp_toml_file() -> Path: - temp_uuid = str(uuid.uuid4()) - root_path = Path(".pydantic_config") - root_path.mkdir(exist_ok=True) - return root_path / f"temp_{temp_uuid}.toml" diff --git a/src/prime_rl/utils/sequence.py b/src/prime_rl/utils/sequence.py new file mode 100644 index 00000000000..5a222e35a68 --- /dev/null +++ b/src/prime_rl/utils/sequence.py @@ -0,0 +1,35 @@ +from __future__ import annotations + +import torch + + +def get_cu_seqlens_from_position_ids(position_ids: torch.Tensor) -> tuple[torch.Tensor, int]: + """Build local-relative cumulative sequence lengths from packed position ids. + + ``position_ids`` may be a full packed sequence or a context-parallel local shard. + In the shard case, the first local token can be a continuation with a non-zero + position id, so cumulative offsets must be anchored to the local tensor rather + than to the absolute position id value. + """ + flat_position_ids = position_ids.view(-1) + total_tokens = flat_position_ids.numel() + assert total_tokens > 0, "Cannot build cu_seqlens for an empty position_ids tensor" + + zero_starts = (flat_position_ids == 0).nonzero(as_tuple=True)[0] + starts = torch.cat( + [torch.zeros(1, dtype=zero_starts.dtype, device=zero_starts.device), zero_starts] + ).unique_consecutive() + + ends = torch.cat( + [ + starts[1:], + torch.tensor([total_tokens], dtype=starts.dtype, device=starts.device), + ] + ) + seqlens = ends - starts + + cu_seqlens = torch.empty(seqlens.numel() + 1, dtype=torch.int32, device=position_ids.device) + cu_seqlens[0] = 0 + cu_seqlens[1:] = seqlens.cumsum(dim=0, dtype=torch.int32) + + return cu_seqlens, seqlens.max().item() diff --git a/src/prime_rl/utils/usage_reporter.py b/src/prime_rl/utils/usage_reporter.py new file mode 100644 index 00000000000..aca348439a0 --- /dev/null +++ b/src/prime_rl/utils/usage_reporter.py @@ -0,0 +1,83 @@ +"""Reports training token usage per step to the platform API for billing. + +Activated entirely via environment variables — not exposed in user-facing +config. Self-hosted users are unaffected. + +The platform API is idempotent on (run_id, step, usage_type), so replays +after crash-resume are safe. +""" + +import asyncio +import os + +from prime_rl.utils.background_async import BackgroundAsync +from prime_rl.utils.logger import get_logger + + +class UsageReporter: + """Fire-and-forget training token usage reporter. + + Activated when ``PI_USAGE_BASE_URL`` is set in the environment. + Wraps a ``BackgroundAsync`` runner so reporting never blocks training. + """ + + def __init__(self) -> None: + self.logger = get_logger() + api_key = os.environ.get("PI_USAGE_API_KEY") + base_url = os.environ.get("PI_USAGE_BASE_URL", "").rstrip("/") + # Both must be set. The orchestrator already gates on this, but + # validate here too so any future caller fails fast at construction + # rather than crashing inside httpx (which rejects None header + # values with an obscure .encode() error on every retry). + if not api_key or not base_url: + raise ValueError( + "UsageReporter requires PI_USAGE_BASE_URL and PI_USAGE_API_KEY to be set in the environment." + ) + self._api_key = api_key + self._base_url = base_url + self._runner = BackgroundAsync(client_timeout=30.0) + self.logger.debug(f"Usage reporter initialized (base_url={self._base_url})") + + def report_training_usage(self, run_id: str, step: int, tokens: int) -> None: + url = f"{self._base_url}/usage" + payload = {"run_id": run_id, "step": step, "usage_type": "training", "tokens": tokens} + self._runner.submit(lambda: self._post(url, payload)) + + def close(self) -> None: + self._runner.close() + + async def _post(self, url: str, data: dict, max_retries: int = 3) -> None: + headers = {"x-api-key": self._api_key, "Content-Type": "application/json"} + for attempt in range(max_retries): + try: + response = await self._runner.client.post(url, headers=headers, json=data) + if response.status_code == 200: + body = response.json() + if body.get("status") == "duplicate": + self.logger.debug( + "Usage already recorded: run=%s step=%s type=%s", + data.get("run_id"), + data.get("step"), + data.get("usage_type"), + ) + return + response.raise_for_status() + return + except Exception as e: + if attempt == max_retries - 1: + self.logger.warning( + "Failed to report usage after %d attempts: %s: %s", + max_retries, + type(e).__name__, + e, + ) + else: + delay = 2**attempt + self.logger.debug( + "Retrying usage report in %ds (attempt %d/%d): %s", + delay, + attempt + 1, + max_retries, + e, + ) + await asyncio.sleep(delay) diff --git a/src/prime_rl/utils/utils.py b/src/prime_rl/utils/utils.py index a4a72253b5a..dec33f625c6 100644 --- a/src/prime_rl/utils/utils.py +++ b/src/prime_rl/utils/utils.py @@ -1,7 +1,9 @@ import asyncio import functools +import importlib import os import subprocess +import sys from collections import defaultdict from contextlib import contextmanager from pathlib import Path @@ -11,7 +13,7 @@ import torch.distributed as dist import wandb -from prime_rl.orchestrator.config import EnvConfig, EvalEnvConfig +from prime_rl.configs.orchestrator import EnvConfig, EvalEnvConfig from prime_rl.utils.logger import get_logger # TODO: Change all imports to use utils.pathing @@ -23,6 +25,7 @@ get_eval_dir, get_log_dir, get_rollout_dir, + get_stable_ckpt_steps, get_step_path, get_weights_dir, resolve_latest_ckpt_step, @@ -31,75 +34,11 @@ ) -def rgetattr(obj: Any, attr_path: str) -> Any: - """ - Try to get a (nested) attribute from an object. For example: - - ```python - class Foo: - bar = "baz" - - class Bar: - foo = Foo() - - foo = Foo() - bar = Bar() - ``` - - Here, the following holds: - - `getattr(foo, "bar")` will return `"baz"`. - - `getattr(bar, "foo)` will return an object of type `Foo`. - - `getattr(bar, "foo.bar")` will error - - This function solves this. `rgetattr(bar, "foo.bar")` will return `"baz"`. - - Args: - obj: The object to get the attribute from. - attr_path: The path to the attribute, nested using `.` as separator. - - Returns: - The attribute - """ - attrs = attr_path.split(".") - current = obj - - for attr in attrs: - if not hasattr(current, attr): - raise AttributeError(f"'{type(current).__name__}' object has no attribute '{attr}'") - current = getattr(current, attr) - - return current - - -def rsetattr(obj: Any, attr_path: str, value: Any) -> None: - """ - Try to set a (nested) attribute from an object. For example: - - ```python - class Foo: - bar = "baz" - - class Bar: - foo = Foo() - - foo = Foo() - bar = Bar() - ``` - - Here, the following holds: - - `rsetattr(bar, "foo.bar", "qux")` will set `bar.foo.bar` to `"qux"`. - - `rsetattr(bar, "foo.bar", "qux")` will set `bar.foo.bar` to `"qux"`. - - Args: - obj: The object to set the attribute on. - attr_path: The path to the attribute, nested using `.` as separator. - value: The value to set the attribute to. - """ - if "." not in attr_path: - return setattr(obj, attr_path, value) - attr_path, attr = attr_path.rsplit(".", 1) - obj = rgetattr(obj, attr_path) - setattr(obj, attr, value) +def import_object(dotted_path: str) -> Any: + """Import an object from a dotted path like 'my_module.submodule.MyClass'.""" + module_path, _, name = dotted_path.rpartition(".") + module = importlib.import_module(module_path) + return getattr(module, name) def capitalize(s: str) -> str: @@ -120,9 +59,13 @@ async def async_wrapper(*args, **kwargs): ret = await func(*args, **kwargs) wandb.finish() return ret - except Exception as e: + except Exception: + get_logger().opt(exception=True).error(f"Fatal error in {func.__name__}") wandb.finish(exit_code=1) - raise e + # sys.exit raises SystemExit so the finally block still runs. + # raise alone doesn't terminate the process in an async context — + # the event loop swallows it and the process hangs indefinitely. + sys.exit(1) finally: if dist.is_initialized(): dist.destroy_process_group() @@ -136,9 +79,11 @@ def sync_wrapper(*args, **kwargs): ret = func(*args, **kwargs) wandb.finish() return ret - except Exception as e: + except Exception: + get_logger().opt(exception=True).error(f"Fatal error in {func.__name__}") wandb.finish(exit_code=1) - raise e + # sys.exit raises SystemExit so the finally block still runs. + sys.exit(1) finally: if dist.is_initialized(): dist.destroy_process_group() @@ -190,27 +135,32 @@ def to_row_format(dict_of_lists: dict[str, list[Any]]) -> list[dict[str, Any]]: return [dict(zip(dict_of_lists.keys(), values)) for values in zip(*dict_of_lists.values())] -def format_time(time_in_seconds: float) -> str: - """Format a time in seconds to a human-readable format.""" - from datetime import timedelta - - td = timedelta(seconds=time_in_seconds) - days = td.days - hours, remainder = divmod(td.seconds, 3600) - minutes, seconds = divmod(remainder, 60) - - # Format based on magnitude - if days > 0: - total_hours = days * 24 + hours - return f"{total_hours + minutes / 60:.2f}h" - elif hours > 0: - return f"{hours + minutes / 60:.2f}h" - elif minutes > 0: - return f"{minutes + seconds / 60:.2f}m" +def format_time(time_s: float) -> str: + """ + Format a time in seconds to a human-readable format: + - >1d -> Xd Yh + - >1h -> Xh Ym + - >1m -> Xm Ys + - <1s -> Xms + - Else: Xs + """ + if time_s >= 86400: + d = time_s // 86400 + h = (time_s % 86400) // 3600 + return f"{d:.0f}d" + (f" {h:.0f}h" if h > 0 else "") + elif time_s >= 3600: + h = time_s // 3600 + m = (time_s % 3600) // 60 + return f"{h:.0f}h" + (f" {m:.0f}m" if m > 0 else "") + elif time_s >= 60: + m = time_s // 60 + s = (time_s % 60) // 1 + return f"{m:.0f}m" + (f" {s:.0f}s" if s > 0 else "") + elif time_s < 1: + ms = time_s * 1e3 + return f"{ms:.0f}ms" else: - # Include microseconds for sub-second precision - total_seconds = seconds + td.microseconds / 1_000_000 - return f"{total_seconds:.2f}s" + return f"{time_s:.0f}s" def format_num(num: float | int, precision: int = 2) -> str: @@ -281,14 +231,22 @@ def default_dtype(dtype): torch.set_default_dtype(prev) -def install_env(env_id: str) -> None: +def install_env(env_id: str, prerelease: bool = False) -> None: """Install an environment in subprocess.""" logger = get_logger() logger.info(f"Installing environment {env_id}") install_cmd = ["uv", "run", "--no-sync", "prime", "env", "install", env_id] + if prerelease: + install_cmd.insert(-1, "--prerelease") result = subprocess.run(install_cmd, capture_output=True, text=True) + for line in result.stdout.splitlines(): + if line.strip(): + logger.info(line) + for line in result.stderr.splitlines(): + if line.strip(): + logger.warning(line) if result.returncode != 0: - raise RuntimeError(f"Failed to install environment {env_id} (stdout={result.stdout}, stderr={result.stderr})") + raise RuntimeError(f"Failed to install environment {env_id} (exit code {result.returncode})") logger.info(f"Successfully installed environment {env_id}") diff --git a/src/prime_rl/utils/vf.py b/src/prime_rl/utils/vf.py deleted file mode 100644 index f67f322a4a1..00000000000 --- a/src/prime_rl/utils/vf.py +++ /dev/null @@ -1,168 +0,0 @@ -import asyncio -from itertools import cycle -from typing import Any, cast - -import verifiers as vf -from openai import AsyncOpenAI -from openai.types.chat.chat_completion import ChatCompletion -from tqdm import tqdm - -from prime_rl.orchestrator.utils import get_semaphore - - -async def generate_group( - client: AsyncOpenAI, - env: vf.Environment, - model_name: str, - example: dict, - rollouts_per_example: int, - sampling_args: dict, -) -> list[vf.State]: - """Asynchronously generate and score rollouts for a single group.""" - semaphore = await get_semaphore() - group_inputs = [vf.RolloutInput(**example) for _ in range(rollouts_per_example)] - return await env.run_group( - group_inputs=group_inputs, - client=client, - model=model_name, - gen_sampling_args=sampling_args, - gen_sem=semaphore, - score_sem=semaphore, - ) - - -async def generate_rollout( - client: AsyncOpenAI, - env: vf.Environment, - model_name: str, - example: dict, - sampling_args: dict, -) -> vf.State: - """Asynchronously generate and score a single rollout.""" - semaphore = await get_semaphore() - rollout_input = vf.RolloutInput(**example) - - state = await env.run_rollout(semaphore, rollout_input, client, model_name, sampling_args) - await env.rubric.score_rollout(state, score_sem=semaphore) - return state - - -async def generate_batch( - clients: list[AsyncOpenAI], - env: vf.Environment, - model_name: str, - examples: list[dict], - rollouts_per_example: int, - sampling_args: dict, - pbar_description: str = "Generating rollouts", -) -> list[vf.State]: - """Asynchronously generate and score rollouts for a list of groups (batch).""" - - total_rollouts = len(examples) * rollouts_per_example - pbar = tqdm(total=total_rollouts, desc=pbar_description) - - async def generate_group_with_progress(client, example): - """Generate rollouts for one problem and update progress.""" - result = await generate_group(client, env, model_name, example, rollouts_per_example, sampling_args) - pbar.update(rollouts_per_example) - return result - - try: - group_states_list: list[list[vf.State]] = await asyncio.gather( - *[generate_group_with_progress(client, example) for client, example in zip(cycle(clients), examples)] - ) - finally: - pbar.close() - - return [state for group_states in group_states_list for state in group_states] - - -def get_prompt_len(state: vf.State) -> int: - """Compute the number of prompt tokens from vf.State. Defined as the number of prompt IDs from the first trajectory step. If raw tokens are not available, falls back to checking the usage of the first response.""" - if not state["trajectory"]: - return 0 - first_step = state["trajectory"][0] - if first_step["tokens"] is not None: - return len(first_step["tokens"]["prompt_ids"]) - first_step_response = cast(ChatCompletion, first_step["response"]) - return getattr(first_step_response.usage, "prompt_tokens", 0) - - -def get_seq_len(state: vf.State) -> int: - """Compute the number of tokens from vf.State. Defined as the sum of prompt and completion tokens from the last trajectory step. If raw tokens are not available, falls back to checking the usage of the last response.""" - if not state["trajectory"]: - return 0 - last_step = state["trajectory"][-1] - if last_step["tokens"] is not None: - return len(last_step["tokens"]["prompt_ids"]) + len(last_step["tokens"]["completion_ids"]) - last_step_response = cast(ChatCompletion, last_step["response"]) - return getattr(last_step_response.usage, "total_tokens", 0) - - -def get_completion_len(state: vf.State) -> int: - """Compute the number of completion tokens from vf.State. Defined as the difference between the total number of tokens and the number of prompt tokens.""" - return get_seq_len(state) - get_prompt_len(state) - - -def get_is_truncated(state: vf.State) -> bool: - """Check if the rollout is truncated. If raw tokens are not available, falls back to checking the finish reason of the last response.""" - if not state["trajectory"]: - return False - last_step = state["trajectory"][-1] - if last_step["tokens"] is not None: - return last_step["tokens"]["is_truncated"] - last_step_response = cast(ChatCompletion, last_step["response"]) - return last_step_response.choices[0].finish_reason == "length" - - -def to_serializable_trajectory_step(step: vf.TrajectoryStep) -> dict: - """Returns a serializable version of vf.TrajectoryStep.""" - serializable_trajectory_step = cast(dict, step.copy()) - if "response" in step and isinstance(step["response"], ChatCompletion): - serializable_trajectory_step["response"] = step["response"].model_dump() - return serializable_trajectory_step - - -def from_serializable_trajectory_step(step: dict) -> vf.TrajectoryStep: - """Inverse of to_serializable_trajectory_step.""" - deserialized_trajectory_step = vf.TrajectoryStep(**step) - if "response" in step and isinstance(step["response"], dict): - deserialized_trajectory_step["response"] = ChatCompletion.model_validate(step["response"]) - return deserialized_trajectory_step - - -def to_serializable_state(state: vf.State) -> dict: - """Returns a serializable copy of vf.State.""" - serializable_state = cast(dict, state.copy()) - - # Flatten input fields to top level for serialization - # This is necessary because the dict object cannot forward access to input fields anymore, e.g. state["prompt"] will automatically forward to state["input"]["prompt"] in vf.State but not in dict. We solve this by populating the inputs into the top-level explicitly - if "input" in state: - input_dict = serializable_state.pop("input") - for field in vf.State.INPUT_FIELDS: - if field in input_dict: - serializable_state[field] = input_dict[field] - - if "trajectory" in state: - serializable_state["trajectory"] = [to_serializable_trajectory_step(step) for step in state["trajectory"]] - - return serializable_state - - -def from_serializable_state(serializable_state: dict) -> vf.State: - """Inverse of to_serializable_state.""" - # Extract input fields from top level and reconstruct input dict - input_dict: dict[str, Any] = {} - for field in vf.State.INPUT_FIELDS: - if field in serializable_state: - input_dict[field] = serializable_state.pop(field) - - if input_dict: - serializable_state["input"] = input_dict - - state = vf.State(**serializable_state) - - if "trajectory" in state: - state["trajectory"] = [from_serializable_trajectory_step(step) for step in state["trajectory"]] - - return state diff --git a/src/prime_rl/utils/vlm.py b/src/prime_rl/utils/vlm.py new file mode 100644 index 00000000000..a099c199251 --- /dev/null +++ b/src/prime_rl/utils/vlm.py @@ -0,0 +1,124 @@ +"""Vision-Language Model (VLM) support utilities. + +Central registry for VLM model families. All model-specific knowledge +lives here. Add new VLM families by extending VLM_REGISTRY. + +For custom models not in the registry, set overrides in config: + [model.vlm] + vision_encoder_attr = "model.my_vision" + language_model_attr = "model.my_lm" +""" + +from dataclasses import dataclass + +import torch.nn as nn +from transformers.configuration_utils import PretrainedConfig + + +@dataclass(frozen=True) +class VLMModelInfo: + """Per-model-family VLM architecture metadata.""" + + vision_encoder_attr: str + language_model_attr: str + + +# Central registry: model_type -> architecture info. +VLM_REGISTRY: dict[str, VLMModelInfo] = { + "qwen3_vl": VLMModelInfo(vision_encoder_attr="model.visual", language_model_attr="model.language_model"), + "qwen3_5": VLMModelInfo(vision_encoder_attr="model.visual", language_model_attr="model.language_model"), + "qwen3_5_moe": VLMModelInfo(vision_encoder_attr="model.visual", language_model_attr="model.language_model"), + "qwen3_vl_moe": VLMModelInfo(vision_encoder_attr="model.visual", language_model_attr="model.language_model"), +} + +# Text-only default +DEFAULT_LAYER_PREFIX = "model.layers." + + +# --------------------------------------------------------------------------- +# Model component access +# --------------------------------------------------------------------------- + + +def get_vision_encoder(model: nn.Module, override: str | None = None) -> nn.Module | None: + """Get the vision encoder module. + + Checks: config override -> registry. Returns None if not found. + Raises ValueError on a bad config override. + """ + if override is not None: + result = _resolve_attr(model, override) + if result is None: + raise ValueError(f"vlm.vision_encoder_attr='{override}' does not resolve on this model") + return result + + info = _get_model_info(model) + if info is not None: + return _resolve_attr(model, info.vision_encoder_attr) + + return None + + +def get_language_model(model: nn.Module, override: str | None = None) -> nn.Module: + """Get the language model module (the part with transformer layers). + + Checks: config override -> registry -> model.model (text-only default). + Raises ValueError on a bad config override. + """ + if override is not None: + result = _resolve_attr(model, override) + if result is None: + raise ValueError(f"vlm.language_model_attr='{override}' does not resolve on this model") + return result + + info = _get_model_info(model) + if info is not None: + result = _resolve_attr(model, info.language_model_attr) + if result is not None: + return result + + # Text-only models: language model is directly at model.model + return model.model + + +def is_vlm_architecture(model_config: PretrainedConfig) -> bool: + """Check if the model config belongs to a known VLM architecture.""" + return _get_model_info_from_config(model_config) is not None + + +def get_layer_prefix(model_config: PretrainedConfig, override: str | None = None) -> str: + """Return the weight key prefix for language model layers. + + Derived from language_model_attr + '.layers.' for registered VLMs, + or 'model.layers.' for text-only / unknown models. + """ + if override is not None: + return override + info = _get_model_info_from_config(model_config) + if info is not None: + return info.language_model_attr + ".layers." + return DEFAULT_LAYER_PREFIX + + +# --------------------------------------------------------------------------- +# Internal +# --------------------------------------------------------------------------- + + +def _get_model_info(model: nn.Module) -> VLMModelInfo | None: + model_type = getattr(getattr(model, "config", None), "model_type", None) + return VLM_REGISTRY.get(model_type) if model_type else None + + +def _get_model_info_from_config(model_config: PretrainedConfig) -> VLMModelInfo | None: + model_type = getattr(model_config, "model_type", None) + return VLM_REGISTRY.get(model_type) if model_type else None + + +def _resolve_attr(obj, dotted_path: str): + """Resolve a dotted attribute path like 'model.visual' on an object.""" + for part in dotted_path.split("."): + obj = getattr(obj, part, None) + if obj is None: + return None + return obj diff --git a/tests/conftest.py b/tests/conftest.py index b23c775671a..413aeba3d3e 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -1,5 +1,6 @@ import os import shutil +import signal import socket import subprocess from pathlib import Path @@ -9,6 +10,7 @@ from prime_rl.trainer.world import reset_world from prime_rl.utils.logger import reset_logger, setup_logger +from prime_rl.utils.process import cleanup_process @pytest.fixture(autouse=True) @@ -116,11 +118,11 @@ def _run_process(command: Command, env: Environment = {}, timeout: int | None = try: process.wait(timeout=timeout) except subprocess.TimeoutExpired: - process.terminate() + cleanup_process(process.pid, signal.SIGTERM) try: - process.wait(timeout=10) # Give it 10 seconds to terminate gracefully + process.wait(timeout=30) except subprocess.TimeoutExpired: - process.kill() + cleanup_process(process.pid, signal.SIGKILL) process.wait() return ProcessResult(process) diff --git a/tests/integration/test_alphabet_sort.py b/tests/integration/test_alphabet_sort.py new file mode 100644 index 00000000000..2e79d78a9d2 --- /dev/null +++ b/tests/integration/test_alphabet_sort.py @@ -0,0 +1,68 @@ +from pathlib import Path +from typing import Callable + +import pytest + +from tests.conftest import ProcessResult +from tests.utils import check_no_error, check_reward_goes_up, check_reward_in_range, strip_escape_codes + +pytestmark = [pytest.mark.gpu, pytest.mark.slow] + + +TIMEOUT = 900 # 15 minutes + + +@pytest.fixture(scope="module") +def rl_output_dir(output_dir: Path) -> Path: + rl_dir = output_dir / "alphabet_sort_start" + rl_dir.mkdir(parents=True, exist_ok=True) + return rl_dir + + +@pytest.fixture(scope="module") +def wandb_name(branch_name: str) -> str: + """Fixture for W&B name for alphabet sort RL CI integration tests.""" + return f"test-alphabet-sort-{branch_name}" + + +@pytest.fixture(scope="module") +def rl_process( + run_process: Callable[..., ProcessResult], + rl_output_dir: Path, + wandb_project: str, + wandb_name: str, +) -> ProcessResult: + cmd = [ + "uv", + "run", + "rl", + "@", + "configs/ci/integration/alphabet_sort/start.toml", + "--wandb.project", + wandb_project, + "--wandb.name", + wandb_name, + "--output-dir", + rl_output_dir.as_posix(), + ] + return run_process(cmd, timeout=TIMEOUT) + + +@pytest.fixture(scope="module") +def test_no_error(rl_process: ProcessResult, rl_output_dir: Path): + """Tests that the RL process does not fail.""" + check_no_error(rl_process, rl_output_dir) + + +def test_reward_goes_up(rl_process: ProcessResult, test_no_error, rl_output_dir: Path): + """Tests that the reward goes up in the RL process.""" + with open(rl_output_dir / "logs" / "orchestrator.log", "r") as f: + orchestrator_stdout = strip_escape_codes(f.read()).splitlines() + check_reward_goes_up(orchestrator_stdout) + + +def test_reward_in_range(rl_process: ProcessResult, test_no_error, rl_output_dir: Path): + """Tests that the reward is in range in the RL process.""" + with open(rl_output_dir / "logs" / "orchestrator.log", "r") as f: + orchestrator_stdout = strip_escape_codes(f.read()).splitlines() + check_reward_in_range(orchestrator_stdout, min_threshold=0.05) diff --git a/tests/integration/test_benchmark_regression.py b/tests/integration/test_benchmark_regression.py new file mode 100644 index 00000000000..5ec0bb27d0c --- /dev/null +++ b/tests/integration/test_benchmark_regression.py @@ -0,0 +1,283 @@ +""" +Integration test for benchmark regression. + +This test runs the benchmark and compares results against a baseline to ensure: +- Peak memory usage is exactly the same +- Other metrics (mfu, throughput, step_time) are within 5% margin +""" + +from __future__ import annotations + +import json +from pathlib import Path +from typing import Callable + +import pytest + +from tests.conftest import ProcessResult + +pytestmark = [pytest.mark.gpu, pytest.mark.slow] + +TIMEOUT = 15 * 60 # 15 minutes +# Note: I would prefer 5% but massed compute A600s seem to be slower than hyperstack A600s (which is the baseline) +METRIC_TOLERANCE = 0.1 # 10% tolerance for mfu, throughput, step_time +MEMORY_TOLERANCE = 0.01 # 1% tolerance for peak memory +TOKEN_CHUNK_SIZE = "1024" + +# Baseline files for the Qwen3-0.6B RL benchmark +BASELINE_FILE_1GPU = Path( + "benchmarks/baselines/benchmark-1xa6000-Qwen--Qwen3-0.6B-rl-full-1gpu-Recompute-flash_attention_2-65536-cp1-ep1.json" +) +BASELINE_FILE_4GPU = Path( + "benchmarks/baselines/benchmark-4xa6000-Qwen--Qwen3-0.6B-rl-full-4gpu-Recompute-flash_attention_2-65536-cp1-ep1.json" +) + + +@pytest.fixture(scope="module") +def baseline_metrics_1gpu() -> dict: + """Load baseline metrics for 1-GPU benchmark.""" + with open(BASELINE_FILE_1GPU) as f: + baseline = json.load(f) + return baseline["metrics"] + + +@pytest.fixture(scope="module") +def baseline_metrics_4gpu() -> dict: + """Load baseline metrics for 4-GPU benchmark.""" + with open(BASELINE_FILE_4GPU) as f: + baseline = json.load(f) + return baseline["metrics"] + + +@pytest.fixture(scope="module") +def benchmark_output_file_1gpu(tmp_path_factory: pytest.TempPathFactory) -> Path: + """Output file path for 1-GPU benchmark results.""" + tmp_dir = tmp_path_factory.mktemp("benchmark_1gpu") + return tmp_dir / "benchmark_result.json" + + +@pytest.fixture(scope="module") +def benchmark_output_file_4gpu(tmp_path_factory: pytest.TempPathFactory) -> Path: + """Output file path for 4-GPU benchmark results.""" + tmp_dir = tmp_path_factory.mktemp("benchmark_4gpu") + return tmp_dir / "benchmark_result.json" + + +@pytest.fixture(scope="module") +def benchmark_process_1gpu( + run_process: Callable[..., ProcessResult], + benchmark_output_file_1gpu: Path, +) -> ProcessResult: + """Run the 1-GPU benchmark and return the process result.""" + cmd = [ + "uv", + "run", + "python", + "benchmarks/scripts/run_single_benchmark.py", + "--type", + "rl", + "--num-gpus", + "1", + "--model-name", + "Qwen/Qwen3-0.6B", + "--seq-len", + "65536", + "--ac", + "Recompute", + "--attention", + "flash_attention_2", + "--output", + str(benchmark_output_file_1gpu), + "--timeout", + str(TIMEOUT - 5 * 60), # Leave 5 min buffer + "--fused-lm-head-token-chunk-size", + TOKEN_CHUNK_SIZE, + ] + return run_process(cmd, timeout=TIMEOUT, env={"PYTORCH_CUDA_ALLOC_CONF": "expandable_segments:True"}) + + +@pytest.fixture(scope="module") +def benchmark_process_4gpu( + run_process: Callable[..., ProcessResult], + benchmark_output_file_4gpu: Path, +) -> ProcessResult: + """Run the 4-GPU benchmark and return the process result.""" + cmd = [ + "uv", + "run", + "python", + "benchmarks/scripts/run_single_benchmark.py", + "--type", + "rl", + "--num-gpus", + "4", + "--model-name", + "Qwen/Qwen3-0.6B", + "--seq-len", + "65536", + "--ac", + "Recompute", + "--attention", + "flash_attention_2", + "--output", + str(benchmark_output_file_4gpu), + "--timeout", + str(TIMEOUT - 5 * 60), # Leave 5 min buffer + "--fused-lm-head-token-chunk-size", + TOKEN_CHUNK_SIZE, + ] + return run_process(cmd, timeout=TIMEOUT, env={"PYTORCH_CUDA_ALLOC_CONF": "expandable_segments:True"}) + + +@pytest.fixture(scope="module") +def benchmark_metrics_1gpu(benchmark_process_1gpu: ProcessResult, benchmark_output_file_1gpu: Path) -> dict: + """Load the 1-GPU benchmark metrics from the output file.""" + assert benchmark_process_1gpu.returncode == 0, ( + f"Benchmark process failed with return code {benchmark_process_1gpu.returncode}" + ) + assert benchmark_output_file_1gpu.exists(), f"Benchmark output file not found at {benchmark_output_file_1gpu}" + with open(benchmark_output_file_1gpu) as f: + result = json.load(f) + assert result.get("config", {}).get("success", False), ( + f"Benchmark did not succeed: {result.get('config', {}).get('error_reason', 'unknown error')}" + ) + print("=== Benchmark metrics (1-GPU) ===") + print(result["metrics"]) + return result["metrics"] + + +@pytest.fixture(scope="module") +def benchmark_metrics_4gpu(benchmark_process_4gpu: ProcessResult, benchmark_output_file_4gpu: Path) -> dict: + """Load the 4-GPU benchmark metrics from the output file.""" + assert benchmark_process_4gpu.returncode == 0, ( + f"Benchmark process failed with return code {benchmark_process_4gpu.returncode}" + ) + assert benchmark_output_file_4gpu.exists(), f"Benchmark output file not found at {benchmark_output_file_4gpu}" + with open(benchmark_output_file_4gpu) as f: + result = json.load(f) + assert result.get("config", {}).get("success", False), ( + f"Benchmark did not succeed: {result.get('config', {}).get('error_reason', 'unknown error')}" + ) + print("=== Benchmark metrics (4-GPU) ===") + print(result["metrics"]) + return result["metrics"] + + +# ============================================================================= +# 1-GPU Benchmark Tests +# ============================================================================= + + +def test_peak_memory_within_tolerance_1gpu(benchmark_metrics_1gpu: dict, baseline_metrics_1gpu: dict): + """Test that peak memory usage does not exceed the baseline by more than 1% (1-GPU).""" + actual_memory = benchmark_metrics_1gpu["peak_memory"]["gib"] + expected_memory = baseline_metrics_1gpu["peak_memory"]["gib"] + + upper_bound = expected_memory * (1 + MEMORY_TOLERANCE) + + assert actual_memory <= upper_bound, ( + f"Peak memory regression! Expected at most {expected_memory:.4f} GiB + 1%, got {actual_memory:.4f} GiB. " + f"Upper bound: {upper_bound:.4f} GiB" + ) + + +def test_mfu_within_tolerance_1gpu(benchmark_metrics_1gpu: dict, baseline_metrics_1gpu: dict): + """Test that MFU (Model FLOPS Utilization) is within 5% of baseline (1-GPU).""" + actual_mfu = benchmark_metrics_1gpu["mfu"]["mean"] + expected_mfu = baseline_metrics_1gpu["mfu"]["mean"] + + lower_bound = expected_mfu * (1 - METRIC_TOLERANCE) + upper_bound = expected_mfu * (1 + METRIC_TOLERANCE) + + assert lower_bound <= actual_mfu <= upper_bound, ( + f"MFU out of tolerance! Expected {expected_mfu:.4f} ± 5%, got {actual_mfu:.4f}. " + f"Acceptable range: [{lower_bound:.4f}, {upper_bound:.4f}]" + ) + + +def test_throughput_within_tolerance_1gpu(benchmark_metrics_1gpu: dict, baseline_metrics_1gpu: dict): + """Test that throughput is within 5% of baseline (1-GPU).""" + actual_throughput = benchmark_metrics_1gpu["throughput"]["mean"] + expected_throughput = baseline_metrics_1gpu["throughput"]["mean"] + + lower_bound = expected_throughput * (1 - METRIC_TOLERANCE) + upper_bound = expected_throughput * (1 + METRIC_TOLERANCE) + + assert lower_bound <= actual_throughput <= upper_bound, ( + f"Throughput out of tolerance! Expected {expected_throughput:.4f} ± 5%, got {actual_throughput:.4f}. " + f"Acceptable range: [{lower_bound:.4f}, {upper_bound:.4f}]" + ) + + +def test_step_time_within_tolerance_1gpu(benchmark_metrics_1gpu: dict, baseline_metrics_1gpu: dict): + """Test that step time is within 5% of baseline (1-GPU).""" + actual_step_time = benchmark_metrics_1gpu["step_time"]["mean"] + expected_step_time = baseline_metrics_1gpu["step_time"]["mean"] + + lower_bound = expected_step_time * (1 - METRIC_TOLERANCE) + upper_bound = expected_step_time * (1 + METRIC_TOLERANCE) + + assert lower_bound <= actual_step_time <= upper_bound, ( + f"Step time out of tolerance! Expected {expected_step_time:.4f} ± 5%, got {actual_step_time:.4f}. " + f"Acceptable range: [{lower_bound:.4f}, {upper_bound:.4f}]" + ) + + +# ============================================================================= +# 4-GPU Benchmark Tests +# ============================================================================= + + +def test_peak_memory_within_tolerance_4gpu(benchmark_metrics_4gpu: dict, baseline_metrics_4gpu: dict): + """Test that peak memory usage does not exceed the baseline by more than 1% (4-GPU).""" + actual_memory = benchmark_metrics_4gpu["peak_memory"]["gib"] + expected_memory = baseline_metrics_4gpu["peak_memory"]["gib"] + + upper_bound = expected_memory * (1 + MEMORY_TOLERANCE) + + assert actual_memory <= upper_bound, ( + f"Peak memory regression! Expected at most {expected_memory:.4f} GiB + 1%, got {actual_memory:.4f} GiB. " + f"Upper bound: {upper_bound:.4f} GiB" + ) + + +def test_mfu_within_tolerance_4gpu(benchmark_metrics_4gpu: dict, baseline_metrics_4gpu: dict): + """Test that MFU (Model FLOPS Utilization) is within 5% of baseline (4-GPU).""" + actual_mfu = benchmark_metrics_4gpu["mfu"]["mean"] + expected_mfu = baseline_metrics_4gpu["mfu"]["mean"] + + lower_bound = expected_mfu * (1 - METRIC_TOLERANCE) + upper_bound = expected_mfu * (1 + METRIC_TOLERANCE) + + assert lower_bound <= actual_mfu <= upper_bound, ( + f"MFU out of tolerance! Expected {expected_mfu:.4f} ± 5%, got {actual_mfu:.4f}. " + f"Acceptable range: [{lower_bound:.4f}, {upper_bound:.4f}]" + ) + + +def test_throughput_within_tolerance_4gpu(benchmark_metrics_4gpu: dict, baseline_metrics_4gpu: dict): + """Test that throughput is within 5% of baseline (4-GPU).""" + actual_throughput = benchmark_metrics_4gpu["throughput"]["mean"] + expected_throughput = baseline_metrics_4gpu["throughput"]["mean"] + + lower_bound = expected_throughput * (1 - METRIC_TOLERANCE) + upper_bound = expected_throughput * (1 + METRIC_TOLERANCE) + + assert lower_bound <= actual_throughput <= upper_bound, ( + f"Throughput out of tolerance! Expected {expected_throughput:.4f} ± 5%, got {actual_throughput:.4f}. " + f"Acceptable range: [{lower_bound:.4f}, {upper_bound:.4f}]" + ) + + +def test_step_time_within_tolerance_4gpu(benchmark_metrics_4gpu: dict, baseline_metrics_4gpu: dict): + """Test that step time is within 5% of baseline (4-GPU).""" + actual_step_time = benchmark_metrics_4gpu["step_time"]["mean"] + expected_step_time = baseline_metrics_4gpu["step_time"]["mean"] + + lower_bound = expected_step_time * (1 - METRIC_TOLERANCE) + upper_bound = expected_step_time * (1 + METRIC_TOLERANCE) + + assert lower_bound <= actual_step_time <= upper_bound, ( + f"Step time out of tolerance! Expected {expected_step_time:.4f} ± 5%, got {actual_step_time:.4f}. " + f"Acceptable range: [{lower_bound:.4f}, {upper_bound:.4f}]" + ) diff --git a/tests/integration/test_eval.py b/tests/integration/test_eval.py deleted file mode 100644 index 77091142017..00000000000 --- a/tests/integration/test_eval.py +++ /dev/null @@ -1,52 +0,0 @@ -from pathlib import Path -from typing import Callable - -import pytest - -from tests.conftest import Command, Environment, ProcessResult - -pytestmark = [pytest.mark.slow] - - -@pytest.fixture(scope="module") -def single_env_eval_process( - run_process: Callable[[Command, Environment], ProcessResult], output_dir: Path -) -> ProcessResult: - """Fixture for running single-env eval CI integration test""" - cmd = [ - "uv", - "run", - "eval", - "@", - "configs/ci/integration/eval/single_env.toml", - "--output-dir", - output_dir.as_posix(), - ] - return run_process(cmd, {}) - - -@pytest.fixture(scope="module") -def multi_env_eval_process( - run_process: Callable[[Command, Environment], ProcessResult], output_dir: Path -) -> ProcessResult: - """Fixture for running multi-env eval CI integration test""" - cmd = [ - "uv", - "run", - "eval", - "@", - "configs/ci/integration/eval/multi_env.toml", - "--output-dir", - output_dir.as_posix(), - ] - return run_process(cmd, {}) - - -def test_no_error_single_env(single_env_eval_process: ProcessResult): - """Tests that the single environment eval process does not fail.""" - assert single_env_eval_process.returncode == 0, f"Process has non-zero return code ({single_env_eval_process})" - - -def test_no_error_multi_env(multi_env_eval_process: ProcessResult): - """Tests that the multi environment eval process does not fail.""" - assert multi_env_eval_process.returncode == 0, f"Process has non-zero return code ({multi_env_eval_process})" diff --git a/tests/integration/test_rl.py b/tests/integration/test_rl.py index 8fa1394378e..54f3508d6b2 100644 --- a/tests/integration/test_rl.py +++ b/tests/integration/test_rl.py @@ -1,11 +1,10 @@ -from functools import partial from pathlib import Path from typing import Callable import pytest from tests.conftest import ProcessResult -from tests.utils import check_no_error, check_number_goes_up_or_down, check_number_in_range, strip_escape_codes +from tests.utils import check_no_error, check_reward_goes_up, check_reward_in_range, strip_escape_codes pytestmark = [pytest.mark.gpu, pytest.mark.slow] @@ -32,6 +31,7 @@ def rl_process( "rl", "@", "configs/ci/integration/rl/start.toml", + "--clean-output-dir", "--wandb.project", wandb_project, "--wandb.name", @@ -70,10 +70,6 @@ def rl_resume_process( return run_process(cmd, timeout=TIMEOUT) -check_reward_goes_up = partial(check_number_goes_up_or_down, go_up=True, pattern=r"Reward:\s*(\d+\.\d{4})") -check_reward_in_range = partial(check_number_in_range, pattern=r"Reward:\s*(\d+\.\d{4})") - - @pytest.fixture(scope="module") def test_no_error(rl_process: ProcessResult, output_dir: Path): """Tests that the RL process does not fail.""" @@ -82,14 +78,14 @@ def test_no_error(rl_process: ProcessResult, output_dir: Path): def test_reward_goes_up(rl_process: ProcessResult, test_no_error, output_dir: Path): """Tests that the reward goes up in the RL process""" - with open(output_dir / "logs" / "orchestrator.stdout", "r") as f: + with open(output_dir / "logs" / "orchestrator.log", "r") as f: orchestrator_stdout = strip_escape_codes(f.read()).splitlines() check_reward_goes_up(orchestrator_stdout) def test_reward_in_range(rl_process: ProcessResult, test_no_error, output_dir: Path): """Tests that the reward is in range in the RL process""" - with open(output_dir / "logs" / "orchestrator.stdout", "r") as f: + with open(output_dir / "logs" / "orchestrator.log", "r") as f: orchestrator_stdout = strip_escape_codes(f.read()).splitlines() check_reward_in_range(orchestrator_stdout, min_threshold=0.65) @@ -102,6 +98,6 @@ def test_no_error_resume(rl_resume_process: ProcessResult, output_dir: Path): def test_reward_in_range_resume(rl_resume_process: ProcessResult, test_no_error_resume, output_dir: Path): """Tests that the reward is in range in the RL resume process""" - with open(output_dir / "logs" / "orchestrator.stdout", "r") as f: + with open(output_dir / "logs" / "orchestrator.log", "r") as f: orchestrator_stdout = strip_escape_codes(f.read()).splitlines() check_reward_in_range(orchestrator_stdout, min_threshold=0.65) diff --git a/tests/integration/test_rl_lora.py b/tests/integration/test_rl_lora.py index bf284f28791..45d86f8a193 100644 --- a/tests/integration/test_rl_lora.py +++ b/tests/integration/test_rl_lora.py @@ -1,11 +1,10 @@ -from functools import partial from pathlib import Path from typing import Callable import pytest from tests.conftest import ProcessResult -from tests.utils import check_no_error, check_number_goes_up_or_down, check_number_in_range, strip_escape_codes +from tests.utils import check_no_error, check_reward_goes_up, check_reward_in_range, strip_escape_codes pytestmark = [pytest.mark.gpu, pytest.mark.slow] @@ -32,6 +31,7 @@ def rl_process( "rl", "@", "configs/ci/integration/rl_lora/start.toml", + "--clean-output-dir", "--wandb.project", wandb_project, "--wandb.name", @@ -70,10 +70,6 @@ def rl_resume_process( return run_process(cmd, env={"VLLM_ALLOW_RUNTIME_LORA_UPDATING": "True"}, timeout=TIMEOUT) -check_reward_goes_up = partial(check_number_goes_up_or_down, go_up=True, pattern=r"Reward:\s*(\d+\.\d{4})") -check_reward_in_range = partial(check_number_in_range, pattern=r"Reward:\s*(\d+\.\d{4})") - - @pytest.fixture(scope="module") def test_no_error(rl_process: ProcessResult, output_dir: Path): """Tests that the RL process does not fail.""" @@ -82,14 +78,14 @@ def test_no_error(rl_process: ProcessResult, output_dir: Path): def test_reward_goes_up(rl_process: ProcessResult, test_no_error, output_dir: Path): """Tests that the reward goes up in the RL process""" - with open(output_dir / "logs" / "orchestrator.stdout", "r") as f: + with open(output_dir / "logs" / "orchestrator.log", "r") as f: orchestrator_stdout = strip_escape_codes(f.read()).splitlines() check_reward_goes_up(orchestrator_stdout) def test_reward_in_range(rl_process: ProcessResult, test_no_error, output_dir: Path): """Tests that the reward is in range in the RL process""" - with open(output_dir / "logs" / "orchestrator.stdout", "r") as f: + with open(output_dir / "logs" / "orchestrator.log", "r") as f: orchestrator_stdout = strip_escape_codes(f.read()).splitlines() check_reward_in_range(orchestrator_stdout, min_threshold=0.65) @@ -102,6 +98,6 @@ def test_no_error_resume(rl_resume_process: ProcessResult, output_dir: Path): def test_reward_in_range_resume(rl_resume_process: ProcessResult, test_no_error_resume, output_dir: Path): """Tests that the reward is in range in the RL resume process""" - with open(output_dir / "logs" / "orchestrator.stdout", "r") as f: + with open(output_dir / "logs" / "orchestrator.log", "r") as f: orchestrator_stdout = strip_escape_codes(f.read()).splitlines() check_reward_in_range(orchestrator_stdout, min_threshold=0.65) diff --git a/tests/integration/test_rl_moe.py b/tests/integration/test_rl_moe.py new file mode 100644 index 00000000000..4614d08c250 --- /dev/null +++ b/tests/integration/test_rl_moe.py @@ -0,0 +1,104 @@ +from pathlib import Path +from typing import Callable + +import pytest + +from tests.conftest import ProcessResult +from tests.utils import check_no_error + +pytestmark = [pytest.mark.gpu, pytest.mark.slow] + +TIMEOUT = 900 # 15 minutes + + +@pytest.fixture(scope="module") +def wandb_name(branch_name: str) -> str: + return f"test-rl-moe-{branch_name}" + + +# --- MoE with HF impl (default) --- + + +@pytest.fixture(scope="module") +def moe_hf_output_dir(output_dir: Path) -> Path: + d = output_dir / "rl_moe_hf" + d.mkdir(parents=True, exist_ok=True) + return d + + +@pytest.fixture(scope="module") +def moe_hf_process( + run_process: Callable[..., ProcessResult], + moe_hf_output_dir: Path, + wandb_project: str, + wandb_name: str, +) -> ProcessResult: + cmd = [ + "uv", + "run", + "rl", + "@", + "configs/ci/integration/rl_moe/start.toml", + "--trainer.model.impl", + "hf", + "--wandb.project", + wandb_project, + "--wandb.name", + f"{wandb_name}-hf", + "--output-dir", + moe_hf_output_dir.as_posix(), + ] + return run_process(cmd, timeout=TIMEOUT) + + +@pytest.fixture(scope="module") +def test_no_error_hf(moe_hf_process: ProcessResult, moe_hf_output_dir: Path): + check_no_error(moe_hf_process, moe_hf_output_dir) + + +def test_moe_hf_runs(moe_hf_process: ProcessResult, test_no_error_hf): + """MoE RL with HF model impl completes without error.""" + + +# --- MoE with custom impl --- + + +@pytest.fixture(scope="module") +def moe_custom_output_dir(output_dir: Path) -> Path: + d = output_dir / "rl_moe_custom" + d.mkdir(parents=True, exist_ok=True) + return d + + +@pytest.fixture(scope="module") +def moe_custom_process( + run_process: Callable[..., ProcessResult], + moe_custom_output_dir: Path, + wandb_project: str, + wandb_name: str, +) -> ProcessResult: + cmd = [ + "uv", + "run", + "rl", + "@", + "configs/ci/integration/rl_moe/start.toml", + "--trainer.model.impl", + "custom", + "--wandb.project", + wandb_project, + "--wandb.name", + f"{wandb_name}-custom", + "--output-dir", + moe_custom_output_dir.as_posix(), + ] + return run_process(cmd, timeout=TIMEOUT) + + +@pytest.fixture(scope="module") +def test_no_error_custom(moe_custom_process: ProcessResult, moe_custom_output_dir: Path): + check_no_error(moe_custom_process, moe_custom_output_dir) + + +def test_moe_custom_runs(moe_custom_process: ProcessResult, test_no_error_custom): + """MoE RL with custom model impl completes without error.""" diff --git a/tests/integration/test_rl_multi_run_lora.py b/tests/integration/test_rl_multi_run_lora.py index cf98069d7d2..7244151b28c 100644 --- a/tests/integration/test_rl_multi_run_lora.py +++ b/tests/integration/test_rl_multi_run_lora.py @@ -5,14 +5,13 @@ import signal import subprocess import time -from functools import partial from pathlib import Path from typing import Generator import pytest from tests.conftest import ProcessResult -from tests.utils import check_number_goes_up_or_down, check_number_in_range, strip_escape_codes +from tests.utils import check_reward_goes_up, check_reward_in_range, strip_escape_codes pytestmark = [pytest.mark.gpu, pytest.mark.slow] @@ -92,21 +91,38 @@ def wandb_name(branch_name: str) -> str: return f"test-rl-multi-run-{branch_name}" +INFERENCE_PORTS = [8000, 8001] +INFERENCE_BASE_URLS = [f"http://localhost:{port}/v1" for port in INFERENCE_PORTS] + + def start_inference_and_trainer( log_dir: Path, output_dir: Path, wandb_project: str, wandb_name: str -) -> tuple[subprocess.Popen, subprocess.Popen]: - # Start inference server - inference_log = log_dir / "inference.stdout" - with open(inference_log, "w") as f: - inference_proc = subprocess.Popen( - ["uv", "run", "inference", "@", "configs/ci/integration/rl_multi_run/inference.toml"], - stdout=f, - stderr=f, - env={**os.environ, "VLLM_ALLOW_RUNTIME_LORA_UPDATING": "True", "CUDA_VISIBLE_DEVICES": "0"}, - ) +) -> tuple[subprocess.Popen, list[subprocess.Popen]]: + # Start inference servers (one per GPU on ports 8000 and 8001) + inference_procs: list[subprocess.Popen] = [] + inference_logs: list[Path] = [] + for i, port in enumerate(INFERENCE_PORTS): + inference_log = log_dir / f"inference_{i}.log" + inference_logs.append(inference_log) + with open(inference_log, "w") as f: + inference_proc = subprocess.Popen( + [ + "uv", + "run", + "inference", + "@", + "configs/ci/integration/rl_multi_run/inference.toml", + "--server.port", + str(port), + ], + stdout=f, + stderr=f, + env={**os.environ, "VLLM_ALLOW_RUNTIME_LORA_UPDATING": "True", "CUDA_VISIBLE_DEVICES": str(i)}, + ) + inference_procs.append(inference_proc) - # Start trainer - trainer_log = log_dir / "trainer.stdout" + # Start trainer with 2 GPUs + trainer_log = log_dir / "trainer.log" with open(trainer_log, "w") as f: trainer_proc = subprocess.Popen( [ @@ -114,7 +130,7 @@ def start_inference_and_trainer( "run", "torchrun", "--nproc-per-node", - "1", + "2", "-m", "prime_rl.trainer.rl.train", "@", @@ -130,22 +146,24 @@ def start_inference_and_trainer( ], stdout=f, stderr=f, - env={**os.environ, "CUDA_VISIBLE_DEVICES": "1"}, + env={**os.environ, "CUDA_VISIBLE_DEVICES": "2,3"}, ) - # Wait for inference to be ready + # Wait for all inference servers to be ready ready_indicators = ["Application startup complete", "Uvicorn running on", "Started server process"] - start_time = time.time() - while time.time() - start_time < 300: - if inference_log.exists(): - content = inference_log.read_text() - if any(ind in content for ind in ready_indicators): - break - time.sleep(2) - else: - trainer_proc.terminate() - inference_proc.terminate() - pytest.fail("Inference server did not start in time") + for i, inference_log in enumerate(inference_logs): + start_time = time.time() + while time.time() - start_time < 300: + if inference_log.exists(): + content = inference_log.read_text() + if any(ind in content for ind in ready_indicators): + break + time.sleep(2) + else: + trainer_proc.terminate() + for proc in inference_procs: + proc.terminate() + pytest.fail(f"Inference server {i} did not start in time") # Wait for trainer to be ready ready_indicators = ["Starting training loop"] @@ -158,10 +176,11 @@ def start_inference_and_trainer( time.sleep(2) else: trainer_proc.terminate() - inference_proc.terminate() + for proc in inference_procs: + proc.terminate() pytest.fail("Trainer did not start in time") - return trainer_proc, inference_proc + return trainer_proc, inference_procs def start_orchestrator( @@ -175,25 +194,30 @@ def start_orchestrator( orch_log_dir = run_dir / "logs" orch_log_dir.mkdir(parents=True, exist_ok=True) - with open(orch_log_dir / "orchestrator.stdout", "w") as f: + # Build command with multiple inference server URLs + cmd = [ + "uv", + "run", + "orchestrator", + "@", + "configs/ci/integration/rl_multi_run/orchestrator.toml", + "--output-dir", + run_dir.as_posix(), + "--max-steps", + str(max_steps), + "--model.lora.name", + name, + "--wandb.project", + wandb_project, + "--wandb.name", + f"{wandb_name}-{proc_name}", + ] + cmd.append("--client.base-url") + cmd.extend(INFERENCE_BASE_URLS) + + with open(orch_log_dir / "orchestrator.log", "w") as f: proc = subprocess.Popen( - [ - "uv", - "run", - "orchestrator", - "@", - "configs/ci/integration/rl_multi_run/orchestrator.toml", - "--output-dir", - run_dir.as_posix(), - "--max-steps", - str(max_steps), - "--model.lora.name", - name, - "--wandb.project", - wandb_project, - "--wandb.name", - f"{wandb_name}-{proc_name}", - ], + cmd, stdout=f, stderr=f, ) @@ -213,9 +237,9 @@ def multi_run_result( processes: list[subprocess.Popen] = [] - trainer_proc, inference_proc = start_inference_and_trainer(log_dir, output_dir, wandb_project, wandb_name) + trainer_proc, inference_procs = start_inference_and_trainer(log_dir, output_dir, wandb_project, wandb_name) processes.append(trainer_proc) - processes.append(inference_proc) + processes.extend(inference_procs) # =========================================== # Start alpha, beta, and gamma orchestrators @@ -231,7 +255,7 @@ def multi_run_result( # Kill alpha orchestrator once it is past step 10 # ------------------------------------------------ # There is a checkpoint at step 10, so we need to wait for step 11 - killed_log = output_dir / "run_alpha" / "logs" / "orchestrator.stdout" + killed_log = output_dir / "run_alpha" / "logs" / "orchestrator.log" wait_for_log( killed_log, conditions=["Step 11", "Step 12", "Step 13"], @@ -244,7 +268,7 @@ def multi_run_result( wait_for_file(alpha_ckpt_dir / "STABLE", timeout=TIMEOUT) # Stash alpha checkpoint and logs - shutil.copy(output_dir / "run_alpha" / "logs" / "orchestrator.stdout", log_dir / "alpha_orchestrator.stdout") + shutil.copy(output_dir / "run_alpha" / "logs" / "orchestrator.log", log_dir / "alpha_orchestrator.log") shutil.copytree(alpha_ckpt_dir, tmp_path / "alpha_ckpt_step_10") print(f"Copied alpha checkpoint to {tmp_path / 'alpha_ckpt_step_10'}") @@ -268,7 +292,7 @@ def multi_run_result( # Clear beta run directory once it saves the final checkpoint # ------------------------------------------------------------ wait_for_log( - output_dir / "run_beta" / "logs" / "orchestrator.stdout", + output_dir / "run_beta" / "logs" / "orchestrator.log", conditions=["Orchestrator finished."], proc=orch_procs["beta"], poll_interval=1, @@ -277,7 +301,7 @@ def multi_run_result( run_dir = output_dir / "run_beta" beta_ckpt_dir = run_dir / "checkpoints" / "step_20" wait_for_file(beta_ckpt_dir / "STABLE", timeout=TIMEOUT) - shutil.copy(run_dir / "logs" / "orchestrator.stdout", log_dir / "beta_orchestrator.stdout") + shutil.copy(run_dir / "logs" / "orchestrator.log", log_dir / "beta_orchestrator.log") shutil.copytree(beta_ckpt_dir, tmp_path / "beta_ckpt_step_20") print(f"Copied {beta_ckpt_dir} to {tmp_path / 'beta_ckpt_step_20'}") shutil.rmtree(run_dir) @@ -298,12 +322,12 @@ def multi_run_result( # Clear gamma run directory once it finishes # ------------------------------------------- wait_for_log( - output_dir / "run_gamma" / "logs" / "orchestrator.stdout", + output_dir / "run_gamma" / "logs" / "orchestrator.log", conditions=["Orchestrator finished."], proc=orch_procs["gamma"], timeout=TIMEOUT, ) - shutil.copy(output_dir / "run_gamma" / "logs" / "orchestrator.stdout", log_dir / "gamma_orchestrator.stdout") + shutil.copy(output_dir / "run_gamma" / "logs" / "orchestrator.log", log_dir / "gamma_orchestrator.log") shutil.rmtree(output_dir / "run_gamma") # ================================================ @@ -316,9 +340,9 @@ def multi_run_result( orch_procs[name].terminate() for name in ["alpha_resume", "beta_resume"]: - src_log = output_dir / f"run_{name}" / "logs" / "orchestrator.stdout" + src_log = output_dir / f"run_{name}" / "logs" / "orchestrator.log" if src_log.exists(): - shutil.copy(src_log, log_dir / f"{name}_orchestrator.stdout") + shutil.copy(src_log, log_dir / f"{name}_orchestrator.log") # =============== # Build results @@ -339,9 +363,6 @@ def multi_run_result( p.kill() -check_reward_goes_up = partial(check_number_goes_up_or_down, go_up=True, pattern=r"Reward:\s*(\d+\.\d{4})") - - def test_remaining_orchestrators_complete( multi_run_result: dict[str, ProcessResult], output_dir: Path, @@ -353,7 +374,7 @@ def test_remaining_orchestrators_complete( if name == "alpha": # We sigtermed alpha continue if result.returncode != 0: - log_file = log_dir / f"{name}_orchestrator.stdout" + log_file = log_dir / f"{name}_orchestrator.log" if log_file.exists(): print(f"=== {name} Orchestrator Outputs ===") print(log_file.read_text()[-5000:]) @@ -369,7 +390,7 @@ def test_reward_goes_up(multi_run_result: dict[str, ProcessResult], output_dir: # The resumes are close to saturation so might not go up if "resume" in name: continue - log_file = log_dir / f"{name}_orchestrator.stdout" + log_file = log_dir / f"{name}_orchestrator.log" with open(log_file, "r") as f: lines = strip_escape_codes(f.read()).splitlines() check_reward_goes_up(lines) @@ -381,17 +402,15 @@ def test_reward_in_range(multi_run_result: dict[str, ProcessResult], output_dir: print("Test reward in range", multi_run_result.keys()) for name in multi_run_result.keys(): - log_file = log_dir / f"{name}_orchestrator.stdout" + log_file = log_dir / f"{name}_orchestrator.log" with open(log_file, "r") as f: lines = strip_escape_codes(f.read()).splitlines() if name in ["beta", "gamma"]: - check_number_in_range( - lines, step=7, min_threshold=0.2, max_threshold=0.6, pattern=r"Reward:\s*(\d+\.\d{4})" - ) - check_number_in_range(lines, min_threshold=0.65, pattern=r"Reward:\s*(\d+\.\d{4})") + check_reward_in_range(lines, step=7, min_threshold=0.2, max_threshold=0.6) + check_reward_in_range(lines, min_threshold=0.65) elif name in ["alpha_resume", "beta_resume"]: - check_number_in_range(lines, min_threshold=0.65, pattern=r"Reward:\s*(\d+\.\d{4})") + check_reward_in_range(lines, min_threshold=0.65) elif name == "alpha": # Only had 10 steps, so it's lower - check_number_in_range(lines, min_threshold=0.4, pattern=r"Reward:\s*(\d+\.\d{4})") + check_reward_in_range(lines, min_threshold=0.4) else: pytest.fail(f"Unknown orchestrator {name}") diff --git a/tests/integration/test_sft.py b/tests/integration/test_sft.py index 8e1eb2094e2..a02c4c24567 100644 --- a/tests/integration/test_sft.py +++ b/tests/integration/test_sft.py @@ -1,11 +1,10 @@ -from functools import partial from pathlib import Path from typing import Callable import pytest from tests.conftest import ProcessResult -from tests.utils import check_number_goes_up_or_down, strip_escape_codes +from tests.utils import check_loss_goes_down, strip_escape_codes pytestmark = [pytest.mark.slow, pytest.mark.gpu] @@ -29,14 +28,12 @@ def sft_process( cmd = [ "uv", "run", - "torchrun", - "--local-ranks-filter", - "0", - "--nproc-per-node", - "2", - "src/prime_rl/trainer/sft/train.py", + "sft", "@", "configs/ci/integration/sft/start.toml", + "--deployment.num-gpus", + "2", + "--clean-output-dir", "--wandb.project", wandb_project, "--wandb.name", @@ -61,14 +58,11 @@ def sft_resume_process( cmd = [ "uv", "run", - "torchrun", - "--local-ranks-filter", - "0", - "--nproc-per-node", - "2", - "src/prime_rl/trainer/sft/train.py", + "sft", "@", "configs/ci/integration/sft/resume.toml", + "--deployment.num-gpus", + "2", "--wandb.project", wandb_project, "--wandb.name", @@ -80,9 +74,6 @@ def sft_resume_process( return run_process(cmd, timeout=TIMEOUT) -check_loss_goes_down = partial(check_number_goes_up_or_down, go_up=False, pattern=r"Loss:\s*(\d+\.\d{4})") - - def test_no_error(sft_process: ProcessResult): """Tests that the SFT process does not fail.""" assert sft_process.returncode == 0, f"Process has non-zero return code ({sft_process})" @@ -90,7 +81,7 @@ def test_no_error(sft_process: ProcessResult): def test_loss_goes_down(sft_process: ProcessResult, output_dir: Path): """Tests that the loss goes down in the SFT process""" - trainer_log_path = output_dir / "logs" / "trainer" / "rank_0.log" + trainer_log_path = output_dir / "logs" / "trainer.log" print(f"Checking trainer path in {trainer_log_path}") with open(trainer_log_path, "r") as f: trainer_stdout = strip_escape_codes(f.read()).splitlines() @@ -104,7 +95,7 @@ def test_no_error_resume(sft_resume_process: ProcessResult): def test_loss_goes_down_resume(sft_resume_process: ProcessResult, output_dir: Path): """Tests that the loss goes down in the SFT resume process""" - trainer_log_path = output_dir / "logs" / "trainer" / "rank_0.log" + trainer_log_path = output_dir / "logs" / "trainer.log" print(f"Checking trainer path in {trainer_log_path}") with open(trainer_log_path, "r") as f: trainer_stdout = strip_escape_codes(f.read()).splitlines() diff --git a/tests/integration/test_sft_lora.py b/tests/integration/test_sft_lora.py new file mode 100644 index 00000000000..ffa9edc976f --- /dev/null +++ b/tests/integration/test_sft_lora.py @@ -0,0 +1,124 @@ +from pathlib import Path +from typing import Callable + +import pytest + +from prime_rl.trainer.weights import load_state_dict +from tests.conftest import ProcessResult +from tests.utils import check_loss_goes_down, strip_escape_codes + +pytestmark = [pytest.mark.slow, pytest.mark.gpu] + +TIMEOUT = 300 # 5 minutes + + +def assert_adapter_checkpoint(adapter_dir: Path) -> None: + assert (adapter_dir / "adapter_config.json").exists() + state_dict = load_state_dict(adapter_dir) + assert state_dict + assert all(".0.weight" not in key for key in state_dict) + assert any(key.endswith("lora_A.weight") for key in state_dict) + assert all(key.startswith("base_model.model.") for key in state_dict) + + +@pytest.fixture(scope="module") +def wandb_name(branch_name: str) -> str: + """Fixture for W&B name for SFT LoRA CI integration tests.""" + return f"test-sft-lora-{branch_name}" + + +@pytest.fixture(scope="module") +def sft_lora_process( + run_process: Callable[..., ProcessResult], + wandb_project: str, + wandb_name: str, + output_dir: Path, +) -> ProcessResult: + """Fixture for running SFT LoRA CI integration test""" + cmd = [ + "uv", + "run", + "sft", + "@", + "configs/ci/integration/sft_lora/start.toml", + "--deployment.num-gpus", + "2", + "--clean-output-dir", + "--wandb.project", + wandb_project, + "--wandb.name", + wandb_name, + "--output-dir", + output_dir.as_posix(), + ] + + return run_process(cmd, timeout=TIMEOUT) + + +@pytest.fixture(scope="module") +def sft_lora_resume_process( + sft_lora_process, # Resume training can only start when regular SFT LoRA process is finished + run_process: Callable[..., ProcessResult], + wandb_project: str, + wandb_name: str, + output_dir: Path, +) -> ProcessResult: + """Fixture for resuming SFT LoRA CI integration test""" + wandb_name += "-resume" + cmd = [ + "uv", + "run", + "sft", + "@", + "configs/ci/integration/sft_lora/resume.toml", + "--deployment.num-gpus", + "2", + "--wandb.project", + wandb_project, + "--wandb.name", + wandb_name, + "--output-dir", + output_dir.as_posix(), + ] + + return run_process(cmd, timeout=TIMEOUT) + + +def test_no_error(sft_lora_process: ProcessResult): + """Tests that the SFT LoRA process does not fail.""" + assert sft_lora_process.returncode == 0, f"Process has non-zero return code ({sft_lora_process})" + + +def test_loss_goes_down(sft_lora_process: ProcessResult, output_dir: Path): + """Tests that the loss goes down in the SFT LoRA process""" + trainer_log_path = output_dir / "logs" / "trainer.log" + print(f"Checking trainer path in {trainer_log_path}") + with open(trainer_log_path, "r") as f: + trainer_stdout = strip_escape_codes(f.read()).splitlines() + check_loss_goes_down(trainer_stdout) + + +def test_adapter_checkpoint_written(sft_lora_process: ProcessResult, output_dir: Path): + """Tests that the adapter checkpoint is written with valid PEFT-compatible keys.""" + adapter_dir = output_dir / "weights" / "step_10" / "lora_adapters" + assert_adapter_checkpoint(adapter_dir) + + +def test_no_error_resume(sft_lora_resume_process: ProcessResult): + """Tests that the SFT LoRA resume process does not fail.""" + assert sft_lora_resume_process.returncode == 0, f"Process has non-zero return code ({sft_lora_resume_process})" + + +def test_loss_goes_down_resume(sft_lora_resume_process: ProcessResult, output_dir: Path): + """Tests that the loss goes down in the SFT LoRA resume process""" + trainer_log_path = output_dir / "logs" / "trainer.log" + print(f"Checking trainer path in {trainer_log_path}") + with open(trainer_log_path, "r") as f: + trainer_stdout = strip_escape_codes(f.read()).splitlines() + check_loss_goes_down(trainer_stdout) + + +def test_adapter_checkpoint_written_resume(sft_lora_resume_process: ProcessResult, output_dir: Path): + """Tests that the adapter checkpoint is written after resuming with valid PEFT-compatible keys.""" + adapter_dir = output_dir / "weights" / "step_20" / "lora_adapters" + assert_adapter_checkpoint(adapter_dir) diff --git a/tests/integration/test_synthesize.py b/tests/integration/test_synthesize.py deleted file mode 100644 index a5378f48adc..00000000000 --- a/tests/integration/test_synthesize.py +++ /dev/null @@ -1,81 +0,0 @@ -from pathlib import Path -from typing import Callable - -import pytest - -from tests.conftest import Command, Environment, ProcessResult - -pytestmark = [pytest.mark.slow] - - -@pytest.fixture(scope="module") -def single_env_synthesize_process( - run_process: Callable[[Command, Environment], ProcessResult], output_dir: Path -) -> ProcessResult: - """Fixture for running single-env synthesize CI integration test""" - cmd = [ - "uv", - "run", - "synthesize", - "@", - "configs/ci/integration/synthesize/single_env.toml", - "--output-dir", - output_dir.as_posix(), - ] - return run_process(cmd, {}) - - -@pytest.fixture(scope="module") -def multi_env_synthesize_process( - run_process: Callable[[Command, Environment], ProcessResult], output_dir: Path -) -> ProcessResult: - """Fixture for running multi-env synthesize CI integration test""" - cmd = [ - "uv", - "run", - "synthesize", - "@", - "configs/ci/integration/synthesize/multi_env.toml", - "--output-dir", - output_dir.as_posix(), - ] - return run_process(cmd, {}) - - -@pytest.fixture(scope="module") -def multi_turn_tool_call_synthesize_process( - run_process: Callable[[Command, Environment], ProcessResult], - output_dir: Path, -) -> ProcessResult: - """Fixture for running multi-turn tool call synthesize CI integration test""" - cmd = [ - "uv", - "run", - "synthesize", - "@", - "configs/ci/integration/synthesize/multi_turn_tool_call.toml", - "--output-dir", - output_dir.as_posix(), - ] - return run_process(cmd, {}) - - -def test_no_error_single_env(single_env_synthesize_process: ProcessResult): - """Tests that the single environment synthesize process does not fail.""" - assert single_env_synthesize_process.returncode == 0, ( - f"Process has non-zero return code ({single_env_synthesize_process})" - ) - - -def test_no_error_multi_env(multi_env_synthesize_process: ProcessResult): - """Tests that the multi-env synthesize process does not fail.""" - assert multi_env_synthesize_process.returncode == 0, ( - f"Process has non-zero return code ({multi_env_synthesize_process})" - ) - - -def test_no_error_multi_turn_tool_call(multi_turn_tool_call_synthesize_process: ProcessResult): - """Tests that the multi-turn tool call synthesize process does not fail.""" - assert multi_turn_tool_call_synthesize_process.returncode == 0, ( - f"Process has non-zero return code ({multi_turn_tool_call_synthesize_process})" - ) diff --git a/tests/nightly/test_acereason_math.py b/tests/nightly/test_acereason_math.py index bb857378f56..57d7ad1f819 100644 --- a/tests/nightly/test_acereason_math.py +++ b/tests/nightly/test_acereason_math.py @@ -1,11 +1,10 @@ -from functools import partial from pathlib import Path from typing import Callable import pytest from tests.conftest import ProcessResult -from tests.utils import check_no_error, check_number_goes_up_or_down, strip_escape_codes +from tests.utils import check_mismatch_kl_in_range, check_no_error, check_reward_goes_up, strip_escape_codes pytestmark = [pytest.mark.gpu, pytest.mark.slow] @@ -39,7 +38,8 @@ def rl_process( return run_process(cmd) -check_reward_goes_up = partial(check_number_goes_up_or_down, go_up=True, pattern=r"Reward:\s*(\d+\.\d{4})") +MISMATCH_KL_MIN = 0.0 +MISMATCH_KL_MAX = 0.0015 @pytest.fixture(scope="module") @@ -50,6 +50,13 @@ def test_no_error(rl_process: ProcessResult, output_dir: Path): def test_reward_goes_up(rl_process: ProcessResult, test_no_error, output_dir: Path): """Tests that the reward goes up in the RL process""" - with open(output_dir / "logs" / "orchestrator.stdout", "r") as f: + with open(output_dir / "logs" / "orchestrator.log", "r") as f: orchestrator_stdout = strip_escape_codes(f.read()).splitlines() check_reward_goes_up(orchestrator_stdout) + + +def test_mismatch_kl_in_band(rl_process: ProcessResult, test_no_error, output_dir: Path): + """Tests that mismatch KL stays within the expected band.""" + with open(output_dir / "logs" / "trainer.log", "r") as f: + trainer_stdout = strip_escape_codes(f.read()).splitlines() + check_mismatch_kl_in_range(trainer_stdout, min_threshold=MISMATCH_KL_MIN, max_threshold=MISMATCH_KL_MAX) diff --git a/tests/nightly/test_alphabet_sort.py b/tests/nightly/test_alphabet_sort.py index 0c9968e8732..07fb6d895f0 100644 --- a/tests/nightly/test_alphabet_sort.py +++ b/tests/nightly/test_alphabet_sort.py @@ -1,11 +1,10 @@ -from functools import partial from pathlib import Path from typing import Callable import pytest from tests.conftest import ProcessResult -from tests.utils import check_no_error, check_number_goes_up_or_down, strip_escape_codes +from tests.utils import check_no_error, check_reward_goes_up, strip_escape_codes pytestmark = [pytest.mark.gpu, pytest.mark.slow] @@ -39,9 +38,6 @@ def rl_process( return run_process(cmd) -check_reward_goes_up = partial(check_number_goes_up_or_down, go_up=True, pattern=r"Reward:\s*(\d+\.\d{4})") - - @pytest.fixture(scope="module") def test_no_error(rl_process: ProcessResult, output_dir: Path): """Tests that the RL process does not fail.""" @@ -50,6 +46,6 @@ def test_no_error(rl_process: ProcessResult, output_dir: Path): def test_reward_goes_up(rl_process: ProcessResult, test_no_error, output_dir: Path): """Tests that the reward goes up in the RL process""" - with open(output_dir / "logs" / "orchestrator.stdout", "r") as f: + with open(output_dir / "logs" / "orchestrator.log", "r") as f: orchestrator_stdout = strip_escape_codes(f.read()).splitlines() check_reward_goes_up(orchestrator_stdout) diff --git a/tests/nightly/test_hendrycks_sanity.py b/tests/nightly/test_hendrycks_sanity.py new file mode 100644 index 00000000000..c4a843a0cfa --- /dev/null +++ b/tests/nightly/test_hendrycks_sanity.py @@ -0,0 +1,78 @@ +from pathlib import Path +from typing import Callable + +import pytest + +from tests.conftest import ProcessResult +from tests.utils import ( + check_mismatch_kl_in_range, + check_no_error, + check_reward_goes_up, + check_reward_in_range, + strip_escape_codes, +) + +pytestmark = [pytest.mark.gpu, pytest.mark.slow] + + +@pytest.fixture(scope="module") +def wandb_name(branch_name: str) -> str: + """Fixture for W&B name for RL CI integration tests.""" + return f"hendrycks-sanity-{branch_name}" + + +@pytest.fixture(scope="module") +def rl_process( + run_process: Callable[..., ProcessResult], + output_dir: Path, + wandb_project: str, + wandb_name: str, +) -> ProcessResult: + cmd = [ + "uv", + "run", + "rl", + "@", + "examples/hendrycks_sanity/rl.toml", + "--wandb.project", + wandb_project, + "--wandb.name", + wandb_name, + "--output-dir", + output_dir.as_posix(), + "--max-steps", + "1000", # do less steps to finish in time + ] + return run_process(cmd) + + +MIN_TRAIN_REWARD = 0.75 +MISMATCH_KL_MIN = 0.0 +MISMATCH_KL_MAX = 0.0005 + + +@pytest.fixture(scope="module") +def test_no_error(rl_process: ProcessResult, output_dir: Path): + """Tests that the RL process does not fail.""" + check_no_error(rl_process, output_dir) + + +def test_reward_goes_up(rl_process: ProcessResult, test_no_error, output_dir: Path): + """Tests that the train reward goes up during training""" + with open(output_dir / "logs" / "orchestrator.log", "r") as f: + orchestrator_stdout = strip_escape_codes(f.read()).splitlines() + check_reward_goes_up(orchestrator_stdout) + + +def test_reward_reaches_threshold(rl_process: ProcessResult, test_no_error, output_dir: Path): + """Tests that the train reward reaches a minimum threshold""" + with open(output_dir / "logs" / "orchestrator.log", "r") as f: + orchestrator_stdout = strip_escape_codes(f.read()).splitlines() + check_reward_in_range(orchestrator_stdout, min_threshold=MIN_TRAIN_REWARD) + + +def test_mismatch_kl_in_band(rl_process: ProcessResult, test_no_error, output_dir: Path): + """Tests that mismatch KL stays within the expected band.""" + with open(output_dir / "logs" / "trainer.log", "r") as f: + trainer_stdout = strip_escape_codes(f.read()).splitlines() + check_mismatch_kl_in_range(trainer_stdout, min_threshold=MISMATCH_KL_MIN, max_threshold=MISMATCH_KL_MAX) diff --git a/tests/nightly/test_hendrycks_math.py b/tests/nightly/test_multimodal_color_codeword.py similarity index 59% rename from tests/nightly/test_hendrycks_math.py rename to tests/nightly/test_multimodal_color_codeword.py index f9a25f2213c..f4855e0ba42 100644 --- a/tests/nightly/test_hendrycks_math.py +++ b/tests/nightly/test_multimodal_color_codeword.py @@ -1,19 +1,21 @@ -from functools import partial from pathlib import Path from typing import Callable import pytest from tests.conftest import ProcessResult -from tests.utils import check_no_error, check_number_goes_up_or_down, strip_escape_codes +from tests.utils import check_avg_reward_in_range, check_no_error, check_reward_goes_up, strip_escape_codes pytestmark = [pytest.mark.gpu, pytest.mark.slow] +REWARD_MIN_THRESHOLD = 0.85 +REWARD_AVG_LAST_N_STEPS = 7 + @pytest.fixture(scope="module") def wandb_name(branch_name: str) -> str: """Fixture for W&B name for RL CI integration tests.""" - return f"hendrycks-math-{branch_name}" + return f"multimodal-color-codeword-{branch_name}" @pytest.fixture(scope="module") @@ -28,7 +30,7 @@ def rl_process( "run", "rl", "@", - "configs/ci/nightly/hendrycks_math.toml", + "configs/ci/nightly/multimodal_color_codeword.toml", "--wandb.project", wandb_project, "--wandb.name", @@ -39,9 +41,6 @@ def rl_process( return run_process(cmd) -check_reward_goes_up = partial(check_number_goes_up_or_down, go_up=True, pattern=r"Reward:\s*(\d+\.\d{4})") - - @pytest.fixture(scope="module") def test_no_error(rl_process: ProcessResult, output_dir: Path): """Tests that the RL process does not fail.""" @@ -50,6 +49,15 @@ def test_no_error(rl_process: ProcessResult, output_dir: Path): def test_reward_goes_up(rl_process: ProcessResult, test_no_error, output_dir: Path): """Tests that the reward goes up in the RL process""" - with open(output_dir / "logs" / "orchestrator.stdout", "r") as f: + with open(output_dir / "logs" / "orchestrator.log", "r") as f: orchestrator_stdout = strip_escape_codes(f.read()).splitlines() check_reward_goes_up(orchestrator_stdout) + + +def test_reward_reaches_threshold(rl_process: ProcessResult, test_no_error, output_dir: Path): + """Tests that the average reward over the last steps exceeds the threshold""" + with open(output_dir / "logs" / "orchestrator.log", "r") as f: + orchestrator_stdout = strip_escape_codes(f.read()).splitlines() + check_avg_reward_in_range( + orchestrator_stdout, last_n_steps=REWARD_AVG_LAST_N_STEPS, min_threshold=REWARD_MIN_THRESHOLD + ) diff --git a/tests/nightly/test_reverse_text.py b/tests/nightly/test_reverse_text.py index 30f288bee1e..7e1ae082ec2 100644 --- a/tests/nightly/test_reverse_text.py +++ b/tests/nightly/test_reverse_text.py @@ -1,11 +1,10 @@ -from functools import partial from pathlib import Path from typing import Callable import pytest from tests.conftest import ProcessResult -from tests.utils import check_no_error, check_number_goes_up_or_down, check_number_in_range, strip_escape_codes +from tests.utils import check_no_error, check_reward_goes_up, check_reward_in_range, strip_escape_codes pytestmark = [pytest.mark.gpu, pytest.mark.slow] @@ -39,10 +38,6 @@ def rl_process( return run_process(cmd) -check_reward_goes_up = partial(check_number_goes_up_or_down, go_up=True, pattern=r"Reward:\s*(\d+\.\d{4})") -check_reward_in_range = partial(check_number_in_range, pattern=r"Reward:\s*(\d+\.\d{4})") - - @pytest.fixture(scope="module") def test_no_error(rl_process: ProcessResult, output_dir: Path): """Tests that the RL process does not fail.""" @@ -51,13 +46,13 @@ def test_no_error(rl_process: ProcessResult, output_dir: Path): def test_reward_goes_up(rl_process: ProcessResult, test_no_error, output_dir: Path): """Tests that the reward goes up in the RL process""" - with open(output_dir / "logs" / "orchestrator.stdout", "r") as f: + with open(output_dir / "logs" / "orchestrator.log", "r") as f: orchestrator_stdout = strip_escape_codes(f.read()).splitlines() check_reward_goes_up(orchestrator_stdout) def test_reward_reaches_threshold(rl_process: ProcessResult, test_no_error, output_dir: Path): """Tests that the reward goes up in the RL process""" - with open(output_dir / "logs" / "orchestrator.stdout", "r") as f: + with open(output_dir / "logs" / "orchestrator.log", "r") as f: orchestrator_stdout = strip_escape_codes(f.read()).splitlines() check_reward_in_range(orchestrator_stdout, min_threshold=0.65) diff --git a/tests/nightly/test_wiki_search.py b/tests/nightly/test_wiki_search.py index 264894eeb13..c9a0598ef26 100644 --- a/tests/nightly/test_wiki_search.py +++ b/tests/nightly/test_wiki_search.py @@ -1,11 +1,10 @@ -from functools import partial from pathlib import Path from typing import Callable import pytest from tests.conftest import ProcessResult -from tests.utils import check_no_error, check_number_goes_up_or_down, strip_escape_codes +from tests.utils import check_no_error, check_reward_goes_up, strip_escape_codes pytestmark = [pytest.mark.gpu, pytest.mark.slow] @@ -39,9 +38,6 @@ def rl_process( return run_process(cmd) -check_reward_goes_up = partial(check_number_goes_up_or_down, go_up=True, pattern=r"Reward:\s*(\d+\.\d{4})") - - @pytest.fixture(scope="module") def test_no_error(rl_process: ProcessResult, output_dir: Path): """Tests that the RL process does not fail.""" @@ -50,6 +46,6 @@ def test_no_error(rl_process: ProcessResult, output_dir: Path): def test_reward_goes_up(rl_process: ProcessResult, test_no_error, output_dir: Path): """Tests that the reward goes up in the RL process""" - with open(output_dir / "logs" / "orchestrator.stdout", "r") as f: + with open(output_dir / "logs" / "orchestrator.log", "r") as f: orchestrator_stdout = strip_escape_codes(f.read()).splitlines() check_reward_goes_up(orchestrator_stdout) diff --git a/tests/nightly/test_wordle.py b/tests/nightly/test_wordle.py index f9dd0d82252..30f2112a030 100644 --- a/tests/nightly/test_wordle.py +++ b/tests/nightly/test_wordle.py @@ -1,11 +1,10 @@ -from functools import partial from pathlib import Path from typing import Callable import pytest from tests.conftest import ProcessResult -from tests.utils import check_no_error, check_number_goes_up_or_down, strip_escape_codes +from tests.utils import check_no_error, check_reward_goes_up, strip_escape_codes pytestmark = [pytest.mark.gpu, pytest.mark.slow] @@ -39,9 +38,6 @@ def rl_process( return run_process(cmd) -check_reward_goes_up = partial(check_number_goes_up_or_down, go_up=True, pattern=r"Reward:\s*(\d+\.\d{4})") - - @pytest.fixture(scope="module") def test_no_error(rl_process: ProcessResult, output_dir: Path): """Tests that the RL process does not fail.""" @@ -50,6 +46,6 @@ def test_no_error(rl_process: ProcessResult, output_dir: Path): def test_reward_goes_up(rl_process: ProcessResult, test_no_error, output_dir: Path): """Tests that the reward goes up in the RL process""" - with open(output_dir / "logs" / "orchestrator.stdout", "r") as f: + with open(output_dir / "logs" / "orchestrator.log", "r") as f: orchestrator_stdout = strip_escape_codes(f.read()).splitlines() check_reward_goes_up(orchestrator_stdout) diff --git a/tests/unit/inference/test_serving_chat_with_tokens.py b/tests/unit/inference/test_serving_chat_with_tokens.py new file mode 100644 index 00000000000..8b137891791 --- /dev/null +++ b/tests/unit/inference/test_serving_chat_with_tokens.py @@ -0,0 +1 @@ + diff --git a/tests/unit/inference/test_serving_tokens.py b/tests/unit/inference/test_serving_tokens.py new file mode 100644 index 00000000000..ac5b52b3d45 --- /dev/null +++ b/tests/unit/inference/test_serving_tokens.py @@ -0,0 +1,107 @@ +"""Sanity tests for the prime-RL ``ServingTokens`` subclass. + +The full happy-path is owned upstream by vLLM 0.20's +``vllm/entrypoints/serve/disagg`` test suite. We only cover the prime-RL +deltas here: + * ``encode_routed_experts`` round-trips a numpy array as expected. + * ``PrimeRlGenerateResponseChoice`` accepts the optional field. + * The subclass attaches its overrides without monkey-patching the parent. + * ``_client_set_max_tokens`` distinguishes raw-body shapes correctly. +""" + +from __future__ import annotations + +import asyncio +import base64 + +import numpy as np + +from prime_rl.inference.vllm.serving_tokens import ( + PrimeRlGenerateResponse, + PrimeRlGenerateResponseChoice, + PrimeRlServingTokens, + _client_set_max_tokens, + encode_routed_experts, +) + + +class _FakeRawRequest: + def __init__(self, body): + self._body = body + self._raise = isinstance(body, Exception) + + async def json(self): + if self._raise: + raise self._body + return self._body + + +def test_encode_routed_experts_roundtrip(): + arr = np.array([[1, 2, 3], [4, 5, 6]], dtype=np.int32) + encoded = encode_routed_experts(arr) + + assert encoded["shape"] == [2, 3] + decoded = np.frombuffer(base64.b85decode(encoded["data"]), dtype=np.int32).reshape(encoded["shape"]) + np.testing.assert_array_equal(decoded, arr) + + +def test_routed_experts_choice_accepts_none_and_dict(): + no_re = PrimeRlGenerateResponseChoice(index=0, finish_reason="stop", token_ids=[1, 2]) + assert no_re.routed_experts is None + + encoded = encode_routed_experts(np.zeros((1, 1), dtype=np.int32)) + with_re = PrimeRlGenerateResponseChoice(index=0, finish_reason="stop", token_ids=[1], routed_experts=encoded) + assert with_re.routed_experts == encoded + + +def test_response_only_serializes_declared_fields(): + # Upstream silently drops id=/created=/model=/usage= because they're not + # declared on GenerateResponse. Our subclass adds nothing to that surface + # — it only widens the choices type — so the JSON shape stays slim. + resp = PrimeRlGenerateResponse( + request_id="gen-x", + choices=[PrimeRlGenerateResponseChoice(index=0, finish_reason="stop", token_ids=[7])], + ) + dumped = resp.model_dump() + assert set(dumped.keys()) == { + "request_id", + "choices", + "prompt_logprobs", + "kv_transfer_params", + } + assert dumped["choices"][0]["routed_experts"] is None + + +def test_subclass_inherits_serve_tokens_full_generator(): + # The subclass adds an override; make sure we didn't accidentally rebind + # ``serve_tokens`` to a parent attribute via __dict__-update tricks later. + assert ( + PrimeRlServingTokens.serve_tokens_full_generator + is not PrimeRlServingTokens.__mro__[1].serve_tokens_full_generator + ) + assert PrimeRlServingTokens.serve_tokens is not PrimeRlServingTokens.__mro__[1].serve_tokens + + +def test_client_set_max_tokens_recognizes_explicit_value(): + body = {"token_ids": [1, 2, 3], "sampling_params": {"max_tokens": 256}} + assert asyncio.run(_client_set_max_tokens(_FakeRawRequest(body))) is True + + +def test_client_set_max_tokens_detects_unset(): + body = {"token_ids": [1, 2, 3], "sampling_params": {}} + assert asyncio.run(_client_set_max_tokens(_FakeRawRequest(body))) is False + + body_without_sp = {"token_ids": [1, 2, 3]} + assert asyncio.run(_client_set_max_tokens(_FakeRawRequest(body_without_sp))) is False + + +def test_client_set_max_tokens_assumes_set_when_body_unreadable(): + # No raw_request → can't tell, don't override. + assert asyncio.run(_client_set_max_tokens(None)) is True + + # body read raises → can't tell, don't override. + err = ValueError("bad json") + assert asyncio.run(_client_set_max_tokens(_FakeRawRequest(err))) is True + + # non-dict body → can't tell, don't override. + assert asyncio.run(_client_set_max_tokens(_FakeRawRequest([1, 2, 3]))) is True diff --git a/tests/unit/orchestrator/test_advantage.py b/tests/unit/orchestrator/test_advantage.py new file mode 100644 index 00000000000..28051ec97d1 --- /dev/null +++ b/tests/unit/orchestrator/test_advantage.py @@ -0,0 +1,369 @@ +import torch + +from prime_rl.configs.orchestrator import ( + CustomAdvantageConfig, + DefaultAdvantageConfig, + TokensLengthPenaltyConfig, + TurnsLengthPenaltyConfig, +) +from prime_rl.orchestrator.advantage import ( + AdvantageInputs, + AdvantageOutputs, + compute_advantages, + default_advantage_fn, + setup_advantage_fn, +) + + +def _make_rollout(reward: float, completion_len: int = 0, num_turns: int = 1) -> dict: + """Create a minimal rollout dict for advantage testing. + + `completion_len` tokens are split across `num_turns` trajectory steps. + """ + per_turn, rem = divmod(completion_len, max(num_turns, 1)) + trajectory = [ + {"tokens": {"prompt_ids": [0], "completion_ids": list(range(per_turn + (rem if i == 0 else 0)))}} + for i in range(num_turns) + ] + return {"reward": reward, "trajectory": trajectory} + + +def _make_inputs(rewards, completion_lengths=None, num_turns=None) -> AdvantageInputs: + """Build AdvantageInputs from 2D arrays of rewards/lengths/turns.""" + rewards_t = torch.as_tensor(rewards, dtype=torch.float32) + num_problems, rollouts_per_example = rewards_t.shape + rollouts = [] + for i in range(num_problems): + group = [] + for j in range(rollouts_per_example): + cl = int(completion_lengths[i][j]) if completion_lengths is not None else 0 + nt = int(num_turns[i][j]) if num_turns is not None else 1 + group.append(_make_rollout(float(rewards_t[i, j]), cl, nt)) + rollouts.append(group) + return AdvantageInputs(rollouts=rollouts) + + +# Helper aliases for readability — completion-only and tool-only token shaping. +_TOKENS_COMPLETION = TokensLengthPenaltyConfig(completion_weight=1.0, tool_response_weight=0.0) +_TOKENS_TOOL_ONLY = TokensLengthPenaltyConfig(completion_weight=0.0, tool_response_weight=1.0) + + +def test_default_advantage_fn_simple_mean(): + inputs = _make_inputs( + rewards=[[1.0, 0.5, 0.8], [0.2, 0.9, 0.1]], + completion_lengths=[[10, 12, 8], [15, 11, 9]], + ) + result = default_advantage_fn(inputs) + + assert result.advantages.shape == (2, 3) + # Check that mean is subtracted per row + assert torch.allclose(result.advantages.mean(dim=1), torch.zeros(2), atol=1e-6) + + +def test_efficiency_mixed_group(): + """Mixed group: reward shaping preserves zero-mean, shorter correct gets higher advantage.""" + inputs = _make_inputs( + rewards=[[1.0, 1.0, 0.0, 1.0]], + completion_lengths=[[10, 30, 20, 20]], + ) + result = default_advantage_fn(inputs, length_penalty=_TOKENS_COMPLETION) + + # mean_correct_len = (10+30+20)/3 = 20 + # bonus = clamp(1 - [10,30,20,20]/20, 0, 1) = [0.5, 0, 0, 0] + # shaped_rewards = R * (1 + bonus * correct_mask) = [1.5, 1, 0, 1] + # baseline = mean(shaped_rewards) = 0.875 + # A = shaped_rewards - baseline = [0.625, 0.125, -0.875, 0.125] + expected = torch.tensor([[0.625, 0.125, -0.875, 0.125]]) + assert torch.allclose(result.advantages, expected, atol=1e-6) + + # Zero-mean per group + assert torch.allclose(result.advantages.mean(dim=1), torch.zeros(1), atol=1e-6) + + # All correct rollouts have positive advantage + rewards = torch.tensor([r["reward"] for r in inputs.rollouts[0]]) + correct_mask = rewards >= 1.0 + assert (result.advantages[0][correct_mask] > 0).all() + + +def test_efficiency_all_correct_group(): + """All-correct group: zero-mean, shorter gets higher advantage.""" + inputs = _make_inputs( + rewards=[[1.0, 1.0, 1.0]], + completion_lengths=[[10, 20, 40]], + ) + result = default_advantage_fn(inputs, length_penalty=_TOKENS_COMPLETION) + + # mean_len = 70/3 ≈ 23.33 + # bonus = clamp(1 - [10, 20, 40] / (70/3), 0, 1) = [4/7, 1/7, 0] + # shaped_rewards = [1+4/7, 1+1/7, 1] = [11/7, 8/7, 1] + # baseline = mean = (11/7 + 8/7 + 1) / 3 = (11+8+7)/(7*3) = 26/21 + # A = shaped - baseline + shaped = torch.tensor([[11.0 / 7, 8.0 / 7, 1.0]]) + baseline = shaped.mean(dim=1, keepdim=True) + expected = shaped - baseline + assert torch.allclose(result.advantages, expected, atol=1e-6) + + # Zero-mean + assert torch.allclose(result.advantages.mean(dim=1), torch.zeros(1), atol=1e-6) + + # Shortest has highest advantage + assert result.advantages[0, 0] > result.advantages[0, 1] > result.advantages[0, 2] + + +def test_efficiency_all_zero_rewards(): + """When all rewards are 0, no length shaping — falls back to standard GRPO.""" + inputs = _make_inputs( + rewards=[[0.0, 0.0, 0.0]], + completion_lengths=[[10, 20, 15]], + ) + result_with = default_advantage_fn(inputs, length_penalty=_TOKENS_COMPLETION) + result_without = default_advantage_fn(inputs) + + assert torch.allclose(result_with.advantages, result_without.advantages, atol=1e-6) + + +def test_efficiency_single_correct(): + """Single correct rollout: bonus=0 (at its own mean), same as standard GRPO.""" + inputs = _make_inputs( + rewards=[[1.0, 0.0, 0.0, 0.0]], + completion_lengths=[[100, 50, 200, 150]], + ) + result = default_advantage_fn(inputs, length_penalty=_TOKENS_COMPLETION) + + expected = torch.tensor([[0.75, -0.25, -0.25, -0.25]]) + assert torch.allclose(result.advantages, expected, atol=1e-6) + + +def test_efficiency_shorter_correct_higher_advantage(): + """Among correct rollouts in a mixed group, shorter always gets higher advantage.""" + inputs = _make_inputs( + rewards=[[1.0, 1.0, 1.0, 0.0, 0.0]], + completion_lengths=[[50, 100, 200, 80, 120]], + ) + result = default_advantage_fn(inputs, length_penalty=_TOKENS_COMPLETION) + + advs = result.advantages[0] + assert advs[0] > advs[1] > advs[2] + assert (advs[:3] > 0).all() + assert (advs[3:] < 0).all() + + +def test_efficiency_zero_mean_per_group(): + """Reward shaping preserves zero-mean advantages per group.""" + inputs = _make_inputs( + rewards=[ + [1.0, 1.0, 0.0, 1.0], # mixed + [1.0, 1.0, 1.0, 1.0], # all correct + ], + completion_lengths=[ + [10, 30, 20, 20], + [10, 20, 40, 80], + ], + ) + result = default_advantage_fn(inputs, length_penalty=_TOKENS_COMPLETION) + + assert torch.allclose(result.advantages.mean(dim=1), torch.zeros(2), atol=1e-6) + + +def test_efficiency_amplification_bounded(): + """Even with extreme length outliers, reward amplification is capped at 2x.""" + inputs = _make_inputs( + rewards=[[1.0, 1.0, 0.0]], + completion_lengths=[[1, 10000, 5000]], + ) + result = default_advantage_fn(inputs, length_penalty=_TOKENS_COMPLETION) + + # Shortest correct gets bonus ≈ 1, so shaped_reward ≈ 2 + # Standard reward = 1, so amplification ≈ 2x + # shaped_rewards ≈ [2, 1, 0], baseline ≈ 1, max advantage ≈ 1 + assert result.advantages[0, 0] < 1.0 + 1e-3 + + +def test_efficiency_multiple_problems(): + """Handles multiple problems independently.""" + inputs = _make_inputs( + rewards=[ + [1.0, 1.0, 0.0], # mixed + [1.0, 1.0, 1.0], # all correct + ], + completion_lengths=[ + [10, 20, 15], + [10, 20, 40], + ], + ) + result = default_advantage_fn(inputs, length_penalty=_TOKENS_COMPLETION) + + # Row 0: mixed group — shorter correct > longer correct + assert result.advantages[0, 0] > result.advantages[0, 1] + assert (result.advantages[0, :2] > 0).all() + assert result.advantages[0, 2] < 0 + + # Row 1: all-correct group — shorter gets higher advantage + assert result.advantages[1, 0] > result.advantages[1, 1] > result.advantages[1, 2] + + # Both rows have zero-mean + assert torch.allclose(result.advantages.mean(dim=1), torch.zeros(2), atol=1e-6) + + +def test_efficiency_tokens_with_tool_response_weight(): + """`tool_response_weight` shifts shaping onto tool-response tokens read from rollout metrics.""" + rollouts = [ + [ + { + "reward": 1.0, + "trajectory": [{"tokens": {"prompt_ids": [0], "completion_ids": list(range(10))}}], + "metrics": {"rlm_total_tool_response_tokens": 200}, + }, + { + "reward": 1.0, + "trajectory": [{"tokens": {"prompt_ids": [0], "completion_ids": list(range(10))}}], + "metrics": {"rlm_total_tool_response_tokens": 0}, + }, + { + "reward": 1.0, + "trajectory": [{"tokens": {"prompt_ids": [0], "completion_ids": list(range(10))}}], + "metrics": {"rlm_total_tool_response_tokens": 100}, + }, + ] + ] + inputs = AdvantageInputs(rollouts=rollouts) + + # completion tokens identical (10 each) → completion-only shaping is a no-op + result_completion_only = default_advantage_fn(inputs, length_penalty=_TOKENS_COMPLETION) + assert torch.allclose(result_completion_only.advantages, torch.zeros(1, 3), atol=1e-6) + + # tool-response only: costs are [200, 0, 100], mean=100, bonus is one-sided + # so only the below-mean rollout (idx 1) gets amplified; the at/above-mean tie. + result_tool_only = default_advantage_fn(inputs, length_penalty=_TOKENS_TOOL_ONLY) + advs = result_tool_only.advantages[0] + assert advs[1] > advs[0] + assert advs[1] > advs[2] + assert torch.allclose(advs[0], advs[2], atol=1e-6) + assert torch.allclose(result_tool_only.advantages.mean(dim=1), torch.zeros(1), atol=1e-6) + + +def test_efficiency_fractional_weight_with_int_rewards(): + """Fractional weights must not truncate when rollout rewards are emitted as ints.""" + rollouts_int = [ + [ + {"reward": 1, "trajectory": [{"tokens": {"prompt_ids": [0], "completion_ids": list(range(7))}}]}, + {"reward": 1, "trajectory": [{"tokens": {"prompt_ids": [0], "completion_ids": list(range(11))}}]}, + {"reward": 0, "trajectory": [{"tokens": {"prompt_ids": [0], "completion_ids": list(range(13))}}]}, + ] + ] + rollouts_float = [[{**r, "reward": float(r["reward"])} for r in g] for g in rollouts_int] + + fractional = TokensLengthPenaltyConfig(completion_weight=0.3, tool_response_weight=0.0) + int_result = default_advantage_fn(AdvantageInputs(rollouts=rollouts_int), length_penalty=fractional) + float_result = default_advantage_fn(AdvantageInputs(rollouts=rollouts_float), length_penalty=fractional) + assert torch.allclose(int_result.advantages, float_result.advantages, atol=1e-6) + + +def test_efficiency_zero_costs_falls_back_to_plain_grpo(): + """When all effective costs are zero, shaping is a no-op (no NaNs from div-by-zero).""" + # tool-only weights but no harness metric → all costs == 0 + rollouts = [ + [ + {"reward": 1.0, "trajectory": [{"tokens": {"prompt_ids": [0], "completion_ids": list(range(10))}}]}, + {"reward": 1.0, "trajectory": [{"tokens": {"prompt_ids": [0], "completion_ids": list(range(10))}}]}, + {"reward": 0.0, "trajectory": [{"tokens": {"prompt_ids": [0], "completion_ids": list(range(10))}}]}, + ] + ] + inputs = AdvantageInputs(rollouts=rollouts) + result = default_advantage_fn(inputs, length_penalty=_TOKENS_TOOL_ONLY) + expected = default_advantage_fn(inputs) # plain GRPO + assert not torch.isnan(result.advantages).any() + assert torch.allclose(result.advantages, expected.advantages, atol=1e-6) + + +def test_efficiency_tokens_default_weights_match_completion_when_no_metric(): + """Default TokensLengthPenaltyConfig (1,1) reduces to completion-only when rollouts lack the metric.""" + inputs = _make_inputs( + rewards=[[1.0, 1.0, 0.0, 1.0]], + completion_lengths=[[10, 30, 20, 20]], + ) + result_default = default_advantage_fn(inputs, length_penalty=TokensLengthPenaltyConfig()) + result_completion = default_advantage_fn(inputs, length_penalty=_TOKENS_COMPLETION) + assert torch.allclose(result_default.advantages, result_completion.advantages, atol=1e-6) + + +def test_efficiency_turns_penalty(): + """`TurnsLengthPenaltyConfig` shapes by trajectory turn count rather than token count.""" + inputs = _make_inputs( + rewards=[[1.0, 1.0, 0.0, 1.0]], + # token counts identical, but turns differ — turns penalty should still differentiate + completion_lengths=[[100, 100, 100, 100]], + num_turns=[[1, 3, 2, 2]], + ) + result = default_advantage_fn(inputs, length_penalty=TurnsLengthPenaltyConfig()) + + # mean_correct_turns = (1+3+2)/3 = 2 + # bonus = clamp(1 - [1,3,2,2]/2, 0, 1) = [0.5, 0, 0, 0] + expected = torch.tensor([[0.625, 0.125, -0.875, 0.125]]) + assert torch.allclose(result.advantages, expected, atol=1e-6) + + +def test_compute_advantages_with_config(): + rewards = [1.0, 0.5, 0.8, 0.2, 0.9, 0.1] + lengths = [10, 12, 8, 15, 11, 9] + rollouts = [_make_rollout(r, l) for r, l in zip(rewards, lengths)] + + compute_advantages(rollouts, samples_per_problem=3, advantage_config=DefaultAdvantageConfig()) + + advantages = [r["advantage"] for r in rollouts] + assert len(advantages) == 6 + assert abs(sum(advantages[:3])) < 1e-5 + assert abs(sum(advantages[3:])) < 1e-5 + + +def test_compute_advantages_no_cross_group_leakage(): + """Per-problem grouping: each problem must be centered against its own mean, in-order. + + Two problems with very different reward scales — cross-group leakage would pull the + small-scale group's advantages toward the large-scale group's mean (and vice versa). + Distinct positional values also catch slicing/transpose bugs in the flat→grouped→flat + round-trip. + """ + rewards = [10.0, 20.0, 30.0, 0.0, 0.1, 0.2] + rollouts = [_make_rollout(r) for r in rewards] + + compute_advantages(rollouts, samples_per_problem=3, advantage_config=DefaultAdvantageConfig()) + + advantages = [r["advantage"] for r in rollouts] + expected = [-10.0, 0.0, 10.0, -0.1, 0.0, 0.1] + for got, want in zip(advantages, expected): + assert abs(got - want) < 1e-5, (advantages, expected) + + +def test_compute_advantages_without_config(): + rewards = [1.0, 0.5, 0.8] + lengths = [10, 12, 8] + rollouts = [_make_rollout(r, l) for r, l in zip(rewards, lengths)] + + compute_advantages(rollouts, samples_per_problem=3, advantage_config=None) + + advantages = [r["advantage"] for r in rollouts] + assert advantages == rewards + + +def test_setup_advantage_fn_with_custom_config(): + config = CustomAdvantageConfig( + import_path="tests.unit.orchestrator.test_advantage._dummy_custom_advantage", + kwargs={"scale": 2.0}, + ) + advantage_fn = setup_advantage_fn(config) + + inputs = _make_inputs( + rewards=[[1.0, 0.5, 0.8]], + completion_lengths=[[10, 12, 8]], + ) + + result = advantage_fn(inputs) + assert isinstance(result, AdvantageOutputs) + assert torch.allclose(result.advantages, torch.tensor([[2.0, 1.0, 1.6]])) + + +def _dummy_custom_advantage(inputs: AdvantageInputs, scale: float = 1.0) -> AdvantageOutputs: + """A simple custom advantage for testing.""" + rewards = torch.tensor([[r["reward"] for r in group] for group in inputs.rollouts]) + return AdvantageOutputs(advantages=rewards * scale) diff --git a/tests/unit/orchestrator/test_batch.py b/tests/unit/orchestrator/test_batch.py index e27ff5bcf9e..a2e2e50079e 100644 --- a/tests/unit/orchestrator/test_batch.py +++ b/tests/unit/orchestrator/test_batch.py @@ -1,20 +1,22 @@ import pytest -from prime_rl.trainer.batch import prepare_batch +from prime_rl.trainer.batch import prepare_batch, prepare_sample from prime_rl.transport.types import TrainingSample @pytest.fixture def make_training_example(): - def _make_training_example() -> TrainingSample: + def _make_training_example(temperature: float = 1.0, sft_loss: bool = False) -> TrainingSample: return TrainingSample( prompt_ids=[1, 2], prompt_mask=[False, False], completion_ids=[3, 4], completion_mask=[True, True], completion_logprobs=[-0.1, -0.2], + completion_temperatures=[temperature, temperature], # Per-token temperatures teacher_logprobs=[0.0, 0.0, 0.0, 0.0], advantage=1.0, + sft_loss=sft_loss, ) return _make_training_example @@ -30,7 +32,6 @@ def test_prepare_batch_balances_micro_batches_across_workers( batches_per_gpu = prepare_batch( rollouts=examples, - temperature=0.5, seq_len=4, num_train_workers=num_train_workers, idxs=[0] * rollout_count, @@ -53,3 +54,109 @@ def test_prepare_batch_balances_micro_batches_across_workers( for batch in flat_batches[len(examples) :]: assert sum(1 for advantage in batch.advantages if advantage != 0.0) == 0 assert sum(1 for loss_mask in batch.loss_mask if loss_mask) == 0 + + +def test_prepare_batch_packs_different_temperatures(make_training_example): + """With per-token temperatures, samples can be packed together regardless of their temperature values.""" + example1 = make_training_example(temperature=0.7) + example2 = make_training_example(temperature=1.1) + + batches_per_gpu = prepare_batch( + rollouts=[example1, example2], + seq_len=16, + num_train_workers=1, + idxs=[0, 0], + num_loras=1, + ) + + flat_batches = [batch for worker_batches in batches_per_gpu for batch in worker_batches] + # With per-token temperatures, samples can now be packed together + assert len(flat_batches) == 1 + # Each sample has 4 tokens (2 prompt + 2 completion), so 8 total tokens + assert len(flat_batches[0].temperatures) == 8 + # First sample (4 tokens): all get temp 0.7 + assert flat_batches[0].temperatures[:4] == [0.7, 0.7, 0.7, 0.7] + # Second sample (4 tokens): all get temp 1.1 + assert flat_batches[0].temperatures[4:8] == [1.1, 1.1, 1.1, 1.1] + + +def test_prepare_sample_propagates_sft_loss(make_training_example): + example = make_training_example(sft_loss=True) + + micro_batch = prepare_sample(example, seq_len=16) + + assert micro_batch.sft_loss is True + + +def test_prepare_batch_does_not_pack_mixed_sft_loss(make_training_example): + rl_example = make_training_example(sft_loss=False) + sft_example = make_training_example(sft_loss=True) + + batches_per_gpu = prepare_batch( + rollouts=[rl_example, sft_example], + seq_len=16, + num_train_workers=1, + idxs=[0, 0], + num_loras=1, + ) + + flat_batches = [batch for worker_batches in batches_per_gpu for batch in worker_batches] + assert len(flat_batches) == 2 + assert {batch.sft_loss for batch in flat_batches} == {False, True} + + +def test_prepare_sample_with_routed_experts(): + """Routed experts are passed through prepare_sample and match input_ids length.""" + # 2 prompt + 2 completion = 4 tokens, 2 layers, topk=2 + routed_experts = [[[0, 1], [2, 3]], [[4, 5], [6, 7]], [[0, 2], [1, 3]], [[1, 0], [3, 2]]] + sample = TrainingSample( + prompt_ids=[1, 2], + prompt_mask=[False, False], + completion_ids=[3, 4], + completion_mask=[True, True], + completion_logprobs=[-0.1, -0.2], + completion_temperatures=[1.0, 1.0], + advantage=1.0, + routed_experts=routed_experts, + ) + + micro_batch = prepare_sample(sample, seq_len=8) + assert micro_batch.routed_experts is not None + assert len(micro_batch.routed_experts) == 4 + assert micro_batch.routed_experts == routed_experts + + +def test_prepare_sample_truncates_routed_experts(): + """Routed experts are truncated to seq_len when input exceeds it.""" + routed_experts = [[[0, 1]], [[2, 3]], [[4, 5]], [[6, 7]]] + sample = TrainingSample( + prompt_ids=[1, 2], + prompt_mask=[False, False], + completion_ids=[3, 4], + completion_mask=[True, True], + completion_logprobs=[-0.1, -0.2], + completion_temperatures=[1.0, 1.0], + advantage=1.0, + routed_experts=routed_experts, + ) + + micro_batch = prepare_sample(sample, seq_len=3) + assert micro_batch.routed_experts is not None + assert len(micro_batch.routed_experts) == 3 + assert micro_batch.routed_experts == routed_experts[:3] + + +def test_prepare_sample_none_routed_experts(): + """When routed_experts is None, micro_batch.routed_experts is None.""" + sample = TrainingSample( + prompt_ids=[1, 2], + prompt_mask=[False, False], + completion_ids=[3, 4], + completion_mask=[True, True], + completion_logprobs=[-0.1, -0.2], + completion_temperatures=[1.0, 1.0], + advantage=1.0, + ) + + micro_batch = prepare_sample(sample, seq_len=8) + assert micro_batch.routed_experts is None diff --git a/tests/unit/orchestrator/test_buffer.py b/tests/unit/orchestrator/test_buffer.py index 10c5f1ab506..e7d2b038c08 100644 --- a/tests/unit/orchestrator/test_buffer.py +++ b/tests/unit/orchestrator/test_buffer.py @@ -5,8 +5,28 @@ import verifiers as vf from datasets import Dataset +from prime_rl.configs.orchestrator import BufferConfig, EnvConfig from prime_rl.orchestrator.buffer import Buffer -from prime_rl.orchestrator.config import BufferConfig +from prime_rl.orchestrator.envs import Envs, TrainEnv + + +def make_env(name: str, vf_env: vf.Environment, **config_kwargs) -> TrainEnv: + """Create a TrainEnv without calling vf.load_environment.""" + config = EnvConfig(id=name, name=name, **config_kwargs) + env = TrainEnv.__new__(TrainEnv) + env.config = config + env._env = vf_env + env._env_client = None + env._env_server_process = None + env.sampling_args = {} + return env + + +def make_envs(env_dict: dict[str, TrainEnv]) -> Envs: + """Create an Envs container from a dict of Env instances.""" + envs = Envs.__new__(Envs) + envs._envs = env_dict + return envs @pytest.fixture(autouse=True) @@ -32,8 +52,8 @@ def dummy_dataset() -> Dataset: @pytest.fixture -def dummy_env_group(mock_openai_client, dummy_dataset) -> vf.EnvGroup: - """Return an EnvGroup with two dummy envs using the same dataset.""" +def dummy_envs(mock_openai_client, dummy_dataset) -> Envs: + """Return an Envs with two dummy envs.""" env_a = vf.SingleTurnEnv( client=mock_openai_client, model="test-model", @@ -46,24 +66,29 @@ def dummy_env_group(mock_openai_client, dummy_dataset) -> vf.EnvGroup: dataset=dummy_dataset, rubric=vf.Rubric(), ) - return vf.EnvGroup(envs=[env_a, env_b], env_names=["env_a", "env_b"]) + return make_envs( + { + "env_a": make_env("env_a", env_a), + "env_b": make_env("env_b", env_b), + } + ) @pytest.fixture def make_rollouts(): - def _make_rollouts(dataset: Dataset, rewards: list[float]) -> list[vf.State]: - rollouts = [] - for i, reward in enumerate(rewards): - task = dataset[i]["task"] - example_id = dataset[i]["example_id"] - prompt = dataset[i]["prompt"] - problem_rollouts = [ - vf.State( - input=vf.RolloutInput( - example_id=example_id, - task=task, - prompt=prompt, - ), + def _make_rollouts( + buffer: Buffer, env_name: str, indices: list[int], rewards: list[float] + ) -> list[vf.RolloutOutput]: + all_rollouts = [] + eb = buffer.env_buffers[env_name] + examples = list(eb.examples.values()) + for idx, reward in zip(indices, rewards): + example = examples[idx] + rollouts = [ + vf.RolloutOutput( + example_id=example["example_id"], + task=example["env_name"], + prompt=example["prompt"], prompt_ids=[0], prompt_mask=[1], completion_ids=[1], @@ -75,104 +100,107 @@ def _make_rollouts(dataset: Dataset, rewards: list[float]) -> list[vf.State]: metrics={}, ) ] * 2 - rollouts.extend(problem_rollouts) - return rollouts + for r in rollouts: + r["env_name"] = env_name + all_rollouts.extend(rollouts) + return all_rollouts return _make_rollouts -def get_normal_ids(buffer: Buffer) -> set[int]: - return {example_id for env in buffer.example_buffer.values() for example_id in env.keys()} +def get_normal_count(buffer: Buffer) -> int: + return sum(eb.num_normal for eb in buffer.env_buffers.values()) -def test_buffer_init_and_sample(dummy_env_group): - dataset = dummy_env_group.get_dataset() - buffer = Buffer(dataset, dummy_env_group.env_names, BufferConfig()) - # Each env has 5 examples, so total is 10 - assert len(buffer.example_buffer["env_a"]) == 5 - assert len(buffer.example_buffer["env_b"]) == 5 +def test_buffer_init_and_sample(dummy_envs): + buffer = Buffer(dummy_envs, BufferConfig()) + assert buffer.env_buffers["env_a"].num_normal == 5 + assert buffer.env_buffers["env_b"].num_normal == 5 samples = buffer.sample_examples(2) assert len(samples) == 2 -def test_buffer_problem_pool_assignment(dummy_env_group, make_rollouts): +def test_buffer_problem_pool_assignment(dummy_envs, make_rollouts): """Problems are moved to easy/hard pools based on reward thresholds.""" - dataset = dummy_env_group.get_dataset() - buffer = Buffer(dataset, dummy_env_group.env_names, BufferConfig(easy_threshold=1.0, hard_threshold=0.0)) - dataset = buffer.dataset - # Use first 5 examples (all from env_a since they come first in concatenated dataset) - buffer.update(make_rollouts(dataset.select(range(5)), rewards=[1.0, 1.0, 0.5, 0.5, 0.0])) + buffer = Buffer(dummy_envs, BufferConfig(easy_threshold=1.0, hard_threshold=0.0)) + buffer.update(make_rollouts(buffer, "env_a", list(range(5)), rewards=[1.0, 1.0, 0.5, 0.5, 0.0])) - assert len(buffer.easy_examples) == 2 - assert len(buffer.hard_examples) == 1 - # 2 normal from first 5, plus 5 from env_b = 7 - assert len(get_normal_ids(buffer)) == 7 + assert len(buffer.env_buffers["env_a"].easy_examples) == 2 + assert len(buffer.env_buffers["env_a"].hard_examples) == 1 + # 2 normal from env_a + 5 from env_b = 7 + assert get_normal_count(buffer) == 7 -def test_buffer_online_difficulty_filtering(dummy_env_group, make_rollouts): +def test_buffer_online_difficulty_filtering(dummy_envs, make_rollouts): """With online_difficulty_filtering=True, only partial reward rollouts are kept.""" - dataset = dummy_env_group.get_dataset() buffer = Buffer( - dataset, - dummy_env_group.env_names, + dummy_envs, BufferConfig(online_difficulty_filtering=True), ) - buffer.update(make_rollouts(dataset.select(range(5)), rewards=[1.0, 0.5, 0.0, 0.5, 0.5])) + buffer.update(make_rollouts(buffer, "env_a", list(range(5)), rewards=[1.0, 0.5, 0.0, 0.5, 0.5])) # Only 3 problems with reward 0.5 -> 6 rollouts kept assert len(buffer.rollout_buffer) == 6 -def test_buffer_no_filtering_by_default(dummy_env_group, make_rollouts): +def test_buffer_no_filtering_by_default(dummy_envs, make_rollouts): """With online_difficulty_filtering=False (default), all rollouts are kept.""" - dataset = dummy_env_group.get_dataset() - buffer = Buffer(dataset, dummy_env_group.env_names, BufferConfig()) - buffer.update(make_rollouts(dataset.select(range(5)), rewards=[1.0, 0.5, 0.0, 0.5, 0.5])) + buffer = Buffer(dummy_envs, BufferConfig()) + buffer.update(make_rollouts(buffer, "env_a", list(range(5)), rewards=[1.0, 0.5, 0.0, 0.5, 0.5])) # All 5 problems -> 10 rollouts kept assert len(buffer.rollout_buffer) == 10 -def test_buffer_save_load_with_conversion(dummy_env_group, make_rollouts, tmp_path): +def test_buffer_save_load_with_conversion(dummy_envs, make_rollouts, tmp_path): """Easy/hard problems are partially converted to normal on load.""" - dataset = dummy_env_group.get_dataset() - buffer = Buffer(dataset, dummy_env_group.env_names, BufferConfig(easy_threshold=1.0, hard_threshold=0.0)) - buffer.update(make_rollouts(dataset.select(range(5)), rewards=[1.0, 1.0, 0.5, 0.5, 0.0])) + buffer = Buffer(dummy_envs, BufferConfig(easy_threshold=1.0, hard_threshold=0.0)) + buffer.update(make_rollouts(buffer, "env_a", list(range(5)), rewards=[1.0, 1.0, 0.5, 0.5, 0.0])) buffer.save(tmp_path / "buffer") - new_buffer = Buffer( - dataset, dummy_env_group.env_names, BufferConfig(easy_fraction=0.5, hash_keys=["prompt", "task"]) - ) + new_buffer = Buffer(dummy_envs, BufferConfig(easy_fraction=0.5, hash_keys=["prompt", "env_name"])) new_buffer.load(tmp_path / "buffer") # 1 of 2 easy problems converted to normal - assert len(new_buffer.easy_examples) == 1 + assert len(new_buffer.env_buffers["env_a"].easy_examples) == 1 # 2 were normal + 5 from env_b + 1 converted from easy = 8 - assert len(get_normal_ids(new_buffer)) == 8 + assert get_normal_count(new_buffer) == 8 -def test_buffer_env_ratios(dummy_env_group): - dataset = dummy_env_group.get_dataset() - buffer = Buffer(dataset, dummy_env_group.env_names, BufferConfig(env_ratios=[0.8, 0.2])) - assert len(buffer.example_buffer["env_a"]) == 5 - assert len(buffer.example_buffer["env_b"]) == 5 +def test_buffer_env_ratios(mock_openai_client, dummy_dataset): + env_a = vf.SingleTurnEnv(client=mock_openai_client, model="test-model", dataset=dummy_dataset, rubric=vf.Rubric()) + env_b = vf.SingleTurnEnv(client=mock_openai_client, model="test-model", dataset=dummy_dataset, rubric=vf.Rubric()) + envs = make_envs( + { + "env_a": make_env("env_a", env_a, ratio=0.8), + "env_b": make_env("env_b", env_b, ratio=0.2), + } + ) + + buffer = Buffer(envs, BufferConfig()) + assert buffer.env_buffers["env_a"].num_normal == 5 + assert buffer.env_buffers["env_b"].num_normal == 5 samples = buffer.sample_examples(100) - env_a_count = sum(1 for p in samples if p["task"] == "env_a") + env_a_count = sum(1 for p in samples if p["env_name"] == "env_a") assert 60 <= env_a_count <= 95 def test_buffer_env_ratios_validation(): - """BufferConfig validates that all env_ratios are positive.""" + """Validates that env ratios must be positive and all-or-nothing.""" from pydantic import ValidationError - with pytest.raises(ValidationError, match="All env_ratios must be positive"): - BufferConfig(env_ratios=[0.5, -0.3, 0.2]) + from prime_rl.configs.orchestrator import TrainConfig, TrainEnvConfig + + with pytest.raises(ValidationError): + EnvConfig(id="env_a", ratio=-0.3) + + with pytest.raises(ValidationError, match="mix of set and unset"): + TrainConfig(env=[TrainEnvConfig(id="a", ratio=0.5), TrainEnvConfig(id="b")]) def test_buffer_no_cross_env_pool_assignment(mock_openai_client, tmp_path): - """Pool assignments don't transfer if example_id exists but task/env changed.""" - # Original: create an env_group with only env_a + """Pool assignments don't transfer if example_id exists but env changed.""" original_dataset = Dataset.from_dict({"question": ["q0"], "answer": ["a0"]}) original_env = vf.SingleTurnEnv( client=mock_openai_client, @@ -180,17 +208,15 @@ def test_buffer_no_cross_env_pool_assignment(mock_openai_client, tmp_path): dataset=original_dataset, rubric=vf.Rubric(), ) - original_env_group = vf.EnvGroup(envs=[original_env], env_names=["env_a"]) - original_env_dataset = original_env_group.get_dataset() - - buffer = Buffer(original_env_dataset, original_env_group.env_names, BufferConfig(easy_threshold=1.0)) - # Manually move the example to easy pool - example_id = list(buffer.example_buffer["env_a"].keys())[0] - example = buffer.example_buffer["env_a"].pop(example_id) - buffer.easy_examples.append(example) + original_env_set = make_envs({"env_a": make_env("env_a", original_env)}) + + buffer = Buffer(original_env_set, BufferConfig(easy_threshold=1.0)) + eb = buffer.env_buffers["env_a"] + example_id = list(eb.examples.keys())[0] + example = eb.examples.pop(example_id) + eb.easy_examples.append(example) buffer.save(tmp_path / "buffer") - # Resume: create a new env_group with different dataset but similar structure new_dataset = Dataset.from_dict({"question": ["different_q"], "answer": ["different_a"]}) new_env = vf.SingleTurnEnv( client=mock_openai_client, @@ -198,13 +224,10 @@ def test_buffer_no_cross_env_pool_assignment(mock_openai_client, tmp_path): dataset=new_dataset, rubric=vf.Rubric(), ) - new_env_group = vf.EnvGroup(envs=[new_env], env_names=["env_b"]) - new_env_dataset = new_env_group.get_dataset() + new_env_set = make_envs({"env_b": make_env("env_b", new_env)}) - new_buffer = Buffer(new_env_dataset, new_env_group.env_names, BufferConfig()) + new_buffer = Buffer(new_env_set, BufferConfig()) new_buffer.load(tmp_path / "buffer") - # Should NOT be in easy pool (different content, different hash) - assert len(new_buffer.easy_examples) == 0 - # Should still be in normal pool for env_b - assert len(new_buffer.example_buffer["env_b"]) == 1 + assert len(new_buffer.env_buffers["env_b"].easy_examples) == 0 + assert new_buffer.env_buffers["env_b"].num_normal == 1 diff --git a/tests/unit/orchestrator/test_env_worker.py b/tests/unit/orchestrator/test_env_worker.py deleted file mode 100644 index 3050b20f93f..00000000000 --- a/tests/unit/orchestrator/test_env_worker.py +++ /dev/null @@ -1,179 +0,0 @@ -"""Tests for EnvWorker auto-restart functionality.""" - -import asyncio -import os -import signal -import time -from multiprocessing import Process, Queue -from unittest.mock import MagicMock, patch - -import pytest - -from prime_rl.orchestrator.env_worker import EnvWorker, WorkerDiedError - - -def dummy_worker_main(request_queue: Queue, response_queue: Queue, **kwargs): - """A minimal worker that just sleeps until killed.""" - while True: - try: - request = request_queue.get(timeout=0.1) - if request is None: - break - except Exception: - pass - time.sleep(0.1) - - -@pytest.fixture -def mock_client_config(): - """Create a mock client config.""" - config = MagicMock() - config.model_dump.return_value = {} - return config - - -@pytest.fixture -def env_worker(mock_client_config): - """Create an EnvWorker with mocked dependencies.""" - worker = EnvWorker( - env_id="test_env", - env_args={}, - client_config=mock_client_config, - model_name="test-model", - seq_len=1024, - interleaved_rollouts=False, - max_concurrent=1, - example_lookup={}, - sampling_args={}, - worker_name="test_worker", - ) - return worker - - -def test_restart_drains_queues(env_worker): - """Test that restart() drains any stale data from queues.""" - # Put some data in the queues - env_worker.request_queue.put("stale_request") - env_worker.response_queue.put("stale_response") - - with patch.object(env_worker, "start"): - env_worker._restart() - - # Queues should be empty after restart - assert env_worker.request_queue.empty() - assert env_worker.response_queue.empty() - - -def test_pending_count_returns_large_when_dead(env_worker): - """Test that pending_count returns large number when worker is dead.""" - env_worker._dead = False - env_worker.pending_futures = {"a": None, "b": None} - assert env_worker.pending_count == 2 - - env_worker._dead = True - assert env_worker.pending_count == 999999 - - -def test_collect_responses_restarts_on_worker_death(env_worker): - """Test that collect_responses auto-restarts worker when it dies.""" - - async def run_test(): - # Create a real process that we can kill - env_worker.process = Process( - target=dummy_worker_main, - args=(env_worker.request_queue, env_worker.response_queue), - daemon=True, - ) - env_worker.process.start() - env_worker._stopping = False - env_worker._dead = False - - # Add a pending future that should be failed - loop = asyncio.get_event_loop() - pending_future = loop.create_future() - env_worker.pending_futures["test_request"] = pending_future - - # Start collect_responses in background - # Use side_effect to actually restart the process so the loop doesn't spin - def real_restart(): - env_worker.process = Process( - target=dummy_worker_main, - args=(env_worker.request_queue, env_worker.response_queue), - daemon=True, - ) - env_worker.process.start() - env_worker._dead = False - - with patch.object(env_worker, "_restart", side_effect=real_restart) as mock_restart: - collect_task = asyncio.create_task(env_worker.collect_responses()) - - # Give it a moment to start - await asyncio.sleep(0.1) - - # Kill the worker process - os.kill(env_worker.process.pid, signal.SIGKILL) - env_worker.process.join(timeout=1) - - # Wait for collect_responses to detect death and restart - await asyncio.sleep(0.2) - - # Cancel the collect task - collect_task.cancel() - try: - await collect_task - except asyncio.CancelledError: - pass - - # Cleanup the restarted process - if env_worker.process and env_worker.process.is_alive(): - env_worker.process.terminate() - env_worker.process.join(timeout=1) - - # Verify restart was called exactly once - mock_restart.assert_called_once() - - # Verify pending future was failed with WorkerDiedError - assert pending_future.done() - with pytest.raises(WorkerDiedError): - pending_future.result() - - asyncio.run(run_test()) - - -def test_full_restart_cycle(mock_client_config): - """Test a full restart cycle with actual process management.""" - worker = EnvWorker( - env_id="test_env", - env_args={}, - client_config=mock_client_config, - model_name="test-model", - seq_len=1024, - interleaved_rollouts=False, - max_concurrent=1, - example_lookup={}, - sampling_args={}, - worker_name="test_worker", - ) - - # Patch worker_main to use our dummy - with patch("prime_rl.orchestrator.env_worker.worker_main", dummy_worker_main): - # Start worker - worker.start() - assert worker.process is not None - assert worker.process.is_alive() - original_pid = worker.process.pid - - # Kill it - os.kill(worker.process.pid, signal.SIGKILL) - worker.process.join(timeout=1) - assert not worker.process.is_alive() - - # Restart - worker._restart() - assert worker.process is not None - assert worker.process.is_alive() - assert worker.process.pid != original_pid # New process - assert not worker._dead - - # Cleanup - worker.stop() diff --git a/tests/unit/orchestrator/test_eval_scheduling.py b/tests/unit/orchestrator/test_eval_scheduling.py new file mode 100644 index 00000000000..f12391ddaf5 --- /dev/null +++ b/tests/unit/orchestrator/test_eval_scheduling.py @@ -0,0 +1,86 @@ +"""Unit tests for eval scheduling logic - specifically the range check +that detects when ckpt_step jumps over eval interval boundaries.""" + +from prime_rl.orchestrator.eval_utils import compute_eval_ckpt_step + + +def test_exact_hit(): + result = compute_eval_ckpt_step(ckpt_step=25, prev_ckpt_step=24, last_eval_step=0, interval=25) + assert result == 25 + + +def test_jump_over_interval(): + result = compute_eval_ckpt_step(ckpt_step=26, prev_ckpt_step=24, last_eval_step=0, interval=25) + assert result == 25 + + +def test_no_interval_crossed(): + result = compute_eval_ckpt_step(ckpt_step=23, prev_ckpt_step=22, last_eval_step=0, interval=25) + assert result is None + + +def test_base_model_eval_at_step_0(): + result = compute_eval_ckpt_step( + ckpt_step=0, prev_ckpt_step=-1, last_eval_step=-1, interval=25, eval_base_model=True + ) + assert result == 0 + + +def test_base_model_eval_disabled(): + result = compute_eval_ckpt_step( + ckpt_step=0, prev_ckpt_step=-1, last_eval_step=-1, interval=25, eval_base_model=False + ) + assert result is None + + +def test_no_double_eval(): + result = compute_eval_ckpt_step(ckpt_step=25, prev_ckpt_step=24, last_eval_step=25, interval=25) + assert result is None + + +def test_no_change_in_ckpt_step(): + result = compute_eval_ckpt_step(ckpt_step=25, prev_ckpt_step=25, last_eval_step=0, interval=25) + assert result is None + + +def test_multiple_intervals_crossed(): + result = compute_eval_ckpt_step(ckpt_step=76, prev_ckpt_step=24, last_eval_step=0, interval=25) + assert result == 75 + + +def test_second_interval(): + result = compute_eval_ckpt_step(ckpt_step=50, prev_ckpt_step=49, last_eval_step=25, interval=25) + assert result == 50 + + +def test_jump_across_second_interval(): + result = compute_eval_ckpt_step(ckpt_step=51, prev_ckpt_step=48, last_eval_step=25, interval=25) + assert result == 50 + + +def test_production_scenario_step25_skipped(): + """Reproduces the bug from run c14miuyha2yhxkw1z3eqgyub.""" + result = compute_eval_ckpt_step(ckpt_step=26, prev_ckpt_step=24, last_eval_step=0, interval=25) + assert result == 25 + + +def test_production_scenario_step50_exact(): + result = compute_eval_ckpt_step(ckpt_step=50, prev_ckpt_step=49, last_eval_step=26, interval=25) + assert result == 50 + + +def test_simulate_full_run(): + ckpt_steps = [0, 0, 3, 5, 10, 15, 20, 24, 26, 30, 35, 40, 48, 51, 60, 70, 74, 76] + interval = 25 + last_eval_step = -1 + prev_ckpt_step = -1 + eval_triggered_at = [] + + for ckpt_step in ckpt_steps: + result = compute_eval_ckpt_step(ckpt_step, prev_ckpt_step, last_eval_step, interval) + if result is not None: + eval_triggered_at.append(result) + last_eval_step = ckpt_step + prev_ckpt_step = ckpt_step + + assert eval_triggered_at == [0, 25, 50, 75] diff --git a/tests/unit/orchestrator/test_filters.py b/tests/unit/orchestrator/test_filters.py new file mode 100644 index 00000000000..e77f51f61f3 --- /dev/null +++ b/tests/unit/orchestrator/test_filters.py @@ -0,0 +1,374 @@ +import math + +from prime_rl.configs.orchestrator import GibberishFilterConfig, RepetitionFilterConfig +from prime_rl.orchestrator.filters import ( + GibberishFilter, + RepetitionFilter, + apply_filters, + setup_filter, + setup_filters, +) + + +def _make_rollout(completion_ids, completion_logprobs, reward=1.0, multi_step=False): + """Create a minimal rollout dict matching the verifiers RolloutOutput structure.""" + if multi_step: + mid = len(completion_ids) // 2 + trajectory = [ + { + "tokens": { + "completion_ids": completion_ids[:mid], + "completion_logprobs": completion_logprobs[:mid], + "completion_mask": [1] * mid, + } + }, + { + "tokens": { + "completion_ids": completion_ids[mid:], + "completion_logprobs": completion_logprobs[mid:], + "completion_mask": [1] * (len(completion_ids) - mid), + } + }, + ] + else: + trajectory = [ + { + "tokens": { + "completion_ids": completion_ids, + "completion_logprobs": completion_logprobs, + "completion_mask": [1] * len(completion_ids), + } + } + ] + return { + "trajectory": trajectory, + "reward": reward, + "stop_condition": None, + "metrics": {}, + } + + +def _make_gibberish_filter(vocab_size=128_000, token_id_threshold=100_000, logprob_offset=2.0, enforce=False): + logprob_threshold = -math.log(vocab_size) - logprob_offset + return GibberishFilter( + name="gibberish", token_id_threshold=token_id_threshold, logprob_threshold=logprob_threshold, enforce=enforce + ) + + +def _make_repetition_filter(window=5, prob_threshold=0.99, enforce=False): + return RepetitionFilter( + name="repetition", window=window, logprob_threshold=math.log(prob_threshold), enforce=enforce + ) + + +# --- GibberishFilter tests --- + + +def test_gibberish_detects_rare_low_prob_token(): + gibberish_filter = _make_gibberish_filter() + + result = gibberish_filter.check( + _make_rollout( + completion_ids=[50, 120_000, 80], + completion_logprobs=[-1.0, gibberish_filter.logprob_threshold - 1.0, -0.5], + ) + ) + assert result.detected is True + assert result.detection_index == 1 + + +def test_gibberish_ignores_normal_tokens(): + gibberish_filter = _make_gibberish_filter() + + result = gibberish_filter.check( + _make_rollout( + completion_ids=[10, 200, 5000], + completion_logprobs=[-1.0, -2.0, -3.0], + ) + ) + assert result.detected is False + assert result.detection_index is None + + +def test_gibberish_ignores_high_prob_rare_token(): + gibberish_filter = _make_gibberish_filter() + + result = gibberish_filter.check( + _make_rollout( + completion_ids=[120_000], + completion_logprobs=[-0.5], + ) + ) + assert result.detected is False + + +def test_gibberish_works_across_trajectory_steps(): + gibberish_filter = _make_gibberish_filter() + + result = gibberish_filter.check( + _make_rollout( + completion_ids=[50, 60, 120_000, 80], + completion_logprobs=[-1.0, -0.5, gibberish_filter.logprob_threshold - 1.0, -0.5], + multi_step=True, + ) + ) + assert result.detected is True + assert result.detection_index == 2 + + +# --- RepetitionFilter tests --- + + +def test_repetition_triggers_after_window(): + repetition_filter = _make_repetition_filter(window=5) + + result = repetition_filter.check( + _make_rollout( + completion_ids=list(range(5)), + completion_logprobs=[-0.001] * 5, + ) + ) + assert result.detected is True + assert result.detection_index == 4 + + +def test_repetition_no_trigger_below_window(): + repetition_filter = _make_repetition_filter(window=5) + + result = repetition_filter.check( + _make_rollout( + completion_ids=list(range(4)), + completion_logprobs=[-0.001] * 4, + ) + ) + assert result.detected is False + + +def test_repetition_resets_on_low_prob(): + repetition_filter = _make_repetition_filter(window=5) + + logprobs = [-0.001] * 3 + [-2.0] + [-0.001] * 3 + result = repetition_filter.check( + _make_rollout( + completion_ids=list(range(7)), + completion_logprobs=logprobs, + ) + ) + assert result.detected is False + + +def test_repetition_varied_probs_no_trigger(): + repetition_filter = _make_repetition_filter(window=3) + + result = repetition_filter.check( + _make_rollout( + completion_ids=list(range(6)), + completion_logprobs=[-0.001, -3.0, -0.001, -3.0, -0.001, -3.0], + ) + ) + assert result.detected is False + + +# --- setup_filter / setup_filters tests --- + + +def test_setup_filter_gibberish(): + config = GibberishFilterConfig(token_id_threshold=100_000, logprob_offset=2.0) + gibberish_filter = setup_filter(config, vocab_size=128_000) + assert isinstance(gibberish_filter, GibberishFilter) + assert gibberish_filter.name == "gibberish" + assert gibberish_filter.token_id_threshold == 100_000 + assert abs(gibberish_filter.logprob_threshold - (-math.log(128_000) - 2.0)) < 1e-10 + assert gibberish_filter.enforce is False + + +def test_setup_filter_gibberish_enforce(): + config = GibberishFilterConfig(enforce=True) + gibberish_filter = setup_filter(config, vocab_size=128_000) + assert gibberish_filter.enforce is True + + +def test_setup_filter_repetition(): + config = RepetitionFilterConfig(window=3_000, prob_threshold=0.99) + repetition_filter = setup_filter(config, vocab_size=128_000) + assert isinstance(repetition_filter, RepetitionFilter) + assert repetition_filter.name == "repetition" + assert repetition_filter.window == 3_000 + assert abs(repetition_filter.logprob_threshold - math.log(0.99)) < 1e-10 + assert repetition_filter.enforce is False + + +def test_setup_filter_repetition_enforce(): + config = RepetitionFilterConfig(enforce=True) + repetition_filter = setup_filter(config, vocab_size=128_000) + assert repetition_filter.enforce is True + + +def test_setup_filters_multiple(): + configs = [ + GibberishFilterConfig(), + RepetitionFilterConfig(), + ] + filters = setup_filters(configs, vocab_size=128_000) + assert len(filters) == 2 + assert filters[0].name == "gibberish" + assert filters[1].name == "repetition" + + +# --- apply_filters tests (enforce=True) --- + + +def test_apply_filters_enforced_flags_rollout(): + gibberish_filter = _make_gibberish_filter(enforce=True) + + rollout = _make_rollout( + completion_ids=[120_000], + completion_logprobs=[gibberish_filter.logprob_threshold - 1.0], + reward=1.0, + ) + + apply_filters([gibberish_filter], [rollout]) + + assert rollout["reward"] == 1.0 + assert rollout["trajectory"][0]["tokens"]["completion_ids"] == [120_000] + assert rollout["trajectory"][0]["tokens"]["completion_mask"] == [1] + assert rollout["stop_condition"] is None + assert rollout["filters"] == {"gibberish": True} + assert rollout["is_filtered"] is True + + +def test_apply_filters_preserves_clean_rollouts(): + gibberish_filter = _make_gibberish_filter(enforce=True) + + rollout = _make_rollout( + completion_ids=[50, 60, 70], + completion_logprobs=[-1.0, -2.0, -1.5], + reward=1.0, + ) + + apply_filters([gibberish_filter], [rollout]) + + assert rollout["reward"] == 1.0 + assert rollout["trajectory"][0]["tokens"]["completion_ids"] == [50, 60, 70] + assert all(m == 1 for m in rollout["trajectory"][0]["tokens"]["completion_mask"]) + assert rollout["stop_condition"] is None + assert rollout["filters"] == {"gibberish": False} + assert rollout["is_filtered"] is False + + +def test_apply_filters_first_filter_wins(): + gibberish_filter = _make_gibberish_filter(enforce=True) + repetition_filter = _make_repetition_filter(window=2, enforce=True) + + rollout = _make_rollout( + completion_ids=[120_000, 1, 2], + completion_logprobs=[gibberish_filter.logprob_threshold - 1.0, -0.001, -0.001], + reward=1.0, + ) + + apply_filters([gibberish_filter, repetition_filter], [rollout]) + + assert rollout["stop_condition"] is None + assert rollout["filters"] == {"gibberish": True, "repetition": False} + assert rollout["is_filtered"] is True + + +def test_apply_filters_empty_list(): + rollout = _make_rollout( + completion_ids=[1, 2, 3], + completion_logprobs=[-1.0, -1.0, -1.0], + ) + apply_filters([], [rollout]) + assert rollout["filters"] == {} + assert rollout["is_filtered"] is False + assert rollout["reward"] == 1.0 + + +def test_apply_filters_mixed_batch(): + gibberish_filter = _make_gibberish_filter(enforce=True) + + clean = _make_rollout(completion_ids=[50], completion_logprobs=[-1.0], reward=1.0) + dirty = _make_rollout( + completion_ids=[120_000], completion_logprobs=[gibberish_filter.logprob_threshold - 1.0], reward=1.0 + ) + + apply_filters([gibberish_filter], [clean, dirty]) + + assert clean["reward"] == 1.0 + assert dirty["reward"] == 1.0 + assert clean["is_filtered"] is False + assert dirty["is_filtered"] is True + + +def test_apply_filters_enforced_preserves_rollout_tokens(): + gibberish_filter = _make_gibberish_filter(enforce=True) + + rollout = _make_rollout( + completion_ids=[10, 120_000, 30], + completion_logprobs=[-1.0, gibberish_filter.logprob_threshold - 1.0, -0.5], + reward=1.0, + ) + + apply_filters([gibberish_filter], [rollout]) + + assert rollout["trajectory"][0]["tokens"]["completion_ids"] == [10, 120_000, 30] + assert rollout["trajectory"][0]["tokens"]["completion_logprobs"] == [ + -1.0, + gibberish_filter.logprob_threshold - 1.0, + -0.5, + ] + assert rollout["trajectory"][0]["tokens"]["completion_mask"] == [1, 1, 1] + assert rollout["is_filtered"] is True + + +def test_apply_filters_preserves_existing_stop_condition(): + gibberish_filter = _make_gibberish_filter(enforce=True) + + rollout = _make_rollout( + completion_ids=[120_000], + completion_logprobs=[gibberish_filter.logprob_threshold - 1.0], + reward=1.0, + ) + rollout["stop_condition"] = "generation_truncated" + + apply_filters([gibberish_filter], [rollout]) + + assert rollout["stop_condition"] == "generation_truncated" + assert rollout["is_filtered"] is True + + +# --- apply_filters tests (monitor-only, enforce=False) --- + + +def test_apply_filters_monitor_only_tracks_detection(): + gibberish_filter = _make_gibberish_filter(enforce=False) + + rollout = _make_rollout( + completion_ids=[120_000], + completion_logprobs=[gibberish_filter.logprob_threshold - 1.0], + reward=1.0, + ) + + apply_filters([gibberish_filter], [rollout]) + + assert rollout["reward"] == 1.0 + assert all(m == 1 for m in rollout["trajectory"][0]["tokens"]["completion_mask"]) + assert rollout["stop_condition"] is None + assert rollout["filters"] == {"gibberish": True} + assert rollout["is_filtered"] is False + + +def test_apply_filters_monitor_only_mixed_batch(): + gibberish_filter = _make_gibberish_filter(enforce=False) + + clean = _make_rollout(completion_ids=[50], completion_logprobs=[-1.0], reward=1.0) + dirty = _make_rollout( + completion_ids=[120_000], completion_logprobs=[gibberish_filter.logprob_threshold - 1.0], reward=1.0 + ) + + apply_filters([gibberish_filter], [clean, dirty]) + + assert clean["reward"] == 1.0 + assert dirty["reward"] == 1.0 + assert clean["is_filtered"] is False + assert dirty["is_filtered"] is False diff --git a/tests/unit/orchestrator/test_orchestrator_setup.py b/tests/unit/orchestrator/test_orchestrator_setup.py new file mode 100644 index 00000000000..ff9bb5b79f9 --- /dev/null +++ b/tests/unit/orchestrator/test_orchestrator_setup.py @@ -0,0 +1,94 @@ +import asyncio +from types import SimpleNamespace +from unittest.mock import AsyncMock, MagicMock, patch + +from prime_rl.orchestrator.orchestrator import setup_rollout_inference_pool + + +def test_setup_rollout_inference_pool_uses_plain_client_for_external_teacher_rollout(): + async def run() -> None: + tokenizer = object() + config = SimpleNamespace( + teacher_rollout_model=SimpleNamespace(), + model=SimpleNamespace(renderer="auto", name="student-model"), + ) + rollout_client_config = SimpleNamespace(base_url=["https://api.pinference.ai/api/v1"]) + logger = MagicMock() + inference_pool = object() + + with ( + patch( + "prime_rl.orchestrator.orchestrator.setup_inference_pool", new=AsyncMock(return_value=inference_pool) + ), + patch("prime_rl.orchestrator.orchestrator.create_renderer") as create_renderer_mock, + ): + renderer, returned_pool = await setup_rollout_inference_pool( + config=config, + rollout_client_config=rollout_client_config, + rollout_model_name="teacher-model", + tokenizer=tokenizer, + logger=logger, + ) + + assert renderer is None + assert returned_pool is inference_pool + create_renderer_mock.assert_not_called() + + asyncio.run(run()) + + +def test_setup_rollout_inference_pool_uses_direct_renderer_client_for_local_vllm(): + async def run() -> None: + tokenizer = object() + config = SimpleNamespace( + teacher_rollout_model=None, + use_renderer=True, + use_token_client=False, + model=SimpleNamespace(name="student-model"), + renderer=SimpleNamespace( + name="qwen3_vl", + tool_parser=None, + reasoning_parser=None, + pool_size=None, + ), + ) + rollout_client_config = SimpleNamespace(base_url=["http://localhost:8000/v1"]) + logger = MagicMock() + renderer = object() + inference_pool = object() + + with ( + patch("prime_rl.orchestrator.orchestrator.create_renderer", return_value=renderer) as create_renderer_mock, + patch( + "prime_rl.orchestrator.orchestrator.setup_inference_pool", + new=AsyncMock(return_value=inference_pool), + ) as setup_pool_mock, + ): + returned_renderer, returned_pool = await setup_rollout_inference_pool( + config=config, + rollout_client_config=rollout_client_config, + rollout_model_name="student-model", + tokenizer=tokenizer, + logger=logger, + ) + + assert returned_renderer is renderer + assert returned_pool is inference_pool + create_renderer_mock.assert_called_once_with( + tokenizer, + renderer="qwen3_vl", + tool_parser=None, + reasoning_parser=None, + ) + setup_pool_mock.assert_awaited_once_with( + rollout_client_config, + model_name="student-model", + train_client_type="renderer", + eval_client_type="openai_chat_completions", + renderer_name="qwen3_vl", + tool_parser=None, + reasoning_parser=None, + renderer_pool_size=None, + ) + + asyncio.run(run()) diff --git a/tests/unit/orchestrator/test_scheduler.py b/tests/unit/orchestrator/test_scheduler.py new file mode 100644 index 00000000000..9e73b5207da --- /dev/null +++ b/tests/unit/orchestrator/test_scheduler.py @@ -0,0 +1,176 @@ +import asyncio +from pathlib import Path +from types import SimpleNamespace +from unittest.mock import AsyncMock, MagicMock, patch + +import verifiers as vf + +from prime_rl.orchestrator.scheduler import InflightRequest, Scheduler +from prime_rl.utils.async_utils import safe_cancel + + +def make_scheduler() -> Scheduler: + scheduler = Scheduler.__new__(Scheduler) + scheduler.max_async_level = 1 + scheduler.strict_async_level = False + scheduler.step = 9 + scheduler.ckpt_step = 7 + scheduler.config = SimpleNamespace(output_dir=Path("/tmp/prime-rl-test")) + scheduler.logger = MagicMock() + scheduler.checkpoint_ready = asyncio.Event() + scheduler.checkpoint_ready.set() + scheduler.lora_name = None + scheduler.model_name = "test-model" + scheduler.update_weights_time = 0 + scheduler.wait_for_ckpt_time = 0 + scheduler.inflight_requests = {} + scheduler.groups = {} + scheduler.max_off_policy_steps = 1 + scheduler.cancelled_rollouts_count = 0 + scheduler.policy_update_lock = asyncio.Lock() + scheduler.inflight_policy_update_task = None + scheduler.update_policy_task = None + scheduler.enable_policy_updates = True + return scheduler + + +def test_update_off_policy_does_not_increment_interleaved_on_policy_tasks(): + async def run() -> None: + scheduler = Scheduler.__new__(Scheduler) + scheduler.max_off_policy_steps = 1 + scheduler.cancelled_rollouts_count = 0 + scheduler.logger = MagicMock() + + client = SimpleNamespace(api_base_url="http://test") + stale_task = asyncio.create_task(asyncio.sleep(60)) + survivor_task = asyncio.create_task(asyncio.sleep(60)) + interleaved_task = None + + scheduler.inflight_requests = { + stale_task: InflightRequest(off_policy_steps=1, client_config=client, env_name="test", group_id=1), + survivor_task: InflightRequest(off_policy_steps=0, client_config=client, env_name="test", group_id=2), + } + + async def drop_group(group_id: int) -> int: + tasks_to_remove = [ + task for task, info in list(scheduler.inflight_requests.items()) if info.group_id == group_id + ] + for task in tasks_to_remove: + scheduler.inflight_requests.pop(task, None) + task.cancel() + + await asyncio.sleep(0) + + nonlocal interleaved_task + if interleaved_task is None: + interleaved_task = asyncio.create_task(asyncio.sleep(60)) + scheduler.inflight_requests[interleaved_task] = InflightRequest( + off_policy_steps=0, + client_config=client, + env_name="test", + group_id=3, + ) + return len(tasks_to_remove) + + scheduler.drop_group = drop_group + + await scheduler._update_off_policy() + + assert stale_task not in scheduler.inflight_requests + assert scheduler.inflight_requests[survivor_task].off_policy_steps == 1 + assert interleaved_task is not None + assert scheduler.inflight_requests[interleaved_task].off_policy_steps == 0 + assert scheduler.cancelled_rollouts_count == 1 + + for task in (stale_task, survivor_task, interleaved_task): + if task is not None and not task.done(): + task.cancel() + await asyncio.sleep(0) + + asyncio.run(run()) + + +def test_maybe_update_policy_reuses_inflight_update_after_cancellation(): + async def run() -> None: + scheduler = make_scheduler() + started = asyncio.Event() + release = asyncio.Event() + applied_steps: list[int] = [] + + async def update_weights(weight_dir, lora_name=None, step=0) -> None: + applied_steps.append(step) + started.set() + await release.wait() + + scheduler.inference_pool = SimpleNamespace( + update_weights=update_weights, + update_model_name=MagicMock(), + ) + scheduler._update_off_policy = AsyncMock() + + with ( + patch("prime_rl.orchestrator.scheduler.get_latest_ckpt_step", return_value=8), + patch("prime_rl.orchestrator.scheduler.wait_for_path", new=AsyncMock()), + ): + first = asyncio.create_task(scheduler.maybe_update_policy()) + await started.wait() + await safe_cancel(first) + + second = asyncio.create_task(scheduler.maybe_update_policy()) + await asyncio.sleep(0) + assert applied_steps == [8] + + release.set() + await second + + assert applied_steps == [8] + assert scheduler.ckpt_step == 8 + + asyncio.run(run()) + + +def test_stop_cancels_inflight_policy_update_task(): + async def run() -> None: + scheduler = make_scheduler() + started = asyncio.Event() + cancelled = asyncio.Event() + + async def update_weights(weight_dir, lora_name=None, step=0) -> None: + started.set() + try: + await asyncio.Future() + finally: + cancelled.set() + + scheduler.inference_pool = SimpleNamespace( + update_weights=update_weights, + update_model_name=MagicMock(), + ) + scheduler._update_off_policy = AsyncMock() + + with ( + patch("prime_rl.orchestrator.scheduler.get_latest_ckpt_step", return_value=8), + patch("prime_rl.orchestrator.scheduler.wait_for_path", new=AsyncMock()), + ): + scheduler.update_policy_task = asyncio.create_task(scheduler.maybe_update_policy()) + await started.wait() + await asyncio.wait_for(scheduler.stop(), timeout=0.2) + + assert cancelled.is_set() + assert scheduler.update_policy_task is None + assert scheduler.inflight_policy_update_task is None + + asyncio.run(run()) + + +def test_client_identity_distinguishes_base_url_and_dp_rank(): + client_a = vf.ClientConfig( + api_base_url="http://worker-a:8000/v1", + extra_headers={"X-data-parallel-rank": "0"}, + ) + client_b = vf.ClientConfig( + api_base_url="http://worker-a:8000/v1", + extra_headers={"X-data-parallel-rank": "1"}, + ) + + assert Scheduler._client_identity(client_a) != Scheduler._client_identity(client_b) diff --git a/tests/unit/orchestrator/test_sft_trajectories.py b/tests/unit/orchestrator/test_sft_trajectories.py new file mode 100644 index 00000000000..be3b30e093f --- /dev/null +++ b/tests/unit/orchestrator/test_sft_trajectories.py @@ -0,0 +1,140 @@ +from unittest.mock import MagicMock + +import verifiers as vf + +from prime_rl.orchestrator.trajectories import interleave_rollout, pretokenize_rollout_trajectory + + +class SimpleChatTokenizer: + def __init__(self): + self._tok2id: dict[str, int] = {} + self._next_id = 1 + + def _id(self, token: str) -> int: + if token not in self._tok2id: + self._tok2id[token] = self._next_id + self._next_id += 1 + return self._tok2id[token] + + def apply_chat_template(self, messages, add_generation_prompt=False, return_dict=False): + del return_dict + ids = [] + for message in messages: + role = message.get("role", "unknown") + ids.append(self._id(f"<|{role}|>")) + content = message.get("content", "") + if isinstance(content, str): + if content: + ids.append(self._id(content)) + else: + ids.append(self._id(str(content))) + if add_generation_prompt: + ids.append(self._id("<|assistant|>")) + return ids + + +def test_interleave_rollout_missing_tokens_returns_none(): + output = vf.RolloutOutput( + example_id=1, + trajectory=[ + vf.TrajectoryStep( + prompt=[{"role": "user", "content": "U1"}], + completion=[{"role": "assistant", "content": "A1"}], + response=MagicMock(), + tokens=None, + reward=None, + advantage=None, + is_truncated=False, + trajectory_id="1", + extras={}, + ) + ], + sampling_args={"temperature": 1.0}, + error=None, + ) + + assert interleave_rollout(output) is None + + +def test_pretokenize_rollout_trajectory_for_sft(): + tokenizer = SimpleChatTokenizer() + output = vf.RolloutOutput( + example_id=42, + trajectory=[ + vf.TrajectoryStep( + prompt=[{"role": "user", "content": "U1"}], + completion=[{"role": "assistant", "content": "A1"}], + response=MagicMock(), + tokens=None, + reward=None, + advantage=None, + is_truncated=False, + trajectory_id="1", + extras={}, + ), + vf.TrajectoryStep( + prompt=[ + {"role": "user", "content": "U1"}, + {"role": "assistant", "content": "A1"}, + {"role": "user", "content": "U2"}, + ], + completion=[{"role": "assistant", "content": "A2"}], + response=MagicMock(), + tokens=None, + reward=None, + advantage=None, + is_truncated=False, + trajectory_id="1", + extras={}, + ), + ], + sampling_args={"temperature": 1.0}, + error=None, + ) + + pretokenize_rollout_trajectory(output, tokenizer) + + rollouts = interleave_rollout(output) + assert rollouts is not None + assert len(rollouts) == 1 + + rollout = rollouts[0] + step1_prompt_ids = tokenizer.apply_chat_template( + [{"role": "user", "content": "U1"}], + add_generation_prompt=True, + ) + step1_full_ids = tokenizer.apply_chat_template( + [{"role": "user", "content": "U1"}, {"role": "assistant", "content": "A1"}], + add_generation_prompt=False, + ) + step2_prompt_ids = tokenizer.apply_chat_template( + [ + {"role": "user", "content": "U1"}, + {"role": "assistant", "content": "A1"}, + {"role": "user", "content": "U2"}, + ], + add_generation_prompt=True, + ) + step2_full_ids = tokenizer.apply_chat_template( + [ + {"role": "user", "content": "U1"}, + {"role": "assistant", "content": "A1"}, + {"role": "user", "content": "U2"}, + {"role": "assistant", "content": "A2"}, + ], + add_generation_prompt=False, + ) + + prefix_len_1 = len(step1_prompt_ids) + prefix_len_2 = len(step2_prompt_ids) + step1_completion_ids = step1_full_ids[prefix_len_1:] + step2_completion_ids = step2_full_ids[prefix_len_2:] + step1_prefix = step1_prompt_ids + step1_completion_ids + step2_new_prompt_ids = step2_prompt_ids[len(step1_prefix) :] + + assert rollout.prompt_ids == step1_prompt_ids + assert rollout.completion_ids == step1_completion_ids + step2_new_prompt_ids + step2_completion_ids + assert rollout.completion_mask == ( + [True] * len(step1_completion_ids) + [False] * len(step2_new_prompt_ids) + [True] * len(step2_completion_ids) + ) + assert rollout.completion_logprobs == [0.0] * len(rollout.completion_ids) diff --git a/tests/unit/orchestrator/test_teacher_logprobs.py b/tests/unit/orchestrator/test_teacher_logprobs.py new file mode 100644 index 00000000000..fb5a04c3b0a --- /dev/null +++ b/tests/unit/orchestrator/test_teacher_logprobs.py @@ -0,0 +1,79 @@ +import asyncio +import json + +import httpx +import verifiers as vf + +from prime_rl.orchestrator import utils as orchestrator_utils +from prime_rl.transport import TrainingSample + + +class _FakeOpenAIClient: + """Stand-in for ``AsyncOpenAI`` that captures the sole ``.post()`` call and + returns a synthesized ``httpx.Response`` so ``cast_to=httpx.Response`` is + handed back verbatim, mirroring the real SDK's short-circuit at + ``AsyncAPIClient._process_response``.""" + + def __init__(self, payload: dict): + # Match what AsyncOpenAI exposes — utils.py reads ``str(client.base_url)``. + self.base_url = "http://fake-host:8000/v1" + self._payload = payload + self.calls: list[dict] = [] + + async def post(self, url, *, cast_to, body): + self.calls.append({"url": url, "cast_to": cast_to, "body": body}) + request = httpx.Request("POST", url, json=body) + return httpx.Response( + status_code=200, + content=json.dumps(self._payload).encode(), + request=request, + ) + + +def test_compute_teacher_logprobs_uses_inference_generate(monkeypatch): + async def _run(): + fake_client = _FakeOpenAIClient( + { + "request_id": "gen-test", + "choices": [], + # Upstream wire shape: list[dict[token_id, Logprob] | None] + "prompt_logprobs": [None, {"11": {"logprob": -0.7}}, {"12": {"logprob": -0.3}}], + "kv_transfer_params": None, + } + ) + monkeypatch.setattr(orchestrator_utils, "setup_openai_client", lambda _: fake_client) + + sample = TrainingSample( + prompt_ids=[1], + prompt_mask=[True], + completion_ids=[2, 3], + completion_mask=[True, True], + completion_logprobs=[-0.1, -0.2], + completion_temperatures=[1.0, 1.0], + ) + + result = await orchestrator_utils.compute_teacher_logprobs( + clients=[vf.ClientConfig()], + model_name="teacher-model", + samples=[sample], + ) + + assert result == [[0.0, -0.7, -0.3]] + assert fake_client.calls == [ + { + "url": "http://fake-host:8000/inference/v1/generate", + "cast_to": httpx.Response, + "body": { + "model": "teacher-model", + "token_ids": [1, 2, 3], + "sampling_params": { + "max_tokens": 1, + "temperature": 1.0, + "top_p": 1.0, + "prompt_logprobs": 1, + }, + }, + } + ] + + asyncio.run(_run()) diff --git a/tests/unit/orchestrator/test_trajectories.py b/tests/unit/orchestrator/test_trajectories.py index 4f44d88c101..6fa169760ca 100644 --- a/tests/unit/orchestrator/test_trajectories.py +++ b/tests/unit/orchestrator/test_trajectories.py @@ -1,14 +1,66 @@ +import base64 +from io import BytesIO from unittest.mock import MagicMock +import numpy as np import pytest import verifiers as vf +from PIL import Image -from prime_rl.orchestrator.trajectories import branch_rollout, interleave_rollout +from prime_rl.orchestrator.trajectories import ( + VLMImageCache, + _align_routed_experts, + _deserialize_tool_calls, + _extract_images_from_examples, + _extract_images_from_messages, + _ImageStore, + build_vlm_image_cache, + interleave_rollout, +) + + +def _pixels(data: list[list[float]]) -> tuple[bytes, list[int]]: + """Convert pixel values list to (bytes, shape) for test cache data.""" + arr = np.array(data, dtype=np.float32) + return arr.tobytes(), list(arr.shape) + + +def _decode_pixels(pixel_bytes: bytes, shape: list[int]) -> list[list[float]]: + """Decode raw pixel bytes back to nested list for assertions.""" + return np.frombuffer(pixel_bytes, dtype=np.float32).reshape(shape).tolist() + + +def test_deserialize_tool_calls_does_not_inject_missing_key(): + messages = [{"role": "assistant", "content": "hello"}] + + deserialized = _deserialize_tool_calls(messages) + + assert "tool_calls" not in deserialized[0] + + +def test_deserialize_tool_calls_parses_arguments_when_present(): + messages = [ + { + "role": "assistant", + "tool_calls": [ + { + "id": "1", + "type": "function", + "function": {"name": "lookup", "arguments": '{"x": 1}'}, + } + ], + } + ] + + deserialized = _deserialize_tool_calls(messages) + + assert deserialized[0]["tool_calls"][0]["function"]["arguments"] == {"x": 1} @pytest.fixture -def single_step_trajectory_state(): - state = vf.State( +def single_step_trajectory_output(): + output = vf.RolloutOutput( + example_id=0, trajectory=[ vf.TrajectoryStep( prompt=[{"role": "user", "content": "U1"}], @@ -25,17 +77,21 @@ def single_step_trajectory_state(): ), reward=None, advantage=None, + is_truncated=False, + trajectory_id="1", extras={}, ) ], + sampling_args={"temperature": 1.0}, error=None, ) - return state + return output @pytest.fixture -def multi_step_trajectory_state(): - state = vf.State( +def multi_step_trajectory_output(): + output = vf.RolloutOutput( + example_id=0, trajectory=[ vf.TrajectoryStep( prompt=[{"role": "user", "content": "U1"}], @@ -52,6 +108,8 @@ def multi_step_trajectory_state(): ), reward=None, advantage=None, + is_truncated=False, + trajectory_id="1", extras={}, ), vf.TrajectoryStep( @@ -73,17 +131,21 @@ def multi_step_trajectory_state(): ), reward=None, advantage=None, + is_truncated=False, + trajectory_id="1", extras={}, ), ], + sampling_args={"temperature": 1.0}, error=None, ) - return state + return output @pytest.fixture -def multi_step_trajectory_with_tool_calls(): - state = vf.State( +def multi_step_trajectory_with_tool_calls_output(): + output = vf.RolloutOutput( + example_id=0, trajectory=[ vf.TrajectoryStep( prompt=[{"role": "user", "content": "U1"}], @@ -100,6 +162,8 @@ def multi_step_trajectory_with_tool_calls(): ), reward=None, advantage=None, + is_truncated=False, + trajectory_id="1", extras={}, ), vf.TrajectoryStep( @@ -121,6 +185,8 @@ def multi_step_trajectory_with_tool_calls(): ), reward=None, advantage=None, + is_truncated=False, + trajectory_id="1", extras={}, ), ], @@ -128,25 +194,135 @@ def multi_step_trajectory_with_tool_calls(): advantage=None, stop_condition=None, metrics={"has_error": 0.0, "tool_calls": 1.0}, + sampling_args={"temperature": 1.0}, error=None, ) - return state + return output -def test_branching_rollout_single_step_trajectory(single_step_trajectory_state): - rollouts = branch_rollout(single_step_trajectory_state) +@pytest.fixture +def multi_step_trajectory_extension_never_holds(): + """ + 2-step trajectory where extension NEVER holds (step 2 has completely different tokens). + This simulates e.g. a chat template that re-renders the entire conversation differently. + """ + output = vf.RolloutOutput( + example_id=0, + trajectory=[ + vf.TrajectoryStep( + prompt=[{"role": "user", "content": "U1"}], + completion=[{"role": "assistant", "content": "A1"}], + response=MagicMock(), + tokens=vf.TrajectoryStepTokens( + prompt_ids=[1, 2], + prompt_mask=[0, 0], + completion_ids=[3, 4], + completion_mask=[1, 1], + completion_logprobs=[-0.1, -0.2], + overlong_prompt=False, + is_truncated=False, + ), + reward=None, + advantage=None, + is_truncated=False, + trajectory_id="1", + extras={}, + ), + vf.TrajectoryStep( + prompt=[ + {"role": "user", "content": "U1"}, + {"role": "assistant", "content": "A1"}, + {"role": "user", "content": "U2"}, + ], + completion=[{"role": "assistant", "content": "A2"}], + response=MagicMock(), + tokens=vf.TrajectoryStepTokens( + # Different tokens - extension breaks (e.g. thinking was stripped) + prompt_ids=[10, 20, 30, 40, 50, 60], + prompt_mask=[0, 0, 0, 0, 0, 0], + completion_ids=[7, 8], + completion_mask=[1, 1], + completion_logprobs=[-0.3, -0.4], + overlong_prompt=False, + is_truncated=False, + ), + reward=None, + advantage=None, + is_truncated=False, + trajectory_id="1", + extras={}, + ), + ], + sampling_args={"temperature": 1.0}, + error=None, + ) + return output + - assert len(rollouts) == 1 - rollout = rollouts[0] - assert rollout.prompt_ids == [1, 2] - assert rollout.prompt_mask == [False, False] - assert rollout.completion_ids == [3, 4] - assert rollout.completion_mask == [True, True] - assert rollout.completion_logprobs == [-0.1, -0.2] +@pytest.fixture +def multi_step_trajectory_with_tool_calls_extension_never_holds(): + """2-step trajectory with tool calls where extension NEVER holds.""" + output = vf.RolloutOutput( + example_id=0, + trajectory=[ + vf.TrajectoryStep( + prompt=[{"role": "user", "content": "U1"}], + completion=[{"role": "assistant", "content": "A1 + TC1"}], + response=MagicMock(), + tokens=vf.TrajectoryStepTokens( + prompt_ids=[1, 2], + prompt_mask=[0, 0], + completion_ids=[3, 4], + completion_mask=[1, 1], + completion_logprobs=[-0.1, -0.2], + overlong_prompt=False, + is_truncated=False, + ), + reward=None, + advantage=None, + extras={}, + is_truncated=False, + trajectory_id="1", + ), + vf.TrajectoryStep( + prompt=[ + {"role": "user", "content": "U1"}, + {"role": "assistant", "content": "A1 + TC1"}, + {"role": "tool", "tool_call_id": "TR1", "content": "TR1"}, + ], + completion=[{"role": "assistant", "content": "A2 + TC2"}], + response=MagicMock(), + tokens=vf.TrajectoryStepTokens( + # Different tokens - extension breaks + prompt_ids=[10, 20, 30, 40, 50, 60], + prompt_mask=[0, 0, 0, 0, 0, 0], + completion_ids=[7, 8], + completion_mask=[1, 1], + completion_logprobs=[-0.3, -0.4], + overlong_prompt=False, + is_truncated=False, + ), + reward=None, + advantage=None, + extras={}, + is_truncated=False, + trajectory_id="1", + ), + ], + reward=1.0, + advantage=None, + stop_condition=None, + sampling_args={"temperature": 1.0}, + metrics={"has_error": 0.0, "tool_calls": 1.0}, + error=None, + ) + return output -def test_branching_rollout_multi_step_trajectory(multi_step_trajectory_state): - rollouts = branch_rollout(multi_step_trajectory_state) +def test_branching_equivalent_multi_step_trajectory(multi_step_trajectory_extension_never_holds): + """When extension never holds, each step becomes its own sample (same as old branching).""" + rollouts = interleave_rollout(multi_step_trajectory_extension_never_holds) + assert rollouts is not None assert len(rollouts) == 2 # first step @@ -156,18 +332,24 @@ def test_branching_rollout_multi_step_trajectory(multi_step_trajectory_state): assert rollout.completion_ids == [3, 4] assert rollout.completion_mask == [True, True] assert rollout.completion_logprobs == [-0.1, -0.2] + assert rollout.completion_temperatures == [1.0, 1.0] # second step rollout = rollouts[1] - assert rollout.prompt_ids == [1, 2, 3, 4, 5, 6] + assert rollout.prompt_ids == [10, 20, 30, 40, 50, 60] assert rollout.prompt_mask == [False, False, False, False, False, False] assert rollout.completion_ids == [7, 8] assert rollout.completion_mask == [True, True] assert rollout.completion_logprobs == [-0.3, -0.4] + assert rollout.completion_temperatures == [1.0, 1.0] -def test_branching_rollout_multi_step_trajectory_with_tool_calls(multi_step_trajectory_with_tool_calls): - rollouts = branch_rollout(multi_step_trajectory_with_tool_calls) +def test_branching_equivalent_multi_step_trajectory_with_tool_calls( + multi_step_trajectory_with_tool_calls_extension_never_holds, +): + """When extension never holds (with tool calls), same as old branching.""" + rollouts = interleave_rollout(multi_step_trajectory_with_tool_calls_extension_never_holds) + assert rollouts is not None assert len(rollouts) == 2 # first step @@ -177,18 +359,21 @@ def test_branching_rollout_multi_step_trajectory_with_tool_calls(multi_step_traj assert rollout.completion_ids == [3, 4] assert rollout.completion_mask == [True, True] assert rollout.completion_logprobs == [-0.1, -0.2] + assert rollout.completion_temperatures == [1.0, 1.0] # second step rollout = rollouts[1] - assert rollout.prompt_ids == [1, 2, 3, 4, 5, 6] + assert rollout.prompt_ids == [10, 20, 30, 40, 50, 60] assert rollout.prompt_mask == [False, False, False, False, False, False] assert rollout.completion_ids == [7, 8] assert rollout.completion_mask == [True, True] assert rollout.completion_logprobs == [-0.3, -0.4] + assert rollout.completion_temperatures == [1.0, 1.0] -def test_interleave_rollout_single_step_trajectory(single_step_trajectory_state): - rollouts = interleave_rollout(single_step_trajectory_state) +def test_interleave_rollout_single_step_trajectory(single_step_trajectory_output): + rollouts = interleave_rollout(single_step_trajectory_output) + assert rollouts is not None assert len(rollouts) == 1 rollout = rollouts[0] @@ -197,10 +382,12 @@ def test_interleave_rollout_single_step_trajectory(single_step_trajectory_state) assert rollout.completion_ids == [3, 4] assert rollout.completion_mask == [True, True] assert rollout.completion_logprobs == [-0.1, -0.2] + assert rollout.completion_temperatures == [1.0, 1.0] -def test_interleave_rollout_multi_step_trajectory(multi_step_trajectory_state): - rollouts = interleave_rollout(multi_step_trajectory_state) +def test_interleave_rollout_multi_step_trajectory(multi_step_trajectory_output): + rollouts = interleave_rollout(multi_step_trajectory_output) + assert rollouts is not None assert len(rollouts) == 1 rollout = rollouts[0] @@ -209,10 +396,13 @@ def test_interleave_rollout_multi_step_trajectory(multi_step_trajectory_state): assert rollout.completion_ids == [3, 4, 5, 6, 7, 8] assert rollout.completion_mask == [True, True, False, False, True, True] assert rollout.completion_logprobs == [-0.1, -0.2, 0, 0, -0.3, -0.4] + # Temperatures: 2 completion tokens at temp 1.0, then 2 prompt tokens at temp 1.0, then 2 completion tokens at temp 1.0 + assert rollout.completion_temperatures == [1.0, 1.0, 1.0, 1.0, 1.0, 1.0] -def test_interleave_rollout_multi_step_trajectory_with_tool_calls(multi_step_trajectory_with_tool_calls): - rollouts = interleave_rollout(multi_step_trajectory_with_tool_calls) +def test_interleave_rollout_multi_step_trajectory_with_tool_calls(multi_step_trajectory_with_tool_calls_output): + rollouts = interleave_rollout(multi_step_trajectory_with_tool_calls_output) + assert rollouts is not None assert len(rollouts) == 1 rollout = rollouts[0] @@ -221,3 +411,1780 @@ def test_interleave_rollout_multi_step_trajectory_with_tool_calls(multi_step_tra assert rollout.completion_ids == [3, 4, 5, 6, 7, 8] assert rollout.completion_mask == [True, True, False, False, True, True] assert rollout.completion_logprobs == [-0.1, -0.2, 0, 0, -0.3, -0.4] + # Temperatures: 2 completion tokens at temp 1.0, then 2 prompt tokens at temp 1.0, then 2 completion tokens at temp 1.0 + assert rollout.completion_temperatures == [1.0, 1.0, 1.0, 1.0, 1.0, 1.0] + + +@pytest.fixture +def five_step_trajectory_with_extension_break(): + """ + 5-step trajectory where extension property breaks at step 4. + + Steps 1-3: extension holds (tokens grow by appending) + Step 4: extension breaks (completely different prefix, e.g. context compaction) + Steps 4-5: extension holds again + + Expected: 2 samples (steps 1-3 merged, steps 4-5 merged) + """ + output = vf.RolloutOutput( + example_id=0, + trajectory=[ + # Step 1: initial prompt and completion + vf.TrajectoryStep( + prompt=[{"role": "user", "content": "U1"}], + completion=[{"role": "assistant", "content": "A1"}], + response=MagicMock(), + tokens=vf.TrajectoryStepTokens( + prompt_ids=[1, 2], + prompt_mask=[0, 0], + completion_ids=[3, 4], + completion_mask=[1, 1], + completion_logprobs=[-0.1, -0.2], + overlong_prompt=False, + is_truncated=False, + ), + reward=None, + advantage=None, + extras={}, + is_truncated=False, + trajectory_id="1", + ), + # Step 2: extends step 1 (prefix [1,2,3,4] matches) + vf.TrajectoryStep( + prompt=[ + {"role": "user", "content": "U1"}, + {"role": "assistant", "content": "A1"}, + {"role": "user", "content": "U2"}, + ], + completion=[{"role": "assistant", "content": "A2"}], + response=MagicMock(), + tokens=vf.TrajectoryStepTokens( + prompt_ids=[1, 2, 3, 4, 5, 6], + prompt_mask=[0, 0, 0, 0, 0, 0], + completion_ids=[7, 8], + completion_mask=[1, 1], + completion_logprobs=[-0.3, -0.4], + overlong_prompt=False, + is_truncated=False, + ), + reward=None, + advantage=None, + extras={}, + is_truncated=False, + trajectory_id="1", + ), + # Step 3: extends step 2 (prefix [1,2,3,4,5,6,7,8] matches) + vf.TrajectoryStep( + prompt=[ + {"role": "user", "content": "U1"}, + {"role": "assistant", "content": "A1"}, + {"role": "user", "content": "U2"}, + {"role": "assistant", "content": "A2"}, + {"role": "user", "content": "U3"}, + ], + completion=[{"role": "assistant", "content": "A3"}], + response=MagicMock(), + tokens=vf.TrajectoryStepTokens( + prompt_ids=[1, 2, 3, 4, 5, 6, 7, 8, 9, 10], + prompt_mask=[0, 0, 0, 0, 0, 0, 0, 0, 0, 0], + completion_ids=[11, 12], + completion_mask=[1, 1], + completion_logprobs=[-0.5, -0.6], + overlong_prompt=False, + is_truncated=False, + ), + reward=None, + advantage=None, + extras={}, + is_truncated=False, + trajectory_id="1", + ), + # Step 4: EXTENSION BREAKS - different prefix (e.g. thinking stripped, context compacted) + vf.TrajectoryStep( + prompt=[ + {"role": "user", "content": "U1"}, + {"role": "assistant", "content": "A1"}, # thinking stripped + {"role": "user", "content": "U2"}, + {"role": "assistant", "content": "A2"}, + {"role": "user", "content": "U3"}, + {"role": "assistant", "content": "A3"}, + {"role": "user", "content": "U4"}, + ], + completion=[{"role": "assistant", "content": "A4"}], + response=MagicMock(), + tokens=vf.TrajectoryStepTokens( + prompt_ids=[100, 101, 102, 103], # completely different tokens (re-rendered) + prompt_mask=[0, 0, 0, 0], + completion_ids=[104, 105], + completion_mask=[1, 1], + completion_logprobs=[-0.7, -0.8], + overlong_prompt=False, + is_truncated=False, + ), + reward=None, + advantage=None, + extras={}, + is_truncated=False, + trajectory_id="1", + ), + # Step 5: extends step 4 (prefix [100,101,102,103,104,105] matches) + vf.TrajectoryStep( + prompt=[ + {"role": "user", "content": "U1"}, + {"role": "assistant", "content": "A1"}, + {"role": "user", "content": "U2"}, + {"role": "assistant", "content": "A2"}, + {"role": "user", "content": "U3"}, + {"role": "assistant", "content": "A3"}, + {"role": "user", "content": "U4"}, + {"role": "assistant", "content": "A4"}, + {"role": "user", "content": "U5"}, + ], + completion=[{"role": "assistant", "content": "A5"}], + response=MagicMock(), + tokens=vf.TrajectoryStepTokens( + prompt_ids=[100, 101, 102, 103, 104, 105, 106, 107], + prompt_mask=[0, 0, 0, 0, 0, 0, 0, 0], + completion_ids=[108, 109], + completion_mask=[1, 1], + completion_logprobs=[-0.9, -1.0], + overlong_prompt=False, + is_truncated=False, + ), + reward=None, + advantage=None, + extras={}, + is_truncated=False, + trajectory_id="1", + ), + ], + sampling_args={"temperature": 1.0}, + error=None, + ) + return output + + +def test_interleave_rollout_extension_break_creates_multiple_samples(five_step_trajectory_with_extension_break): + """ + When extension property breaks mid-trajectory, interleave_rollout should: + - Merge steps 1-3 into first sample (extension held) + - Start new sample at step 4 (extension broke) + - Merge steps 4-5 into second sample (extension held again) + """ + rollouts = interleave_rollout(five_step_trajectory_with_extension_break) + + assert rollouts is not None + assert len(rollouts) == 2, "Should produce 2 samples when extension breaks at step 4" + + # First sample: steps 1-3 merged + sample1 = rollouts[0] + assert sample1.prompt_ids == [1, 2] + assert sample1.prompt_mask == [False, False] + # completion_ids: step1 completion [3,4] + step2 new prompt [5,6] + step2 completion [7,8] + # + step3 new prompt [9,10] + step3 completion [11,12] + assert sample1.completion_ids == [3, 4, 5, 6, 7, 8, 9, 10, 11, 12] + # completion_mask: step1 [T,T] + step2 prompt [F,F] + step2 completion [T,T] + # + step3 prompt [F,F] + step3 completion [T,T] + assert sample1.completion_mask == [True, True, False, False, True, True, False, False, True, True] + assert sample1.completion_logprobs == [-0.1, -0.2, 0, 0, -0.3, -0.4, 0, 0, -0.5, -0.6] + + # Second sample: steps 4-5 merged (fresh start after extension break) + sample2 = rollouts[1] + assert sample2.prompt_ids == [100, 101, 102, 103] + assert sample2.prompt_mask == [False, False, False, False] + # completion_ids: step4 completion [104,105] + step5 new prompt [106,107] + step5 completion [108,109] + assert sample2.completion_ids == [104, 105, 106, 107, 108, 109] + # completion_mask: step4 [T,T] + step5 prompt [F,F] + step5 completion [T,T] + assert sample2.completion_mask == [True, True, False, False, True, True] + assert sample2.completion_logprobs == [-0.7, -0.8, 0, 0, -0.9, -1.0] + + +@pytest.fixture +def interleaved_agents_trajectory(): + """ + Trajectory with interleaved agents: agent1 steps, then agent2 step, then agent1 continues. + This tests multi-prefix tracking where agent1-step3 should merge back with agent1 sample. + + agent1-step1: prompt=[1,2], completion=[3,4] + agent1-step2: prompt=[1,2,3,4,5,6], completion=[7,8] (extends agent1-step1) + agent2-step1: prompt=[100,101], completion=[102,103] (different prefix, new sample) + agent1-step3: prompt=[1,2,3,4,5,6,7,8,9,10], completion=[11,12] (extends agent1-step2!) + """ + output = vf.RolloutOutput( + example_id=1, + task="test", + trajectory=[ + # agent1-step1 + vf.TrajectoryStep( + prompt="agent1 turn 1", + completion="response 1", + response=None, + tokens=vf.TrajectoryStepTokens( + prompt_ids=[1, 2], + prompt_mask=[0, 0], + completion_ids=[3, 4], + completion_mask=[1, 1], + completion_logprobs=[-0.1, -0.2], + overlong_prompt=False, + is_truncated=False, + ), + reward=None, + advantage=None, + is_truncated=False, + trajectory_id="traj1", + extras={}, + ), + # agent1-step2 (extends agent1-step1) + vf.TrajectoryStep( + prompt="agent1 turn 2", + completion="response 2", + response=None, + tokens=vf.TrajectoryStepTokens( + prompt_ids=[1, 2, 3, 4, 5, 6], + prompt_mask=[0, 0, 0, 0, 0, 0], + completion_ids=[7, 8], + completion_mask=[1, 1], + completion_logprobs=[-0.3, -0.4], + overlong_prompt=False, + is_truncated=False, + ), + reward=None, + advantage=None, + is_truncated=False, + trajectory_id="traj1", + extras={}, + ), + # agent2-step1 (different prefix, starts new sample) + vf.TrajectoryStep( + prompt="agent2 turn 1", + completion="agent2 response", + response=None, + tokens=vf.TrajectoryStepTokens( + prompt_ids=[100, 101], + prompt_mask=[0, 0], + completion_ids=[102, 103], + completion_mask=[1, 1], + completion_logprobs=[-0.5, -0.6], + overlong_prompt=False, + is_truncated=False, + ), + reward=None, + advantage=None, + is_truncated=False, + trajectory_id="traj2", + extras={}, + ), + # agent1-step3 (extends agent1-step2, should merge back!) + vf.TrajectoryStep( + prompt="agent1 turn 3", + completion="response 3", + response=None, + tokens=vf.TrajectoryStepTokens( + prompt_ids=[1, 2, 3, 4, 5, 6, 7, 8, 9, 10], + prompt_mask=[0, 0, 0, 0, 0, 0, 0, 0, 0, 0], + completion_ids=[11, 12], + completion_mask=[1, 1], + completion_logprobs=[-0.7, -0.8], + overlong_prompt=False, + is_truncated=False, + ), + reward=None, + advantage=None, + is_truncated=False, + trajectory_id="traj1", + extras={}, + ), + ], + sampling_args={"temperature": 1.0}, + error=None, + ) + return output + + +def test_interleave_rollout_interleaved_agents(interleaved_agents_trajectory): + """ + When agents are interleaved (agent1, agent1, agent2, agent1), the multi-prefix + tracking should merge agent1-step3 back into the agent1 sample, not start a new one. + """ + rollouts = interleave_rollout(interleaved_agents_trajectory) + + assert rollouts is not None + assert len(rollouts) == 2, "Should produce 2 samples (agent1 merged, agent2 separate)" + + # First sample: agent1 steps 1, 2, 3 merged + agent1_sample = rollouts[0] + assert agent1_sample.prompt_ids == [1, 2] + assert agent1_sample.prompt_mask == [False, False] + # completion_ids: step1 [3,4] + step2 new prompt [5,6] + step2 completion [7,8] + # + step3 new prompt [9,10] + step3 completion [11,12] + assert agent1_sample.completion_ids == [3, 4, 5, 6, 7, 8, 9, 10, 11, 12] + assert agent1_sample.completion_mask == [True, True, False, False, True, True, False, False, True, True] + assert agent1_sample.completion_logprobs == [-0.1, -0.2, 0, 0, -0.3, -0.4, 0, 0, -0.7, -0.8] + + # Second sample: agent2 step 1 only + agent2_sample = rollouts[1] + assert agent2_sample.prompt_ids == [100, 101] + assert agent2_sample.prompt_mask == [False, False] + assert agent2_sample.completion_ids == [102, 103] + assert agent2_sample.completion_mask == [True, True] + assert agent2_sample.completion_logprobs == [-0.5, -0.6] + + +# ============================================================================= +# VLM Multi-Turn Tests +# ============================================================================= + + +def _create_test_image(color: str = "red") -> str: + """Create a small test image and return its base64 data URL.""" + colors = {"red": (255, 0, 0), "green": (0, 255, 0), "blue": (0, 0, 255)} + img = Image.new("RGB", (10, 10), colors.get(color, (255, 255, 255))) + buffer = BytesIO() + img.save(buffer, format="PNG") + b64 = base64.b64encode(buffer.getvalue()).decode() + return f"data:image/png;base64,{b64}" + + +def _create_image_message(image_url: str, text: str = "What is this?") -> dict: + """Create an OpenAI-style user message with an image.""" + return { + "role": "user", + "content": [ + {"type": "image_url", "image_url": {"url": image_url}}, + {"type": "text", "text": text}, + ], + } + + +def test_extract_images_from_messages_no_images(): + messages = [{"role": "user", "content": "Hello"}] + images = _extract_images_from_messages(messages) + assert images == [] + + +def test_extract_images_from_messages_single_image(): + image_url = _create_test_image("red") + messages = [_create_image_message(image_url)] + images = _extract_images_from_messages(messages) + assert len(images) == 1 + assert isinstance(images[0][0], Image.Image) + + +def test_extract_images_from_messages_multiple_images(): + messages = [ + _create_image_message(_create_test_image("red")), + {"role": "assistant", "content": "I see a red image"}, + _create_image_message(_create_test_image("green")), + ] + images = _extract_images_from_messages(messages) + assert len(images) == 2 + + +def test_extract_images_from_examples_single_turn(): + image_url = _create_test_image("red") + output = vf.RolloutOutput( + example_id=1, + trajectory=[ + vf.TrajectoryStep( + prompt=[_create_image_message(image_url)], + completion=[{"role": "assistant", "content": "A red square"}], + response=MagicMock(), + tokens=MagicMock(), + reward=None, + advantage=None, + is_truncated=False, + trajectory_id="1", + extras={}, + ), + ], + sampling_args={"temperature": 1.0}, + error=None, + ) + + all_images, images_per_step = _extract_images_from_examples([(1, output)]) + + assert len(all_images) == 1 + assert images_per_step == {1: [[0]]} # step 0 has image at index 0 + + +def test_extract_images_from_examples_multi_turn_new_image_each_turn(): + """Test that new images in later turns are correctly extracted.""" + red_url = _create_test_image("red") + green_url = _create_test_image("green") + + output = vf.RolloutOutput( + example_id=1, + trajectory=[ + # Turn 1: just the red image + vf.TrajectoryStep( + prompt=[_create_image_message(red_url, "What color is this?")], + completion=[{"role": "assistant", "content": "Red"}], + response=MagicMock(), + tokens=MagicMock(), + reward=None, + advantage=None, + is_truncated=False, + trajectory_id="1", + extras={}, + ), + # Turn 2: cumulative prompt with red image + green image + vf.TrajectoryStep( + prompt=[ + _create_image_message(red_url, "What color is this?"), + {"role": "assistant", "content": "Red"}, + _create_image_message(green_url, "And this one?"), + ], + completion=[{"role": "assistant", "content": "Green"}], + response=MagicMock(), + tokens=MagicMock(), + reward=None, + advantage=None, + is_truncated=False, + trajectory_id="1", + extras={}, + ), + ], + sampling_args={"temperature": 1.0}, + error=None, + ) + + all_images, images_per_step = _extract_images_from_examples([(1, output)]) + + assert len(all_images) == 2 # 2 unique images total + assert images_per_step == {1: [[0], [0, 1]]} # step 0: [red], step 1: [red, green] + + +def test_extract_images_from_examples_multi_turn_no_new_images(): + """Test turns where no new images are added.""" + red_url = _create_test_image("red") + + output = vf.RolloutOutput( + example_id=1, + trajectory=[ + vf.TrajectoryStep( + prompt=[_create_image_message(red_url)], + completion=[{"role": "assistant", "content": "Red"}], + response=MagicMock(), + tokens=MagicMock(), + reward=None, + advantage=None, + is_truncated=False, + trajectory_id="1", + extras={}, + ), + # Turn 2: same image, no new ones + vf.TrajectoryStep( + prompt=[ + _create_image_message(red_url), + {"role": "assistant", "content": "Red"}, + {"role": "user", "content": "Are you sure?"}, # text only + ], + completion=[{"role": "assistant", "content": "Yes"}], + response=MagicMock(), + tokens=MagicMock(), + reward=None, + advantage=None, + is_truncated=False, + trajectory_id="1", + extras={}, + ), + ], + sampling_args={"temperature": 1.0}, + error=None, + ) + + all_images, images_per_step = _extract_images_from_examples([(1, output)]) + + assert len(all_images) == 1 # Only 1 unique image (deduped) + assert images_per_step == {1: [[0], [0]]} # both steps reference the same image + + +def test_extract_images_from_examples_step_with_fewer_images_than_prior_steps(): + """Test that image counts reflect the prompt's actual images, not a monotonically increasing total. + + When a later step's prompt contains fewer images than the cumulative total from prior steps + (i.e. the prompt is not strictly cumulative), the count for that step should match + the number of images actually present in that step's prompt. + """ + red_url = _create_test_image("red") + green_url = _create_test_image("green") + + output = vf.RolloutOutput( + example_id=1, + trajectory=[ + # Step 0: red only + vf.TrajectoryStep( + prompt=[_create_image_message(red_url)], + completion=[], + response=MagicMock(), + tokens=MagicMock(), + reward=None, + advantage=None, + is_truncated=False, + trajectory_id="1", + extras={}, + ), + # Step 1: cumulative — red + green + vf.TrajectoryStep( + prompt=[_create_image_message(red_url), _create_image_message(green_url)], + completion=[], + response=MagicMock(), + tokens=MagicMock(), + reward=None, + advantage=None, + is_truncated=False, + trajectory_id="1", + extras={}, + ), + # Step 2: only green — fewer images than cumulative total + vf.TrajectoryStep( + prompt=[_create_image_message(green_url)], + completion=[], + response=MagicMock(), + tokens=MagicMock(), + reward=None, + advantage=None, + is_truncated=False, + trajectory_id="1", + extras={}, + ), + ], + sampling_args={"temperature": 1.0}, + error=None, + ) + + _, images_per_step = _extract_images_from_examples([(1, output)]) + + # Step 0: [red] → index 0; Step 1: [red, green] → indices [0, 1]; Step 2: [green] → index [1] + assert images_per_step == {1: [[0], [0, 1], [1]]} + + +def test_vlm_image_cache_get_for_step(): + cache_data = { + 1: [ + (*_pixels([[1.0, 2.0]]), [[1, 2, 3]]), # Step 0: 1 image + (*_pixels([[1.0, 2.0], [3.0, 4.0]]), [[1, 2, 3], [1, 4, 4]]), # Step 1: 2 images cumulative + ], + } + cache = VLMImageCache(cache_data, num_unique_examples=1, extract_time=0.0, preprocess_time=0.0) + + # Step 0 should have 1 image + pv, shape, grid = cache.get_for_step(1, 0) + assert _decode_pixels(pv, shape) == [[1.0, 2.0]] + assert grid == [[1, 2, 3]] + + # Step 1 should have 2 images + pv, shape, grid = cache.get_for_step(1, 1) + assert _decode_pixels(pv, shape) == [[1.0, 2.0], [3.0, 4.0]] + assert grid == [[1, 2, 3], [1, 4, 4]] + + +def test_vlm_image_cache_get_all(): + cache_data = { + 1: [ + (*_pixels([[1.0]]), [[1, 2, 3]]), + (*_pixels([[1.0], [2.0]]), [[1, 2, 3], [1, 4, 4]]), + ], + } + cache = VLMImageCache(cache_data, num_unique_examples=1, extract_time=0.0, preprocess_time=0.0) + + # get_all should return the last step's data + pv, shape, grid = cache.get_all(1) + assert _decode_pixels(pv, shape) == [[1.0], [2.0]] + assert grid == [[1, 2, 3], [1, 4, 4]] + + +def test_vlm_image_cache_step_out_of_range(): + cache_data = { + 1: [ + (*_pixels([[1.0]]), [[1, 2, 3]]), + ], + } + cache = VLMImageCache(cache_data, num_unique_examples=1, extract_time=0.0, preprocess_time=0.0) + + pv, shape, grid = cache.get_for_step(1, 2) + assert pv is None + assert shape is None + assert grid is None + + +def test_vlm_image_cache_missing_example(): + cache = VLMImageCache({}, num_unique_examples=0, extract_time=0.0, preprocess_time=0.0) + + pv, shape, grid = cache.get_for_step(999, 0) + assert pv is None + assert shape is None + assert grid is None + + pv, shape, grid = cache.get_all(999) + assert pv is None + assert shape is None + assert grid is None + + +def test_interleave_rollout_with_vlm_cache(): + """Test that interleave_rollout correctly uses per-step images from VLM cache.""" + cache_data = { + 1: [ + (*_pixels([[1.0]]), [[1, 2, 3]]), # Step 0 + (*_pixels([[1.0], [2.0]]), [[1, 2, 3], [1, 4, 4]]), # Step 1 + ], + } + cache = VLMImageCache(cache_data, num_unique_examples=1, extract_time=0.0, preprocess_time=0.0) + + output = vf.RolloutOutput( + example_id=1, + trajectory=[ + vf.TrajectoryStep( + prompt=[{"role": "user", "content": "Turn 1"}], + completion=[{"role": "assistant", "content": "Response 1"}], + response=MagicMock(), + tokens=vf.TrajectoryStepTokens( + prompt_ids=[1, 2], + prompt_mask=[0, 0], + completion_ids=[3, 4], + completion_mask=[1, 1], + completion_logprobs=[-0.1, -0.2], + overlong_prompt=False, + is_truncated=False, + ), + reward=None, + advantage=None, + is_truncated=False, + trajectory_id="1", + extras={}, + ), + vf.TrajectoryStep( + prompt=[{"role": "user", "content": "Turn 2"}], + completion=[{"role": "assistant", "content": "Response 2"}], + response=MagicMock(), + tokens=vf.TrajectoryStepTokens( + prompt_ids=[1, 2, 3, 4, 5], + prompt_mask=[0, 0, 0, 0, 0], + completion_ids=[6, 7], + completion_mask=[1, 1], + completion_logprobs=[-0.3, -0.4], + overlong_prompt=False, + is_truncated=False, + ), + reward=None, + advantage=None, + is_truncated=False, + trajectory_id="1", + extras={}, + ), + ], + sampling_args={"temperature": 1.0}, + error=None, + ) + + # Token 2 is an image token, token 5 is a video token + mm_mapping = {2: 1, 5: 2} + rollouts = interleave_rollout(output, vlm_cache=cache, mm_token_type_ids_mapping=mm_mapping) + + # Extension holds (step 1 prompt [1,2,3,4,5] extends prefix [1,2,3,4]) + # so both steps merge into a single sample with cumulative images from step 1 + assert rollouts is not None + assert len(rollouts) == 1 + rollout = rollouts[0] + assert rollout.prompt_ids == [1, 2] + assert rollout.completion_ids == [3, 4, 5, 6, 7] + assert rollout.completion_mask == [True, True, False, True, True] + assert rollout.completion_logprobs == [-0.1, -0.2, 0.0, -0.3, -0.4] + # Images: cumulative from last merged step (step 1 has 2 images) + assert _decode_pixels(rollout.pixel_values, rollout.pixel_values_shape) == [[1.0], [2.0]] + assert rollout.image_grid_thw == [[1, 2, 3], [1, 4, 4]] + # mm_token_type_ids: full sequence [1,2,3,4,5,6,7] → [0,1,0,0,2,0,0] + assert rollout.mm_token_type_ids == [0, 1, 0, 0, 2, 0, 0] + + +def test_interleave_rollout_uses_cache_key_override(): + cache_data = { + 7: [ + (*_pixels([[9.0]]), [[1, 2, 3]]), + ], + } + cache = VLMImageCache(cache_data, num_unique_examples=1, extract_time=0.0, preprocess_time=0.0) + + output = vf.RolloutOutput( + example_id=123, + trajectory=[ + vf.TrajectoryStep( + prompt=[{"role": "user", "content": "Turn 1"}], + completion=[{"role": "assistant", "content": "Response 1"}], + response=MagicMock(), + tokens=vf.TrajectoryStepTokens( + prompt_ids=[1, 2], + prompt_mask=[0, 0], + completion_ids=[3, 4], + completion_mask=[1, 1], + completion_logprobs=[-0.1, -0.2], + overlong_prompt=False, + is_truncated=False, + ), + reward=None, + advantage=None, + is_truncated=False, + trajectory_id="1", + extras={}, + ), + ], + sampling_args={"temperature": 1.0}, + error=None, + ) + + rollouts = interleave_rollout(output, vlm_cache=cache, cache_key=7) + + assert rollouts is not None + assert len(rollouts) == 1 + assert _decode_pixels(rollouts[0].pixel_values, rollouts[0].pixel_values_shape) == [[9.0]] + assert rollouts[0].image_grid_thw == [[1, 2, 3]] + + +def test_interleave_rollout_vlm_image_then_text_turns(): + """ + VLM 3-step trajectory: image in step 0, text-only in steps 1 and 2. + Extension holds throughout so all steps merge into 1 sample carrying + step 0's pixel_values (no new images added in later steps). + """ + cache_data = { + 1: [ + (*_pixels([[1.0, 2.0]]), [[1, 3, 3]]), # Step 0: 1 image + (*_pixels([[1.0, 2.0]]), [[1, 3, 3]]), # Step 1: same 1 image (no new) + (*_pixels([[1.0, 2.0]]), [[1, 3, 3]]), # Step 2: same 1 image (no new) + ], + } + cache = VLMImageCache(cache_data, num_unique_examples=1, extract_time=0.0, preprocess_time=0.0) + + output = vf.RolloutOutput( + example_id=1, + trajectory=[ + # Step 0: user sends image + vf.TrajectoryStep( + prompt=[{"role": "user", "content": "Describe"}], + completion=[{"role": "assistant", "content": "A cat"}], + response=MagicMock(), + tokens=vf.TrajectoryStepTokens( + prompt_ids=[1, 2], + prompt_mask=[0, 0], + completion_ids=[3, 4], + completion_mask=[1, 1], + completion_logprobs=[-0.1, -0.2], + overlong_prompt=False, + is_truncated=False, + ), + reward=None, + advantage=None, + is_truncated=False, + trajectory_id="1", + extras={}, + ), + # Step 1: text-only follow-up (extension holds) + vf.TrajectoryStep( + prompt=[{"role": "user", "content": "More detail"}], + completion=[{"role": "assistant", "content": "Fluffy"}], + response=MagicMock(), + tokens=vf.TrajectoryStepTokens( + prompt_ids=[1, 2, 3, 4, 5, 6], + prompt_mask=[0, 0, 0, 0, 0, 0], + completion_ids=[7, 8], + completion_mask=[1, 1], + completion_logprobs=[-0.3, -0.4], + overlong_prompt=False, + is_truncated=False, + ), + reward=None, + advantage=None, + is_truncated=False, + trajectory_id="1", + extras={}, + ), + # Step 2: another text-only follow-up (extension holds) + vf.TrajectoryStep( + prompt=[{"role": "user", "content": "Color?"}], + completion=[{"role": "assistant", "content": "Orange"}], + response=MagicMock(), + tokens=vf.TrajectoryStepTokens( + prompt_ids=[1, 2, 3, 4, 5, 6, 7, 8, 9, 10], + prompt_mask=[0, 0, 0, 0, 0, 0, 0, 0, 0, 0], + completion_ids=[11, 12], + completion_mask=[1, 1], + completion_logprobs=[-0.5, -0.6], + overlong_prompt=False, + is_truncated=False, + ), + reward=None, + advantage=None, + is_truncated=False, + trajectory_id="1", + extras={}, + ), + ], + sampling_args={"temperature": 1.0}, + error=None, + ) + + rollouts = interleave_rollout(output, vlm_cache=cache) + + # All 3 steps merge into 1 sample (extension always holds) + assert rollouts is not None + assert len(rollouts) == 1 + rollout = rollouts[0] + assert rollout.prompt_ids == [1, 2] + # completion: step0 [3,4] + step1 new prompt [5,6] + step1 completion [7,8] + # + step2 new prompt [9,10] + step2 completion [11,12] + assert rollout.completion_ids == [3, 4, 5, 6, 7, 8, 9, 10, 11, 12] + assert rollout.completion_mask == [True, True, False, False, True, True, False, False, True, True] + # pixel_values from step 2 (cumulative = same 1 image throughout) + assert _decode_pixels(rollout.pixel_values, rollout.pixel_values_shape) == [[1.0, 2.0]] + assert rollout.image_grid_thw == [[1, 3, 3]] + + +def test_interleave_rollout_vlm_new_image_mid_conversation(): + """ + VLM 3-step trajectory: image in step 0, text in step 1, NEW image in step 2. + Extension holds throughout, so 1 merged sample with cumulative images from step 2. + """ + cache_data = { + 1: [ + (*_pixels([[1.0]]), [[1, 2, 3]]), # Step 0: 1 image + (*_pixels([[1.0]]), [[1, 2, 3]]), # Step 1: still 1 image + (*_pixels([[1.0], [2.0]]), [[1, 2, 3], [1, 4, 4]]), # Step 2: 2 images + ], + } + cache = VLMImageCache(cache_data, num_unique_examples=1, extract_time=0.0, preprocess_time=0.0) + + output = vf.RolloutOutput( + example_id=1, + trajectory=[ + vf.TrajectoryStep( + prompt=[{"role": "user", "content": "Image 1"}], + completion=[{"role": "assistant", "content": "A"}], + response=MagicMock(), + tokens=vf.TrajectoryStepTokens( + prompt_ids=[1, 2], + prompt_mask=[0, 0], + completion_ids=[3, 4], + completion_mask=[1, 1], + completion_logprobs=[-0.1, -0.2], + overlong_prompt=False, + is_truncated=False, + ), + reward=None, + advantage=None, + is_truncated=False, + trajectory_id="1", + extras={}, + ), + vf.TrajectoryStep( + prompt=[{"role": "user", "content": "Text only"}], + completion=[{"role": "assistant", "content": "B"}], + response=MagicMock(), + tokens=vf.TrajectoryStepTokens( + prompt_ids=[1, 2, 3, 4, 5, 6], + prompt_mask=[0, 0, 0, 0, 0, 0], + completion_ids=[7, 8], + completion_mask=[1, 1], + completion_logprobs=[-0.3, -0.4], + overlong_prompt=False, + is_truncated=False, + ), + reward=None, + advantage=None, + is_truncated=False, + trajectory_id="1", + extras={}, + ), + vf.TrajectoryStep( + prompt=[{"role": "user", "content": "Image 2"}], + completion=[{"role": "assistant", "content": "C"}], + response=MagicMock(), + tokens=vf.TrajectoryStepTokens( + prompt_ids=[1, 2, 3, 4, 5, 6, 7, 8, 9, 10], + prompt_mask=[0, 0, 0, 0, 0, 0, 0, 0, 0, 0], + completion_ids=[11, 12], + completion_mask=[1, 1], + completion_logprobs=[-0.5, -0.6], + overlong_prompt=False, + is_truncated=False, + ), + reward=None, + advantage=None, + is_truncated=False, + trajectory_id="1", + extras={}, + ), + ], + sampling_args={"temperature": 1.0}, + error=None, + ) + + rollouts = interleave_rollout(output, vlm_cache=cache) + + assert rollouts is not None + assert len(rollouts) == 1 + rollout = rollouts[0] + assert rollout.prompt_ids == [1, 2] + assert rollout.completion_ids == [3, 4, 5, 6, 7, 8, 9, 10, 11, 12] + # Cumulative images from last merged step (step 2): both images + assert _decode_pixels(rollout.pixel_values, rollout.pixel_values_shape) == [[1.0], [2.0]] + assert rollout.image_grid_thw == [[1, 2, 3], [1, 4, 4]] + + +def test_interleave_rollout_vlm_extension_break(): + """ + VLM 3-step trajectory where extension breaks at step 2. + Step 0 has image, step 1 extends (text-only), step 2 breaks (different prefix). + Should produce 2 samples, each with their own cumulative images. + """ + cache_data = { + 1: [ + (*_pixels([[1.0]]), [[1, 2, 3]]), # Step 0: 1 image + (*_pixels([[1.0]]), [[1, 2, 3]]), # Step 1: still 1 image + (*_pixels([[1.0], [2.0]]), [[1, 2, 3], [1, 4, 4]]), # Step 2: 2 images (new image added) + ], + } + cache = VLMImageCache(cache_data, num_unique_examples=1, extract_time=0.0, preprocess_time=0.0) + + output = vf.RolloutOutput( + example_id=1, + trajectory=[ + vf.TrajectoryStep( + prompt=[{"role": "user", "content": "Image 1"}], + completion=[{"role": "assistant", "content": "A"}], + response=MagicMock(), + tokens=vf.TrajectoryStepTokens( + prompt_ids=[1, 2], + prompt_mask=[0, 0], + completion_ids=[3, 4], + completion_mask=[1, 1], + completion_logprobs=[-0.1, -0.2], + overlong_prompt=False, + is_truncated=False, + ), + reward=None, + advantage=None, + is_truncated=False, + trajectory_id="1", + extras={}, + ), + # Step 1: extends step 0 + vf.TrajectoryStep( + prompt=[{"role": "user", "content": "Follow-up"}], + completion=[{"role": "assistant", "content": "B"}], + response=MagicMock(), + tokens=vf.TrajectoryStepTokens( + prompt_ids=[1, 2, 3, 4, 5, 6], + prompt_mask=[0, 0, 0, 0, 0, 0], + completion_ids=[7, 8], + completion_mask=[1, 1], + completion_logprobs=[-0.3, -0.4], + overlong_prompt=False, + is_truncated=False, + ), + reward=None, + advantage=None, + is_truncated=False, + trajectory_id="1", + extras={}, + ), + # Step 2: extension breaks (different prefix, e.g. context compaction) + vf.TrajectoryStep( + prompt=[{"role": "user", "content": "Image 2"}], + completion=[{"role": "assistant", "content": "C"}], + response=MagicMock(), + tokens=vf.TrajectoryStepTokens( + prompt_ids=[100, 101, 102, 103], + prompt_mask=[0, 0, 0, 0], + completion_ids=[104, 105], + completion_mask=[1, 1], + completion_logprobs=[-0.5, -0.6], + overlong_prompt=False, + is_truncated=False, + ), + reward=None, + advantage=None, + is_truncated=False, + trajectory_id="1", + extras={}, + ), + ], + sampling_args={"temperature": 1.0}, + error=None, + ) + + rollouts = interleave_rollout(output, vlm_cache=cache) + + assert rollouts is not None + assert len(rollouts) == 2 + + # Sample 1: steps 0-1 merged, images from step 1 (still 1 image) + assert rollouts[0].prompt_ids == [1, 2] + assert rollouts[0].completion_ids == [3, 4, 5, 6, 7, 8] + assert _decode_pixels(rollouts[0].pixel_values, rollouts[0].pixel_values_shape) == [[1.0]] + assert rollouts[0].image_grid_thw == [[1, 2, 3]] + + # Sample 2: step 2 alone (extension broke), images from step 2 (2 images) + assert rollouts[1].prompt_ids == [100, 101, 102, 103] + assert rollouts[1].completion_ids == [104, 105] + assert _decode_pixels(rollouts[1].pixel_values, rollouts[1].pixel_values_shape) == [[1.0], [2.0]] + assert rollouts[1].image_grid_thw == [[1, 2, 3], [1, 4, 4]] + + +def test_interleave_rollout_vlm_image_appears_late(): + """ + VLM 3-step trajectory: text-only in steps 0 and 1, first image in step 2. + Extension holds throughout so all steps merge into 1 sample. + The sample should have pixel_values=None until step 2 sets them. + """ + cache_data = { + 1: [ + (None, None, None), # Step 0: no images + (None, None, None), # Step 1: no images + (*_pixels([[5.0, 6.0]]), [[1, 3, 3]]), # Step 2: first image appears + ], + } + cache = VLMImageCache(cache_data, num_unique_examples=1, extract_time=0.0, preprocess_time=0.0) + + output = vf.RolloutOutput( + example_id=1, + trajectory=[ + # Step 0: text-only + vf.TrajectoryStep( + prompt=[{"role": "user", "content": "Hello"}], + completion=[{"role": "assistant", "content": "Hi"}], + response=MagicMock(), + tokens=vf.TrajectoryStepTokens( + prompt_ids=[1, 2], + prompt_mask=[0, 0], + completion_ids=[3, 4], + completion_mask=[1, 1], + completion_logprobs=[-0.1, -0.2], + overlong_prompt=False, + is_truncated=False, + ), + reward=None, + advantage=None, + is_truncated=False, + trajectory_id="1", + extras={}, + ), + # Step 1: text-only (extension holds) + vf.TrajectoryStep( + prompt=[{"role": "user", "content": "Question"}], + completion=[{"role": "assistant", "content": "Answer"}], + response=MagicMock(), + tokens=vf.TrajectoryStepTokens( + prompt_ids=[1, 2, 3, 4, 5, 6], + prompt_mask=[0, 0, 0, 0, 0, 0], + completion_ids=[7, 8], + completion_mask=[1, 1], + completion_logprobs=[-0.3, -0.4], + overlong_prompt=False, + is_truncated=False, + ), + reward=None, + advantage=None, + is_truncated=False, + trajectory_id="1", + extras={}, + ), + # Step 2: user sends image (extension holds) + vf.TrajectoryStep( + prompt=[{"role": "user", "content": "Describe this"}], + completion=[{"role": "assistant", "content": "A photo"}], + response=MagicMock(), + tokens=vf.TrajectoryStepTokens( + prompt_ids=[1, 2, 3, 4, 5, 6, 7, 8, 9, 10], + prompt_mask=[0, 0, 0, 0, 0, 0, 0, 0, 0, 0], + completion_ids=[11, 12], + completion_mask=[1, 1], + completion_logprobs=[-0.5, -0.6], + overlong_prompt=False, + is_truncated=False, + ), + reward=None, + advantage=None, + is_truncated=False, + trajectory_id="1", + extras={}, + ), + ], + sampling_args={"temperature": 1.0}, + error=None, + ) + + rollouts = interleave_rollout(output, vlm_cache=cache) + + assert rollouts is not None + assert len(rollouts) == 1 + rollout = rollouts[0] + assert rollout.prompt_ids == [1, 2] + assert rollout.completion_ids == [3, 4, 5, 6, 7, 8, 9, 10, 11, 12] + assert rollout.completion_mask == [True, True, False, False, True, True, False, False, True, True] + # pixel_values from step 2 (the first step with an image) + assert _decode_pixels(rollout.pixel_values, rollout.pixel_values_shape) == [[5.0, 6.0]] + assert rollout.image_grid_thw == [[1, 3, 3]] + + +def test_interleave_rollout_empty_trajectory(): + """Empty trajectory returns None.""" + output = vf.RolloutOutput( + example_id=1, + trajectory=[], + error=None, + ) + assert interleave_rollout(output) is None + + +def test_interleave_rollout_error_masks_all_false(): + """ + When rollout output has an error, all completion_mask values should be False + across both make_sample (step 0) and extend_sample (step 1). + """ + output = vf.RolloutOutput( + example_id=1, + trajectory=[ + vf.TrajectoryStep( + prompt=[{"role": "user", "content": "U1"}], + completion=[{"role": "assistant", "content": "A1"}], + response=MagicMock(), + tokens=vf.TrajectoryStepTokens( + prompt_ids=[1, 2], + prompt_mask=[0, 0], + completion_ids=[3, 4], + completion_mask=[1, 1], + completion_logprobs=[-0.1, -0.2], + overlong_prompt=False, + is_truncated=False, + ), + reward=None, + advantage=None, + is_truncated=False, + trajectory_id="1", + extras={}, + ), + vf.TrajectoryStep( + prompt=[{"role": "user", "content": "U2"}], + completion=[{"role": "assistant", "content": "A2"}], + response=MagicMock(), + tokens=vf.TrajectoryStepTokens( + prompt_ids=[1, 2, 3, 4, 5, 6], + prompt_mask=[0, 0, 0, 0, 0, 0], + completion_ids=[7, 8], + completion_mask=[1, 1], + completion_logprobs=[-0.3, -0.4], + overlong_prompt=False, + is_truncated=False, + ), + reward=None, + advantage=None, + is_truncated=False, + trajectory_id="1", + extras={}, + ), + ], + error="timeout: environment exceeded time limit", + sampling_args={"temperature": 0.8}, + ) + + rollouts = interleave_rollout(output) + + assert rollouts is not None + assert len(rollouts) == 1 + rollout = rollouts[0] + # Extension holds so tokens merge, but ALL completion_mask should be False + assert rollout.completion_ids == [3, 4, 5, 6, 7, 8] + assert rollout.completion_mask == [False, False, False, False, False, False] + # Logprobs and temperatures still present + assert rollout.completion_logprobs == [-0.1, -0.2, 0.0, 0.0, -0.3, -0.4] + assert rollout.completion_temperatures == [0.8] * 6 + + +def test_interleave_rollout_vlm_interleaved_agents(): + """ + VLM + interleaved agents: agent1 and agent2 interleaved, each with images. + agent1 gets cumulative images from its own steps, agent2 from its step. + + Steps (0-indexed): + 0: agent1-step1 (image A) + 1: agent1-step2 (extends step 0, image A still) + 2: agent2-step1 (different prefix, image B) + 3: agent1-step3 (extends step 0+1, image A + new image C) + + Expected: 2 samples + - agent1: merged steps 0,1,3 → pixel_values from step 3 (images A+C) + - agent2: step 2 alone → pixel_values from step 2 (image B) + """ + cache_data = { + 1: [ + (*_pixels([[1.0]]), [[1, 2, 2]]), # Step 0: image A + (*_pixels([[1.0]]), [[1, 2, 2]]), # Step 1: still image A + (*_pixels([[9.0]]), [[1, 5, 5]]), # Step 2: image B (agent2) + (*_pixels([[1.0], [3.0]]), [[1, 2, 2], [1, 3, 3]]), # Step 3: images A+C (agent1) + ], + } + cache = VLMImageCache(cache_data, num_unique_examples=1, extract_time=0.0, preprocess_time=0.0) + + output = vf.RolloutOutput( + example_id=1, + trajectory=[ + # Step 0: agent1-step1 + vf.TrajectoryStep( + prompt=[{"role": "user", "content": "Image A"}], + completion=[{"role": "assistant", "content": "A1"}], + response=MagicMock(), + tokens=vf.TrajectoryStepTokens( + prompt_ids=[1, 2], + prompt_mask=[0, 0], + completion_ids=[3, 4], + completion_mask=[1, 1], + completion_logprobs=[-0.1, -0.2], + overlong_prompt=False, + is_truncated=False, + ), + reward=None, + advantage=None, + is_truncated=False, + trajectory_id="1", + extras={}, + ), + # Step 1: agent1-step2 (extends step 0) + vf.TrajectoryStep( + prompt=[{"role": "user", "content": "Follow-up"}], + completion=[{"role": "assistant", "content": "A2"}], + response=MagicMock(), + tokens=vf.TrajectoryStepTokens( + prompt_ids=[1, 2, 3, 4, 5, 6], + prompt_mask=[0, 0, 0, 0, 0, 0], + completion_ids=[7, 8], + completion_mask=[1, 1], + completion_logprobs=[-0.3, -0.4], + overlong_prompt=False, + is_truncated=False, + ), + reward=None, + advantage=None, + is_truncated=False, + trajectory_id="1", + extras={}, + ), + # Step 2: agent2-step1 (different prefix) + vf.TrajectoryStep( + prompt=[{"role": "user", "content": "Image B"}], + completion=[{"role": "assistant", "content": "B1"}], + response=MagicMock(), + tokens=vf.TrajectoryStepTokens( + prompt_ids=[100, 101], + prompt_mask=[0, 0], + completion_ids=[102, 103], + completion_mask=[1, 1], + completion_logprobs=[-0.5, -0.6], + overlong_prompt=False, + is_truncated=False, + ), + reward=None, + advantage=None, + is_truncated=False, + trajectory_id="1", + extras={}, + ), + # Step 3: agent1-step3 (extends agent1, merges back) + vf.TrajectoryStep( + prompt=[{"role": "user", "content": "Image C added"}], + completion=[{"role": "assistant", "content": "A3"}], + response=MagicMock(), + tokens=vf.TrajectoryStepTokens( + prompt_ids=[1, 2, 3, 4, 5, 6, 7, 8, 9, 10], + prompt_mask=[0, 0, 0, 0, 0, 0, 0, 0, 0, 0], + completion_ids=[11, 12], + completion_mask=[1, 1], + completion_logprobs=[-0.7, -0.8], + overlong_prompt=False, + is_truncated=False, + ), + reward=None, + advantage=None, + is_truncated=False, + trajectory_id="1", + extras={}, + ), + ], + error=None, + sampling_args={"temperature": 1.0}, + ) + + rollouts = interleave_rollout(output, vlm_cache=cache) + + assert rollouts is not None + assert len(rollouts) == 2 + + # Agent1: steps 0,1,3 merged → images from step 3 (A+C) + agent1 = rollouts[0] + assert agent1.prompt_ids == [1, 2] + assert agent1.completion_ids == [3, 4, 5, 6, 7, 8, 9, 10, 11, 12] + assert agent1.completion_mask == [True, True, False, False, True, True, False, False, True, True] + assert _decode_pixels(agent1.pixel_values, agent1.pixel_values_shape) == [[1.0], [3.0]] + assert agent1.image_grid_thw == [[1, 2, 2], [1, 3, 3]] + + # Agent2: step 2 alone → images from step 2 (B) + agent2 = rollouts[1] + assert agent2.prompt_ids == [100, 101] + assert agent2.completion_ids == [102, 103] + assert agent2.completion_mask == [True, True] + assert _decode_pixels(agent2.pixel_values, agent2.pixel_values_shape) == [[9.0]] + assert agent2.image_grid_thw == [[1, 5, 5]] + + +def test_build_vlm_image_cache_handles_divergent_rollouts(): + """Test that build_vlm_image_cache keys images per rollout when trajectories diverge.""" + import torch + + red_url = _create_test_image("red") + blue_url = _create_test_image("blue") + green_url = _create_test_image("green") + + rollout_a = vf.RolloutOutput( + example_id=1, + trajectory=[ + vf.TrajectoryStep( + prompt=[_create_image_message(red_url, "What color?")], + completion=[{"role": "assistant", "content": "Red"}], + response=MagicMock(), + tokens=MagicMock(), + reward=None, + advantage=None, + is_truncated=False, + trajectory_id="1", + extras={}, + ), + ], + sampling_args={"temperature": 1.0}, + error=None, + ) + + rollout_b = vf.RolloutOutput( + example_id=1, + trajectory=[ + vf.TrajectoryStep( + prompt=[_create_image_message(blue_url, "What color?")], + completion=[{"role": "assistant", "content": "Blue"}], + response=MagicMock(), + tokens=MagicMock(), + reward=None, + advantage=None, + is_truncated=False, + trajectory_id="1", + extras={}, + ), + vf.TrajectoryStep( + prompt=[ + _create_image_message(blue_url, "What color?"), + {"role": "assistant", "content": "Blue"}, + _create_image_message(green_url, "And this one?"), + ], + completion=[{"role": "assistant", "content": "Green"}], + response=MagicMock(), + tokens=MagicMock(), + reward=None, + advantage=None, + is_truncated=False, + trajectory_id="1", + extras={}, + ), + ], + sampling_args={"temperature": 1.0}, + error=None, + ) + + # Mock processor that returns predictable tensors + mock_processor = MagicMock() + mock_processor.image_processor = MagicMock( + side_effect=lambda images, return_tensors: { + "pixel_values": torch.arange(len(images), dtype=torch.float32).view(-1, 1), + "image_grid_thw": torch.tensor([[1, 1, 1]] * len(images)), + } + ) + + rollouts = [rollout_a, rollout_b] + cache = build_vlm_image_cache(rollouts, mock_processor) + + assert cache.num_unique_examples == 1 + + pv, shape, grid = cache.get_for_step(0, 0) + assert _decode_pixels(pv, shape) == [[0.0]] + assert grid == [[1, 1, 1]] + + pv, shape, grid = cache.get_for_step(1, 0) + assert _decode_pixels(pv, shape) == [[1.0]] + assert grid == [[1, 1, 1]] + + pv, shape, grid = cache.get_for_step(1, 1) + assert _decode_pixels(pv, shape) == [[1.0], [2.0]] + assert grid == [[1, 1, 1], [1, 1, 1]] + + +def test_build_vlm_image_cache_no_images(): + output = vf.RolloutOutput( + example_id=1, + trajectory=[ + vf.TrajectoryStep( + prompt=[{"role": "user", "content": "Hello"}], + completion=[{"role": "assistant", "content": "Hi"}], + response=MagicMock(), + tokens=MagicMock(), + reward=None, + advantage=None, + is_truncated=False, + trajectory_id="1", + extras={}, + ) + ], + sampling_args={"temperature": 1.0}, + error=None, + ) + + cache = build_vlm_image_cache([output], MagicMock()) + + pv, shape, grid = cache.get_for_step(0, 0) + assert pv is None + assert shape is None + assert grid is None + + +def test_align_routed_experts_none(): + assert _align_routed_experts(None, 10) is None + + +def test_align_routed_experts_empty(): + result = _align_routed_experts([], 10) + assert result == [] + + +def test_align_routed_experts_no_deficit(): + # 3 tokens, 2 layers, topk=2 + experts = [[[0, 1], [2, 3]], [[4, 5], [6, 7]], [[0, 2], [1, 3]]] + result = _align_routed_experts(experts, expected_len=3) + assert result == experts + + +def test_align_routed_experts_with_deficit(): + # 2 tokens but expected 4 (deficit of 2) + experts = [[[1, 2], [3, 4]], [[5, 6], [7, 0]]] + result = _align_routed_experts(experts, expected_len=4) + assert len(result) == 4 + assert result[:2] == experts + # Padded entries should be zero-filled with same shape [layers=2, topk=2] + assert result[2] == [[0, 0], [0, 0]] + assert result[3] == [[0, 0], [0, 0]] + + +def test_align_routed_experts_excess_length(): + experts = [[[1, 2]], [[3, 4]], [[5, 6]]] + result = _align_routed_experts(experts, expected_len=2) + # No truncation, just returns as-is + assert result == experts + + +def test_interleave_rollout_single_step_with_routed_experts(): + """Routed experts are aligned and passed through for a single-step trajectory.""" + # prompt_ids=[1,2], completion_ids=[3,4] -> total 4 tokens + # vLLM returns num_tokens-1 = 3 routed expert entries + routed_experts_from_vllm = [[[0, 1]], [[2, 3]], [[4, 5]]] # 3 entries, 1 layer, topk=2 + output = vf.RolloutOutput( + example_id=0, + trajectory=[ + vf.TrajectoryStep( + prompt=[{"role": "user", "content": "U1"}], + completion=[{"role": "assistant", "content": "A1"}], + response=MagicMock(), + tokens=vf.TrajectoryStepTokens( + prompt_ids=[1, 2], + prompt_mask=[0, 0], + completion_ids=[3, 4], + completion_mask=[1, 1], + completion_logprobs=[-0.1, -0.2], + overlong_prompt=False, + is_truncated=False, + routed_experts=routed_experts_from_vllm, + ), + reward=None, + advantage=None, + is_truncated=False, + trajectory_id="1", + extras={}, + ) + ], + sampling_args={"temperature": 1.0}, + error=None, + ) + + rollouts = interleave_rollout(output) + assert rollouts is not None + assert len(rollouts) == 1 + sample = rollouts[0] + + # Should be aligned to 4 tokens (2 prompt + 2 completion) + assert sample.routed_experts is not None + assert len(sample.routed_experts) == 4 + # First 3 are original, last one is zero-padded + assert sample.routed_experts[:3] == routed_experts_from_vllm + assert sample.routed_experts[3] == [[0, 0]] + + +def test_interleave_rollout_multi_step_with_routed_experts(): + """Routed experts are extended and aligned across multi-step trajectories.""" + # Step 1: prompt=[1,2], completion=[3,4] -> 4 tokens, vLLM returns 3 + step1_experts = [[[1, 2]], [[3, 4]], [[5, 6]]] + # Step 2: prompt=[1,2,3,4,5,6], completion=[7,8] -> 8 tokens, vLLM returns 7 + step2_experts = [[[1, 0]], [[2, 0]], [[3, 0]], [[4, 0]], [[5, 0]], [[6, 0]], [[7, 0]]] + + output = vf.RolloutOutput( + example_id=0, + trajectory=[ + vf.TrajectoryStep( + prompt=[{"role": "user", "content": "U1"}], + completion=[{"role": "assistant", "content": "A1"}], + response=MagicMock(), + tokens=vf.TrajectoryStepTokens( + prompt_ids=[1, 2], + prompt_mask=[0, 0], + completion_ids=[3, 4], + completion_mask=[1, 1], + completion_logprobs=[-0.1, -0.2], + overlong_prompt=False, + is_truncated=False, + routed_experts=step1_experts, + ), + reward=None, + advantage=None, + is_truncated=False, + trajectory_id="1", + extras={}, + ), + vf.TrajectoryStep( + prompt=[ + {"role": "user", "content": "U1"}, + {"role": "assistant", "content": "A1"}, + {"role": "user", "content": "U2"}, + ], + completion=[{"role": "assistant", "content": "A2"}], + response=MagicMock(), + tokens=vf.TrajectoryStepTokens( + prompt_ids=[1, 2, 3, 4, 5, 6], + prompt_mask=[0, 0, 0, 0, 0, 0], + completion_ids=[7, 8], + completion_mask=[1, 1], + completion_logprobs=[-0.3, -0.4], + overlong_prompt=False, + is_truncated=False, + routed_experts=step2_experts, + ), + reward=None, + advantage=None, + is_truncated=False, + trajectory_id="1", + extras={}, + ), + ], + sampling_args={"temperature": 1.0}, + error=None, + ) + + rollouts = interleave_rollout(output) + assert rollouts is not None + assert len(rollouts) == 1 + sample = rollouts[0] + + # Merged sample: prompt=[1,2], completion=[3,4,5,6,7,8] -> 8 tokens total + assert len(sample.prompt_ids) + len(sample.completion_ids) == 8 + assert sample.routed_experts is not None + assert len(sample.routed_experts) == 8 + + +def test_interleave_rollout_none_routed_experts_stays_none(): + """When routed_experts is None, sample.routed_experts remains None.""" + output = vf.RolloutOutput( + example_id=0, + trajectory=[ + vf.TrajectoryStep( + prompt=[{"role": "user", "content": "U1"}], + completion=[{"role": "assistant", "content": "A1"}], + response=MagicMock(), + tokens=vf.TrajectoryStepTokens( + prompt_ids=[1, 2], + prompt_mask=[0, 0], + completion_ids=[3, 4], + completion_mask=[1, 1], + completion_logprobs=[-0.1, -0.2], + overlong_prompt=False, + is_truncated=False, + routed_experts=None, + ), + reward=None, + advantage=None, + is_truncated=False, + trajectory_id="1", + extras={}, + ) + ], + sampling_args={"temperature": 1.0}, + error=None, + ) + + rollouts = interleave_rollout(output) + assert rollouts is not None + assert rollouts[0].routed_experts is None + + +# ============================================================================= +# _ImageStore and store-backed VLMImageCache tests +# ============================================================================= + + +def test_image_store_assemble(): + """_ImageStore.assemble joins per-image bytes and computes correct shape/grids.""" + # 2 images: image 0 has 3 patches, image 1 has 2 patches, patch_dim=4 + patch_dim = 4 + img0 = np.arange(3 * patch_dim, dtype=np.float32).tobytes() + img1 = np.arange(2 * patch_dim, dtype=np.float32).tobytes() + + store = _ImageStore( + image_bytes=[img0, img1], + image_num_patches=[3, 2], + patch_dim=patch_dim, + image_grids=[[1, 1, 3], [1, 1, 2]], + ) + + # Assemble both images + pixel_bytes, shape, grids = store.assemble([0, 1]) + assert shape == [5, 4] + assert grids == [[1, 1, 3], [1, 1, 2]] + assert pixel_bytes == img0 + img1 + + # Assemble single image + pixel_bytes, shape, grids = store.assemble([1]) + assert shape == [2, 4] + assert grids == [[1, 1, 2]] + assert pixel_bytes == img1 + + # Assemble in reverse order + pixel_bytes, shape, grids = store.assemble([1, 0]) + assert shape == [5, 4] + assert grids == [[1, 1, 2], [1, 1, 3]] + assert pixel_bytes == img1 + img0 + + +def test_vlm_image_cache_from_store(): + """VLMImageCache.from_store provides correct get_for_step/get_all via lazy assembly.""" + patch_dim = 2 + img0_data = np.array([[1.0, 2.0]], dtype=np.float32) + img1_data = np.array([[3.0, 4.0]], dtype=np.float32) + + store = _ImageStore( + image_bytes=[img0_data.tobytes(), img1_data.tobytes()], + image_num_patches=[1, 1], + patch_dim=patch_dim, + image_grids=[[1, 2, 3], [1, 4, 4]], + ) + + step_indices = { + 1: [[0], [0, 1]], # step 0: image 0; step 1: images 0+1 + } + + cache = VLMImageCache.from_store( + store=store, + step_indices=step_indices, + num_unique_examples=1, + num_unique_images=2, + extract_time=0.0, + preprocess_time=0.0, + ) + + # Step 0: just image 0 + pv, shape, grid = cache.get_for_step(1, 0) + assert _decode_pixels(pv, shape) == [[1.0, 2.0]] + assert grid == [[1, 2, 3]] + + # Step 1: images 0 + 1 + pv, shape, grid = cache.get_for_step(1, 1) + assert _decode_pixels(pv, shape) == [[1.0, 2.0], [3.0, 4.0]] + assert grid == [[1, 2, 3], [1, 4, 4]] + + # get_all returns last step + pv, shape, grid = cache.get_all(1) + assert _decode_pixels(pv, shape) == [[1.0, 2.0], [3.0, 4.0]] + assert grid == [[1, 2, 3], [1, 4, 4]] + + # Missing key + pv, shape, grid = cache.get_for_step(999, 0) + assert pv is None + + # Out of range step + pv, shape, grid = cache.get_for_step(1, 5) + assert pv is None + + +def test_vlm_image_cache_from_store_no_images(): + """from_store with store=None returns (None, None, None) for all queries.""" + step_indices = {0: [[], []]} # 2 steps with no images + + cache = VLMImageCache.from_store( + store=None, + step_indices=step_indices, + num_unique_examples=1, + num_unique_images=0, + extract_time=0.0, + preprocess_time=0.0, + ) + + pv, shape, grid = cache.get_for_step(0, 0) + assert pv is None + assert shape is None + assert grid is None + + +def test_build_vlm_image_cache_uses_store(): + """build_vlm_image_cache returns a store-backed cache.""" + import torch + + red_url = _create_test_image("red") + + output = vf.RolloutOutput( + example_id=1, + trajectory=[ + vf.TrajectoryStep( + prompt=[_create_image_message(red_url, "What color?")], + completion=[{"role": "assistant", "content": "Red"}], + response=MagicMock(), + tokens=MagicMock(), + reward=None, + advantage=None, + is_truncated=False, + trajectory_id="1", + extras={}, + ), + ], + sampling_args={"temperature": 1.0}, + error=None, + ) + + mock_processor = MagicMock() + mock_processor.image_processor = MagicMock( + side_effect=lambda images, return_tensors: { + "pixel_values": torch.arange(len(images), dtype=torch.float32).view(-1, 1), + "image_grid_thw": torch.tensor([[1, 1, 1]] * len(images)), + } + ) + + cache = build_vlm_image_cache([output], mock_processor) + + # Should be store-backed + assert cache._store is not None + assert cache._step_indices is not None + + # Should still work correctly + pv, shape, grid = cache.get_for_step(0, 0) + assert pv is not None + assert shape == [1, 1] + assert grid == [[1, 1, 1]] diff --git a/tests/unit/orchestrator/test_vf.py b/tests/unit/orchestrator/test_vf.py deleted file mode 100644 index 7210e7fedbe..00000000000 --- a/tests/unit/orchestrator/test_vf.py +++ /dev/null @@ -1,69 +0,0 @@ -import json -from copy import deepcopy - -import pytest -import verifiers as vf -from openai.types.chat.chat_completion import ChatCompletion, Choice -from openai.types.chat.chat_completion_message import ChatCompletionMessage -from openai.types.completion_usage import CompletionUsage - -from prime_rl.utils.vf import from_serializable_state, to_serializable_state - - -@pytest.fixture -def single_step_trajectory_state(): - state = vf.State( - trajectory=[ - vf.TrajectoryStep( - prompt=[{"role": "user", "content": "U1"}], - completion=[{"role": "assistant", "content": "A1"}], - response=ChatCompletion( - id="cm123", - object="chat.completion", - created=1716806400, - model="gpt-4o", - choices=[ - Choice( - index=0, - message=ChatCompletionMessage( - role="assistant", - content="A1", - ), - finish_reason="stop", - ) - ], - usage=CompletionUsage(prompt_tokens=2, completion_tokens=1, total_tokens=3), - ), - tokens=vf.TrajectoryStepTokens( - prompt_ids=[1, 2], - prompt_mask=[0, 0], - completion_ids=[3, 4], - completion_mask=[1, 1], - completion_logprobs=[-0.1, -0.2], - overlong_prompt=False, - is_truncated=False, - ), - reward=None, - advantage=None, - extras={}, - ) - ], - ) - return state - - -def test_serialize_state(single_step_trajectory_state): - # Regular state is not JSON serializable - with pytest.raises(Exception): - json.dumps(single_step_trajectory_state) - - # Serialized state is JSON serializable - serializable_state = to_serializable_state(single_step_trajectory_state) - json.dumps(serializable_state) - - -def test_deserialize_state(single_step_trajectory_state): - original_state = deepcopy(single_step_trajectory_state) - serialized_state = to_serializable_state(single_step_trajectory_state) - deserialized_state = from_serializable_state(serialized_state) - assert deserialized_state == original_state diff --git a/tests/unit/test_configs.py b/tests/unit/test_configs.py index d398dbd3b63..ffcc18f2705 100644 --- a/tests/unit/test_configs.py +++ b/tests/unit/test_configs.py @@ -1,19 +1,27 @@ -import sys from pathlib import Path +from typing import Annotated, Literal import pytest -from pydantic import ValidationError +import tomli_w +from pydantic import BaseModel, Field, ValidationError +from pydantic_config import ConfigFileError -from prime_rl.eval.config import OfflineEvalConfig -from prime_rl.inference.config import InferenceConfig -from prime_rl.orchestrator.config import OrchestratorConfig -from prime_rl.rl import RLConfig -from prime_rl.trainer.rl.config import RLTrainerConfig -from prime_rl.trainer.sft.config import SFTTrainerConfig -from prime_rl.utils.pydantic_config import parse_argv +from prime_rl.configs.inference import InferenceConfig +from prime_rl.configs.orchestrator import OrchestratorConfig +from prime_rl.configs.rl import RLConfig +from prime_rl.configs.sft import SFTConfig +from prime_rl.configs.trainer import ModelConfig as TrainerModelConfig +from prime_rl.configs.trainer import TrainerConfig +from prime_rl.utils.config import BaseConfig, cli # All config config classes -CONFIG_CLASSES = [RLConfig, RLTrainerConfig, SFTTrainerConfig, OrchestratorConfig, InferenceConfig, OfflineEvalConfig] +CONFIG_CLASSES = [ + RLConfig, + TrainerConfig, + SFTConfig, + OrchestratorConfig, + InferenceConfig, +] def get_config_files() -> list[Path]: @@ -25,22 +33,129 @@ def get_config_files() -> list[Path]: @pytest.mark.parametrize("config_file", get_config_files(), ids=lambda x: x.as_posix()) -def test_load_configs(config_file: Path, monkeypatch: pytest.MonkeyPatch): +def test_load_configs(config_file: Path): """Tests that all config files can be loaded by at least one config class.""" - if "intellect_3/evals" in config_file.as_posix(): - pytest.skip("Skipped because uses partial configs, which are not supported by this test.") - - monkeypatch.setattr( - sys, - "argv", - ["dummy.py", "@", config_file.as_posix()], - raising=False, - ) could_parse = [] for config_cls in CONFIG_CLASSES: try: - parse_argv(config_cls) + cli(config_cls, args=["@", config_file.as_posix()]) could_parse.append(True) - except ValidationError: + except (ValidationError, ConfigFileError, SystemExit): could_parse.append(False) assert any(could_parse), f"No config class could be parsed from {config_file}" + + +class NestedConfig(BaseConfig): + lr: float = 1e-4 + weight_decay: float = 0.01 + name: str = "default" + + +class VariantA(BaseModel): + type: Literal["a"] = "a" + alpha: float = 0.1 + shared: int = 1 + + +class VariantB(BaseModel): + type: Literal["b"] = "b" + beta: float = 0.2 + shared: int = 1 + + +VariantType = Annotated[VariantA | VariantB, Field(discriminator="type")] + + +class DummyConfig(BaseConfig): + name: str = "experiment" + seed: int = 42 + nested: NestedConfig = NestedConfig() + variant: VariantType = VariantA() + + +def write_toml(path: Path, data: dict) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + with open(path, "wb") as f: + tomli_w.dump(data, f) + + +def test_defaults(): + """All defaults are applied when no TOML or CLI args are given.""" + config = cli(DummyConfig, args=[]) + assert config.name == "experiment" + assert config.seed == 42 + assert config.nested.lr == 1e-4 + assert config.nested.weight_decay == 0.01 + assert config.variant.type == "a" + assert config.variant.alpha == 0.1 + + +def test_toml_partial_nested_override(tmp_path): + """Partially overriding a nested model preserves unset field defaults.""" + write_toml(tmp_path / "cfg.toml", {"nested": {"lr": 3e-4}}) + config = cli(DummyConfig, args=["@", str(tmp_path / "cfg.toml")]) + assert config.nested.lr == 3e-4 + assert config.nested.weight_decay == 0.01 + assert config.nested.name == "default" + + +def test_toml_discriminated_union_default_type(tmp_path): + """Overriding a discriminated union field without 'type' uses the default variant.""" + write_toml(tmp_path / "cfg.toml", {"variant": {"alpha": 0.9}}) + config = cli(DummyConfig, args=["@", str(tmp_path / "cfg.toml")]) + assert config.variant.type == "a" + assert config.variant.alpha == 0.9 + assert config.variant.shared == 1 + + +def test_toml_discriminated_union_switch_variant(tmp_path): + """Providing an explicit 'type' switches to that variant.""" + write_toml(tmp_path / "cfg.toml", {"variant": {"type": "b"}}) + config = cli(DummyConfig, args=["@", str(tmp_path / "cfg.toml")]) + assert config.variant.type == "b" + assert config.variant.beta == 0.2 + + +def test_toml_discriminated_union_override_switch_variant(tmp_path): + """Providing an explicit 'type' overrides the default variant.""" + write_toml(tmp_path / "cfg.toml", {"variant": {"type": "b", "beta": 0.5}}) + config = cli(DummyConfig, args=["@", str(tmp_path / "cfg.toml")]) + assert config.variant.type == "b" + assert config.variant.beta == 0.5 + + +def test_cli_overrides_defaults(): + """CLI args override defaults.""" + config = cli(DummyConfig, args=["--name", "my-run", "--seed", "7"]) + assert config.name == "my-run" + assert config.seed == 7 + assert config.nested.lr == 1e-4 + + +def test_toml_overrides_defaults(tmp_path): + """TOML overrides defaults.""" + write_toml(tmp_path / "cfg.toml", {"name": "my-run", "seed": 7, "nested": {"lr": 3e-4}}) + config = cli(DummyConfig, args=["@", str(tmp_path / "cfg.toml")]) + assert config.name == "my-run" + assert config.seed == 7 + assert config.nested.lr == 3e-4 + + +def test_cli_overrides_toml(tmp_path): + """CLI args override TOML.""" + write_toml(tmp_path / "cfg.toml", {"seed": 1, "nested": {"lr": 3e-4}}) + config = cli(DummyConfig, args=["@", str(tmp_path / "cfg.toml"), "--seed", "99", "--nested.lr", "5e-5"]) + assert config.seed == 99 + assert config.nested.lr == 5e-5 + # TOML value not overridden by CLI should still be applied (not reverted to class default) + assert config.nested.weight_decay == 0.01 + + +def test_removed_fused_lm_head_chunk_size_field_is_rejected(): + with pytest.raises(ValidationError, match="fused_lm_head_chunk_size"): + TrainerModelConfig.model_validate({"fused_lm_head_chunk_size": "auto"}) + + +def test_selective_activation_checkpointing_requires_custom_impl(): + with pytest.raises(ValidationError, match="Selective activation checkpointing requires model.impl='custom'"): + TrainerModelConfig.model_validate({"impl": "hf", "ac": {"mode": "selective"}}) diff --git a/tests/unit/test_tool_call_parser.py b/tests/unit/test_tool_call_parser.py new file mode 100644 index 00000000000..113f22114e3 --- /dev/null +++ b/tests/unit/test_tool_call_parser.py @@ -0,0 +1,49 @@ +import pytest + +from prime_rl.inference.vllm.server import resolve_tool_call_parser + + +@pytest.mark.parametrize( + "model_name,expected_parser", + [ + # GLM-4.5 + ("zai-org/GLM-4.5", "glm45"), + ("zai-org/GLM-4.5-Air", "glm45"), + ("zai-org/GLM-4.5V", "glm45"), + # GLM-4.7 + ("zai-org/GLM-4.7", "glm47"), + ("zai-org/GLM-4.7-Flash", "glm47"), + # GLM-5.1 + ("zai-org/GLM-5.1", "glm47"), + # MiniMax + ("MiniMaxAI/MiniMax-M2", "minimax_m2"), + ("MiniMaxAI/MiniMax-M2.1", "minimax_m2"), + ("MiniMaxAI/MiniMax-M2.5", "minimax_m2"), + # INTELLECT + ("PrimeIntellect/INTELLECT-3", "hermes"), + ("PrimeIntellect/INTELLECT-3-FP8", "hermes"), + ("PrimeIntellect/INTELLECT-3.1", "hermes"), + # Qwen3 + ("Qwen/Qwen3-0.6B", "hermes"), + ("Qwen/Qwen3-32B", "hermes"), + ("Qwen/Qwen3-235B-A22B", "hermes"), + ("Qwen/Qwen3-4B-Instruct-2507", "hermes"), + ("Qwen/Qwen3-Coder-480B-A35B-Instruct", "hermes"), + ("Qwen/Qwen3-Next-80B-A3B-Instruct", "hermes"), + ("Qwen/Qwen3.5-397B-A17B", "qwen3_coder"), + ], +) +def test_auto_detect_tool_call_parser(model_name: str, expected_parser: str): + assert resolve_tool_call_parser(model_name, "auto") == expected_parser + + +def test_explicit_parser_not_overridden(): + assert resolve_tool_call_parser("Qwen/Qwen3-4B-Instruct-2507", "qwen3_xml") == "qwen3_xml" + + +def test_auto_unknown_model(): + assert resolve_tool_call_parser("some/unknown-model", "auto") is None + + +def test_none_skips_resolution(): + assert resolve_tool_call_parser("Qwen/Qwen3-0.6B", None) is None diff --git a/tests/unit/train/models/afmoe_hf_modeling/configuration_afmoe.py b/tests/unit/train/models/afmoe_hf_modeling/configuration_afmoe.py index 762cb60ba29..d1338b24883 100644 --- a/tests/unit/train/models/afmoe_hf_modeling/configuration_afmoe.py +++ b/tests/unit/train/models/afmoe_hf_modeling/configuration_afmoe.py @@ -13,11 +13,11 @@ # See the License for the specific language governing permissions and # limitations under the License. from transformers.configuration_utils import PretrainedConfig, layer_type_validation -from transformers.modeling_rope_utils import rope_config_validation from transformers.utils import logging logger = logging.get_logger(__name__) + class AfmoeConfig(PretrainedConfig): """ n_group (`int`, *optional*, defaults to 1): @@ -25,6 +25,7 @@ class AfmoeConfig(PretrainedConfig): topk_group (`int`, *optional*, defaults to 1): Number of selected groups for each token(for each token, ensuring the selected experts is only within `topk_group` groups). """ + model_type = "afmoe" base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), @@ -68,6 +69,7 @@ def __init__( attention_dropout: float = 0.0, n_group: int = 1, topk_group: int = 1, + pad_token_id: int | None = None, **kwargs, ): self.vocab_size = vocab_size @@ -84,8 +86,7 @@ def __init__( self.use_cache = use_cache self.rope_theta = rope_theta self.rope_scaling = rope_scaling - - + # MoE specific self.moe_intermediate_size = moe_intermediate_size self.num_experts_per_tok = num_experts_per_tok @@ -101,7 +102,6 @@ def __init__( self.load_balance_coeff = load_balance_coeff self.use_grouped_mm = use_grouped_mm - # Attention specific self.attention_dropout = attention_dropout self.global_attn_every_n_layers = global_attn_every_n_layers @@ -109,7 +109,8 @@ def __init__( self.layer_types = layer_types if self.layer_types is None: self.layer_types = [ - "sliding_attention" if bool((i + 1) % global_attn_every_n_layers) else "full_attention" for i in range(self.num_hidden_layers) + "sliding_attention" if bool((i + 1) % global_attn_every_n_layers) else "full_attention" + for i in range(self.num_hidden_layers) ] layer_type_validation(self.layer_types) @@ -120,12 +121,7 @@ def __init__( num_key_value_heads = num_attention_heads self.num_key_value_heads = num_key_value_heads - - - # Validate rope configs - if self.rope_scaling is not None and "type" in self.rope_scaling: - self.rope_scaling["rope_type"] = self.rope_scaling["type"] - rope_config_validation(self) + self.pad_token_id = pad_token_id super().__init__( tie_word_embeddings=tie_word_embeddings, diff --git a/tests/unit/train/models/afmoe_hf_modeling/modeling_afmoe.py b/tests/unit/train/models/afmoe_hf_modeling/modeling_afmoe.py index 156f903a734..e5fd36c96e4 100644 --- a/tests/unit/train/models/afmoe_hf_modeling/modeling_afmoe.py +++ b/tests/unit/train/models/afmoe_hf_modeling/modeling_afmoe.py @@ -26,8 +26,8 @@ except: from configuration_afmoe import AfmoeConfig -class AfmoeRotaryEmbedding(nn.Module): +class AfmoeRotaryEmbedding(nn.Module): def __init__(self, config: AfmoeConfig, device=None): super().__init__() # BC: "rope_type" was originally "type" @@ -45,6 +45,20 @@ def __init__(self, config: AfmoeConfig, device=None): self.register_buffer("inv_freq", inv_freq, persistent=False) self.original_inv_freq = self.inv_freq + def compute_default_rope_parameters(self, config=None, device=None, seq_len=None, layer_type=None): + """Required by transformers 5.0.0 for weight initialization when rope_type is 'default'.""" + cfg = config or self.config + cfg.standardize_rope_params() + rope_params = cfg.rope_parameters[layer_type] if layer_type is not None else cfg.rope_parameters + base = rope_params["rope_theta"] + partial_rotary_factor = rope_params.get("partial_rotary_factor", 1.0) + head_dim = getattr(cfg, "head_dim", None) or cfg.hidden_size // cfg.num_attention_heads + dim = int(head_dim * partial_rotary_factor) + inv_freq = 1.0 / ( + base ** (torch.arange(0, dim, 2, dtype=torch.int64).to(device=device, dtype=torch.float) / dim) + ) + return inv_freq, 1.0 + def _dynamic_frequency_update(self, position_ids, device): """ dynamic RoPE layers should recompute `inv_freq` in the following situations: @@ -129,11 +143,10 @@ def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor: batch, num_key_value_heads, slen, head_dim = hidden_states.shape if n_rep == 1: return hidden_states - hidden_states = hidden_states[:, :, None, :, :].expand( - batch, num_key_value_heads, n_rep, slen, head_dim - ) + hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim) return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim) + @use_kernel_forward_from_hub("RMSNorm") class AfmoeRMSNorm(nn.Module): def __init__(self, hidden_size: int, eps: float): @@ -155,7 +168,6 @@ def extra_repr(self): return f"{tuple(self.weight.shape)}, eps={self.variance_epsilon}" - def eager_attention_forward( module: nn.Module, query: torch.Tensor, @@ -174,12 +186,8 @@ def eager_attention_forward( causal_mask = attention_mask[:, :, :, : key_states.shape[-2]] attn_weights = attn_weights + causal_mask - attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to( - query.dtype - ) - attn_weights = nn.functional.dropout( - attn_weights, p=dropout, training=module.training - ) + attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query.dtype) + attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training) attn_output = torch.matmul(attn_weights, value_states) attn_output = attn_output.transpose(1, 2).contiguous() @@ -212,7 +220,7 @@ def __init__(self, config): self.score_func = config.score_func self.route_norm = config.route_norm self.route_scale = config.route_scale - self.gate = nn.Linear(config.hidden_size, config.num_experts, bias=False) + self.gate = nn.Linear(config.hidden_size, config.num_experts, bias=False) def forward(self, hidden_states, expert_bias: torch.Tensor | None): _, _, hidden_dim = hidden_states.shape @@ -240,6 +248,7 @@ def forward(self, hidden_states, expert_bias: torch.Tensor | None): top_scores = top_scores * self.route_scale return top_scores, selected_experts + class AfmoeMoE(nn.Module): def __init__(self, config): super().__init__() @@ -248,16 +257,11 @@ def __init__(self, config): self.shared_experts = None if config.num_shared_experts > 0: - self.shared_experts = AfmoeMLP( - config, config.moe_intermediate_size * config.num_shared_experts - ) + self.shared_experts = AfmoeMLP(config, config.moe_intermediate_size * config.num_shared_experts) self.experts = nn.ModuleList( - [AfmoeMLP( - config, intermediate_size=config.moe_intermediate_size - ) for _ in range(config.num_experts)] + [AfmoeMLP(config, intermediate_size=config.moe_intermediate_size) for _ in range(config.num_experts)] ) self.expert_bias = nn.Parameter(torch.zeros(config.num_experts, dtype=torch.float32), requires_grad=False) - def forward(self, hidden_states): batch_size, seq_len, hidden_dim = hidden_states.shape @@ -279,12 +283,8 @@ def forward(self, hidden_states): token_indices_sorted = token_indices_sorted // self.config.num_experts_per_tok # Gather input tokens - token_indices_expanded = token_indices_sorted.unsqueeze(-1).expand( - -1, hidden_dim - ) - routed_input = torch.gather( - hidden_states_flat, dim=0, index=token_indices_expanded - ) + token_indices_expanded = token_indices_sorted.unsqueeze(-1).expand(-1, hidden_dim) + routed_input = torch.gather(hidden_states_flat, dim=0, index=token_indices_expanded) routed_output = torch.zeros_like(routed_input) for expert_id in range(self.config.num_experts): @@ -293,15 +293,11 @@ def forward(self, hidden_states): expert_input = routed_input[mask] expert_out = self.experts[expert_id](expert_input) routed_output[mask] = expert_out - - routed_output = ( - routed_output.to(torch.float32) * top_scores_sorted.unsqueeze(-1) - ).to(hidden_states.dtype) + + routed_output = (routed_output.to(torch.float32) * top_scores_sorted.unsqueeze(-1)).to(hidden_states.dtype) # Scatter back to original positions - output = shared_output.scatter_add( - dim=0, index=token_indices_expanded, src=routed_output - ) + output = shared_output.scatter_add(dim=0, index=token_indices_expanded, src=routed_output) return output.view(batch_size, seq_len, hidden_dim) @@ -323,25 +319,15 @@ def __init__(self, config: AfmoeConfig, layer_idx: int): self.is_local_attention = config.layer_types[layer_idx] == "sliding_attention" self.sliding_window = config.sliding_window if self.is_local_attention else None - self.q_proj = nn.Linear( - config.hidden_size, self.num_heads * self.head_dim, bias=False - ) - self.k_proj = nn.Linear( - config.hidden_size, self.num_key_value_heads * self.head_dim, bias=False - ) - self.v_proj = nn.Linear( - config.hidden_size, self.num_key_value_heads * self.head_dim, bias=False - ) - self.o_proj = nn.Linear( - self.num_heads * self.head_dim, config.hidden_size, bias=False - ) + self.q_proj = nn.Linear(config.hidden_size, self.num_heads * self.head_dim, bias=False) + self.k_proj = nn.Linear(config.hidden_size, self.num_key_value_heads * self.head_dim, bias=False) + self.v_proj = nn.Linear(config.hidden_size, self.num_key_value_heads * self.head_dim, bias=False) + self.o_proj = nn.Linear(self.num_heads * self.head_dim, config.hidden_size, bias=False) self.q_norm = AfmoeRMSNorm(self.head_dim, eps=config.rms_norm_eps) self.k_norm = AfmoeRMSNorm(self.head_dim, eps=config.rms_norm_eps) - self.gate_proj = nn.Linear( - config.hidden_size, self.num_heads * self.head_dim, bias=False - ) + self.gate_proj = nn.Linear(config.hidden_size, self.num_heads * self.head_dim, bias=False) def forward( self, @@ -351,8 +337,7 @@ def forward( past_key_value: Optional[Cache] = None, cache_position: Optional[torch.LongTensor] = None, **kwargs: Unpack[TransformersKwargs], - ) -> torch.Tensor: - + ) -> torch.Tensor: input_shape = hidden_states.shape[:-1] hidden_shape = (*input_shape, -1, self.head_dim) @@ -363,7 +348,7 @@ def forward( query_states = self.q_norm(query_states) key_states = self.k_norm(key_states) - + query_states = query_states.transpose(1, 2) key_states = key_states.transpose(1, 2) value_states = value_states.transpose(1, 2) @@ -378,9 +363,7 @@ def forward( attention_interface: Callable = eager_attention_forward if self.config._attn_implementation != "eager": - attention_interface = ALL_ATTENTION_FUNCTIONS[ - self.config._attn_implementation - ] + attention_interface = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation] output, _ = attention_interface( self, @@ -492,14 +475,9 @@ def __init__(self, config: AfmoeConfig): self.padding_idx = config.pad_token_id self.vocab_size = config.vocab_size - self.embed_tokens = nn.Embedding( - config.vocab_size, config.hidden_size, self.padding_idx - ) + self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx) self.layers = nn.ModuleList( - [ - AfmoeDecoderLayer(config, layer_idx) - for layer_idx in range(config.num_hidden_layers) - ] + [AfmoeDecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)] ) self.norm = AfmoeRMSNorm(config.hidden_size, eps=config.rms_norm_eps) self.rotary_emb = AfmoeRotaryEmbedding(config=config) @@ -513,7 +491,6 @@ def get_input_embeddings(self): def set_input_embeddings(self, value): self.embed_tokens = value - def forward( self, input_ids: torch.LongTensor, @@ -526,9 +503,7 @@ def forward( **kwargs: Unpack[TransformersKwargs], ) -> MoeModelOutputWithPast: if (input_ids is None) ^ (inputs_embeds is not None): - raise ValueError( - "You must specify exactly one of input_ids or inputs_embeds" - ) + raise ValueError("You must specify exactly one of input_ids or inputs_embeds") if use_cache and past_key_values is None: past_key_values = DynamicCache() @@ -537,9 +512,7 @@ def forward( inputs_embeds = self.embed_tokens(input_ids) if cache_position is None: - past_seen_tokens = ( - past_key_values.get_seq_length() if past_key_values is not None else 0 - ) + past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0 cache_position = torch.arange( past_seen_tokens, past_seen_tokens + inputs_embeds.shape[1], @@ -552,7 +525,7 @@ def forward( if not isinstance(causal_mask_mapping := attention_mask, dict): mask_kwargs = { "config": self.config, - "input_embeds": inputs_embeds, + "inputs_embeds": inputs_embeds, "attention_mask": attention_mask, "cache_position": cache_position, "past_key_values": past_key_values, @@ -590,7 +563,7 @@ def forward( class AfmoeForCausalLM(AfmoePreTrainedModel, GenerationMixin): - _tied_weights_keys = ["lm_head.weight"] + _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": "colwise_rep"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} @@ -648,18 +621,13 @@ def forward( hidden_states = outputs.last_hidden_state # Only compute necessary logits - slice_indices = ( - slice(-logits_to_keep, None) - if isinstance(logits_to_keep, int) - else logits_to_keep - ) + slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep logits = self.lm_head(hidden_states[:, slice_indices, :]) loss = None if labels is not None: loss = self.loss_function(logits, labels, self.vocab_size, **kwargs) - return MoeCausalLMOutputWithPast( loss=loss, logits=logits, diff --git a/tests/unit/train/models/test_afmoe.py b/tests/unit/train/models/test_afmoe.py index 6c4e2b751ee..036e24d195a 100644 --- a/tests/unit/train/models/test_afmoe.py +++ b/tests/unit/train/models/test_afmoe.py @@ -70,12 +70,17 @@ def get_model_pairs(): return hf_model, prime_model +class _IdentityMLP(nn.Identity): + def forward(self, x, **kwargs): + return super().forward(x) + + def test_afmoe_attn_only() -> None: hf_model, prime_model = get_model_pairs() for layer in hf_model.model.layers: layer.mlp = nn.Identity() for layer in prime_model.model.layers: - layer.mlp = nn.Identity() + layer.mlp = _IdentityMLP() with torch.device("cuda"), default_dtype(torch.float32): input_ids = torch.randint(0, hf_model.config.vocab_size, (1, 100)) diff --git a/tests/unit/train/models/test_checkpointing.py b/tests/unit/train/models/test_checkpointing.py new file mode 100644 index 00000000000..90fe3f32c3a --- /dev/null +++ b/tests/unit/train/models/test_checkpointing.py @@ -0,0 +1,76 @@ +import torch.nn as nn + +from prime_rl.trainer.models.layers.checkpointing import ( + get_supported_targets, + set_selective_activation_checkpointing, +) + +_PATCHED_METHODS_ATTR = "_prime_rl_selective_ac_patched_methods" + + +class DummySelfAttention(nn.Module): + def attn_projections(self, hidden_states, position_embeddings=None): + return hidden_states + + def output_proj(self, attn_output): + return attn_output + + def forward(self, hidden_states): + return hidden_states + + +class DummySlidingAttentionLayer(nn.Module): + def __init__(self): + super().__init__() + self.self_attn = DummySelfAttention() + self.attention_type = "sliding_attention" + + +class DummyMamba(nn.Module): + def forward(self, hidden_states): + return hidden_states + + +class DummyMambaLayer(nn.Module): + def __init__(self): + super().__init__() + self.mamba = DummyMamba() + + +class DummyMoEMlp(nn.Module): + def forward(self, hidden_states): + return hidden_states + + def _run_routed_experts(self, hidden_states, *args): + return hidden_states + + def _run_local_routed_experts(self, hidden_states, num_tokens_per_expert): + return hidden_states + + +class DummyMoELayer(nn.Module): + def __init__(self): + super().__init__() + self.mlp = DummyMoEMlp() + + +def test_get_supported_targets_treats_mamba_as_linear_attention(): + assert get_supported_targets(DummyMambaLayer()) == frozenset({"norm", "linear_attn"}) + + +def test_sliding_attention_linear_attn_subsumes_attn_proj_hooks(): + layer = DummySlidingAttentionLayer() + + set_selective_activation_checkpointing(layer, ["attn_proj", "linear_attn"]) + + assert getattr(layer.self_attn, _PATCHED_METHODS_ATTR) == frozenset({"forward"}) + + +def test_routed_experts_checkpointing_patches_local_and_global_helpers(): + layer = DummyMoELayer() + + assert "routed_experts" in get_supported_targets(layer) + + set_selective_activation_checkpointing(layer, ["routed_experts"]) + + assert getattr(layer.mlp, _PATCHED_METHODS_ATTR) == frozenset({"_run_local_routed_experts", "_run_routed_experts"}) diff --git a/tests/unit/train/models/test_glm4_moe.py b/tests/unit/train/models/test_glm4_moe.py index fd50b0d17c5..67aaf0be73f 100644 --- a/tests/unit/train/models/test_glm4_moe.py +++ b/tests/unit/train/models/test_glm4_moe.py @@ -47,12 +47,17 @@ def get_model_pairs() -> tuple[HFGlm4MoeForCausalLM, PrimeRLGlm4MoeForCausalLM]: return hf_model, prime_model +class _IdentityMLP(nn.Identity): + def forward(self, x, **kwargs): + return super().forward(x) + + def test_glm4_moe_attn_only() -> None: hf_model, prime_model = get_model_pairs() for layer in hf_model.model.layers: layer.mlp = nn.Identity() for layer in prime_model.model.layers: - layer.mlp = nn.Identity() + layer.mlp = _IdentityMLP() with torch.device("cuda"), default_dtype(torch.float32): input_ids = torch.randint(0, hf_model.config.vocab_size, (1, 100)) @@ -118,21 +123,6 @@ def test_glm4_moe() -> None: grad_diff = hf_model.model.embed_tokens.weight.grad - prime_model.model.embed_tokens.weight.grad assert torch.allclose(grad_diff, torch.zeros_like(grad_diff), atol=2), f"Max grad diff: {grad_diff.abs().max()}" - with torch.device("cuda"), default_dtype(torch.float32): - hf_from_prime_model = HFGlm4MoeForCausalLM._from_config(hf_model.config) - converted_state_dict = prime_model.convert_to_hf(prime_model.state_dict()) - hf_from_prime_model.load_state_dict(converted_state_dict) - - hf_from_prime_output = hf_from_prime_model(input_ids, position_ids) - hf_from_prime_output.logits.sum().backward() - - logits_diff = hf_from_prime_output.logits - hf_output.logits - assert torch.allclose(logits_diff, torch.zeros_like(logits_diff), atol=2e-2), ( - f"Max logits diff: {logits_diff.abs().max()}" - ) - grad_diff = hf_from_prime_model.model.embed_tokens.weight.grad - hf_model.model.embed_tokens.weight.grad - assert torch.allclose(grad_diff, torch.zeros_like(grad_diff), atol=2), f"Max grad diff: {grad_diff.abs().max()}" - if __name__ == "__main__": test_glm4_moe_mlp_only() diff --git a/tests/unit/train/models/test_glm_moe_dsa_conversion.py b/tests/unit/train/models/test_glm_moe_dsa_conversion.py new file mode 100644 index 00000000000..88018032147 --- /dev/null +++ b/tests/unit/train/models/test_glm_moe_dsa_conversion.py @@ -0,0 +1,68 @@ +import torch + +from prime_rl.trainer.models.glm_moe_dsa.converting_glm_moe_dsa import convert_tt_layer_to_vllm_kernel + + +def _build_prime_layer_state(layer_idx: int = 0) -> dict[str, torch.Tensor]: + prefix = f"model.layers.{layer_idx}" + return { + f"{prefix}.self_attn.q_a_proj.weight": torch.randn(4, 6), + f"{prefix}.self_attn.kv_a_proj_with_mqa.weight": torch.randn(3, 6), + f"{prefix}.self_attn.q_b_proj.weight": torch.randn(5, 6), + f"{prefix}.self_attn.kv_b_proj.weight": torch.randn(7, 6), + f"{prefix}.self_attn.o_proj.weight": torch.randn(6, 5), + f"{prefix}.self_attn.indexer.wq_b.weight": torch.randn(4, 6), + f"{prefix}.self_attn.indexer.wk.weight": torch.randn(4, 6), + f"{prefix}.self_attn.indexer.k_norm.weight": torch.randn(6), + f"{prefix}.self_attn.indexer.k_norm.bias": torch.randn(6), + f"{prefix}.self_attn.indexer.weights_proj.weight": torch.randn(2, 6), + f"{prefix}.mlp.gate_proj.weight": torch.randn(8, 6), + f"{prefix}.mlp.up_proj.weight": torch.randn(8, 6), + f"{prefix}.mlp.down_proj.weight": torch.randn(6, 8), + f"{prefix}.mlp.router.gate.weight": torch.randn(4, 6), + f"{prefix}.mlp.expert_bias": torch.randn(4), + f"{prefix}.mlp.experts.w1": torch.randn(2, 3, 6), + f"{prefix}.mlp.experts.w2": torch.randn(2, 6, 3), + f"{prefix}.mlp.experts.w3": torch.randn(2, 3, 6), + f"{prefix}.mlp.shared_expert.w1": torch.randn(1, 3, 6), + f"{prefix}.mlp.shared_expert.w2": torch.randn(1, 6, 3), + f"{prefix}.mlp.shared_expert.w3": torch.randn(1, 3, 6), + } + + +def test_convert_tt_layer_to_vllm_kernel_no_fp8(): + state = _build_prime_layer_state() + out = convert_tt_layer_to_vllm_kernel(state, layer_idx=0, quantize_fp8=False) + + assert "model.layers.0.self_attn.fused_qkv_a_proj.weight" in out + assert out["model.layers.0.self_attn.fused_qkv_a_proj.weight"].shape == (7, 6) + assert "model.layers.0.self_attn.indexer.wk_weights_proj.weight" in out + assert out["model.layers.0.self_attn.indexer.wk_weights_proj.weight"].shape == (6, 6) + assert "model.layers.0.self_attn.indexer.wk.weight" not in out + assert "model.layers.0.self_attn.indexer.weights_proj.weight" not in out + + assert "model.layers.0.mlp.gate_up_proj.weight" in out + assert out["model.layers.0.mlp.gate_up_proj.weight"].shape == (16, 6) + + assert "model.layers.0.mlp.experts.w13_weight" in out + assert out["model.layers.0.mlp.experts.w13_weight"].shape == (2, 6, 6) + assert "model.layers.0.mlp.experts.w2_weight" in out + assert out["model.layers.0.mlp.experts.w2_weight"].shape == (2, 6, 3) + + assert "model.layers.0.mlp.gate.weight" in out + assert "model.layers.0.mlp.gate.e_score_correction_bias" in out + + +def test_convert_tt_layer_to_vllm_kernel_with_fp8(): + state = _build_prime_layer_state() + out = convert_tt_layer_to_vllm_kernel(state, layer_idx=0, quantize_fp8=True) + + assert out["model.layers.0.self_attn.fused_qkv_a_proj.weight"].dtype == torch.float8_e4m3fn + assert out["model.layers.0.self_attn.fused_qkv_a_proj.weight_scale_inv"].dtype == torch.float32 + assert out["model.layers.0.self_attn.indexer.wk_weights_proj.weight"].dtype == torch.float32 + assert "model.layers.0.self_attn.indexer.wk.weight_scale_inv" not in out + + assert out["model.layers.0.mlp.experts.w13_weight"].dtype == torch.float8_e4m3fn + assert out["model.layers.0.mlp.experts.w13_weight_scale_inv"].dtype == torch.float32 + assert out["model.layers.0.mlp.experts.w2_weight"].dtype == torch.float8_e4m3fn + assert out["model.layers.0.mlp.experts.w2_weight_scale_inv"].dtype == torch.float32 diff --git a/tests/unit/train/models/test_nemotron_h.py b/tests/unit/train/models/test_nemotron_h.py new file mode 100644 index 00000000000..404594a183f --- /dev/null +++ b/tests/unit/train/models/test_nemotron_h.py @@ -0,0 +1,240 @@ +import pytest +import torch +from transformers.models.nemotron_h.configuration_nemotron_h import NemotronHConfig as HFNemotronHConfig +from transformers.models.nemotron_h.modeling_nemotron_h import ( + NemotronHAttention, +) +from transformers.models.nemotron_h.modeling_nemotron_h import ( + NemotronHForCausalLM as HFNemotronHForCausalLM, +) + +from prime_rl.trainer.models.layers.lm_head import inject_prime_lm_head +from prime_rl.trainer.models.nemotron_h import NemotronHConfig, NemotronHForCausalLM +from prime_rl.trainer.models.nemotron_h.modeling_nemotron_h import NemotronHAttentionLayer +from prime_rl.utils.utils import default_dtype + +pytestmark = [pytest.mark.gpu] + +# Shared small-model hyperparams (satisfy mamba_expand * hidden_size == mamba_num_heads * mamba_head_dim) +_BASE = dict( + vocab_size=256, + hidden_size=256, + num_attention_heads=4, + num_key_value_heads=2, + head_dim=64, + max_position_embeddings=128, + intermediate_size=512, + mamba_expand=2, + mamba_num_heads=8, + mamba_head_dim=64, + ssm_state_size=64, + mamba_n_groups=1, + mamba_d_conv=4, + mamba_chunk_size=64, + n_routed_experts=4, + n_shared_experts=1, + moe_intermediate_size=256, + moe_shared_expert_intermediate_size=256, + moe_latent_size=128, + num_experts_per_tok=2, + n_group=1, + topk_group=1, + norm_topk_prob=True, + routed_scaling_factor=1.0, +) + + +def get_model_pairs(): + """Create an HF model and a PrimeRL model with shared weights.""" + hf_config = HFNemotronHConfig(**_BASE, hybrid_override_pattern="ME*E") + hf_config._attn_implementation = "sdpa" + + prime_config = NemotronHConfig( + **_BASE, + layers_block_type=["mamba", "moe", "attention", "moe"], + use_grouped_mm=False, + ) + prime_config._attn_implementation = "sdpa" + + with torch.device("cuda"), default_dtype(torch.float32): + hf_model = HFNemotronHForCausalLM._from_config(hf_config) + prime_model = NemotronHForCausalLM._from_config(prime_config) + + with torch.no_grad(): + state_dict = hf_model.state_dict() + prime_state_keys = prime_model.state_dict().keys() + NemotronHForCausalLM.convert_to_prime(state_dict) + prime_model.load_state_dict(state_dict) + + inject_prime_lm_head(prime_model, chunk_size=None) + assert set(prime_state_keys) - set(state_dict.keys()) == set() + return hf_model, prime_model + + +def test_nemotron_h_mamba_moe_only(): + """Test Mamba and MoE layers produce identical outputs (attention bypassed).""" + hf_model, prime_model = get_model_pairs() + + # Bypass attention layers in both models so only Mamba+MoE are exercised + for layer in hf_model.model.layers: + if isinstance(layer.mixer, NemotronHAttention): + layer.forward = lambda hidden_states, **kwargs: hidden_states + + for layer in prime_model.model.layers: + if isinstance(layer, NemotronHAttentionLayer): + layer.forward = lambda hidden_states, **kwargs: hidden_states + + torch.manual_seed(42) + with torch.device("cuda"), default_dtype(torch.float32): + input_ids = torch.randint(0, 256, (1, 32)) + position_ids = torch.arange(0, 32).unsqueeze(0) + + hf_output = hf_model(input_ids, position_ids=position_ids) + prime_output = prime_model(input_ids, position_ids=position_ids) + hf_output.logits.sum().backward() + prime_output["logits"].sum().backward() + + logits_diff = prime_output["logits"] - hf_output.logits + assert torch.allclose(logits_diff, torch.zeros_like(logits_diff), atol=1e-2), ( + f"Max logits diff: {logits_diff.abs().max()}" + ) + grad_diff = hf_model.model.embeddings.weight.grad - prime_model.model.embed_tokens.weight.grad + assert torch.allclose(grad_diff, torch.zeros_like(grad_diff), atol=2), f"Max grad diff: {grad_diff.abs().max()}" + + +@pytest.mark.xfail(reason="HF NemotronH now uses fused expert tensors; convert_to_hf produces individual expert format") +def test_nemotron_h_reverse(): + """Test reverse: PrimeRL weights loaded into HF model produce identical outputs.""" + prime_config = NemotronHConfig( + **_BASE, + layers_block_type=["mamba", "moe", "attention", "moe"], + use_grouped_mm=False, + ) + prime_config._attn_implementation = "sdpa" + + hf_config = HFNemotronHConfig(**_BASE, hybrid_override_pattern="ME*E") + hf_config._attn_implementation = "sdpa" + + with torch.device("cuda"), default_dtype(torch.float32): + prime_model = NemotronHForCausalLM._from_config(prime_config) + hf_model = HFNemotronHForCausalLM._from_config(hf_config) + + inject_prime_lm_head(prime_model, chunk_size=None) + + with torch.no_grad(): + sd = prime_model.state_dict() + NemotronHForCausalLM.convert_to_hf(sd) + # convert_to_hf produces checkpoint format with "backbone." prefix; + # the HF model uses "model." prefix for its state dict + keys_to_rename = [k for k in sd if k.startswith("backbone.")] + for key in keys_to_rename: + sd["model." + key[len("backbone.") :]] = sd.pop(key) + hf_model.load_state_dict(sd) + + # Bypass attention to isolate Mamba+MoE matching + for layer in hf_model.model.layers: + if isinstance(layer.mixer, NemotronHAttention): + layer.forward = lambda hidden_states, **kwargs: hidden_states + for layer in prime_model.model.layers: + if isinstance(layer, NemotronHAttentionLayer): + layer.forward = lambda hidden_states, **kwargs: hidden_states + + with torch.device("cuda"), default_dtype(torch.float32): + input_ids = torch.randint(0, 256, (1, 32)) + position_ids = torch.arange(0, 32).unsqueeze(0) + + hf_output = hf_model(input_ids, position_ids=position_ids) + prime_output = prime_model(input_ids, position_ids=position_ids) + + logits_diff = prime_output["logits"] - hf_output.logits + assert torch.allclose(logits_diff, torch.zeros_like(logits_diff), atol=1e-2), ( + f"Max logits diff: {logits_diff.abs().max()}" + ) + + +def test_nemotron_h(): + """Test full model (Mamba + MoE + Attention) produces close outputs.""" + hf_model, prime_model = get_model_pairs() + + with torch.device("cuda"), default_dtype(torch.float32): + input_ids = torch.randint(0, 256, (1, 32)) + position_ids = torch.arange(0, 32).unsqueeze(0) + + hf_output = hf_model(input_ids, position_ids=position_ids) + prime_output = prime_model(input_ids, position_ids=position_ids) + # Slightly larger tolerance due to different SDPA attention implementations + logits_diff = prime_output["logits"] - hf_output.logits + assert torch.allclose(logits_diff, torch.zeros_like(logits_diff), atol=5e-2), ( + f"Max logits diff: {logits_diff.abs().max()}" + ) + + +def test_nemotron_h_backward(): + """Verify all parameters receive non-zero gradients.""" + prime_config = NemotronHConfig( + **_BASE, + layers_block_type=["mamba", "moe", "attention", "moe"], + use_grouped_mm=False, + ) + model = NemotronHForCausalLM(prime_config).to("cuda") + inject_prime_lm_head(model) + + input_ids = torch.randint(0, 256, (2, 16), device="cuda") + output = model(input_ids) + output["logits"].sum().backward() + + zero_grads = [] + for name, p in model.named_parameters(): + if p.numel() == 0: + continue + if p.grad is None or p.grad.norm().item() == 0: + zero_grads.append(name) + assert not zero_grads, f"Parameters with zero/no gradients: {zero_grads}" + + +def test_nemotron_h_weight_conversion_roundtrip(): + """Verify PrimeRL -> HF -> PrimeRL conversion preserves all weights.""" + prime_config = NemotronHConfig( + **_BASE, + layers_block_type=["mamba", "moe", "attention", "moe"], + use_grouped_mm=False, + ) + model = NemotronHForCausalLM(prime_config).to("cuda") + original_sd = {k: v.clone() for k, v in model.state_dict().items()} + + sd = model.state_dict() + NemotronHForCausalLM.convert_to_hf(sd) + assert NemotronHForCausalLM.is_hf_state_dict(sd) + NemotronHForCausalLM.convert_to_prime(sd) + assert NemotronHForCausalLM.is_prime_state_dict(sd) + + for key in original_sd: + assert key in sd, f"Missing key after roundtrip: {key}" + assert torch.equal(original_sd[key], sd[key]), f"Value mismatch for {key}" + + +def test_nemotron_h_hybrid_override_pattern(): + """Verify hybrid_override_pattern correctly maps to layers_block_type.""" + config = NemotronHConfig(**_BASE, hybrid_override_pattern="ME*E") + assert config.layers_block_type == ["mamba", "moe", "attention", "moe"] + assert config.num_hidden_layers == 4 + + +def test_nemotron_h_no_latent_projection(): + """Verify model works without latent projections (moe_latent_size=None).""" + prime_config = NemotronHConfig( + **{**_BASE, "moe_latent_size": None}, + layers_block_type=["mamba", "moe", "attention", "moe"], + use_grouped_mm=False, + ) + model = NemotronHForCausalLM(prime_config).to("cuda") + inject_prime_lm_head(model) + + input_ids = torch.randint(0, 256, (2, 16), device="cuda") + output = model(input_ids) + assert output["logits"].shape == (2, 16, 256) + + output["logits"].sum().backward() + for name, p in model.named_parameters(): + if "experts.w1" in name and p.numel() > 0: + assert p.grad is not None and p.grad.norm().item() > 0, f"Zero grad for {name}" diff --git a/tests/unit/train/models/test_nemotron_h_kl.py b/tests/unit/train/models/test_nemotron_h_kl.py new file mode 100644 index 00000000000..7e590cadf9b --- /dev/null +++ b/tests/unit/train/models/test_nemotron_h_kl.py @@ -0,0 +1,316 @@ +"""Test KL mismatch computation for NemotronH models. + +Creates two small NemotronH models (policy + reference), perturbs +the policy to simulate SFT drift, computes logprobs from both, and +verifies KL mismatch through the loss pipeline. +""" + +import pytest +import torch + +from prime_rl.configs.trainer import DefaultLossConfig +from prime_rl.trainer.models.layers.lm_head import inject_prime_lm_head +from prime_rl.trainer.models.nemotron_h import NemotronHConfig, NemotronHForCausalLM +from prime_rl.trainer.rl.loss import ( + LossInputs, + default_loss_fn, + selective_log_softmax, + shift_tensor_right, +) +from prime_rl.utils.utils import default_dtype + +pytestmark = [pytest.mark.gpu] + +_BASE = dict( + vocab_size=256, + hidden_size=256, + num_attention_heads=4, + num_key_value_heads=2, + head_dim=64, + max_position_embeddings=128, + intermediate_size=512, + mamba_expand=2, + mamba_num_heads=8, + mamba_head_dim=64, + ssm_state_size=64, + mamba_n_groups=1, + mamba_d_conv=4, + mamba_chunk_size=64, + n_routed_experts=4, + n_shared_experts=1, + moe_intermediate_size=256, + moe_shared_expert_intermediate_size=256, + moe_latent_size=128, + num_experts_per_tok=2, + n_group=1, + topk_group=1, + norm_topk_prob=True, + routed_scaling_factor=1.0, +) + + +def _make_model(device="cuda"): + config = NemotronHConfig( + **_BASE, + layers_block_type=["mamba", "moe", "attention", "moe"], + use_grouped_mm=False, + ) + config._attn_implementation = "sdpa" + with torch.device(device), default_dtype(torch.float32): + model = NemotronHForCausalLM._from_config(config) + inject_prime_lm_head(model, chunk_size=None) + return model + + +def _get_logprobs_vanilla(model, input_ids): + """Get logprobs using VanillaOutputLinear (returns logits, we compute logprobs).""" + with torch.no_grad(): + out = model(input_ids) + logits = out["logits"] + labels = torch.cat( + [input_ids[:, 1:], torch.zeros(input_ids.shape[0], 1, dtype=torch.long, device=input_ids.device)], dim=1 + ) + logprobs = selective_log_softmax(logits, labels) + logprobs = shift_tensor_right(logprobs, pad_value=torch.log(torch.tensor(1.0 / model.config.vocab_size)).item()) + return logprobs + + +def _perturb_model(model, scale=0.1): + """Add small noise to model weights to simulate SFT drift.""" + with torch.no_grad(): + for p in model.parameters(): + if p.numel() > 0: + p.add_(torch.randn_like(p) * scale) + + +def test_kl_zero_when_identical(): + """Two identical models should have zero KL mismatch.""" + model = _make_model() + + input_ids = torch.randint(0, 256, (2, 32), device="cuda") + logprobs = _get_logprobs_vanilla(model, input_ids) + + loss_mask = torch.ones(32, dtype=torch.bool, device="cuda") + # mask out first token (shifted right padding) + loss_mask[0] = False + advantages = torch.ones(32, device="cuda") + + for i in range(logprobs.shape[0]): + inputs = LossInputs( + trainer_logprobs=logprobs[i], + inference_logprobs=logprobs[i], + teacher_logprobs=None, + advantages=advantages, + loss_mask=loss_mask, + ) + result = default_loss_fn(inputs, DefaultLossConfig()) + + assert result.metrics["mismatch_kl"].item() == pytest.approx(0.0, abs=1e-6), ( + f"Expected zero KL for identical models, got {result.metrics['mismatch_kl'].item()}" + ) + + +def test_kl_positive_after_perturbation(): + """Perturbing the policy model should produce positive KL.""" + ref_model = _make_model() + policy_model = _make_model() + + # Copy reference weights to policy, then perturb + with torch.no_grad(): + policy_model.load_state_dict(ref_model.state_dict()) + _perturb_model(policy_model, scale=0.05) + + input_ids = torch.randint(0, 256, (2, 32), device="cuda") + + ref_logprobs = _get_logprobs_vanilla(ref_model, input_ids) + policy_logprobs = _get_logprobs_vanilla(policy_model, input_ids) + + loss_mask = torch.ones(32, dtype=torch.bool, device="cuda") + loss_mask[0] = False + advantages = torch.ones(32, device="cuda") + + for i in range(ref_logprobs.shape[0]): + inputs = LossInputs( + trainer_logprobs=policy_logprobs[i], + inference_logprobs=ref_logprobs[i], + teacher_logprobs=None, + advantages=advantages, + loss_mask=loss_mask, + ) + result = default_loss_fn(inputs, DefaultLossConfig()) + kl = result.metrics["mismatch_kl"].item() + + assert kl > 0, f"Expected positive KL after perturbation, got {kl}" + assert kl < 100, f"KL unexpectedly large: {kl}" + + +def test_kl_increases_with_larger_perturbation(): + """Larger perturbation should produce larger KL.""" + ref_model = _make_model() + policy_small = _make_model() + policy_large = _make_model() + + with torch.no_grad(): + policy_small.load_state_dict(ref_model.state_dict()) + policy_large.load_state_dict(ref_model.state_dict()) + + _perturb_model(policy_small, scale=0.01) + _perturb_model(policy_large, scale=0.1) + + input_ids = torch.randint(0, 256, (1, 32), device="cuda") + + ref_logprobs = _get_logprobs_vanilla(ref_model, input_ids) + small_logprobs = _get_logprobs_vanilla(policy_small, input_ids) + large_logprobs = _get_logprobs_vanilla(policy_large, input_ids) + + loss_mask = torch.ones(32, dtype=torch.bool, device="cuda") + loss_mask[0] = False + + log_ratio_small = small_logprobs[0] - ref_logprobs[0] + ratio_small = torch.exp(log_ratio_small) + kl_small = ((ratio_small - log_ratio_small - 1) * loss_mask).sum() / loss_mask.sum() + + log_ratio_large = large_logprobs[0] - ref_logprobs[0] + ratio_large = torch.exp(log_ratio_large) + kl_large = ((ratio_large - log_ratio_large - 1) * loss_mask).sum() / loss_mask.sum() + + assert kl_large > kl_small, ( + f"Larger perturbation should produce larger KL: small={kl_small.item()}, large={kl_large.item()}" + ) + + +def test_kl_with_fused_lm_head(): + """FusedOutputLinear should produce same logprobs as VanillaOutputLinear.""" + config = NemotronHConfig( + **_BASE, + layers_block_type=["mamba", "moe", "attention", "moe"], + use_grouped_mm=False, + ) + config._attn_implementation = "sdpa" + + with torch.device("cuda"), default_dtype(torch.float32): + model = NemotronHForCausalLM._from_config(config) + + # Get logits from vanilla head + inject_prime_lm_head(model, chunk_size=None) + input_ids = torch.randint(0, 256, (1, 16), device="cuda") + labels = torch.cat([input_ids[:, 1:], torch.zeros(1, 1, dtype=torch.long, device="cuda")], dim=1) + + with torch.no_grad(): + vanilla_out = model(input_ids) + vanilla_logits = vanilla_out["logits"] + temperature = torch.ones(1, 16, device="cuda") + vanilla_logprobs = selective_log_softmax(vanilla_logits / temperature.unsqueeze(-1), labels) + + # Now switch to fused head and compare + inject_prime_lm_head(model, chunk_size=16) + with torch.no_grad(): + fused_out = model(input_ids, labels=labels, temperature=temperature) + fused_logprobs = fused_out["logprobs"] + + diff = (vanilla_logprobs - fused_logprobs).abs().max() + assert diff < 1e-3, f"Vanilla vs fused logprob diff: {diff.item()}" + + +def test_kl_logprob_alignment(): + """Verify logprobs are properly aligned between trainer and inference after shift.""" + model = _make_model() + + input_ids = torch.randint(0, 256, (1, 16), device="cuda") + + # Simulate what the training loop does + with torch.no_grad(): + out = model(input_ids) + logits = out["logits"] + + # Labels = shifted input_ids (predict next token) + labels = torch.cat([input_ids[:, 1:], torch.zeros(1, 1, dtype=torch.long, device="cuda")], dim=1) + + logprobs = selective_log_softmax(logits, labels) + logprobs_shifted = shift_tensor_right( + logprobs, pad_value=torch.log(torch.tensor(1.0 / model.config.vocab_size)).item() + ) + + # Position 0 should be the pad value (uniform distribution logprob) + expected_pad = torch.log(torch.tensor(1.0 / model.config.vocab_size)).item() + assert logprobs_shifted[0, 0].item() == pytest.approx(expected_pad, abs=1e-5), ( + f"Position 0 should be pad value {expected_pad}, got {logprobs_shifted[0, 0].item()}" + ) + + # Positions 1..N-1 should contain actual logprobs (negative values) + actual_logprobs = logprobs_shifted[0, 1:] + assert (actual_logprobs < 0).all(), "Logprobs should be negative" + assert (actual_logprobs > -20).all(), "Logprobs shouldn't be extremely negative" + + +def test_kl_backward_through_policy(): + """KL loss should produce gradients for the policy model.""" + ref_model = _make_model() + policy_model = _make_model() + + with torch.no_grad(): + policy_model.load_state_dict(ref_model.state_dict()) + _perturb_model(policy_model, scale=0.05) + + input_ids = torch.randint(0, 256, (1, 16), device="cuda") + labels = torch.cat([input_ids[:, 1:], torch.zeros(1, 1, dtype=torch.long, device="cuda")], dim=1) + + # Reference logprobs (detached) + with torch.no_grad(): + ref_out = ref_model(input_ids) + ref_logprobs = selective_log_softmax(ref_out["logits"], labels) + ref_logprobs = shift_tensor_right( + ref_logprobs, pad_value=torch.log(torch.tensor(1.0 / ref_model.config.vocab_size)).item() + ).detach() + + # Policy logprobs (with grad) + policy_out = policy_model(input_ids) + policy_logprobs = selective_log_softmax(policy_out["logits"], labels) + policy_logprobs = shift_tensor_right( + policy_logprobs, pad_value=torch.log(torch.tensor(1.0 / policy_model.config.vocab_size)).item() + ) + + loss_mask = torch.ones(16, dtype=torch.bool, device="cuda") + loss_mask[0] = False + + log_ratio = policy_logprobs[0] - ref_logprobs[0] + kl_loss = (loss_mask * log_ratio**2).sum() + kl_loss.backward() + + # Check gradients flow to policy model parameters + has_grad = False + for name, p in policy_model.named_parameters(): + if p.grad is not None and p.grad.norm().item() > 0: + has_grad = True + break + + assert has_grad, "KL loss should produce non-zero gradients for policy model" + + +def test_kl_mismatch_formula_sanity(): + """Verify the mismatch_kl formula: r - log(r) - 1 >= 0 for all r > 0.""" + model = _make_model() + ref_model = _make_model() + + with torch.no_grad(): + ref_model.load_state_dict(model.state_dict()) + _perturb_model(model, scale=0.1) + + input_ids = torch.randint(0, 256, (2, 32), device="cuda") + + trainer_logprobs = _get_logprobs_vanilla(model, input_ids) + inference_logprobs = _get_logprobs_vanilla(ref_model, input_ids) + + log_ratio = trainer_logprobs - inference_logprobs + ratio = torch.exp(log_ratio) + mismatch_kl = ratio - log_ratio - 1 + + # KL divergence should always be non-negative (up to numerical precision) + assert (mismatch_kl >= -1e-5).all(), f"KL should be non-negative, min={mismatch_kl.min().item()}" + + # For identical models (log_ratio=0, ratio=1), KL = 1 - 0 - 1 = 0 + # For divergent models, KL > 0 + loss_mask = torch.ones(32, dtype=torch.bool, device="cuda") + loss_mask[0] = False + mean_kl = mismatch_kl[0][loss_mask].mean() + assert mean_kl > 0, f"Mean KL should be positive for perturbed model, got {mean_kl.item()}" diff --git a/tests/unit/train/models/test_qwen3_5_moe.py b/tests/unit/train/models/test_qwen3_5_moe.py new file mode 100644 index 00000000000..6ab8d7f1bef --- /dev/null +++ b/tests/unit/train/models/test_qwen3_5_moe.py @@ -0,0 +1,135 @@ +import pytest +import torch +from transformers import Qwen3_5MoeForCausalLM as HFQwen3_5MoeForCausalLM + +from prime_rl.trainer.models.layers.lm_head import inject_prime_lm_head +from prime_rl.trainer.models.qwen3_5_moe import Qwen3_5MoeConfig +from prime_rl.trainer.models.qwen3_5_moe import Qwen3_5MoeForCausalLM as PrimeRLQwen3_5MoeForCausalLM +from prime_rl.utils.utils import default_dtype + +pytestmark = [pytest.mark.gpu] + + +def get_model_pairs(): + config = Qwen3_5MoeConfig( + vocab_size=256, + hidden_size=256, + num_hidden_layers=4, + num_attention_heads=4, + num_key_value_heads=2, + head_dim=64, + moe_intermediate_size=128, + shared_expert_intermediate_size=128, + num_experts=8, + num_experts_per_tok=2, + max_position_embeddings=512, + rms_norm_eps=1e-6, + linear_conv_kernel_dim=4, + linear_key_head_dim=32, + linear_value_head_dim=32, + linear_num_key_heads=4, + linear_num_value_heads=8, + use_grouped_mm=False, + ) + config._attn_implementation = "sdpa" + with torch.device("cuda"), default_dtype(torch.float32): + hf_model = HFQwen3_5MoeForCausalLM._from_config(config) + prime_model = PrimeRLQwen3_5MoeForCausalLM._from_config(config) + with torch.no_grad(): + state_dict = hf_model.state_dict() + prime_state_keys = prime_model.state_dict().keys() + prime_model.convert_to_prime(state_dict) + prime_model.load_state_dict(state_dict) + inject_prime_lm_head(prime_model, chunk_size=None) + assert set(prime_state_keys) - set(state_dict.keys()) == set() + return hf_model, prime_model + + +def test_qwen3_5_moe(): + hf_model, prime_model = get_model_pairs() + + with torch.device("cuda"), default_dtype(torch.float32): + input_ids = torch.randint(0, hf_model.config.vocab_size, (1, 100)) + position_ids = torch.arange(1, 101).unsqueeze(0) + + hf_output = hf_model(input_ids, position_ids=position_ids) + prime_output = prime_model(input_ids, position_ids=position_ids) + hf_output.logits.sum().backward() + prime_output["logits"].sum().backward() + + logits_diff = prime_output["logits"] - hf_output.logits + assert torch.allclose(logits_diff, torch.zeros_like(logits_diff), atol=2e-2), ( + f"Max logits diff: {logits_diff.abs().max()}" + ) + grad_diff = hf_model.model.embed_tokens.weight.grad - prime_model.model.embed_tokens.weight.grad + assert torch.allclose(grad_diff, torch.zeros_like(grad_diff), atol=2), f"Max grad diff: {grad_diff.abs().max()}" + + +def test_qwen3_5_moe_roundtrip(): + """Verify HF → PrimeRL → HF weight conversion is lossless at the state_dict level.""" + hf_model, prime_model = get_model_pairs() + + # Get original HF state_dict and the PrimeRL-converted version + original_hf_sd = hf_model.state_dict() + prime_sd = prime_model.state_dict() + + # Convert PrimeRL → per-expert HF format + converted_hf_sd = PrimeRLQwen3_5MoeForCausalLM.convert_to_hf(dict(prime_sd)) + + # Also convert original HF (fused) to per-expert format for comparison + + # First convert original HF → PrimeRL, then back to per-expert HF + orig_prime_sd = dict(original_hf_sd) + PrimeRLQwen3_5MoeForCausalLM.convert_to_prime(orig_prime_sd) + orig_roundtripped = dict(orig_prime_sd) + PrimeRLQwen3_5MoeForCausalLM.convert_to_hf(orig_roundtripped) + + # All non-expert keys should match exactly, expert keys should match after roundtrip + for key in orig_roundtripped: + assert key in converted_hf_sd, f"Missing key: {key}" + assert torch.equal(orig_roundtripped[key], converted_hf_sd[key]), f"Mismatch at {key}" + + +def test_qwen3_5_moe_router_replay(): + """When routed_experts are provided, the model uses them instead of computing routing.""" + _, prime_model = get_model_pairs() + + with torch.device("cuda"), default_dtype(torch.float32): + input_ids = torch.randint(0, prime_model.config.vocab_size, (1, 100)) + position_ids = torch.arange(1, 101).unsqueeze(0) + + out_normal = prime_model(input_ids, position_ids=position_ids) + + num_layers = prime_model.config.num_hidden_layers + topk = prime_model.config.num_experts_per_tok + routed_experts = torch.randint(0, prime_model.config.num_experts, (1, 100, num_layers, topk), device="cuda") + + prime_model.zero_grad() + out_replay = prime_model(input_ids, position_ids=position_ids, routed_experts=routed_experts) + + assert out_replay["logits"].shape == out_normal["logits"].shape + + out_replay["logits"].sum().backward() + assert prime_model.model.embed_tokens.weight.grad is not None + + +def test_qwen3_5_moe_cp_patching(): + """Verify substitute_ring_attn patches Qwen3_5MoeGatedFlashAttention._compute_attention.""" + from unittest.mock import MagicMock + + from prime_rl.trainer.models.layers.attn import substitute_ring_attn + from prime_rl.trainer.models.qwen3_5_moe.modeling_qwen3_5_moe import Qwen3_5MoeGatedFlashAttention + + original_method = Qwen3_5MoeGatedFlashAttention._compute_attention + + mock_group = MagicMock() + substitute_ring_attn(process_group=mock_group, heads_k_stride=1) + + assert Qwen3_5MoeGatedFlashAttention._compute_attention is not original_method + + # Restore to avoid polluting other tests + Qwen3_5MoeGatedFlashAttention._compute_attention = original_method + + +if __name__ == "__main__": + test_qwen3_5_moe() diff --git a/tests/unit/train/models/test_qwen3_5_moe_vlm.py b/tests/unit/train/models/test_qwen3_5_moe_vlm.py new file mode 100644 index 00000000000..c629f6be255 --- /dev/null +++ b/tests/unit/train/models/test_qwen3_5_moe_vlm.py @@ -0,0 +1,216 @@ +import pytest +import torch +from transformers import AutoConfig +from transformers.models.qwen3_5_moe.modeling_qwen3_5_moe import ( + Qwen3_5MoeForConditionalGeneration as HFQwen3_5MoeVLM, +) + +from prime_rl.trainer.model import can_reinit_empty_buffers +from prime_rl.trainer.models.layers.lm_head import inject_prime_lm_head +from prime_rl.trainer.models.qwen3_5_moe import Qwen3_5MoeForCausalLM +from prime_rl.utils.utils import default_dtype + +pytestmark = [pytest.mark.gpu] + + +def _tiny_vlm_config(): + """HF composite config shrunk for unit testing.""" + config = AutoConfig.from_pretrained("Qwen/Qwen3.5-35B-A3B", trust_remote_code=True, attn_implementation="sdpa") + config.use_cache = False + tc = config.text_config + tc.vocab_size = 256 + tc.hidden_size = 256 + tc.num_hidden_layers = 2 + tc.layer_types = ["linear_attention", "full_attention"] + tc.num_attention_heads = 4 + tc.num_key_value_heads = 2 + tc.head_dim = 64 + tc.moe_intermediate_size = 128 + tc.shared_expert_intermediate_size = 128 + tc.num_experts = 4 + tc.num_experts_per_tok = 2 + tc.max_position_embeddings = 512 + tc.linear_key_head_dim = 32 + tc.linear_value_head_dim = 32 + tc.linear_num_key_heads = 4 + tc.linear_num_value_heads = 8 + tc.use_cache = False + + vc = config.vision_config + vc.depth = 2 + vc.hidden_size = 128 + vc.intermediate_size = 256 + vc.num_heads = 4 + vc.out_hidden_size = tc.hidden_size + + # Special token IDs must fit within the tiny vocab + config.image_token_id = 250 + config.video_token_id = 251 + config.vision_start_token_id = 252 + config.vision_end_token_id = 253 + return config + + +def _make_image_inputs(config, device="cuda", dtype=torch.float32): + """Create minimal image inputs matching the vision config.""" + vc = config.vision_config + patch_dim = vc.in_channels * vc.temporal_patch_size * vc.patch_size * vc.patch_size + image_grid_thw = torch.tensor([[1, 2, 2]], device=device) + num_patches = int(image_grid_thw.prod().item()) + pixel_values = torch.randn(num_patches, patch_dim, device=device, dtype=dtype) + num_image_tokens = num_patches // (vc.spatial_merge_size**2) + return pixel_values, image_grid_thw, num_image_tokens + + +def test_vlm_forward(): + """Custom VLM produces logits for both text-only and multimodal inputs.""" + config = _tiny_vlm_config() + with torch.device("cuda"), default_dtype(torch.float32): + model = Qwen3_5MoeForCausalLM(config) + inject_prime_lm_head(model) + + vocab = config.text_config.vocab_size + + # Text-only (avoid special token range 250-253) + input_ids = torch.randint(0, 200, (1, 20), device="cuda") + position_ids = torch.arange(1, 21, device="cuda").unsqueeze(0) + out_text = model(input_ids=input_ids, position_ids=position_ids) + assert out_text["logits"].shape == (1, 20, vocab) + + # Multimodal + pixel_values, image_grid_thw, n_img_tokens = _make_image_inputs(config) + text_part = torch.randint(0, 200, (1, 10), device="cuda") + img_part = torch.full((1, n_img_tokens), config.image_token_id, device="cuda") + input_ids_mm = torch.cat([text_part[:, :5], img_part, text_part[:, 5:]], dim=1) + + out_mm = model(input_ids=input_ids_mm, pixel_values=pixel_values, image_grid_thw=image_grid_thw) + assert out_mm["logits"].shape == (1, input_ids_mm.shape[1], vocab) + + +def test_vlm_backward(): + """Gradients flow through both vision scatter and text model.""" + config = _tiny_vlm_config() + with torch.device("cuda"), default_dtype(torch.float32): + model = Qwen3_5MoeForCausalLM(config) + inject_prime_lm_head(model) + + pixel_values, image_grid_thw, n_img_tokens = _make_image_inputs(config) + text_part = torch.randint(0, 200, (1, 10), device="cuda") + img_part = torch.full((1, n_img_tokens), config.image_token_id, device="cuda") + input_ids = torch.cat([text_part[:, :5], img_part, text_part[:, 5:]], dim=1) + + out = model(input_ids=input_ids, pixel_values=pixel_values, image_grid_thw=image_grid_thw) + out["logits"].sum().backward() + + assert model.model.language_model.embed_tokens.weight.grad is not None + assert model.model.visual.patch_embed.proj.weight.grad is not None + + +def test_vlm_weight_load_from_hf(): + """Weights from HF VLM checkpoint load correctly into custom VLM after conversion. + + Text model numerical match is already validated by test_qwen3_5_moe.py::test_qwen3_5_moe. + This test verifies that VLM weight conversion + loading produces a working model. + """ + config = _tiny_vlm_config() + with torch.device("cuda"), default_dtype(torch.float32): + hf_model = HFQwen3_5MoeVLM._from_config(config) + prime_model = Qwen3_5MoeForCausalLM(config) + + # Copy weights: HF -> PrimeRL (with MoE conversion) + with torch.no_grad(): + hf_sd = hf_model.state_dict() + prime_model.convert_to_prime(hf_sd) + prime_model.load_state_dict(hf_sd) + inject_prime_lm_head(prime_model) + + # Verify vision encoder weights match exactly (should be untouched by conversion) + for name, param in hf_model.model.visual.named_parameters(): + prime_param = dict(prime_model.model.visual.named_parameters())[name] + assert torch.equal(param, prime_param), f"Vision weight mismatch: {name}" + + # Verify model produces output after weight loading + input_ids = torch.randint(0, 200, (1, 20), device="cuda") + position_ids = torch.arange(1, 21, device="cuda").unsqueeze(0) + out = prime_model(input_ids=input_ids, position_ids=position_ids) + assert out["logits"].shape[2] == config.text_config.vocab_size + assert not torch.isnan(out["logits"]).any() + + +def test_vlm_weight_roundtrip(): + """HF -> PrimeRL -> HF weight conversion is lossless (vision keys untouched, text keys converted).""" + config = _tiny_vlm_config() + with torch.device("cuda"), default_dtype(torch.float32): + hf_model = HFQwen3_5MoeVLM._from_config(config) + + hf_sd = hf_model.state_dict() + original_vision_key = "model.visual.blocks.0.mlp.linear_fc1.weight" + original_vision_weight = hf_sd[original_vision_key].clone() + + # HF -> PrimeRL + prime_sd = dict(hf_sd) + Qwen3_5MoeForCausalLM.convert_to_prime(prime_sd) + assert any("language_model" in k and "mlp.experts.w1" in k for k in prime_sd) + assert original_vision_key in prime_sd + + # PrimeRL -> HF + roundtripped = dict(prime_sd) + Qwen3_5MoeForCausalLM.convert_to_hf(roundtripped) + + # Original HF also needs roundtrip for expert format normalization + orig_rt = dict(hf_sd) + Qwen3_5MoeForCausalLM.convert_to_prime(orig_rt) + Qwen3_5MoeForCausalLM.convert_to_hf(orig_rt) + + for key in orig_rt: + assert key in roundtripped, f"Missing key: {key}" + assert torch.equal(orig_rt[key], roundtripped[key]), f"Mismatch at {key}" + + # Vision weights preserved through the whole roundtrip + assert torch.equal(roundtripped[original_vision_key], original_vision_weight) + + +def test_vlm_router_replay(): + """routed_experts bypasses router computation in VLM multimodal forward.""" + config = _tiny_vlm_config() + with torch.device("cuda"), default_dtype(torch.float32): + model = Qwen3_5MoeForCausalLM(config) + inject_prime_lm_head(model) + + vocab = config.text_config.vocab_size + pixel_values, image_grid_thw, n_img_tokens = _make_image_inputs(config) + text_part = torch.randint(0, 200, (1, 10), device="cuda") + img_part = torch.full((1, n_img_tokens), config.image_token_id, device="cuda") + input_ids = torch.cat([text_part[:, :5], img_part, text_part[:, 5:]], dim=1) + seq_len = input_ids.shape[1] + + num_layers = config.text_config.num_hidden_layers + topk = config.text_config.num_experts_per_tok + routed_experts = torch.randint(0, config.text_config.num_experts, (1, seq_len, num_layers, topk), device="cuda") + + out = model( + input_ids=input_ids, pixel_values=pixel_values, image_grid_thw=image_grid_thw, routed_experts=routed_experts + ) + assert out["logits"].shape == (1, seq_len, vocab) + + out["logits"].sum().backward() + assert model.model.language_model.embed_tokens.weight.grad is not None + + +def test_vlm_meta_device_and_buffer_reinit(): + """Model can be created on meta device and buffers reinitialized.""" + config = _tiny_vlm_config() + with torch.device("meta"): + model = Qwen3_5MoeForCausalLM.from_config(config) + + assert can_reinit_empty_buffers(model) + + model.to_empty(device="cuda") + model.init_buffers_post_meta() + + lm_inv = model.model.language_model.rotary_emb.inv_freq + vis_inv = model.model.visual.rotary_pos_emb.inv_freq + assert lm_inv.device.type == "cuda" + assert vis_inv.device.type == "cuda" + assert lm_inv.abs().sum() > 0 + assert vis_inv.abs().sum() > 0 diff --git a/tests/unit/train/models/test_qwen3_moe.py b/tests/unit/train/models/test_qwen3_moe.py index 3838bb6bac5..00835fd3d33 100644 --- a/tests/unit/train/models/test_qwen3_moe.py +++ b/tests/unit/train/models/test_qwen3_moe.py @@ -46,12 +46,17 @@ def get_model_pairs(): return hf_model, prime_model +class _IdentityMLP(nn.Identity): + def forward(self, x, **kwargs): + return super().forward(x) + + def test_qwen3_moe_attn_only(): hf_model, prime_model = get_model_pairs() for layer in hf_model.model.layers: layer.mlp = nn.Identity() for layer in prime_model.model.layers: - layer.mlp = nn.Identity() + layer.mlp = _IdentityMLP() with torch.device("cuda"), default_dtype(torch.float32): input_ids = torch.randint(0, hf_model.config.vocab_size, (1, 100)) @@ -117,20 +122,34 @@ def test_qwen3_moe(): grad_diff = hf_model.model.embed_tokens.weight.grad - prime_model.model.embed_tokens.weight.grad assert torch.allclose(grad_diff, torch.zeros_like(grad_diff), atol=2), f"Max grad diff: {grad_diff.abs().max()}" + +def test_qwen3_moe_router_replay(): + """When routed_experts are provided, the model uses them instead of computing routing.""" + _, prime_model = get_model_pairs() + with torch.device("cuda"), default_dtype(torch.float32): - hf_from_prime_model = HFQwen3MoeForCausalLM._from_config(hf_model.config) - converted_state_dict = prime_model.convert_to_hf(prime_model.state_dict()) - hf_from_prime_model.load_state_dict(converted_state_dict) + input_ids = torch.randint(0, prime_model.config.vocab_size, (1, 100)) + position_ids = torch.arange(1, 101).unsqueeze(0) - hf_from_prime_output = hf_from_prime_model(input_ids, position_ids) - hf_from_prime_output.logits.sum().backward() + # Forward without router replay + out_normal = prime_model(input_ids, position_ids) - logits_diff = hf_from_prime_output.logits - hf_output.logits - assert torch.allclose(logits_diff, torch.zeros_like(logits_diff), atol=2e-2), ( - f"Max logits diff: {logits_diff.abs().max()}" - ) - grad_diff = hf_from_prime_model.model.embed_tokens.weight.grad - hf_model.model.embed_tokens.weight.grad - assert torch.allclose(grad_diff, torch.zeros_like(grad_diff), atol=2), f"Max grad diff: {grad_diff.abs().max()}" + # Construct routed_experts with fixed expert indices + # Shape: [batch=1, seq_len=100, num_hidden_layers=3, num_experts_per_tok=4] + num_layers = prime_model.config.num_hidden_layers + topk = prime_model.config.num_experts_per_tok + routed_experts = torch.randint(0, prime_model.config.num_experts, (1, 100, num_layers, topk), device="cuda") + + # Forward with router replay + prime_model.zero_grad() + out_replay = prime_model(input_ids, position_ids, routed_experts=routed_experts) + + # Outputs should differ because routing is forced to different experts + assert out_replay["logits"].shape == out_normal["logits"].shape + + # Verify gradients flow through the model with router replay + out_replay["logits"].sum().backward() + assert prime_model.model.embed_tokens.weight.grad is not None if __name__ == "__main__": diff --git a/tests/unit/train/rl/test_fused_lm_head.py b/tests/unit/train/rl/test_fused_lm_head.py index 776b2453eb3..e3149a7948f 100644 --- a/tests/unit/train/rl/test_fused_lm_head.py +++ b/tests/unit/train/rl/test_fused_lm_head.py @@ -11,10 +11,12 @@ def _baseline_logprobs_and_entropy( - hidden: torch.Tensor, weight: torch.Tensor, labels: torch.Tensor, *, temperature: float + hidden: torch.Tensor, weight: torch.Tensor, labels: torch.Tensor, *, temperature: torch.Tensor ) -> tuple[torch.Tensor, torch.Tensor]: + """Baseline logprobs and entropy with per-token temperature tensor.""" logits = hidden @ weight.t() - logits = logits / float(temperature) + # temperature is [b, s], logits is [b, s, v] + logits = logits / temperature.unsqueeze(-1) logp = torch.log_softmax(logits, dim=-1).gather(dim=-1, index=labels.unsqueeze(-1)).squeeze(-1) ent = compute_entropy(logits) return logp, ent @@ -23,8 +25,8 @@ def _baseline_logprobs_and_entropy( def test_fused_lm_head_matches_full_logits_forward_and_backward_cpu(): torch.manual_seed(0) b, s, h, v = 2, 4, 8, 37 - temperature = 1.7 - chunk_size = 11 + temperature = torch.full((b, s), 1.7, dtype=torch.float32) + chunk_size = 3 hidden0 = torch.randn(b, s, h, dtype=torch.float32, requires_grad=True) labels = torch.randint(0, v, (b, s), dtype=torch.long) @@ -66,12 +68,13 @@ def test_fused_lm_head_requires_labels(): hidden = torch.randn(b, s, h, dtype=torch.float32) weight = torch.randn(v, h, dtype=torch.float32) + temperature = torch.full((b, s), 1.0, dtype=torch.float32) lm = FusedOutputLinear(in_features=h, out_features=v, chunk_size=5) lm.weight = torch.nn.Parameter(weight) with pytest.raises(AssertionError, match="FusedOutputLinear requires labels"): - lm(hidden, labels=None, temperature=1.0) + lm(hidden, labels=None, temperature=temperature) def test_vanilla_lm_head_returns_logits(): @@ -85,7 +88,8 @@ def test_vanilla_lm_head_returns_logits(): lm = VanillaOutputLinear(in_features=h, out_features=v) lm.weight = torch.nn.Parameter(weight) - out = lm(hidden, labels=None, temperature=1.0) + # VanillaOutputLinear doesn't use temperature - it just returns logits + out = lm(hidden, labels=None, temperature=None) assert out.get("logits") is not None assert out.get("logprobs") is None assert out.get("entropy") is None @@ -98,8 +102,9 @@ def test_fused_vs_vanilla_integration(): """Integration test comparing fused and vanilla outputs after postprocessing.""" torch.manual_seed(42) b, s, h, v = 2, 4, 8, 37 - temperature = 1.7 - chunk_size = 11 + temp_value = 1.7 + temperature = torch.full((b, s), temp_value, dtype=torch.float32) + chunk_size = 3 hidden = torch.randn(b, s, h, dtype=torch.float16) labels = torch.randint(0, v, (b, s), dtype=torch.long) @@ -108,10 +113,10 @@ def test_fused_vs_vanilla_integration(): # Vanilla path: get logits, compute logprobs manually vanilla_lm = VanillaOutputLinear(in_features=h, out_features=v) vanilla_lm.weight = torch.nn.Parameter(weight.clone()) - vanilla_out = cast_float_and_contiguous(vanilla_lm(hidden, labels=None, temperature=temperature)) + vanilla_out = cast_float_and_contiguous(vanilla_lm(hidden, labels=None, temperature=None)) assert vanilla_out.get("logits") is not None - logits = vanilla_out["logits"] / float(temperature) + logits = vanilla_out["logits"] / temp_value vanilla_logprobs = torch.log_softmax(logits, dim=-1).gather(dim=-1, index=labels.unsqueeze(-1)).squeeze(-1) vanilla_entropy = compute_entropy(logits) @@ -158,7 +163,7 @@ def test_full_model_fused_vs_vanilla(): # Wrap with different LM heads inject_prime_lm_head(model_vanilla, chunk_size=None) # Vanilla - inject_prime_lm_head(model_fused, chunk_size=256) # Fused with chunking + inject_prime_lm_head(model_fused, chunk_size=32) # Fused with chunking # Setup optimizers optimizer_vanilla = torch.optim.AdamW(model_vanilla.parameters(), lr=1e-4) @@ -167,13 +172,14 @@ def test_full_model_fused_vs_vanilla(): # Run a few training steps num_steps = 3 batch_size, seq_len = 2, 64 - temperature = 1.5 + temp_value = 1.5 for step in range(num_steps): # Generate random batch with torch.device("cuda"): labels = torch.randint(0, config.vocab_size, (batch_size, seq_len)) position_ids = torch.arange(seq_len).unsqueeze(0).expand(batch_size, -1) + temperature = torch.full((batch_size, seq_len), temp_value, dtype=torch.float32, device="cuda") # Vanilla forward (returns logits, compute logprobs/entropy using RL train functions) optimizer_vanilla.zero_grad() @@ -182,7 +188,7 @@ def test_full_model_fused_vs_vanilla(): ) if out_vanilla.get("logprobs") is None: assert out_vanilla.get("logits") is not None - logits = out_vanilla["logits"] / float(temperature) + logits = out_vanilla["logits"] / temp_value out_vanilla["logprobs"] = selective_log_softmax(logits, labels) out_vanilla["entropy"] = compute_entropy(logits) loss_vanilla = -out_vanilla["logprobs"].mean() @@ -194,7 +200,7 @@ def test_full_model_fused_vs_vanilla(): out_fused = cast_float_and_contiguous(model_fused(labels, position_ids, labels=labels, temperature=temperature)) if out_fused.get("logprobs") is None: assert out_fused.get("logits") is not None - logits = out_fused["logits"] / float(temperature) + logits = out_fused["logits"] / temp_value out_fused["logprobs"] = selective_log_softmax(logits, labels) out_fused["entropy"] = compute_entropy(logits) loss_fused = -out_fused["logprobs"].mean() @@ -222,7 +228,8 @@ def test_fused_lm_head_correct_shift(): """ torch.manual_seed(999) b, s, h, v = 2, 16, 32, 50 - temperature = 1.5 + temp_value = 1.5 + temperature = torch.full((b, s), temp_value, dtype=torch.float32) chunk_size = 13 hidden = torch.randn(b, s, h, dtype=torch.float32) @@ -240,7 +247,7 @@ def test_fused_lm_head_correct_shift(): # === Inference convention (baseline) === logits = hidden @ weight.t() - logits = logits / temperature + logits = logits / temp_value # Shift logits right (prepend zeros, drop last) to get inference convention shifted_logits = torch.cat([torch.zeros(b, 1, v, dtype=logits.dtype), logits[:, :-1, :]], dim=1) inference_logprobs = ( @@ -289,12 +296,12 @@ def test_inject_prime_lm_head_vanilla(): # Test forward with labels and temperature batch_size, seq_len = 2, 64 - temperature = 1.5 with torch.device("cuda"): input_ids = torch.randint(0, config.vocab_size, (batch_size, seq_len)) position_ids = torch.arange(seq_len).unsqueeze(0).expand(batch_size, -1) labels = torch.randint(0, config.vocab_size, (batch_size, seq_len)) + temperature = torch.full((batch_size, seq_len), 1.5, dtype=torch.float32) with torch.autocast("cuda", dtype=torch.bfloat16): out = model(input_ids=input_ids, position_ids=position_ids, labels=labels, temperature=temperature) @@ -329,19 +336,19 @@ def test_inject_prime_lm_head_fused(): with torch.device("cuda"), default_dtype(torch.float32): model = AutoModelForCausalLM.from_config(config) - # Wrap with FusedOutputLinear (chunk_size=512) + # Wrap with FusedOutputLinear inject_prime_lm_head(model, chunk_size=512) assert isinstance(model.lm_head, FusedOutputLinear), "lm_head should be FusedOutputLinear" # Test forward with labels and temperature batch_size, seq_len = 2, 64 - temperature = 1.5 with torch.device("cuda"): input_ids = torch.randint(0, config.vocab_size, (batch_size, seq_len)) position_ids = torch.arange(seq_len).unsqueeze(0).expand(batch_size, -1) labels = torch.randint(0, config.vocab_size, (batch_size, seq_len)) + temperature = torch.full((batch_size, seq_len), 1.5, dtype=torch.float32) with torch.autocast("cuda", dtype=torch.bfloat16): out = model(input_ids=input_ids, position_ids=position_ids, labels=labels, temperature=temperature) @@ -383,16 +390,17 @@ def test_hf_model_fused_vs_vanilla_matches(): # Wrap models inject_prime_lm_head(model_vanilla, chunk_size=None) - inject_prime_lm_head(model_fused, chunk_size=512) + inject_prime_lm_head(model_fused, chunk_size=32) # Test data batch_size, seq_len = 2, 64 - temperature = 1.5 + temp_value = 1.5 with torch.device("cuda"): input_ids = torch.randint(0, config.vocab_size, (batch_size, seq_len)) position_ids = torch.arange(seq_len).unsqueeze(0).expand(batch_size, -1) labels = torch.randint(0, config.vocab_size, (batch_size, seq_len)) + temperature = torch.full((batch_size, seq_len), temp_value, dtype=torch.float32) # Run vanilla model with torch.autocast("cuda", dtype=torch.bfloat16): @@ -401,7 +409,7 @@ def test_hf_model_fused_vs_vanilla_matches(): ) # Compute logprobs and entropy from vanilla logits - logits = out_vanilla["logits"].float() / float(temperature) + logits = out_vanilla["logits"].float() / temp_value vanilla_logprobs = selective_log_softmax(logits, labels) vanilla_entropy = compute_entropy(logits) diff --git a/tests/unit/train/rl/test_loss.py b/tests/unit/train/rl/test_loss.py index c299d51d5ce..696897e3689 100644 --- a/tests/unit/train/rl/test_loss.py +++ b/tests/unit/train/rl/test_loss.py @@ -1,8 +1,8 @@ import pytest import torch -from prime_rl.trainer.rl.config import LossConfig -from prime_rl.trainer.rl.loss import compute_entropy, compute_loss +from prime_rl.configs.trainer import CustomLossConfig, DefaultLossConfig, SFTLossConfig +from prime_rl.trainer.rl.loss import LossInputs, LossOutputs, compute_entropy, compute_loss, setup_loss_fn pytestmark = [pytest.mark.gpu] @@ -14,33 +14,34 @@ def test_grpo_loss(): advantages = [torch.randn(50).cuda(), torch.randn(30).cuda()] loss_mask = [torch.ones(50, dtype=torch.bool).cuda(), torch.ones(30, dtype=torch.bool).cuda()] + loss_fn = setup_loss_fn(DefaultLossConfig(dppo_mask_high=10.0)) loss, _ = compute_loss( trainer_logprobs, inference_logprobs, teacher_logprobs, advantages, loss_mask=loss_mask, - loss_config=LossConfig(ratio_type="token", token_mask_high=10.0), + loss_fn=loss_fn, loss_scale=1.0, ) assert loss.shape == () def test_gspo_loss(): - # Create list of tensors as expected by compute_loss (simulating split sequences) trainer_logprobs = [torch.randn(40, dtype=torch.float32).cuda(), torch.randn(60, dtype=torch.float32).cuda()] inference_logprobs = [torch.randn(40, dtype=torch.float32).cuda(), torch.randn(60, dtype=torch.float32).cuda()] teacher_logprobs = [torch.randn(40, dtype=torch.float32).cuda(), torch.randn(60, dtype=torch.float32).cuda()] advantages = [torch.randn(40).cuda(), torch.randn(60).cuda()] loss_mask = [torch.ones(40, dtype=torch.bool).cuda(), torch.ones(60, dtype=torch.bool).cuda()] + loss_fn = setup_loss_fn(DefaultLossConfig(dppo_mask_high=10.0)) loss, _ = compute_loss( trainer_logprobs, inference_logprobs, teacher_logprobs, advantages, loss_mask=loss_mask, - loss_config=LossConfig(ratio_type="sequence", token_mask_high=10.0), + loss_fn=loss_fn, loss_scale=1.0, ) assert loss.shape == () @@ -50,3 +51,79 @@ def test_entropy_loss(): shifted_logits = torch.randn(10, 10, 10, dtype=torch.float32).cuda() entropy = compute_entropy(shifted_logits) assert entropy.shape == (10, 10) + + +def test_setup_loss_fn_with_custom_config(): + """Test setup_loss_fn with CustomLossConfig importing a custom loss.""" + loss_config = CustomLossConfig( + import_path="tests.unit.train.rl.test_loss._dummy_custom_loss", + kwargs={"multiplier": 2.0}, + ) + loss_fn = setup_loss_fn(loss_config) + + inputs = LossInputs( + trainer_logprobs=torch.randn(50, dtype=torch.float32).cuda(), + inference_logprobs=torch.randn(50, dtype=torch.float32).cuda(), + teacher_logprobs=None, + advantages=torch.randn(50).cuda(), + loss_mask=torch.ones(50, dtype=torch.bool).cuda(), + ) + + result = loss_fn(inputs) + assert isinstance(result, LossOutputs) + assert result.loss.shape == () + assert "custom_metric" in result.metrics + + +def test_sft_loss_matches_masked_nll(): + trainer_logprobs = [torch.tensor([-0.1, -0.5, -0.2], dtype=torch.float32).cuda()] + inference_logprobs = [torch.zeros(3, dtype=torch.float32).cuda()] + advantages = [torch.zeros(3, dtype=torch.float32).cuda()] + loss_mask = [torch.tensor([True, False, True], dtype=torch.bool).cuda()] + + loss_fn = setup_loss_fn(SFTLossConfig()) + loss, metrics = compute_loss( + trainer_logprobs=trainer_logprobs, + inference_logprobs=inference_logprobs, + teacher_logprobs=None, + advantages=advantages, + loss_mask=loss_mask, + loss_fn=loss_fn, + loss_scale=2, + ) + + # loss = -sum(masked logprobs) / loss_scale = -(-0.1 - 0.2) / 2 = 0.15 + assert torch.isclose(loss, torch.tensor(0.15, device=loss.device), atol=1e-6) + assert "nll" in metrics + + +def test_sft_loss_override_uses_masked_nll_with_default_loss_config(): + trainer_logprobs = [torch.tensor([-0.1, -0.5, -0.2], dtype=torch.float32).cuda()] + inference_logprobs = [torch.zeros(3, dtype=torch.float32).cuda()] + advantages = [torch.ones(3, dtype=torch.float32).cuda()] + loss_mask = [torch.tensor([True, False, True], dtype=torch.bool).cuda()] + + loss_fn = setup_loss_fn(DefaultLossConfig()) + loss, metrics = compute_loss( + trainer_logprobs=trainer_logprobs, + inference_logprobs=inference_logprobs, + teacher_logprobs=None, + advantages=advantages, + loss_mask=loss_mask, + loss_fn=loss_fn, + loss_scale=2, + sft_loss=True, + ) + + assert torch.isclose(loss, torch.tensor(0.15, device=loss.device), atol=1e-6) + assert "nll" in metrics + assert "mismatch_kl" not in metrics + + +def _dummy_custom_loss(inputs: LossInputs, multiplier: float = 1.0) -> LossOutputs: + """A simple custom loss for testing.""" + loss = (inputs.trainer_logprobs[inputs.loss_mask].sum() * multiplier).abs() + return LossOutputs( + loss=loss, + metrics={"custom_metric": torch.tensor(multiplier)}, + ) diff --git a/tests/unit/train/rl/test_packer.py b/tests/unit/train/rl/test_packer.py new file mode 100644 index 00000000000..c661ec0df53 --- /dev/null +++ b/tests/unit/train/rl/test_packer.py @@ -0,0 +1,109 @@ +from pathlib import Path +from typing import Generator + +import pytest +import tomli_w +import torch +import torch.distributed as dist + +import prime_rl.trainer.runs as runs +from prime_rl.configs.shared import FileSystemTransportConfig +from prime_rl.trainer.rl.packer import MultiPacker +from prime_rl.trainer.runs import setup_multi_run_manager +from prime_rl.trainer.world import reset_world +from prime_rl.transport.types import TrainingSample + + +@pytest.fixture(autouse=True, scope="module") +def init_process_group() -> Generator[None, None, None]: + dist.init_process_group(backend="gloo", init_method="tcp://localhost:12356", rank=0, world_size=1) + yield + dist.destroy_process_group() + + +def create_run_with_config(output_dir: Path, run_name: str) -> Path: + run_dir = output_dir / run_name + run_dir.mkdir() + control_dir = run_dir / "control" + control_dir.mkdir() + config = { + "model": {"name": "test-model"}, + "batch_size": 2, + "rollouts_per_example": 1, + "env": [{"id": "test-env"}], + "sampling": {"temperature": 1.0}, + } + with open(control_dir / "orch.toml", "wb") as f: + tomli_w.dump(config, f) + return run_dir + + +def make_training_sample() -> TrainingSample: + return TrainingSample( + prompt_ids=[1], + prompt_mask=[False], + completion_ids=[2], + completion_mask=[True], + completion_logprobs=[-0.1], + completion_temperatures=[1.0], + ) + + +def test_packer_progress_updates_once_per_run(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + reset_world() + runs._MULTI_RUN_MANAGER = None + manager = setup_multi_run_manager(output_dir=tmp_path, max_runs=1, device=torch.device("cpu")) + + create_run_with_config(tmp_path, "run_test123") + manager.discover_runs() + run_idx = manager.id_2_idx["run_test123"] + + class DummyReceiver: + def receive(self): + return [] + + def reset_run(self, idx: int) -> None: + pass + + class DummySender: + def __init__(self): + self.sent = [] + + def send(self, micro_batch_grid): + self.sent.append(micro_batch_grid) + + sender_holder: dict[str, DummySender] = {} + + def fake_receiver(_config): + return DummyReceiver() + + def fake_sender(_output_dir, _data_world_size, _current_step, _config): + sender = DummySender() + sender_holder["sender"] = sender + return sender + + monkeypatch.setattr("prime_rl.trainer.rl.packer.setup_training_batch_receiver", fake_receiver) + monkeypatch.setattr("prime_rl.trainer.rl.packer.setup_micro_batch_sender", fake_sender) + + packer = MultiPacker( + dp_world_size=1, + seq_len=4, + pad_to_multiple_of=1, + tokenizer=None, + config=FileSystemTransportConfig(), + start_step=0, + ) + + packer.buffers[run_idx].append((make_training_sample(), 0)) + packer.buffers[run_idx].append((make_training_sample(), 0)) + + packer.pack() + + progress = manager.progress[run_idx] + assert progress.total_samples == 2 + assert progress.total_tokens == 4 + assert progress.step == 1 + + sender = sender_holder["sender"] + assert len(sender.sent) == 1 + assert len(sender.sent[0][0]) == 1 diff --git a/tests/unit/train/sft/test_dataloader.py b/tests/unit/train/sft/test_dataloader.py index 3445b4c0b3a..bf263ee4cf2 100644 --- a/tests/unit/train/sft/test_dataloader.py +++ b/tests/unit/train/sft/test_dataloader.py @@ -3,7 +3,7 @@ import pytest from transformers import AutoTokenizer -from prime_rl.trainer.sft.config import FakeDataConfig +from prime_rl.configs.sft import FakeDataConfig from prime_rl.trainer.sft.data import setup_dataloader, setup_dataset from prime_rl.trainer.world import reset_world diff --git a/tests/unit/train/sft/test_sft_dataset.py b/tests/unit/train/sft/test_sft_dataset.py index 0bd5c34a031..b8465e59d70 100644 --- a/tests/unit/train/sft/test_sft_dataset.py +++ b/tests/unit/train/sft/test_sft_dataset.py @@ -21,7 +21,7 @@ def test_init_sft_dataset(build_dummy_dataset): def test_raise_error_if_no_prompt_and_completion(build_dummy_dataset): - """Tests that an error is raised if no prompt and completion are provided but a tokenizer is provided.""" + """Tests that an error is raised if no supported SFT message fields are provided.""" dataset = build_dummy_dataset("a", 1) tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen3-0.6B") sft_dataset = SFTDataset(dataset, tokenizer=tokenizer) @@ -260,3 +260,55 @@ def test_multiturn_loss_mask_with_tools(): dataset = SFTDataset(dataset, tokenizer=tokenizer, max_examples=1) sample = next(iter(dataset)) print_sample(sample["input_ids"], sample["loss_mask"], tokenizer) + + +def test_messages_rows_are_equivalent_to_empty_prompt_completion(): + messages = [ + {"role": "system", "content": "You are a helpful assistant with access to tools."}, + {"role": "user", "content": "What's the weather in San Francisco?"}, + { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "call_1", + "type": "function", + "function": {"name": "get_weather", "arguments": '{"location": "San Francisco, CA"}'}, + } + ], + }, + {"role": "tool", "content": '{"temperature": 65, "condition": "Sunny"}', "tool_call_id": "call_1"}, + {"role": "assistant", "content": "It is 65F and sunny in San Francisco."}, + ] + + tokenizer = AutoTokenizer.from_pretrained("PrimeIntellect/Qwen3-0.6B") + messages_dataset = SFTDataset(Dataset.from_list([{"messages": messages}]), tokenizer=tokenizer, max_examples=1) + split_dataset = SFTDataset( + Dataset.from_list([{"prompt": [], "completion": messages}]), + tokenizer=tokenizer, + max_examples=1, + ) + + assert next(iter(messages_dataset)) == next(iter(split_dataset)) + + +def test_messages_take_precedence_over_prompt_and_completion(): + tokenizer = AutoTokenizer.from_pretrained("PrimeIntellect/Qwen3-0.6B") + row = { + "messages": [ + {"role": "system", "content": "System from messages"}, + {"role": "user", "content": "Prompt from messages"}, + {"role": "assistant", "content": "Completion from messages"}, + ], + "prompt": [{"role": "user", "content": "Ignored prompt"}], + "completion": [{"role": "assistant", "content": "Ignored completion"}], + } + + messages_dataset = SFTDataset(Dataset.from_list([row]), tokenizer=tokenizer, max_examples=1) + expected_dataset = SFTDataset( + Dataset.from_list([{"prompt": [], "completion": row["messages"]}]), + tokenizer=tokenizer, + max_examples=1, + ) + + assert next(iter(messages_dataset)) == next(iter(expected_dataset)) diff --git a/tests/unit/train/test_model.py b/tests/unit/train/test_model.py index fe65f3f2402..df66ab4fc9c 100644 --- a/tests/unit/train/test_model.py +++ b/tests/unit/train/test_model.py @@ -1,7 +1,7 @@ import pytest import torch -from prime_rl.trainer.config import AttnImplementation, ModelConfig +from prime_rl.configs.trainer import AttnImplementation, ModelConfig from prime_rl.trainer.model import get_model from prime_rl.trainer.models.layers.lm_head import inject_prime_lm_head diff --git a/tests/unit/train/test_model_forward.py b/tests/unit/train/test_model_forward.py new file mode 100644 index 00000000000..b62805bf3d4 --- /dev/null +++ b/tests/unit/train/test_model_forward.py @@ -0,0 +1,50 @@ +from types import SimpleNamespace + +import torch +import torch.nn as nn + +from prime_rl.trainer.model import forward + + +class _CaptureModel(nn.Module): + def __init__(self, config: SimpleNamespace): + super().__init__() + self.config = config + self.kwargs = None + + def forward(self, **kwargs): + self.kwargs = kwargs + input_ids = kwargs["input_ids"] + return {"logits": torch.zeros(*input_ids.shape, 4)} + + +def test_forward_adds_qwen3_vl_mm_token_type_ids(): + image_token_id = 10 + video_token_id = 20 + model = _CaptureModel( + SimpleNamespace(model_type="qwen3_vl", image_token_id=image_token_id, video_token_id=video_token_id) + ) + input_ids = torch.tensor([[1, image_token_id, image_token_id, 2, video_token_id]]) + position_ids = torch.arange(input_ids.shape[1]).unsqueeze(0) + pixel_values = torch.ones(2, 3) + image_grid_thw = torch.tensor([[1, 1, 2]]) + + forward(model, input_ids, position_ids, pixel_values=pixel_values, image_grid_thw=image_grid_thw) + + assert model.kwargs is not None + assert "position_ids" not in model.kwargs + torch.testing.assert_close(model.kwargs["pixel_values"], pixel_values) + torch.testing.assert_close(model.kwargs["image_grid_thw"], image_grid_thw) + torch.testing.assert_close(model.kwargs["mm_token_type_ids"], torch.tensor([[0, 1, 1, 0, 2]])) + + +def test_forward_skips_mm_token_type_ids_for_other_vlm_models(): + model = _CaptureModel(SimpleNamespace(model_type="qwen3_5_moe", image_token_id=10)) + input_ids = torch.tensor([[1, 10, 10, 2]]) + position_ids = torch.arange(input_ids.shape[1]).unsqueeze(0) + + forward(model, input_ids, position_ids, pixel_values=torch.ones(2, 3), image_grid_thw=torch.tensor([[1, 1, 2]])) + + assert model.kwargs is not None + assert "position_ids" not in model.kwargs + assert "mm_token_type_ids" not in model.kwargs diff --git a/tests/unit/train/test_perf.py b/tests/unit/train/test_perf.py index 3b670f36912..bb090832b0a 100644 --- a/tests/unit/train/test_perf.py +++ b/tests/unit/train/test_perf.py @@ -1,6 +1,6 @@ import pytest import torch -from transformers import AutoConfig, AutoModelForCausalLM +from transformers import AutoConfig, AutoModelForCausalLM, GenerationConfig from prime_rl.trainer.perf import PerfCounter @@ -13,8 +13,19 @@ def test_perf_counter(model_name: str, active_params: int, flops_per_token: int) # This speeds up the model loading as its a fake device with torch.device("meta"): config = AutoConfig.from_pretrained(model_name) + if getattr(config, "pad_token_id", None) is None: + gen_config = GenerationConfig.from_model_config(config) + # Use `is not None` instead of truthiness: token ID 0 is valid. + config.pad_token_id = next( + ( + v + for v in [gen_config.pad_token_id, gen_config.eos_token_id, getattr(config, "eos_token_id", None)] + if v is not None + ), + None, + ) model = AutoModelForCausalLM.from_config(config) - perf_counter = PerfCounter(model, seq_len=1024, window_size=10) + perf_counter = PerfCounter(model, seq_len=1024) assert perf_counter.get_active_mm_params(config) == active_params, ( f"Expected {active_params:,} active parameters, got {perf_counter.get_active_mm_params(config):,} active parameters" diff --git a/tests/unit/train/test_runs.py b/tests/unit/train/test_runs.py index 8a8e79e68c5..10b939ed182 100644 --- a/tests/unit/train/test_runs.py +++ b/tests/unit/train/test_runs.py @@ -5,7 +5,7 @@ import tomli_w import torch.distributed as dist -from prime_rl.trainer.runs import Runs +from prime_rl.trainer.runs import MultiRunManager @pytest.fixture(autouse=True, scope="module") @@ -33,7 +33,7 @@ def create_run_with_config( """ run_dir = output_dir / run_name run_dir.mkdir() - config_dir = run_dir / "configs" + config_dir = run_dir / "control" config_dir.mkdir() if config is None: @@ -51,121 +51,121 @@ def create_run_with_config( def test_initial_state(tmp_path: Path) -> None: - """Test that Runs initializes correctly.""" - runs = Runs(output_dir=tmp_path, max_runs=5) + """Test that MultiRunManager initializes correctly.""" + multi_run_manager = MultiRunManager(output_dir=tmp_path, max_runs=5) - assert runs.output_dir == tmp_path - assert runs.max_runs == 5 - assert len(runs.idx_2_id) == 0 - assert len(runs.id_2_idx) == 0 - assert len(runs.unused_idxs) == 5 - assert runs.run_dirs() == [] + assert multi_run_manager.output_dir == tmp_path + assert multi_run_manager.max_runs == 5 + assert len(multi_run_manager.idx_2_id) == 0 + assert len(multi_run_manager.id_2_idx) == 0 + assert len(multi_run_manager.unused_idxs) == 5 + assert multi_run_manager.run_dirs() == [] def test_detect_new_runs(tmp_path: Path) -> None: """Test that new runs are detected correctly.""" - runs = Runs(output_dir=tmp_path, max_runs=5) + multi_run_manager = MultiRunManager(output_dir=tmp_path, max_runs=5) # Create some run directories with valid configs for run_name in ["run_abc123", "run_def456"]: create_run_with_config(tmp_path, run_name) # Check for changes - runs.check_for_changes() + multi_run_manager.discover_runs() # Verify runs were detected - assert len(runs.id_2_idx) == 2 - assert len(runs.idx_2_id) == 2 - assert "run_abc123" in runs.id_2_idx - assert "run_def456" in runs.id_2_idx + assert len(multi_run_manager.id_2_idx) == 2 + assert len(multi_run_manager.idx_2_id) == 2 + assert "run_abc123" in multi_run_manager.id_2_idx + assert "run_def456" in multi_run_manager.id_2_idx # Verify indices are assigned from available pool - assert len(runs.unused_idxs) == 3 # 5 - 2 = 3 - assert runs.id_2_idx["run_abc123"] in range(5) - assert runs.id_2_idx["run_def456"] in range(5) + assert len(multi_run_manager.unused_idxs) == 3 # 5 - 2 = 3 + assert multi_run_manager.id_2_idx["run_abc123"] in range(5) + assert multi_run_manager.id_2_idx["run_def456"] in range(5) # Verify bidirectional mapping - idx1 = runs.id_2_idx["run_abc123"] - idx2 = runs.id_2_idx["run_def456"] - assert runs.idx_2_id[idx1] == "run_abc123" - assert runs.idx_2_id[idx2] == "run_def456" + idx1 = multi_run_manager.id_2_idx["run_abc123"] + idx2 = multi_run_manager.id_2_idx["run_def456"] + assert multi_run_manager.idx_2_id[idx1] == "run_abc123" + assert multi_run_manager.idx_2_id[idx2] == "run_def456" def test_detect_deleted_runs(tmp_path: Path) -> None: """Test that deleted runs are detected correctly.""" - runs = Runs(output_dir=tmp_path, max_runs=5) + multi_run_manager = MultiRunManager(output_dir=tmp_path, max_runs=5) # Create run directories with valid configs for run_name in ["run_abc123", "run_def456"]: create_run_with_config(tmp_path, run_name) # Detect initial runs - runs.check_for_changes() - initial_idx1 = runs.id_2_idx["run_abc123"] - initial_idx2 = runs.id_2_idx["run_def456"] + multi_run_manager.discover_runs() + initial_idx1 = multi_run_manager.id_2_idx["run_abc123"] + initial_idx2 = multi_run_manager.id_2_idx["run_def456"] - assert len(runs.id_2_idx) == 2 - assert len(runs.unused_idxs) == 3 + assert len(multi_run_manager.id_2_idx) == 2 + assert len(multi_run_manager.unused_idxs) == 3 # Delete one run run1 = tmp_path / "run_abc123" import shutil shutil.rmtree(run1) - runs.check_for_changes() + multi_run_manager.discover_runs() # Verify run was removed - assert len(runs.id_2_idx) == 1 - assert len(runs.idx_2_id) == 1 - assert "run_abc123" not in runs.id_2_idx - assert "run_def456" in runs.id_2_idx - assert initial_idx1 not in runs.idx_2_id + assert len(multi_run_manager.id_2_idx) == 1 + assert len(multi_run_manager.idx_2_id) == 1 + assert "run_abc123" not in multi_run_manager.id_2_idx + assert "run_def456" in multi_run_manager.id_2_idx + assert initial_idx1 not in multi_run_manager.idx_2_id # Verify index was returned to unused pool - assert len(runs.unused_idxs) == 4 - assert initial_idx1 in runs.unused_idxs - assert initial_idx2 not in runs.unused_idxs + assert len(multi_run_manager.unused_idxs) == 4 + assert initial_idx1 in multi_run_manager.unused_idxs + assert initial_idx2 not in multi_run_manager.unused_idxs def test_max_runs_limit(tmp_path: Path) -> None: """Test that only max_runs are tracked.""" - runs = Runs(output_dir=tmp_path, max_runs=2) + multi_run_manager = MultiRunManager(output_dir=tmp_path, max_runs=2) # Create more runs than max_runs with valid configs for run_name in ["run_001", "run_002", "run_003"]: create_run_with_config(tmp_path, run_name) - runs.check_for_changes() + multi_run_manager.discover_runs() # Only max_runs should be tracked - assert len(runs.id_2_idx) == 2 - assert len(runs.idx_2_id) == 2 - assert len(runs.unused_idxs) == 0 + assert len(multi_run_manager.id_2_idx) == 2 + assert len(multi_run_manager.idx_2_id) == 2 + assert len(multi_run_manager.unused_idxs) == 0 - to_delete_run = runs.get_run_dir(0) + to_delete_run = multi_run_manager.get_run_dir(0) import shutil shutil.rmtree(to_delete_run) - runs.check_for_changes() + multi_run_manager.discover_runs() - assert len(runs.id_2_idx) == 2 - assert len(runs.idx_2_id) == 2 - assert len(runs.unused_idxs) == 0 - assert to_delete_run not in runs.run_dirs() + assert len(multi_run_manager.id_2_idx) == 2 + assert len(multi_run_manager.idx_2_id) == 2 + assert len(multi_run_manager.unused_idxs) == 0 + assert to_delete_run not in multi_run_manager.run_dirs() def test_run_dirs(tmp_path: Path) -> None: """Test that run_dirs returns correct paths.""" - runs = Runs(output_dir=tmp_path, max_runs=5) + multi_run_manager = MultiRunManager(output_dir=tmp_path, max_runs=5) # Create run directories with valid configs for run_name in ["run_abc", "run_def"]: create_run_with_config(tmp_path, run_name) - runs.check_for_changes() + multi_run_manager.discover_runs() - run_dirs = runs.run_dirs() + run_dirs = multi_run_manager.run_dirs() assert len(run_dirs) == 2 assert tmp_path / "run_abc" in run_dirs assert tmp_path / "run_def" in run_dirs @@ -173,7 +173,7 @@ def test_run_dirs(tmp_path: Path) -> None: def test_non_run_directories_ignored(tmp_path: Path) -> None: """Test that non-run directories are ignored.""" - runs = Runs(output_dir=tmp_path, max_runs=5) + multi_run_manager = MultiRunManager(output_dir=tmp_path, max_runs=5) # Create mix of run and non-run directories create_run_with_config(tmp_path, "run_abc") @@ -181,18 +181,18 @@ def test_non_run_directories_ignored(tmp_path: Path) -> None: (tmp_path / "other_dir").mkdir() (tmp_path / "random").mkdir() - runs.check_for_changes() + multi_run_manager.discover_runs() # Only run_* directories should be tracked - assert len(runs.id_2_idx) == 1 - assert "run_abc" in runs.id_2_idx - assert "other_dir" not in runs.id_2_idx - assert "random" not in runs.id_2_idx + assert len(multi_run_manager.id_2_idx) == 1 + assert "run_abc" in multi_run_manager.id_2_idx + assert "other_dir" not in multi_run_manager.id_2_idx + assert "random" not in multi_run_manager.id_2_idx def test_config_loading(tmp_path: Path) -> None: """Test that orchestrator configs are loaded correctly.""" - runs = Runs(output_dir=tmp_path, max_runs=5) + multi_run_manager = MultiRunManager(output_dir=tmp_path, max_runs=5) # Create a run directory with config test_config = { @@ -205,45 +205,39 @@ def test_config_loading(tmp_path: Path) -> None: create_run_with_config(tmp_path, "run_test123", config=test_config) # Detect the run - runs.check_for_changes() + multi_run_manager.discover_runs() # Verify config was loaded and parsed as OrchestratorConfig - assert len(runs.config) == 1 - run_idx = runs.id_2_idx["run_test123"] - assert run_idx in runs.config + assert len(multi_run_manager.config) == 1 + run_idx = multi_run_manager.id_2_idx["run_test123"] + assert run_idx in multi_run_manager.config # Access config as OrchestratorConfig object - config = runs.config[run_idx] + config = multi_run_manager.config[run_idx] assert config.model.name == "test-model" assert config.batch_size == 32 assert config.max_steps == 1000 def test_config_missing(tmp_path: Path) -> None: - """Test that runs without configs are skipped and error.txt is created.""" - runs = Runs(output_dir=tmp_path, max_runs=5) + """Test that runs without configs are skipped and config_validation_error.txt is created.""" + multi_run_manager = MultiRunManager(output_dir=tmp_path, max_runs=5) # Create a run directory without config run_dir = tmp_path / "run_noconfig" run_dir.mkdir() # Detect the run - runs.check_for_changes() + multi_run_manager.discover_runs() # Verify run was not added - assert len(runs.config) == 0 - assert "run_noconfig" not in runs.id_2_idx - - # Verify error.txt was created - error_path = run_dir / "configs" / "error.txt" - assert error_path.exists() - error_content = error_path.read_text() - assert "No orchestrator config found" in error_content + assert len(multi_run_manager.config) == 0 + assert "run_noconfig" not in multi_run_manager.id_2_idx def test_config_cleanup_on_deletion(tmp_path: Path) -> None: """Test that configs are cleaned up when runs are deleted.""" - runs = Runs(output_dir=tmp_path, max_runs=5) + multi_run_manager = MultiRunManager(output_dir=tmp_path, max_runs=5) # Create a run directory with valid config test_config = { @@ -255,24 +249,24 @@ def test_config_cleanup_on_deletion(tmp_path: Path) -> None: run_dir = create_run_with_config(tmp_path, "run_delete_me", config=test_config) # Detect the run - runs.check_for_changes() - run_idx = runs.id_2_idx["run_delete_me"] - assert run_idx in runs.config + multi_run_manager.discover_runs() + run_idx = multi_run_manager.id_2_idx["run_delete_me"] + assert run_idx in multi_run_manager.config # Delete the run directory import shutil shutil.rmtree(run_dir) - runs.check_for_changes() + multi_run_manager.discover_runs() # Verify config was cleaned up - assert run_idx not in runs.config - assert "run_delete_me" not in runs.id_2_idx + assert run_idx not in multi_run_manager.config + assert "run_delete_me" not in multi_run_manager.id_2_idx def test_config_invalid(tmp_path: Path) -> None: - """Test that runs with invalid configs are skipped and error.txt is created.""" - runs = Runs(output_dir=tmp_path, max_runs=5) + """Test that runs with invalid configs are skipped and config_validation_error.txt is created.""" + multi_run_manager = MultiRunManager(output_dir=tmp_path, max_runs=5) # Create a run directory with invalid config (invalid type for a field) # Invalid config - batch_size should be int, not string @@ -283,17 +277,59 @@ def test_config_invalid(tmp_path: Path) -> None: "env": [{"id": "test-env"}], } run_dir = create_run_with_config(tmp_path, "run_invalid", config=invalid_config) - config_dir = run_dir / "configs" + config_dir = run_dir / "control" # Detect the run - runs.check_for_changes() + multi_run_manager.discover_runs() # Verify run was not added - assert len(runs.config) == 0 - assert "run_invalid" not in runs.id_2_idx + assert len(multi_run_manager.config) == 0 + assert "run_invalid" not in multi_run_manager.id_2_idx - # Verify error.txt was created with error details - error_path = config_dir / "error.txt" + # Verify config_validation_error.txt was created with error details + error_path = config_dir / "config_validation_error.txt" assert error_path.exists() error_content = error_path.read_text() assert "Error parsing orchestrator config" in error_content + + +def test_evict_run_writes_file(tmp_path: Path) -> None: + """Test that evict_run writes the eviction reason to the correct file.""" + multi_run_manager = MultiRunManager(output_dir=tmp_path, max_runs=5) + + # Create a run directory with valid config + create_run_with_config(tmp_path, "run_to_evict") + multi_run_manager.discover_runs() + + # Get the run index + run_idx = multi_run_manager.id_2_idx["run_to_evict"] + + # Evict the run + eviction_reason = "Test eviction reason" + multi_run_manager.evict_run(run_idx, eviction_reason) + + # Verify evicted.txt was created with the reason + evicted_path = tmp_path / "run_to_evict" / "control" / "evicted.txt" + assert evicted_path.exists() + assert evicted_path.read_text() == eviction_reason + + +def test_discover_runs_ignores_evicted(tmp_path: Path) -> None: + """Test that discover_runs ignores runs with evicted.txt.""" + multi_run_manager = MultiRunManager(output_dir=tmp_path, max_runs=5) + + # Create two run directories + create_run_with_config(tmp_path, "run_normal") + create_run_with_config(tmp_path, "run_evicted") + + # Mark one as evicted by creating evicted.txt + evicted_path = tmp_path / "run_evicted" / "control" / "evicted.txt" + evicted_path.write_text("Previously evicted") + + # Discover runs + multi_run_manager.discover_runs() + + # Only the normal run should be discovered + assert len(multi_run_manager.id_2_idx) == 1 + assert "run_normal" in multi_run_manager.id_2_idx + assert "run_evicted" not in multi_run_manager.id_2_idx diff --git a/tests/unit/utils/test_client.py b/tests/unit/utils/test_client.py index 9a71a4c3db8..6b48790ef36 100644 --- a/tests/unit/utils/test_client.py +++ b/tests/unit/utils/test_client.py @@ -3,9 +3,10 @@ from unittest.mock import AsyncMock, MagicMock import httpx -import pytest +import verifiers as vf -from prime_rl.utils.client import _is_retryable_lora_error, load_lora_adapter +from prime_rl.configs.shared import ClientConfig +from prime_rl.utils.client import _is_retryable_lora_error, load_lora_adapter, setup_clients def test_is_retryable_lora_error_returns_true_for_404(): @@ -42,44 +43,74 @@ def test_load_lora_adapter_succeeds_on_first_attempt(): asyncio.run(load_lora_adapter([mock_client], "test-lora", Path("/test/path"))) mock_client.post.assert_called_once_with( - "/v1/load_lora_adapter", + "/load_lora_adapter", json={"lora_name": "test-lora", "lora_path": "/test/path"}, + timeout=httpx.Timeout(connect=10.0, read=30.0, write=60.0, pool=10.0), ) -def test_load_lora_adapter_retries_on_404_then_succeeds(): - mock_client = AsyncMock() - - error_response = MagicMock() - error_response.status_code = 404 - success_response = MagicMock() - success_response.raise_for_status = MagicMock() - - call_count = 0 +def test_setup_clients_assigns_renderer_and_dp_rank_headers(): + client_config = ClientConfig( + base_url=["http://worker-a:8000/v1"], + api_key_var="PRIME_API_KEY", + headers={"X-Test": "test"}, + dp_rank_count=2, + extra_headers_from_state={"X-Session-ID": "session_id"}, + ) - async def mock_post(*args, **kwargs): - nonlocal call_count - call_count += 1 - if call_count == 1: - raise httpx.HTTPStatusError("Not found", request=MagicMock(), response=error_response) - return success_response + clients = setup_clients( + client_config, + client_type="renderer", + renderer_name="qwen3_vl", + ) - mock_client.post = mock_post + assert [client.client_type for client in clients] == ["renderer", "renderer"] + assert [client.renderer for client in clients] == ["qwen3_vl", "qwen3_vl"] + assert [client.renderer_model_name for client in clients] == [None, None] + assert [client.api_base_url for client in clients] == ["http://worker-a:8000/v1"] * 2 + assert [client.extra_headers["X-data-parallel-rank"] for client in clients] == ["0", "1"] + assert clients[0].extra_headers["X-Test"] == "test" + assert clients[0].extra_headers_from_state == {"X-Session-ID": "session_id"} - asyncio.run(load_lora_adapter([mock_client], "test-lora", Path("/test/path"))) - assert call_count == 2 +def test_setup_clients_assigns_renderer_model_name(): + client_config = ClientConfig( + base_url=["http://worker-a:8000/v1"], + api_key_var="PRIME_API_KEY", + ) + clients = setup_clients( + client_config, + client_type="renderer", + renderer_name="qwen3_vl", + renderer_model_name="Qwen/Qwen3-VL-4B-Instruct", + ) -def test_load_lora_adapter_raises_non_retryable_error_immediately(): - mock_client = AsyncMock() + assert clients[0].renderer_model_name == "Qwen/Qwen3-VL-4B-Instruct" - error_response = MagicMock() - error_response.status_code = 400 - mock_client.post.side_effect = httpx.HTTPStatusError("Bad request", request=MagicMock(), response=error_response) - with pytest.raises(httpx.HTTPStatusError) as exc_info: - asyncio.run(load_lora_adapter([mock_client], "test-lora", Path("/test/path"))) +def test_setup_clients_preserves_chat_client_defaults(): + client_config = ClientConfig( + base_url=["http://worker-a:8000/v1"], + api_key_var="PRIME_API_KEY", + ) - assert exc_info.value.response.status_code == 400 - assert mock_client.post.call_count == 1 + clients = setup_clients(client_config) + + assert clients == [ + vf.ClientConfig( + client_idx=0, + client_type="openai_chat_completions", + renderer="auto", + renderer_model_name=None, + api_key_var="PRIME_API_KEY", + api_base_url="http://worker-a:8000/v1", + timeout=client_config.timeout, + connect_timeout=client_config.connect_timeout, + max_connections=8192, + max_keepalive_connections=8192, + max_retries=10, + extra_headers={}, + extra_headers_from_state={}, + ) + ] diff --git a/tests/unit/utils/test_elastic.py b/tests/unit/utils/test_elastic.py new file mode 100644 index 00000000000..a2cd40d1d17 --- /dev/null +++ b/tests/unit/utils/test_elastic.py @@ -0,0 +1,450 @@ +import asyncio +import socket +from pathlib import Path +from unittest.mock import AsyncMock, MagicMock, patch + +import httpx +import verifiers as vf + +from prime_rl.utils.elastic import ( + AdapterState, + ElasticInferencePool, + check_server_model, + discover_ready_servers, + discover_server_ips, +) + +# discover_server_ips tests + + +def test_discover_server_ips_returns_sorted_ips(): + with patch("socket.gethostbyname_ex") as mock_dns: + mock_dns.return_value = ("hostname", [], ["10.0.0.3", "10.0.0.1", "10.0.0.2"]) + result = discover_server_ips("test.hostname") + assert result == ["10.0.0.1", "10.0.0.2", "10.0.0.3"] + + +def test_discover_server_ips_returns_empty_list_on_dns_failure(): + with patch("socket.gethostbyname_ex") as mock_dns: + mock_dns.side_effect = socket.gaierror("DNS lookup failed") + result = discover_server_ips("nonexistent.hostname") + assert result == [] + + +def test_discover_server_ips_returns_single_ip(): + with patch("socket.gethostbyname_ex") as mock_dns: + mock_dns.return_value = ("hostname", [], ["10.0.0.1"]) + result = discover_server_ips("single.hostname") + assert result == ["10.0.0.1"] + + +# check_server_model tests + + +def test_check_server_model_returns_true_when_model_found(): + with patch("httpx.AsyncClient") as mock_client_class: + mock_client = AsyncMock() + mock_client_class.return_value.__aenter__.return_value = mock_client + + mock_response = MagicMock() + mock_response.raise_for_status = MagicMock() + mock_response.json.return_value = {"data": [{"id": "model-a"}, {"id": "model-b"}]} + mock_client.get.return_value = mock_response + + has_model, is_healthy = asyncio.run(check_server_model("http://10.0.0.1:8000", "model-a")) + + assert has_model is True + assert is_healthy is True + + +def test_check_server_model_returns_false_when_model_not_found(): + with patch("httpx.AsyncClient") as mock_client_class: + mock_client = AsyncMock() + mock_client_class.return_value.__aenter__.return_value = mock_client + + mock_response = MagicMock() + mock_response.raise_for_status = MagicMock() + mock_response.json.return_value = {"data": [{"id": "other-model"}]} + mock_client.get.return_value = mock_response + + has_model, is_healthy = asyncio.run(check_server_model("http://10.0.0.1:8000", "model-a")) + + assert has_model is False + assert is_healthy is True + + +def test_check_server_model_returns_false_on_connection_error(): + with patch("httpx.AsyncClient") as mock_client_class: + mock_client = AsyncMock() + mock_client_class.return_value.__aenter__.return_value = mock_client + mock_client.get.side_effect = httpx.ConnectError("Connection refused") + + has_model, is_healthy = asyncio.run(check_server_model("http://10.0.0.1:8000", "model-a")) + + assert has_model is False + assert is_healthy is False + + +def test_check_server_model_returns_false_on_http_error(): + with patch("httpx.AsyncClient") as mock_client_class: + mock_client = AsyncMock() + mock_client_class.return_value.__aenter__.return_value = mock_client + + mock_response = MagicMock() + mock_response.status_code = 500 + mock_response.raise_for_status.side_effect = httpx.HTTPStatusError( + "Server error", request=MagicMock(), response=mock_response + ) + mock_client.get.return_value = mock_response + + has_model, is_healthy = asyncio.run(check_server_model("http://10.0.0.1:8000", "model-a")) + + assert has_model is False + assert is_healthy is False + + +# discover_ready_servers tests + + +def test_discover_ready_servers_returns_servers_with_model_when_any_have_it(): + with ( + patch("prime_rl.utils.elastic.discover_server_ips") as mock_discover, + patch("prime_rl.utils.elastic.check_server_model") as mock_check, + ): + mock_discover.return_value = ["10.0.0.1", "10.0.0.2", "10.0.0.3"] + + async def mock_check_impl(url, model_name): + if "10.0.0.1" in url: + return True, True # has model, healthy + elif "10.0.0.2" in url: + return False, True # no model, healthy + else: + return False, True # no model, healthy + + mock_check.side_effect = mock_check_impl + + result = asyncio.run(discover_ready_servers("test.hostname", 8000, "my-lora")) + + assert result == ["http://10.0.0.1:8000/v1"] + + +def test_discover_ready_servers_returns_empty_when_none_have_model(): + with ( + patch("prime_rl.utils.elastic.discover_server_ips") as mock_discover, + patch("prime_rl.utils.elastic.check_server_model") as mock_check, + ): + mock_discover.return_value = ["10.0.0.1", "10.0.0.2"] + mock_check.return_value = (False, True) # no model, but healthy + + result = asyncio.run(discover_ready_servers("test.hostname", 8000, "my-lora")) + + assert result == [] # No servers have the model, so return empty + + +def test_discover_ready_servers_returns_empty_when_no_dns_records(): + with patch("prime_rl.utils.elastic.discover_server_ips") as mock_discover: + mock_discover.return_value = [] + + result = asyncio.run(discover_ready_servers("test.hostname", 8000, "my-lora")) + + assert result == [] + + +def test_discover_ready_servers_only_returns_servers_with_model(): + with ( + patch("prime_rl.utils.elastic.discover_server_ips") as mock_discover, + patch("prime_rl.utils.elastic.check_server_model") as mock_check, + ): + mock_discover.return_value = ["10.0.0.1", "10.0.0.2"] + + async def mock_check_impl(url, model_name): + if "10.0.0.1" in url: + return True, True # has model, healthy + else: + return False, True # no model, healthy + + mock_check.side_effect = mock_check_impl + + result = asyncio.run(discover_ready_servers("test.hostname", 8000, "my-lora")) + + # Only returns servers that have the model + assert result == ["http://10.0.0.1:8000/v1"] + + +# AdapterState tests + + +def test_adapter_state_creation(): + adapter = AdapterState(name="my-lora", path=Path("/weights/step_100"), step=100) + assert adapter.name == "my-lora" + assert adapter.path == Path("/weights/step_100") + assert adapter.step == 100 + + +# ElasticInferencePool adapter matching tests + + +def test_adapter_matches_when_no_adapter_desired(): + with patch("prime_rl.utils.elastic.get_logger"): + mock_config = MagicMock() + mock_config.elastic.hostname = "test.hostname" + mock_config.elastic.port = 8000 + mock_config.elastic.sync_interval = 5.0 + mock_config.router_url = None + pool = ElasticInferencePool(client_config=mock_config, model_name="base-model") + # No adapter desired (base model inference) + assert pool._adapter_matches_desired(None) is True + assert pool._adapter_matches_desired(AdapterState("x", Path("/x"), 0)) is True + + +def test_adapter_matches_by_path(): + with patch("prime_rl.utils.elastic.get_logger"): + mock_config = MagicMock() + mock_config.elastic.hostname = "test.hostname" + mock_config.elastic.port = 8000 + mock_config.elastic.sync_interval = 5.0 + mock_config.router_url = None + pool = ElasticInferencePool(client_config=mock_config, model_name="base-model") + pool._desired.path = Path("/weights/step_100") + pool._desired.step = 100 + + loaded = AdapterState(name="lora", path=Path("/weights/step_100"), step=100) + assert pool._adapter_matches_desired(loaded) is True + + loaded_wrong_path = AdapterState(name="lora", path=Path("/weights/step_50"), step=50) + assert pool._adapter_matches_desired(loaded_wrong_path) is False + + +def test_adapter_matches_by_step_when_nonzero(): + with patch("prime_rl.utils.elastic.get_logger"): + mock_config = MagicMock() + mock_config.elastic.hostname = "test.hostname" + mock_config.elastic.port = 8000 + mock_config.elastic.sync_interval = 5.0 + mock_config.router_url = None + pool = ElasticInferencePool(client_config=mock_config, model_name="base-model") + pool._desired.path = Path("/weights/step_100") + pool._desired.step = 100 + + # Different path but same step + loaded = AdapterState(name="lora", path=Path("/other/path"), step=100) + assert pool._adapter_matches_desired(loaded) is True + + +def test_adapter_does_not_match_by_zero_step(): + with patch("prime_rl.utils.elastic.get_logger"): + mock_config = MagicMock() + mock_config.elastic.hostname = "test.hostname" + mock_config.elastic.port = 8000 + mock_config.elastic.sync_interval = 5.0 + mock_config.router_url = None + pool = ElasticInferencePool(client_config=mock_config, model_name="base-model") + pool._desired.path = Path("/weights/step_0") + pool._desired.step = 0 + + # Step 0 should not match by step alone (avoid false positives) + loaded = AdapterState(name="lora", path=Path("/other/path"), step=0) + assert pool._adapter_matches_desired(loaded) is False + + +def test_adapter_returns_false_when_no_adapter_loaded(): + with patch("prime_rl.utils.elastic.get_logger"): + mock_config = MagicMock() + mock_config.elastic.hostname = "test.hostname" + mock_config.elastic.port = 8000 + mock_config.elastic.sync_interval = 5.0 + mock_config.router_url = None + pool = ElasticInferencePool(client_config=mock_config, model_name="base-model") + pool._desired.path = Path("/weights/step_100") + pool._desired.step = 100 + + assert pool._adapter_matches_desired(None) is False + + +# _get_loaded_adapter tests + + +def test_get_loaded_adapter_finds_correct_adapter_when_multiple_loaded(): + """Test that _get_loaded_adapter returns the adapter matching desired name, not the first one.""" + with patch("prime_rl.utils.elastic.get_logger"): + mock_config = MagicMock() + mock_config.elastic.hostname = "test.hostname" + mock_config.elastic.port = 8000 + mock_config.elastic.sync_interval = 5.0 + mock_config.router_url = None + pool = ElasticInferencePool(client_config=mock_config, model_name="base-model") + + # Set the desired adapter name + pool._desired.name = "rft-target-run" + pool._desired.path = Path("/data/outputs/target_run/broadcasts/step_10") + pool._desired.step = 10 + + # Mock admin client with multiple adapters (like the real scenario) + mock_admin = AsyncMock() + mock_response = MagicMock() + mock_response.raise_for_status = MagicMock() + # Simulates vLLM response with multiple adapters from different runs + mock_response.json.return_value = { + "data": [ + { + "id": "base-model", + "parent": None, + "root": "base-model", + }, + { + "id": "rft-other-run", # Different run's adapter (comes first!) + "parent": "base-model", + "root": "/data/outputs/other_run/broadcasts/step_48", + }, + { + "id": "rft-target-run", # Our desired adapter + "parent": "base-model", + "root": "/data/outputs/target_run/broadcasts/step_10", + }, + ] + } + mock_admin.get.return_value = mock_response + pool._admin_clients["10.0.0.1"] = mock_admin + + result = asyncio.run(pool._get_loaded_adapter("10.0.0.1")) + + # Should return our target adapter, not the first one found + assert result is not None + assert result.name == "rft-target-run" + assert result.path == Path("/data/outputs/target_run/broadcasts/step_10") + assert result.step == 10 + + +def test_get_loaded_adapter_returns_none_when_desired_adapter_not_found(): + """Test that _get_loaded_adapter returns None when desired adapter is not in the list.""" + with patch("prime_rl.utils.elastic.get_logger"): + mock_config = MagicMock() + mock_config.elastic.hostname = "test.hostname" + mock_config.elastic.port = 8000 + mock_config.elastic.sync_interval = 5.0 + mock_config.router_url = None + pool = ElasticInferencePool(client_config=mock_config, model_name="base-model") + + pool._desired.name = "rft-missing-run" + pool._desired.path = Path("/data/outputs/missing_run/broadcasts/step_5") + pool._desired.step = 5 + + mock_admin = AsyncMock() + mock_response = MagicMock() + mock_response.raise_for_status = MagicMock() + mock_response.json.return_value = { + "data": [ + {"id": "base-model", "parent": None, "root": "base-model"}, + {"id": "rft-other-run", "parent": "base-model", "root": "/data/outputs/other_run/step_10"}, + ] + } + mock_admin.get.return_value = mock_response + pool._admin_clients["10.0.0.1"] = mock_admin + + result = asyncio.run(pool._get_loaded_adapter("10.0.0.1")) + + assert result is None + + +def test_get_loaded_adapter_parses_step_from_path(): + """Test that _get_loaded_adapter correctly parses step number from path.""" + with patch("prime_rl.utils.elastic.get_logger"): + mock_config = MagicMock() + mock_config.elastic.hostname = "test.hostname" + mock_config.elastic.port = 8000 + mock_config.elastic.sync_interval = 5.0 + mock_config.router_url = None + pool = ElasticInferencePool(client_config=mock_config, model_name="base-model") + + pool._desired.name = "my-lora" + + mock_admin = AsyncMock() + mock_response = MagicMock() + mock_response.raise_for_status = MagicMock() + mock_response.json.return_value = { + "data": [ + {"id": "my-lora", "parent": "base", "root": "/weights/step_42"}, + ] + } + mock_admin.get.return_value = mock_response + pool._admin_clients["10.0.0.1"] = mock_admin + + result = asyncio.run(pool._get_loaded_adapter("10.0.0.1")) + + assert result is not None + assert result.step == 42 + + +def test_get_loaded_adapter_handles_step_dash_format(): + """Test that _get_loaded_adapter parses step-N format (with dash).""" + with patch("prime_rl.utils.elastic.get_logger"): + mock_config = MagicMock() + mock_config.elastic.hostname = "test.hostname" + mock_config.elastic.port = 8000 + mock_config.elastic.sync_interval = 5.0 + mock_config.router_url = None + pool = ElasticInferencePool(client_config=mock_config, model_name="base-model") + + pool._desired.name = "my-lora" + + mock_admin = AsyncMock() + mock_response = MagicMock() + mock_response.raise_for_status = MagicMock() + mock_response.json.return_value = { + "data": [ + {"id": "my-lora", "parent": "base", "root": "/weights/step-99"}, + ] + } + mock_admin.get.return_value = mock_response + pool._admin_clients["10.0.0.1"] = mock_admin + + result = asyncio.run(pool._get_loaded_adapter("10.0.0.1")) + + assert result is not None + assert result.step == 99 + + +def test_elastic_clients_preserve_renderer_model_name_when_model_name_updates(): + with patch("prime_rl.utils.elastic.get_logger"): + client_config = MagicMock() + client_config.elastic.hostname = "test.hostname" + client_config.elastic.port = 8000 + client_config.elastic.sync_interval = 5.0 + client_config.router_url = None + client_config.timeout = 1200 + client_config.connect_timeout = 30.0 + client_config.api_key_var = "PRIME_API_KEY" + client_config.headers = {} + client_config.extra_headers_from_state = {} + client_config.dp_rank_count = 1 + + pool = ElasticInferencePool( + client_config=client_config, + model_name="Qwen/Qwen3-VL-4B-Instruct", + train_client_type="renderer", + renderer_name="qwen3_vl", + ) + pool._servers = { + "10.0.0.1": MagicMock(status="ready"), + } + + pool.update_model_name("r8-smoke") + clients = pool.train_clients + + assert clients == [ + vf.ClientConfig( + client_idx=0, + client_type="renderer", + renderer="qwen3_vl", + renderer_model_name="Qwen/Qwen3-VL-4B-Instruct", + api_key_var="PRIME_API_KEY", + api_base_url="http://10.0.0.1:8000/v1", + timeout=1200, + connect_timeout=30.0, + max_connections=8192, + max_keepalive_connections=8192, + max_retries=10, + extra_headers={}, + extra_headers_from_state={}, + ) + ] diff --git a/tests/unit/utils/test_metrics_server.py b/tests/unit/utils/test_metrics_server.py index 63cd95dbb27..11d81b87158 100644 --- a/tests/unit/utils/test_metrics_server.py +++ b/tests/unit/utils/test_metrics_server.py @@ -7,8 +7,8 @@ import pytest -from prime_rl.utils.config import MetricsServerConfig -from prime_rl.utils.metrics_server import PROMETHEUS_AVAILABLE, HealthServer, MetricsServer, RunStats +from prime_rl.configs.shared import MetricsServerConfig +from prime_rl.utils.metrics_server import HealthServer, MetricsServer, RunStats def find_free_port() -> int: @@ -88,22 +88,23 @@ def test_server_update_metrics(): response = urllib.request.urlopen(f"http://localhost:{port}/metrics", timeout=2) content = response.read().decode() - if PROMETHEUS_AVAILABLE: - assert "trainer_step" in content - assert "trainer_loss" in content - assert "trainer_last_step_timestamp_seconds" in content - assert "trainer_mismatch_kl" in content + assert "trainer_step" in content + assert "trainer_loss" in content + assert "trainer_last_step_timestamp_seconds" in content + assert "trainer_mismatch_kl" in content + assert "trainer_kl_ent_ratio" in content server.stop() -@pytest.mark.skipif(not PROMETHEUS_AVAILABLE, reason="prometheus_client not installed") def test_server_metrics_values_are_correct(): port = find_free_port() server = MetricsServer(MetricsServerConfig(port=port)) server.start() time.sleep(0.1) + entropy = 1.5 + mismatch_kl = 0.045 server.update( step=100, loss=0.123, @@ -111,7 +112,8 @@ def test_server_metrics_values_are_correct(): grad_norm=2.5, peak_memory_gib=16.0, learning_rate=3e-5, - mismatch_kl=0.045, + entropy=entropy, + mismatch_kl=mismatch_kl, ) response = urllib.request.urlopen(f"http://localhost:{port}/metrics", timeout=2) @@ -120,6 +122,7 @@ def test_server_metrics_values_are_correct(): assert "trainer_step 100.0" in content assert "trainer_loss 0.123" in content assert "trainer_mismatch_kl 0.045" in content + assert f"trainer_kl_ent_ratio {mismatch_kl / entropy}" in content server.stop() @@ -138,11 +141,10 @@ def test_server_isolated_registry(): server1.update(step=999, loss=0.1, throughput=100, grad_norm=1, peak_memory_gib=1, learning_rate=1e-3) - if PROMETHEUS_AVAILABLE: - resp1 = urllib.request.urlopen(f"http://localhost:{port1}/metrics", timeout=2).read().decode() - resp2 = urllib.request.urlopen(f"http://localhost:{port2}/metrics", timeout=2).read().decode() - assert "999.0" in resp1 - assert "999.0" not in resp2 + resp1 = urllib.request.urlopen(f"http://localhost:{port1}/metrics", timeout=2).read().decode() + resp2 = urllib.request.urlopen(f"http://localhost:{port2}/metrics", timeout=2).read().decode() + assert "999.0" in resp1 + assert "999.0" not in resp2 server1.stop() server2.stop() @@ -161,7 +163,6 @@ def test_server_port_conflict_raises(): server1.stop() -@pytest.mark.skipif(not PROMETHEUS_AVAILABLE, reason="prometheus_client not installed") def test_server_update_runs(): port = find_free_port() server = MetricsServer(MetricsServerConfig(port=port)) @@ -194,7 +195,6 @@ def test_server_update_runs(): server.stop() -@pytest.mark.skipif(not PROMETHEUS_AVAILABLE, reason="prometheus_client not installed") def test_server_update_runs_cleanup(): """Test that removed runs are cleaned up from metrics.""" port = find_free_port() diff --git a/tests/unit/utils/test_pathing.py b/tests/unit/utils/test_pathing.py new file mode 100644 index 00000000000..fb419c812c1 --- /dev/null +++ b/tests/unit/utils/test_pathing.py @@ -0,0 +1,57 @@ +import pytest + +from prime_rl.utils.pathing import validate_output_dir + + +def test_nonexistent_dir_passes(tmp_path): + output_dir = tmp_path / "does_not_exist" + validate_output_dir(output_dir, resuming=False, clean=False) + + +def test_empty_dir_passes(tmp_path): + output_dir = tmp_path / "empty" + output_dir.mkdir() + validate_output_dir(output_dir, resuming=False, clean=False) + + +def test_dir_with_only_logs_passes(tmp_path): + output_dir = tmp_path / "has_logs" + output_dir.mkdir() + (output_dir / "logs").mkdir() + (output_dir / "logs" / "trainer.log").touch() + validate_output_dir(output_dir, resuming=False, clean=False) + + +def test_dir_with_checkpoints_raises(tmp_path): + output_dir = tmp_path / "has_ckpt" + output_dir.mkdir() + (output_dir / "checkpoints").mkdir() + (output_dir / "checkpoints" / "step_0").mkdir() + with pytest.raises(FileExistsError, match="already contains checkpoints"): + validate_output_dir(output_dir, resuming=False, clean=False) + + +def test_dir_with_checkpoints_passes_when_resuming(tmp_path): + output_dir = tmp_path / "has_ckpt" + output_dir.mkdir() + (output_dir / "checkpoints").mkdir() + (output_dir / "checkpoints" / "step_0").mkdir() + validate_output_dir(output_dir, resuming=True, clean=False) + + +def test_dir_with_checkpoints_cleaned_when_flag_set(tmp_path): + output_dir = tmp_path / "has_ckpt" + output_dir.mkdir() + (output_dir / "checkpoints").mkdir() + (output_dir / "checkpoints" / "step_0").mkdir() + (output_dir / "logs").mkdir() + + validate_output_dir(output_dir, resuming=False, clean=True) + + assert not output_dir.exists() + + +def test_clean_on_nonexistent_dir_is_noop(tmp_path): + output_dir = tmp_path / "does_not_exist" + validate_output_dir(output_dir, resuming=False, clean=True) + assert not output_dir.exists() diff --git a/tests/unit/utils/test_prime_monitor.py b/tests/unit/utils/test_prime_monitor.py new file mode 100644 index 00000000000..f44065a7c68 --- /dev/null +++ b/tests/unit/utils/test_prime_monitor.py @@ -0,0 +1,111 @@ +import io +import json +from unittest.mock import Mock + +import pyarrow.parquet as pq + +from prime_rl.utils.monitor.prime import PrimeMonitor + + +def _new_monitor() -> PrimeMonitor: + monitor = PrimeMonitor.__new__(PrimeMonitor) + monitor._closed = True + return monitor + + +def _build_rollout(*, example_id: int, reward: float, task: str) -> dict: + return { + "example_id": example_id, + "prompt": [{"role": "user", "content": f"prompt-{example_id}"}], + "completion": [{"role": "assistant", "content": f"completion-{example_id}"}], + "trajectory": [ + { + "prompt": [{"role": "user", "content": f"prompt-{example_id}"}], + "completion": [{"role": "assistant", "content": f"completion-{example_id}"}], + "reward": reward, + "advantage": reward / 2, + "extras": {"source": "test"}, + "tokens": { + "prompt_ids": [1, 2, 3], + "completion_ids": [4, 5], + }, + } + ], + "answer": f"answer-{example_id}", + "task": task, + "info": {"difficulty": "easy"}, + "reward": reward, + "advantage": reward / 2, + "metrics": {"accuracy": reward}, + "timing": {"generation": {"start": 0.0, "end": 12.5, "duration": 12.5}}, + } + + +def test_rollouts_to_parquet_bytes_preserves_all_rollouts_and_ids(): + monitor = _new_monitor() + monitor.run_id = "run-123" + + parquet_bytes = monitor._rollouts_to_parquet_bytes( + [ + _build_rollout(example_id=101, reward=1.0, task="task-a"), + _build_rollout(example_id=202, reward=0.0, task="task-b"), + ], + step=7, + ) + + assert parquet_bytes is not None + + table = pq.read_table(io.BytesIO(parquet_bytes)) + rows = table.to_pylist() + + assert len(rows) == 2 + assert [row["problem_id"] for row in rows] == [101, 202] + assert [row["sample_id"] for row in rows] == [0, 1] + assert all(row["run_id"] == "run-123" for row in rows) + assert all(row["step"] == 7 for row in rows) + assert json.loads(rows[0]["prompt"])[0]["content"] == "prompt-101" + assert json.loads(rows[1]["completion"])[0]["content"] == "completion-202" + + +def test_rollouts_to_parquet_bytes_skips_rollouts_without_trajectory(): + monitor = _new_monitor() + monitor.run_id = "run-456" + + parquet_bytes = monitor._rollouts_to_parquet_bytes( + [ + _build_rollout(example_id=1, reward=1.0, task="task-a"), + { + "example_id": 2, + "prompt": [{"role": "user", "content": "missing-trajectory"}], + "completion": [{"role": "assistant", "content": "ignored"}], + "trajectory": [], + }, + ], + step=3, + ) + + assert parquet_bytes is not None + + table = pq.read_table(io.BytesIO(parquet_bytes)) + rows = table.to_pylist() + + assert len(rows) == 1 + assert rows[0]["problem_id"] == 1 + assert rows[0]["sample_id"] == 0 + + +def test_sanitize_json_payload_drops_non_finite_values_and_logs_paths(): + monitor = _new_monitor() + monitor.logger = Mock() + + payload = { + "metrics": {"finite": 1.0, "nan": float("nan")}, + "distributions": [0.5, float("inf")], + } + + sanitized = monitor._sanitize_json_payload("metrics", payload) + + assert sanitized == {"metrics": {"finite": 1.0}, "distributions": [0.5]} + monitor.logger.warning.assert_called_once_with( + "Dropping 2 non-finite value(s) from Prime monitor metrics payload: metrics.nan, distributions[1]" + ) diff --git a/tests/unit/utils/test_sequence.py b/tests/unit/utils/test_sequence.py new file mode 100644 index 00000000000..cd622ea24a3 --- /dev/null +++ b/tests/unit/utils/test_sequence.py @@ -0,0 +1,25 @@ +import pytest +import torch + +from prime_rl.utils.sequence import get_cu_seqlens_from_position_ids + + +@pytest.mark.parametrize( + ("position_ids", "expected_cu_seqlens", "expected_max_seqlen"), + [ + (torch.arange(8).unsqueeze(0), [0, 8], 8), + (torch.arange(8, 16).unsqueeze(0), [0, 8], 8), + (torch.tensor([[0, 1, 2, 3, 0, 1, 2]]), [0, 4, 7], 4), + (torch.tensor([[5, 6, 7, 0, 1, 2]]), [0, 3, 6], 3), + ], +) +def test_get_cu_seqlens_from_position_ids_is_local_relative( + position_ids: torch.Tensor, + expected_cu_seqlens: list[int], + expected_max_seqlen: int, +) -> None: + cu_seqlens, max_seqlen = get_cu_seqlens_from_position_ids(position_ids) + + assert cu_seqlens.dtype == torch.int32 + assert cu_seqlens.tolist() == expected_cu_seqlens + assert max_seqlen == expected_max_seqlen diff --git a/tests/utils.py b/tests/utils.py index bfbdcdac1ce..27f6739b0ba 100644 --- a/tests/utils.py +++ b/tests/utils.py @@ -13,10 +13,10 @@ def check_no_error(process: ProcessResult, output_dir: Path) -> None: """Helper to assert that a process did not error""" if process.returncode != 0: print("=== Inference Outputs ===") - with open(output_dir / "logs" / "inference.stdout", "r") as f: + with open(output_dir / "logs" / "inference.log", "r") as f: print(*f.readlines()[-100:], sep="\n") print("=== Orchestrator Outputs ===") - with open(output_dir / "logs" / "orchestrator.stdout", "r") as f: + with open(output_dir / "logs" / "orchestrator.log", "r") as f: print(*f.readlines()[-1000:], sep="\n") assert process.returncode == 0, f"Process has non-zero return code ({process})" @@ -29,7 +29,7 @@ def check_number_goes_up_or_down( go_up: bool = True, ): """Helper to assert that a number in lines goes up from a specified start to end step""" - step_lines = [line for line in lines if "Step" in line] + step_lines = [line for line in lines if "SUCCESS" in line and "Step" in line and re.search(pattern, line)] assert len(step_lines) > 0, f"No step lines found in output ({lines})" try: start_step_line = step_lines[start_step] @@ -63,17 +63,18 @@ def check_number_goes_up_or_down( ) -def check_number_in_range( +def check_metric_in_range( lines: list[str], + metric_name: str, + pattern: str, step: int = -1, - min_threshold: float | None = 0.0, + min_threshold: float | None = None, max_threshold: float | None = None, - pattern: str = r"Reward:\s*(\d+\.\d{4})", ): - """Helper to assert that a number in step logs is within a threshold""" - step_lines = [line for line in lines if "Step" in line] + """Helper to assert that a metric in step logs is within a threshold""" + step_lines = [line for line in lines if "SUCCESS" in line and "Step" in line and re.search(pattern, line)] assert len(step_lines) > 0, f"No step lines found in output ({lines})" - + # Search for the specific step number in the lines if step == -1: step_line = step_lines[-1] @@ -82,17 +83,92 @@ def check_number_in_range( matching_lines = [line for line in step_lines if re.search(step_pattern, line)] assert len(matching_lines) > 0, f"Could not find step {step} in output ({step_lines})" step_line = matching_lines[0] - - step_reward = re.search(pattern, step_line) - assert step_reward is not None, f"Could not find reward for step {step}. Line: {step_line} ({lines})" - step_reward = float(step_reward.group(1)) + + metric_match = re.search(pattern, step_line) + assert metric_match is not None, f"Could not find {metric_name} for step {step}. Line: {step_line} ({lines})" + metric_value = float(metric_match.group(1)) if min_threshold is not None: - assert step_reward >= min_threshold, ( - f"Reward did not reach minimum threshold. Found reward={step_reward} < {min_threshold} " + assert metric_value >= min_threshold, ( + f"{metric_name} did not reach minimum threshold. Found {metric_name}={metric_value} < {min_threshold} " f"(line: {step_line}) ({lines})" ) if max_threshold is not None: - assert step_reward <= max_threshold, ( - f"Reward did not reach maximum threshold. Found reward={step_reward} > {max_threshold} " + assert metric_value <= max_threshold, ( + f"{metric_name} exceeded maximum threshold. Found {metric_name}={metric_value} > {max_threshold} " f"(line: {step_line}) ({lines})" ) + + +def check_reward_goes_up(lines: list[str]): + return check_number_goes_up_or_down(lines, go_up=True, pattern=r"Reward:\s*(\d+\.\d{4})") + + +def check_loss_goes_down(lines: list[str]): + return check_number_goes_up_or_down(lines, go_up=False, pattern=r"Loss:\s*(\d+\.\d{4})") + + +def check_reward_in_range( + lines: list[str], + step: int = -1, + min_threshold: float | None = 0.0, + max_threshold: float | None = None, +): + """Helper to assert that reward in step logs is within a threshold""" + check_metric_in_range( + lines, + metric_name="Reward", + pattern=r"Reward:\s*(\d+\.\d{4})", + step=step, + min_threshold=min_threshold, + max_threshold=max_threshold, + ) + + +def check_avg_reward_in_range( + lines: list[str], + last_n_steps: int, + min_threshold: float | None = None, + max_threshold: float | None = None, +): + """Helper to assert that the average reward over the last N steps is within a threshold""" + pattern = r"Reward:\s*(\d+\.\d{4})" + step_lines = [line for line in lines if "SUCCESS" in line and "Step" in line and re.search(pattern, line)] + assert len(step_lines) >= last_n_steps, ( + f"Not enough step lines found. Expected at least {last_n_steps}, got {len(step_lines)}" + ) + + recent_lines = step_lines[-last_n_steps:] + rewards = [] + for line in recent_lines: + match = re.search(pattern, line) + assert match is not None, f"Could not find Reward in line: {line}" + rewards.append(float(match.group(1))) + + avg_reward = sum(rewards) / len(rewards) + if min_threshold is not None: + assert avg_reward >= min_threshold, ( + f"Average reward over last {last_n_steps} steps did not reach minimum threshold. " + f"Found avg_reward={avg_reward:.4f} < {min_threshold} (rewards={rewards})" + ) + if max_threshold is not None: + assert avg_reward <= max_threshold, ( + f"Average reward over last {last_n_steps} steps exceeded maximum threshold. " + f"Found avg_reward={avg_reward:.4f} > {max_threshold} (rewards={rewards})" + ) + + +def check_mismatch_kl_in_range( + lines: list[str], + step: int = -1, + min_threshold: float | None = None, + max_threshold: float | None = None, +): + """Helper to assert that mismatch KL in step logs is within a threshold""" + check_metric_in_range( + lines, + metric_name="Mismatch KL", + pattern=r"Mismatch KL:\s*(\d+\.\d{4})", + step=step, + min_threshold=min_threshold, + max_threshold=max_threshold, + ) diff --git a/uv.lock b/uv.lock index 3447c1c2d46..0b2d2268eed 100644 --- a/uv.lock +++ b/uv.lock @@ -2,12 +2,58 @@ version = 1 revision = 3 requires-python = "==3.12.*" resolution-markers = [ - "sys_platform == 'linux'", - "sys_platform != 'linux'", + "platform_machine == 'x86_64' and sys_platform == 'linux'", + "platform_machine == 'aarch64' and sys_platform == 'linux'", +] +supported-markers = [ + "platform_machine == 'x86_64' and sys_platform == 'linux'", + "platform_machine == 'aarch64' and sys_platform == 'linux'", ] [options] -prerelease-mode = "allow" +exclude-newer = "0001-01-01T00:00:00Z" # This has no effect and is included for backwards compatibility when using relative exclude-newer values. +exclude-newer-span = "P7D" + +[options.exclude-newer-package] +vllm = false +verifiers = false +vllm-router = false +dion = false +alphabet-sort = false +science-env = false +color-codeword = false +nixl-cu12 = false +flash-attn-3 = false +prime-tunnel = false +prime-sandboxes = false +deep-gemm = false +aime2024 = false +prime-evals = false +deepdive = false +prime = false +reverse-text = false +code-env = false +mini-swe-agent-plus = false +deep-ep = false +pydantic-config = false +renderers = false +math-env = false +logic-env = false +wiki-search = false +math-python = false +math500 = false +aime2025 = false + +[manifest] +members = [ + "prime-rl", + "prime-rl-configs", +] +overrides = [ + { name = "nvidia-cudnn-cu12", specifier = ">=9.15" }, + { name = "nvidia-cutlass-dsl", specifier = ">=4.4.1" }, + { name = "transformers", git = "https://github.com/huggingface/transformers.git?rev=c1c3424" }, +] [[package]] name = "absl-py" @@ -19,21 +65,27 @@ wheels = [ ] [[package]] -name = "accelerate" -version = "1.10.1" -source = { registry = "https://pypi.org/simple" } +name = "aime2024" +version = "0.1.18" +source = { registry = "https://hub.primeintellect.ai/primeintellect/simple/" } dependencies = [ - { name = "huggingface-hub" }, - { name = "numpy" }, - { name = "packaging" }, - { name = "psutil" }, - { name = "pyyaml" }, - { name = "safetensors" }, - { name = "torch" }, + { name = "math-verify", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "verifiers", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/b1/72/ff3961c19ee395c3d30ac630ee77bfb0e1b46b87edc504d4f83bb4a89705/accelerate-1.10.1.tar.gz", hash = "sha256:3dea89e433420e4bfac0369cae7e36dcd6a56adfcfd38cdda145c6225eab5df8", size = 392446, upload-time = "2025-08-25T13:57:06.21Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/5f/a0/d9ef19f780f319c21ee90ecfef4431cbeeca95bec7f14071785c17b6029b/accelerate-1.10.1-py3-none-any.whl", hash = "sha256:3621cff60b9a27ce798857ece05e2b9f56fcc71631cfb31ccf71f0359c311f11", size = 374909, upload-time = "2025-08-25T13:57:04.55Z" }, + { url = "https://hub.primeintellect.ai/primeintellect/aime2024/@e6768e58/aime2024-0.1.18-py3-none-any.whl", hash = "sha256:75b4cfdb46108e5b2483c809a82bfa57802c2e476ea800708900f7512e8405f6" }, +] + +[[package]] +name = "aime2025" +version = "0.1.18" +source = { registry = "https://hub.primeintellect.ai/primeintellect/simple/" } +dependencies = [ + { name = "math-verify", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "verifiers", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, +] +wheels = [ + { url = "https://hub.primeintellect.ai/primeintellect/aime2025/@ec6af76f/aime2025-0.1.18-py3-none-any.whl", hash = "sha256:cda06c86accc3bcd5881b7888951e9d7d6573df122814257f7eefa9e6dc3ec16" }, ] [[package]] @@ -56,36 +108,32 @@ wheels = [ [[package]] name = "aiohttp" -version = "3.12.15" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "aiohappyeyeballs" }, - { name = "aiosignal" }, - { name = "attrs" }, - { name = "frozenlist" }, - { name = "multidict" }, - { name = "propcache" }, - { name = "yarl" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/9b/e7/d92a237d8802ca88483906c388f7c201bbe96cd80a165ffd0ac2f6a8d59f/aiohttp-3.12.15.tar.gz", hash = "sha256:4fc61385e9c98d72fcdf47e6dd81833f47b2f77c114c29cd64a361be57a763a2", size = 7823716, upload-time = "2025-07-29T05:52:32.215Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/63/97/77cb2450d9b35f517d6cf506256bf4f5bda3f93a66b4ad64ba7fc917899c/aiohttp-3.12.15-cp312-cp312-macosx_10_13_universal2.whl", hash = "sha256:802d3868f5776e28f7bf69d349c26fc0efadb81676d0afa88ed00d98a26340b7", size = 702333, upload-time = "2025-07-29T05:50:46.507Z" }, - { url = "https://files.pythonhosted.org/packages/83/6d/0544e6b08b748682c30b9f65640d006e51f90763b41d7c546693bc22900d/aiohttp-3.12.15-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:f2800614cd560287be05e33a679638e586a2d7401f4ddf99e304d98878c29444", size = 476948, upload-time = "2025-07-29T05:50:48.067Z" }, - { url = "https://files.pythonhosted.org/packages/3a/1d/c8c40e611e5094330284b1aea8a4b02ca0858f8458614fa35754cab42b9c/aiohttp-3.12.15-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:8466151554b593909d30a0a125d638b4e5f3836e5aecde85b66b80ded1cb5b0d", size = 469787, upload-time = "2025-07-29T05:50:49.669Z" }, - { url = "https://files.pythonhosted.org/packages/38/7d/b76438e70319796bfff717f325d97ce2e9310f752a267bfdf5192ac6082b/aiohttp-3.12.15-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:2e5a495cb1be69dae4b08f35a6c4579c539e9b5706f606632102c0f855bcba7c", size = 1716590, upload-time = "2025-07-29T05:50:51.368Z" }, - { url = "https://files.pythonhosted.org/packages/79/b1/60370d70cdf8b269ee1444b390cbd72ce514f0d1cd1a715821c784d272c9/aiohttp-3.12.15-cp312-cp312-manylinux_2_17_armv7l.manylinux2014_armv7l.manylinux_2_31_armv7l.whl", hash = "sha256:6404dfc8cdde35c69aaa489bb3542fb86ef215fc70277c892be8af540e5e21c0", size = 1699241, upload-time = "2025-07-29T05:50:53.628Z" }, - { url = "https://files.pythonhosted.org/packages/a3/2b/4968a7b8792437ebc12186db31523f541943e99bda8f30335c482bea6879/aiohttp-3.12.15-cp312-cp312-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:3ead1c00f8521a5c9070fcb88f02967b1d8a0544e6d85c253f6968b785e1a2ab", size = 1754335, upload-time = "2025-07-29T05:50:55.394Z" }, - { url = "https://files.pythonhosted.org/packages/fb/c1/49524ed553f9a0bec1a11fac09e790f49ff669bcd14164f9fab608831c4d/aiohttp-3.12.15-cp312-cp312-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:6990ef617f14450bc6b34941dba4f12d5613cbf4e33805932f853fbd1cf18bfb", size = 1800491, upload-time = "2025-07-29T05:50:57.202Z" }, - { url = "https://files.pythonhosted.org/packages/de/5e/3bf5acea47a96a28c121b167f5ef659cf71208b19e52a88cdfa5c37f1fcc/aiohttp-3.12.15-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:fd736ed420f4db2b8148b52b46b88ed038d0354255f9a73196b7bbce3ea97545", size = 1719929, upload-time = "2025-07-29T05:50:59.192Z" }, - { url = "https://files.pythonhosted.org/packages/39/94/8ae30b806835bcd1cba799ba35347dee6961a11bd507db634516210e91d8/aiohttp-3.12.15-cp312-cp312-manylinux_2_5_i686.manylinux1_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:3c5092ce14361a73086b90c6efb3948ffa5be2f5b6fbcf52e8d8c8b8848bb97c", size = 1635733, upload-time = "2025-07-29T05:51:01.394Z" }, - { url = "https://files.pythonhosted.org/packages/7a/46/06cdef71dd03acd9da7f51ab3a9107318aee12ad38d273f654e4f981583a/aiohttp-3.12.15-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:aaa2234bb60c4dbf82893e934d8ee8dea30446f0647e024074237a56a08c01bd", size = 1696790, upload-time = "2025-07-29T05:51:03.657Z" }, - { url = "https://files.pythonhosted.org/packages/02/90/6b4cfaaf92ed98d0ec4d173e78b99b4b1a7551250be8937d9d67ecb356b4/aiohttp-3.12.15-cp312-cp312-musllinux_1_2_armv7l.whl", hash = "sha256:6d86a2fbdd14192e2f234a92d3b494dd4457e683ba07e5905a0b3ee25389ac9f", size = 1718245, upload-time = "2025-07-29T05:51:05.911Z" }, - { url = "https://files.pythonhosted.org/packages/2e/e6/2593751670fa06f080a846f37f112cbe6f873ba510d070136a6ed46117c6/aiohttp-3.12.15-cp312-cp312-musllinux_1_2_i686.whl", hash = "sha256:a041e7e2612041a6ddf1c6a33b883be6a421247c7afd47e885969ee4cc58bd8d", size = 1658899, upload-time = "2025-07-29T05:51:07.753Z" }, - { url = "https://files.pythonhosted.org/packages/8f/28/c15bacbdb8b8eb5bf39b10680d129ea7410b859e379b03190f02fa104ffd/aiohttp-3.12.15-cp312-cp312-musllinux_1_2_ppc64le.whl", hash = "sha256:5015082477abeafad7203757ae44299a610e89ee82a1503e3d4184e6bafdd519", size = 1738459, upload-time = "2025-07-29T05:51:09.56Z" }, - { url = "https://files.pythonhosted.org/packages/00/de/c269cbc4faa01fb10f143b1670633a8ddd5b2e1ffd0548f7aa49cb5c70e2/aiohttp-3.12.15-cp312-cp312-musllinux_1_2_s390x.whl", hash = "sha256:56822ff5ddfd1b745534e658faba944012346184fbfe732e0d6134b744516eea", size = 1766434, upload-time = "2025-07-29T05:51:11.423Z" }, - { url = "https://files.pythonhosted.org/packages/52/b0/4ff3abd81aa7d929b27d2e1403722a65fc87b763e3a97b3a2a494bfc63bc/aiohttp-3.12.15-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:b2acbbfff69019d9014508c4ba0401822e8bae5a5fdc3b6814285b71231b60f3", size = 1726045, upload-time = "2025-07-29T05:51:13.689Z" }, - { url = "https://files.pythonhosted.org/packages/71/16/949225a6a2dd6efcbd855fbd90cf476052e648fb011aa538e3b15b89a57a/aiohttp-3.12.15-cp312-cp312-win32.whl", hash = "sha256:d849b0901b50f2185874b9a232f38e26b9b3d4810095a7572eacea939132d4e1", size = 423591, upload-time = "2025-07-29T05:51:15.452Z" }, - { url = "https://files.pythonhosted.org/packages/2b/d8/fa65d2a349fe938b76d309db1a56a75c4fb8cc7b17a398b698488a939903/aiohttp-3.12.15-cp312-cp312-win_amd64.whl", hash = "sha256:b390ef5f62bb508a9d67cb3bba9b8356e23b3996da7062f1a57ce1a79d2b3d34", size = 450266, upload-time = "2025-07-29T05:51:17.239Z" }, +version = "3.13.3" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "aiohappyeyeballs", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "aiosignal", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "attrs", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "frozenlist", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "multidict", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "propcache", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "yarl", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/50/42/32cf8e7704ceb4481406eb87161349abb46a57fee3f008ba9cb610968646/aiohttp-3.13.3.tar.gz", hash = "sha256:a949eee43d3782f2daae4f4a2819b2cb9b0c5d3b7f7a927067cc84dafdbb9f88", size = 7844556, upload-time = "2026-01-03T17:33:05.204Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/8d/a8/5a35dc56a06a2c90d4742cbf35294396907027f80eea696637945a106f25/aiohttp-3.13.3-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:d32764c6c9aafb7fb55366a224756387cd50bfa720f32b88e0e6fa45b27dcf29", size = 1737839, upload-time = "2026-01-03T17:30:19.422Z" }, + { url = "https://files.pythonhosted.org/packages/86/f6/a62cbbf13f0ac80a70f71b1672feba90fdb21fd7abd8dbf25c0105fb6fa3/aiohttp-3.13.3-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:9ae8dd55c8e6c4257eae3a20fd2c8f41edaea5992ed67156642493b8daf3cecc", size = 1755181, upload-time = "2026-01-03T17:30:27.554Z" }, + { url = "https://files.pythonhosted.org/packages/de/95/8fd69a66682012f6716e1bc09ef8a1a2a91922c5725cb904689f112309c4/aiohttp-3.13.3-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:bb4f7475e359992b580559e008c598091c45b5088f28614e855e42d39c2f1033", size = 1697900, upload-time = "2026-01-03T17:30:31.033Z" }, + { url = "https://files.pythonhosted.org/packages/9d/f2/a07a75173124f31f11ea6f863dc44e6f09afe2bca45dd4e64979490deab1/aiohttp-3.13.3-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:87797e645d9d8e222e04160ee32aa06bc5c163e8499f24db719e7852ec23093a", size = 1722379, upload-time = "2026-01-03T17:30:41.081Z" }, +] + +[[package]] +name = "aiolimiter" +version = "1.2.1" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/f1/23/b52debf471f7a1e42e362d959a3982bdcb4fe13a5d46e63d28868807a79c/aiolimiter-1.2.1.tar.gz", hash = "sha256:e02a37ea1a855d9e832252a105420ad4d15011505512a1a1d814647451b5cca9", size = 7185, upload-time = "2024-12-08T15:31:51.496Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/f3/ba/df6e8e1045aebc4778d19b8a3a9bc1808adb1619ba94ca354d9ba17d86c3/aiolimiter-1.2.1-py3-none-any.whl", hash = "sha256:d3f249e9059a20badcb56b61601a83556133655c11d1eb3dd3e04ff069e5f3c7", size = 6711, upload-time = "2024-12-08T15:31:49.874Z" }, ] [[package]] @@ -93,14 +141,25 @@ name = "aiosignal" version = "1.4.0" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "frozenlist" }, - { name = "typing-extensions" }, + { name = "frozenlist", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "typing-extensions", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, ] sdist = { url = "https://files.pythonhosted.org/packages/61/62/06741b579156360248d1ec624842ad0edf697050bbaf7c3e46394e106ad1/aiosignal-1.4.0.tar.gz", hash = "sha256:f47eecd9468083c2029cc99945502cb7708b082c232f9aca65da147157b251c7", size = 25007, upload-time = "2025-07-03T22:54:43.528Z" } wheels = [ { url = "https://files.pythonhosted.org/packages/fb/76/641ae371508676492379f16e2fa48f4e2c11741bd63c48be4b12a6b09cba/aiosignal-1.4.0-py3-none-any.whl", hash = "sha256:053243f8b92b990551949e63930a839ff0cf0b0ebbe0597b0f3fb19e1a0fe82e", size = 7490, upload-time = "2025-07-03T22:54:42.156Z" }, ] +[[package]] +name = "alphabet-sort" +version = "0.1.12" +source = { registry = "https://hub.primeintellect.ai/primeintellect/simple/" } +dependencies = [ + { name = "verifiers", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, +] +wheels = [ + { url = "https://hub.primeintellect.ai/primeintellect/alphabet-sort/@c6c2fbe1/alphabet_sort-0.1.12-py2.py3-none-any.whl", hash = "sha256:a5c14b6eb8334e1876f843d753231c200bb7d8d821c8ab3c750b87b383dc9cc8" }, +] + [[package]] name = "annotated-doc" version = "0.0.4" @@ -121,21 +180,21 @@ wheels = [ [[package]] name = "anthropic" -version = "0.71.0" +version = "0.79.0" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "anyio" }, - { name = "distro" }, - { name = "docstring-parser" }, - { name = "httpx" }, - { name = "jiter" }, - { name = "pydantic" }, - { name = "sniffio" }, - { name = "typing-extensions" }, + { name = "anyio", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "distro", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "docstring-parser", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "httpx", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "jiter", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "pydantic", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "sniffio", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "typing-extensions", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/82/4f/70682b068d897841f43223df82d96ec1d617435a8b759c4a2d901a50158b/anthropic-0.71.0.tar.gz", hash = "sha256:eb8e6fa86d049061b3ef26eb4cbae0174ebbff21affa6de7b3098da857d8de6a", size = 489102, upload-time = "2025-10-16T15:54:40.08Z" } +sdist = { url = "https://files.pythonhosted.org/packages/15/b1/91aea3f8fd180d01d133d931a167a78a3737b3fd39ccef2ae8d6619c24fd/anthropic-0.79.0.tar.gz", hash = "sha256:8707aafb3b1176ed6c13e2b1c9fb3efddce90d17aee5d8b83a86c70dcdcca871", size = 509825, upload-time = "2026-02-07T18:06:18.388Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/5d/77/073e8ac488f335aec7001952825275582fb8f433737e90f24eeef9d878f6/anthropic-0.71.0-py3-none-any.whl", hash = "sha256:85c5015fcdbdc728390f11b17642a65a4365d03b12b799b18b6cc57e71fdb327", size = 355035, upload-time = "2025-10-16T15:54:38.238Z" }, + { url = "https://files.pythonhosted.org/packages/95/b2/cc0b8e874a18d7da50b0fda8c99e4ac123f23bf47b471827c5f6f3e4a767/anthropic-0.79.0-py3-none-any.whl", hash = "sha256:04cbd473b6bbda4ca2e41dd670fe2f829a911530f01697d0a1e37321eb75f3cf", size = 405918, upload-time = "2026-02-07T18:06:20.246Z" }, ] [[package]] @@ -149,9 +208,9 @@ name = "anyio" version = "4.10.0" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "idna" }, - { name = "sniffio" }, - { name = "typing-extensions" }, + { name = "idna", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "sniffio", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "typing-extensions", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, ] sdist = { url = "https://files.pythonhosted.org/packages/f1/b4/636b3b65173d3ce9a38ef5f0522789614e590dab6a8d505340a4efe4c567/anyio-4.10.0.tar.gz", hash = "sha256:3f3fae35c96039744587aa5b8371e7e8e603c0702999535961dd336026973ba6", size = 213252, upload-time = "2025-08-04T08:54:26.451Z" } wheels = [ @@ -160,28 +219,17 @@ wheels = [ [[package]] name = "apache-tvm-ffi" -version = "0.1.5" +version = "0.1.9" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "typing-extensions" }, + { name = "typing-extensions", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/2f/21/adcd83d08482fc297ece9a6f0e3820588a71b7bb196d188a2d6a2ee71884/apache_tvm_ffi-0.1.5.tar.gz", hash = "sha256:21bc35cb9d9fdec54061b802a2b432fd1155d2733da94df8678ff21bab1d9a2f", size = 1340655, upload-time = "2025-12-07T19:39:15.306Z" } +sdist = { url = "https://files.pythonhosted.org/packages/6f/60/1e787a0b5ebf318483235be2a689ee367173983067e441b8379564f667c0/apache_tvm_ffi-0.1.9.tar.gz", hash = "sha256:d2d402587e8906de0a07f4746aa78f3d452c7efe3625d4bb39ac2ad693bce530", size = 2513731, upload-time = "2026-02-27T19:28:06.602Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/04/43/7d9c821670beb6e365b4f056f97411db7a6fa976dc982557e699f5dc853c/apache_tvm_ffi-0.1.5-cp312-abi3-macosx_11_0_arm64.whl", hash = "sha256:868a59786c405dcee01846788cdb228069c974ec76967268d9da3e2386cc51c7", size = 1767737, upload-time = "2025-12-07T19:38:34.172Z" }, - { url = "https://files.pythonhosted.org/packages/ad/10/792243025104d1371105c11f5c3d7e018e03545bda0e219fa1c8a45ecb16/apache_tvm_ffi-0.1.5-cp312-abi3-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:405a4a9181fa0f03cb28ee72ce26d48b9108a234b0f0d9b0516e21e5cad0807b", size = 1924320, upload-time = "2025-12-07T19:38:35.574Z" }, - { url = "https://files.pythonhosted.org/packages/bb/20/15403acb9cc06f7e32f31f197531f3c9184392b63ebd9933be48b5fd75af/apache_tvm_ffi-0.1.5-cp312-abi3-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:78eebce47d10d6ebaf749f202b914cdc0819d5bbaaac76cd54aff58103ed0110", size = 1998767, upload-time = "2025-12-07T19:38:37.183Z" }, - { url = "https://files.pythonhosted.org/packages/cc/4d/73922e5a3e9014d84f3b98e901e0339be9a041ca276cf3a37f5457b04842/apache_tvm_ffi-0.1.5-cp312-abi3-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:1ce92887ad67e773702b9536682a4ea5f9362c5fd43abe40abcc8266ab26e074", size = 1870775, upload-time = "2025-12-07T19:38:38.679Z" }, - { url = "https://files.pythonhosted.org/packages/17/2d/a5b857f8ce4af36355744e6f713c81725edf5ab538398c55eb13b1c84fd0/apache_tvm_ffi-0.1.5-cp312-abi3-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:e32cdc18ca33ab1a7d17707bf2c5ad8c1d90ccd7ed822d1f80a9f694768cfef4", size = 1980726, upload-time = "2025-12-07T19:38:40.062Z" }, - { url = "https://files.pythonhosted.org/packages/a0/72/581e314c7e308f9fc363afb4c96fe43de9253477f0b7a3236da9aa40ec8a/apache_tvm_ffi-0.1.5-cp312-abi3-win_amd64.whl", hash = "sha256:6a8dec4372e533b9d4fc09c78bcd95a1222b4c992e6f3d4b34ad41092bbd5435", size = 1748073, upload-time = "2025-12-07T19:38:42.915Z" }, -] - -[[package]] -name = "appnope" -version = "0.1.4" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/35/5d/752690df9ef5b76e169e68d6a129fa6d08a7100ca7f754c89495db3c6019/appnope-0.1.4.tar.gz", hash = "sha256:1de3860566df9caf38f01f86f65e0e13e379af54f9e4bee1e66b48f2efffd1ee", size = 4170, upload-time = "2024-02-06T09:43:11.258Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/81/29/5ecc3a15d5a33e31b26c11426c45c501e439cb865d0bff96315d86443b78/appnope-0.1.4-py2.py3-none-any.whl", hash = "sha256:502575ee11cd7a28c0205f379b525beefebab9d161b7c964670864014ed7213c", size = 4321, upload-time = "2024-02-06T09:43:09.663Z" }, + { url = "https://files.pythonhosted.org/packages/a7/c0/6d3d54f50012255b41bc3e24944c086f63c4707c8686c7c6780e9283eb96/apache_tvm_ffi-0.1.9-cp312-abi3-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:7d503029e66c43b1a1cb1a42a1e9bb428c8a28dcbdec31c28e705472ca648a3a", size = 2203712, upload-time = "2026-02-27T19:27:25.867Z" }, + { url = "https://files.pythonhosted.org/packages/c6/dd/2bab4c6cd86257dbf99e93452a1af833113f8dc3e25a25579f6e4e4c8a94/apache_tvm_ffi-0.1.9-cp312-abi3-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:28241371934ea8af10d5067087ba1229ebddded7b2c02d33a258ec2a96df8c46", size = 2299704, upload-time = "2026-02-27T19:27:27.477Z" }, + { url = "https://files.pythonhosted.org/packages/7a/4a/b469bcb2e1014cb84d336d2a59f42958a058251c577a4c2680cacad346e2/apache_tvm_ffi-0.1.9-cp312-abi3-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:87cacce81df55685fc6a76e1e3c5db1200e85e87bf5974b692c59d131b7bc622", size = 2130865, upload-time = "2026-02-27T19:27:29.092Z" }, + { url = "https://files.pythonhosted.org/packages/70/ef/5402da5d37f5270fd88ea0348acca78dba9be8bdbf6c2bcae0935eb03ef1/apache_tvm_ffi-0.1.9-cp312-abi3-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:f45eb43499acac45ff6c93564f0ff2d3ca27b69656d540fd56ce59d51c0b4c65", size = 2278991, upload-time = "2026-02-27T19:27:30.729Z" }, ] [[package]] @@ -211,6 +259,34 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/77/06/bb80f5f86020c4551da315d78b3ab75e8228f89f0162f2c3a819e407941a/attrs-25.3.0-py3-none-any.whl", hash = "sha256:427318ce031701fea540783410126f03899a97ffc6f61596ad581ac2e40e3bc3", size = 63815, upload-time = "2025-03-13T11:10:21.14Z" }, ] +[[package]] +name = "bcrypt" +version = "5.0.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/d4/36/3329e2518d70ad8e2e5817d5a4cac6bba05a47767ec416c7d020a965f408/bcrypt-5.0.0.tar.gz", hash = "sha256:f748f7c2d6fd375cc93d3fba7ef4a9e3a092421b8dbf34d8d4dc06be9492dfdd", size = 25386, upload-time = "2025-09-25T19:50:47.829Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/45/b6/4c1205dde5e464ea3bd88e8742e19f899c16fa8916fb8510a851fae985b5/bcrypt-5.0.0-cp38-abi3-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:c2388ca94ffee269b6038d48747f4ce8df0ffbea43f31abfa18ac72f0218effb", size = 275009, upload-time = "2025-09-25T19:49:50.581Z" }, + { url = "https://files.pythonhosted.org/packages/3b/71/427945e6ead72ccffe77894b2655b695ccf14ae1866cd977e185d606dd2f/bcrypt-5.0.0-cp38-abi3-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:560ddb6ec730386e7b3b26b8b4c88197aaed924430e7b74666a586ac997249ef", size = 278029, upload-time = "2025-09-25T19:49:52.533Z" }, + { url = "https://files.pythonhosted.org/packages/17/72/c344825e3b83c5389a369c8a8e58ffe1480b8a699f46c127c34580c4666b/bcrypt-5.0.0-cp38-abi3-manylinux_2_28_aarch64.whl", hash = "sha256:d79e5c65dcc9af213594d6f7f1fa2c98ad3fc10431e7aa53c176b441943efbdd", size = 275907, upload-time = "2025-09-25T19:49:54.709Z" }, + { url = "https://files.pythonhosted.org/packages/0f/c3/0ae57a68be2039287ec28bc463b82e4b8dc23f9d12c0be331f4782e19108/bcrypt-5.0.0-cp38-abi3-manylinux_2_28_x86_64.whl", hash = "sha256:0c8e093ea2532601a6f686edbc2c6b2ec24131ff5c52f7610dd64fa4553b5464", size = 278412, upload-time = "2025-09-25T19:49:57.356Z" }, + { url = "https://files.pythonhosted.org/packages/45/2b/77424511adb11e6a99e3a00dcc7745034bee89036ad7d7e255a7e47be7d8/bcrypt-5.0.0-cp38-abi3-manylinux_2_34_aarch64.whl", hash = "sha256:5b1589f4839a0899c146e8892efe320c0fa096568abd9b95593efac50a87cb75", size = 275486, upload-time = "2025-09-25T19:49:59.116Z" }, + { url = "https://files.pythonhosted.org/packages/43/0a/405c753f6158e0f3f14b00b462d8bca31296f7ecfc8fc8bc7919c0c7d73a/bcrypt-5.0.0-cp38-abi3-manylinux_2_34_x86_64.whl", hash = "sha256:89042e61b5e808b67daf24a434d89bab164d4de1746b37a8d173b6b14f3db9ff", size = 277940, upload-time = "2025-09-25T19:50:00.869Z" }, + { url = "https://files.pythonhosted.org/packages/62/83/b3efc285d4aadc1fa83db385ec64dcfa1707e890eb42f03b127d66ac1b7b/bcrypt-5.0.0-cp38-abi3-musllinux_1_1_aarch64.whl", hash = "sha256:e3cf5b2560c7b5a142286f69bde914494b6d8f901aaa71e453078388a50881c4", size = 310776, upload-time = "2025-09-25T19:50:02.393Z" }, + { url = "https://files.pythonhosted.org/packages/95/7d/47ee337dacecde6d234890fe929936cb03ebc4c3a7460854bbd9c97780b8/bcrypt-5.0.0-cp38-abi3-musllinux_1_1_x86_64.whl", hash = "sha256:f632fd56fc4e61564f78b46a2269153122db34988e78b6be8b32d28507b7eaeb", size = 312922, upload-time = "2025-09-25T19:50:04.232Z" }, + { url = "https://files.pythonhosted.org/packages/d6/3a/43d494dfb728f55f4e1cf8fd435d50c16a2d75493225b54c8d06122523c6/bcrypt-5.0.0-cp38-abi3-musllinux_1_2_aarch64.whl", hash = "sha256:801cad5ccb6b87d1b430f183269b94c24f248dddbbc5c1f78b6ed231743e001c", size = 341367, upload-time = "2025-09-25T19:50:05.559Z" }, + { url = "https://files.pythonhosted.org/packages/55/ab/a0727a4547e383e2e22a630e0f908113db37904f58719dc48d4622139b5c/bcrypt-5.0.0-cp38-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:3cf67a804fc66fc217e6914a5635000259fbbbb12e78a99488e4d5ba445a71eb", size = 359187, upload-time = "2025-09-25T19:50:06.916Z" }, + { url = "https://files.pythonhosted.org/packages/ac/ee/2f4985dbad090ace5ad1f7dd8ff94477fe089b5fab2040bd784a3d5f187b/bcrypt-5.0.0-cp39-abi3-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:ddb4e1500f6efdd402218ffe34d040a1196c072e07929b9820f363a1fd1f4191", size = 275290, upload-time = "2025-09-25T19:50:13.673Z" }, + { url = "https://files.pythonhosted.org/packages/e4/6e/b77ade812672d15cf50842e167eead80ac3514f3beacac8902915417f8b7/bcrypt-5.0.0-cp39-abi3-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:7aeef54b60ceddb6f30ee3db090351ecf0d40ec6e2abf41430997407a46d2254", size = 278253, upload-time = "2025-09-25T19:50:15.089Z" }, + { url = "https://files.pythonhosted.org/packages/36/c4/ed00ed32f1040f7990dac7115f82273e3c03da1e1a1587a778d8cea496d8/bcrypt-5.0.0-cp39-abi3-manylinux_2_28_aarch64.whl", hash = "sha256:f0ce778135f60799d89c9693b9b398819d15f1921ba15fe719acb3178215a7db", size = 276084, upload-time = "2025-09-25T19:50:16.699Z" }, + { url = "https://files.pythonhosted.org/packages/24/b4/11f8a31d8b67cca3371e046db49baa7c0594d71eb40ac8121e2fc0888db0/bcrypt-5.0.0-cp39-abi3-manylinux_2_28_x86_64.whl", hash = "sha256:f8429e1c410b4073944f03bd778a9e066e7fad723564a52ff91841d278dfc822", size = 278656, upload-time = "2025-09-25T19:50:19.809Z" }, + { url = "https://files.pythonhosted.org/packages/ac/31/79f11865f8078e192847d2cb526e3fa27c200933c982c5b2869720fa5fce/bcrypt-5.0.0-cp39-abi3-manylinux_2_34_aarch64.whl", hash = "sha256:edfcdcedd0d0f05850c52ba3127b1fce70b9f89e0fe5ff16517df7e81fa3cbb8", size = 275662, upload-time = "2025-09-25T19:50:21.567Z" }, + { url = "https://files.pythonhosted.org/packages/d4/8d/5e43d9584b3b3591a6f9b68f755a4da879a59712981ef5ad2a0ac1379f7a/bcrypt-5.0.0-cp39-abi3-manylinux_2_34_x86_64.whl", hash = "sha256:611f0a17aa4a25a69362dcc299fda5c8a3d4f160e2abb3831041feb77393a14a", size = 278240, upload-time = "2025-09-25T19:50:23.305Z" }, + { url = "https://files.pythonhosted.org/packages/89/48/44590e3fc158620f680a978aafe8f87a4c4320da81ed11552f0323aa9a57/bcrypt-5.0.0-cp39-abi3-musllinux_1_1_aarch64.whl", hash = "sha256:db99dca3b1fdc3db87d7c57eac0c82281242d1eabf19dcb8a6b10eb29a2e72d1", size = 311152, upload-time = "2025-09-25T19:50:24.597Z" }, + { url = "https://files.pythonhosted.org/packages/5f/85/e4fbfc46f14f47b0d20493669a625da5827d07e8a88ee460af6cd9768b44/bcrypt-5.0.0-cp39-abi3-musllinux_1_1_x86_64.whl", hash = "sha256:5feebf85a9cefda32966d8171f5db7e3ba964b77fdfe31919622256f80f9cf42", size = 313284, upload-time = "2025-09-25T19:50:26.268Z" }, + { url = "https://files.pythonhosted.org/packages/25/ae/479f81d3f4594456a01ea2f05b132a519eff9ab5768a70430fa1132384b1/bcrypt-5.0.0-cp39-abi3-musllinux_1_2_aarch64.whl", hash = "sha256:3ca8a166b1140436e058298a34d88032ab62f15aae1c598580333dc21d27ef10", size = 341643, upload-time = "2025-09-25T19:50:28.02Z" }, + { url = "https://files.pythonhosted.org/packages/df/d2/36a086dee1473b14276cd6ea7f61aef3b2648710b5d7f1c9e032c29b859f/bcrypt-5.0.0-cp39-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:61afc381250c3182d9078551e3ac3a41da14154fbff647ddf52a769f588c4172", size = 359698, upload-time = "2025-09-25T19:50:31.347Z" }, +] + [[package]] name = "beartype" version = "0.21.0" @@ -220,39 +296,29 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/94/31/87045d1c66ee10a52486c9d2047bc69f00f2689f69401bb1e998afb4b205/beartype-0.21.0-py3-none-any.whl", hash = "sha256:b6a1bd56c72f31b0a496a36cc55df6e2f475db166ad07fa4acc7e74f4c7f34c0", size = 1191340, upload-time = "2025-05-22T05:09:24.606Z" }, ] +[[package]] +name = "beautifulsoup4" +version = "4.14.3" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "soupsieve", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "typing-extensions", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/c3/b0/1c6a16426d389813b48d95e26898aff79abbde42ad353958ad95cc8c9b21/beautifulsoup4-4.14.3.tar.gz", hash = "sha256:6292b1c5186d356bba669ef9f7f051757099565ad9ada5dd630bd9de5fa7fb86", size = 627737, upload-time = "2025-11-30T15:08:26.084Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/1a/39/47f9197bdd44df24d67ac8893641e16f386c984a0619ef2ee4c51fbbc019/beautifulsoup4-4.14.3-py3-none-any.whl", hash = "sha256:0918bfe44902e6ad8d57732ba310582e98da931428d231a5ecb9e7c703a735bb", size = 107721, upload-time = "2025-11-30T15:08:24.087Z" }, +] + [[package]] name = "blake3" version = "1.0.5" source = { registry = "https://pypi.org/simple" } sdist = { url = "https://files.pythonhosted.org/packages/e7/08/22b6326dbe002ca77c92082b37b14a935003897b0e3eed025da92c700751/blake3-1.0.5.tar.gz", hash = "sha256:7bac73f393a67ea6d5ac32e4a45d39c184487c89c712ab3ed839c1a51ed82259", size = 115140, upload-time = "2025-05-19T20:08:29.911Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/8f/a4/7ea6cb45d8ce36b05dd01cc35a1bf9921c07d36dc56869e461f0e832ca76/blake3-1.0.5-cp312-cp312-macosx_10_12_x86_64.whl", hash = "sha256:73dd1bfc802e2343113805d104b9600e794bf700c844f05dda86a9a05c0e7c41", size = 345971, upload-time = "2025-05-19T20:07:03.913Z" }, - { url = "https://files.pythonhosted.org/packages/13/09/87c56b1d3113e1381178e2ff386ac58d32b23c65b20054ce4b8de59be93d/blake3-1.0.5-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:d4e53332a5db53a652395f5e56c72fb81c7e584a192e6931a4eb3f9b32edcf0a", size = 328272, upload-time = "2025-05-19T20:07:05.158Z" }, { url = "https://files.pythonhosted.org/packages/c1/40/b81a25077df6fa1722be8c268732205281e12a244f9d5a15e9e72c2baa04/blake3-1.0.5-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:abe84cc2db3172bbade48dbf7b6029decb82e9cd382bc3cb783b8624a3ee55d8", size = 374599, upload-time = "2025-05-19T20:07:06.951Z" }, - { url = "https://files.pythonhosted.org/packages/58/1b/8fc14c7b7ae116edc42f8e8cd5c21a99d8b68ab761e31347c4c9c6bbedf6/blake3-1.0.5-cp312-cp312-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:ca8935b4a733968a463d6445dc7cb0dcc09759c280df4847f020deec8fcaff27", size = 375221, upload-time = "2025-05-19T20:07:08.39Z" }, - { url = "https://files.pythonhosted.org/packages/26/fa/879c74815dbb39e9b91d35b672b25c3547435e479b9aaf1a80191a86f3f4/blake3-1.0.5-cp312-cp312-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:12e5c722ef966f2b8df0d4024e6f4afd4c466bb0dcd3f8f671fad6cb5dab6a3e", size = 445913, upload-time = "2025-05-19T20:07:09.698Z" }, - { url = "https://files.pythonhosted.org/packages/ce/91/e335f22765d7e80fd5aa6a25b2f2f5f0c5d649049f88d0c8ac1f6a8c431d/blake3-1.0.5-cp312-cp312-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:15ecd628f824d5591a1958babd4217749f1facd3945f33a14c3e5fbb52ffb922", size = 509907, upload-time = "2025-05-19T20:07:11.023Z" }, - { url = "https://files.pythonhosted.org/packages/9b/ec/c1676c275592efdb3a6e4489d0f5e029d38565593466ba70c42b73e76b1a/blake3-1.0.5-cp312-cp312-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:a12b12df3c40089bf2785c333f8f1161b2a66ecacb44828de9fbf2868037934b", size = 395611, upload-time = "2025-05-19T20:07:12.815Z" }, { url = "https://files.pythonhosted.org/packages/5c/04/a86bfb3c20e859e43ead0b13be59afd98feb166ea929e76fa3d190f65f6e/blake3-1.0.5-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:f39e8d36e33f413938230683f192f0565f44ee2b050ad92fb94b343706f3df55", size = 384757, upload-time = "2025-05-19T20:07:14.122Z" }, { url = "https://files.pythonhosted.org/packages/6b/bf/93ce719f88b48d5bcdf2f765789a5a955ea6a02a33f310321508c8421ad6/blake3-1.0.5-cp312-cp312-musllinux_1_1_aarch64.whl", hash = "sha256:7083e1b2cfb737c812e20d790c232c38045c7bfe37ef02526f395d491f90f213", size = 551032, upload-time = "2025-05-19T20:07:15.56Z" }, { url = "https://files.pythonhosted.org/packages/13/99/a2e644e0a2039977beb67abbc1f48f6f6c7e0f0c345665811cfa2880b196/blake3-1.0.5-cp312-cp312-musllinux_1_1_x86_64.whl", hash = "sha256:21240932fc914fd719e2d33297f29742c28a31d8a96cb666ec4679bf2c35aa48", size = 555543, upload-time = "2025-05-19T20:07:17.056Z" }, - { url = "https://files.pythonhosted.org/packages/45/15/80d9b2866af5d7ec4c665bb961b16d3db9a9527a80de78e44b828129d51f/blake3-1.0.5-cp312-cp312-win32.whl", hash = "sha256:cba3e6d12bd310b5ff4970daddd7e77a0ca383678e1f0a1ec414d4c7cb083f9d", size = 234714, upload-time = "2025-05-19T20:07:18.321Z" }, - { url = "https://files.pythonhosted.org/packages/09/a5/76cd4402c685ad1d336351f22483bc2ecd48e5604ba5f5ad340e22b8703a/blake3-1.0.5-cp312-cp312-win_amd64.whl", hash = "sha256:adb54b8bfe4fb2e8106b3a1bddc3614d2de555d2b657861068160176ff723eb0", size = 222127, upload-time = "2025-05-19T20:07:19.579Z" }, -] - -[[package]] -name = "blobfile" -version = "3.1.0" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "filelock" }, - { name = "lxml" }, - { name = "pycryptodomex" }, - { name = "urllib3" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/f0/6d/2e7567da75ddbb24fe979f52284b708da349d67a41042635af36071a5a6b/blobfile-3.1.0.tar.gz", hash = "sha256:d45b6b1fa3b0920732314c23ddbdb4f494ca12f787c2b6eb6bba6faa51382671", size = 77229, upload-time = "2025-09-06T00:36:15.583Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/77/a7/51af11120d75af2828f8eede0b13a4caff650d708ac50e62d000aefe1ffb/blobfile-3.1.0-py3-none-any.whl", hash = "sha256:2b4c5e766ebb7dfa20e4990cf6ec3d2106bdc91d632fb9377f170a234c5a5c6a", size = 75741, upload-time = "2025-09-06T00:36:14.11Z" }, ] [[package]] @@ -260,9 +326,8 @@ name = "build" version = "1.3.0" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "colorama", marker = "os_name == 'nt' and sys_platform != 'linux'" }, - { name = "packaging" }, - { name = "pyproject-hooks" }, + { name = "packaging", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "pyproject-hooks", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, ] sdist = { url = "https://files.pythonhosted.org/packages/25/1c/23e33405a7c9eac261dff640926b8b5adaed6a6eb3e1767d441ed611d0c0/build-1.3.0.tar.gz", hash = "sha256:698edd0ea270bde950f53aed21f3a0135672206f3911e0176261a31e0e07b397", size = 48544, upload-time = "2025-08-01T21:27:09.268Z" } wheels = [ @@ -284,14 +349,10 @@ version = "5.7.0" source = { registry = "https://pypi.org/simple" } sdist = { url = "https://files.pythonhosted.org/packages/3a/89/01df16cdc9c60c07956756c90fe92c684021003079e358a78e213bce45a2/cbor2-5.7.0.tar.gz", hash = "sha256:3f6d843f4db4d0ec501c46453c22a4fbebb1abfb5b740e1bcab34c615cd7406b", size = 102374, upload-time = "2025-08-14T08:59:47.294Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/b1/b1/d54c41b1bc71b8dea0bad3409d2a497df35f7b5ae5db70c1cc9ebc8d556d/cbor2-5.7.0-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:7ad36f0537b75c1aa2c7a462cbdbeec5e8ba02802ea985e0b9fe5deee3b946f4", size = 69020, upload-time = "2025-08-14T08:59:02.276Z" }, - { url = "https://files.pythonhosted.org/packages/f4/e0/45368d5d78b520caaa9ca5a09f55365bc9933d43bce978a528922654ca9f/cbor2-5.7.0-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:5fc9b335cf28e63d9eed4ae03d1e8f90f1a6b287cabc8d29bfddf73fa70643e9", size = 68950, upload-time = "2025-08-14T08:59:03.882Z" }, { url = "https://files.pythonhosted.org/packages/1e/6a/9aed5b716407c1d48425ba55c6022a01a9abdbf58a691f50416461fa371d/cbor2-5.7.0-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:16bea83598a1eeedbd50c2e9fdf3685bae78ca9d9ec8cd8010777db14a315578", size = 285685, upload-time = "2025-08-14T08:59:05.165Z" }, { url = "https://files.pythonhosted.org/packages/a8/6e/3499eaa2b858c7695a447b6311303f06ffc90fc2c45851337121661f1f5c/cbor2-5.7.0-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:8e837825a16e60ace6e284095aa9fbe504bf87a8f4494bf7d95931e37fb01a70", size = 284948, upload-time = "2025-08-14T08:59:06.64Z" }, { url = "https://files.pythonhosted.org/packages/d1/3e/ae67866ef65717665e0acf2873d466c5d4a1d965b0d0348f2269b73f28fb/cbor2-5.7.0-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:27396c5e275ff7c7cd87fe8aaadf781e6194903921f250934af7c86d5efec82e", size = 276375, upload-time = "2025-08-14T08:59:08.845Z" }, { url = "https://files.pythonhosted.org/packages/b6/3d/2f8e9671111661dd571de206344ecd7706f6d458aab191e06834c89aa58e/cbor2-5.7.0-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:c84bfef78c4e9c81eb0a10cec340222ba4e39498a63fc2e3d5f982a3f4efa4a7", size = 277680, upload-time = "2025-08-14T08:59:10.292Z" }, - { url = "https://files.pythonhosted.org/packages/85/03/27a9fefa4e084c1129d7180727791a166629fdae39e0609508401d322626/cbor2-5.7.0-cp312-cp312-win_amd64.whl", hash = "sha256:f64270a24aaadb15dd31cbd64a98d99fca8e0398a65b1570ba07f3c259eb5516", size = 68354, upload-time = "2025-08-14T08:59:11.561Z" }, - { url = "https://files.pythonhosted.org/packages/25/d9/b856d078696542a0d7486d1ece5c936e937bebe5b114674db18d76feb131/cbor2-5.7.0-cp312-cp312-win_arm64.whl", hash = "sha256:73ef321d7b580f08c9fadc41c3d2a218aa3f01e163be9793c6969aadee07f57a", size = 63896, upload-time = "2025-08-14T08:59:12.977Z" }, { url = "https://files.pythonhosted.org/packages/41/cc/0ce73676d2a0c9e5a9330c301940c50eb325dacf5f6d9690fd43a8817fe9/cbor2-5.7.0-py3-none-any.whl", hash = "sha256:a871e7a6f7cba1ddb02503ea974f15f6524c95078fbfe0b860fd4193d7c8f27a", size = 23828, upload-time = "2025-08-14T08:59:46.129Z" }, ] @@ -309,22 +370,14 @@ name = "cffi" version = "2.0.0" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "pycparser", marker = "implementation_name != 'PyPy'" }, + { name = "pycparser", marker = "(implementation_name != 'PyPy' and platform_machine == 'aarch64' and sys_platform == 'linux') or (implementation_name != 'PyPy' and platform_machine == 'x86_64' and sys_platform == 'linux')" }, ] sdist = { url = "https://files.pythonhosted.org/packages/eb/56/b1ba7935a17738ae8453301356628e8147c79dbb825bcbc73dc7401f9846/cffi-2.0.0.tar.gz", hash = "sha256:44d1b5909021139fe36001ae048dbdde8214afa20200eda0f64c068cac5d5529", size = 523588, upload-time = "2025-09-08T23:24:04.541Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/ea/47/4f61023ea636104d4f16ab488e268b93008c3d0bb76893b1b31db1f96802/cffi-2.0.0-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:6d02d6655b0e54f54c4ef0b94eb6be0607b70853c45ce98bd278dc7de718be5d", size = 185271, upload-time = "2025-09-08T23:22:44.795Z" }, - { url = "https://files.pythonhosted.org/packages/df/a2/781b623f57358e360d62cdd7a8c681f074a71d445418a776eef0aadb4ab4/cffi-2.0.0-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:8eca2a813c1cb7ad4fb74d368c2ffbbb4789d377ee5bb8df98373c2cc0dee76c", size = 181048, upload-time = "2025-09-08T23:22:45.938Z" }, - { url = "https://files.pythonhosted.org/packages/ff/df/a4f0fbd47331ceeba3d37c2e51e9dfc9722498becbeec2bd8bc856c9538a/cffi-2.0.0-cp312-cp312-manylinux1_i686.manylinux2014_i686.manylinux_2_17_i686.manylinux_2_5_i686.whl", hash = "sha256:21d1152871b019407d8ac3985f6775c079416c282e431a4da6afe7aefd2bccbe", size = 212529, upload-time = "2025-09-08T23:22:47.349Z" }, { url = "https://files.pythonhosted.org/packages/d5/72/12b5f8d3865bf0f87cf1404d8c374e7487dcf097a1c91c436e72e6badd83/cffi-2.0.0-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:b21e08af67b8a103c71a250401c78d5e0893beff75e28c53c98f4de42f774062", size = 220097, upload-time = "2025-09-08T23:22:48.677Z" }, - { url = "https://files.pythonhosted.org/packages/c2/95/7a135d52a50dfa7c882ab0ac17e8dc11cec9d55d2c18dda414c051c5e69e/cffi-2.0.0-cp312-cp312-manylinux2014_ppc64le.manylinux_2_17_ppc64le.whl", hash = "sha256:1e3a615586f05fc4065a8b22b8152f0c1b00cdbc60596d187c2a74f9e3036e4e", size = 207983, upload-time = "2025-09-08T23:22:50.06Z" }, - { url = "https://files.pythonhosted.org/packages/3a/c8/15cb9ada8895957ea171c62dc78ff3e99159ee7adb13c0123c001a2546c1/cffi-2.0.0-cp312-cp312-manylinux2014_s390x.manylinux_2_17_s390x.whl", hash = "sha256:81afed14892743bbe14dacb9e36d9e0e504cd204e0b165062c488942b9718037", size = 206519, upload-time = "2025-09-08T23:22:51.364Z" }, { url = "https://files.pythonhosted.org/packages/78/2d/7fa73dfa841b5ac06c7b8855cfc18622132e365f5b81d02230333ff26e9e/cffi-2.0.0-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:3e17ed538242334bf70832644a32a7aae3d83b57567f9fd60a26257e992b79ba", size = 219572, upload-time = "2025-09-08T23:22:52.902Z" }, { url = "https://files.pythonhosted.org/packages/07/e0/267e57e387b4ca276b90f0434ff88b2c2241ad72b16d31836adddfd6031b/cffi-2.0.0-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:3925dd22fa2b7699ed2617149842d2e6adde22b262fcbfada50e3d195e4b3a94", size = 222963, upload-time = "2025-09-08T23:22:54.518Z" }, { url = "https://files.pythonhosted.org/packages/b6/75/1f2747525e06f53efbd878f4d03bac5b859cbc11c633d0fb81432d98a795/cffi-2.0.0-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:2c8f814d84194c9ea681642fd164267891702542f028a15fc97d4674b6206187", size = 221361, upload-time = "2025-09-08T23:22:55.867Z" }, - { url = "https://files.pythonhosted.org/packages/7b/2b/2b6435f76bfeb6bbf055596976da087377ede68df465419d192acf00c437/cffi-2.0.0-cp312-cp312-win32.whl", hash = "sha256:da902562c3e9c550df360bfa53c035b2f241fed6d9aef119048073680ace4a18", size = 172932, upload-time = "2025-09-08T23:22:57.188Z" }, - { url = "https://files.pythonhosted.org/packages/f8/ed/13bd4418627013bec4ed6e54283b1959cf6db888048c7cf4b4c3b5b36002/cffi-2.0.0-cp312-cp312-win_amd64.whl", hash = "sha256:da68248800ad6320861f129cd9c1bf96ca849a2771a59e0344e88681905916f5", size = 183557, upload-time = "2025-09-08T23:22:58.351Z" }, - { url = "https://files.pythonhosted.org/packages/95/31/9f7f93ad2f8eff1dbc1c3656d7ca5bfd8fb52c9d786b4dcf19b2d02217fa/cffi-2.0.0-cp312-cp312-win_arm64.whl", hash = "sha256:4671d9dd5ec934cb9a73e7ee9676f9362aba54f7f34910956b84d727b0d73fb6", size = 177762, upload-time = "2025-09-08T23:22:59.668Z" }, ] [[package]] @@ -336,39 +389,72 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/c5/55/51844dd50c4fc7a33b653bfaba4c2456f06955289ca770a5dbd5fd267374/cfgv-3.4.0-py2.py3-none-any.whl", hash = "sha256:b7265b1f29fd3316bfcd2b330d63d024f2bfd8bcb8b0272f8e19a504856c48f9", size = 7249, upload-time = "2023-08-12T20:38:16.269Z" }, ] +[[package]] +name = "chardet" +version = "7.4.0.post1" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/29/59/133001a6a7549dd34a3c28d1358122b0e68a59f27840efb2102b72eb73cf/chardet-7.4.0.post1-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:45adca296eddec38b803ce352bffcc5fff40576246e74fcd2aa358f9c813ffe0", size = 859436, upload-time = "2026-03-26T19:17:30.2Z" }, + { url = "https://files.pythonhosted.org/packages/c3/00/2eec7b47263524f204b3225d023f7d75e9c06a0a75c06a4c85faf2aec246/chardet-7.4.0.post1-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:f03aef24d91168232db8b01e4d6726676708be503e6aa07e28ab808e6d0fe606", size = 868655, upload-time = "2026-03-26T19:17:31.871Z" }, + { url = "https://files.pythonhosted.org/packages/91/d7/47988d40231b41376f5a66346ef3b322c81091dfd4c0f84df5a1e3bb06b5/chardet-7.4.0.post1-py3-none-any.whl", hash = "sha256:57a62ef50f69bc2fb3a3ea1ffffec6d10f3d2112d3b05d6e3cb15c2c9b55f6cc", size = 624666, upload-time = "2026-03-26T19:17:51.248Z" }, +] + [[package]] name = "charset-normalizer" version = "3.4.3" source = { registry = "https://pypi.org/simple" } sdist = { url = "https://files.pythonhosted.org/packages/83/2d/5fd176ceb9b2fc619e63405525573493ca23441330fcdaee6bef9460e924/charset_normalizer-3.4.3.tar.gz", hash = "sha256:6fce4b8500244f6fcb71465d4a4930d132ba9ab8e71a7859e6a5d59851068d14", size = 122371, upload-time = "2025-08-09T07:57:28.46Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/e9/5e/14c94999e418d9b87682734589404a25854d5f5d0408df68bc15b6ff54bb/charset_normalizer-3.4.3-cp312-cp312-macosx_10_13_universal2.whl", hash = "sha256:e28e334d3ff134e88989d90ba04b47d84382a828c061d0d1027b1b12a62b39b1", size = 205655, upload-time = "2025-08-09T07:56:08.475Z" }, { url = "https://files.pythonhosted.org/packages/7d/a8/c6ec5d389672521f644505a257f50544c074cf5fc292d5390331cd6fc9c3/charset_normalizer-3.4.3-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:0cacf8f7297b0c4fcb74227692ca46b4a5852f8f4f24b3c766dd94a1075c4884", size = 146223, upload-time = "2025-08-09T07:56:09.708Z" }, - { url = "https://files.pythonhosted.org/packages/fc/eb/a2ffb08547f4e1e5415fb69eb7db25932c52a52bed371429648db4d84fb1/charset_normalizer-3.4.3-cp312-cp312-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:c6fd51128a41297f5409deab284fecbe5305ebd7e5a1f959bee1c054622b7018", size = 159366, upload-time = "2025-08-09T07:56:11.326Z" }, - { url = "https://files.pythonhosted.org/packages/82/10/0fd19f20c624b278dddaf83b8464dcddc2456cb4b02bb902a6da126b87a1/charset_normalizer-3.4.3-cp312-cp312-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:3cfb2aad70f2c6debfbcb717f23b7eb55febc0bb23dcffc0f076009da10c6392", size = 157104, upload-time = "2025-08-09T07:56:13.014Z" }, { url = "https://files.pythonhosted.org/packages/16/ab/0233c3231af734f5dfcf0844aa9582d5a1466c985bbed6cedab85af9bfe3/charset_normalizer-3.4.3-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:1606f4a55c0fd363d754049cdf400175ee96c992b1f8018b993941f221221c5f", size = 151830, upload-time = "2025-08-09T07:56:14.428Z" }, { url = "https://files.pythonhosted.org/packages/ae/02/e29e22b4e02839a0e4a06557b1999d0a47db3567e82989b5bb21f3fbbd9f/charset_normalizer-3.4.3-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:027b776c26d38b7f15b26a5da1044f376455fb3766df8fc38563b4efbc515154", size = 148854, upload-time = "2025-08-09T07:56:16.051Z" }, - { url = "https://files.pythonhosted.org/packages/05/6b/e2539a0a4be302b481e8cafb5af8792da8093b486885a1ae4d15d452bcec/charset_normalizer-3.4.3-cp312-cp312-musllinux_1_2_ppc64le.whl", hash = "sha256:42e5088973e56e31e4fa58eb6bd709e42fc03799c11c42929592889a2e54c491", size = 160670, upload-time = "2025-08-09T07:56:17.314Z" }, - { url = "https://files.pythonhosted.org/packages/31/e7/883ee5676a2ef217a40ce0bffcc3d0dfbf9e64cbcfbdf822c52981c3304b/charset_normalizer-3.4.3-cp312-cp312-musllinux_1_2_s390x.whl", hash = "sha256:cc34f233c9e71701040d772aa7490318673aa7164a0efe3172b2981218c26d93", size = 158501, upload-time = "2025-08-09T07:56:18.641Z" }, { url = "https://files.pythonhosted.org/packages/c1/35/6525b21aa0db614cf8b5792d232021dca3df7f90a1944db934efa5d20bb1/charset_normalizer-3.4.3-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:320e8e66157cc4e247d9ddca8e21f427efc7a04bbd0ac8a9faf56583fa543f9f", size = 153173, upload-time = "2025-08-09T07:56:20.289Z" }, - { url = "https://files.pythonhosted.org/packages/50/ee/f4704bad8201de513fdc8aac1cabc87e38c5818c93857140e06e772b5892/charset_normalizer-3.4.3-cp312-cp312-win32.whl", hash = "sha256:fb6fecfd65564f208cbf0fba07f107fb661bcd1a7c389edbced3f7a493f70e37", size = 99822, upload-time = "2025-08-09T07:56:21.551Z" }, - { url = "https://files.pythonhosted.org/packages/39/f5/3b3836ca6064d0992c58c7561c6b6eee1b3892e9665d650c803bd5614522/charset_normalizer-3.4.3-cp312-cp312-win_amd64.whl", hash = "sha256:86df271bf921c2ee3818f0522e9a5b8092ca2ad8b065ece5d7d9d0e9f4849bcc", size = 107543, upload-time = "2025-08-09T07:56:23.115Z" }, { url = "https://files.pythonhosted.org/packages/8a/1f/f041989e93b001bc4e44bb1669ccdcf54d3f00e628229a85b08d330615c5/charset_normalizer-3.4.3-py3-none-any.whl", hash = "sha256:ce571ab16d890d23b5c278547ba694193a45011ff86a9162a71307ed9f86759a", size = 53175, upload-time = "2025-08-09T07:57:26.864Z" }, ] [[package]] -name = "chess" -version = "1.11.2" +name = "chromadb" +version = "1.5.4" source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/93/09/7d04d7581ae3bb8b598017941781bceb7959dd1b13e3ebf7b6a2cd843bc9/chess-1.11.2.tar.gz", hash = "sha256:a8b43e5678fdb3000695bdaa573117ad683761e5ca38e591c4826eba6d25bb39", size = 6131385, upload-time = "2025-02-25T19:10:27.328Z" } +dependencies = [ + { name = "bcrypt", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "build", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "grpcio", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "httpx", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "importlib-resources", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "jsonschema", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "kubernetes", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "mmh3", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "numpy", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "onnxruntime", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "opentelemetry-api", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "opentelemetry-exporter-otlp-proto-grpc", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "opentelemetry-sdk", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "orjson", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "overrides", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "pybase64", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "pydantic", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "pydantic-settings", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "pypika", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "pyyaml", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "rich", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "tenacity", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "tokenizers", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "tqdm", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "typer", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "typing-extensions", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "uvicorn", extra = ["standard"], marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/6a/0e/fe9f016a49dcf4bfa318b2106f3564ffa95baa955725f7ae8e074b2b2980/chromadb-1.5.4.tar.gz", hash = "sha256:d9e3284004b5222e97207f8a69366e684189c7de73d3ce2b5003f9043c8d22dd", size = 2410734, upload-time = "2026-03-08T20:08:25.91Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/1c/76/c747fd05b752ef5f489d01946c77082416141999af57bba484d15eae2586/chromadb-1.5.4-cp39-abi3-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:f37204b826b98f8561244382abf99141bc4b7bf21015dbbcf2280d771afeeb94", size = 20732448, upload-time = "2026-03-08T20:08:13.448Z" }, + { url = "https://files.pythonhosted.org/packages/21/7f/7525519b4eb0d2940d6202b40b21eb9f17df8061fb1b280f8569d0378ae7/chromadb-1.5.4-cp39-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:be11c6d522319ee4fbe1cb8a9a11f6ba9c7f2db625a0e2ea377843bf8e5f99d5", size = 21595920, upload-time = "2026-03-08T20:08:17.741Z" }, +] [[package]] name = "click" version = "8.2.1" source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "colorama", marker = "sys_platform == 'win32'" }, -] sdist = { url = "https://files.pythonhosted.org/packages/60/6c/8ca2efa64cf75a977a0d7fac081354553ebe483345c734fb6b6515d96bbc/click-8.2.1.tar.gz", hash = "sha256:27c491cc05d968d271d5a1db13e3b5a184636d9d930f148c50b038f0d0646202", size = 286342, upload-time = "2025-05-20T23:19:49.832Z" } wheels = [ { url = "https://files.pythonhosted.org/packages/85/32/10bb5764d90a8eee674e9dc6f4db6a0ab47c8c4d0d83c27f7c39ac415a4d/click-8.2.1-py3-none-any.whl", hash = "sha256:61a3265b914e850b85317d0b3109c7f8cd35a670f963866005d6ef1d5175a12b", size = 102215, upload-time = "2025-05-20T23:19:47.796Z" }, @@ -383,6 +469,30 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/7e/e8/64c37fadfc2816a7701fa8a6ed8d87327c7d54eacfbfb6edab14a2f2be75/cloudpickle-3.1.1-py3-none-any.whl", hash = "sha256:c8c5a44295039331ee9dad40ba100a9c7297b6f988e50e87ccdf3765a668350e", size = 20992, upload-time = "2025-01-14T17:02:02.417Z" }, ] +[[package]] +name = "code-env" +version = "0.1.1" +source = { registry = "https://hub.primeintellect.ai/primeintellect/simple/" } +dependencies = [ + { name = "prime-sandboxes", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "verifiers", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, +] +wheels = [ + { url = "https://hub.primeintellect.ai/primeintellect/code-env/@2e4fb479/code_env-0.1.1-py3-none-any.whl", hash = "sha256:42a5c0b9bf6d858ff1f71a779d38e5268fb03f63ed347a736eb6938859303177" }, +] + +[[package]] +name = "color-codeword" +version = "0.1.0" +source = { registry = "https://hub.primeintellect.ai/primeintellect/simple/" } +dependencies = [ + { name = "pillow", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "verifiers", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, +] +wheels = [ + { url = "https://hub.primeintellect.ai/primeintellect/color-codeword/@828f24c5/color_codeword-0.1.0-py3-none-any.whl", hash = "sha256:5bf72f74d6ee001b6be96768dd269e8762617dd3644238bda7623f58a6bc069a" }, +] + [[package]] name = "colorama" version = "0.4.6" @@ -403,39 +513,30 @@ wheels = [ [[package]] name = "compressed-tensors" -version = "0.12.2" +version = "0.15.0.1" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "loguru" }, - { name = "pydantic" }, - { name = "torch" }, - { name = "transformers" }, + { name = "loguru", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "pydantic", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "torch", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "transformers", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/a2/79/4c5c1cd14266f8cf2650bdb940f986ce7fcaeb56aad8cfa9e9afedf14e2f/compressed_tensors-0.12.2.tar.gz", hash = "sha256:5bb40856dd17f128ab73557ecc73799f80db4dd82fab6de875f1e6899b9ea0c4", size = 190409, upload-time = "2025-10-07T14:30:59.302Z" } +sdist = { url = "https://files.pythonhosted.org/packages/41/1b/c3c4a98ec5f2727656336f07a0c35862195c310d8eb0b2fa5b4be6848680/compressed_tensors-0.15.0.1.tar.gz", hash = "sha256:a8e93054e8a5ec49c980b09ed36c4c1249b4a8ee167920a8e461c4da26e78d99", size = 229412, upload-time = "2026-04-10T14:23:54.708Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/f0/c0/1695b87d369e6652ec0d650912e02eca2151c5e9c29244f94d2afccfe970/compressed_tensors-0.12.2-py3-none-any.whl", hash = "sha256:e554ea761710ca2b0c0ea49276a4ef8e08658624f1591e6a7368817106b48fbe", size = 183049, upload-time = "2025-10-07T14:30:56.523Z" }, + { url = "https://files.pythonhosted.org/packages/a8/52/93833dc1610e017ac5b7dcd59b8304d8ef67d1114c2d124e728a2cbbea12/compressed_tensors-0.15.0.1-py3-none-any.whl", hash = "sha256:e1b1f322e82e475715e242bad46925a304ea8e5c98b5055a15b8eb22fb6bfea9", size = 194260, upload-time = "2026-04-10T14:23:53.098Z" }, ] [[package]] -name = "contourpy" -version = "1.3.3" +name = "connect-python" +version = "0.8.1" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "numpy" }, + { name = "protobuf", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "pyqwest", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/58/01/1253e6698a07380cd31a736d248a3f2a50a7c88779a1813da27503cadc2a/contourpy-1.3.3.tar.gz", hash = "sha256:083e12155b210502d0bca491432bb04d56dc3432f95a979b429f2848c3dbe880", size = 13466174, upload-time = "2025-07-26T12:03:12.549Z" } +sdist = { url = "https://files.pythonhosted.org/packages/8f/8d/e04954dacc3c32c2f858a115b2b989ed6cccefc3835fb440bb12dcb468c2/connect_python-0.8.1.tar.gz", hash = "sha256:0d36c9cfe050661f8f167afcfff8df622805a4f348961d0d8fab630c96f0bdab", size = 39292, upload-time = "2026-01-27T05:00:02.751Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/be/45/adfee365d9ea3d853550b2e735f9d66366701c65db7855cd07621732ccfc/contourpy-1.3.3-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:b08a32ea2f8e42cf1d4be3169a98dd4be32bafe4f22b6c4cb4ba810fa9e5d2cb", size = 293419, upload-time = "2025-07-26T12:01:21.16Z" }, - { url = "https://files.pythonhosted.org/packages/53/3e/405b59cfa13021a56bba395a6b3aca8cec012b45bf177b0eaf7a202cde2c/contourpy-1.3.3-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:556dba8fb6f5d8742f2923fe9457dbdd51e1049c4a43fd3986a0b14a1d815fc6", size = 273979, upload-time = "2025-07-26T12:01:22.448Z" }, - { url = "https://files.pythonhosted.org/packages/d4/1c/a12359b9b2ca3a845e8f7f9ac08bdf776114eb931392fcad91743e2ea17b/contourpy-1.3.3-cp312-cp312-manylinux_2_26_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:92d9abc807cf7d0e047b95ca5d957cf4792fcd04e920ca70d48add15c1a90ea7", size = 332653, upload-time = "2025-07-26T12:01:24.155Z" }, - { url = "https://files.pythonhosted.org/packages/63/12/897aeebfb475b7748ea67b61e045accdfcf0d971f8a588b67108ed7f5512/contourpy-1.3.3-cp312-cp312-manylinux_2_26_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:b2e8faa0ed68cb29af51edd8e24798bb661eac3bd9f65420c1887b6ca89987c8", size = 379536, upload-time = "2025-07-26T12:01:25.91Z" }, - { url = "https://files.pythonhosted.org/packages/43/8a/a8c584b82deb248930ce069e71576fc09bd7174bbd35183b7943fb1064fd/contourpy-1.3.3-cp312-cp312-manylinux_2_26_s390x.manylinux_2_28_s390x.whl", hash = "sha256:626d60935cf668e70a5ce6ff184fd713e9683fb458898e4249b63be9e28286ea", size = 384397, upload-time = "2025-07-26T12:01:27.152Z" }, - { url = "https://files.pythonhosted.org/packages/cc/8f/ec6289987824b29529d0dfda0d74a07cec60e54b9c92f3c9da4c0ac732de/contourpy-1.3.3-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:4d00e655fcef08aba35ec9610536bfe90267d7ab5ba944f7032549c55a146da1", size = 362601, upload-time = "2025-07-26T12:01:28.808Z" }, - { url = "https://files.pythonhosted.org/packages/05/0a/a3fe3be3ee2dceb3e615ebb4df97ae6f3828aa915d3e10549ce016302bd1/contourpy-1.3.3-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:451e71b5a7d597379ef572de31eeb909a87246974d960049a9848c3bc6c41bf7", size = 1331288, upload-time = "2025-07-26T12:01:31.198Z" }, - { url = "https://files.pythonhosted.org/packages/33/1d/acad9bd4e97f13f3e2b18a3977fe1b4a37ecf3d38d815333980c6c72e963/contourpy-1.3.3-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:459c1f020cd59fcfe6650180678a9993932d80d44ccde1fa1868977438f0b411", size = 1403386, upload-time = "2025-07-26T12:01:33.947Z" }, - { url = "https://files.pythonhosted.org/packages/cf/8f/5847f44a7fddf859704217a99a23a4f6417b10e5ab1256a179264561540e/contourpy-1.3.3-cp312-cp312-win32.whl", hash = "sha256:023b44101dfe49d7d53932be418477dba359649246075c996866106da069af69", size = 185018, upload-time = "2025-07-26T12:01:35.64Z" }, - { url = "https://files.pythonhosted.org/packages/19/e8/6026ed58a64563186a9ee3f29f41261fd1828f527dd93d33b60feca63352/contourpy-1.3.3-cp312-cp312-win_amd64.whl", hash = "sha256:8153b8bfc11e1e4d75bcb0bff1db232f9e10b274e0929de9d608027e0d34ff8b", size = 226567, upload-time = "2025-07-26T12:01:36.804Z" }, - { url = "https://files.pythonhosted.org/packages/d1/e2/f05240d2c39a1ed228d8328a78b6f44cd695f7ef47beb3e684cf93604f86/contourpy-1.3.3-cp312-cp312-win_arm64.whl", hash = "sha256:07ce5ed73ecdc4a03ffe3e1b3e3c1166db35ae7584be76f65dbbe28a7791b0cc", size = 193655, upload-time = "2025-07-26T12:01:37.999Z" }, + { url = "https://files.pythonhosted.org/packages/7d/e4/e36a6d0450493e2b940086ebe94394381803b2b09c97e8f8e34230bcaa99/connect_python-0.8.1-py3-none-any.whl", hash = "sha256:bfe31fb3d90c2a6715fc2996fa4bd5c98dfdccd81e47ac232c9d03aadd666be1", size = 54902, upload-time = "2026-01-27T05:00:01.18Z" }, ] [[package]] @@ -443,53 +544,38 @@ name = "cryptography" version = "46.0.3" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "cffi", marker = "platform_python_implementation != 'PyPy'" }, + { name = "cffi", marker = "(platform_machine == 'aarch64' and platform_python_implementation != 'PyPy' and sys_platform == 'linux') or (platform_machine == 'x86_64' and platform_python_implementation != 'PyPy' and sys_platform == 'linux')" }, ] sdist = { url = "https://files.pythonhosted.org/packages/9f/33/c00162f49c0e2fe8064a62cb92b93e50c74a72bc370ab92f86112b33ff62/cryptography-46.0.3.tar.gz", hash = "sha256:a8b17438104fed022ce745b362294d9ce35b4c2e45c1d958ad4a4b019285f4a1", size = 749258, upload-time = "2025-10-15T23:18:31.74Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/1d/42/9c391dd801d6cf0d561b5890549d4b27bafcc53b39c31a817e69d87c625b/cryptography-46.0.3-cp311-abi3-macosx_10_9_universal2.whl", hash = "sha256:109d4ddfadf17e8e7779c39f9b18111a09efb969a301a31e987416a0191ed93a", size = 7225004, upload-time = "2025-10-15T23:16:52.239Z" }, { url = "https://files.pythonhosted.org/packages/1c/67/38769ca6b65f07461eb200e85fc1639b438bdc667be02cf7f2cd6a64601c/cryptography-46.0.3-cp311-abi3-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:09859af8466b69bc3c27bdf4f5d84a665e0f7ab5088412e9e2ec49758eca5cbc", size = 4296667, upload-time = "2025-10-15T23:16:54.369Z" }, { url = "https://files.pythonhosted.org/packages/5c/49/498c86566a1d80e978b42f0d702795f69887005548c041636df6ae1ca64c/cryptography-46.0.3-cp311-abi3-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:01ca9ff2885f3acc98c29f1860552e37f6d7c7d013d7334ff2a9de43a449315d", size = 4450807, upload-time = "2025-10-15T23:16:56.414Z" }, { url = "https://files.pythonhosted.org/packages/4b/0a/863a3604112174c8624a2ac3c038662d9e59970c7f926acdcfaed8d61142/cryptography-46.0.3-cp311-abi3-manylinux_2_28_aarch64.whl", hash = "sha256:6eae65d4c3d33da080cff9c4ab1f711b15c1d9760809dad6ea763f3812d254cb", size = 4299615, upload-time = "2025-10-15T23:16:58.442Z" }, - { url = "https://files.pythonhosted.org/packages/64/02/b73a533f6b64a69f3cd3872acb6ebc12aef924d8d103133bb3ea750dc703/cryptography-46.0.3-cp311-abi3-manylinux_2_28_armv7l.manylinux_2_31_armv7l.whl", hash = "sha256:e5bf0ed4490068a2e72ac03d786693adeb909981cc596425d09032d372bcc849", size = 4016800, upload-time = "2025-10-15T23:17:00.378Z" }, - { url = "https://files.pythonhosted.org/packages/25/d5/16e41afbfa450cde85a3b7ec599bebefaef16b5c6ba4ec49a3532336ed72/cryptography-46.0.3-cp311-abi3-manylinux_2_28_ppc64le.whl", hash = "sha256:5ecfccd2329e37e9b7112a888e76d9feca2347f12f37918facbb893d7bb88ee8", size = 4984707, upload-time = "2025-10-15T23:17:01.98Z" }, { url = "https://files.pythonhosted.org/packages/c9/56/e7e69b427c3878352c2fb9b450bd0e19ed552753491d39d7d0a2f5226d41/cryptography-46.0.3-cp311-abi3-manylinux_2_28_x86_64.whl", hash = "sha256:a2c0cd47381a3229c403062f764160d57d4d175e022c1df84e168c6251a22eec", size = 4482541, upload-time = "2025-10-15T23:17:04.078Z" }, { url = "https://files.pythonhosted.org/packages/78/f6/50736d40d97e8483172f1bb6e698895b92a223dba513b0ca6f06b2365339/cryptography-46.0.3-cp311-abi3-manylinux_2_34_aarch64.whl", hash = "sha256:549e234ff32571b1f4076ac269fcce7a808d3bf98b76c8dd560e42dbc66d7d91", size = 4299464, upload-time = "2025-10-15T23:17:05.483Z" }, - { url = "https://files.pythonhosted.org/packages/00/de/d8e26b1a855f19d9994a19c702fa2e93b0456beccbcfe437eda00e0701f2/cryptography-46.0.3-cp311-abi3-manylinux_2_34_ppc64le.whl", hash = "sha256:c0a7bb1a68a5d3471880e264621346c48665b3bf1c3759d682fc0864c540bd9e", size = 4950838, upload-time = "2025-10-15T23:17:07.425Z" }, { url = "https://files.pythonhosted.org/packages/8f/29/798fc4ec461a1c9e9f735f2fc58741b0daae30688f41b2497dcbc9ed1355/cryptography-46.0.3-cp311-abi3-manylinux_2_34_x86_64.whl", hash = "sha256:10b01676fc208c3e6feeb25a8b83d81767e8059e1fe86e1dc62d10a3018fa926", size = 4481596, upload-time = "2025-10-15T23:17:09.343Z" }, { url = "https://files.pythonhosted.org/packages/15/8d/03cd48b20a573adfff7652b76271078e3045b9f49387920e7f1f631d125e/cryptography-46.0.3-cp311-abi3-musllinux_1_2_aarch64.whl", hash = "sha256:0abf1ffd6e57c67e92af68330d05760b7b7efb243aab8377e583284dbab72c71", size = 4426782, upload-time = "2025-10-15T23:17:11.22Z" }, { url = "https://files.pythonhosted.org/packages/fa/b1/ebacbfe53317d55cf33165bda24c86523497a6881f339f9aae5c2e13e57b/cryptography-46.0.3-cp311-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:a04bee9ab6a4da801eb9b51f1b708a1b5b5c9eb48c03f74198464c66f0d344ac", size = 4698381, upload-time = "2025-10-15T23:17:12.829Z" }, - { url = "https://files.pythonhosted.org/packages/96/92/8a6a9525893325fc057a01f654d7efc2c64b9de90413adcf605a85744ff4/cryptography-46.0.3-cp311-abi3-win32.whl", hash = "sha256:f260d0d41e9b4da1ed1e0f1ce571f97fe370b152ab18778e9e8f67d6af432018", size = 3055988, upload-time = "2025-10-15T23:17:14.65Z" }, - { url = "https://files.pythonhosted.org/packages/7e/bf/80fbf45253ea585a1e492a6a17efcb93467701fa79e71550a430c5e60df0/cryptography-46.0.3-cp311-abi3-win_amd64.whl", hash = "sha256:a9a3008438615669153eb86b26b61e09993921ebdd75385ddd748702c5adfddb", size = 3514451, upload-time = "2025-10-15T23:17:16.142Z" }, - { url = "https://files.pythonhosted.org/packages/2e/af/9b302da4c87b0beb9db4e756386a7c6c5b8003cd0e742277888d352ae91d/cryptography-46.0.3-cp311-abi3-win_arm64.whl", hash = "sha256:5d7f93296ee28f68447397bf5198428c9aeeab45705a55d53a6343455dcb2c3c", size = 2928007, upload-time = "2025-10-15T23:17:18.04Z" }, - { url = "https://files.pythonhosted.org/packages/fd/23/45fe7f376a7df8daf6da3556603b36f53475a99ce4faacb6ba2cf3d82021/cryptography-46.0.3-cp38-abi3-macosx_10_9_universal2.whl", hash = "sha256:cb3d760a6117f621261d662bccc8ef5bc32ca673e037c83fbe565324f5c46936", size = 7218248, upload-time = "2025-10-15T23:17:46.294Z" }, { url = "https://files.pythonhosted.org/packages/27/32/b68d27471372737054cbd34c84981f9edbc24fe67ca225d389799614e27f/cryptography-46.0.3-cp38-abi3-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:4b7387121ac7d15e550f5cb4a43aef2559ed759c35df7336c402bb8275ac9683", size = 4294089, upload-time = "2025-10-15T23:17:48.269Z" }, { url = "https://files.pythonhosted.org/packages/26/42/fa8389d4478368743e24e61eea78846a0006caffaf72ea24a15159215a14/cryptography-46.0.3-cp38-abi3-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:15ab9b093e8f09daab0f2159bb7e47532596075139dd74365da52ecc9cb46c5d", size = 4440029, upload-time = "2025-10-15T23:17:49.837Z" }, { url = "https://files.pythonhosted.org/packages/5f/eb/f483db0ec5ac040824f269e93dd2bd8a21ecd1027e77ad7bdf6914f2fd80/cryptography-46.0.3-cp38-abi3-manylinux_2_28_aarch64.whl", hash = "sha256:46acf53b40ea38f9c6c229599a4a13f0d46a6c3fa9ef19fc1a124d62e338dfa0", size = 4297222, upload-time = "2025-10-15T23:17:51.357Z" }, - { url = "https://files.pythonhosted.org/packages/fd/cf/da9502c4e1912cb1da3807ea3618a6829bee8207456fbbeebc361ec38ba3/cryptography-46.0.3-cp38-abi3-manylinux_2_28_armv7l.manylinux_2_31_armv7l.whl", hash = "sha256:10ca84c4668d066a9878890047f03546f3ae0a6b8b39b697457b7757aaf18dbc", size = 4012280, upload-time = "2025-10-15T23:17:52.964Z" }, - { url = "https://files.pythonhosted.org/packages/6b/8f/9adb86b93330e0df8b3dcf03eae67c33ba89958fc2e03862ef1ac2b42465/cryptography-46.0.3-cp38-abi3-manylinux_2_28_ppc64le.whl", hash = "sha256:36e627112085bb3b81b19fed209c05ce2a52ee8b15d161b7c643a7d5a88491f3", size = 4978958, upload-time = "2025-10-15T23:17:54.965Z" }, { url = "https://files.pythonhosted.org/packages/d1/a0/5fa77988289c34bdb9f913f5606ecc9ada1adb5ae870bd0d1054a7021cc4/cryptography-46.0.3-cp38-abi3-manylinux_2_28_x86_64.whl", hash = "sha256:1000713389b75c449a6e979ffc7dcc8ac90b437048766cef052d4d30b8220971", size = 4473714, upload-time = "2025-10-15T23:17:56.754Z" }, { url = "https://files.pythonhosted.org/packages/14/e5/fc82d72a58d41c393697aa18c9abe5ae1214ff6f2a5c18ac470f92777895/cryptography-46.0.3-cp38-abi3-manylinux_2_34_aarch64.whl", hash = "sha256:b02cf04496f6576afffef5ddd04a0cb7d49cf6be16a9059d793a30b035f6b6ac", size = 4296970, upload-time = "2025-10-15T23:17:58.588Z" }, - { url = "https://files.pythonhosted.org/packages/78/06/5663ed35438d0b09056973994f1aec467492b33bd31da36e468b01ec1097/cryptography-46.0.3-cp38-abi3-manylinux_2_34_ppc64le.whl", hash = "sha256:71e842ec9bc7abf543b47cf86b9a743baa95f4677d22baa4c7d5c69e49e9bc04", size = 4940236, upload-time = "2025-10-15T23:18:00.897Z" }, { url = "https://files.pythonhosted.org/packages/fc/59/873633f3f2dcd8a053b8dd1d38f783043b5fce589c0f6988bf55ef57e43e/cryptography-46.0.3-cp38-abi3-manylinux_2_34_x86_64.whl", hash = "sha256:402b58fc32614f00980b66d6e56a5b4118e6cb362ae8f3fda141ba4689bd4506", size = 4472642, upload-time = "2025-10-15T23:18:02.749Z" }, { url = "https://files.pythonhosted.org/packages/3d/39/8e71f3930e40f6877737d6f69248cf74d4e34b886a3967d32f919cc50d3b/cryptography-46.0.3-cp38-abi3-musllinux_1_2_aarch64.whl", hash = "sha256:ef639cb3372f69ec44915fafcd6698b6cc78fbe0c2ea41be867f6ed612811963", size = 4423126, upload-time = "2025-10-15T23:18:04.85Z" }, { url = "https://files.pythonhosted.org/packages/cd/c7/f65027c2810e14c3e7268353b1681932b87e5a48e65505d8cc17c99e36ae/cryptography-46.0.3-cp38-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:3b51b8ca4f1c6453d8829e1eb7299499ca7f313900dd4d89a24b8b87c0a780d4", size = 4686573, upload-time = "2025-10-15T23:18:06.908Z" }, - { url = "https://files.pythonhosted.org/packages/0a/6e/1c8331ddf91ca4730ab3086a0f1be19c65510a33b5a441cb334e7a2d2560/cryptography-46.0.3-cp38-abi3-win32.whl", hash = "sha256:6276eb85ef938dc035d59b87c8a7dc559a232f954962520137529d77b18ff1df", size = 3036695, upload-time = "2025-10-15T23:18:08.672Z" }, - { url = "https://files.pythonhosted.org/packages/90/45/b0d691df20633eff80955a0fc7695ff9051ffce8b69741444bd9ed7bd0db/cryptography-46.0.3-cp38-abi3-win_amd64.whl", hash = "sha256:416260257577718c05135c55958b674000baef9a1c7d9e8f306ec60d71db850f", size = 3501720, upload-time = "2025-10-15T23:18:10.632Z" }, - { url = "https://files.pythonhosted.org/packages/e8/cb/2da4cc83f5edb9c3257d09e1e7ab7b23f049c7962cae8d842bbef0a9cec9/cryptography-46.0.3-cp38-abi3-win_arm64.whl", hash = "sha256:d89c3468de4cdc4f08a57e214384d0471911a3830fcdaf7a8cc587e42a866372", size = 2918740, upload-time = "2025-10-15T23:18:12.277Z" }, ] [[package]] name = "cuda-bindings" -version = "13.1.1" +version = "12.9.4" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "cuda-pathfinder" }, + { name = "cuda-pathfinder", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, ] wheels = [ - { url = "https://files.pythonhosted.org/packages/53/3d/c8ed9d169843091f3f0d6b8218e826fd59520a37e0434c204feada597988/cuda_bindings-13.1.1-cp312-cp312-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:1e75ad0cb863330df784236d289612d71ca855c013d19ae00e5693574abd6915", size = 15530160, upload-time = "2025-12-09T22:05:55.386Z" }, - { url = "https://files.pythonhosted.org/packages/4a/8e/368295623ee43fba622909d780fbb6863efc1638dff55f67a0f04eac6470/cuda_bindings-13.1.1-cp312-cp312-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:25785d1a3cdcd98f151240fd5efd025609319a6720a217dee2a929241749d488", size = 16110386, upload-time = "2025-12-09T22:05:57.71Z" }, - { url = "https://files.pythonhosted.org/packages/60/1f/ecc4701ade3e85f091c625a920574527b9daf7fb354189fbfbc5516af6cd/cuda_bindings-13.1.1-cp312-cp312-win_amd64.whl", hash = "sha256:ccde9c95c0e953b31fe7731bb08da9d0a34b1770498df9a3c156fdfdbe3951ad", size = 15250028, upload-time = "2025-12-09T22:06:00.346Z" }, + { url = "https://files.pythonhosted.org/packages/0c/c2/65bfd79292b8ff18be4dd7f7442cea37bcbc1a228c1886f1dea515c45b67/cuda_bindings-12.9.4-cp312-cp312-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:694ba35023846625ef471257e6b5a4bc8af690f961d197d77d34b1d1db393f56", size = 11760260, upload-time = "2025-10-21T14:51:40.79Z" }, + { url = "https://files.pythonhosted.org/packages/a9/c1/dabe88f52c3e3760d861401bb994df08f672ec893b8f7592dc91626adcf3/cuda_bindings-12.9.4-cp312-cp312-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:fda147a344e8eaeca0c6ff113d2851ffca8f7dfc0a6c932374ee5c47caa649c8", size = 12151019, upload-time = "2025-10-21T14:51:43.167Z" }, ] [[package]] @@ -502,60 +588,68 @@ wheels = [ [[package]] name = "cuda-python" -version = "13.1.1" +version = "12.9.4" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "cuda-bindings" }, - { name = "cuda-pathfinder" }, + { name = "cuda-bindings", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, ] wheels = [ - { url = "https://files.pythonhosted.org/packages/cd/08/b5e3b9822662d72d540d830531e3ab6a7cabbda3dd56175696aabccfeb76/cuda_python-13.1.1-py3-none-any.whl", hash = "sha256:944cc4fe6482673d28dd545797a28840945a1668739328fa2ad1e9be4f7050d9", size = 8038, upload-time = "2025-12-09T22:13:10.719Z" }, + { url = "https://files.pythonhosted.org/packages/af/f3/6b032a554019cfb3447e671798c1bd3e79b5f1af20d10253f56cea269ef2/cuda_python-12.9.4-py3-none-any.whl", hash = "sha256:d2cacea882a69863f1e7d27ee71d75f0684f4c76910aff839067e4f89c902279", size = 7594, upload-time = "2025-10-21T14:55:12.846Z" }, ] [[package]] -name = "cupy-cuda12x" -version = "13.6.0" +name = "cuda-tile" +version = "1.3.0" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "fastrlock" }, - { name = "numpy" }, + { name = "typing-extensions", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, ] wheels = [ - { url = "https://files.pythonhosted.org/packages/12/c5/7e7fc4816d0de0154e5d9053242c3a08a0ca8b43ee656a6f7b3b95055a7b/cupy_cuda12x-13.6.0-cp312-cp312-manylinux2014_aarch64.whl", hash = "sha256:a6970ceefe40f9acbede41d7fe17416bd277b1bd2093adcde457b23b578c5a59", size = 127334633, upload-time = "2025-08-18T08:24:43.065Z" }, - { url = "https://files.pythonhosted.org/packages/e0/95/d7e1295141e7d530674a3cc567e13ed0eb6b81524cb122d797ed996b5bea/cupy_cuda12x-13.6.0-cp312-cp312-manylinux2014_x86_64.whl", hash = "sha256:79b0cacb5e8b190ef409f9e03f06ac8de1b021b0c0dda47674d446f5557e0eb1", size = 112886268, upload-time = "2025-08-18T08:24:49.294Z" }, - { url = "https://files.pythonhosted.org/packages/ae/8c/14555b63fd78cfac7b88af0094cea0a3cb845d243661ec7da69f7b3ea0de/cupy_cuda12x-13.6.0-cp312-cp312-win_amd64.whl", hash = "sha256:ca06fede7b8b83ca9ad80062544ef2e5bb8d4762d1c4fc3ac8349376de9c8a5e", size = 89785108, upload-time = "2025-08-18T08:24:54.527Z" }, + { url = "https://files.pythonhosted.org/packages/f3/49/4592bc94ca05a07c7947ea114fd12734c8497f2daffee9faa79a03e39fb5/cuda_tile-1.3.0-cp312-cp312-manylinux2014_aarch64.whl", hash = "sha256:375316b64c51ee7cfadb2f170a30c1547bc41eb39f1e233a6556713857d2e81f", size = 245744, upload-time = "2026-04-20T15:52:09.621Z" }, + { url = "https://files.pythonhosted.org/packages/40/76/84cb68be463c827bf79da9fa0aa5140838de6455ef6f438bbe0ffa75d378/cuda_tile-1.3.0-cp312-cp312-manylinux2014_x86_64.whl", hash = "sha256:e4865acbff1172aaee304bf9c550586088d8b4545a384423597a590899386709", size = 247301, upload-time = "2026-04-20T15:51:04.042Z" }, ] [[package]] -name = "cycler" -version = "0.12.1" +name = "cuda-toolkit" +version = "12.8.1" source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/a9/95/a3dbbb5028f35eafb79008e7522a75244477d2838f38cbb722248dabc2a8/cycler-0.12.1.tar.gz", hash = "sha256:88bb128f02ba341da8ef447245a9e138fae777f6a23943da4540077d3601eb1c", size = 7615, upload-time = "2023-10-07T05:32:18.335Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/e7/05/c19819d5e3d95294a6f5947fb9b9629efb316b96de511b418c53d245aae6/cycler-0.12.1-py3-none-any.whl", hash = "sha256:85cef7cff222d8644161529808465972e51340599459b8ac3ccbac5a854e0d30", size = 8321, upload-time = "2023-10-07T05:32:16.783Z" }, + { url = "https://files.pythonhosted.org/packages/d4/c8/7dce3a0b15b42a3b58e7d96eb22a687d3bf2c44e01d149a6874629cd9938/cuda_toolkit-12.8.1-py2.py3-none-any.whl", hash = "sha256:adc7906af4ecbf9a352f9dca5734eceb21daec281ccfcf5675e1d2f724fc2cba", size = 2283, upload-time = "2025-08-13T02:03:07.842Z" }, ] -[[package]] -name = "cydifflib" -version = "1.2.0" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/16/3d/37bd6a166e657259470ec12ad4c6fe8f544e67529c9b046f6450d6d890ed/cydifflib-1.2.0.tar.gz", hash = "sha256:b8fb1bd1a1ac4360aacd9210baed6e342a1e1a3972e033b0aeaad7f6c536d96d", size = 340192, upload-time = "2025-04-11T13:33:15.338Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/1b/ce/608f6814fc6f1f35c4848a0f0a3c182250a1d9002c81a39a6dbe3bcd954c/cydifflib-1.2.0-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:63aaffca90b20fff61e934507d9fe266a3fc7ee9675c4e33ded2d65e131976cc", size = 225861, upload-time = "2025-04-11T13:31:52.851Z" }, - { url = "https://files.pythonhosted.org/packages/91/15/45d5ec8c0db329d729a3705ffed709ca9e2f1d461ad98b5ba247e98f79c2/cydifflib-1.2.0-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:7aab77aa93fa765f7208e0eb4b77686f0df69fc4bc640c113d0cdd3bd6cb3ca7", size = 207121, upload-time = "2025-04-11T13:31:53.994Z" }, - { url = "https://files.pythonhosted.org/packages/7e/ab/8611841c72316fefd3cc85caeac9b3f56a2abed2b203ed74fd815a3188f5/cydifflib-1.2.0-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:2b60d81c2c6bab138a5ebf79d4a83239f3d82007d1afc34f22a128e676765d01", size = 242467, upload-time = "2025-04-11T13:31:55.513Z" }, - { url = "https://files.pythonhosted.org/packages/4b/76/6ffdf1cd92f36d4e049e3e948dc9f03f2faf2f6c4a2273f89de064a1dba7/cydifflib-1.2.0-cp312-cp312-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:ff63069684a81047efb32ecefe7d526efac9719527da849e772f418642c4c3cd", size = 272096, upload-time = "2025-04-11T13:31:56.676Z" }, - { url = "https://files.pythonhosted.org/packages/ab/4d/1a0ee21b52364ef2d29573c3eaee7a93acc15e379a1e91c1632b0d566d68/cydifflib-1.2.0-cp312-cp312-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:34eb03bbb7c48d844ae132a6b5d50150283c151aeae14e549df08ae574cafc4e", size = 239453, upload-time = "2025-04-11T13:31:57.81Z" }, - { url = "https://files.pythonhosted.org/packages/b0/64/f1f180e6c00ced87cdcd7be8f6bee79ae32e6ac3598ea631f19a4ddb0dc4/cydifflib-1.2.0-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:81126e2c2d2fecd926e25cc117d5c6bc71ce4ceb570c54ccd33aff790dd9e9be", size = 260592, upload-time = "2025-04-11T13:31:59.316Z" }, - { url = "https://files.pythonhosted.org/packages/d7/98/76e5be3e2eea1e1af24cca647f23c9d9064eca8cc2272938835041db7ecf/cydifflib-1.2.0-cp312-cp312-manylinux_2_5_i686.manylinux1_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:be8566b7cebfc10ee855c11abfdded6016a804cae9dda2e3ad7fbd5bcacaf485", size = 263467, upload-time = "2025-04-11T13:32:00.431Z" }, - { url = "https://files.pythonhosted.org/packages/e5/65/55f1b211345950050fa72c3d8ca44aa61afc77a9408bd24dcb3dbeb9fa13/cydifflib-1.2.0-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:67050f67e8731c5583f5e09e6672eacd00a66743a745e498abda7aea19cdcea6", size = 1231399, upload-time = "2025-04-11T13:32:01.564Z" }, - { url = "https://files.pythonhosted.org/packages/38/e6/4dfae8528e9b1da0193dc8778241e693123758a189e4757e31eb75872477/cydifflib-1.2.0-cp312-cp312-musllinux_1_2_i686.whl", hash = "sha256:12505c8e1f23ec6c128dded1de1f0b36f47db5a2b38f9b4f5cb0eb0c6a56170a", size = 1417724, upload-time = "2025-04-11T13:32:03.065Z" }, - { url = "https://files.pythonhosted.org/packages/52/81/28545f382274480b1e8265273de609330fa549a4b131f5c0dc4ad360495c/cydifflib-1.2.0-cp312-cp312-musllinux_1_2_ppc64le.whl", hash = "sha256:3e0233af573d722613f9a3a7d7d72dfe50e4a49da02ff37ece1cd9bffafe8e86", size = 1340796, upload-time = "2025-04-11T13:32:04.796Z" }, - { url = "https://files.pythonhosted.org/packages/27/da/b93c5adbf8d85bd750d290d0e31dcc360a25072723e10fdc568f9cba539b/cydifflib-1.2.0-cp312-cp312-musllinux_1_2_s390x.whl", hash = "sha256:66940e09be6440cfd1eb2bec0bb7f3b811002aa40eb0844def4932d5ccbe60cb", size = 1463537, upload-time = "2025-04-11T13:32:06.187Z" }, - { url = "https://files.pythonhosted.org/packages/28/38/5a1f5991248a17a19a7863e01fc991eef0fd42ea3e9626751ea866158f8b/cydifflib-1.2.0-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:6b512dc94cfcc6b5d7084e5fb03d3fe619b5fb948e79f96b728763a90b49869c", size = 1312173, upload-time = "2025-04-11T13:32:07.87Z" }, - { url = "https://files.pythonhosted.org/packages/e8/f8/28a41bd6623a1f67f4d84d788844b23b71515ded97a5859a4e9f6a07e6c3/cydifflib-1.2.0-cp312-cp312-win32.whl", hash = "sha256:c968f90ab7b77ad5013eecdc4eb67db76a1277626e533e264403266224f7f3ee", size = 177525, upload-time = "2025-04-11T13:32:09.147Z" }, - { url = "https://files.pythonhosted.org/packages/26/e2/d10d0a0ef37ec82b95cee6e9fdd400808677bcf6af3f0d9d2a748481661d/cydifflib-1.2.0-cp312-cp312-win_amd64.whl", hash = "sha256:7c314f5bd727a57154e269741b9441082434f08c069929aa0a98625775cf3470", size = 203363, upload-time = "2025-04-11T13:32:10.766Z" }, - { url = "https://files.pythonhosted.org/packages/cc/e1/24bd1590fc637bdac966c7530a1d4dd1a54fefc2674d6177f8384b90ad16/cydifflib-1.2.0-cp312-cp312-win_arm64.whl", hash = "sha256:e7a7f53d89e18b746f9b309bce1372274fffdf25c4dc3de2d7238be17cf4c673", size = 167227, upload-time = "2025-04-11T13:32:11.785Z" }, +[package.optional-dependencies] +cublas = [ + { name = "nvidia-cublas-cu12", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, +] +cudart = [ + { name = "nvidia-cuda-runtime-cu12", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, +] +cufft = [ + { name = "nvidia-cufft-cu12", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, +] +cufile = [ + { name = "nvidia-cufile-cu12", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, +] +cupti = [ + { name = "nvidia-cuda-cupti-cu12", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, +] +curand = [ + { name = "nvidia-curand-cu12", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, +] +cusolver = [ + { name = "nvidia-cusolver-cu12", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, +] +cusparse = [ + { name = "nvidia-cusparse-cu12", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, +] +nvjitlink = [ + { name = "nvidia-nvjitlink-cu12", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, +] +nvrtc = [ + { name = "nvidia-cuda-nvrtc-cu12", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, +] +nvtx = [ + { name = "nvidia-nvtx-cu12", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, ] [[package]] @@ -563,19 +657,19 @@ name = "datasets" version = "4.0.0" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "dill" }, - { name = "filelock" }, - { name = "fsspec", extra = ["http"] }, - { name = "huggingface-hub" }, - { name = "multiprocess" }, - { name = "numpy" }, - { name = "packaging" }, - { name = "pandas" }, - { name = "pyarrow" }, - { name = "pyyaml" }, - { name = "requests" }, - { name = "tqdm" }, - { name = "xxhash" }, + { name = "dill", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "filelock", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "fsspec", extra = ["http"], marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "huggingface-hub", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "multiprocess", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "numpy", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "packaging", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "pandas", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "pyarrow", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "pyyaml", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "requests", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "tqdm", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "xxhash", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, ] sdist = { url = "https://files.pythonhosted.org/packages/e3/9d/348ed92110ba5f9b70b51ca1078d4809767a835aa2b7ce7e74ad2b98323d/datasets-4.0.0.tar.gz", hash = "sha256:9657e7140a9050db13443ba21cb5de185af8af944479b00e7ff1e00a61c8dbf1", size = 569566, upload-time = "2025-07-09T14:35:52.431Z" } wheels = [ @@ -588,10 +682,7 @@ version = "1.8.16" source = { registry = "https://pypi.org/simple" } sdist = { url = "https://files.pythonhosted.org/packages/ca/d4/722d0bcc7986172ac2ef3c979ad56a1030e3afd44ced136d45f8142b1f4a/debugpy-1.8.16.tar.gz", hash = "sha256:31e69a1feb1cf6b51efbed3f6c9b0ef03bc46ff050679c4be7ea6d2e23540870", size = 1643809, upload-time = "2025-08-06T18:00:02.647Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/61/fb/0387c0e108d842c902801bc65ccc53e5b91d8c169702a9bbf4f7efcedf0c/debugpy-1.8.16-cp312-cp312-macosx_14_0_universal2.whl", hash = "sha256:b202e2843e32e80b3b584bcebfe0e65e0392920dc70df11b2bfe1afcb7a085e4", size = 2511822, upload-time = "2025-08-06T18:00:18.526Z" }, { url = "https://files.pythonhosted.org/packages/37/44/19e02745cae22bf96440141f94e15a69a1afaa3a64ddfc38004668fcdebf/debugpy-1.8.16-cp312-cp312-manylinux_2_5_x86_64.manylinux1_x86_64.manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:64473c4a306ba11a99fe0bb14622ba4fbd943eb004847d9b69b107bde45aa9ea", size = 4230135, upload-time = "2025-08-06T18:00:19.997Z" }, - { url = "https://files.pythonhosted.org/packages/f3/0b/19b1ba5ee4412f303475a2c7ad5858efb99c90eae5ec627aa6275c439957/debugpy-1.8.16-cp312-cp312-win32.whl", hash = "sha256:833a61ed446426e38b0dd8be3e9d45ae285d424f5bf6cd5b2b559c8f12305508", size = 5281271, upload-time = "2025-08-06T18:00:21.281Z" }, - { url = "https://files.pythonhosted.org/packages/b1/e0/bc62e2dc141de53bd03e2c7cb9d7011de2e65e8bdcdaa26703e4d28656ba/debugpy-1.8.16-cp312-cp312-win_amd64.whl", hash = "sha256:75f204684581e9ef3dc2f67687c3c8c183fde2d6675ab131d94084baf8084121", size = 5323149, upload-time = "2025-08-06T18:00:23.033Z" }, { url = "https://files.pythonhosted.org/packages/52/57/ecc9ae29fa5b2d90107cd1d9bf8ed19aacb74b2264d986ae9d44fe9bdf87/debugpy-1.8.16-py2.py3-none-any.whl", hash = "sha256:19c9521962475b87da6f673514f7fd610328757ec993bf7ec0d8c96f9a325f9e", size = 5287700, upload-time = "2025-08-06T18:00:42.333Z" }, ] @@ -604,13 +695,44 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/4e/8c/f3147f5c4b73e7550fe5f9352eaa956ae838d5c51eb58e7a25b9f3e2643b/decorator-5.2.1-py3-none-any.whl", hash = "sha256:d316bb415a2d9e2d2b3abcc4084c6502fc09240e292cd76a76afc106a1c8e04a", size = 9190, upload-time = "2025-02-24T04:41:32.565Z" }, ] +[[package]] +name = "deep-ep" +version = "1.2.1+29d31c0" +source = { url = "https://github.com/PrimeIntellect-ai/prime-rl/releases/download/v0.5.0/deep_ep-1.2.1+29d31c0-cp312-cp312-linux_x86_64.whl" } +wheels = [ + { url = "https://github.com/PrimeIntellect-ai/prime-rl/releases/download/v0.5.0/deep_ep-1.2.1+29d31c0-cp312-cp312-linux_x86_64.whl", hash = "sha256:a52b709d62c4dfc46a018b988724a7c10894125091aef596bf6d68e94c2b564e" }, +] + +[[package]] +name = "deep-gemm" +version = "2.5.0+891d57b" +source = { url = "https://github.com/PrimeIntellect-ai/prime-rl/releases/download/v0.5.0/deep_gemm-2.5.0+891d57b-cp312-cp312-linux_x86_64.whl" } +wheels = [ + { url = "https://github.com/PrimeIntellect-ai/prime-rl/releases/download/v0.5.0/deep_gemm-2.5.0+891d57b-cp312-cp312-linux_x86_64.whl", hash = "sha256:3a9e4472b86e8e7e2c77cc76951c886bbe7a23a1f5908958969ab5b9c707a2f3" }, +] + +[[package]] +name = "deepdive" +version = "0.2.4" +source = { registry = "https://hub.primeintellect.ai/primeintellect/simple/" } +dependencies = [ + { name = "aiohttp", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "datasets", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "diskcache", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "pdfminer-six", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "verifiers", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, +] +wheels = [ + { url = "https://hub.primeintellect.ai/primeintellect/deepdive/@ea82b01c/deepdive-0.2.4-py2.py3-none-any.whl", hash = "sha256:3d95c5242a78cba0baa087a4ac128568617a2da30e2d3c603a2649a89a4998ab" }, +] + [[package]] name = "depyf" version = "0.20.0" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "astor" }, - { name = "dill" }, + { name = "astor", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "dill", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, ] sdist = { url = "https://files.pythonhosted.org/packages/88/35/83fb0178212279aa0af031031905804c6de5618435d229f41ed21bb9ad2c/depyf-0.20.0.tar.gz", hash = "sha256:fb7683bd72c44f67b56029df2c47721e9a02ffa4d7b19095f1c54c4ebf797a98", size = 6168761, upload-time = "2025-10-13T12:33:38.589Z" } wheels = [ @@ -629,12 +751,11 @@ wheels = [ [[package]] name = "dion" version = "0.1.0" -source = { git = "https://github.com/samsja/dion.git?rev=d891eebbd950a6ff11d9b5f806282e33df83df9d#d891eebbd950a6ff11d9b5f806282e33df83df9d" } +source = { git = "https://github.com/samsja/dion.git?rev=d891eeb#d891eebbd950a6ff11d9b5f806282e33df83df9d" } dependencies = [ - { name = "numpy" }, - { name = "torch" }, - { name = "triton", version = "3.3.1", source = { registry = "https://pypi.org/simple" }, marker = "sys_platform != 'linux'" }, - { name = "triton", version = "3.5.0", source = { registry = "https://pypi.org/simple" }, marker = "sys_platform == 'linux'" }, + { name = "numpy", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "torch", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "triton", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, ] [[package]] @@ -673,6 +794,19 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/ba/5a/18ad964b0086c6e62e2e7500f7edc89e3faa45033c71c1893d34eed2b2de/dnspython-2.8.0-py3-none-any.whl", hash = "sha256:01d9bbc4a2d76bf0db7c1f729812ded6d912bd318d3b1cf81d30c0f845dbf3af", size = 331094, upload-time = "2025-09-07T18:57:58.071Z" }, ] +[[package]] +name = "docker" +version = "7.1.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "requests", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "urllib3", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/91/9b/4a2ea29aeba62471211598dac5d96825bb49348fa07e906ea930394a83ce/docker-7.1.0.tar.gz", hash = "sha256:ad8c70e6e3f8926cb8a92619b832b4ea5299e2831c14284663184e200546fa6c", size = 117834, upload-time = "2024-05-23T11:13:57.216Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/e3/26/57c6fb270950d476074c087527a558ccb6f4436657314bfb6cdf484114c4/docker-7.1.0-py3-none-any.whl", hash = "sha256:c96b93b7f0a746f9e77d325bcfb87422a3d8bd4f03136ae8a85b37f1898d5fc0", size = 147774, upload-time = "2024-05-23T11:13:55.01Z" }, +] + [[package]] name = "docstring-parser" version = "0.17.0" @@ -682,6 +816,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/55/e2/2537ebcff11c1ee1ff17d8d0b6f4db75873e3b0fb32c2d4a2ee31ecb310a/docstring_parser-0.17.0-py3-none-any.whl", hash = "sha256:cf2569abd23dce8099b300f9b4fa8191e9582dda731fd533daf54c4551658708", size = 36896, upload-time = "2025-07-21T07:35:00.684Z" }, ] +[[package]] +name = "durationpy" +version = "0.10" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/9d/a4/e44218c2b394e31a6dd0d6b095c4e1f32d0be54c2a4b250032d717647bab/durationpy-0.10.tar.gz", hash = "sha256:1fa6893409a6e739c9c72334fc65cca1f355dbdd93405d30f726deb5bde42fba", size = 3335, upload-time = "2025-05-17T13:52:37.26Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/b0/0d/9feae160378a3553fa9a339b0e9c1a048e147a4127210e286ef18b730f03/durationpy-0.10-py3-none-any.whl", hash = "sha256:3b41e1b601234296b4fb368338fdcd3e13e0b4fb5b67345948f4f2bf9868b286", size = 3922, upload-time = "2025-05-17T13:52:36.463Z" }, +] + [[package]] name = "einops" version = "0.8.1" @@ -696,8 +839,8 @@ name = "email-validator" version = "2.3.0" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "dnspython" }, - { name = "idna" }, + { name = "dnspython", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "idna", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, ] sdist = { url = "https://files.pythonhosted.org/packages/f5/22/900cb125c76b7aaa450ce02fd727f452243f2e91a61af068b40adba60ea9/email_validator-2.3.0.tar.gz", hash = "sha256:9fc05c37f2f6cf439ff414f8fc46d917929974a82244c20eb10231ba60c54426", size = 51238, upload-time = "2025-08-26T13:09:06.831Z" } wheels = [ @@ -718,10 +861,10 @@ name = "fastapi" version = "0.124.4" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "annotated-doc" }, - { name = "pydantic" }, - { name = "starlette" }, - { name = "typing-extensions" }, + { name = "annotated-doc", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "pydantic", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "starlette", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "typing-extensions", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, ] sdist = { url = "https://files.pythonhosted.org/packages/cd/21/ade3ff6745a82ea8ad88552b4139d27941549e4f19125879f848ac8f3c3d/fastapi-0.124.4.tar.gz", hash = "sha256:0e9422e8d6b797515f33f500309f6e1c98ee4e85563ba0f2debb282df6343763", size = 378460, upload-time = "2025-12-12T15:00:43.891Z" } wheels = [ @@ -730,12 +873,12 @@ wheels = [ [package.optional-dependencies] standard = [ - { name = "email-validator" }, - { name = "fastapi-cli", extra = ["standard"] }, - { name = "httpx" }, - { name = "jinja2" }, - { name = "python-multipart" }, - { name = "uvicorn", extra = ["standard"] }, + { name = "email-validator", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "fastapi-cli", extra = ["standard"], marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "httpx", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "jinja2", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "python-multipart", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "uvicorn", extra = ["standard"], marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, ] [[package]] @@ -743,9 +886,9 @@ name = "fastapi-cli" version = "0.0.11" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "rich-toolkit" }, - { name = "typer" }, - { name = "uvicorn", extra = ["standard"] }, + { name = "rich-toolkit", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "typer", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "uvicorn", extra = ["standard"], marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, ] sdist = { url = "https://files.pythonhosted.org/packages/23/08/0af729f6231ebdc17a0356397f966838cbe2efa38529951e24017c7435d5/fastapi_cli-0.0.11.tar.gz", hash = "sha256:4f01d751c14d3d2760339cca0f45e81d816218cae8174d1dc757b5375868cde5", size = 17550, upload-time = "2025-09-09T12:50:38.917Z" } wheels = [ @@ -754,8 +897,8 @@ wheels = [ [package.optional-dependencies] standard = [ - { name = "fastapi-cloud-cli" }, - { name = "uvicorn", extra = ["standard"] }, + { name = "fastapi-cloud-cli", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "uvicorn", extra = ["standard"], marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, ] [[package]] @@ -763,13 +906,13 @@ name = "fastapi-cloud-cli" version = "0.1.5" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "httpx" }, - { name = "pydantic", extra = ["email"] }, - { name = "rich-toolkit" }, - { name = "rignore" }, - { name = "sentry-sdk" }, - { name = "typer" }, - { name = "uvicorn", extra = ["standard"] }, + { name = "httpx", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "pydantic", extra = ["email"], marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "rich-toolkit", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "rignore", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "sentry-sdk", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "typer", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "uvicorn", extra = ["standard"], marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, ] sdist = { url = "https://files.pythonhosted.org/packages/a9/2e/3b6e5016affc310e5109bc580f760586eabecea0c8a7ab067611cd849ac0/fastapi_cloud_cli-0.1.5.tar.gz", hash = "sha256:341ee585eb731a6d3c3656cb91ad38e5f39809bf1a16d41de1333e38635a7937", size = 22710, upload-time = "2025-07-28T13:30:48.216Z" } wheels = [ @@ -778,27 +921,23 @@ wheels = [ [[package]] name = "fastcore" -version = "1.8.8" +version = "1.12.31" source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "packaging" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/bf/15/2303dcaea2d11a6997a47a6aa9b0e8dcae5b4c8941c7694e1f9e97990cae/fastcore-1.8.8.tar.gz", hash = "sha256:939f7b79c5510b059ba9fc512dfe8aab30f5a752130a640070f8a1b494493b56", size = 76675, upload-time = "2025-08-23T22:35:56.378Z" } +sdist = { url = "https://files.pythonhosted.org/packages/1b/72/acab01c6b68c7f94080f79049b3d67bbec3af98b65ef101775e08cd11f41/fastcore-1.12.31.tar.gz", hash = "sha256:684a7b05b9f49d94d9ebbf7e104bbe54fa001aa8c78fa681cc7aabbd12d1bbde", size = 94319, upload-time = "2026-03-22T02:21:21.211Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/7b/ec/bedc0f8b644deb023fcbc8e2879a8f84935f9109ebc61c6614b0bb78f90c/fastcore-1.8.8-py3-none-any.whl", hash = "sha256:f38853245e5ae5abb16275daac92fca00edeced020871b1d2416fb22cde70df8", size = 79436, upload-time = "2025-08-23T22:35:54.433Z" }, + { url = "https://files.pythonhosted.org/packages/1e/5c/9516828037b1680de14bae4d93b2868b26417059e1e5cb7d8ad7be649197/fastcore-1.12.31-py3-none-any.whl", hash = "sha256:2634f117fe3a2f6c250f8500f70211c96158399936e90fa6d1010d2d516bf03a", size = 98509, upload-time = "2026-03-22T02:21:19.274Z" }, ] [[package]] -name = "fastrlock" -version = "0.8.3" +name = "fastsafetensors" +version = "0.3" source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/73/b1/1c3d635d955f2b4bf34d45abf8f35492e04dbd7804e94ce65d9f928ef3ec/fastrlock-0.8.3.tar.gz", hash = "sha256:4af6734d92eaa3ab4373e6c9a1dd0d5ad1304e172b1521733c6c3b3d73c8fa5d", size = 79327, upload-time = "2024-12-17T11:03:39.638Z" } +dependencies = [ + { name = "typer", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/39/98/053c622e61bb766d31327a88215082320a4ba8bd6a62c4c5435221844103/fastsafetensors-0.3.tar.gz", hash = "sha256:89f392569d2281d1a966d3b64f99a6386149116e37eef4f4890168c87a8c4f19", size = 57500, upload-time = "2026-04-22T07:16:30.84Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/57/21/ea1511b0ef0d5457efca3bf1823effb9c5cad4fc9dca86ce08e4d65330ce/fastrlock-0.8.3-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:85a49a1f1e020097d087e1963e42cea6f307897d5ebe2cb6daf4af47ffdd3eed", size = 52201, upload-time = "2024-12-17T11:02:19.512Z" }, - { url = "https://files.pythonhosted.org/packages/80/07/cdecb7aa976f34328372f1c4efd6c9dc1b039b3cc8d3f38787d640009a25/fastrlock-0.8.3-cp312-cp312-manylinux_2_5_x86_64.manylinux1_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:5f13ec08f1adb1aa916c384b05ecb7dbebb8df9ea81abd045f60941c6283a670", size = 53924, upload-time = "2024-12-17T11:02:20.85Z" }, - { url = "https://files.pythonhosted.org/packages/88/6d/59c497f8db9a125066dd3a7442fab6aecbe90d6fec344c54645eaf311666/fastrlock-0.8.3-cp312-cp312-musllinux_1_1_aarch64.whl", hash = "sha256:0ea4e53a04980d646def0f5e4b5e8bd8c7884288464acab0b37ca0c65c482bfe", size = 52140, upload-time = "2024-12-17T11:02:22.263Z" }, - { url = "https://files.pythonhosted.org/packages/62/04/9138943c2ee803d62a48a3c17b69de2f6fa27677a6896c300369e839a550/fastrlock-0.8.3-cp312-cp312-musllinux_1_1_x86_64.whl", hash = "sha256:38340f6635bd4ee2a4fb02a3a725759fe921f2ca846cb9ca44531ba739cc17b4", size = 53261, upload-time = "2024-12-17T11:02:24.418Z" }, - { url = "https://files.pythonhosted.org/packages/e2/4b/db35a52589764c7745a613b6943bbd018f128d42177ab92ee7dde88444f6/fastrlock-0.8.3-cp312-cp312-win_amd64.whl", hash = "sha256:da06d43e1625e2ffddd303edcd6d2cd068e1c486f5fd0102b3f079c44eb13e2c", size = 31235, upload-time = "2024-12-17T11:02:25.708Z" }, + { url = "https://files.pythonhosted.org/packages/0a/06/bca80663bf8136f273643d149953dd29ca2c52aa4faac4b67506b871a5ec/fastsafetensors-0.3-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:3ce38241c5afedf94ef37531b8b8703016b2ea39350cfd33e819e65d4d5305e0", size = 1855661, upload-time = "2026-04-22T07:16:25.833Z" }, ] [[package]] @@ -812,14 +951,14 @@ wheels = [ [[package]] name = "flash-attn" -version = "2.8.3+cu128torch2.9" -source = { url = "https://github.com/mjun0812/flash-attention-prebuild-wheels/releases/download/v0.6.8/flash_attn-2.8.3+cu128torch2.9-cp312-cp312-linux_x86_64.whl" } +version = "2.8.3+cu128torch2.11" +source = { url = "https://github.com/mjun0812/flash-attention-prebuild-wheels/releases/download/v0.9.4/flash_attn-2.8.3+cu128torch2.11-cp312-cp312-linux_x86_64.whl" } dependencies = [ - { name = "einops" }, - { name = "torch" }, + { name = "einops", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "torch", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, ] wheels = [ - { url = "https://github.com/mjun0812/flash-attention-prebuild-wheels/releases/download/v0.6.8/flash_attn-2.8.3+cu128torch2.9-cp312-cp312-linux_x86_64.whl", hash = "sha256:95c7c64416c1be3f3f6d17468e49c343a716d1df681f30fb10148a1441339401" }, + { url = "https://github.com/mjun0812/flash-attention-prebuild-wheels/releases/download/v0.9.4/flash_attn-2.8.3+cu128torch2.11-cp312-cp312-linux_x86_64.whl", hash = "sha256:a16162f436286cc03ebbfb174c0853343ed98ae13c37abf1042947668ec40549" }, ] [package.metadata] @@ -830,65 +969,81 @@ requires-dist = [ [[package]] name = "flash-attn-3" -version = "3.0.0b1" -source = { url = "https://github.com/samsja/flash-attn-builds/releases/download/v0.1/flash_attn_3-3.0.0b1-cp39-abi3-linux_x86_64.whl" } +version = "3.0.0" +source = { registry = "https://download.pytorch.org/whl/test/cu128" } dependencies = [ - { name = "einops" }, - { name = "ninja" }, - { name = "packaging" }, - { name = "torch" }, + { name = "einops", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "ninja", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "packaging", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "torch", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, ] wheels = [ - { url = "https://github.com/samsja/flash-attn-builds/releases/download/v0.1/flash_attn_3-3.0.0b1-cp39-abi3-linux_x86_64.whl", hash = "sha256:4878b0e81704972109479f0667b2ace65f62e62a1161c9dbc23fa746889c79e4" }, + { url = "https://download.pytorch.org/whl/test/cu128/flash_attn_3-3.0.0-cp39-abi3-manylinux_2_28_x86_64.whl", hash = "sha256:10a730d5e1a1c23afd930ce419ee425cd1925a16b7eac789dfe98c3ca39bc009", upload-time = "2026-02-14T04:28:12Z" }, + { url = "https://download.pytorch.org/whl/test/cu128/flash_attn_3-3.0.0-cp39-abi3-manylinux_2_34_aarch64.whl", hash = "sha256:bbb51d13d7aa1bfe04ef8875ed1bf630ebb116775084a5df612495e10609cd76", upload-time = "2026-02-14T04:28:05Z" }, ] -[package.metadata] -requires-dist = [ - { name = "einops" }, - { name = "ninja" }, - { name = "packaging" }, - { name = "torch" }, +[[package]] +name = "flash-attn-4" +version = "4.0.0b11.dev11+g96bd151" +source = { git = "https://github.com/Dao-AILab/flash-attention.git?subdirectory=flash_attn%2Fcute&rev=96bd151#96bd151b00add1dc5744115cfc350c0a809c3fa2" } +dependencies = [ + { name = "apache-tvm-ffi", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "einops", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "nvidia-cutlass-dsl", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "quack-kernels", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "torch-c-dlpack-ext", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "typing-extensions", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, +] + +[[package]] +name = "flash-linear-attention" +version = "0.4.2" +source = { git = "https://github.com/fla-org/flash-linear-attention#02af88ef8aa7f7043a899c9ca6fde168e1cf8c7e" } +dependencies = [ + { name = "einops", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "torch", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "transformers", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, +] + +[[package]] +name = "flashinfer-cubin" +version = "0.6.8.post1" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/11/b7/5e3b1a8c67031b421a8bd29c2bc29b900a550bb3392e8bda18bb15b5e476/flashinfer_cubin-0.6.8.post1-py3-none-any.whl", hash = "sha256:43636d4cd39e694a83d76a89f87fefcdf4cecb4c4f7dd22dac25ec368c1e901f", size = 295154113, upload-time = "2026-04-18T18:28:21.738Z" }, ] [[package]] name = "flashinfer-python" -version = "0.5.3" +version = "0.6.8.post1" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "apache-tvm-ffi" }, - { name = "click" }, - { name = "einops" }, - { name = "ninja" }, - { name = "numpy" }, - { name = "nvidia-cudnn-frontend" }, - { name = "nvidia-cutlass-dsl" }, - { name = "nvidia-ml-py" }, - { name = "packaging" }, - { name = "requests" }, - { name = "tabulate" }, - { name = "torch" }, - { name = "tqdm" }, + { name = "apache-tvm-ffi", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "click", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "cuda-tile", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "einops", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "ninja", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "numpy", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "nvidia-cudnn-frontend", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "nvidia-cutlass-dsl", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "nvidia-ml-py", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "packaging", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "requests", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "tabulate", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "torch", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "tqdm", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/b4/91/cca69baeff24bb3efd12c7479a026432c8717ee47193694010494c528b22/flashinfer_python-0.5.3.tar.gz", hash = "sha256:100d59b0ede47878d2808cd3a1b9039d7a952d66338bc9f68dac192ae1b2e3f1", size = 4682367, upload-time = "2025-11-20T21:22:46.976Z" } +sdist = { url = "https://files.pythonhosted.org/packages/53/1e/2760fef9e74abc4480961048e5790b4c9e955872fb4d7d97900cfddced5a/flashinfer_python-0.6.8.post1.tar.gz", hash = "sha256:b18e4121baf9b93fa9a9f368ba9b981a0342895f50ab9dddc224aeb964ed346f", size = 6675885, upload-time = "2026-04-18T18:28:13.299Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/76/78/6dc7e7da8cb87c9965644ea0d2439457a1bc9256c45ceda0044595be4143/flashinfer_python-0.5.3-py3-none-any.whl", hash = "sha256:b601293b72f9138bad173edc28df84b9f239a013be974e2e79d4ba98aeb38cf5", size = 6998069, upload-time = "2025-11-20T21:22:45.104Z" }, + { url = "https://files.pythonhosted.org/packages/73/6d/1e8a8533913e33a50a486332ce0673f4fdb860f6eb9ed450327c5c1762cb/flashinfer_python-0.6.8.post1-py3-none-any.whl", hash = "sha256:818f9b8cc2fe66c42a1f6264be4841ac8821ada703685a02cfccb2b5124a710b", size = 9385316, upload-time = "2026-04-18T18:28:10.285Z" }, ] [[package]] -name = "fonttools" -version = "4.59.2" +name = "flatbuffers" +version = "25.12.19" source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/0d/a5/fba25f9fbdab96e26dedcaeeba125e5f05a09043bf888e0305326e55685b/fonttools-4.59.2.tar.gz", hash = "sha256:e72c0749b06113f50bcb80332364c6be83a9582d6e3db3fe0b280f996dc2ef22", size = 3540889, upload-time = "2025-08-27T16:40:30.97Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/ba/3d/1f45db2df51e7bfa55492e8f23f383d372200be3a0ded4bf56a92753dd1f/fonttools-4.59.2-cp312-cp312-macosx_10_13_universal2.whl", hash = "sha256:82906d002c349cad647a7634b004825a7335f8159d0d035ae89253b4abf6f3ea", size = 2769711, upload-time = "2025-08-27T16:39:04.423Z" }, - { url = "https://files.pythonhosted.org/packages/29/df/cd236ab32a8abfd11558f296e064424258db5edefd1279ffdbcfd4fd8b76/fonttools-4.59.2-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:a10c1bd7644dc58f8862d8ba0cf9fb7fef0af01ea184ba6ce3f50ab7dfe74d5a", size = 2340225, upload-time = "2025-08-27T16:39:06.143Z" }, - { url = "https://files.pythonhosted.org/packages/98/12/b6f9f964fe6d4b4dd4406bcbd3328821c3de1f909ffc3ffa558fe72af48c/fonttools-4.59.2-cp312-cp312-manylinux1_x86_64.manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:738f31f23e0339785fd67652a94bc69ea49e413dfdb14dcb8c8ff383d249464e", size = 4912766, upload-time = "2025-08-27T16:39:08.138Z" }, - { url = "https://files.pythonhosted.org/packages/73/78/82bde2f2d2c306ef3909b927363170b83df96171f74e0ccb47ad344563cd/fonttools-4.59.2-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:0ec99f9bdfee9cdb4a9172f9e8fd578cce5feb231f598909e0aecf5418da4f25", size = 4955178, upload-time = "2025-08-27T16:39:10.094Z" }, - { url = "https://files.pythonhosted.org/packages/92/77/7de766afe2d31dda8ee46d7e479f35c7d48747e558961489a2d6e3a02bd4/fonttools-4.59.2-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:0476ea74161322e08c7a982f83558a2b81b491509984523a1a540baf8611cc31", size = 4897898, upload-time = "2025-08-27T16:39:12.087Z" }, - { url = "https://files.pythonhosted.org/packages/c5/77/ce0e0b905d62a06415fda9f2b2e109a24a5db54a59502b769e9e297d2242/fonttools-4.59.2-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:95922a922daa1f77cc72611747c156cfb38030ead72436a2c551d30ecef519b9", size = 5049144, upload-time = "2025-08-27T16:39:13.84Z" }, - { url = "https://files.pythonhosted.org/packages/d9/ea/870d93aefd23fff2e07cbeebdc332527868422a433c64062c09d4d5e7fe6/fonttools-4.59.2-cp312-cp312-win32.whl", hash = "sha256:39ad9612c6a622726a6a130e8ab15794558591f999673f1ee7d2f3d30f6a3e1c", size = 2206473, upload-time = "2025-08-27T16:39:15.854Z" }, - { url = "https://files.pythonhosted.org/packages/61/c4/e44bad000c4a4bb2e9ca11491d266e857df98ab6d7428441b173f0fe2517/fonttools-4.59.2-cp312-cp312-win_amd64.whl", hash = "sha256:980fd7388e461b19a881d35013fec32c713ffea1fc37aef2f77d11f332dfd7da", size = 2254706, upload-time = "2025-08-27T16:39:17.893Z" }, - { url = "https://files.pythonhosted.org/packages/65/a4/d2f7be3c86708912c02571db0b550121caab8cd88a3c0aacb9cfa15ea66e/fonttools-4.59.2-py3-none-any.whl", hash = "sha256:8bd0f759020e87bb5d323e6283914d9bf4ae35a7307dafb2cbd1e379e720ad37", size = 1132315, upload-time = "2025-08-27T16:40:28.984Z" }, + { url = "https://files.pythonhosted.org/packages/e8/2d/d2a548598be01649e2d46231d151a6c56d10b964d94043a335ae56ea2d92/flatbuffers-25.12.19-py2.py3-none-any.whl", hash = "sha256:7634f50c427838bb021c2d66a3d1168e9d199b0607e6329399f04846d42e20b4", size = 26661, upload-time = "2025-12-19T23:16:13.622Z" }, ] [[package]] @@ -897,23 +1052,10 @@ version = "1.7.0" source = { registry = "https://pypi.org/simple" } sdist = { url = "https://files.pythonhosted.org/packages/79/b1/b64018016eeb087db503b038296fd782586432b9c077fc5c7839e9cb6ef6/frozenlist-1.7.0.tar.gz", hash = "sha256:2e310d81923c2437ea8670467121cc3e9b0f76d3043cc1d2331d56c7fb7a3a8f", size = 45078, upload-time = "2025-06-09T23:02:35.538Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/ef/a2/c8131383f1e66adad5f6ecfcce383d584ca94055a34d683bbb24ac5f2f1c/frozenlist-1.7.0-cp312-cp312-macosx_10_13_universal2.whl", hash = "sha256:3dbf9952c4bb0e90e98aec1bd992b3318685005702656bc6f67c1a32b76787f2", size = 81424, upload-time = "2025-06-09T23:00:42.24Z" }, - { url = "https://files.pythonhosted.org/packages/4c/9d/02754159955088cb52567337d1113f945b9e444c4960771ea90eb73de8db/frozenlist-1.7.0-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:1f5906d3359300b8a9bb194239491122e6cf1444c2efb88865426f170c262cdb", size = 47952, upload-time = "2025-06-09T23:00:43.481Z" }, - { url = "https://files.pythonhosted.org/packages/01/7a/0046ef1bd6699b40acd2067ed6d6670b4db2f425c56980fa21c982c2a9db/frozenlist-1.7.0-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:3dabd5a8f84573c8d10d8859a50ea2dec01eea372031929871368c09fa103478", size = 46688, upload-time = "2025-06-09T23:00:44.793Z" }, { url = "https://files.pythonhosted.org/packages/d6/a2/a910bafe29c86997363fb4c02069df4ff0b5bc39d33c5198b4e9dd42d8f8/frozenlist-1.7.0-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:aa57daa5917f1738064f302bf2626281a1cb01920c32f711fbc7bc36111058a8", size = 243084, upload-time = "2025-06-09T23:00:46.125Z" }, - { url = "https://files.pythonhosted.org/packages/64/3e/5036af9d5031374c64c387469bfcc3af537fc0f5b1187d83a1cf6fab1639/frozenlist-1.7.0-cp312-cp312-manylinux_2_17_armv7l.manylinux2014_armv7l.manylinux_2_31_armv7l.whl", hash = "sha256:c193dda2b6d49f4c4398962810fa7d7c78f032bf45572b3e04dd5249dff27e08", size = 233524, upload-time = "2025-06-09T23:00:47.73Z" }, - { url = "https://files.pythonhosted.org/packages/06/39/6a17b7c107a2887e781a48ecf20ad20f1c39d94b2a548c83615b5b879f28/frozenlist-1.7.0-cp312-cp312-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:bfe2b675cf0aaa6d61bf8fbffd3c274b3c9b7b1623beb3809df8a81399a4a9c4", size = 248493, upload-time = "2025-06-09T23:00:49.742Z" }, - { url = "https://files.pythonhosted.org/packages/be/00/711d1337c7327d88c44d91dd0f556a1c47fb99afc060ae0ef66b4d24793d/frozenlist-1.7.0-cp312-cp312-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:8fc5d5cda37f62b262405cf9652cf0856839c4be8ee41be0afe8858f17f4c94b", size = 244116, upload-time = "2025-06-09T23:00:51.352Z" }, - { url = "https://files.pythonhosted.org/packages/24/fe/74e6ec0639c115df13d5850e75722750adabdc7de24e37e05a40527ca539/frozenlist-1.7.0-cp312-cp312-manylinux_2_5_i686.manylinux1_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:b0d5ce521d1dd7d620198829b87ea002956e4319002ef0bc8d3e6d045cb4646e", size = 224557, upload-time = "2025-06-09T23:00:52.855Z" }, { url = "https://files.pythonhosted.org/packages/8d/db/48421f62a6f77c553575201e89048e97198046b793f4a089c79a6e3268bd/frozenlist-1.7.0-cp312-cp312-manylinux_2_5_x86_64.manylinux1_x86_64.manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:488d0a7d6a0008ca0db273c542098a0fa9e7dfaa7e57f70acef43f32b3f69dca", size = 241820, upload-time = "2025-06-09T23:00:54.43Z" }, { url = "https://files.pythonhosted.org/packages/1d/fa/cb4a76bea23047c8462976ea7b7a2bf53997a0ca171302deae9d6dd12096/frozenlist-1.7.0-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:15a7eaba63983d22c54d255b854e8108e7e5f3e89f647fc854bd77a237e767df", size = 236542, upload-time = "2025-06-09T23:00:56.409Z" }, - { url = "https://files.pythonhosted.org/packages/5d/32/476a4b5cfaa0ec94d3f808f193301debff2ea42288a099afe60757ef6282/frozenlist-1.7.0-cp312-cp312-musllinux_1_2_armv7l.whl", hash = "sha256:1eaa7e9c6d15df825bf255649e05bd8a74b04a4d2baa1ae46d9c2d00b2ca2cb5", size = 249350, upload-time = "2025-06-09T23:00:58.468Z" }, - { url = "https://files.pythonhosted.org/packages/8d/ba/9a28042f84a6bf8ea5dbc81cfff8eaef18d78b2a1ad9d51c7bc5b029ad16/frozenlist-1.7.0-cp312-cp312-musllinux_1_2_i686.whl", hash = "sha256:e4389e06714cfa9d47ab87f784a7c5be91d3934cd6e9a7b85beef808297cc025", size = 225093, upload-time = "2025-06-09T23:01:00.015Z" }, - { url = "https://files.pythonhosted.org/packages/bc/29/3a32959e68f9cf000b04e79ba574527c17e8842e38c91d68214a37455786/frozenlist-1.7.0-cp312-cp312-musllinux_1_2_ppc64le.whl", hash = "sha256:73bd45e1488c40b63fe5a7df892baf9e2a4d4bb6409a2b3b78ac1c6236178e01", size = 245482, upload-time = "2025-06-09T23:01:01.474Z" }, - { url = "https://files.pythonhosted.org/packages/80/e8/edf2f9e00da553f07f5fa165325cfc302dead715cab6ac8336a5f3d0adc2/frozenlist-1.7.0-cp312-cp312-musllinux_1_2_s390x.whl", hash = "sha256:99886d98e1643269760e5fe0df31e5ae7050788dd288947f7f007209b8c33f08", size = 249590, upload-time = "2025-06-09T23:01:02.961Z" }, { url = "https://files.pythonhosted.org/packages/1c/80/9a0eb48b944050f94cc51ee1c413eb14a39543cc4f760ed12657a5a3c45a/frozenlist-1.7.0-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:290a172aae5a4c278c6da8a96222e6337744cd9c77313efe33d5670b9f65fc43", size = 237785, upload-time = "2025-06-09T23:01:05.095Z" }, - { url = "https://files.pythonhosted.org/packages/f3/74/87601e0fb0369b7a2baf404ea921769c53b7ae00dee7dcfe5162c8c6dbf0/frozenlist-1.7.0-cp312-cp312-win32.whl", hash = "sha256:426c7bc70e07cfebc178bc4c2bf2d861d720c4fff172181eeb4a4c41d4ca2ad3", size = 39487, upload-time = "2025-06-09T23:01:06.54Z" }, - { url = "https://files.pythonhosted.org/packages/0b/15/c026e9a9fc17585a9d461f65d8593d281fedf55fbf7eb53f16c6df2392f9/frozenlist-1.7.0-cp312-cp312-win_amd64.whl", hash = "sha256:563b72efe5da92e02eb68c59cb37205457c977aa7a449ed1b37e6939e5c47c6a", size = 43874, upload-time = "2025-06-09T23:01:07.752Z" }, { url = "https://files.pythonhosted.org/packages/ee/45/b82e3c16be2182bff01179db177fe144d58b5dc787a7d4492c6ed8b9317f/frozenlist-1.7.0-py3-none-any.whl", hash = "sha256:9a5af342e34f7e97caf8c995864c7a396418ae2859cc6fdf1b1073020d516a7e", size = 13106, upload-time = "2025-06-09T23:02:34.204Z" }, ] @@ -928,7 +1070,16 @@ wheels = [ [package.optional-dependencies] http = [ - { name = "aiohttp" }, + { name = "aiohttp", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, +] + +[[package]] +name = "gepa" +version = "0.0.27" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/4e/99/6840f84498f2dcbfd27e8a15eeb4e637e84e9dbbd6331977b71f6ffad9c7/gepa-0.0.27.tar.gz", hash = "sha256:02ecb19e4aa6a1f5bb2994cd54b2057f25cd5e8cd662fee514303b2a8ba0a738", size = 155106, upload-time = "2026-01-28T00:33:51.72Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/a3/bd/f9e83519099a9fd5d090520ab0c51a6278e473b914a9b32d1e812662dd3b/gepa-0.0.27-py3-none-any.whl", hash = "sha256:592beb084fd638d525edae84ca75ad8e9e9c3758e5498b933c17153addf5dd2d", size = 146454, upload-time = "2026-01-28T00:33:50.262Z" }, ] [[package]] @@ -936,45 +1087,75 @@ name = "gguf" version = "0.17.1" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "numpy" }, - { name = "pyyaml" }, - { name = "tqdm" }, + { name = "numpy", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "pyyaml", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "tqdm", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, ] sdist = { url = "https://files.pythonhosted.org/packages/08/08/7de1ca4b71e7bf33b547f82bb22505e221b5fa42f67d635e200e0ad22ad6/gguf-0.17.1.tar.gz", hash = "sha256:36ad71aad900a3e75fc94ebe96ea6029f03a4e44be7627ef7ad3d03e8c7bcb53", size = 89338, upload-time = "2025-06-19T14:00:33.705Z" } wheels = [ { url = "https://files.pythonhosted.org/packages/fc/31/6a93a887617ee7deeaa602ca3d02d1c12a6cb8a742a695de5d128f5fa46a/gguf-0.17.1-py3-none-any.whl", hash = "sha256:7bc5aa7eeb1931f7d39b48fdc5b38fda6b294b9dca75cf607ac69557840a3943", size = 96224, upload-time = "2025-06-19T14:00:32.88Z" }, ] +[[package]] +name = "ghapi" +version = "1.0.13" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "fastcore", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/62/09/1b88f97e8599cda096d42dac830bb2e28ddf202d71843f61bda52bbe99ce/ghapi-1.0.13.tar.gz", hash = "sha256:fb46f5e101efa33bd12a0ae7694de761eec5be1de90f48847699f1e00128f928", size = 72914, upload-time = "2026-02-28T02:21:01.892Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/67/ac/e1960ec21cfd5a0fd9b329822c04d0b5f91abb688c3b1acd7e8ff3390432/ghapi-1.0.13-py3-none-any.whl", hash = "sha256:49d7e336e5664e4d4f92b1d442dfe80f31ecccbee4370bd1d271bd63a1ccf18e", size = 71409, upload-time = "2026-02-28T02:21:00.457Z" }, +] + [[package]] name = "gitdb" version = "4.0.12" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "smmap" }, + { name = "smmap", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, ] sdist = { url = "https://files.pythonhosted.org/packages/72/94/63b0fc47eb32792c7ba1fe1b694daec9a63620db1e313033d18140c2320a/gitdb-4.0.12.tar.gz", hash = "sha256:5ef71f855d191a3326fcfbc0d5da835f26b13fbcba60c32c21091c349ffdb571", size = 394684, upload-time = "2025-01-02T07:20:46.413Z" } wheels = [ { url = "https://files.pythonhosted.org/packages/a0/61/5c78b91c3143ed5c14207f463aecfc8f9dbb5092fb2869baf37c273b2705/gitdb-4.0.12-py3-none-any.whl", hash = "sha256:67073e15955400952c6565cc3e707c554a4eea2e428946f7a4c162fab9bd9bcf", size = 62794, upload-time = "2025-01-02T07:20:43.624Z" }, ] +[[package]] +name = "gitignore-parser" +version = "0.1.13" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/e5/51/e391a1a4238f18d0abb47be479b07af265ad4519022cf51b7da47ef82487/gitignore_parser-0.1.13.tar.gz", hash = "sha256:c7e10c8190accb8ae57fb3711889e73a9c0dbc04d4222b91ace8a4bf64d2f746", size = 5603, upload-time = "2025-08-25T06:33:22.704Z" } + [[package]] name = "gitpython" version = "3.1.45" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "gitdb" }, + { name = "gitdb", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, ] sdist = { url = "https://files.pythonhosted.org/packages/9a/c8/dd58967d119baab745caec2f9d853297cec1989ec1d63f677d3880632b88/gitpython-3.1.45.tar.gz", hash = "sha256:85b0ee964ceddf211c41b9f27a49086010a190fd8132a24e21f362a4b36a791c", size = 215076, upload-time = "2025-07-24T03:45:54.871Z" } wheels = [ { url = "https://files.pythonhosted.org/packages/01/61/d4b89fec821f72385526e1b9d9a3a0385dda4a72b206d28049e2c7cd39b8/gitpython-3.1.45-py3-none-any.whl", hash = "sha256:8908cb2e02fb3b93b7eb0f2827125cb699869470432cc885f019b8fd0fccff77", size = 208168, upload-time = "2025-07-24T03:45:52.517Z" }, ] +[[package]] +name = "googleapis-common-protos" +version = "1.72.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "protobuf", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/e5/7b/adfd75544c415c487b33061fe7ae526165241c1ea133f9a9125a56b39fd8/googleapis_common_protos-1.72.0.tar.gz", hash = "sha256:e55a601c1b32b52d7a3e65f43563e2aa61bcd737998ee672ac9b951cd49319f5", size = 147433, upload-time = "2025-11-06T18:29:24.087Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/c4/ab/09169d5a4612a5f92490806649ac8d41e3ec9129c636754575b3553f4ea4/googleapis_common_protos-1.72.0-py3-none-any.whl", hash = "sha256:4299c5a82d5ae1a9702ada957347726b167f9f8d1fc352477702a1e851ff4038", size = 297515, upload-time = "2025-11-06T18:29:13.14Z" }, +] + [[package]] name = "griffe" version = "1.14.0" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "colorama" }, + { name = "colorama", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, ] sdist = { url = "https://files.pythonhosted.org/packages/ec/d7/6c09dd7ce4c7837e4cdb11dce980cb45ae3cd87677298dc3b781b6bce7d3/griffe-1.14.0.tar.gz", hash = "sha256:9d2a15c1eca966d68e00517de5d69dd1bc5c9f2335ef6c1775362ba5b8651a13", size = 424684, upload-time = "2025-09-05T15:02:29.167Z" } wheels = [ @@ -983,23 +1164,30 @@ wheels = [ [[package]] name = "grpcio" -version = "1.75.0" +version = "1.78.0rc2" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "typing-extensions", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/0b/bb/d89b2f8ed062af360e872746cab9d5a98acf80f9fec537536203695cce63/grpcio-1.78.0rc2.tar.gz", hash = "sha256:d624592c82a19a5898c5576fbda43c28d7062bac04ea6f33bbd8871bc0639e64", size = 12831859, upload-time = "2026-01-16T07:30:23.727Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/2e/e7/94e4c7ae7fdbaf7adc8af47eb6e3b53166c184b281a470210df8bf0dcb96/grpcio-1.78.0rc2-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:7fe343a2ccaa3ca48a933e81f4c0a9de37057cf5bc5567864a98775cce570456", size = 6476173, upload-time = "2026-01-16T07:28:38.599Z" }, + { url = "https://files.pythonhosted.org/packages/c4/d7/60e821443c044365222a7fcd6630344016ea3e31bf4903f3a22f93e5b3a1/grpcio-1.78.0rc2-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:e7bde54ad7bee2d4dbc6d4a351d5b62cc2bfa87c58e9db911ed8a0489192ca9a", size = 6690812, upload-time = "2026-01-16T07:28:43.887Z" }, + { url = "https://files.pythonhosted.org/packages/80/47/b19c67ca6e0622fccb88558452e6f7458551ef365456585968dcd84a1db3/grpcio-1.78.0rc2-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:4e99bbce4f509eb6af4b523109152258043b92bbb945950456a9318eca71ef2e", size = 7266122, upload-time = "2026-01-16T07:28:46.109Z" }, + { url = "https://files.pythonhosted.org/packages/df/81/dcf11b5915a39024d4a98ef14b16cb0c636a4f2f26ef657982d3144c6544/grpcio-1.78.0rc2-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:ea66e360e5ea032a1f6dde926915ba683edca811ed6f0e0620c52e264ba364e4", size = 7698266, upload-time = "2026-01-16T07:28:51.342Z" }, +] + +[[package]] +name = "grpclib" +version = "0.4.9" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "typing-extensions" }, + { name = "h2", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "multidict", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/91/88/fe2844eefd3d2188bc0d7a2768c6375b46dfd96469ea52d8aeee8587d7e0/grpcio-1.75.0.tar.gz", hash = "sha256:b989e8b09489478c2d19fecc744a298930f40d8b27c3638afbfe84d22f36ce4e", size = 12722485, upload-time = "2025-09-16T09:20:21.731Z" } +sdist = { url = "https://files.pythonhosted.org/packages/5b/28/5a2c299ec82a876a252c5919aa895a6f1d1d35c96417c5ce4a4660dc3a80/grpclib-0.4.9.tar.gz", hash = "sha256:cc589c330fa81004c6400a52a566407574498cb5b055fa927013361e21466c46", size = 84798, upload-time = "2025-12-14T22:23:14.349Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/0d/93/a1b29c2452d15cecc4a39700fbf54721a3341f2ddbd1bd883f8ec0004e6e/grpcio-1.75.0-cp312-cp312-linux_armv7l.whl", hash = "sha256:fa35ccd9501ffdd82b861809cbfc4b5b13f4b4c5dc3434d2d9170b9ed38a9054", size = 5661861, upload-time = "2025-09-16T09:18:58.748Z" }, - { url = "https://files.pythonhosted.org/packages/b8/ce/7280df197e602d14594e61d1e60e89dfa734bb59a884ba86cdd39686aadb/grpcio-1.75.0-cp312-cp312-macosx_11_0_universal2.whl", hash = "sha256:0fcb77f2d718c1e58cc04ef6d3b51e0fa3b26cf926446e86c7eba105727b6cd4", size = 11459982, upload-time = "2025-09-16T09:19:01.211Z" }, - { url = "https://files.pythonhosted.org/packages/7c/9b/37e61349771f89b543a0a0bbc960741115ea8656a2414bfb24c4de6f3dd7/grpcio-1.75.0-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:36764a4ad9dc1eb891042fab51e8cdf7cc014ad82cee807c10796fb708455041", size = 6239680, upload-time = "2025-09-16T09:19:04.443Z" }, - { url = "https://files.pythonhosted.org/packages/a6/66/f645d9d5b22ca307f76e71abc83ab0e574b5dfef3ebde4ec8b865dd7e93e/grpcio-1.75.0-cp312-cp312-manylinux2014_i686.manylinux_2_17_i686.whl", hash = "sha256:725e67c010f63ef17fc052b261004942763c0b18dcd84841e6578ddacf1f9d10", size = 6908511, upload-time = "2025-09-16T09:19:07.884Z" }, - { url = "https://files.pythonhosted.org/packages/e6/9a/34b11cd62d03c01b99068e257595804c695c3c119596c7077f4923295e19/grpcio-1.75.0-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:91fbfc43f605c5ee015c9056d580a70dd35df78a7bad97e05426795ceacdb59f", size = 6429105, upload-time = "2025-09-16T09:19:10.085Z" }, - { url = "https://files.pythonhosted.org/packages/1a/46/76eaceaad1f42c1e7e6a5b49a61aac40fc5c9bee4b14a1630f056ac3a57e/grpcio-1.75.0-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:7a9337ac4ce61c388e02019d27fa837496c4b7837cbbcec71b05934337e51531", size = 7060578, upload-time = "2025-09-16T09:19:12.283Z" }, - { url = "https://files.pythonhosted.org/packages/3d/82/181a0e3f1397b6d43239e95becbeb448563f236c0db11ce990f073b08d01/grpcio-1.75.0-cp312-cp312-musllinux_1_2_i686.whl", hash = "sha256:ee16e232e3d0974750ab5f4da0ab92b59d6473872690b5e40dcec9a22927f22e", size = 8003283, upload-time = "2025-09-16T09:19:15.601Z" }, - { url = "https://files.pythonhosted.org/packages/de/09/a335bca211f37a3239be4b485e3c12bf3da68d18b1f723affdff2b9e9680/grpcio-1.75.0-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:55dfb9122973cc69520b23d39867726722cafb32e541435707dc10249a1bdbc6", size = 7460319, upload-time = "2025-09-16T09:19:18.409Z" }, - { url = "https://files.pythonhosted.org/packages/aa/59/6330105cdd6bc4405e74c96838cd7e148c3653ae3996e540be6118220c79/grpcio-1.75.0-cp312-cp312-win32.whl", hash = "sha256:fb64dd62face3d687a7b56cd881e2ea39417af80f75e8b36f0f81dfd93071651", size = 3934011, upload-time = "2025-09-16T09:19:21.013Z" }, - { url = "https://files.pythonhosted.org/packages/ff/14/e1309a570b7ebdd1c8ca24c4df6b8d6690009fa8e0d997cb2c026ce850c9/grpcio-1.75.0-cp312-cp312-win_amd64.whl", hash = "sha256:6b365f37a9c9543a9e91c6b4103d68d38d5bcb9965b11d5092b3c157bd6a5ee7", size = 4637934, upload-time = "2025-09-16T09:19:23.19Z" }, + { url = "https://files.pythonhosted.org/packages/5c/90/b0cbbd9efcc82816c58f31a34963071aa19fb792a212a5d9caf8e0fc3097/grpclib-0.4.9-py3-none-any.whl", hash = "sha256:7762ec1c8ed94dfad597475152dd35cbd11aecaaca2f243e29702435ca24cf0e", size = 77063, upload-time = "2025-12-14T22:23:13.224Z" }, ] [[package]] @@ -1011,19 +1199,38 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/04/4b/29cac41a4d98d144bf5f6d33995617b185d14b22401f75ca86f384e87ff1/h11-0.16.0-py3-none-any.whl", hash = "sha256:63cf8bbe7522de3bf65932fda1d9c2772064ffb3dae62d55932da54b31cb6c86", size = 37515, upload-time = "2025-04-24T03:35:24.344Z" }, ] +[[package]] +name = "h2" +version = "4.3.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "hpack", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "hyperframe", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/1d/17/afa56379f94ad0fe8defd37d6eb3f89a25404ffc71d4d848893d270325fc/h2-4.3.0.tar.gz", hash = "sha256:6c59efe4323fa18b47a632221a1888bd7fde6249819beda254aeca909f221bf1", size = 2152026, upload-time = "2025-08-23T18:12:19.778Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/69/b2/119f6e6dcbd96f9069ce9a2665e0146588dc9f88f29549711853645e736a/h2-4.3.0-py3-none-any.whl", hash = "sha256:c438f029a25f7945c69e0ccf0fb951dc3f73a5f6412981daee861431b70e2bdd", size = 61779, upload-time = "2025-08-23T18:12:17.779Z" }, +] + [[package]] name = "hf-xet" -version = "1.1.9" +version = "1.4.3" source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/23/0f/5b60fc28ee7f8cc17a5114a584fd6b86e11c3e0a6e142a7f97a161e9640a/hf_xet-1.1.9.tar.gz", hash = "sha256:c99073ce404462e909f1d5839b2d14a3827b8fe75ed8aed551ba6609c026c803", size = 484242, upload-time = "2025-08-27T23:05:19.441Z" } +sdist = { url = "https://files.pythonhosted.org/packages/53/92/ec9ad04d0b5728dca387a45af7bc98fbb0d73b2118759f5f6038b61a57e8/hf_xet-1.4.3.tar.gz", hash = "sha256:8ddedb73c8c08928c793df2f3401ec26f95be7f7e516a7bee2fbb546f6676113", size = 670477, upload-time = "2026-03-31T22:40:07.874Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/de/12/56e1abb9a44cdef59a411fe8a8673313195711b5ecce27880eb9c8fa90bd/hf_xet-1.1.9-cp37-abi3-macosx_10_12_x86_64.whl", hash = "sha256:a3b6215f88638dd7a6ff82cb4e738dcbf3d863bf667997c093a3c990337d1160", size = 2762553, upload-time = "2025-08-27T23:05:15.153Z" }, - { url = "https://files.pythonhosted.org/packages/3a/e6/2d0d16890c5f21b862f5df3146519c182e7f0ae49b4b4bf2bd8a40d0b05e/hf_xet-1.1.9-cp37-abi3-macosx_11_0_arm64.whl", hash = "sha256:9b486de7a64a66f9a172f4b3e0dfe79c9f0a93257c501296a2521a13495a698a", size = 2623216, upload-time = "2025-08-27T23:05:13.778Z" }, - { url = "https://files.pythonhosted.org/packages/81/42/7e6955cf0621e87491a1fb8cad755d5c2517803cea174229b0ec00ff0166/hf_xet-1.1.9-cp37-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:a4c5a840c2c4e6ec875ed13703a60e3523bc7f48031dfd750923b2a4d1a5fc3c", size = 3186789, upload-time = "2025-08-27T23:05:12.368Z" }, - { url = "https://files.pythonhosted.org/packages/df/8b/759233bce05457f5f7ec062d63bbfd2d0c740b816279eaaa54be92aa452a/hf_xet-1.1.9-cp37-abi3-manylinux_2_28_aarch64.whl", hash = "sha256:96a6139c9e44dad1c52c52520db0fffe948f6bce487cfb9d69c125f254bb3790", size = 3088747, upload-time = "2025-08-27T23:05:10.439Z" }, - { url = "https://files.pythonhosted.org/packages/6c/3c/28cc4db153a7601a996985bcb564f7b8f5b9e1a706c7537aad4b4809f358/hf_xet-1.1.9-cp37-abi3-musllinux_1_2_aarch64.whl", hash = "sha256:ad1022e9a998e784c97b2173965d07fe33ee26e4594770b7785a8cc8f922cd95", size = 3251429, upload-time = "2025-08-27T23:05:16.471Z" }, - { url = "https://files.pythonhosted.org/packages/84/17/7caf27a1d101bfcb05be85850d4aa0a265b2e1acc2d4d52a48026ef1d299/hf_xet-1.1.9-cp37-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:86754c2d6d5afb11b0a435e6e18911a4199262fe77553f8c50d75e21242193ea", size = 3354643, upload-time = "2025-08-27T23:05:17.828Z" }, - { url = "https://files.pythonhosted.org/packages/cd/50/0c39c9eed3411deadcc98749a6699d871b822473f55fe472fad7c01ec588/hf_xet-1.1.9-cp37-abi3-win_amd64.whl", hash = "sha256:5aad3933de6b725d61d51034e04174ed1dce7a57c63d530df0014dea15a40127", size = 2804797, upload-time = "2025-08-27T23:05:20.77Z" }, + { url = "https://files.pythonhosted.org/packages/df/9a/a24b26dc8a65f0ecc0fe5be981a19e61e7ca963b85e062c083f3a9100529/hf_xet-1.4.3-cp37-abi3-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:fc360b70c815bf340ed56c7b8c63aacf11762a4b099b2fe2c9bd6d6068668c08", size = 4212320, upload-time = "2026-03-31T22:39:42.922Z" }, + { url = "https://files.pythonhosted.org/packages/53/60/46d493db155d2ee2801b71fb1b0fd67696359047fdd8caee2c914cc50c79/hf_xet-1.4.3-cp37-abi3-manylinux_2_28_aarch64.whl", hash = "sha256:39f2d2e9654cd9b4319885733993807aab6de9dfbd34c42f0b78338d6617421f", size = 3991546, upload-time = "2026-03-31T22:39:41.335Z" }, + { url = "https://files.pythonhosted.org/packages/bc/f5/067363e1c96c6b17256910830d1b54099d06287e10f4ec6ec4e7e08371fc/hf_xet-1.4.3-cp37-abi3-musllinux_1_2_aarch64.whl", hash = "sha256:49ad8a8cead2b56051aa84d7fce3e1335efe68df3cf6c058f22a65513885baac", size = 4193200, upload-time = "2026-03-31T22:40:01.936Z" }, + { url = "https://files.pythonhosted.org/packages/42/4b/53951592882d9c23080c7644542fda34a3813104e9e11fa1a7d82d419cb8/hf_xet-1.4.3-cp37-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:7716d62015477a70ea272d2d68cd7cad140f61c52ee452e133e139abfe2c17ba", size = 4429392, upload-time = "2026-03-31T22:40:03.492Z" }, +] + +[[package]] +name = "hpack" +version = "4.1.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/2c/48/71de9ed269fdae9c8057e5a4c0aa7402e8bb16f2c6e90b3aa53327b113f8/hpack-4.1.0.tar.gz", hash = "sha256:ec5eca154f7056aa06f196a557655c5b009b382873ac8d1e66e79e87535f1dca", size = 51276, upload-time = "2025-01-22T21:44:58.347Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/07/c6/80c95b1b2b94682a72cbdbfb85b81ae2daffa4291fbfa1b1464502ede10d/hpack-4.1.0-py3-none-any.whl", hash = "sha256:157ac792668d995c657d93111f46b4535ed114f0c9c8d672271bbec7eae1b496", size = 34357, upload-time = "2025-01-22T21:44:56.92Z" }, ] [[package]] @@ -1031,8 +1238,8 @@ name = "httpcore" version = "1.0.9" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "certifi" }, - { name = "h11" }, + { name = "certifi", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "h11", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, ] sdist = { url = "https://files.pythonhosted.org/packages/06/94/82699a10bca87a5556c9c59b5963f2d039dbd239f25bc2a63907a05a14cb/httpcore-1.0.9.tar.gz", hash = "sha256:6e34463af53fd2ab5d807f399a9b45ea31c3dfa2276f15a2c3f00afff6e176e8", size = 85484, upload-time = "2025-04-24T22:06:22.219Z" } wheels = [ @@ -1045,13 +1252,10 @@ version = "0.6.4" source = { registry = "https://pypi.org/simple" } sdist = { url = "https://files.pythonhosted.org/packages/a7/9a/ce5e1f7e131522e6d3426e8e7a490b3a01f39a6696602e1c4f33f9e94277/httptools-0.6.4.tar.gz", hash = "sha256:4e93eee4add6493b59a5c514da98c939b244fce4a0d8879cd3f466562f4b7d5c", size = 240639, upload-time = "2024-10-16T19:45:08.902Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/bb/0e/d0b71465c66b9185f90a091ab36389a7352985fe857e352801c39d6127c8/httptools-0.6.4-cp312-cp312-macosx_10_13_universal2.whl", hash = "sha256:df017d6c780287d5c80601dafa31f17bddb170232d85c066604d8558683711a2", size = 200683, upload-time = "2024-10-16T19:44:30.175Z" }, - { url = "https://files.pythonhosted.org/packages/e2/b8/412a9bb28d0a8988de3296e01efa0bd62068b33856cdda47fe1b5e890954/httptools-0.6.4-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:85071a1e8c2d051b507161f6c3e26155b5c790e4e28d7f236422dbacc2a9cc44", size = 104337, upload-time = "2024-10-16T19:44:31.786Z" }, { url = "https://files.pythonhosted.org/packages/9b/01/6fb20be3196ffdc8eeec4e653bc2a275eca7f36634c86302242c4fbb2760/httptools-0.6.4-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:69422b7f458c5af875922cdb5bd586cc1f1033295aa9ff63ee196a87519ac8e1", size = 508796, upload-time = "2024-10-16T19:44:32.825Z" }, { url = "https://files.pythonhosted.org/packages/f7/d8/b644c44acc1368938317d76ac991c9bba1166311880bcc0ac297cb9d6bd7/httptools-0.6.4-cp312-cp312-manylinux_2_5_x86_64.manylinux1_x86_64.manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:16e603a3bff50db08cd578d54f07032ca1631450ceb972c2f834c2b860c28ea2", size = 510837, upload-time = "2024-10-16T19:44:33.974Z" }, { url = "https://files.pythonhosted.org/packages/52/d8/254d16a31d543073a0e57f1c329ca7378d8924e7e292eda72d0064987486/httptools-0.6.4-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:ec4f178901fa1834d4a060320d2f3abc5c9e39766953d038f1458cb885f47e81", size = 485289, upload-time = "2024-10-16T19:44:35.111Z" }, { url = "https://files.pythonhosted.org/packages/5f/3c/4aee161b4b7a971660b8be71a92c24d6c64372c1ab3ae7f366b3680df20f/httptools-0.6.4-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:f9eb89ecf8b290f2e293325c646a211ff1c2493222798bb80a530c5e7502494f", size = 489779, upload-time = "2024-10-16T19:44:36.253Z" }, - { url = "https://files.pythonhosted.org/packages/12/b7/5cae71a8868e555f3f67a50ee7f673ce36eac970f029c0c5e9d584352961/httptools-0.6.4-cp312-cp312-win_amd64.whl", hash = "sha256:db78cb9ca56b59b016e64b6031eda5653be0589dba2b1b43453f6e8b405a0970", size = 88634, upload-time = "2024-10-16T19:44:37.357Z" }, ] [[package]] @@ -1059,10 +1263,10 @@ name = "httpx" version = "0.28.1" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "anyio" }, - { name = "certifi" }, - { name = "httpcore" }, - { name = "idna" }, + { name = "anyio", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "certifi", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "httpcore", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "idna", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, ] sdist = { url = "https://files.pythonhosted.org/packages/b1/df/48c586a5fe32a0f01324ee087459e112ebb7224f646c0b5023f5e79e9956/httpx-0.28.1.tar.gz", hash = "sha256:75e98c5f16b0f35b567856f597f06ff2270a374470a5c2392242528e3e3e42fc", size = 141406, upload-time = "2024-12-06T15:37:23.222Z" } wheels = [ @@ -1080,21 +1284,31 @@ wheels = [ [[package]] name = "huggingface-hub" -version = "0.35.0rc0" +version = "1.9.0" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "filelock" }, - { name = "fsspec" }, - { name = "hf-xet", marker = "platform_machine == 'aarch64' or platform_machine == 'amd64' or platform_machine == 'arm64' or platform_machine == 'x86_64'" }, - { name = "packaging" }, - { name = "pyyaml" }, - { name = "requests" }, - { name = "tqdm" }, - { name = "typing-extensions" }, + { name = "filelock", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "fsspec", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "hf-xet", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "httpx", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "packaging", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "pyyaml", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "tqdm", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "typer", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "typing-extensions", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/3c/6b/d2234ceb0842fb36345e0c36290662f30429cfca3e9726e705f6f079e7ee/huggingface_hub-0.35.0rc0.tar.gz", hash = "sha256:1a8d599a28f8071d45e82f94b5e3437880afda153934198b91ce2d7eea724722", size = 456850, upload-time = "2025-07-24T14:02:10.538Z" } +sdist = { url = "https://files.pythonhosted.org/packages/88/bb/62c7aa86f63a05e2f9b96642fdef9b94526a23979820b09f5455deff4983/huggingface_hub-1.9.0.tar.gz", hash = "sha256:0ea5be7a56135c91797cae6ad726e38eaeb6eb4b77cefff5c9d38ba0ecf874f7", size = 750326, upload-time = "2026-04-03T08:35:55.888Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/73/37/0d15d16150e1829f3e90962c99f28257f6de9e526a680b4c6f5acdb54fd2/huggingface_hub-1.9.0-py3-none-any.whl", hash = "sha256:2999328c058d39fd19ab748dd09bd4da2fbaa4f4c1ddea823eab103051e14a1f", size = 637355, upload-time = "2026-04-03T08:35:53.897Z" }, +] + +[[package]] +name = "hyperframe" +version = "6.1.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/02/e7/94f8232d4a74cc99514c13a9f995811485a6903d48e5d952771ef6322e30/hyperframe-6.1.0.tar.gz", hash = "sha256:f630908a00854a7adeabd6382b43923a4c4cd4b821fcb527e6ab9e15382a3b08", size = 26566, upload-time = "2025-01-22T21:41:49.302Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/e4/56/4d882ae1577879616e620bc5a4c20c05bd8e4a2367ca630331c48b3193ed/huggingface_hub-0.35.0rc0-py3-none-any.whl", hash = "sha256:812de58926a5ff57b1a1393fff26f26ea8e5aa9de126da62607094018d811296", size = 558670, upload-time = "2025-07-24T14:02:08.744Z" }, + { url = "https://files.pythonhosted.org/packages/48/30/47d0bf6072f7252e6521f3447ccfa40b421b6824517f82854703d0f5a98b/hyperframe-6.1.0-py3-none-any.whl", hash = "sha256:b03380493a519fce58ea5af42e4a42317bf9bd425596f7a0835ffce80f1a42e5", size = 13007, upload-time = "2025-01-22T21:41:47.295Z" }, ] [[package]] @@ -1115,18 +1329,39 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/76/c6/c88e154df9c4e1a2a66ccf0005a88dfb2650c1dffb6f5ce603dfbd452ce3/idna-3.10-py3-none-any.whl", hash = "sha256:946d195a0d259cbba61165e88e65941f16e9b36ea6ddb97f00452bae8b1287d3", size = 70442, upload-time = "2024-09-15T18:07:37.964Z" }, ] +[[package]] +name = "ijson" +version = "3.4.0.post0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/2d/30/7ab4b9e88e7946f6beef419f74edcc541df3ea562c7882257b4eaa82417d/ijson-3.4.0.post0.tar.gz", hash = "sha256:9aa02dc70bb245670a6ca7fba737b992aeeb4895360980622f7e568dbf23e41e", size = 67216, upload-time = "2025-10-10T05:29:25.62Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/c3/d3/22e3cc806fcdda7ad4c8482ed74db7a017d4a1d49b4300c7bc07052fb561/ijson-3.4.0.post0-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:54a0e3e05d9a0c95ecba73d9579f146cf6d5c5874116c849dba2d39a5f30380e", size = 149110, upload-time = "2025-10-10T05:28:12.263Z" }, + { url = "https://files.pythonhosted.org/packages/3e/04/efb30f413648b9267f5a33920ac124d7ebef3bc4063af8f6ffc8ca11ddcb/ijson-3.4.0.post0-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:05807edc0bcbd222dc6ea32a2b897f0c81dc7f12c8580148bc82f6d7f5e7ec7b", size = 149026, upload-time = "2025-10-10T05:28:13.557Z" }, + { url = "https://files.pythonhosted.org/packages/2d/cf/481165f7046ade32488719300a3994a437020bc41cfbb54334356348f513/ijson-3.4.0.post0-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:a5269af16f715855d9864937f9dd5c348ca1ac49cee6a2c7a1b7091c159e874f", size = 150012, upload-time = "2025-10-10T05:28:14.859Z" }, + { url = "https://files.pythonhosted.org/packages/0f/f5/fd2f038abe95e553e1c3ee207cda19db9196eb416e63c7c89699a8cf0db7/ijson-3.4.0.post0-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:6458bd8e679cdff459a0a5e555b107c3bbacb1f382da3fe0f40e392871eb518d", size = 150904, upload-time = "2025-10-10T05:28:17.401Z" }, +] + [[package]] name = "importlib-metadata" version = "8.7.0" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "zipp" }, + { name = "zipp", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, ] sdist = { url = "https://files.pythonhosted.org/packages/76/66/650a33bd90f786193e4de4b3ad86ea60b53c89b669a5c7be931fac31cdb0/importlib_metadata-8.7.0.tar.gz", hash = "sha256:d13b81ad223b890aa16c5471f2ac3056cf76c5f10f82d6f9292f0b415f389000", size = 56641, upload-time = "2025-04-27T15:29:01.736Z" } wheels = [ { url = "https://files.pythonhosted.org/packages/20/b0/36bd937216ec521246249be3bf9855081de4c5e06a0c9b4219dbeda50373/importlib_metadata-8.7.0-py3-none-any.whl", hash = "sha256:e5dd1551894c77868a30651cef00984d50e1002d06942a7101d34870c5f02afd", size = 27656, upload-time = "2025-04-27T15:29:00.214Z" }, ] +[[package]] +name = "importlib-resources" +version = "6.5.2" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/cf/8c/f834fbf984f691b4f7ff60f50b514cc3de5cc08abfc3295564dd89c5e2e7/importlib_resources-6.5.2.tar.gz", hash = "sha256:185f87adef5bcc288449d98fb4fba07cea78bc036455dd44c5fc4a2fe78fed2c", size = 44693, upload-time = "2025-01-03T18:51:56.698Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/a4/ed/1f1afb2e9e7f38a545d628f864d562a5ae64fe6f7a10e28ffb9b185b4e89/importlib_resources-6.5.2-py3-none-any.whl", hash = "sha256:789cfdc3ed28c78b67a06acb8126751ced69a3d5f79c095a98298cd8a760ccec", size = 37461, upload-time = "2025-01-03T18:51:54.306Z" }, +] + [[package]] name = "iniconfig" version = "2.1.0" @@ -1150,19 +1385,18 @@ name = "ipykernel" version = "7.0.0a2" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "appnope", marker = "sys_platform == 'darwin'" }, - { name = "comm" }, - { name = "debugpy" }, - { name = "ipython" }, - { name = "jupyter-client" }, - { name = "jupyter-core" }, - { name = "matplotlib-inline" }, - { name = "nest-asyncio" }, - { name = "packaging" }, - { name = "psutil" }, - { name = "pyzmq" }, - { name = "tornado" }, - { name = "traitlets" }, + { name = "comm", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "debugpy", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "ipython", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "jupyter-client", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "jupyter-core", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "matplotlib-inline", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "nest-asyncio", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "packaging", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "psutil", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "pyzmq", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "tornado", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "traitlets", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, ] sdist = { url = "https://files.pythonhosted.org/packages/c1/1a/baaccb506ef660e9e3d420088e3a68f7873b1bcc152c0d069541641fd230/ipykernel-7.0.0a2.tar.gz", hash = "sha256:b3bd996d94ca09012b6df6654a086ddb3a01a40bcd1ad8cc1f5e80b327e0f120", size = 170826, upload-time = "2025-08-13T09:51:57.264Z" } wheels = [ @@ -1174,16 +1408,15 @@ name = "ipython" version = "9.5.0" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "colorama", marker = "sys_platform == 'win32'" }, - { name = "decorator" }, - { name = "ipython-pygments-lexers" }, - { name = "jedi" }, - { name = "matplotlib-inline" }, - { name = "pexpect", marker = "sys_platform != 'emscripten' and sys_platform != 'win32'" }, - { name = "prompt-toolkit" }, - { name = "pygments" }, - { name = "stack-data" }, - { name = "traitlets" }, + { name = "decorator", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "ipython-pygments-lexers", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "jedi", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "matplotlib-inline", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "pexpect", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "prompt-toolkit", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "pygments", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "stack-data", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "traitlets", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, ] sdist = { url = "https://files.pythonhosted.org/packages/6e/71/a86262bf5a68bf211bcc71fe302af7e05f18a2852fdc610a854d20d085e6/ipython-9.5.0.tar.gz", hash = "sha256:129c44b941fe6d9b82d36fc7a7c18127ddb1d6f02f78f867f402e2e3adde3113", size = 4389137, upload-time = "2025-08-29T12:15:21.519Z" } wheels = [ @@ -1195,7 +1428,7 @@ name = "ipython-pygments-lexers" version = "1.1.1" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "pygments" }, + { name = "pygments", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, ] sdist = { url = "https://files.pythonhosted.org/packages/ef/4c/5dd1d8af08107f88c7f741ead7a40854b8ac24ddf9ae850afbcf698aa552/ipython_pygments_lexers-1.1.1.tar.gz", hash = "sha256:09c0138009e56b6854f9535736f4171d855c8c08a563a0dcd8022f78355c7e81", size = 8393, upload-time = "2025-01-17T11:24:34.505Z" } wheels = [ @@ -1207,11 +1440,11 @@ name = "ipywidgets" version = "8.1.7" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "comm" }, - { name = "ipython" }, - { name = "jupyterlab-widgets" }, - { name = "traitlets" }, - { name = "widgetsnbextension" }, + { name = "comm", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "ipython", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "jupyterlab-widgets", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "traitlets", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "widgetsnbextension", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, ] sdist = { url = "https://files.pythonhosted.org/packages/3e/48/d3dbac45c2814cb73812f98dd6b38bbcc957a4e7bb31d6ea9c03bf94ed87/ipywidgets-8.1.7.tar.gz", hash = "sha256:15f1ac050b9ccbefd45dccfbb2ef6bed0029d8278682d569d71b8dd96bee0376", size = 116721, upload-time = "2025-05-05T12:42:03.489Z" } wheels = [ @@ -1223,7 +1456,7 @@ name = "jaxtyping" version = "0.3.2" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "wadler-lindig" }, + { name = "wadler-lindig", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, ] sdist = { url = "https://files.pythonhosted.org/packages/88/a8/416bd7ea110ec6b68e8868b0f99c985c735adcf7badc491d3c343937260a/jaxtyping-0.3.2.tar.gz", hash = "sha256:f30483fac4b42e915db8ad2414a85c3b63284aa7d3c100b96b59f755cf4a86ad", size = 44989, upload-time = "2025-04-23T09:48:16.907Z" } wheels = [ @@ -1235,7 +1468,7 @@ name = "jedi" version = "0.19.2" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "parso" }, + { name = "parso", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, ] sdist = { url = "https://files.pythonhosted.org/packages/72/3a/79a912fbd4d8dd6fbb02bf69afd3bb72cf0c729bb3063c6f4498603db17a/jedi-0.19.2.tar.gz", hash = "sha256:4770dc3de41bde3966b02eb84fbcf557fb33cce26ad23da12c742fb50ecb11f0", size = 1231287, upload-time = "2024-11-11T01:41:42.873Z" } wheels = [ @@ -1247,7 +1480,7 @@ name = "jinja2" version = "3.1.6" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "markupsafe" }, + { name = "markupsafe", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, ] sdist = { url = "https://files.pythonhosted.org/packages/df/bf/f7da0350254c0ed7c72f3e33cef02e048281fec7ecec5f032d4aac52226b/jinja2-3.1.6.tar.gz", hash = "sha256:0137fb05990d35f1275a587e9aee6d56da821fc83491a0fb838183be43f66d6d", size = 245115, upload-time = "2025-03-05T20:05:02.478Z" } wheels = [ @@ -1260,18 +1493,10 @@ version = "0.10.0" source = { registry = "https://pypi.org/simple" } sdist = { url = "https://files.pythonhosted.org/packages/ee/9d/ae7ddb4b8ab3fb1b51faf4deb36cb48a4fbbd7cb36bad6a5fca4741306f7/jiter-0.10.0.tar.gz", hash = "sha256:07a7142c38aacc85194391108dc91b5b57093c978a9932bd86a36862759d9500", size = 162759, upload-time = "2025-05-18T19:04:59.73Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/6d/b5/348b3313c58f5fbfb2194eb4d07e46a35748ba6e5b3b3046143f3040bafa/jiter-0.10.0-cp312-cp312-macosx_10_12_x86_64.whl", hash = "sha256:1e274728e4a5345a6dde2d343c8da018b9d4bd4350f5a472fa91f66fda44911b", size = 312262, upload-time = "2025-05-18T19:03:44.637Z" }, - { url = "https://files.pythonhosted.org/packages/9c/4a/6a2397096162b21645162825f058d1709a02965606e537e3304b02742e9b/jiter-0.10.0-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:7202ae396446c988cb2a5feb33a543ab2165b786ac97f53b59aafb803fef0744", size = 320124, upload-time = "2025-05-18T19:03:46.341Z" }, { url = "https://files.pythonhosted.org/packages/2a/85/1ce02cade7516b726dd88f59a4ee46914bf79d1676d1228ef2002ed2f1c9/jiter-0.10.0-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:23ba7722d6748b6920ed02a8f1726fb4b33e0fd2f3f621816a8b486c66410ab2", size = 345330, upload-time = "2025-05-18T19:03:47.596Z" }, - { url = "https://files.pythonhosted.org/packages/75/d0/bb6b4f209a77190ce10ea8d7e50bf3725fc16d3372d0a9f11985a2b23eff/jiter-0.10.0-cp312-cp312-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:371eab43c0a288537d30e1f0b193bc4eca90439fc08a022dd83e5e07500ed026", size = 369670, upload-time = "2025-05-18T19:03:49.334Z" }, - { url = "https://files.pythonhosted.org/packages/a0/f5/a61787da9b8847a601e6827fbc42ecb12be2c925ced3252c8ffcb56afcaf/jiter-0.10.0-cp312-cp312-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:6c675736059020365cebc845a820214765162728b51ab1e03a1b7b3abb70f74c", size = 489057, upload-time = "2025-05-18T19:03:50.66Z" }, - { url = "https://files.pythonhosted.org/packages/12/e4/6f906272810a7b21406c760a53aadbe52e99ee070fc5c0cb191e316de30b/jiter-0.10.0-cp312-cp312-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:0c5867d40ab716e4684858e4887489685968a47e3ba222e44cde6e4a2154f959", size = 389372, upload-time = "2025-05-18T19:03:51.98Z" }, { url = "https://files.pythonhosted.org/packages/e2/ba/77013b0b8ba904bf3762f11e0129b8928bff7f978a81838dfcc958ad5728/jiter-0.10.0-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:395bb9a26111b60141757d874d27fdea01b17e8fac958b91c20128ba8f4acc8a", size = 352038, upload-time = "2025-05-18T19:03:53.703Z" }, - { url = "https://files.pythonhosted.org/packages/67/27/c62568e3ccb03368dbcc44a1ef3a423cb86778a4389e995125d3d1aaa0a4/jiter-0.10.0-cp312-cp312-manylinux_2_5_i686.manylinux1_i686.whl", hash = "sha256:6842184aed5cdb07e0c7e20e5bdcfafe33515ee1741a6835353bb45fe5d1bd95", size = 391538, upload-time = "2025-05-18T19:03:55.046Z" }, { url = "https://files.pythonhosted.org/packages/c0/72/0d6b7e31fc17a8fdce76164884edef0698ba556b8eb0af9546ae1a06b91d/jiter-0.10.0-cp312-cp312-musllinux_1_1_aarch64.whl", hash = "sha256:62755d1bcea9876770d4df713d82606c8c1a3dca88ff39046b85a048566d56ea", size = 523557, upload-time = "2025-05-18T19:03:56.386Z" }, { url = "https://files.pythonhosted.org/packages/2f/09/bc1661fbbcbeb6244bd2904ff3a06f340aa77a2b94e5a7373fd165960ea3/jiter-0.10.0-cp312-cp312-musllinux_1_1_x86_64.whl", hash = "sha256:533efbce2cacec78d5ba73a41756beff8431dfa1694b6346ce7af3a12c42202b", size = 514202, upload-time = "2025-05-18T19:03:57.675Z" }, - { url = "https://files.pythonhosted.org/packages/1b/84/5a5d5400e9d4d54b8004c9673bbe4403928a00d28529ff35b19e9d176b19/jiter-0.10.0-cp312-cp312-win32.whl", hash = "sha256:8be921f0cadd245e981b964dfbcd6fd4bc4e254cdc069490416dd7a2632ecc01", size = 211781, upload-time = "2025-05-18T19:03:59.025Z" }, - { url = "https://files.pythonhosted.org/packages/9b/52/7ec47455e26f2d6e5f2ea4951a0652c06e5b995c291f723973ae9e724a65/jiter-0.10.0-cp312-cp312-win_amd64.whl", hash = "sha256:a7c7d785ae9dda68c2678532a5a1581347e9c15362ae9f6e68f3fdbfb64f2e49", size = 206176, upload-time = "2025-05-18T19:04:00.305Z" }, ] [[package]] @@ -1283,24 +1508,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/31/b4/b9b800c45527aadd64d5b442f9b932b00648617eb5d63d2c7a6587b7cafc/jmespath-1.0.1-py3-none-any.whl", hash = "sha256:02e2e4cc71b5bcab88332eebf907519190dd9e6e82107fa7f83b1003a6252980", size = 20256, upload-time = "2022-06-17T18:00:10.251Z" }, ] -[[package]] -name = "joblib" -version = "1.5.2" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/e8/5d/447af5ea094b9e4c4054f82e223ada074c552335b9b4b2d14bd9b35a67c4/joblib-1.5.2.tar.gz", hash = "sha256:3faa5c39054b2f03ca547da9b2f52fde67c06240c31853f306aea97f13647b55", size = 331077, upload-time = "2025-08-27T12:15:46.575Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/1e/e8/685f47e0d754320684db4425a0967f7d3fa70126bffd76110b7009a0090f/joblib-1.5.2-py3-none-any.whl", hash = "sha256:4e1f0bdbb987e6d843c70cf43714cb276623def372df3c22fe5266b2670bc241", size = 308396, upload-time = "2025-08-27T12:15:45.188Z" }, -] - [[package]] name = "jsonschema" version = "4.25.1" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "attrs" }, - { name = "jsonschema-specifications" }, - { name = "referencing" }, - { name = "rpds-py" }, + { name = "attrs", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "jsonschema-specifications", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "referencing", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "rpds-py", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, ] sdist = { url = "https://files.pythonhosted.org/packages/74/69/f7185de793a29082a9f3c7728268ffb31cb5095131a9c139a74078e27336/jsonschema-4.25.1.tar.gz", hash = "sha256:e4a9655ce0da0c0b67a085847e00a3a51449e1157f4f75e9fb5aa545e122eb85", size = 357342, upload-time = "2025-08-18T17:03:50.038Z" } wheels = [ @@ -1312,7 +1528,7 @@ name = "jsonschema-specifications" version = "2025.9.1" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "referencing" }, + { name = "referencing", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, ] sdist = { url = "https://files.pythonhosted.org/packages/19/74/a633ee74eb36c44aa6d1095e7cc5569bebf04342ee146178e2d36600708b/jsonschema_specifications-2025.9.1.tar.gz", hash = "sha256:b540987f239e745613c7a9176f3edb72b832a4ac465cf02712288397832b5e8d", size = 32855, upload-time = "2025-09-08T01:34:59.186Z" } wheels = [ @@ -1324,11 +1540,11 @@ name = "jupyter-client" version = "8.6.3" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "jupyter-core" }, - { name = "python-dateutil" }, - { name = "pyzmq" }, - { name = "tornado" }, - { name = "traitlets" }, + { name = "jupyter-core", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "python-dateutil", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "pyzmq", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "tornado", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "traitlets", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, ] sdist = { url = "https://files.pythonhosted.org/packages/71/22/bf9f12fdaeae18019a468b68952a60fe6dbab5d67cd2a103cac7659b41ca/jupyter_client-8.6.3.tar.gz", hash = "sha256:35b3a0947c4a6e9d589eb97d7d4cd5e90f910ee73101611f01283732bd6d9419", size = 342019, upload-time = "2024-09-17T10:44:17.613Z" } wheels = [ @@ -1340,9 +1556,8 @@ name = "jupyter-core" version = "5.8.1" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "platformdirs" }, - { name = "pywin32", marker = "platform_python_implementation != 'PyPy' and sys_platform == 'win32'" }, - { name = "traitlets" }, + { name = "platformdirs", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "traitlets", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, ] sdist = { url = "https://files.pythonhosted.org/packages/99/1b/72906d554acfeb588332eaaa6f61577705e9ec752ddb486f302dafa292d9/jupyter_core-5.8.1.tar.gz", hash = "sha256:0a5f9706f70e64786b75acba995988915ebd4601c8a52e534a40b51c95f59941", size = 88923, upload-time = "2025-05-27T07:38:16.655Z" } wheels = [ @@ -1359,24 +1574,37 @@ wheels = [ ] [[package]] -name = "kiwisolver" -version = "1.4.10rc0" +name = "kernels" +version = "0.12.3" source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/bf/de/354c903d772c1cc0a9310344e077b31c6c893cc5a664019b907a04997099/kiwisolver-1.4.10rc0.tar.gz", hash = "sha256:d321718aaa2583577be9836e8cc0ed9fd0863e57a85b1b73b328aac063bc9903", size = 97614, upload-time = "2025-08-10T20:22:27.702Z" } +dependencies = [ + { name = "huggingface-hub", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "packaging", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "pyyaml", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/b3/84/9f68f355f6ce99e977872021fbdbafadcf2820f51d3f7bd697ec3801cb7a/kernels-0.12.3.tar.gz", hash = "sha256:87e29716578e7e71dc5a7578e0132bfdae305bedaeb602698f87c88ca6c60e32", size = 57407, upload-time = "2026-03-20T10:20:42.166Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/e7/3e/778e4a86830e9139df2d16d86c4488fce426ec19daa83cbd2854ef389030/kernels-0.12.3-py3-none-any.whl", hash = "sha256:5d1d33fcb774e03bb7f0688ac24d91ef6b963692f80f0a85ddd2286e69f3cf2f", size = 55501, upload-time = "2026-03-20T10:20:40.643Z" }, +] + +[[package]] +name = "kubernetes" +version = "35.0.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "certifi", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "durationpy", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "python-dateutil", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "pyyaml", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "requests", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "requests-oauthlib", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "six", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "urllib3", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "websocket-client", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/2c/8f/85bf51ad4150f64e8c665daf0d9dfe9787ae92005efb9a4d1cba592bd79d/kubernetes-35.0.0.tar.gz", hash = "sha256:3d00d344944239821458b9efd484d6df9f011da367ecb155dadf9513f05f09ee", size = 1094642, upload-time = "2026-01-16T01:05:27.76Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/dc/87/3df31abf12db3ccabfa52a96dc49e6defe233d8ffca1361091a1ea3a109c/kiwisolver-1.4.10rc0-cp312-cp312-macosx_10_13_universal2.whl", hash = "sha256:1cb9ae443b2dba2229ac3b8a771420ee76bfce56f610dcb4998676cebed79346", size = 123742, upload-time = "2025-08-10T20:20:44.391Z" }, - { url = "https://files.pythonhosted.org/packages/7b/62/fc9adfd88082b95971969736d777762f7940f3d49f5ffae37c439699156d/kiwisolver-1.4.10rc0-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:2727d842ff63050315d2a69334d64ed4ada261c04155c09d9534fd4033d38641", size = 66522, upload-time = "2025-08-10T20:20:45.81Z" }, - { url = "https://files.pythonhosted.org/packages/7d/d1/3802735c705ffa861bbb568ed4226936fcfc917a179bf3998fdf97e48e57/kiwisolver-1.4.10rc0-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:f9563b4c98c23f52d98e15be40cbf0c215c824e7268d5222e9ad1803e8c1156d", size = 65016, upload-time = "2025-08-10T20:20:46.911Z" }, - { url = "https://files.pythonhosted.org/packages/72/76/29b4d717f5614fc91cb4542cd67f42b5bcb6c946201a9cf9d4a34231efc2/kiwisolver-1.4.10rc0-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:75cef7209e3c71dc81d1d5bc8d5eb1b34be7a8d2cd94b83f0eb15533512a45ef", size = 1474820, upload-time = "2025-08-10T20:20:48.447Z" }, - { url = "https://files.pythonhosted.org/packages/f1/9b/828761ad3841b0e7d464514939a980eef7547b5952fb9a33232b17ed1540/kiwisolver-1.4.10rc0-cp312-cp312-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:374c11179eaacd3b8bfef677aa28a0a6d703b3474ea399f3b08b8e4d67522016", size = 1276468, upload-time = "2025-08-10T20:20:49.775Z" }, - { url = "https://files.pythonhosted.org/packages/92/6d/cb780c0ebad56e2b7b3c2c0376e9d5c25e90680d2b3b254afaec2507a62f/kiwisolver-1.4.10rc0-cp312-cp312-manylinux_2_24_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:9abc647f46a322fd19e0564ce59fcc0f14b9934540ddc7404ed7f3eea54d0d11", size = 1294484, upload-time = "2025-08-10T20:20:51.207Z" }, - { url = "https://files.pythonhosted.org/packages/55/70/44814b447c4e38da6f466b6a3e992d330c3e2c1c9c29731c436997b78f68/kiwisolver-1.4.10rc0-cp312-cp312-manylinux_2_24_s390x.manylinux_2_28_s390x.whl", hash = "sha256:f7e3c788da96b50e60f484d9dd790554a80800c22a468cd059b9a7a9c753d273", size = 1343677, upload-time = "2025-08-10T20:20:52.906Z" }, - { url = "https://files.pythonhosted.org/packages/5c/54/44eee9dc53be9c4c6ac3b099aedc482f1a1a6b193d0f258ccfa955c291df/kiwisolver-1.4.10rc0-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:468814c0e8f41b8f7b537da0c77a05a70f89aa4b7cfff96aa47f7936e57add9b", size = 2225010, upload-time = "2025-08-10T20:20:54.365Z" }, - { url = "https://files.pythonhosted.org/packages/50/5e/e05e24f858352e6985ace1f9ab8ece32b0962f4c5074ddb38fc91617809a/kiwisolver-1.4.10rc0-cp312-cp312-musllinux_1_2_ppc64le.whl", hash = "sha256:aa62a0aa711b1f94f80f1407a668362718c64f27e176ac6952d983ec1a5cd745", size = 2321356, upload-time = "2025-08-10T20:20:55.803Z" }, - { url = "https://files.pythonhosted.org/packages/f1/a7/1d21b9aa468bdbd4be190043d73ad07756c078003e43fdc4af5ccb3df75a/kiwisolver-1.4.10rc0-cp312-cp312-musllinux_1_2_s390x.whl", hash = "sha256:6e8697067105c536bcc14a33f5c9c0f0c155cf909818d01475f45d012ac441c1", size = 2488062, upload-time = "2025-08-10T20:20:57.371Z" }, - { url = "https://files.pythonhosted.org/packages/88/9b/17512a47070d022499f19078b980531b7be5d50eb9990dfc4ec29aa554ca/kiwisolver-1.4.10rc0-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:4c168de06cd8a4947a8af1e49d63341face78aca8e9e6b7685625701147ab22d", size = 2291890, upload-time = "2025-08-10T20:20:59.237Z" }, - { url = "https://files.pythonhosted.org/packages/85/c1/084d9b537e33555d8bf5d41ffaee88cf0ee49fa42587fdee181d31a40b61/kiwisolver-1.4.10rc0-cp312-cp312-win_amd64.whl", hash = "sha256:0b8bb7b6b3964d0454f8504e003097f2ae628679a1054ecb63578feeb7671cab", size = 73949, upload-time = "2025-08-10T20:21:00.845Z" }, - { url = "https://files.pythonhosted.org/packages/ff/fa/538442202d639add2f52a814bdfc58207ee6fbb6d1ecd1a6e867f48ec1af/kiwisolver-1.4.10rc0-cp312-cp312-win_arm64.whl", hash = "sha256:44cd6dfea8a6c2becac4f3d60ebdcfe4fed858bbf7fe9cd38ffea7b58df66435", size = 65077, upload-time = "2025-08-10T20:21:01.882Z" }, + { url = "https://files.pythonhosted.org/packages/0c/70/05b685ea2dffcb2adbf3cdcea5d8865b7bc66f67249084cf845012a0ff13/kubernetes-35.0.0-py2.py3-none-any.whl", hash = "sha256:39e2b33b46e5834ef6c3985ebfe2047ab39135d41de51ce7641a7ca5b372a13d", size = 2017602, upload-time = "2026-01-16T01:05:25.991Z" }, ] [[package]] @@ -1393,8 +1621,8 @@ name = "latex2sympy2-extended" version = "1.10.2" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "antlr4-python3-runtime" }, - { name = "sympy" }, + { name = "antlr4-python3-runtime", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "sympy", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, ] sdist = { url = "https://files.pythonhosted.org/packages/f4/de/472f9115c14c6f6d8a5889cabe3418283d708bde62ce00402c29441deed4/latex2sympy2_extended-1.10.2.tar.gz", hash = "sha256:41a517ffcc5a140e910a7d1646ce6ff440817e5f9d48fc8279d88bd0925bc389", size = 206188, upload-time = "2025-07-02T15:26:06.225Z" } wheels = [ @@ -1406,9 +1634,8 @@ name = "liger-kernel" version = "0.6.2" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "torch" }, - { name = "triton", version = "3.3.1", source = { registry = "https://pypi.org/simple" }, marker = "sys_platform != 'linux'" }, - { name = "triton", version = "3.5.0", source = { registry = "https://pypi.org/simple" }, marker = "sys_platform == 'linux'" }, + { name = "torch", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "triton", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, ] sdist = { url = "https://files.pythonhosted.org/packages/31/23/be0b4dcac42d77f99406c906567cde22a7a3d71b3f3ffdfda2ac6153ec36/liger_kernel-0.6.2.tar.gz", hash = "sha256:5c5bcffffa769bc26ae838f5a4954170dd5cacde036abb1b383039f39fa5fd69", size = 3679495, upload-time = "2025-08-22T00:15:28.456Z" } wheels = [ @@ -1420,7 +1647,7 @@ name = "linkify-it-py" version = "2.0.3" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "uc-micro-py" }, + { name = "uc-micro-py", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, ] sdist = { url = "https://files.pythonhosted.org/packages/2a/ae/bb56c6828e4797ba5a4821eec7c43b8bf40f69cda4d4f5f8c8a2810ec96a/linkify-it-py-2.0.3.tar.gz", hash = "sha256:68cda27e162e9215c17d786649d1da0021a451bdc436ef9e0fa0ba5234b9b048", size = 27946, upload-time = "2024-02-04T14:48:04.179Z" } wheels = [ @@ -1433,24 +1660,18 @@ version = "1.3.0" source = { registry = "https://pypi.org/simple" } sdist = { url = "https://files.pythonhosted.org/packages/95/48/3f7a9d3ff1b36bba92b5107a3a21286821227afe9ea464736133994d61fb/llguidance-1.3.0.tar.gz", hash = "sha256:861249afd51dc325646834462ea827e57a5c2b2042e108e6aae7059fdad9104d", size = 1070460, upload-time = "2025-10-20T19:58:44.164Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/3b/33/be5acb85cd8cdc4afde33d9c234eece9f318e087920255af3c05864cd3e7/llguidance-1.3.0-cp39-abi3-macosx_10_12_x86_64.whl", hash = "sha256:f7685222660a762e481ac633d49cc559c64980fe2ee59c8f932a5bb5cbc0c2c2", size = 3220647, upload-time = "2025-10-20T19:58:42.542Z" }, - { url = "https://files.pythonhosted.org/packages/82/e6/b48bda5b15efeaeb62bd0dba8fc6a01d4ae5457a85dbb5d18632385fe15c/llguidance-1.3.0-cp39-abi3-macosx_11_0_arm64.whl", hash = "sha256:098030ff0687261a3f1bd54cf21fe951fc861d56d37a0671250dd36677eaf224", size = 3099830, upload-time = "2025-10-20T19:58:40.826Z" }, { url = "https://files.pythonhosted.org/packages/aa/11/44389d3d1526d7a5c38ffd587a5ebc61d7bee443ac1dea95f2089ad58f5f/llguidance-1.3.0-cp39-abi3-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:6f6caca5d78db7f76e1fbb0fff8607b861c32d47fa3d5dee2fc49de27ee269df", size = 2835242, upload-time = "2025-10-20T19:58:34.518Z" }, { url = "https://files.pythonhosted.org/packages/83/a8/1ff2bedb8f9acb46a2d2d603415d272bb622c142ea86f5b95445cc6e366c/llguidance-1.3.0-cp39-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:bc17e9dd602c3879bf91664a64bf72f54c74dbfbeb24ccfab6a5fe435b12f7aa", size = 3033133, upload-time = "2025-10-20T19:58:38.721Z" }, - { url = "https://files.pythonhosted.org/packages/5a/7e/809349638231f469b9056c0e1bfd924d5ef5558b3b3ec72d093b6fad33b1/llguidance-1.3.0-cp39-abi3-win_amd64.whl", hash = "sha256:1d1cd1c8618d1a13605d3e057c978651e551c8c469b481ee4041f1d6c436002d", size = 2789946, upload-time = "2025-10-20T19:58:45.958Z" }, ] [[package]] name = "llvmlite" -version = "0.44.0" +version = "0.47.0" source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/89/6a/95a3d3610d5c75293d5dbbb2a76480d5d4eeba641557b69fe90af6c5b84e/llvmlite-0.44.0.tar.gz", hash = "sha256:07667d66a5d150abed9157ab6c0b9393c9356f229784a4385c02f99e94fc94d4", size = 171880, upload-time = "2025-01-20T11:14:41.342Z" } +sdist = { url = "https://files.pythonhosted.org/packages/01/88/a8952b6d5c21e74cbf158515b779666f692846502623e9e3c39d8e8ba25f/llvmlite-0.47.0.tar.gz", hash = "sha256:62031ce968ec74e95092184d4b0e857e444f8fdff0b8f9213707699570c33ccc", size = 193614, upload-time = "2026-03-31T18:29:53.497Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/15/86/e3c3195b92e6e492458f16d233e58a1a812aa2bfbef9bdd0fbafcec85c60/llvmlite-0.44.0-cp312-cp312-macosx_10_14_x86_64.whl", hash = "sha256:1d671a56acf725bf1b531d5ef76b86660a5ab8ef19bb6a46064a705c6ca80aad", size = 28132297, upload-time = "2025-01-20T11:13:32.57Z" }, - { url = "https://files.pythonhosted.org/packages/d6/53/373b6b8be67b9221d12b24125fd0ec56b1078b660eeae266ec388a6ac9a0/llvmlite-0.44.0-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:5f79a728e0435493611c9f405168682bb75ffd1fbe6fc360733b850c80a026db", size = 26201105, upload-time = "2025-01-20T11:13:38.744Z" }, - { url = "https://files.pythonhosted.org/packages/cb/da/8341fd3056419441286c8e26bf436923021005ece0bff5f41906476ae514/llvmlite-0.44.0-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:c0143a5ef336da14deaa8ec26c5449ad5b6a2b564df82fcef4be040b9cacfea9", size = 42361901, upload-time = "2025-01-20T11:13:46.711Z" }, - { url = "https://files.pythonhosted.org/packages/53/ad/d79349dc07b8a395a99153d7ce8b01d6fcdc9f8231355a5df55ded649b61/llvmlite-0.44.0-cp312-cp312-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:d752f89e31b66db6f8da06df8b39f9b91e78c5feea1bf9e8c1fba1d1c24c065d", size = 41184247, upload-time = "2025-01-20T11:13:56.159Z" }, - { url = "https://files.pythonhosted.org/packages/e2/3b/a9a17366af80127bd09decbe2a54d8974b6d8b274b39bf47fbaedeec6307/llvmlite-0.44.0-cp312-cp312-win_amd64.whl", hash = "sha256:eae7e2d4ca8f88f89d315b48c6b741dcb925d6a1042da694aa16ab3dd4cbd3a1", size = 30332380, upload-time = "2025-01-20T11:14:02.442Z" }, + { url = "https://files.pythonhosted.org/packages/e6/4b/e3f2cd17822cf772a4a51a0a8080b0032e6d37b2dbe8cfb724eac4e31c52/llvmlite-0.47.0-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:5853bf26160857c0c2573415ff4efe01c4c651e59e2c55c2a088740acfee51cd", size = 56275178, upload-time = "2026-03-31T18:28:48.342Z" }, + { url = "https://files.pythonhosted.org/packages/b6/55/a3b4a543185305a9bdf3d9759d53646ed96e55e7dfd43f53e7a421b8fbae/llvmlite-0.47.0-cp312-cp312-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:003bcf7fa579e14db59c1a1e113f93ab8a06b56a4be31c7f08264d1d4072d077", size = 55128632, upload-time = "2026-03-31T18:28:52.901Z" }, ] [[package]] @@ -1458,10 +1679,10 @@ name = "lm-format-enforcer" version = "0.11.3" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "interegular" }, - { name = "packaging" }, - { name = "pydantic" }, - { name = "pyyaml" }, + { name = "interegular", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "packaging", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "pydantic", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "pyyaml", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, ] sdist = { url = "https://files.pythonhosted.org/packages/84/d5/41cd417ba7dfdbbcfe46cebf81fb3dfd7c591b89897560ad05bb410a465d/lm_format_enforcer-0.11.3.tar.gz", hash = "sha256:e68081c108719cce284a9bcc889709b26ffb085a1945b5eba3a12cfa96d528da", size = 40258, upload-time = "2025-08-24T19:37:47.527Z" } wheels = [ @@ -1469,71 +1690,42 @@ wheels = [ ] [[package]] -name = "loguru" -version = "0.7.3" -source = { registry = "https://pypi.org/simple" } +name = "logic-env" +version = "0.1.0" +source = { registry = "https://hub.primeintellect.ai/primeintellect/simple/" } dependencies = [ - { name = "colorama", marker = "sys_platform == 'win32'" }, - { name = "win32-setctime", marker = "sys_platform == 'win32'" }, + { name = "markdown", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "math-verify", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "sympy", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "verifiers", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/3a/05/a1dae3dffd1116099471c643b8924f5aa6524411dc6c63fdae648c4f1aca/loguru-0.7.3.tar.gz", hash = "sha256:19480589e77d47b8d85b2c827ad95d49bf31b0dcde16593892eb51dd18706eb6", size = 63559, upload-time = "2024-12-06T11:20:56.608Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/0c/29/0348de65b8cc732daa3e33e67806420b2ae89bdce2b04af740289c5c6c8c/loguru-0.7.3-py3-none-any.whl", hash = "sha256:31a33c10c8e1e10422bfd431aeb5d351c7cf7fa671e3c4df004162264b28220c", size = 61595, upload-time = "2024-12-06T11:20:54.538Z" }, + { url = "https://hub.primeintellect.ai/primeintellect/logic-env/@7bded58d/logic_env-0.1.0-py3-none-any.whl", hash = "sha256:3ecc5f31feb13fb5bd8d9a091d1d7cf0e6a8687c758f9b2a04e5f1b3b0c235a8" }, ] [[package]] -name = "lovely-numpy" -version = "0.2.13" +name = "loguru" +version = "0.7.3" source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "fastcore" }, - { name = "ipython" }, - { name = "matplotlib" }, - { name = "numpy" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/f7/46/ba51024fbf97d8370afa8a764fe18bf5c663e64badf129dfb2949a7ccbdd/lovely_numpy-0.2.13.tar.gz", hash = "sha256:0ed56660986731db3d3d7ff130e85e8d162e96c96042a0bd9a3992d32a3b34e2", size = 23920, upload-time = "2024-06-28T06:37:27.63Z" } +sdist = { url = "https://files.pythonhosted.org/packages/3a/05/a1dae3dffd1116099471c643b8924f5aa6524411dc6c63fdae648c4f1aca/loguru-0.7.3.tar.gz", hash = "sha256:19480589e77d47b8d85b2c827ad95d49bf31b0dcde16593892eb51dd18706eb6", size = 63559, upload-time = "2024-12-06T11:20:56.608Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/92/f7/a74742c9f9823a6a1fb18eb11fe77da61ea9426770ede6406f2f97a9622a/lovely_numpy-0.2.13-py3-none-any.whl", hash = "sha256:2e696d145d301264390f790cd602560ca5e525534d9059e0fcd3f313b63b5786", size = 24357, upload-time = "2024-06-28T06:37:25.206Z" }, + { url = "https://files.pythonhosted.org/packages/0c/29/0348de65b8cc732daa3e33e67806420b2ae89bdce2b04af740289c5c6c8c/loguru-0.7.3-py3-none-any.whl", hash = "sha256:31a33c10c8e1e10422bfd431aeb5d351c7cf7fa671e3c4df004162264b28220c", size = 61595, upload-time = "2024-12-06T11:20:54.538Z" }, ] [[package]] -name = "lovely-tensors" -version = "0.1.18" +name = "mamba-ssm" +version = "2.3.1" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "lovely-numpy" }, - { name = "torch" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/20/c4/b109b6c912929a471cf03c6a2f02d8c5ab2889ec57177692d7b971c4ec70/lovely_tensors-0.1.18.tar.gz", hash = "sha256:afb07d52de9ec6e560e77d9b3e01758bbc921a6f25c661b1c51ebbb3f75ee0c5", size = 21962, upload-time = "2024-11-20T22:00:29.209Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/71/fb/688604e72694c597cf236e5ccd9a4bcc34bacc5f8e5dbb2cf7d7ddccb25b/lovely_tensors-0.1.18-py3-none-any.whl", hash = "sha256:91dc30f0d6224364851e6f14497e1677076cd3a59bae4d92c78bbe8684f6f22b", size = 19303, upload-time = "2024-11-20T22:00:27.147Z" }, -] - -[[package]] -name = "lxml" -version = "6.0.1" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/8f/bd/f9d01fd4132d81c6f43ab01983caea69ec9614b913c290a26738431a015d/lxml-6.0.1.tar.gz", hash = "sha256:2b3a882ebf27dd026df3801a87cf49ff791336e0f94b0fad195db77e01240690", size = 4070214, upload-time = "2025-08-22T10:37:53.525Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/b0/a9/82b244c8198fcdf709532e39a1751943a36b3e800b420adc739d751e0299/lxml-6.0.1-cp312-cp312-macosx_10_13_universal2.whl", hash = "sha256:c03ac546adaabbe0b8e4a15d9ad815a281afc8d36249c246aecf1aaad7d6f200", size = 8422788, upload-time = "2025-08-22T10:32:56.612Z" }, - { url = "https://files.pythonhosted.org/packages/c9/8d/1ed2bc20281b0e7ed3e6c12b0a16e64ae2065d99be075be119ba88486e6d/lxml-6.0.1-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:33b862c7e3bbeb4ba2c96f3a039f925c640eeba9087a4dc7a572ec0f19d89392", size = 4593547, upload-time = "2025-08-22T10:32:59.016Z" }, - { url = "https://files.pythonhosted.org/packages/76/53/d7fd3af95b72a3493bf7fbe842a01e339d8f41567805cecfecd5c71aa5ee/lxml-6.0.1-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:7a3ec1373f7d3f519de595032d4dcafae396c29407cfd5073f42d267ba32440d", size = 4948101, upload-time = "2025-08-22T10:33:00.765Z" }, - { url = "https://files.pythonhosted.org/packages/9d/51/4e57cba4d55273c400fb63aefa2f0d08d15eac021432571a7eeefee67bed/lxml-6.0.1-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:03b12214fb1608f4cffa181ec3d046c72f7e77c345d06222144744c122ded870", size = 5108090, upload-time = "2025-08-22T10:33:03.108Z" }, - { url = "https://files.pythonhosted.org/packages/f6/6e/5f290bc26fcc642bc32942e903e833472271614e24d64ad28aaec09d5dae/lxml-6.0.1-cp312-cp312-manylinux_2_26_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:207ae0d5f0f03b30f95e649a6fa22aa73f5825667fee9c7ec6854d30e19f2ed8", size = 5021791, upload-time = "2025-08-22T10:33:06.972Z" }, - { url = "https://files.pythonhosted.org/packages/13/d4/2e7551a86992ece4f9a0f6eebd4fb7e312d30f1e372760e2109e721d4ce6/lxml-6.0.1-cp312-cp312-manylinux_2_26_i686.manylinux_2_28_i686.whl", hash = "sha256:32297b09ed4b17f7b3f448de87a92fb31bb8747496623483788e9f27c98c0f00", size = 5358861, upload-time = "2025-08-22T10:33:08.967Z" }, - { url = "https://files.pythonhosted.org/packages/8a/5f/cb49d727fc388bf5fd37247209bab0da11697ddc5e976ccac4826599939e/lxml-6.0.1-cp312-cp312-manylinux_2_26_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:7e18224ea241b657a157c85e9cac82c2b113ec90876e01e1f127312006233756", size = 5652569, upload-time = "2025-08-22T10:33:10.815Z" }, - { url = "https://files.pythonhosted.org/packages/ca/b8/66c1ef8c87ad0f958b0a23998851e610607c74849e75e83955d5641272e6/lxml-6.0.1-cp312-cp312-manylinux_2_26_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:a07a994d3c46cd4020c1ea566345cf6815af205b1e948213a4f0f1d392182072", size = 5252262, upload-time = "2025-08-22T10:33:12.673Z" }, - { url = "https://files.pythonhosted.org/packages/1a/ef/131d3d6b9590e64fdbb932fbc576b81fcc686289da19c7cb796257310e82/lxml-6.0.1-cp312-cp312-manylinux_2_31_armv7l.whl", hash = "sha256:2287fadaa12418a813b05095485c286c47ea58155930cfbd98c590d25770e225", size = 4710309, upload-time = "2025-08-22T10:33:14.952Z" }, - { url = "https://files.pythonhosted.org/packages/bc/3f/07f48ae422dce44902309aa7ed386c35310929dc592439c403ec16ef9137/lxml-6.0.1-cp312-cp312-manylinux_2_38_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:b4e597efca032ed99f418bd21314745522ab9fa95af33370dcee5533f7f70136", size = 5265786, upload-time = "2025-08-22T10:33:16.721Z" }, - { url = "https://files.pythonhosted.org/packages/11/c7/125315d7b14ab20d9155e8316f7d287a4956098f787c22d47560b74886c4/lxml-6.0.1-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:9696d491f156226decdd95d9651c6786d43701e49f32bf23715c975539aa2b3b", size = 5062272, upload-time = "2025-08-22T10:33:18.478Z" }, - { url = "https://files.pythonhosted.org/packages/8b/c3/51143c3a5fc5168a7c3ee626418468ff20d30f5a59597e7b156c1e61fba8/lxml-6.0.1-cp312-cp312-musllinux_1_2_armv7l.whl", hash = "sha256:e4e3cd3585f3c6f87cdea44cda68e692cc42a012f0131d25957ba4ce755241a7", size = 4786955, upload-time = "2025-08-22T10:33:20.34Z" }, - { url = "https://files.pythonhosted.org/packages/11/86/73102370a420ec4529647b31c4a8ce8c740c77af3a5fae7a7643212d6f6e/lxml-6.0.1-cp312-cp312-musllinux_1_2_ppc64le.whl", hash = "sha256:45cbc92f9d22c28cd3b97f8d07fcefa42e569fbd587dfdac76852b16a4924277", size = 5673557, upload-time = "2025-08-22T10:33:22.282Z" }, - { url = "https://files.pythonhosted.org/packages/d7/2d/aad90afaec51029aef26ef773b8fd74a9e8706e5e2f46a57acd11a421c02/lxml-6.0.1-cp312-cp312-musllinux_1_2_riscv64.whl", hash = "sha256:f8c9bcfd2e12299a442fba94459adf0b0d001dbc68f1594439bfa10ad1ecb74b", size = 5254211, upload-time = "2025-08-22T10:33:24.15Z" }, - { url = "https://files.pythonhosted.org/packages/63/01/c9e42c8c2d8b41f4bdefa42ab05448852e439045f112903dd901b8fbea4d/lxml-6.0.1-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:1e9dc2b9f1586e7cd77753eae81f8d76220eed9b768f337dc83a3f675f2f0cf9", size = 5275817, upload-time = "2025-08-22T10:33:26.007Z" }, - { url = "https://files.pythonhosted.org/packages/bc/1f/962ea2696759abe331c3b0e838bb17e92224f39c638c2068bf0d8345e913/lxml-6.0.1-cp312-cp312-win32.whl", hash = "sha256:987ad5c3941c64031f59c226167f55a04d1272e76b241bfafc968bdb778e07fb", size = 3610889, upload-time = "2025-08-22T10:33:28.169Z" }, - { url = "https://files.pythonhosted.org/packages/41/e2/22c86a990b51b44442b75c43ecb2f77b8daba8c4ba63696921966eac7022/lxml-6.0.1-cp312-cp312-win_amd64.whl", hash = "sha256:abb05a45394fd76bf4a60c1b7bec0e6d4e8dfc569fc0e0b1f634cd983a006ddc", size = 4010925, upload-time = "2025-08-22T10:33:29.874Z" }, - { url = "https://files.pythonhosted.org/packages/b2/21/dc0c73325e5eb94ef9c9d60dbb5dcdcb2e7114901ea9509735614a74e75a/lxml-6.0.1-cp312-cp312-win_arm64.whl", hash = "sha256:c4be29bce35020d8579d60aa0a4e95effd66fcfce31c46ffddf7e5422f73a299", size = 3671922, upload-time = "2025-08-22T10:33:31.535Z" }, + { name = "einops", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "ninja", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "packaging", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "setuptools", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "torch", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "transformers", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "triton", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, ] +sdist = { url = "https://files.pythonhosted.org/packages/34/67/ec89aa703da194a813e35d2ea2de8f74a7ce6991a120a29f3a0c5e30d4b9/mamba_ssm-2.3.1.tar.gz", hash = "sha256:4d529477ad94753962216d583fc8f1c127c717b7d7c875d6bbb9376366d0d761", size = 121707, upload-time = "2026-03-10T09:27:34.798Z" } [[package]] name = "markdown" @@ -1549,7 +1741,7 @@ name = "markdown-it-py" version = "4.0.0" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "mdurl" }, + { name = "mdurl", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, ] sdist = { url = "https://files.pythonhosted.org/packages/5b/f5/4ec618ed16cc4f8fb3b701563655a69816155e79e24a17b651541804721d/markdown_it_py-4.0.0.tar.gz", hash = "sha256:cb0a2b4aa34f932c007117b194e945bd74e0ec24133ceb5bac59009cda1cb9f3", size = 73070, upload-time = "2025-08-11T12:57:52.854Z" } wheels = [ @@ -1558,10 +1750,10 @@ wheels = [ [package.optional-dependencies] linkify = [ - { name = "linkify-it-py" }, + { name = "linkify-it-py", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, ] plugins = [ - { name = "mdit-py-plugins" }, + { name = "mdit-py-plugins", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, ] [[package]] @@ -1570,16 +1762,36 @@ version = "3.0.2" source = { registry = "https://pypi.org/simple" } sdist = { url = "https://files.pythonhosted.org/packages/b2/97/5d42485e71dfc078108a86d6de8fa46db44a1a9295e89c5d6d4a06e23a62/markupsafe-3.0.2.tar.gz", hash = "sha256:ee55d3edf80167e48ea11a923c7386f4669df67d7994554387f84e7d8b0a2bf0", size = 20537, upload-time = "2024-10-18T15:21:54.129Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/22/09/d1f21434c97fc42f09d290cbb6350d44eb12f09cc62c9476effdb33a18aa/MarkupSafe-3.0.2-cp312-cp312-macosx_10_13_universal2.whl", hash = "sha256:9778bd8ab0a994ebf6f84c2b949e65736d5575320a17ae8984a77fab08db94cf", size = 14274, upload-time = "2024-10-18T15:21:13.777Z" }, - { url = "https://files.pythonhosted.org/packages/6b/b0/18f76bba336fa5aecf79d45dcd6c806c280ec44538b3c13671d49099fdd0/MarkupSafe-3.0.2-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:846ade7b71e3536c4e56b386c2a47adf5741d2d8b94ec9dc3e92e5e1ee1e2225", size = 12348, upload-time = "2024-10-18T15:21:14.822Z" }, { url = "https://files.pythonhosted.org/packages/e0/25/dd5c0f6ac1311e9b40f4af06c78efde0f3b5cbf02502f8ef9501294c425b/MarkupSafe-3.0.2-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:1c99d261bd2d5f6b59325c92c73df481e05e57f19837bdca8413b9eac4bd8028", size = 24149, upload-time = "2024-10-18T15:21:15.642Z" }, { url = "https://files.pythonhosted.org/packages/f3/f0/89e7aadfb3749d0f52234a0c8c7867877876e0a20b60e2188e9850794c17/MarkupSafe-3.0.2-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:e17c96c14e19278594aa4841ec148115f9c7615a47382ecb6b82bd8fea3ab0c8", size = 23118, upload-time = "2024-10-18T15:21:17.133Z" }, - { url = "https://files.pythonhosted.org/packages/d5/da/f2eeb64c723f5e3777bc081da884b414671982008c47dcc1873d81f625b6/MarkupSafe-3.0.2-cp312-cp312-manylinux_2_5_i686.manylinux1_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:88416bd1e65dcea10bc7569faacb2c20ce071dd1f87539ca2ab364bf6231393c", size = 22993, upload-time = "2024-10-18T15:21:18.064Z" }, { url = "https://files.pythonhosted.org/packages/da/0e/1f32af846df486dce7c227fe0f2398dc7e2e51d4a370508281f3c1c5cddc/MarkupSafe-3.0.2-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:2181e67807fc2fa785d0592dc2d6206c019b9502410671cc905d132a92866557", size = 24178, upload-time = "2024-10-18T15:21:18.859Z" }, - { url = "https://files.pythonhosted.org/packages/c4/f6/bb3ca0532de8086cbff5f06d137064c8410d10779c4c127e0e47d17c0b71/MarkupSafe-3.0.2-cp312-cp312-musllinux_1_2_i686.whl", hash = "sha256:52305740fe773d09cffb16f8ed0427942901f00adedac82ec8b67752f58a1b22", size = 23319, upload-time = "2024-10-18T15:21:19.671Z" }, { url = "https://files.pythonhosted.org/packages/a2/82/8be4c96ffee03c5b4a034e60a31294daf481e12c7c43ab8e34a1453ee48b/MarkupSafe-3.0.2-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:ad10d3ded218f1039f11a75f8091880239651b52e9bb592ca27de44eed242a48", size = 23352, upload-time = "2024-10-18T15:21:20.971Z" }, - { url = "https://files.pythonhosted.org/packages/51/ae/97827349d3fcffee7e184bdf7f41cd6b88d9919c80f0263ba7acd1bbcb18/MarkupSafe-3.0.2-cp312-cp312-win32.whl", hash = "sha256:0f4ca02bea9a23221c0182836703cbf8930c5e9454bacce27e767509fa286a30", size = 15097, upload-time = "2024-10-18T15:21:22.646Z" }, - { url = "https://files.pythonhosted.org/packages/c1/80/a61f99dc3a936413c3ee4e1eecac96c0da5ed07ad56fd975f1a9da5bc630/MarkupSafe-3.0.2-cp312-cp312-win_amd64.whl", hash = "sha256:8e06879fc22a25ca47312fbe7c8264eb0b662f6db27cb2d3bbbc74b1df4b9b87", size = 15601, upload-time = "2024-10-18T15:21:23.499Z" }, +] + +[[package]] +name = "math-env" +version = "0.1.2" +source = { registry = "https://hub.primeintellect.ai/primeintellect/simple/" } +dependencies = [ + { name = "math-verify", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "prime-sandboxes", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "verifiers", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, +] +wheels = [ + { url = "https://hub.primeintellect.ai/primeintellect/math-env/@07fdeaab/math_env-0.1.2-py3-none-any.whl", hash = "sha256:87003ca56c65818e7a7707f12e80511180038190abe4df5ec66cb2012d38b6d9" }, +] + +[[package]] +name = "math-python" +version = "0.1.10" +source = { registry = "https://hub.primeintellect.ai/primeintellect/simple/" } +dependencies = [ + { name = "math-verify", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "prime-sandboxes", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "verifiers", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, +] +wheels = [ + { url = "https://hub.primeintellect.ai/primeintellect/math-python/@113ba8c2/math_python-0.1.10-py3-none-any.whl", hash = "sha256:772d176938cedafe3b11af068452b3e413005f41fd57e27abe622e7c4ebba316" }, ] [[package]] @@ -1587,7 +1799,7 @@ name = "math-verify" version = "0.8.0" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "latex2sympy2-extended" }, + { name = "latex2sympy2-extended", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, ] sdist = { url = "https://files.pythonhosted.org/packages/35/b5/b1db6fa6b6c28ebbe1889ee11a4703a72a2ca7750ec415f4559c758cf01a/math_verify-0.8.0.tar.gz", hash = "sha256:3295e0adb94bfe553ff6e3189c44f1916a85aa24ab5d1900f2086a706e28f7c4", size = 60191, upload-time = "2025-07-02T15:52:07.209Z" } wheels = [ @@ -1595,29 +1807,15 @@ wheels = [ ] [[package]] -name = "matplotlib" -version = "3.10.6" -source = { registry = "https://pypi.org/simple" } +name = "math500" +version = "0.1.16" +source = { registry = "https://hub.primeintellect.ai/primeintellect/simple/" } dependencies = [ - { name = "contourpy" }, - { name = "cycler" }, - { name = "fonttools" }, - { name = "kiwisolver" }, - { name = "numpy" }, - { name = "packaging" }, - { name = "pillow" }, - { name = "pyparsing" }, - { name = "python-dateutil" }, + { name = "math-verify", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "verifiers", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/a0/59/c3e6453a9676ffba145309a73c462bb407f4400de7de3f2b41af70720a3c/matplotlib-3.10.6.tar.gz", hash = "sha256:ec01b645840dd1996df21ee37f208cd8ba57644779fa20464010638013d3203c", size = 34804264, upload-time = "2025-08-30T00:14:25.137Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/ea/1a/7042f7430055d567cc3257ac409fcf608599ab27459457f13772c2d9778b/matplotlib-3.10.6-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:31ca662df6a80bd426f871105fdd69db7543e28e73a9f2afe80de7e531eb2347", size = 8272404, upload-time = "2025-08-30T00:12:59.112Z" }, - { url = "https://files.pythonhosted.org/packages/a9/5d/1d5f33f5b43f4f9e69e6a5fe1fb9090936ae7bc8e2ff6158e7a76542633b/matplotlib-3.10.6-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:1678bb61d897bb4ac4757b5ecfb02bfb3fddf7f808000fb81e09c510712fda75", size = 8128262, upload-time = "2025-08-30T00:13:01.141Z" }, - { url = "https://files.pythonhosted.org/packages/67/c3/135fdbbbf84e0979712df58e5e22b4f257b3f5e52a3c4aacf1b8abec0d09/matplotlib-3.10.6-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:56cd2d20842f58c03d2d6e6c1f1cf5548ad6f66b91e1e48f814e4fb5abd1cb95", size = 8697008, upload-time = "2025-08-30T00:13:03.24Z" }, - { url = "https://files.pythonhosted.org/packages/9c/be/c443ea428fb2488a3ea7608714b1bd85a82738c45da21b447dc49e2f8e5d/matplotlib-3.10.6-cp312-cp312-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:662df55604a2f9a45435566d6e2660e41efe83cd94f4288dfbf1e6d1eae4b0bb", size = 9530166, upload-time = "2025-08-30T00:13:05.951Z" }, - { url = "https://files.pythonhosted.org/packages/a9/35/48441422b044d74034aea2a3e0d1a49023f12150ebc58f16600132b9bbaf/matplotlib-3.10.6-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:08f141d55148cd1fc870c3387d70ca4df16dee10e909b3b038782bd4bda6ea07", size = 9593105, upload-time = "2025-08-30T00:13:08.356Z" }, - { url = "https://files.pythonhosted.org/packages/45/c3/994ef20eb4154ab84cc08d033834555319e4af970165e6c8894050af0b3c/matplotlib-3.10.6-cp312-cp312-win_amd64.whl", hash = "sha256:590f5925c2d650b5c9d813c5b3b5fc53f2929c3f8ef463e4ecfa7e052044fb2b", size = 8122784, upload-time = "2025-08-30T00:13:10.367Z" }, - { url = "https://files.pythonhosted.org/packages/57/b8/5c85d9ae0e40f04e71bedb053aada5d6bab1f9b5399a0937afb5d6b02d98/matplotlib-3.10.6-cp312-cp312-win_arm64.whl", hash = "sha256:f44c8d264a71609c79a78d50349e724f5d5fc3684ead7c2a473665ee63d868aa", size = 7992823, upload-time = "2025-08-30T00:13:12.24Z" }, + { url = "https://hub.primeintellect.ai/primeintellect/math500/@4951d227/math500-0.1.16-py3-none-any.whl", hash = "sha256:08c6b96755dd64d90efa5bc4e488e8232abddacedbc238267ecb09632ed3f5cb" }, ] [[package]] @@ -1625,7 +1823,7 @@ name = "matplotlib-inline" version = "0.1.7" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "traitlets" }, + { name = "traitlets", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, ] sdist = { url = "https://files.pythonhosted.org/packages/99/5b/a36a337438a14116b16480db471ad061c36c3694df7c2084a0da7ba538b7/matplotlib_inline-0.1.7.tar.gz", hash = "sha256:8423b23ec666be3d16e16b60bdd8ac4e86e840ebd1dd11a30b9f117f2fa0ab90", size = 8159, upload-time = "2024-04-15T13:44:44.803Z" } wheels = [ @@ -1637,20 +1835,19 @@ name = "mcp" version = "1.25.0" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "anyio" }, - { name = "httpx" }, - { name = "httpx-sse" }, - { name = "jsonschema" }, - { name = "pydantic" }, - { name = "pydantic-settings" }, - { name = "pyjwt", extra = ["crypto"] }, - { name = "python-multipart" }, - { name = "pywin32", marker = "sys_platform == 'win32'" }, - { name = "sse-starlette" }, - { name = "starlette" }, - { name = "typing-extensions" }, - { name = "typing-inspection" }, - { name = "uvicorn", marker = "sys_platform != 'emscripten'" }, + { name = "anyio", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "httpx", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "httpx-sse", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "jsonschema", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "pydantic", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "pydantic-settings", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "pyjwt", extra = ["crypto"], marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "python-multipart", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "sse-starlette", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "starlette", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "typing-extensions", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "typing-inspection", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "uvicorn", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, ] sdist = { url = "https://files.pythonhosted.org/packages/d5/2d/649d80a0ecf6a1f82632ca44bec21c0461a9d9fc8934d38cb5b319f2db5e/mcp-1.25.0.tar.gz", hash = "sha256:56310361ebf0364e2d438e5b45f7668cbb124e158bb358333cd06e49e83a6802", size = 605387, upload-time = "2025-12-19T10:19:56.985Z" } wheels = [ @@ -1662,7 +1859,7 @@ name = "mdit-py-plugins" version = "0.5.0" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "markdown-it-py" }, + { name = "markdown-it-py", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, ] sdist = { url = "https://files.pythonhosted.org/packages/b2/fd/a756d36c0bfba5f6e39a1cdbdbfdd448dc02692467d83816dff4592a1ebc/mdit_py_plugins-0.5.0.tar.gz", hash = "sha256:f4918cb50119f50446560513a8e311d574ff6aaed72606ddae6d35716fe809c6", size = 44655, upload-time = "2025-08-11T07:25:49.083Z" } wheels = [ @@ -1678,86 +1875,110 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/b3/38/89ba8ad64ae25be8de66a6d463314cf1eb366222074cfda9ee839c56a4b4/mdurl-0.1.2-py3-none-any.whl", hash = "sha256:84008a41e51615a49fc9966191ff91509e3c40b939176e643fd50a5c2196b8f8", size = 9979, upload-time = "2022-08-14T12:40:09.779Z" }, ] +[[package]] +name = "mini-swe-agent-plus" +version = "0.2.17" +source = { registry = "https://hub.primeintellect.ai/primeintellect/simple/" } +dependencies = [ + { name = "prime-sandboxes", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "swebench", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "tenacity", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "verifiers", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, +] +wheels = [ + { url = "https://hub.primeintellect.ai/primeintellect/mini-swe-agent-plus/@aa6f9500/mini_swe_agent_plus-0.2.17-py3-none-any.whl", hash = "sha256:af4f6e10412b86dd4be0623393fc35275fe95080c95ac032bf39e109a2dd390b" }, +] + [[package]] name = "mistral-common" -version = "1.8.6" +version = "1.11.0" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "jsonschema" }, - { name = "numpy" }, - { name = "pillow" }, - { name = "pydantic" }, - { name = "pydantic-extra-types", extra = ["pycountry"] }, - { name = "requests" }, - { name = "tiktoken" }, - { name = "typing-extensions" }, + { name = "jsonschema", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "numpy", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "pillow", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "pydantic", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "pydantic-extra-types", extra = ["pycountry"], marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "requests", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "tiktoken", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "typing-extensions", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/23/3d/2ce630343f4574ac5150e355dc15bf95e8ce18954aa0f5949e195ffde382/mistral_common-1.8.6.tar.gz", hash = "sha256:c61702720093f7a06508e81923917b04e35062b9ff396b8512b9c4d1139767ee", size = 6332366, upload-time = "2025-11-30T22:03:37.698Z" } +sdist = { url = "https://files.pythonhosted.org/packages/61/97/753c85b5c0a19f4331ac99e0300ac8da06d4b29b629c9cb03064b38561bd/mistral_common-1.11.0.tar.gz", hash = "sha256:439b7fa38f9c3f020154af51bdf30eb81def507643017d8ce9f798384ec47ec3", size = 6355512, upload-time = "2026-04-01T13:54:12.36Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/e6/eb/56ba11528da93b890fe985cbbaa17955064ff4b7019dd18566e3c3559b9e/mistral_common-1.8.6-py3-none-any.whl", hash = "sha256:dd8c0e55b397e8167751eb3da147cf23fd970824673ca0e260aa58c888be1b0a", size = 6515009, upload-time = "2025-11-30T22:03:35.108Z" }, + { url = "https://files.pythonhosted.org/packages/60/e4/73ad3c27e3fb613c3ce0953c928202c46cddebac3989b87be1b6f305a9f6/mistral_common-1.11.0-py3-none-any.whl", hash = "sha256:1d3ecaf7c3aa7338cb37b596fd0fb294485753958ee8e7254a6cc23eb30b249b", size = 6531513, upload-time = "2026-04-01T13:54:16.536Z" }, ] [package.optional-dependencies] image = [ - { name = "opencv-python-headless" }, + { name = "opencv-python-headless", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, ] [[package]] -name = "mlx" -version = "0.29.0" +name = "ml-dtypes" +version = "0.5.4" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "mlx-metal", marker = "sys_platform == 'darwin'" }, + { name = "numpy", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, ] +sdist = { url = "https://files.pythonhosted.org/packages/0e/4a/c27b42ed9b1c7d13d9ba8b6905dece787d6259152f2309338aed29b2447b/ml_dtypes-0.5.4.tar.gz", hash = "sha256:8ab06a50fb9bf9666dd0fe5dfb4676fa2b0ac0f31ecff72a6c3af8e22c063453", size = 692314, upload-time = "2025-11-17T22:32:31.031Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/60/c0/a95f24f4b78fdf96f57aaa1a8919f5688603fd4fa0294c2f868f60bc98af/mlx-0.29.0-cp312-cp312-macosx_13_0_arm64.whl", hash = "sha256:a9f9f760a179d96fa7fecc0b31a9885a9ab4ce9e2914f7f34c2980edd7f012fa", size = 546443, upload-time = "2025-08-29T17:14:17.997Z" }, - { url = "https://files.pythonhosted.org/packages/61/c7/9f6a5cf1a0b0eb1ace3892c930f64fa84351cecc75221ce4ca5abd0ad64d/mlx-0.29.0-cp312-cp312-macosx_14_0_arm64.whl", hash = "sha256:446ea784a4717c9b06eee9c91ee038245ebb231760a2cf61b677dc3d256cbd83", size = 546442, upload-time = "2025-08-29T17:14:19.035Z" }, - { url = "https://files.pythonhosted.org/packages/85/40/32b1fe76d80f559ff3af18d104a22c4d07ecad62a88e3c313083722868e4/mlx-0.29.0-cp312-cp312-macosx_15_0_arm64.whl", hash = "sha256:593a3d71c5859f83bc55a9bdc612fe5f8b495c14112e0e7a96c027c11ce5ab52", size = 546444, upload-time = "2025-08-29T17:14:58.946Z" }, + { url = "https://files.pythonhosted.org/packages/54/0f/428ef6881782e5ebb7eca459689448c0394fa0a80bea3aa9262cba5445ea/ml_dtypes-0.5.4-cp312-cp312-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:a7f7c643e8b1320fd958bf098aa7ecf70623a42ec5154e3be3be673f4c34d900", size = 5028464, upload-time = "2025-11-17T22:31:50.135Z" }, + { url = "https://files.pythonhosted.org/packages/3a/cb/28ce52eb94390dda42599c98ea0204d74799e4d8047a0eb559b6fd648056/ml_dtypes-0.5.4-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:9ad459e99793fa6e13bd5b7e6792c8f9190b4e5a1b45c63aba14a4d0a7f1d5ff", size = 5009002, upload-time = "2025-11-17T22:31:52.001Z" }, ] [[package]] -name = "mlx-lm" -version = "0.27.1" +name = "mmh3" +version = "5.2.1" source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "jinja2", marker = "sys_platform != 'linux'" }, - { name = "mlx", marker = "sys_platform != 'linux'" }, - { name = "numpy", marker = "sys_platform != 'linux'" }, - { name = "protobuf", marker = "sys_platform != 'linux'" }, - { name = "pyyaml", marker = "sys_platform != 'linux'" }, - { name = "transformers", marker = "sys_platform != 'linux'" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/41/77/e8d3a82658a2070bc392a583dd08c8d24088433e920eac4905bf882255ad/mlx_lm-0.27.1.tar.gz", hash = "sha256:36640fb64c909cfd9baddf37b16e7d3b94a1a141033e6b7ea7a0ef5a965fb4ae", size = 185170, upload-time = "2025-09-04T16:06:57.949Z" } +sdist = { url = "https://files.pythonhosted.org/packages/91/1a/edb23803a168f070ded7a3014c6d706f63b90c84ccc024f89d794a3b7a6d/mmh3-5.2.1.tar.gz", hash = "sha256:bbea5b775f0ac84945191fb83f845a6fd9a21a03ea7f2e187defac7e401616ad", size = 33775, upload-time = "2026-03-05T15:55:57.716Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/e1/54/5f35831d208cbf81572e9a0ae8ac6d595ca7c59f3e1da57c367894b0a75b/mlx_lm-0.27.1-py3-none-any.whl", hash = "sha256:300da6f63d8d392483b62b2abda794730fa04343dcb28a1f6a712f4c3ab60f3c", size = 255687, upload-time = "2025-09-04T16:06:54.904Z" }, + { url = "https://files.pythonhosted.org/packages/e5/e2/51ed62063b44d10b06d975ac87af287729eeb5e3ed9772f7584a17983e90/mmh3-5.2.1-cp312-cp312-manylinux1_x86_64.manylinux_2_28_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:8e6c219e375f6341d0959af814296372d265a8ca1af63825f65e2e87c618f006", size = 103274, upload-time = "2026-03-05T15:54:26.44Z" }, + { url = "https://files.pythonhosted.org/packages/75/ce/12a7524dca59eec92e5b31fdb13ede1e98eda277cf2b786cf73bfbc24e81/mmh3-5.2.1-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:26fb5b9c3946bf7f1daed7b37e0c03898a6f062149127570f8ede346390a0825", size = 106158, upload-time = "2026-03-05T15:54:28.578Z" }, + { url = "https://files.pythonhosted.org/packages/76/b3/70b73923fd0284c439860ff5c871b20210dfdbe9a6b9dd0ee6496d77f174/mmh3-5.2.1-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:b3f99e1756fc48ad507b95e5d86f2fb21b3d495012ff13e6592ebac14033f166", size = 99111, upload-time = "2026-03-05T15:54:32.353Z" }, + { url = "https://files.pythonhosted.org/packages/3b/ac/ca8e0c19a34f5b71390171d2ff0b9f7f187550d66801a731bb68925126a4/mmh3-5.2.1-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:3d74a03fb57757ece25aa4b3c1c60157a1cece37a020542785f942e2f827eed5", size = 97507, upload-time = "2026-03-05T15:54:37.804Z" }, ] [[package]] -name = "mlx-metal" -version = "0.29.0" +name = "modal" +version = "1.4.0" source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "aiohttp", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "cbor2", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "certifi", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "click", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "grpclib", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "protobuf", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "rich", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "synchronicity", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "toml", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "typer", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "types-certifi", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "types-toml", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "typing-extensions", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "watchfiles", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/b8/31/70526ac8d7546253edec8d3f7f99ffb2ae00519a0c68d394eaa48b96c4a6/modal-1.4.0.tar.gz", hash = "sha256:340743f7284fded151c84d54bb6fc545f68f93e064656364c198dadc9d0421df", size = 682472, upload-time = "2026-03-25T21:41:35.946Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/08/28/3fb942b27d53061cacd12883629e5c48821f9e661f01210f55d89dabe3a6/mlx_metal-0.29.0-py3-none-macosx_13_0_arm64.whl", hash = "sha256:7e8279851504a8394122497d491276db37ef85d29f2f54d8a711ac90edfef7bf", size = 34964891, upload-time = "2025-08-29T17:16:35.72Z" }, - { url = "https://files.pythonhosted.org/packages/6d/82/582555232044bbdc0d74186077bcd79f1ce45f07a79e49b36afc25f9e313/mlx_metal-0.29.0-py3-none-macosx_14_0_arm64.whl", hash = "sha256:674fc7df4b9502dffe56c92cc95bae3a44950c70aa50f5b53291f228cfbbb631", size = 34694819, upload-time = "2025-08-29T17:16:32.636Z" }, - { url = "https://files.pythonhosted.org/packages/30/d6/b3fc6d07b9bc814551078e53438ebfae8e44c47e4e5f762df244150e4e3a/mlx_metal-0.29.0-py3-none-macosx_15_0_arm64.whl", hash = "sha256:f69eb71128dcc55b5d9baa7606c31641bdd214b859d809deac51be4248d6fddb", size = 34685946, upload-time = "2025-08-29T17:15:25.028Z" }, + { url = "https://files.pythonhosted.org/packages/a4/eb/e0f8154f4f9a80552b4d2ff962994714f907c1298c60783d80697bc12e66/modal-1.4.0-py3-none-any.whl", hash = "sha256:eb41eb52d3cc472ad04cff0f7957ca41d9a53961eca46f2058b3039934d6ff49", size = 784602, upload-time = "2026-03-25T21:41:34.136Z" }, ] [[package]] name = "model-hosting-container-standards" -version = "0.1.11" +version = "0.1.13" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "fastapi" }, - { name = "httpx" }, - { name = "jmespath" }, - { name = "pydantic" }, - { name = "setuptools" }, - { name = "starlette" }, - { name = "supervisor" }, + { name = "fastapi", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "httpx", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "jmespath", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "pydantic", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "setuptools", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "starlette", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "supervisor", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/2b/cd/1d70e84cd8bde6227ae42e13d02fe39f7ea7d2216b496082372746f5c010/model_hosting_container_standards-0.1.11.tar.gz", hash = "sha256:dbb54b966a435db247f204872b22e1e98141822e1460a17459011f2cdf12d6b2", size = 78140, upload-time = "2025-12-10T20:39:08.111Z" } +sdist = { url = "https://files.pythonhosted.org/packages/d7/b7/a6a31b4dfd30d14b1019dc358f09c9d88ca38e555ba7c976e7d3e6b593fe/model_hosting_container_standards-0.1.13.tar.gz", hash = "sha256:27a1333410dde2719286a300a2803e24fdde407baa91894eb845c0f268aa194d", size = 79116, upload-time = "2026-01-09T21:45:20.683Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/9c/70/9d1cd3730309cb47f59923a0a0bd5832552c5f3725113d4a15f80eaacda0/model_hosting_container_standards-0.1.11-py3-none-any.whl", hash = "sha256:50b8d97ab3fddf2e22ffb9d489ef5a26369b96aaf850db344549da1795311dc6", size = 104688, upload-time = "2025-12-10T20:39:06.651Z" }, + { url = "https://files.pythonhosted.org/packages/8c/37/6dc61971ba31450bbed460b5f40543f0915e352680534e3bcaf57116d8d7/model_hosting_container_standards-0.1.13-py3-none-any.whl", hash = "sha256:be307d4a988cc660df4e6bd8bdedb7917844bac940e332f9fd001cb385d7994c", size = 105738, upload-time = "2026-01-09T21:45:18.959Z" }, ] [[package]] @@ -1771,20 +1992,14 @@ wheels = [ [[package]] name = "msgpack" -version = "1.1.1" +version = "1.1.2" source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/45/b1/ea4f68038a18c77c9467400d166d74c4ffa536f34761f7983a104357e614/msgpack-1.1.1.tar.gz", hash = "sha256:77b79ce34a2bdab2594f490c8e80dd62a02d650b91a75159a63ec413b8d104cd", size = 173555, upload-time = "2025-06-13T06:52:51.324Z" } +sdist = { url = "https://files.pythonhosted.org/packages/4d/f2/bfb55a6236ed8725a96b0aa3acbd0ec17588e6a2c3b62a93eb513ed8783f/msgpack-1.1.2.tar.gz", hash = "sha256:3b60763c1373dd60f398488069bcdc703cd08a711477b5d480eecc9f9626f47e", size = 173581, upload-time = "2025-10-08T09:15:56.596Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/e3/26/389b9c593eda2b8551b2e7126ad3a06af6f9b44274eb3a4f054d48ff7e47/msgpack-1.1.1-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:ae497b11f4c21558d95de9f64fff7053544f4d1a17731c866143ed6bb4591238", size = 82359, upload-time = "2025-06-13T06:52:03.909Z" }, - { url = "https://files.pythonhosted.org/packages/ab/65/7d1de38c8a22cf8b1551469159d4b6cf49be2126adc2482de50976084d78/msgpack-1.1.1-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:33be9ab121df9b6b461ff91baac6f2731f83d9b27ed948c5b9d1978ae28bf157", size = 79172, upload-time = "2025-06-13T06:52:05.246Z" }, - { url = "https://files.pythonhosted.org/packages/0f/bd/cacf208b64d9577a62c74b677e1ada005caa9b69a05a599889d6fc2ab20a/msgpack-1.1.1-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:6f64ae8fe7ffba251fecb8408540c34ee9df1c26674c50c4544d72dbf792e5ce", size = 425013, upload-time = "2025-06-13T06:52:06.341Z" }, - { url = "https://files.pythonhosted.org/packages/4d/ec/fd869e2567cc9c01278a736cfd1697941ba0d4b81a43e0aa2e8d71dab208/msgpack-1.1.1-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:a494554874691720ba5891c9b0b39474ba43ffb1aaf32a5dac874effb1619e1a", size = 426905, upload-time = "2025-06-13T06:52:07.501Z" }, - { url = "https://files.pythonhosted.org/packages/55/2a/35860f33229075bce803a5593d046d8b489d7ba2fc85701e714fc1aaf898/msgpack-1.1.1-cp312-cp312-manylinux_2_5_i686.manylinux1_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:cb643284ab0ed26f6957d969fe0dd8bb17beb567beb8998140b5e38a90974f6c", size = 407336, upload-time = "2025-06-13T06:52:09.047Z" }, - { url = "https://files.pythonhosted.org/packages/8c/16/69ed8f3ada150bf92745fb4921bd621fd2cdf5a42e25eb50bcc57a5328f0/msgpack-1.1.1-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:d275a9e3c81b1093c060c3837e580c37f47c51eca031f7b5fb76f7b8470f5f9b", size = 409485, upload-time = "2025-06-13T06:52:10.382Z" }, - { url = "https://files.pythonhosted.org/packages/c6/b6/0c398039e4c6d0b2e37c61d7e0e9d13439f91f780686deb8ee64ecf1ae71/msgpack-1.1.1-cp312-cp312-musllinux_1_2_i686.whl", hash = "sha256:4fd6b577e4541676e0cc9ddc1709d25014d3ad9a66caa19962c4f5de30fc09ef", size = 412182, upload-time = "2025-06-13T06:52:11.644Z" }, - { url = "https://files.pythonhosted.org/packages/b8/d0/0cf4a6ecb9bc960d624c93effaeaae75cbf00b3bc4a54f35c8507273cda1/msgpack-1.1.1-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:bb29aaa613c0a1c40d1af111abf025f1732cab333f96f285d6a93b934738a68a", size = 419883, upload-time = "2025-06-13T06:52:12.806Z" }, - { url = "https://files.pythonhosted.org/packages/62/83/9697c211720fa71a2dfb632cad6196a8af3abea56eece220fde4674dc44b/msgpack-1.1.1-cp312-cp312-win32.whl", hash = "sha256:870b9a626280c86cff9c576ec0d9cbcc54a1e5ebda9cd26dab12baf41fee218c", size = 65406, upload-time = "2025-06-13T06:52:14.271Z" }, - { url = "https://files.pythonhosted.org/packages/c0/23/0abb886e80eab08f5e8c485d6f13924028602829f63b8f5fa25a06636628/msgpack-1.1.1-cp312-cp312-win_amd64.whl", hash = "sha256:5692095123007180dca3e788bb4c399cc26626da51629a31d40207cb262e67f4", size = 72558, upload-time = "2025-06-13T06:52:15.252Z" }, + { url = "https://files.pythonhosted.org/packages/f2/60/a064b0345fc36c4c3d2c743c82d9100c40388d77f0b48b2f04d6041dbec1/msgpack-1.1.2-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:c63eea553c69ab05b6747901b97d620bb2a690633c77f23feb0c6a947a8a7b8f", size = 417131, upload-time = "2025-10-08T09:15:05.136Z" }, + { url = "https://files.pythonhosted.org/packages/65/92/a5100f7185a800a5d29f8d14041f61475b9de465ffcc0f3b9fba606e4505/msgpack-1.1.2-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:372839311ccf6bdaf39b00b61288e0557916c3729529b301c52c2d88842add42", size = 427556, upload-time = "2025-10-08T09:15:06.837Z" }, + { url = "https://files.pythonhosted.org/packages/f5/87/ffe21d1bf7d9991354ad93949286f643b2bb6ddbeab66373922b44c3b8cc/msgpack-1.1.2-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:2929af52106ca73fcb28576218476ffbb531a036c2adbcf54a3664de124303e9", size = 404920, upload-time = "2025-10-08T09:15:08.179Z" }, + { url = "https://files.pythonhosted.org/packages/ff/41/8543ed2b8604f7c0d89ce066f42007faac1eaa7d79a81555f206a5cdb889/msgpack-1.1.2-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:be52a8fc79e45b0364210eef5234a7cf8d330836d0a64dfbb878efa903d84620", size = 415013, upload-time = "2025-10-08T09:15:09.83Z" }, ] [[package]] @@ -1793,13 +2008,10 @@ version = "0.19.0" source = { registry = "https://pypi.org/simple" } sdist = { url = "https://files.pythonhosted.org/packages/cf/9b/95d8ce458462b8b71b8a70fa94563b2498b89933689f3a7b8911edfae3d7/msgspec-0.19.0.tar.gz", hash = "sha256:604037e7cd475345848116e89c553aa9a233259733ab51986ac924ab1b976f8e", size = 216934, upload-time = "2024-12-27T17:40:28.597Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/b2/5f/a70c24f075e3e7af2fae5414c7048b0e11389685b7f717bb55ba282a34a7/msgspec-0.19.0-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:f98bd8962ad549c27d63845b50af3f53ec468b6318400c9f1adfe8b092d7b62f", size = 190485, upload-time = "2024-12-27T17:39:44.974Z" }, - { url = "https://files.pythonhosted.org/packages/89/b0/1b9763938cfae12acf14b682fcf05c92855974d921a5a985ecc197d1c672/msgspec-0.19.0-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:43bbb237feab761b815ed9df43b266114203f53596f9b6e6f00ebd79d178cdf2", size = 183910, upload-time = "2024-12-27T17:39:46.401Z" }, { url = "https://files.pythonhosted.org/packages/87/81/0c8c93f0b92c97e326b279795f9c5b956c5a97af28ca0fbb9fd86c83737a/msgspec-0.19.0-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:4cfc033c02c3e0aec52b71710d7f84cb3ca5eb407ab2ad23d75631153fdb1f12", size = 210633, upload-time = "2024-12-27T17:39:49.099Z" }, { url = "https://files.pythonhosted.org/packages/d0/ef/c5422ce8af73928d194a6606f8ae36e93a52fd5e8df5abd366903a5ca8da/msgspec-0.19.0-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:d911c442571605e17658ca2b416fd8579c5050ac9adc5e00c2cb3126c97f73bc", size = 213594, upload-time = "2024-12-27T17:39:51.204Z" }, { url = "https://files.pythonhosted.org/packages/19/2b/4137bc2ed45660444842d042be2cf5b18aa06efd2cda107cff18253b9653/msgspec-0.19.0-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:757b501fa57e24896cf40a831442b19a864f56d253679f34f260dcb002524a6c", size = 214053, upload-time = "2024-12-27T17:39:52.866Z" }, { url = "https://files.pythonhosted.org/packages/9d/e6/8ad51bdc806aac1dc501e8fe43f759f9ed7284043d722b53323ea421c360/msgspec-0.19.0-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:5f0f65f29b45e2816d8bded36e6b837a4bf5fb60ec4bc3c625fa2c6da4124537", size = 219081, upload-time = "2024-12-27T17:39:55.142Z" }, - { url = "https://files.pythonhosted.org/packages/b1/ef/27dd35a7049c9a4f4211c6cd6a8c9db0a50647546f003a5867827ec45391/msgspec-0.19.0-cp312-cp312-win_amd64.whl", hash = "sha256:067f0de1c33cfa0b6a8206562efdf6be5985b988b53dd244a8e06f993f27c8c0", size = 187467, upload-time = "2024-12-27T17:39:56.531Z" }, ] [[package]] @@ -1808,24 +2020,10 @@ version = "6.6.4" source = { registry = "https://pypi.org/simple" } sdist = { url = "https://files.pythonhosted.org/packages/69/7f/0652e6ed47ab288e3756ea9c0df8b14950781184d4bd7883f4d87dd41245/multidict-6.6.4.tar.gz", hash = "sha256:d2d4e4787672911b48350df02ed3fa3fffdc2f2e8ca06dd6afdf34189b76a9dd", size = 101843, upload-time = "2025-08-11T12:08:48.217Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/05/f6/512ffd8fd8b37fb2680e5ac35d788f1d71bbaf37789d21a820bdc441e565/multidict-6.6.4-cp312-cp312-macosx_10_13_universal2.whl", hash = "sha256:0ffb87be160942d56d7b87b0fdf098e81ed565add09eaa1294268c7f3caac4c8", size = 76516, upload-time = "2025-08-11T12:06:53.393Z" }, - { url = "https://files.pythonhosted.org/packages/99/58/45c3e75deb8855c36bd66cc1658007589662ba584dbf423d01df478dd1c5/multidict-6.6.4-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:d191de6cbab2aff5de6c5723101705fd044b3e4c7cfd587a1929b5028b9714b3", size = 45394, upload-time = "2025-08-11T12:06:54.555Z" }, - { url = "https://files.pythonhosted.org/packages/fd/ca/e8c4472a93a26e4507c0b8e1f0762c0d8a32de1328ef72fd704ef9cc5447/multidict-6.6.4-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:38a0956dd92d918ad5feff3db8fcb4a5eb7dba114da917e1a88475619781b57b", size = 43591, upload-time = "2025-08-11T12:06:55.672Z" }, - { url = "https://files.pythonhosted.org/packages/05/51/edf414f4df058574a7265034d04c935aa84a89e79ce90fcf4df211f47b16/multidict-6.6.4-cp312-cp312-manylinux1_i686.manylinux2014_i686.manylinux_2_17_i686.manylinux_2_5_i686.whl", hash = "sha256:6865f6d3b7900ae020b495d599fcf3765653bc927951c1abb959017f81ae8287", size = 237215, upload-time = "2025-08-11T12:06:57.213Z" }, { url = "https://files.pythonhosted.org/packages/c8/45/8b3d6dbad8cf3252553cc41abea09ad527b33ce47a5e199072620b296902/multidict-6.6.4-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:0a2088c126b6f72db6c9212ad827d0ba088c01d951cee25e758c450da732c138", size = 258299, upload-time = "2025-08-11T12:06:58.946Z" }, - { url = "https://files.pythonhosted.org/packages/3c/e8/8ca2e9a9f5a435fc6db40438a55730a4bf4956b554e487fa1b9ae920f825/multidict-6.6.4-cp312-cp312-manylinux2014_armv7l.manylinux_2_17_armv7l.manylinux_2_31_armv7l.whl", hash = "sha256:0f37bed7319b848097085d7d48116f545985db988e2256b2e6f00563a3416ee6", size = 242357, upload-time = "2025-08-11T12:07:00.301Z" }, - { url = "https://files.pythonhosted.org/packages/0f/84/80c77c99df05a75c28490b2af8f7cba2a12621186e0a8b0865d8e745c104/multidict-6.6.4-cp312-cp312-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:01368e3c94032ba6ca0b78e7ccb099643466cf24f8dc8eefcfdc0571d56e58f9", size = 268369, upload-time = "2025-08-11T12:07:01.638Z" }, - { url = "https://files.pythonhosted.org/packages/0d/e9/920bfa46c27b05fb3e1ad85121fd49f441492dca2449c5bcfe42e4565d8a/multidict-6.6.4-cp312-cp312-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:8fe323540c255db0bffee79ad7f048c909f2ab0edb87a597e1c17da6a54e493c", size = 269341, upload-time = "2025-08-11T12:07:02.943Z" }, { url = "https://files.pythonhosted.org/packages/af/65/753a2d8b05daf496f4a9c367fe844e90a1b2cac78e2be2c844200d10cc4c/multidict-6.6.4-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:b8eb3025f17b0a4c3cd08cda49acf312a19ad6e8a4edd9dbd591e6506d999402", size = 256100, upload-time = "2025-08-11T12:07:04.564Z" }, { url = "https://files.pythonhosted.org/packages/09/54/655be13ae324212bf0bc15d665a4e34844f34c206f78801be42f7a0a8aaa/multidict-6.6.4-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:bbc14f0365534d35a06970d6a83478b249752e922d662dc24d489af1aa0d1be7", size = 253584, upload-time = "2025-08-11T12:07:05.914Z" }, - { url = "https://files.pythonhosted.org/packages/5c/74/ab2039ecc05264b5cec73eb018ce417af3ebb384ae9c0e9ed42cb33f8151/multidict-6.6.4-cp312-cp312-musllinux_1_2_armv7l.whl", hash = "sha256:75aa52fba2d96bf972e85451b99d8e19cc37ce26fd016f6d4aa60da9ab2b005f", size = 251018, upload-time = "2025-08-11T12:07:08.301Z" }, - { url = "https://files.pythonhosted.org/packages/af/0a/ccbb244ac848e56c6427f2392741c06302bbfba49c0042f1eb3c5b606497/multidict-6.6.4-cp312-cp312-musllinux_1_2_i686.whl", hash = "sha256:4fefd4a815e362d4f011919d97d7b4a1e566f1dde83dc4ad8cfb5b41de1df68d", size = 251477, upload-time = "2025-08-11T12:07:10.248Z" }, - { url = "https://files.pythonhosted.org/packages/0e/b0/0ed49bba775b135937f52fe13922bc64a7eaf0a3ead84a36e8e4e446e096/multidict-6.6.4-cp312-cp312-musllinux_1_2_ppc64le.whl", hash = "sha256:db9801fe021f59a5b375ab778973127ca0ac52429a26e2fd86aa9508f4d26eb7", size = 263575, upload-time = "2025-08-11T12:07:11.928Z" }, - { url = "https://files.pythonhosted.org/packages/3e/d9/7fb85a85e14de2e44dfb6a24f03c41e2af8697a6df83daddb0e9b7569f73/multidict-6.6.4-cp312-cp312-musllinux_1_2_s390x.whl", hash = "sha256:a650629970fa21ac1fb06ba25dabfc5b8a2054fcbf6ae97c758aa956b8dba802", size = 259649, upload-time = "2025-08-11T12:07:13.244Z" }, { url = "https://files.pythonhosted.org/packages/03/9e/b3a459bcf9b6e74fa461a5222a10ff9b544cb1cd52fd482fb1b75ecda2a2/multidict-6.6.4-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:452ff5da78d4720d7516a3a2abd804957532dd69296cb77319c193e3ffb87e24", size = 251505, upload-time = "2025-08-11T12:07:14.57Z" }, - { url = "https://files.pythonhosted.org/packages/86/a2/8022f78f041dfe6d71e364001a5cf987c30edfc83c8a5fb7a3f0974cff39/multidict-6.6.4-cp312-cp312-win32.whl", hash = "sha256:8c2fcb12136530ed19572bbba61b407f655e3953ba669b96a35036a11a485793", size = 41888, upload-time = "2025-08-11T12:07:15.904Z" }, - { url = "https://files.pythonhosted.org/packages/c7/eb/d88b1780d43a56db2cba24289fa744a9d216c1a8546a0dc3956563fd53ea/multidict-6.6.4-cp312-cp312-win_amd64.whl", hash = "sha256:047d9425860a8c9544fed1b9584f0c8bcd31bcde9568b047c5e567a1025ecd6e", size = 46072, upload-time = "2025-08-11T12:07:17.045Z" }, - { url = "https://files.pythonhosted.org/packages/9f/16/b929320bf5750e2d9d4931835a4c638a19d2494a5b519caaaa7492ebe105/multidict-6.6.4-cp312-cp312-win_arm64.whl", hash = "sha256:14754eb72feaa1e8ae528468f24250dd997b8e2188c3d2f593f9eba259e4b364", size = 43222, upload-time = "2025-08-11T12:07:18.328Z" }, { url = "https://files.pythonhosted.org/packages/fd/69/b547032297c7e63ba2af494edba695d781af8a0c6e89e4d06cf848b21d80/multidict-6.6.4-py3-none-any.whl", hash = "sha256:27d8f8e125c07cb954e54d75d04905a9bba8a439c1d84aca94949d4d03d8601c", size = 12313, upload-time = "2025-08-11T12:08:46.891Z" }, ] @@ -1834,7 +2032,7 @@ name = "multiprocess" version = "0.70.16" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "dill" }, + { name = "dill", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, ] sdist = { url = "https://files.pythonhosted.org/packages/b5/ae/04f39c5d0d0def03247c2893d6f2b83c136bf3320a2154d7b8858f2ba72d/multiprocess-0.70.16.tar.gz", hash = "sha256:161af703d4652a0e1410be6abccecde4a7ddffd19341be0a7011b94aeb171ac1", size = 1772603, upload-time = "2024-01-28T18:52:34.85Z" } wheels = [ @@ -1869,39 +2067,58 @@ version = "1.13.0" source = { registry = "https://pypi.org/simple" } sdist = { url = "https://files.pythonhosted.org/packages/43/73/79a0b22fc731989c708068427579e840a6cf4e937fe7ae5c5d0b7356ac22/ninja-1.13.0.tar.gz", hash = "sha256:4a40ce995ded54d9dc24f8ea37ff3bf62ad192b547f6c7126e7e25045e76f978", size = 242558, upload-time = "2025-08-11T15:10:19.421Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/3c/74/d02409ed2aa865e051b7edda22ad416a39d81a84980f544f8de717cab133/ninja-1.13.0-py3-none-macosx_10_9_universal2.whl", hash = "sha256:fa2a8bfc62e31b08f83127d1613d10821775a0eb334197154c4d6067b7068ff1", size = 310125, upload-time = "2025-08-11T15:09:50.971Z" }, { url = "https://files.pythonhosted.org/packages/8e/de/6e1cd6b84b412ac1ef327b76f0641aeb5dcc01e9d3f9eee0286d0c34fd93/ninja-1.13.0-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:3d00c692fb717fd511abeb44b8c5d00340c36938c12d6538ba989fe764e79630", size = 177467, upload-time = "2025-08-11T15:09:52.767Z" }, - { url = "https://files.pythonhosted.org/packages/c8/83/49320fb6e58ae3c079381e333575fdbcf1cca3506ee160a2dcce775046fa/ninja-1.13.0-py3-none-manylinux2014_i686.manylinux_2_17_i686.whl", hash = "sha256:be7f478ff9f96a128b599a964fc60a6a87b9fa332ee1bd44fa243ac88d50291c", size = 187834, upload-time = "2025-08-11T15:09:54.115Z" }, - { url = "https://files.pythonhosted.org/packages/56/c7/ba22748fb59f7f896b609cd3e568d28a0a367a6d953c24c461fe04fc4433/ninja-1.13.0-py3-none-manylinux2014_ppc64le.manylinux_2_17_ppc64le.whl", hash = "sha256:60056592cf495e9a6a4bea3cd178903056ecb0943e4de45a2ea825edb6dc8d3e", size = 202736, upload-time = "2025-08-11T15:09:55.745Z" }, - { url = "https://files.pythonhosted.org/packages/79/22/d1de07632b78ac8e6b785f41fa9aad7a978ec8c0a1bf15772def36d77aac/ninja-1.13.0-py3-none-manylinux2014_s390x.manylinux_2_17_s390x.whl", hash = "sha256:1c97223cdda0417f414bf864cfb73b72d8777e57ebb279c5f6de368de0062988", size = 179034, upload-time = "2025-08-11T15:09:57.394Z" }, { url = "https://files.pythonhosted.org/packages/ed/de/0e6edf44d6a04dabd0318a519125ed0415ce437ad5a1ec9b9be03d9048cf/ninja-1.13.0-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:fb46acf6b93b8dd0322adc3a4945452a4e774b75b91293bafcc7b7f8e6517dfa", size = 180716, upload-time = "2025-08-11T15:09:58.696Z" }, - { url = "https://files.pythonhosted.org/packages/54/28/938b562f9057aaa4d6bfbeaa05e81899a47aebb3ba6751e36c027a7f5ff7/ninja-1.13.0-py3-none-manylinux_2_28_armv7l.manylinux_2_31_armv7l.whl", hash = "sha256:4be9c1b082d244b1ad7ef41eb8ab088aae8c109a9f3f0b3e56a252d3e00f42c1", size = 146843, upload-time = "2025-08-11T15:10:00.046Z" }, - { url = "https://files.pythonhosted.org/packages/2a/fb/d06a3838de4f8ab866e44ee52a797b5491df823901c54943b2adb0389fbb/ninja-1.13.0-py3-none-manylinux_2_31_riscv64.whl", hash = "sha256:6739d3352073341ad284246f81339a384eec091d9851a886dfa5b00a6d48b3e2", size = 154402, upload-time = "2025-08-11T15:10:01.657Z" }, { url = "https://files.pythonhosted.org/packages/31/bf/0d7808af695ceddc763cf251b84a9892cd7f51622dc8b4c89d5012779f06/ninja-1.13.0-py3-none-musllinux_1_2_aarch64.whl", hash = "sha256:11be2d22027bde06f14c343f01d31446747dbb51e72d00decca2eb99be911e2f", size = 552388, upload-time = "2025-08-11T15:10:03.349Z" }, - { url = "https://files.pythonhosted.org/packages/9d/70/c99d0c2c809f992752453cce312848abb3b1607e56d4cd1b6cded317351a/ninja-1.13.0-py3-none-musllinux_1_2_armv7l.whl", hash = "sha256:aa45b4037b313c2f698bc13306239b8b93b4680eb47e287773156ac9e9304714", size = 472501, upload-time = "2025-08-11T15:10:04.735Z" }, - { url = "https://files.pythonhosted.org/packages/9f/43/c217b1153f0e499652f5e0766da8523ce3480f0a951039c7af115e224d55/ninja-1.13.0-py3-none-musllinux_1_2_i686.whl", hash = "sha256:5f8e1e8a1a30835eeb51db05cf5a67151ad37542f5a4af2a438e9490915e5b72", size = 638280, upload-time = "2025-08-11T15:10:06.512Z" }, - { url = "https://files.pythonhosted.org/packages/8c/45/9151bba2c8d0ae2b6260f71696330590de5850e5574b7b5694dce6023e20/ninja-1.13.0-py3-none-musllinux_1_2_ppc64le.whl", hash = "sha256:3d7d7779d12cb20c6d054c61b702139fd23a7a964ec8f2c823f1ab1b084150db", size = 642420, upload-time = "2025-08-11T15:10:08.35Z" }, - { url = "https://files.pythonhosted.org/packages/3c/fb/95752eb635bb8ad27d101d71bef15bc63049de23f299e312878fc21cb2da/ninja-1.13.0-py3-none-musllinux_1_2_riscv64.whl", hash = "sha256:d741a5e6754e0bda767e3274a0f0deeef4807f1fec6c0d7921a0244018926ae5", size = 585106, upload-time = "2025-08-11T15:10:09.818Z" }, - { url = "https://files.pythonhosted.org/packages/c1/31/aa56a1a286703800c0cbe39fb4e82811c277772dc8cd084f442dd8e2938a/ninja-1.13.0-py3-none-musllinux_1_2_s390x.whl", hash = "sha256:e8bad11f8a00b64137e9b315b137d8bb6cbf3086fbdc43bf1f90fd33324d2e96", size = 707138, upload-time = "2025-08-11T15:10:11.366Z" }, { url = "https://files.pythonhosted.org/packages/34/6f/5f5a54a1041af945130abdb2b8529cbef0cdcbbf9bcf3f4195378319d29a/ninja-1.13.0-py3-none-musllinux_1_2_x86_64.whl", hash = "sha256:b4f2a072db3c0f944c32793e91532d8948d20d9ab83da9c0c7c15b5768072200", size = 581758, upload-time = "2025-08-11T15:10:13.295Z" }, - { url = "https://files.pythonhosted.org/packages/95/97/51359c77527d45943fe7a94d00a3843b81162e6c4244b3579fe8fc54cb9c/ninja-1.13.0-py3-none-win32.whl", hash = "sha256:8cfbb80b4a53456ae8a39f90ae3d7a2129f45ea164f43fadfa15dc38c4aef1c9", size = 267201, upload-time = "2025-08-11T15:10:15.158Z" }, - { url = "https://files.pythonhosted.org/packages/29/45/c0adfbfb0b5895aa18cec400c535b4f7ff3e52536e0403602fc1a23f7de9/ninja-1.13.0-py3-none-win_amd64.whl", hash = "sha256:fb8ee8719f8af47fed145cced4a85f0755dd55d45b2bddaf7431fa89803c5f3e", size = 309975, upload-time = "2025-08-11T15:10:16.697Z" }, - { url = "https://files.pythonhosted.org/packages/df/93/a7b983643d1253bb223234b5b226e69de6cda02b76cdca7770f684b795f5/ninja-1.13.0-py3-none-win_arm64.whl", hash = "sha256:3c0b40b1f0bba764644385319028650087b4c1b18cdfa6f45cb39a3669b81aa9", size = 290806, upload-time = "2025-08-11T15:10:18.018Z" }, ] [[package]] -name = "nltk" -version = "3.9.1" +name = "nixl" +version = "0.10.1" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "click" }, - { name = "joblib" }, - { name = "regex" }, - { name = "tqdm" }, + { name = "nixl-cu12", version = "0.10.1", source = { registry = "https://pypi.org/simple" }, marker = "platform_machine == 'aarch64' and sys_platform == 'linux'" }, + { name = "nixl-cu12", version = "0.10.1", source = { url = "https://github.com/PrimeIntellect-ai/prime-rl/releases/download/v0.5.0/nixl_cu12-0.10.1-cp312-cp312-linux_x86_64.whl" }, marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/bb/a4/10e80e623790d3fb070e966b36bced009419d483c92ae0df6645230f606f/nixl-0.10.1-py3-none-any.whl", hash = "sha256:616465673dae5180d296525a03237af4cd5f2c00c3228d185bc06dbe621509b7", size = 6680, upload-time = "2026-03-03T19:55:44.848Z" }, +] + +[[package]] +name = "nixl-cu12" +version = "0.10.1" +source = { registry = "https://pypi.org/simple" } +resolution-markers = [ + "platform_machine == 'aarch64' and sys_platform == 'linux'", +] +dependencies = [ + { name = "numpy", marker = "platform_machine == 'aarch64' and sys_platform == 'linux'" }, + { name = "torch", marker = "platform_machine == 'aarch64' and sys_platform == 'linux'" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/3c/87/db8be88ad32c2d042420b6fd9ffd4a149f9a0d7f0e86b3f543be2eeeedd2/nltk-3.9.1.tar.gz", hash = "sha256:87d127bd3de4bd89a4f81265e5fa59cb1b199b27440175370f7417d2bc7ae868", size = 2904691, upload-time = "2024-08-18T19:48:37.769Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/4d/66/7d9e26593edda06e8cb531874633f7c2372279c3b0f46235539fe546df8b/nltk-3.9.1-py3-none-any.whl", hash = "sha256:4fa26829c5b00715afe3061398a8989dc643b92ce7dd93fb4585a70930d168a1", size = 1505442, upload-time = "2024-08-18T19:48:21.909Z" }, + { url = "https://files.pythonhosted.org/packages/d4/21/c1f34212a3e612b5c0da389c1aa3557588d9cfc2adf864914b32e270847b/nixl_cu12-0.10.1-cp312-cp312-manylinux_2_28_aarch64.whl", hash = "sha256:b4712c6e0f18f57fee34cd970faac01480f0caf12da33d4a40ef4e9096a4caf7", size = 50227261, upload-time = "2026-03-03T19:55:17.177Z" }, +] + +[[package]] +name = "nixl-cu12" +version = "0.10.1" +source = { url = "https://github.com/PrimeIntellect-ai/prime-rl/releases/download/v0.5.0/nixl_cu12-0.10.1-cp312-cp312-linux_x86_64.whl" } +resolution-markers = [ + "platform_machine == 'x86_64' and sys_platform == 'linux'", +] +dependencies = [ + { name = "numpy", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "torch", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, +] +wheels = [ + { url = "https://github.com/PrimeIntellect-ai/prime-rl/releases/download/v0.5.0/nixl_cu12-0.10.1-cp312-cp312-linux_x86_64.whl", hash = "sha256:bd01eb9e52f1affee3330062baa2f3bea7d711664f7cb602bf5d3f694daeb796" }, +] + +[package.metadata] +requires-dist = [ + { name = "numpy" }, + { name = "torch" }, ] [[package]] @@ -1915,19 +2132,16 @@ wheels = [ [[package]] name = "numba" -version = "0.61.2" +version = "0.65.0" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "llvmlite" }, - { name = "numpy" }, + { name = "llvmlite", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "numpy", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/1c/a0/e21f57604304aa03ebb8e098429222722ad99176a4f979d34af1d1ee80da/numba-0.61.2.tar.gz", hash = "sha256:8750ee147940a6637b80ecf7f95062185ad8726c8c28a2295b8ec1160a196f7d", size = 2820615, upload-time = "2025-04-09T02:58:07.659Z" } +sdist = { url = "https://files.pythonhosted.org/packages/49/61/7299643b9c18d669e04be7c5bcb64d985070d07553274817b45b049e7bfe/numba-0.65.0.tar.gz", hash = "sha256:edad0d9f6682e93624c00125a471ae4df186175d71fd604c983c377cdc03e68b", size = 2764131, upload-time = "2026-04-01T03:52:01.946Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/b4/a0/c6b7b9c615cfa3b98c4c63f4316e3f6b3bbe2387740277006551784218cd/numba-0.61.2-cp312-cp312-macosx_10_14_x86_64.whl", hash = "sha256:34fba9406078bac7ab052efbf0d13939426c753ad72946baaa5bf9ae0ebb8dd2", size = 2776626, upload-time = "2025-04-09T02:57:51.857Z" }, - { url = "https://files.pythonhosted.org/packages/92/4a/fe4e3c2ecad72d88f5f8cd04e7f7cff49e718398a2fac02d2947480a00ca/numba-0.61.2-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:4ddce10009bc097b080fc96876d14c051cc0c7679e99de3e0af59014dab7dfe8", size = 2779287, upload-time = "2025-04-09T02:57:53.658Z" }, - { url = "https://files.pythonhosted.org/packages/9a/2d/e518df036feab381c23a624dac47f8445ac55686ec7f11083655eb707da3/numba-0.61.2-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:5b1bb509d01f23d70325d3a5a0e237cbc9544dd50e50588bc581ba860c213546", size = 3885928, upload-time = "2025-04-09T02:57:55.206Z" }, - { url = "https://files.pythonhosted.org/packages/10/0f/23cced68ead67b75d77cfcca3df4991d1855c897ee0ff3fe25a56ed82108/numba-0.61.2-cp312-cp312-manylinux_2_28_aarch64.whl", hash = "sha256:48a53a3de8f8793526cbe330f2a39fe9a6638efcbf11bd63f3d2f9757ae345cd", size = 3577115, upload-time = "2025-04-09T02:57:56.818Z" }, - { url = "https://files.pythonhosted.org/packages/68/1d/ddb3e704c5a8fb90142bf9dc195c27db02a08a99f037395503bfbc1d14b3/numba-0.61.2-cp312-cp312-win_amd64.whl", hash = "sha256:97cf4f12c728cf77c9c1d7c23707e4d8fb4632b46275f8f3397de33e5877af18", size = 2831929, upload-time = "2025-04-09T02:57:58.45Z" }, + { url = "https://files.pythonhosted.org/packages/73/36/88406bd58600cc696417b8e5dd6a056478da808f3eaf48d18e2421e0c2d9/numba-0.65.0-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:a52d92ffd297c10364bce60cd1fcb88f99284ab5df085f2c6bcd1cb33b529a6f", size = 3801411, upload-time = "2026-04-01T03:51:34.321Z" }, + { url = "https://files.pythonhosted.org/packages/0c/61/ce753a1d7646dd477e16d15e89473703faebb8995d2f71d7ad69a540b565/numba-0.65.0-cp312-cp312-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:da8e371e328c06d0010c3d8b44b21858652831b85bcfba78cb22c042e22dbd8e", size = 3501622, upload-time = "2026-04-01T03:51:36.348Z" }, ] [[package]] @@ -1936,16 +2150,10 @@ version = "2.2.6" source = { registry = "https://pypi.org/simple" } sdist = { url = "https://files.pythonhosted.org/packages/76/21/7d2a95e4bba9dc13d043ee156a356c0a8f0c6309dff6b21b4d71a073b8a8/numpy-2.2.6.tar.gz", hash = "sha256:e29554e2bef54a90aa5cc07da6ce955accb83f21ab5de01a62c8478897b264fd", size = 20276440, upload-time = "2025-05-17T22:38:04.611Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/82/5d/c00588b6cf18e1da539b45d3598d3557084990dcc4331960c15ee776ee41/numpy-2.2.6-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:41c5a21f4a04fa86436124d388f6ed60a9343a6f767fced1a8a71c3fbca038ff", size = 20875348, upload-time = "2025-05-17T21:34:39.648Z" }, - { url = "https://files.pythonhosted.org/packages/66/ee/560deadcdde6c2f90200450d5938f63a34b37e27ebff162810f716f6a230/numpy-2.2.6-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:de749064336d37e340f640b05f24e9e3dd678c57318c7289d222a8a2f543e90c", size = 14119362, upload-time = "2025-05-17T21:35:01.241Z" }, - { url = "https://files.pythonhosted.org/packages/3c/65/4baa99f1c53b30adf0acd9a5519078871ddde8d2339dc5a7fde80d9d87da/numpy-2.2.6-cp312-cp312-macosx_14_0_arm64.whl", hash = "sha256:894b3a42502226a1cac872f840030665f33326fc3dac8e57c607905773cdcde3", size = 5084103, upload-time = "2025-05-17T21:35:10.622Z" }, - { url = "https://files.pythonhosted.org/packages/cc/89/e5a34c071a0570cc40c9a54eb472d113eea6d002e9ae12bb3a8407fb912e/numpy-2.2.6-cp312-cp312-macosx_14_0_x86_64.whl", hash = "sha256:71594f7c51a18e728451bb50cc60a3ce4e6538822731b2933209a1f3614e9282", size = 6625382, upload-time = "2025-05-17T21:35:21.414Z" }, { url = "https://files.pythonhosted.org/packages/f8/35/8c80729f1ff76b3921d5c9487c7ac3de9b2a103b1cd05e905b3090513510/numpy-2.2.6-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:f2618db89be1b4e05f7a1a847a9c1c0abd63e63a1607d892dd54668dd92faf87", size = 14018462, upload-time = "2025-05-17T21:35:42.174Z" }, { url = "https://files.pythonhosted.org/packages/8c/3d/1e1db36cfd41f895d266b103df00ca5b3cbe965184df824dec5c08c6b803/numpy-2.2.6-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:fd83c01228a688733f1ded5201c678f0c53ecc1006ffbc404db9f7a899ac6249", size = 16527618, upload-time = "2025-05-17T21:36:06.711Z" }, { url = "https://files.pythonhosted.org/packages/61/c6/03ed30992602c85aa3cd95b9070a514f8b3c33e31124694438d88809ae36/numpy-2.2.6-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:37c0ca431f82cd5fa716eca9506aefcabc247fb27ba69c5062a6d3ade8cf8f49", size = 15505511, upload-time = "2025-05-17T21:36:29.965Z" }, { url = "https://files.pythonhosted.org/packages/b7/25/5761d832a81df431e260719ec45de696414266613c9ee268394dd5ad8236/numpy-2.2.6-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:fe27749d33bb772c80dcd84ae7e8df2adc920ae8297400dabec45f0dedb3f6de", size = 18313783, upload-time = "2025-05-17T21:36:56.883Z" }, - { url = "https://files.pythonhosted.org/packages/57/0a/72d5a3527c5ebffcd47bde9162c39fae1f90138c961e5296491ce778e682/numpy-2.2.6-cp312-cp312-win32.whl", hash = "sha256:4eeaae00d789f66c7a25ac5f34b71a7035bb474e679f410e5e1a94deb24cf2d4", size = 6246506, upload-time = "2025-05-17T21:37:07.368Z" }, - { url = "https://files.pythonhosted.org/packages/36/fa/8c9210162ca1b88529ab76b41ba02d433fd54fecaf6feb70ef9f124683f1/numpy-2.2.6-cp312-cp312-win_amd64.whl", hash = "sha256:c1f9540be57940698ed329904db803cf7a402f3fc200bfe599334c9bd84a40b2", size = 12614190, upload-time = "2025-05-17T21:37:26.213Z" }, ] [[package]] @@ -1986,14 +2194,14 @@ wheels = [ [[package]] name = "nvidia-cudnn-cu12" -version = "9.10.2.21" +version = "9.18.1.3" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "nvidia-cublas-cu12", marker = "sys_platform == 'linux'" }, + { name = "nvidia-cublas-cu12", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, ] wheels = [ - { url = "https://files.pythonhosted.org/packages/fa/41/e79269ce215c857c935fd86bcfe91a451a584dfc27f1e068f568b9ad1ab7/nvidia_cudnn_cu12-9.10.2.21-py3-none-manylinux_2_27_aarch64.whl", hash = "sha256:c9132cc3f8958447b4910a1720036d9eff5928cc3179b0a51fb6d167c6cc87d8", size = 705026878, upload-time = "2025-06-06T21:52:51.348Z" }, - { url = "https://files.pythonhosted.org/packages/ba/51/e123d997aa098c61d029f76663dedbfb9bc8dcf8c60cbd6adbe42f76d049/nvidia_cudnn_cu12-9.10.2.21-py3-none-manylinux_2_27_x86_64.whl", hash = "sha256:949452be657fa16687d0930933f032835951ef0892b37d2d53824d1a84dc97a8", size = 706758467, upload-time = "2025-06-06T21:54:08.597Z" }, + { url = "https://files.pythonhosted.org/packages/34/d3/f957d4af442aae223544288d417ac44e6276cf81e11979ef8b90e0017cf8/nvidia_cudnn_cu12-9.18.1.3-py3-none-manylinux_2_27_aarch64.whl", hash = "sha256:6966836f89c41fb4ad67361b5c5f3bf81eccf2b45d46f3b9eb84e8b7557263a7", size = 647470574, upload-time = "2026-01-23T20:29:24.813Z" }, + { url = "https://files.pythonhosted.org/packages/ce/a0/c2fe112e2dd28e97ae7e58364795f0aba6e474c22bb29a1dada729f95f65/nvidia_cudnn_cu12-9.18.1.3-py3-none-manylinux_2_27_x86_64.whl", hash = "sha256:733df57fc960e18f8dd27168e0e5495ef80a9d01d34647fb6774acc0cbee86e5", size = 648588125, upload-time = "2026-01-23T20:31:04.197Z" }, ] [[package]] @@ -2003,7 +2211,6 @@ source = { registry = "https://pypi.org/simple" } wheels = [ { url = "https://files.pythonhosted.org/packages/00/39/79b606e805abd67ab4fa72f752a5413a496159f10d94fbdb1d67bb5ae86c/nvidia_cudnn_frontend-1.16.0-cp312-cp312-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:dd6fdd71c0896ff2ca1809d914cbd17f2904d55863f8881f47946e1d634c7a88", size = 1839271, upload-time = "2025-11-07T01:29:53.06Z" }, { url = "https://files.pythonhosted.org/packages/09/21/a0e0d50ba8d7b639fe635500fee0d9c0319561b1ae72176d7024ec04b439/nvidia_cudnn_frontend-1.16.0-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:16efb069d4bda4d3b99134f59f376cfd4d09558298bd96af778fdc7f2851e696", size = 1954062, upload-time = "2025-11-07T01:32:18.556Z" }, - { url = "https://files.pythonhosted.org/packages/ce/d6/30ae67bb9c010e9459d1211c56d73373eb4e3dd9f57f4c3c1fe0966efcb1/nvidia_cudnn_frontend-1.16.0-cp312-cp312-win_amd64.whl", hash = "sha256:7b7860db03767c158accbe0b4e9c9553506513cc970ff08ed28c7761681ac466", size = 1368435, upload-time = "2025-11-07T01:26:28.022Z" }, ] [[package]] @@ -2011,7 +2218,7 @@ name = "nvidia-cufft-cu12" version = "11.3.3.83" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "nvidia-nvjitlink-cu12", marker = "sys_platform == 'linux'" }, + { name = "nvidia-nvjitlink-cu12", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, ] wheels = [ { url = "https://files.pythonhosted.org/packages/60/bc/7771846d3a0272026c416fbb7e5f4c1f146d6d80704534d0b187dd6f4800/nvidia_cufft_cu12-11.3.3.83-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:848ef7224d6305cdb2a4df928759dca7b1201874787083b6e7550dd6765ce69a", size = 193109211, upload-time = "2025-03-07T01:44:56.873Z" }, @@ -2041,9 +2248,9 @@ name = "nvidia-cusolver-cu12" version = "11.7.3.90" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "nvidia-cublas-cu12", marker = "sys_platform == 'linux'" }, - { name = "nvidia-cusparse-cu12", marker = "sys_platform == 'linux'" }, - { name = "nvidia-nvjitlink-cu12", marker = "sys_platform == 'linux'" }, + { name = "nvidia-cublas-cu12", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "nvidia-cusparse-cu12", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "nvidia-nvjitlink-cu12", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, ] wheels = [ { url = "https://files.pythonhosted.org/packages/c8/32/f7cd6ce8a7690544d084ea21c26e910a97e077c9b7f07bf5de623ee19981/nvidia_cusolver_cu12-11.7.3.90-py3-none-manylinux_2_27_aarch64.whl", hash = "sha256:db9ed69dbef9715071232caa9b69c52ac7de3a95773c2db65bdba85916e4e5c0", size = 267229841, upload-time = "2025-03-07T01:46:54.356Z" }, @@ -2055,7 +2262,7 @@ name = "nvidia-cusparse-cu12" version = "12.5.8.93" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "nvidia-nvjitlink-cu12", marker = "sys_platform == 'linux'" }, + { name = "nvidia-nvjitlink-cu12", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, ] wheels = [ { url = "https://files.pythonhosted.org/packages/bc/f7/cd777c4109681367721b00a106f491e0d0d15cfa1fd59672ce580ce42a97/nvidia_cusparse_cu12-12.5.8.93-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:9b6c161cb130be1a07a27ea6923df8141f3c295852f4b260c65f18f3e0a091dc", size = 288117129, upload-time = "2025-03-07T01:47:40.407Z" }, @@ -2073,16 +2280,27 @@ wheels = [ [[package]] name = "nvidia-cutlass-dsl" -version = "4.3.3" +version = "4.4.1" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "cuda-python" }, - { name = "numpy" }, - { name = "typing-extensions" }, + { name = "nvidia-cutlass-dsl-libs-base", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/0d/a8/d9f2b82bf6f6b48502267fcf2fa7b229392affb68a6092da92b0edef7476/nvidia_cutlass_dsl-4.4.1-py3-none-any.whl", hash = "sha256:7b8ffa0117be35ef6c9a88f4462ee2a794efd0f7d9f65090e10a953e434fbfce", size = 10167, upload-time = "2026-02-27T09:37:34.551Z" }, +] + +[[package]] +name = "nvidia-cutlass-dsl-libs-base" +version = "4.4.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "cuda-python", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "numpy", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "typing-extensions", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, ] wheels = [ - { url = "https://files.pythonhosted.org/packages/c6/d4/7c5ef53ccf75d7f99a9ea29cae9f9c0233229b75b3b22f85a4ef4f52e6ab/nvidia_cutlass_dsl-4.3.3-cp312-cp312-manylinux_2_28_aarch64.whl", hash = "sha256:3278526f54bddd920d8e539771e5820c6166c549a1e67813375025f39417dec6", size = 58734009, upload-time = "2025-12-10T09:23:29.305Z" }, - { url = "https://files.pythonhosted.org/packages/88/a8/a27562194cc4182c67793cd21c5dbf9468cd5a49c775a487153c6f28364c/nvidia_cutlass_dsl-4.3.3-cp312-cp312-manylinux_2_28_x86_64.whl", hash = "sha256:f2b25816b8bb8bc332bcbf6fc341347b5d728344cf185c65af0dd73e8503d5c7", size = 58596724, upload-time = "2025-12-10T11:01:07.228Z" }, + { url = "https://files.pythonhosted.org/packages/f5/cd/d09f6c998a9d52372d97d85d6561392d745ca00cf46de956d7cd7ec608cf/nvidia_cutlass_dsl_libs_base-4.4.1-cp312-cp312-manylinux_2_28_aarch64.whl", hash = "sha256:74192716b18c1825382723891842f87fa2a045b4b100c5c0f474042731e21e86", size = 75458464, upload-time = "2026-02-27T09:45:01.155Z" }, + { url = "https://files.pythonhosted.org/packages/b4/c1/acca814bc209562ef6cefbdec2ca36520f9a0380cdc7c6feaa69874bb50d/nvidia_cutlass_dsl_libs_base-4.4.1-cp312-cp312-manylinux_2_28_x86_64.whl", hash = "sha256:ba5e3d7148f7882911bb3cb453c313c790d1c2096bdfdd2d96da2123cf562201", size = 74347149, upload-time = "2026-02-27T09:45:56.602Z" }, ] [[package]] @@ -2096,11 +2314,11 @@ wheels = [ [[package]] name = "nvidia-nccl-cu12" -version = "2.27.5" +version = "2.28.9" source = { registry = "https://pypi.org/simple" } wheels = [ - { url = "https://files.pythonhosted.org/packages/bb/1c/857979db0ef194ca5e21478a0612bcdbbe59458d7694361882279947b349/nvidia_nccl_cu12-2.27.5-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:31432ad4d1fb1004eb0c56203dc9bc2178a1ba69d1d9e02d64a6938ab5e40e7a", size = 322400625, upload-time = "2025-06-26T04:11:04.496Z" }, - { url = "https://files.pythonhosted.org/packages/6e/89/f7a07dc961b60645dbbf42e80f2bc85ade7feb9a491b11a1e973aa00071f/nvidia_nccl_cu12-2.27.5-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:ad730cf15cb5d25fe849c6e6ca9eb5b76db16a80f13f425ac68d8e2e55624457", size = 322348229, upload-time = "2025-06-26T04:11:28.385Z" }, + { url = "https://files.pythonhosted.org/packages/08/c4/120d2dfd92dff2c776d68f361ff8705fdea2ca64e20b612fab0fd3f581ac/nvidia_nccl_cu12-2.28.9-py3-none-manylinux_2_18_aarch64.whl", hash = "sha256:50a36e01c4a090b9f9c47d92cec54964de6b9fcb3362d0e19b8ffc6323c21b60", size = 296766525, upload-time = "2025-11-18T05:49:16.094Z" }, + { url = "https://files.pythonhosted.org/packages/4a/4e/44dbb46b3d1b0ec61afda8e84837870f2f9ace33c564317d59b70bc19d3e/nvidia_nccl_cu12-2.28.9-py3-none-manylinux_2_18_x86_64.whl", hash = "sha256:485776daa8447da5da39681af455aa3b2c2586ddcf4af8772495e7c532c7e5ab", size = 296782137, upload-time = "2025-11-18T05:49:34.248Z" }, ] [[package]] @@ -2114,11 +2332,11 @@ wheels = [ [[package]] name = "nvidia-nvshmem-cu12" -version = "3.3.20" +version = "3.4.5" source = { registry = "https://pypi.org/simple" } wheels = [ - { url = "https://files.pythonhosted.org/packages/92/9d/3dd98852568fb845ec1f7902c90a22b240fe1cbabda411ccedf2fd737b7b/nvidia_nvshmem_cu12-3.3.20-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:0b0b960da3842212758e4fa4696b94f129090b30e5122fea3c5345916545cff0", size = 124484616, upload-time = "2025-08-04T20:24:59.172Z" }, - { url = "https://files.pythonhosted.org/packages/3b/6c/99acb2f9eb85c29fc6f3a7ac4dccfd992e22666dd08a642b303311326a97/nvidia_nvshmem_cu12-3.3.20-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:d00f26d3f9b2e3c3065be895e3059d6479ea5c638a3f38c9fec49b1b9dd7c1e5", size = 124657145, upload-time = "2025-08-04T20:25:19.995Z" }, + { url = "https://files.pythonhosted.org/packages/1d/6a/03aa43cc9bd3ad91553a88b5f6fb25ed6a3752ae86ce2180221962bc2aa5/nvidia_nvshmem_cu12-3.4.5-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:0b48363fc6964dede448029434c6abed6c5e37f823cb43c3bcde7ecfc0457e15", size = 138936938, upload-time = "2025-09-06T00:32:05.589Z" }, + { url = "https://files.pythonhosted.org/packages/b5/09/6ea3ea725f82e1e76684f0708bbedd871fc96da89945adeba65c3835a64c/nvidia_nvshmem_cu12-3.4.5-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:042f2500f24c021db8a06c5eec2539027d57460e1c1a762055a6554f72c369bd", size = 139103095, upload-time = "2025-09-06T00:32:31.266Z" }, ] [[package]] @@ -2130,23 +2348,48 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/a2/eb/86626c1bbc2edb86323022371c39aa48df6fd8b0a1647bc274577f72e90b/nvidia_nvtx_cu12-12.8.90-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:5b17e2001cc0d751a5bc2c6ec6d26ad95913324a4adb86788c944f8ce9ba441f", size = 89954, upload-time = "2025-03-07T01:42:44.131Z" }, ] +[[package]] +name = "oauthlib" +version = "3.3.1" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/0b/5f/19930f824ffeb0ad4372da4812c50edbd1434f678c90c2733e1188edfc63/oauthlib-3.3.1.tar.gz", hash = "sha256:0f0f8aa759826a193cf66c12ea1af1637f87b9b4622d46e866952bb022e538c9", size = 185918, upload-time = "2025-06-19T22:48:08.269Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/be/9c/92789c596b8df838baa98fa71844d84283302f7604ed565dafe5a6b5041a/oauthlib-3.3.1-py3-none-any.whl", hash = "sha256:88119c938d2b8fb88561af5f6ee0eec8cc8d552b7bb1f712743136eb7523b7a1", size = 160065, upload-time = "2025-06-19T22:48:06.508Z" }, +] + +[[package]] +name = "onnxruntime" +version = "1.24.3" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "flatbuffers", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "numpy", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "packaging", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "protobuf", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "sympy", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/1c/dc/1f5489f7b21817d4ad352bf7a92a252bd5b438bcbaa7ad20ea50814edc79/onnxruntime-1.24.3-cp312-cp312-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:0a9847b870b6cb462652b547bc98c49e0efb67553410a082fde1918a38707452", size = 15150105, upload-time = "2026-03-05T16:34:56.897Z" }, + { url = "https://files.pythonhosted.org/packages/28/7c/fd253da53594ab8efbefdc85b3638620ab1a6aab6eb7028a513c853559ce/onnxruntime-1.24.3-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:b354afce3333f2859c7e8706d84b6c552beac39233bcd3141ce7ab77b4cabb5d", size = 17237101, upload-time = "2026-03-05T17:18:02.561Z" }, +] + [[package]] name = "openai" -version = "2.3.0" +version = "2.24.0" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "anyio" }, - { name = "distro" }, - { name = "httpx" }, - { name = "jiter" }, - { name = "pydantic" }, - { name = "sniffio" }, - { name = "tqdm" }, - { name = "typing-extensions" }, + { name = "anyio", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "distro", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "httpx", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "jiter", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "pydantic", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "sniffio", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "tqdm", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "typing-extensions", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/de/90/8f26554d24d63ed4f94d33c24271559863223a67e624f4d2e65ba8e48dca/openai-2.3.0.tar.gz", hash = "sha256:8d213ee5aaf91737faea2d7fc1cd608657a5367a18966372a3756ceaabfbd812", size = 589616, upload-time = "2025-10-10T01:12:50.851Z" } +sdist = { url = "https://files.pythonhosted.org/packages/55/13/17e87641b89b74552ed408a92b231283786523edddc95f3545809fab673c/openai-2.24.0.tar.gz", hash = "sha256:1e5769f540dbd01cb33bc4716a23e67b9d695161a734aff9c5f925e2bf99a673", size = 658717, upload-time = "2026-02-24T20:02:07.958Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/9c/5b/4be258ff072ed8ee15f6bfd8d5a1a4618aa4704b127c0c5959212ad177d6/openai-2.3.0-py3-none-any.whl", hash = "sha256:a7aa83be6f7b0ab2e4d4d7bcaf36e3d790874c0167380c5d0afd0ed99a86bd7b", size = 999768, upload-time = "2025-10-10T01:12:48.647Z" }, + { url = "https://files.pythonhosted.org/packages/c9/30/844dc675ee6902579b8eef01ed23917cc9319a1c9c0c14ec6e39340c96d0/openai-2.24.0-py3-none-any.whl", hash = "sha256:fed30480d7d6c884303287bde864980a4b137b60553ffbcf9ab4a233b7a73d94", size = 1120122, upload-time = "2026-02-24T20:02:05.669Z" }, ] [[package]] @@ -2154,13 +2397,13 @@ name = "openai-agents" version = "0.1.0" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "griffe" }, - { name = "mcp" }, - { name = "openai" }, - { name = "pydantic" }, - { name = "requests" }, - { name = "types-requests" }, - { name = "typing-extensions" }, + { name = "griffe", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "mcp", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "openai", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "pydantic", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "requests", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "types-requests", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "typing-extensions", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, ] sdist = { url = "https://files.pythonhosted.org/packages/99/f8/a292d8f506997355755d88db619966539ec838ce18f070c5a101e5a430ec/openai_agents-0.1.0.tar.gz", hash = "sha256:a697a4fdd881a7a16db8c0dcafba0f17d9e90b6236a4b79923bd043b6ae86d80", size = 1379588, upload-time = "2025-06-27T20:58:03.186Z" } wheels = [ @@ -2169,43 +2412,31 @@ wheels = [ [[package]] name = "openai-harmony" -version = "0.0.4" +version = "0.0.8" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "pydantic" }, + { name = "pydantic", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/92/94/01509d510bebf6606614e51113e5a415ced15b8f34aa98a8bf2539314650/openai_harmony-0.0.4.tar.gz", hash = "sha256:5c67ac6df349236fb7b64f57c3dbb0273efcdca24314daa108f2a482c427106c", size = 279848, upload-time = "2025-08-09T01:43:24.974Z" } +sdist = { url = "https://files.pythonhosted.org/packages/3e/92/2d038d096f29179c7c9571b431f9e739f87a487121901725e23fe338dd9d/openai_harmony-0.0.8.tar.gz", hash = "sha256:6e43f98e6c242fa2de6f8ea12eab24af63fa2ed3e89c06341fb9d92632c5cbdf", size = 284777, upload-time = "2025-11-05T19:07:06.727Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/a2/3e/6bb75a4d15a6aad0ba1b23193ca0d2c202cc1f3364ba840833374b7c9c1a/openai_harmony-0.0.4-cp38-abi3-macosx_10_12_x86_64.whl", hash = "sha256:3586d90c899cd41f8624e7b82a48c289f6e4be56c66304ecaf3a0ba88963a73f", size = 2772770, upload-time = "2025-08-09T01:43:14.839Z" }, - { url = "https://files.pythonhosted.org/packages/34/41/2f256fba6762d028ed6f935f0015f71d81927a52b9a1c873679a409b72bf/openai_harmony-0.0.4-cp38-abi3-macosx_11_0_arm64.whl", hash = "sha256:ef21a1e2384a65c62d5ec5e1cded9fe026f1d032d5c5d725110d1a8d330d8f54", size = 2633682, upload-time = "2025-08-09T01:43:12.681Z" }, - { url = "https://files.pythonhosted.org/packages/05/88/ade63bd8f36603610040e7cc086bc134d57a99a742e05f7fcddfdf822ee1/openai_harmony-0.0.4-cp38-abi3-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:3cf2344366f10981bbc0f6d9949a0b2bb87151d209ed295943ed6ad8eda37932", size = 2963206, upload-time = "2025-08-09T01:43:02.433Z" }, - { url = "https://files.pythonhosted.org/packages/8e/ef/a65a0ff177fdf67bc0afd18bb9e7ad690d1b553a8eb5ebf27f601b22dbd0/openai_harmony-0.0.4-cp38-abi3-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:2d8d16d84702059833fb03b841b28c25600c54e83cadccef79af44e1c81166b1", size = 2724854, upload-time = "2025-08-09T01:43:04.606Z" }, - { url = "https://files.pythonhosted.org/packages/8a/a1/ebaf0f55601a98609641283884d52dbfe9a1cf34b04f1cf80acb1560ab74/openai_harmony-0.0.4-cp38-abi3-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:97f1fe3909733212cc6b36f0f199b1421a9c57b79ec665f0322bd604cec47340", size = 2984312, upload-time = "2025-08-09T01:43:08.908Z" }, - { url = "https://files.pythonhosted.org/packages/45/24/246f6f470bfbc89a117714b68f27cdaee12b31166237a227cc657780cc1d/openai_harmony-0.0.4-cp38-abi3-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:567cc568b6bf7b4d041b0c9aa7d6b2c9394f8af6065bc87fa6d23f207b5af9a7", size = 3447870, upload-time = "2025-08-09T01:43:06.734Z" }, - { url = "https://files.pythonhosted.org/packages/1f/ec/dcdcace0ffcf3a532cca910e0c351b62d3a7decf0b091ea8cf856d2a67a6/openai_harmony-0.0.4-cp38-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:31e9bcac0902a309e2fc688e52f247eec7fffcd00d17e958b9a83a8fea6519c2", size = 3049306, upload-time = "2025-08-09T01:43:11.019Z" }, - { url = "https://files.pythonhosted.org/packages/ad/39/172f1048d935db1523a82b45fee5231ad6c622645e566706e6bcf3731da8/openai_harmony-0.0.4-cp38-abi3-musllinux_1_2_aarch64.whl", hash = "sha256:96a63199c0d81095b5d5d1ae8ca82b64c1c13d18d4e30323ae9e8ab31bc80a3d", size = 3121347, upload-time = "2025-08-09T01:43:16.705Z" }, - { url = "https://files.pythonhosted.org/packages/6b/36/8ee4ca5d0b25587121fd3621e6a6106fba80218cb6d159e1670aeb2b22ef/openai_harmony-0.0.4-cp38-abi3-musllinux_1_2_armv7l.whl", hash = "sha256:d38f2639f6bf7c3c34a5dfd79e29075811ae2fa9b895a63e76767f74a47a971e", size = 2952326, upload-time = "2025-08-09T01:43:18.841Z" }, - { url = "https://files.pythonhosted.org/packages/ae/a0/ec8906393968679e269e23e957e11ff419978d1d077fb9af9561b161c988/openai_harmony-0.0.4-cp38-abi3-musllinux_1_2_i686.whl", hash = "sha256:038f1d6772d1be5213b36ae76e5d042022395ec35c428a73ccb8b839b2cecf6a", size = 3015832, upload-time = "2025-08-09T01:43:21.076Z" }, - { url = "https://files.pythonhosted.org/packages/a8/bd/aa9e6e5cf140716dbcae17402fac2a81a9ebb3f934059ac0eec61cb447fc/openai_harmony-0.0.4-cp38-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:15e6d53a66502491a3675a536df30e271f976e6c5efe68250a65191efcb85c4f", size = 3221129, upload-time = "2025-08-09T01:43:23.146Z" }, - { url = "https://files.pythonhosted.org/packages/5a/22/2c7e1728689c7fa98a259ca2d14e718ea7af964516a617a9784f0d35d88a/openai_harmony-0.0.4-cp38-abi3-win32.whl", hash = "sha256:b9ee9e9ab6a237cebbe16563c787a6e83f3fcc034075c3d321dab94448426282", size = 2077125, upload-time = "2025-08-09T01:43:28.91Z" }, - { url = "https://files.pythonhosted.org/packages/e7/93/3a08a06ff3bde7f4c264f86d437e6a5c49792a6e362383b3a669f39c9690/openai_harmony-0.0.4-cp38-abi3-win_amd64.whl", hash = "sha256:746f751de5033b3dbcfcd4a726a4c56ce452c593ad3d54472d8597ce8d8b6d44", size = 2444821, upload-time = "2025-08-09T01:43:26.846Z" }, + { url = "https://files.pythonhosted.org/packages/d3/d2/ce6953ca87db9cae3e775024184da7d1c5cb88cead19a2d75b42f00a959c/openai_harmony-0.0.8-cp38-abi3-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:e4f709815924ec325b9a890e6ab2bbb0ceec8e319a4e257328eb752cf36b2efc", size = 2948463, upload-time = "2025-11-05T19:06:48.17Z" }, + { url = "https://files.pythonhosted.org/packages/25/3f/1a192b93bb47c6b44cd98ba8cc1d3d2a9308f1bb700c3017e6352da11bda/openai_harmony-0.0.8-cp38-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:c007d277218a50db8839e599ed78e0fffe5130f614c3f6d93ae257f282071a29", size = 2953260, upload-time = "2025-11-05T19:06:55.406Z" }, + { url = "https://files.pythonhosted.org/packages/5b/f8/93b582cad3531797c3db7c2db5400fd841538ccddfd9f5e3df61be99a630/openai_harmony-0.0.8-cp38-abi3-musllinux_1_2_aarch64.whl", hash = "sha256:8565d4f5a0638da1bffde29832ed63c9e695c558611053add3b2dc0b56c92dbc", size = 3127044, upload-time = "2025-11-05T19:06:59.553Z" }, + { url = "https://files.pythonhosted.org/packages/60/c3/3d1e01e2dba517a91760e4a03e4f20ffc75039a6fe584d0e6f9b5c78fd15/openai_harmony-0.0.8-cp38-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:007b0476a1f331f8130783f901f1da6f5a7057af1a4891f1b6a31dec364189b5", size = 3205080, upload-time = "2025-11-05T19:07:05.078Z" }, ] [[package]] name = "opencv-python-headless" -version = "4.12.0.88" +version = "4.13.0.92" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "numpy" }, + { name = "numpy", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/a4/63/6861102ec149c3cd298f4d1ea7ce9d6adbc7529221606ff1dab991a19adb/opencv-python-headless-4.12.0.88.tar.gz", hash = "sha256:cfdc017ddf2e59b6c2f53bc12d74b6b0be7ded4ec59083ea70763921af2b6c09", size = 95379675, upload-time = "2025-07-07T09:21:06.815Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/f7/7d/414e243c5c8216a5277afd104a319cc1291c5e23f5eeef512db5629ee7f4/opencv_python_headless-4.12.0.88-cp37-abi3-macosx_13_0_arm64.whl", hash = "sha256:1e58d664809b3350c1123484dd441e1667cd7bed3086db1b9ea1b6f6cb20b50e", size = 37877864, upload-time = "2025-07-07T09:14:41.693Z" }, - { url = "https://files.pythonhosted.org/packages/05/14/7e162714beed1cd5e7b5eb66fcbcba2f065c51b1d9da2463024c84d2f7c0/opencv_python_headless-4.12.0.88-cp37-abi3-macosx_13_0_x86_64.whl", hash = "sha256:365bb2e486b50feffc2d07a405b953a8f3e8eaa63865bc650034e5c71e7a5154", size = 57326608, upload-time = "2025-07-07T09:14:51.885Z" }, - { url = "https://files.pythonhosted.org/packages/69/4e/116720df7f1f7f3b59abc608ca30fbec9d2b3ae810afe4e4d26483d9dfa0/opencv_python_headless-4.12.0.88-cp37-abi3-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:aeb4b13ecb8b4a0beb2668ea07928160ea7c2cd2d9b5ef571bbee6bafe9cc8d0", size = 33145800, upload-time = "2025-07-07T09:15:00.367Z" }, - { url = "https://files.pythonhosted.org/packages/89/53/e19c21e0c4eb1275c3e2c97b081103b6dfb3938172264d283a519bf728b9/opencv_python_headless-4.12.0.88-cp37-abi3-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:236c8df54a90f4d02076e6f9c1cc763d794542e886c576a6fee46ec8ff75a7a9", size = 54023419, upload-time = "2025-07-07T09:15:10.164Z" }, - { url = "https://files.pythonhosted.org/packages/bf/9c/a76fd5414de6ec9f21f763a600058a0c3e290053cea87e0275692b1375c0/opencv_python_headless-4.12.0.88-cp37-abi3-win32.whl", hash = "sha256:fde2cf5c51e4def5f2132d78e0c08f9c14783cd67356922182c6845b9af87dbd", size = 30225230, upload-time = "2025-07-07T09:15:17.045Z" }, - { url = "https://files.pythonhosted.org/packages/f2/35/0858e9e71b36948eafbc5e835874b63e515179dc3b742cbe3d76bc683439/opencv_python_headless-4.12.0.88-cp37-abi3-win_amd64.whl", hash = "sha256:86b413bdd6c6bf497832e346cd5371995de148e579b9774f8eba686dee3f5528", size = 38923559, upload-time = "2025-07-07T09:15:25.229Z" }, + { url = "https://files.pythonhosted.org/packages/21/76/9417a6aef9def70e467a5bf560579f816148a4c658b7d525581b356eda9e/opencv_python_headless-4.13.0.92-cp37-abi3-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:5c8cfc8e87ed452b5cecb9419473ee5560a989859fe1d10d1ce11ae87b09a2cb", size = 33703709, upload-time = "2026-02-05T10:24:46.469Z" }, + { url = "https://files.pythonhosted.org/packages/92/ce/bd17ff5772938267fd49716e94ca24f616ff4cb1ff4c6be13085108037be/opencv_python_headless-4.13.0.92-cp37-abi3-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:0525a3d2c0b46c611e2130b5fdebc94cf404845d8fa64d2f3a3b679572a5bd22", size = 56016764, upload-time = "2026-02-05T10:26:48.904Z" }, + { url = "https://files.pythonhosted.org/packages/8f/b4/b7bcbf7c874665825a8c8e1097e93ea25d1f1d210a3e20d4451d01da30aa/opencv_python_headless-4.13.0.92-cp37-abi3-manylinux_2_28_aarch64.whl", hash = "sha256:eb60e36b237b1ebd40a912da5384b348df8ed534f6f644d8e0b4f103e272ba7d", size = 35010236, upload-time = "2026-02-05T10:28:11.031Z" }, + { url = "https://files.pythonhosted.org/packages/4b/33/b5db29a6c00eb8f50708110d8d453747ca125c8b805bc437b289dbdcc057/opencv_python_headless-4.13.0.92-cp37-abi3-manylinux_2_28_x86_64.whl", hash = "sha256:0bd48544f77c68b2941392fcdf9bcd2b9cdf00e98cb8c29b2455d194763cf99e", size = 60391106, upload-time = "2026-02-05T10:30:14.236Z" }, ] [[package]] @@ -2213,22 +2444,95 @@ name = "opentelemetry-api" version = "1.36.0" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "importlib-metadata" }, - { name = "typing-extensions" }, + { name = "importlib-metadata", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "typing-extensions", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, ] sdist = { url = "https://files.pythonhosted.org/packages/27/d2/c782c88b8afbf961d6972428821c302bd1e9e7bc361352172f0ca31296e2/opentelemetry_api-1.36.0.tar.gz", hash = "sha256:9a72572b9c416d004d492cbc6e61962c0501eaf945ece9b5a0f56597d8348aa0", size = 64780, upload-time = "2025-07-29T15:12:06.02Z" } wheels = [ { url = "https://files.pythonhosted.org/packages/bb/ee/6b08dde0a022c463b88f55ae81149584b125a42183407dc1045c486cc870/opentelemetry_api-1.36.0-py3-none-any.whl", hash = "sha256:02f20bcacf666e1333b6b1f04e647dc1d5111f86b8e510238fcc56d7762cda8c", size = 65564, upload-time = "2025-07-29T15:11:47.998Z" }, ] +[[package]] +name = "opentelemetry-exporter-otlp" +version = "1.36.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "opentelemetry-exporter-otlp-proto-grpc", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "opentelemetry-exporter-otlp-proto-http", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/95/7f/d31294ac28d567a14aefd855756bab79fed69c5a75df712f228f10c47e04/opentelemetry_exporter_otlp-1.36.0.tar.gz", hash = "sha256:72f166ea5a8923ac42889337f903e93af57db8893de200369b07401e98e4e06b", size = 6144, upload-time = "2025-07-29T15:12:07.153Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/a0/a2/8966111a285124f3d6156a663ddf2aeddd52843c1a3d6b56cbd9b6c3fd0e/opentelemetry_exporter_otlp-1.36.0-py3-none-any.whl", hash = "sha256:de93b7c45bcc78296998775d52add7c63729e83ef2cd6560730a6b336d7f6494", size = 7018, upload-time = "2025-07-29T15:11:50.498Z" }, +] + +[[package]] +name = "opentelemetry-exporter-otlp-proto-common" +version = "1.36.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "opentelemetry-proto", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/34/da/7747e57eb341c59886052d733072bc878424bf20f1d8cf203d508bbece5b/opentelemetry_exporter_otlp_proto_common-1.36.0.tar.gz", hash = "sha256:6c496ccbcbe26b04653cecadd92f73659b814c6e3579af157d8716e5f9f25cbf", size = 20302, upload-time = "2025-07-29T15:12:07.71Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/d0/ed/22290dca7db78eb32e0101738366b5bbda00d0407f00feffb9bf8c3fdf87/opentelemetry_exporter_otlp_proto_common-1.36.0-py3-none-any.whl", hash = "sha256:0fc002a6ed63eac235ada9aa7056e5492e9a71728214a61745f6ad04b923f840", size = 18349, upload-time = "2025-07-29T15:11:51.327Z" }, +] + +[[package]] +name = "opentelemetry-exporter-otlp-proto-grpc" +version = "1.36.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "googleapis-common-protos", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "grpcio", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "opentelemetry-api", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "opentelemetry-exporter-otlp-proto-common", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "opentelemetry-proto", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "opentelemetry-sdk", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "typing-extensions", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/72/6f/6c1b0bdd0446e5532294d1d41bf11fbaea39c8a2423a4cdfe4fe6b708127/opentelemetry_exporter_otlp_proto_grpc-1.36.0.tar.gz", hash = "sha256:b281afbf7036b325b3588b5b6c8bb175069e3978d1bd24071f4a59d04c1e5bbf", size = 23822, upload-time = "2025-07-29T15:12:08.292Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/0c/67/5f6bd188d66d0fd8e81e681bbf5822e53eb150034e2611dd2b935d3ab61a/opentelemetry_exporter_otlp_proto_grpc-1.36.0-py3-none-any.whl", hash = "sha256:734e841fc6a5d6f30e7be4d8053adb703c70ca80c562ae24e8083a28fadef211", size = 18828, upload-time = "2025-07-29T15:11:52.235Z" }, +] + +[[package]] +name = "opentelemetry-exporter-otlp-proto-http" +version = "1.36.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "googleapis-common-protos", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "opentelemetry-api", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "opentelemetry-exporter-otlp-proto-common", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "opentelemetry-proto", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "opentelemetry-sdk", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "requests", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "typing-extensions", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/25/85/6632e7e5700ba1ce5b8a065315f92c1e6d787ccc4fb2bdab15139eaefc82/opentelemetry_exporter_otlp_proto_http-1.36.0.tar.gz", hash = "sha256:dd3637f72f774b9fc9608ab1ac479f8b44d09b6fb5b2f3df68a24ad1da7d356e", size = 16213, upload-time = "2025-07-29T15:12:08.932Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/7f/41/a680d38b34f8f5ddbd78ed9f0042e1cc712d58ec7531924d71cb1e6c629d/opentelemetry_exporter_otlp_proto_http-1.36.0-py3-none-any.whl", hash = "sha256:3d769f68e2267e7abe4527f70deb6f598f40be3ea34c6adc35789bea94a32902", size = 18752, upload-time = "2025-07-29T15:11:53.164Z" }, +] + +[[package]] +name = "opentelemetry-proto" +version = "1.36.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "protobuf", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/fd/02/f6556142301d136e3b7e95ab8ea6a5d9dc28d879a99f3dd673b5f97dca06/opentelemetry_proto-1.36.0.tar.gz", hash = "sha256:0f10b3c72f74c91e0764a5ec88fd8f1c368ea5d9c64639fb455e2854ef87dd2f", size = 46152, upload-time = "2025-07-29T15:12:15.717Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/b3/57/3361e06136225be8180e879199caea520f38026f8071366241ac458beb8d/opentelemetry_proto-1.36.0-py3-none-any.whl", hash = "sha256:151b3bf73a09f94afc658497cf77d45a565606f62ce0c17acb08cd9937ca206e", size = 72537, upload-time = "2025-07-29T15:12:02.243Z" }, +] + [[package]] name = "opentelemetry-sdk" version = "1.36.0" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "opentelemetry-api" }, - { name = "opentelemetry-semantic-conventions" }, - { name = "typing-extensions" }, + { name = "opentelemetry-api", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "opentelemetry-semantic-conventions", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "typing-extensions", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, ] sdist = { url = "https://files.pythonhosted.org/packages/4c/85/8567a966b85a2d3f971c4d42f781c305b2b91c043724fa08fd37d158e9dc/opentelemetry_sdk-1.36.0.tar.gz", hash = "sha256:19c8c81599f51b71670661ff7495c905d8fdf6976e41622d5245b791b06fa581", size = 162557, upload-time = "2025-07-29T15:12:16.76Z" } wheels = [ @@ -2240,28 +2544,52 @@ name = "opentelemetry-semantic-conventions" version = "0.57b0" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "opentelemetry-api" }, - { name = "typing-extensions" }, + { name = "opentelemetry-api", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "typing-extensions", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, ] sdist = { url = "https://files.pythonhosted.org/packages/7e/31/67dfa252ee88476a29200b0255bda8dfc2cf07b56ad66dc9a6221f7dc787/opentelemetry_semantic_conventions-0.57b0.tar.gz", hash = "sha256:609a4a79c7891b4620d64c7aac6898f872d790d75f22019913a660756f27ff32", size = 124225, upload-time = "2025-07-29T15:12:17.873Z" } wheels = [ { url = "https://files.pythonhosted.org/packages/05/75/7d591371c6c39c73de5ce5da5a2cc7b72d1d1cd3f8f4638f553c01c37b11/opentelemetry_semantic_conventions-0.57b0-py3-none-any.whl", hash = "sha256:757f7e76293294f124c827e514c2a3144f191ef175b069ce8d1211e1e38e9e78", size = 201627, upload-time = "2025-07-29T15:12:04.174Z" }, ] +[[package]] +name = "opentelemetry-semantic-conventions-ai" +version = "0.4.13" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/ba/e6/40b59eda51ac47009fb47afcdf37c6938594a0bd7f3b9fadcbc6058248e3/opentelemetry_semantic_conventions_ai-0.4.13.tar.gz", hash = "sha256:94efa9fb4ffac18c45f54a3a338ffeb7eedb7e1bb4d147786e77202e159f0036", size = 5368, upload-time = "2025-08-22T10:14:17.387Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/35/b5/cf25da2218910f0d6cdf7f876a06bed118c4969eacaf60a887cbaef44f44/opentelemetry_semantic_conventions_ai-0.4.13-py3-none-any.whl", hash = "sha256:883a30a6bb5deaec0d646912b5f9f6dcbb9f6f72557b73d0f2560bf25d13e2d5", size = 6080, upload-time = "2025-08-22T10:14:16.477Z" }, +] + +[[package]] +name = "orjson" +version = "3.11.7" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/53/45/b268004f745ede84e5798b48ee12b05129d19235d0e15267aa57dcdb400b/orjson-3.11.7.tar.gz", hash = "sha256:9b1a67243945819ce55d24a30b59d6a168e86220452d2c96f4d1f093e71c0c49", size = 6144992, upload-time = "2026-02-02T15:38:49.29Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/2c/83/10e48852865e5dd151bdfe652c06f7da484578ed02c5fca938e3632cb0b8/orjson-3.11.7-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:5ede977b5fe5ac91b1dffc0a517ca4542d2ec8a6a4ff7b2652d94f640796342a", size = 128154, upload-time = "2026-02-02T15:37:48.954Z" }, + { url = "https://files.pythonhosted.org/packages/dc/91/98a52415059db3f374757d0b7f0f16e3b5cd5976c90d1c2b56acaea039e6/orjson-3.11.7-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:e7745312efa9e11c17fbd3cb3097262d079da26930ae9ae7ba28fb738367cbad", size = 133440, upload-time = "2026-02-02T15:37:55.615Z" }, + { url = "https://files.pythonhosted.org/packages/dc/b6/cb540117bda61791f46381f8c26c8f93e802892830a6055748d3bb1925ab/orjson-3.11.7-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:f904c24bdeabd4298f7a977ef14ca2a022ca921ed670b92ecd16ab6f3d01f867", size = 138386, upload-time = "2026-02-02T15:37:56.814Z" }, + { url = "https://files.pythonhosted.org/packages/0f/fe/d605d700c35dd55f51710d159fc54516a280923cd1b7e47508982fbb387d/orjson-3.11.7-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:4682d1db3bcebd2b64757e0ddf9e87ae5f00d29d16c5cdf3a62f561d08cc3dd2", size = 134818, upload-time = "2026-02-02T15:38:01.507Z" }, +] + [[package]] name = "outlines-core" -version = "0.2.11" +version = "0.2.14" source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/1a/d3/e04e9145f8f806723dec9b9e5227ad695a3efcd3ced7794cf7c22b15df5e/outlines_core-0.2.11.tar.gz", hash = "sha256:dfce56f717ff5083e54cbcfdb66cad243365437fccbb5509adaa7e31e030f1d8", size = 197263, upload-time = "2025-05-19T10:12:51.719Z" } +sdist = { url = "https://files.pythonhosted.org/packages/6a/04/4a0812eb27c086cfd2e66e7ec9150f33e105912a9b7f8b335e3479f03a06/outlines_core-0.2.14.tar.gz", hash = "sha256:64808deed1591ca3029ff64346ceb974cd5d780c916ea82504951fe83523039e", size = 191539, upload-time = "2026-01-09T15:59:10.016Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/5f/2c/c7636823244c70e2960060bf9bd978248dffb55c5e7c91c46d18354b2a24/outlines_core-0.2.11-cp312-cp312-macosx_14_0_arm64.whl", hash = "sha256:4a9db4872bae083631d720994f4cee603bce0536b33d5a988814576863b657cf", size = 1957668, upload-time = "2025-05-19T10:12:18.29Z" }, - { url = "https://files.pythonhosted.org/packages/c7/09/5c62047da139d722317a444a4d01cd5f11943a8c2eaecce784341dd0844a/outlines_core-0.2.11-cp312-cp312-macosx_14_0_x86_64.whl", hash = "sha256:8359a45c59f6a8f2eb717245806501a59044c75f6ea8bd08faaa131cc8cdec45", size = 2130493, upload-time = "2025-05-19T10:12:19.537Z" }, - { url = "https://files.pythonhosted.org/packages/89/7a/d6a2810f90e37d550168e0c0a9a915086ea721444727e3ca2c630898d1ef/outlines_core-0.2.11-cp312-cp312-macosx_15_0_arm64.whl", hash = "sha256:5d26a46591377340e0b870b8a96ea8341058341a62ee0bded9098e0c88dd24f4", size = 1956804, upload-time = "2025-05-19T10:12:20.755Z" }, - { url = "https://files.pythonhosted.org/packages/ca/ea/339e6c273b5581128c3b7ca27d428d8993c3085912af1a467aa32ef0e9d1/outlines_core-0.2.11-cp312-cp312-macosx_15_0_x86_64.whl", hash = "sha256:ae460a34675fb11d92a5c605a480fbae4cd6c1b2d11b3698da64a7fcaba64dcf", size = 2127085, upload-time = "2025-05-19T10:12:22.02Z" }, - { url = "https://files.pythonhosted.org/packages/92/c7/a65d1fddf49830ebc41422294eacde35286d9f68994a8aa905cb14f5aade/outlines_core-0.2.11-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:86df9740368866295077346440d911df4972da2b3f1f54b8125e6f329e8a8891", size = 2287677, upload-time = "2025-05-19T10:12:24.24Z" }, - { url = "https://files.pythonhosted.org/packages/23/79/8795aed8be9b77dd69d78e7cfbfcf28c179e6b08da6e56bbbf48a09fe55f/outlines_core-0.2.11-cp312-cp312-manylinux_2_28_aarch64.whl", hash = "sha256:96ce4dd78f106799be4a0a5795cefd1352806162973756a4b6fce4bb6eddd7e4", size = 2113000, upload-time = "2025-05-19T10:12:25.446Z" }, - { url = "https://files.pythonhosted.org/packages/59/e3/cbe9294b06d92ee1892dbb6f2125d833d68e8629d45d080d6daba54eec2d/outlines_core-0.2.11-cp312-cp312-win32.whl", hash = "sha256:358db161cce3650ba822e118dcf0a1efa571c7deb4864ab9d64ca2c9cca7425d", size = 1765703, upload-time = "2025-05-19T10:12:26.693Z" }, - { url = "https://files.pythonhosted.org/packages/1d/c9/ed3cf362515fac16e313368b9b2f2497051f4ded88679205830b6f889f54/outlines_core-0.2.11-cp312-cp312-win_amd64.whl", hash = "sha256:231f9d20d2630c70665345821780d7808b29539620a75c99f65113b518c51032", size = 2060945, upload-time = "2025-05-19T10:12:28.294Z" }, + { url = "https://files.pythonhosted.org/packages/29/29/3a04944407207a5d214879ca5ca33c2bd3e65199a4e927051c1bdaaa4d50/outlines_core-0.2.14-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:3bb2060c240c4507f334965a8948dbeeb22007560d797f6debd92346c0b620cb", size = 2341426, upload-time = "2026-01-09T15:58:33.553Z" }, + { url = "https://files.pythonhosted.org/packages/b2/a7/a77f746272504bac3f628047d56ea1731b61549a3e1d9bbfd226f2968246/outlines_core-0.2.14-cp312-cp312-manylinux_2_28_aarch64.whl", hash = "sha256:1de34681c7e0e7e1551fc9036e4fa3c57986336c905a10536591ceb6d869c258", size = 2236941, upload-time = "2026-01-09T15:58:35.118Z" }, +] + +[[package]] +name = "overrides" +version = "7.7.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/36/86/b585f53236dec60aba864e050778b25045f857e17f6e5ea0ae95fe80edd2/overrides-7.7.0.tar.gz", hash = "sha256:55158fa3d93b98cc75299b1e67078ad9003ca27945c76162c1c0766d6f91820a", size = 22812, upload-time = "2024-01-27T21:01:33.423Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/2c/ab/fc8290c6a4c722e5514d80f62b2dc4c4df1a68a41d1364e625c35990fcf3/overrides-7.7.0-py3-none-any.whl", hash = "sha256:c7ed9d062f78b8e4c1a7b70bd8796b35ead4d9f510227ef9c5dc7626c60d7e49", size = 17832, upload-time = "2024-01-27T21:01:31.393Z" }, ] [[package]] @@ -2278,20 +2606,17 @@ name = "pandas" version = "2.3.2" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "numpy" }, - { name = "python-dateutil" }, - { name = "pytz" }, - { name = "tzdata" }, + { name = "numpy", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "python-dateutil", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "pytz", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "tzdata", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, ] sdist = { url = "https://files.pythonhosted.org/packages/79/8e/0e90233ac205ad182bd6b422532695d2b9414944a280488105d598c70023/pandas-2.3.2.tar.gz", hash = "sha256:ab7b58f8f82706890924ccdfb5f48002b83d2b5a3845976a9fb705d36c34dcdb", size = 4488684, upload-time = "2025-08-21T10:28:29.257Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/ec/db/614c20fb7a85a14828edd23f1c02db58a30abf3ce76f38806155d160313c/pandas-2.3.2-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:3fbb977f802156e7a3f829e9d1d5398f6192375a3e2d1a9ee0803e35fe70a2b9", size = 11587652, upload-time = "2025-08-21T10:27:15.888Z" }, - { url = "https://files.pythonhosted.org/packages/99/b0/756e52f6582cade5e746f19bad0517ff27ba9c73404607c0306585c201b3/pandas-2.3.2-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:1b9b52693123dd234b7c985c68b709b0b009f4521000d0525f2b95c22f15944b", size = 10717686, upload-time = "2025-08-21T10:27:18.486Z" }, { url = "https://files.pythonhosted.org/packages/37/4c/dd5ccc1e357abfeee8353123282de17997f90ff67855f86154e5a13b81e5/pandas-2.3.2-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:0bd281310d4f412733f319a5bc552f86d62cddc5f51d2e392c8787335c994175", size = 11278722, upload-time = "2025-08-21T10:27:21.149Z" }, { url = "https://files.pythonhosted.org/packages/d3/a4/f7edcfa47e0a88cda0be8b068a5bae710bf264f867edfdf7b71584ace362/pandas-2.3.2-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:96d31a6b4354e3b9b8a2c848af75d31da390657e3ac6f30c05c82068b9ed79b9", size = 11987803, upload-time = "2025-08-21T10:27:23.767Z" }, { url = "https://files.pythonhosted.org/packages/f6/61/1bce4129f93ab66f1c68b7ed1c12bac6a70b1b56c5dab359c6bbcd480b52/pandas-2.3.2-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:df4df0b9d02bb873a106971bb85d448378ef14b86ba96f035f50bbd3688456b4", size = 12766345, upload-time = "2025-08-21T10:27:26.6Z" }, { url = "https://files.pythonhosted.org/packages/8e/46/80d53de70fee835531da3a1dae827a1e76e77a43ad22a8cd0f8142b61587/pandas-2.3.2-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:213a5adf93d020b74327cb2c1b842884dbdd37f895f42dcc2f09d451d949f811", size = 13439314, upload-time = "2025-08-21T10:27:29.213Z" }, - { url = "https://files.pythonhosted.org/packages/28/30/8114832daff7489f179971dbc1d854109b7f4365a546e3ea75b6516cea95/pandas-2.3.2-cp312-cp312-win_amd64.whl", hash = "sha256:8c13b81a9347eb8c7548f53fd9a4f08d4dfe996836543f805c987bafa03317ae", size = 10983326, upload-time = "2025-08-21T10:27:31.901Z" }, ] [[package]] @@ -2307,9 +2632,22 @@ wheels = [ name = "partial-json-parser" version = "0.2.1.1.post6" source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/86/13/459e86c9c67a006651803a3df3d0b08f7708bc5483fdc482582d75562949/partial_json_parser-0.2.1.1.post6.tar.gz", hash = "sha256:43896b68929678224cbbe4884a6a5fe9251ded4b30b8b7d7eb569e5feea93afc", size = 10299, upload-time = "2025-06-23T17:51:45.372Z" } +sdist = { url = "https://files.pythonhosted.org/packages/86/13/459e86c9c67a006651803a3df3d0b08f7708bc5483fdc482582d75562949/partial_json_parser-0.2.1.1.post6.tar.gz", hash = "sha256:43896b68929678224cbbe4884a6a5fe9251ded4b30b8b7d7eb569e5feea93afc", size = 10299, upload-time = "2025-06-23T17:51:45.372Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/cb/40/1f922794af3dc7503f19319a8804b398a161a2cd54183cff8b12225b8d85/partial_json_parser-0.2.1.1.post6-py3-none-any.whl", hash = "sha256:abc332f09b13ef5233384dbfe7128a0e9ea3fa4b8f8be9b37ac1b433c810e99e", size = 10876, upload-time = "2025-06-23T17:51:44.332Z" }, +] + +[[package]] +name = "pdfminer-six" +version = "20260107" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "charset-normalizer", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "cryptography", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/34/a4/5cec1112009f0439a5ca6afa8ace321f0ab2f48da3255b7a1c8953014670/pdfminer_six-20260107.tar.gz", hash = "sha256:96bfd431e3577a55a0efd25676968ca4ce8fd5b53f14565f85716ff363889602", size = 8512094, upload-time = "2026-01-07T13:29:12.937Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/cb/40/1f922794af3dc7503f19319a8804b398a161a2cd54183cff8b12225b8d85/partial_json_parser-0.2.1.1.post6-py3-none-any.whl", hash = "sha256:abc332f09b13ef5233384dbfe7128a0e9ea3fa4b8f8be9b37ac1b433c810e99e", size = 10876, upload-time = "2025-06-23T17:51:44.332Z" }, + { url = "https://files.pythonhosted.org/packages/20/8b/28c4eaec9d6b036a52cb44720408f26b1a143ca9bce76cc19e8f5de00ab4/pdfminer_six-20260107-py3-none-any.whl", hash = "sha256:366585ba97e80dffa8f00cebe303d2f381884d8637af4ce422f1df3ef38111a9", size = 6592252, upload-time = "2026-01-07T13:29:10.742Z" }, ] [[package]] @@ -2317,7 +2655,7 @@ name = "pexpect" version = "4.9.0" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "ptyprocess" }, + { name = "ptyprocess", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, ] sdist = { url = "https://files.pythonhosted.org/packages/42/92/cc564bf6381ff43ce1f4d06852fc19a2f11d180f23dc32d9588bee2f149d/pexpect-4.9.0.tar.gz", hash = "sha256:ee7d41123f3c9911050ea2c2dac107568dc43b2d3b0c7557a33212c398ead30f", size = 166450, upload-time = "2023-11-25T09:07:26.339Z" } wheels = [ @@ -2330,17 +2668,12 @@ version = "11.3.0" source = { registry = "https://pypi.org/simple" } sdist = { url = "https://files.pythonhosted.org/packages/f3/0d/d0d6dea55cd152ce3d6767bb38a8fc10e33796ba4ba210cbab9354b6d238/pillow-11.3.0.tar.gz", hash = "sha256:3828ee7586cd0b2091b6209e5ad53e20d0649bbe87164a459d0676e035e8f523", size = 47113069, upload-time = "2025-07-01T09:16:30.666Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/40/fe/1bc9b3ee13f68487a99ac9529968035cca2f0a51ec36892060edcc51d06a/pillow-11.3.0-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:fdae223722da47b024b867c1ea0be64e0df702c5e0a60e27daad39bf960dd1e4", size = 5278800, upload-time = "2025-07-01T09:14:17.648Z" }, - { url = "https://files.pythonhosted.org/packages/2c/32/7e2ac19b5713657384cec55f89065fb306b06af008cfd87e572035b27119/pillow-11.3.0-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:921bd305b10e82b4d1f5e802b6850677f965d8394203d182f078873851dada69", size = 4686296, upload-time = "2025-07-01T09:14:19.828Z" }, { url = "https://files.pythonhosted.org/packages/8e/1e/b9e12bbe6e4c2220effebc09ea0923a07a6da1e1f1bfbc8d7d29a01ce32b/pillow-11.3.0-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:eb76541cba2f958032d79d143b98a3a6b3ea87f0959bbe256c0b5e416599fd5d", size = 5871726, upload-time = "2025-07-03T13:10:04.448Z" }, { url = "https://files.pythonhosted.org/packages/8d/33/e9200d2bd7ba00dc3ddb78df1198a6e80d7669cce6c2bdbeb2530a74ec58/pillow-11.3.0-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:67172f2944ebba3d4a7b54f2e95c786a3a50c21b88456329314caaa28cda70f6", size = 7644652, upload-time = "2025-07-03T13:10:10.391Z" }, { url = "https://files.pythonhosted.org/packages/41/f1/6f2427a26fc683e00d985bc391bdd76d8dd4e92fac33d841127eb8fb2313/pillow-11.3.0-cp312-cp312-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:97f07ed9f56a3b9b5f49d3661dc9607484e85c67e27f3e8be2c7d28ca032fec7", size = 5977787, upload-time = "2025-07-01T09:14:21.63Z" }, { url = "https://files.pythonhosted.org/packages/e4/c9/06dd4a38974e24f932ff5f98ea3c546ce3f8c995d3f0985f8e5ba48bba19/pillow-11.3.0-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:676b2815362456b5b3216b4fd5bd89d362100dc6f4945154ff172e206a22c024", size = 6645236, upload-time = "2025-07-01T09:14:23.321Z" }, { url = "https://files.pythonhosted.org/packages/40/e7/848f69fb79843b3d91241bad658e9c14f39a32f71a301bcd1d139416d1be/pillow-11.3.0-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:3e184b2f26ff146363dd07bde8b711833d7b0202e27d13540bfe2e35a323a809", size = 6086950, upload-time = "2025-07-01T09:14:25.237Z" }, { url = "https://files.pythonhosted.org/packages/0b/1a/7cff92e695a2a29ac1958c2a0fe4c0b2393b60aac13b04a4fe2735cad52d/pillow-11.3.0-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:6be31e3fc9a621e071bc17bb7de63b85cbe0bfae91bb0363c893cbe67247780d", size = 6723358, upload-time = "2025-07-01T09:14:27.053Z" }, - { url = "https://files.pythonhosted.org/packages/26/7d/73699ad77895f69edff76b0f332acc3d497f22f5d75e5360f78cbcaff248/pillow-11.3.0-cp312-cp312-win32.whl", hash = "sha256:7b161756381f0918e05e7cb8a371fff367e807770f8fe92ecb20d905d0e1c149", size = 6275079, upload-time = "2025-07-01T09:14:30.104Z" }, - { url = "https://files.pythonhosted.org/packages/8c/ce/e7dfc873bdd9828f3b6e5c2bbb74e47a98ec23cc5c74fc4e54462f0d9204/pillow-11.3.0-cp312-cp312-win_amd64.whl", hash = "sha256:a6444696fce635783440b7f7a9fc24b3ad10a9ea3f0ab66c5905be1c19ccf17d", size = 6986324, upload-time = "2025-07-01T09:14:31.899Z" }, - { url = "https://files.pythonhosted.org/packages/16/8f/b13447d1bf0b1f7467ce7d86f6e6edf66c0ad7cf44cf5c87a37f9bed9936/pillow-11.3.0-cp312-cp312-win_arm64.whl", hash = "sha256:2aceea54f957dd4448264f9bf40875da0415c83eb85f55069d89c0ed436e3542", size = 2423067, upload-time = "2025-07-01T09:14:33.709Z" }, ] [[package]] @@ -2366,11 +2699,11 @@ name = "pre-commit" version = "4.3.0" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "cfgv" }, - { name = "identify" }, - { name = "nodeenv" }, - { name = "pyyaml" }, - { name = "virtualenv" }, + { name = "cfgv", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "identify", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "nodeenv", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "pyyaml", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "virtualenv", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, ] sdist = { url = "https://files.pythonhosted.org/packages/ff/29/7cf5bbc236333876e4b41f56e06857a87937ce4bf91e117a6991a2dbb02a/pre_commit-4.3.0.tar.gz", hash = "sha256:499fe450cc9d42e9d58e606262795ecb64dd05438943c62b66f6a8673da30b16", size = 193792, upload-time = "2025-08-09T18:56:14.651Z" } wheels = [ @@ -2379,23 +2712,25 @@ wheels = [ [[package]] name = "prime" -version = "0.5.0" +version = "0.5.73" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "build" }, - { name = "cryptography" }, - { name = "httpx" }, - { name = "prime-evals" }, - { name = "prime-sandboxes" }, - { name = "pydantic" }, - { name = "rich" }, - { name = "toml" }, - { name = "typer" }, - { name = "verifiers" }, + { name = "build", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "cryptography", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "gitignore-parser", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "httpx", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "prime-evals", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "prime-sandboxes", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "prime-tunnel", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "pydantic", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "rich", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "toml", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "typer", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "verifiers", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/5d/6f/0cd7ee178b600b95b110f88c4c9892dac6668ea5c4eab917bed6bae27c49/prime-0.5.0.tar.gz", hash = "sha256:bb7bade9a0bb36c72a7a6420b5138144c29e8a56c497a9521d133337fbbeca4a", size = 226349, upload-time = "2025-12-02T03:48:24.942Z" } +sdist = { url = "https://files.pythonhosted.org/packages/ef/c9/36ad6f85aebfab3767e2b837a40e75fbfac65d3f2d00aabc5cd139973215/prime-0.5.73.tar.gz", hash = "sha256:2ce07e2a41eaabbb96e3e2cf4d18545413a5fcdd5806489345a7e7c5c41ba3e3", size = 349341, upload-time = "2026-04-25T23:19:40.098Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/7a/a2/e4ac530ad0853e8c190c9ccb17345604c7c6098ed626e8ca8c7c0f9e7d8a/prime-0.5.0-py3-none-any.whl", hash = "sha256:dbf806e65d241446eaf92030d63384bd602b8395fdb7489d5917231f582b416b", size = 78332, upload-time = "2025-12-02T03:48:23.785Z" }, + { url = "https://files.pythonhosted.org/packages/55/8f/2d1cead03c4bc9be60c61e968efef9359725457af266ce3c585b32b8bca9/prime-0.5.73-py3-none-any.whl", hash = "sha256:5267d43d6d3ae02098b3a59919fb9b3cbb83b737d0ef05fa69bed69aaabdb2d0", size = 178264, upload-time = "2026-04-25T23:19:38.721Z" }, ] [[package]] @@ -2403,8 +2738,8 @@ name = "prime-evals" version = "0.1.5" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "httpx" }, - { name = "pydantic" }, + { name = "httpx", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "pydantic", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, ] sdist = { url = "https://files.pythonhosted.org/packages/20/71/d91715f6ca6253c670ddfb6cda90e9fd47c2788f4db97fe835f395ac19be/prime_evals-0.1.5.tar.gz", hash = "sha256:d3822c4ca40cd7c1721d105bbef10df9f65cbb8aa928ac6de646243941bf5fc5", size = 11216, upload-time = "2025-12-02T03:48:18.425Z" } wheels = [ @@ -2413,109 +2748,172 @@ wheels = [ [[package]] name = "prime-rl" -version = "0.1.0" +version = "0.4.0" source = { editable = "." } dependencies = [ - { name = "accelerate" }, - { name = "beartype" }, - { name = "blobfile" }, - { name = "cydifflib" }, - { name = "datasets" }, - { name = "dion" }, - { name = "jaxtyping" }, - { name = "liger-kernel" }, - { name = "loguru" }, - { name = "lovely-tensors" }, - { name = "math-verify" }, - { name = "nltk" }, - { name = "numpy" }, - { name = "openai" }, - { name = "prime" }, - { name = "prime-evals" }, - { name = "pydantic" }, - { name = "pydantic-settings" }, - { name = "pylatexenc" }, - { name = "pyzmq" }, - { name = "reverse-text" }, - { name = "rich" }, - { name = "ring-flash-attn" }, - { name = "tenacity" }, - { name = "textarena" }, - { name = "tomli" }, - { name = "tomli-w" }, - { name = "torch" }, - { name = "torchdata" }, - { name = "torchtitan" }, - { name = "transformers" }, - { name = "uvloop" }, - { name = "verifiers" }, - { name = "vllm" }, - { name = "wandb" }, + { name = "aiolimiter", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "beartype", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "datasets", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "dion", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "flash-linear-attention", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "jaxtyping", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "liger-kernel", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "loguru", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "numpy", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "nvidia-ml-py", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "openai", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "prime", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "prime-rl-configs", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "pyarrow", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "pyzmq", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "renderers", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "rich", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "ring-flash-attn", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "setproctitle", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "tenacity", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "tilelang", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "torch", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "torchaudio", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "torchdata", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "torchtitan", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "torchvision", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "transformers", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "uvloop", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "verifiers", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "vllm", version = "0.20.2+cu129", source = { url = "https://github.com/vllm-project/vllm/releases/download/v0.20.2/vllm-0.20.2+cu129-cp38-abi3-manylinux_2_31_aarch64.whl" }, marker = "platform_machine == 'aarch64' and sys_platform == 'linux'" }, + { name = "vllm", version = "0.20.2+cu129", source = { url = "https://github.com/vllm-project/vllm/releases/download/v0.20.2/vllm-0.20.2+cu129-cp38-abi3-manylinux_2_31_x86_64.whl" }, marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "wandb", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, ] [package.optional-dependencies] +all = [ + { name = "deep-ep", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "deep-gemm", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "flash-attn", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "flash-attn-3", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "flash-attn-4", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "nixl", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "nixl-cu12", version = "0.10.1", source = { url = "https://github.com/PrimeIntellect-ai/prime-rl/releases/download/v0.5.0/nixl_cu12-0.10.1-cp312-cp312-linux_x86_64.whl" }, marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "quack-kernels", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "vllm-router", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, +] +disagg = [ + { name = "deep-ep", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "deep-gemm", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "nixl", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "nixl-cu12", version = "0.10.1", source = { url = "https://github.com/PrimeIntellect-ai/prime-rl/releases/download/v0.5.0/nixl_cu12-0.10.1-cp312-cp312-linux_x86_64.whl" }, marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "vllm-router", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, +] +envs = [ + { name = "aime2024", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "aime2025", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "alphabet-sort", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "code-env", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "color-codeword", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "deepdive", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "logic-env", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "math-env", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "math-python", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "math500", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "mini-swe-agent-plus", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "reverse-text", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "science-env", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "wiki-search", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, +] flash-attn = [ - { name = "flash-attn" }, + { name = "flash-attn", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, ] flash-attn-3 = [ - { name = "flash-attn-3" }, + { name = "flash-attn-3", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, ] -metrics = [ - { name = "prometheus-client" }, +flash-attn-cute = [ + { name = "flash-attn-4", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, +] +gpt-oss = [ + { name = "kernels", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, +] +quack = [ + { name = "quack-kernels", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, ] [package.dev-dependencies] dev = [ - { name = "ipykernel" }, - { name = "ipywidgets" }, - { name = "pre-commit" }, - { name = "pytest" }, - { name = "ruff" }, + { name = "ipykernel", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "ipywidgets", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "pre-commit", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "pytest", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "ruff", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, +] +mamba-ssm = [ + { name = "mamba-ssm", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, ] [package.metadata] requires-dist = [ - { name = "accelerate", specifier = ">=1.10.1" }, + { name = "aime2024", marker = "extra == 'envs'", index = "https://hub.primeintellect.ai/primeintellect/simple/" }, + { name = "aime2025", marker = "extra == 'envs'", index = "https://hub.primeintellect.ai/primeintellect/simple/" }, + { name = "aiolimiter", specifier = ">=1.2.1" }, + { name = "alphabet-sort", marker = "extra == 'envs'", index = "https://hub.primeintellect.ai/primeintellect/simple/" }, { name = "beartype", specifier = ">=0.21.0" }, - { name = "blobfile", specifier = ">=3.0.0" }, - { name = "cydifflib", specifier = ">=1.2.0" }, + { name = "code-env", marker = "extra == 'envs'", index = "https://hub.primeintellect.ai/primeintellect/simple/" }, + { name = "color-codeword", marker = "extra == 'envs'", index = "https://hub.primeintellect.ai/primeintellect/simple/" }, { name = "datasets", specifier = ">=4.0.0" }, - { name = "dion", git = "https://github.com/samsja/dion.git?rev=d891eebbd950a6ff11d9b5f806282e33df83df9d" }, - { name = "flash-attn", marker = "extra == 'flash-attn'", url = "https://github.com/mjun0812/flash-attention-prebuild-wheels/releases/download/v0.6.8/flash_attn-2.8.3+cu128torch2.9-cp312-cp312-linux_x86_64.whl" }, - { name = "flash-attn-3", marker = "extra == 'flash-attn-3'", url = "https://github.com/samsja/flash-attn-builds/releases/download/v0.1/flash_attn_3-3.0.0b1-cp39-abi3-linux_x86_64.whl" }, + { name = "deep-ep", marker = "platform_machine == 'x86_64' and extra == 'disagg'", url = "https://github.com/PrimeIntellect-ai/prime-rl/releases/download/v0.5.0/deep_ep-1.2.1+29d31c0-cp312-cp312-linux_x86_64.whl" }, + { name = "deep-gemm", marker = "platform_machine == 'x86_64' and extra == 'disagg'", url = "https://github.com/PrimeIntellect-ai/prime-rl/releases/download/v0.5.0/deep_gemm-2.5.0+891d57b-cp312-cp312-linux_x86_64.whl" }, + { name = "deepdive", marker = "extra == 'envs'", index = "https://hub.primeintellect.ai/primeintellect/simple/" }, + { name = "dion", git = "https://github.com/samsja/dion.git?rev=d891eeb" }, + { name = "flash-attn", marker = "platform_machine == 'x86_64' and extra == 'flash-attn'", url = "https://github.com/mjun0812/flash-attention-prebuild-wheels/releases/download/v0.9.4/flash_attn-2.8.3+cu128torch2.11-cp312-cp312-linux_x86_64.whl" }, + { name = "flash-attn-3", marker = "extra == 'flash-attn-3'", index = "https://download.pytorch.org/whl/test/cu128" }, + { name = "flash-attn-4", marker = "extra == 'flash-attn-cute'", git = "https://github.com/Dao-AILab/flash-attention.git?subdirectory=flash_attn%2Fcute&rev=96bd151" }, + { name = "flash-linear-attention", git = "https://github.com/fla-org/flash-linear-attention" }, { name = "jaxtyping", specifier = ">=0.3.2" }, + { name = "kernels", marker = "extra == 'gpt-oss'" }, { name = "liger-kernel", specifier = ">=0.5.10" }, + { name = "logic-env", marker = "extra == 'envs'", index = "https://hub.primeintellect.ai/primeintellect/simple/" }, { name = "loguru", specifier = ">=0.7.3" }, - { name = "lovely-tensors", specifier = ">=0.1.18" }, - { name = "math-verify", specifier = ">=0.8.0" }, - { name = "nltk", specifier = ">=3.9.1" }, + { name = "math-env", marker = "extra == 'envs'", index = "https://hub.primeintellect.ai/primeintellect/simple/" }, + { name = "math-python", marker = "extra == 'envs'", index = "https://hub.primeintellect.ai/primeintellect/simple/" }, + { name = "math500", marker = "extra == 'envs'", index = "https://hub.primeintellect.ai/primeintellect/simple/" }, + { name = "mini-swe-agent-plus", marker = "extra == 'envs'", index = "https://hub.primeintellect.ai/primeintellect/simple/" }, + { name = "nixl", marker = "extra == 'disagg'" }, + { name = "nixl-cu12", marker = "platform_machine == 'x86_64' and extra == 'disagg'", url = "https://github.com/PrimeIntellect-ai/prime-rl/releases/download/v0.5.0/nixl_cu12-0.10.1-cp312-cp312-linux_x86_64.whl" }, { name = "numpy", specifier = ">=2.2.6" }, + { name = "nvidia-ml-py", specifier = ">=12.575.51" }, { name = "openai", specifier = ">=1.106.1" }, - { name = "prime", specifier = ">=0.5.0" }, - { name = "prime-evals", specifier = ">=0.1.5" }, - { name = "prometheus-client", marker = "extra == 'metrics'", specifier = ">=0.20.0" }, - { name = "pydantic", specifier = ">=1.10.13" }, - { name = "pydantic-settings", specifier = ">=2.10.1" }, - { name = "pylatexenc", specifier = ">=2.10" }, + { name = "prime", specifier = ">=0.5.73" }, + { name = "prime-rl", extras = ["disagg"], marker = "extra == 'all'" }, + { name = "prime-rl", extras = ["flash-attn"], marker = "extra == 'all'" }, + { name = "prime-rl", extras = ["flash-attn-3"], marker = "extra == 'all'" }, + { name = "prime-rl", extras = ["flash-attn-cute"], marker = "extra == 'all'" }, + { name = "prime-rl", extras = ["quack"], marker = "extra == 'all'" }, + { name = "prime-rl-configs", editable = "packages/prime-rl-configs" }, + { name = "pyarrow", specifier = ">=21.0.0" }, { name = "pyzmq", specifier = ">=27.1.0" }, - { name = "reverse-text", index = "https://hub.primeintellect.ai/primeintellect/simple/" }, - { name = "reverse-text", specifier = ">=0.1.4", index = "https://hub.primeintellect.ai/primeintellect/simple/" }, + { name = "quack-kernels", marker = "extra == 'quack'", specifier = ">=0.3.3" }, + { name = "renderers", specifier = "==0.1.6" }, + { name = "reverse-text", marker = "extra == 'envs'", index = "https://hub.primeintellect.ai/primeintellect/simple/" }, { name = "rich", specifier = ">=14.0.0" }, { name = "ring-flash-attn", specifier = ">=0.1.8" }, + { name = "science-env", marker = "extra == 'envs'", index = "https://hub.primeintellect.ai/primeintellect/simple/" }, + { name = "setproctitle", specifier = ">=1.3.0" }, { name = "tenacity", specifier = ">=8.2.0" }, - { name = "textarena", specifier = ">=0.6.16" }, - { name = "tomli", specifier = ">=2.2.1" }, - { name = "tomli-w", specifier = ">=1.2.0" }, - { name = "torch", specifier = ">=2.9.0", index = "https://download.pytorch.org/whl/test/cu128" }, + { name = "tilelang", specifier = ">=0.1.8" }, + { name = "torch", specifier = ">=2.9.0", index = "https://download.pytorch.org/whl/cu128" }, + { name = "torchaudio", index = "https://download.pytorch.org/whl/cu128" }, { name = "torchdata", specifier = ">=0.11.0" }, { name = "torchtitan", git = "https://github.com/pytorch/torchtitan?rev=a1fdd7e" }, - { name = "transformers", specifier = ">=4.56.0" }, + { name = "torchvision", index = "https://download.pytorch.org/whl/cu128" }, + { name = "transformers", git = "https://github.com/huggingface/transformers.git?rev=c1c3424" }, { name = "uvloop", specifier = ">=0.21.0" }, - { name = "verifiers", git = "https://github.com/PrimeIntellect-ai/verifiers.git?rev=209774a" }, - { name = "vllm", specifier = "==0.12.0" }, - { name = "wandb", specifier = ">=0.20.1" }, + { name = "verifiers", git = "https://github.com/PrimeIntellect-ai/verifiers.git?rev=aa428f3" }, + { name = "vllm", marker = "platform_machine != 'aarch64' and platform_machine != 'x86_64'", specifier = ">=0.20.2" }, + { name = "vllm", marker = "platform_machine == 'aarch64'", url = "https://github.com/vllm-project/vllm/releases/download/v0.20.2/vllm-0.20.2+cu129-cp38-abi3-manylinux_2_31_aarch64.whl" }, + { name = "vllm", marker = "platform_machine == 'x86_64'", url = "https://github.com/vllm-project/vllm/releases/download/v0.20.2/vllm-0.20.2+cu129-cp38-abi3-manylinux_2_31_x86_64.whl" }, + { name = "vllm-router", marker = "platform_machine == 'x86_64' and extra == 'disagg'", url = "https://github.com/PrimeIntellect-ai/router/releases/download/v0.1.22/vllm_router-0.1.22-cp38-abi3-manylinux_2_28_x86_64.whl" }, + { name = "wandb", specifier = ">=0.26.1" }, + { name = "wiki-search", marker = "extra == 'envs'", index = "https://hub.primeintellect.ai/primeintellect/simple/" }, ] -provides-extras = ["flash-attn", "flash-attn-3", "metrics"] +provides-extras = ["flash-attn", "flash-attn-3", "flash-attn-cute", "envs", "disagg", "gpt-oss", "quack", "all"] [package.metadata.requires-dev] dev = [ @@ -2525,20 +2923,56 @@ dev = [ { name = "pytest", specifier = ">=8.4.1" }, { name = "ruff", specifier = ">=0.12.1" }, ] +mamba-ssm = [{ name = "mamba-ssm", specifier = ">=2.3.0" }] + +[[package]] +name = "prime-rl-configs" +version = "0.4.0" +source = { editable = "packages/prime-rl-configs" } +dependencies = [ + { name = "pydantic", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "pydantic-config", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "tomli", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "tomli-w", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, +] + +[package.metadata] +requires-dist = [ + { name = "pydantic", specifier = ">=1.10.13" }, + { name = "pydantic-config", git = "https://github.com/samsja/pydantic_config.git?branch=main" }, + { name = "tomli", specifier = ">=2.2.1" }, + { name = "tomli-w", specifier = ">=1.2.0" }, +] [[package]] name = "prime-sandboxes" -version = "0.2.9" +version = "0.2.22" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "aiofiles", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "connect-python", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "httpx", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "protobuf", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "pydantic", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "tenacity", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/dc/a7/eb9cb01ce7fa812022ec7c2ffaeb379e1103e07f093a4b85e879b2d0143d/prime_sandboxes-0.2.22.tar.gz", hash = "sha256:7901adca4ff7fff1e7f08c06c4b76765cd90791965216655eb3a87d459588f5e", size = 64970, upload-time = "2026-04-22T18:32:12.973Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/90/b8/46bf98c06826fc26d982c53e10d9d4f086a4c322993a6c51188e5bd6fd30/prime_sandboxes-0.2.22-py3-none-any.whl", hash = "sha256:0a75bb02049c16b55aae7475cc4082e8559ffa6d65181b238077771a95b57855", size = 33182, upload-time = "2026-04-22T18:32:11.739Z" }, +] + +[[package]] +name = "prime-tunnel" +version = "0.1.6" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "aiofiles" }, - { name = "httpx" }, - { name = "pydantic" }, - { name = "tenacity" }, + { name = "httpx", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "pydantic", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "tenacity", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/59/55/e6c855074af9ef027f3601396d4e5e0b464f26f724c074253dd0dd2f47b7/prime_sandboxes-0.2.9.tar.gz", hash = "sha256:af4ad06d205a95f4c54778ecac8492f3f382c850528072df79252e7fd0f3fa08", size = 45538, upload-time = "2026-01-07T20:24:15.352Z" } +sdist = { url = "https://files.pythonhosted.org/packages/d1/e8/fec186cfbcc670deed94a7b5ace67688feed38b114b1c602d35a8f8d7564/prime_tunnel-0.1.6.tar.gz", hash = "sha256:5e371bc3557d8a5257a43ef1aeea1f8918da966a8565dacd33a647b6c82a0eaa", size = 13019, upload-time = "2026-04-13T15:13:17.24Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/88/ae/5d38212ff0cbad2c468e9b5cf1e3235af04d123b55c4adaa8812abdeabbb/prime_sandboxes-0.2.9-py3-none-any.whl", hash = "sha256:dab4dd9c1f4f633386593cfb08f7a3dacb2de4eb6e75f4a2ec7f5f1c6ea2207b", size = 18824, upload-time = "2026-01-07T20:24:14.133Z" }, + { url = "https://files.pythonhosted.org/packages/52/35/adfac88e5880418f095758729b73e395c81be869d0be786fcd4b31938b2f/prime_tunnel-0.1.6-py3-none-any.whl", hash = "sha256:bc096bca509621379f671e2401253ab2ebe6c8c143d797b6ef85d8290c462e0d", size = 14914, upload-time = "2026-04-13T15:13:16.08Z" }, ] [[package]] @@ -2555,8 +2989,8 @@ name = "prometheus-fastapi-instrumentator" version = "7.1.0" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "prometheus-client" }, - { name = "starlette" }, + { name = "prometheus-client", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "starlette", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, ] sdist = { url = "https://files.pythonhosted.org/packages/69/6d/24d53033cf93826aa7857699a4450c1c67e5b9c710e925b1ed2b320c04df/prometheus_fastapi_instrumentator-7.1.0.tar.gz", hash = "sha256:be7cd61eeea4e5912aeccb4261c6631b3f227d8924542d79eaf5af3f439cbe5e", size = 20220, upload-time = "2025-03-19T19:35:05.351Z" } wheels = [ @@ -2568,7 +3002,7 @@ name = "prompt-toolkit" version = "3.0.52" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "wcwidth" }, + { name = "wcwidth", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, ] sdist = { url = "https://files.pythonhosted.org/packages/a1/96/06e01a7b38dce6fe1db213e061a4602dd6032a8a97ef6c1a862537732421/prompt_toolkit-3.0.52.tar.gz", hash = "sha256:28cde192929c8e7321de85de1ddbe736f1375148b02f2e17edd840042b1be855", size = 434198, upload-time = "2025-08-27T15:24:02.057Z" } wheels = [ @@ -2581,37 +3015,22 @@ version = "0.3.2" source = { registry = "https://pypi.org/simple" } sdist = { url = "https://files.pythonhosted.org/packages/a6/16/43264e4a779dd8588c21a70f0709665ee8f611211bdd2c87d952cfa7c776/propcache-0.3.2.tar.gz", hash = "sha256:20d7d62e4e7ef05f221e0db2856b979540686342e7dd9973b815599c7057e168", size = 44139, upload-time = "2025-06-09T22:56:06.081Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/a8/42/9ca01b0a6f48e81615dca4765a8f1dd2c057e0540f6116a27dc5ee01dfb6/propcache-0.3.2-cp312-cp312-macosx_10_13_universal2.whl", hash = "sha256:8de106b6c84506b31c27168582cd3cb3000a6412c16df14a8628e5871ff83c10", size = 73674, upload-time = "2025-06-09T22:54:30.551Z" }, - { url = "https://files.pythonhosted.org/packages/af/6e/21293133beb550f9c901bbece755d582bfaf2176bee4774000bd4dd41884/propcache-0.3.2-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:28710b0d3975117239c76600ea351934ac7b5ff56e60953474342608dbbb6154", size = 43570, upload-time = "2025-06-09T22:54:32.296Z" }, - { url = "https://files.pythonhosted.org/packages/0c/c8/0393a0a3a2b8760eb3bde3c147f62b20044f0ddac81e9d6ed7318ec0d852/propcache-0.3.2-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:ce26862344bdf836650ed2487c3d724b00fbfec4233a1013f597b78c1cb73615", size = 43094, upload-time = "2025-06-09T22:54:33.929Z" }, { url = "https://files.pythonhosted.org/packages/37/2c/489afe311a690399d04a3e03b069225670c1d489eb7b044a566511c1c498/propcache-0.3.2-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:bca54bd347a253af2cf4544bbec232ab982f4868de0dd684246b67a51bc6b1db", size = 226958, upload-time = "2025-06-09T22:54:35.186Z" }, - { url = "https://files.pythonhosted.org/packages/9d/ca/63b520d2f3d418c968bf596839ae26cf7f87bead026b6192d4da6a08c467/propcache-0.3.2-cp312-cp312-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:55780d5e9a2ddc59711d727226bb1ba83a22dd32f64ee15594b9392b1f544eb1", size = 234894, upload-time = "2025-06-09T22:54:36.708Z" }, - { url = "https://files.pythonhosted.org/packages/11/60/1d0ed6fff455a028d678df30cc28dcee7af77fa2b0e6962ce1df95c9a2a9/propcache-0.3.2-cp312-cp312-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:035e631be25d6975ed87ab23153db6a73426a48db688070d925aa27e996fe93c", size = 233672, upload-time = "2025-06-09T22:54:38.062Z" }, { url = "https://files.pythonhosted.org/packages/37/7c/54fd5301ef38505ab235d98827207176a5c9b2aa61939b10a460ca53e123/propcache-0.3.2-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:ee6f22b6eaa39297c751d0e80c0d3a454f112f5c6481214fcf4c092074cecd67", size = 224395, upload-time = "2025-06-09T22:54:39.634Z" }, - { url = "https://files.pythonhosted.org/packages/ee/1a/89a40e0846f5de05fdc6779883bf46ba980e6df4d2ff8fb02643de126592/propcache-0.3.2-cp312-cp312-manylinux_2_5_i686.manylinux1_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:7ca3aee1aa955438c4dba34fc20a9f390e4c79967257d830f137bd5a8a32ed3b", size = 212510, upload-time = "2025-06-09T22:54:41.565Z" }, { url = "https://files.pythonhosted.org/packages/5e/33/ca98368586c9566a6b8d5ef66e30484f8da84c0aac3f2d9aec6d31a11bd5/propcache-0.3.2-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:7a4f30862869fa2b68380d677cc1c5fcf1e0f2b9ea0cf665812895c75d0ca3b8", size = 222949, upload-time = "2025-06-09T22:54:43.038Z" }, - { url = "https://files.pythonhosted.org/packages/ba/11/ace870d0aafe443b33b2f0b7efdb872b7c3abd505bfb4890716ad7865e9d/propcache-0.3.2-cp312-cp312-musllinux_1_2_armv7l.whl", hash = "sha256:b77ec3c257d7816d9f3700013639db7491a434644c906a2578a11daf13176251", size = 217258, upload-time = "2025-06-09T22:54:44.376Z" }, - { url = "https://files.pythonhosted.org/packages/5b/d2/86fd6f7adffcfc74b42c10a6b7db721d1d9ca1055c45d39a1a8f2a740a21/propcache-0.3.2-cp312-cp312-musllinux_1_2_i686.whl", hash = "sha256:cab90ac9d3f14b2d5050928483d3d3b8fb6b4018893fc75710e6aa361ecb2474", size = 213036, upload-time = "2025-06-09T22:54:46.243Z" }, - { url = "https://files.pythonhosted.org/packages/07/94/2d7d1e328f45ff34a0a284cf5a2847013701e24c2a53117e7c280a4316b3/propcache-0.3.2-cp312-cp312-musllinux_1_2_ppc64le.whl", hash = "sha256:0b504d29f3c47cf6b9e936c1852246c83d450e8e063d50562115a6be6d3a2535", size = 227684, upload-time = "2025-06-09T22:54:47.63Z" }, - { url = "https://files.pythonhosted.org/packages/b7/05/37ae63a0087677e90b1d14710e532ff104d44bc1efa3b3970fff99b891dc/propcache-0.3.2-cp312-cp312-musllinux_1_2_s390x.whl", hash = "sha256:ce2ac2675a6aa41ddb2a0c9cbff53780a617ac3d43e620f8fd77ba1c84dcfc06", size = 234562, upload-time = "2025-06-09T22:54:48.982Z" }, { url = "https://files.pythonhosted.org/packages/a4/7c/3f539fcae630408d0bd8bf3208b9a647ccad10976eda62402a80adf8fc34/propcache-0.3.2-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:62b4239611205294cc433845b914131b2a1f03500ff3c1ed093ed216b82621e1", size = 222142, upload-time = "2025-06-09T22:54:50.424Z" }, - { url = "https://files.pythonhosted.org/packages/7c/d2/34b9eac8c35f79f8a962546b3e97e9d4b990c420ee66ac8255d5d9611648/propcache-0.3.2-cp312-cp312-win32.whl", hash = "sha256:df4a81b9b53449ebc90cc4deefb052c1dd934ba85012aa912c7ea7b7e38b60c1", size = 37711, upload-time = "2025-06-09T22:54:52.072Z" }, - { url = "https://files.pythonhosted.org/packages/19/61/d582be5d226cf79071681d1b46b848d6cb03d7b70af7063e33a2787eaa03/propcache-0.3.2-cp312-cp312-win_amd64.whl", hash = "sha256:7046e79b989d7fe457bb755844019e10f693752d169076138abf17f31380800c", size = 41479, upload-time = "2025-06-09T22:54:53.234Z" }, { url = "https://files.pythonhosted.org/packages/cc/35/cc0aaecf278bb4575b8555f2b137de5ab821595ddae9da9d3cd1da4072c7/propcache-0.3.2-py3-none-any.whl", hash = "sha256:98f1ec44fb675f5052cccc8e609c46ed23a35a1cfd18545ad4e29002d858a43f", size = 12663, upload-time = "2025-06-09T22:56:04.484Z" }, ] [[package]] name = "protobuf" -version = "6.32.0" +version = "6.33.5" source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/c0/df/fb4a8eeea482eca989b51cffd274aac2ee24e825f0bf3cbce5281fa1567b/protobuf-6.32.0.tar.gz", hash = "sha256:a81439049127067fc49ec1d36e25c6ee1d1a2b7be930675f919258d03c04e7d2", size = 440614, upload-time = "2025-08-14T21:21:25.015Z" } +sdist = { url = "https://files.pythonhosted.org/packages/ba/25/7c72c307aafc96fa87062aa6291d9f7c94836e43214d43722e86037aac02/protobuf-6.33.5.tar.gz", hash = "sha256:6ddcac2a081f8b7b9642c09406bc6a4290128fce5f471cddd165960bb9119e5c", size = 444465, upload-time = "2026-01-29T21:51:33.494Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/33/18/df8c87da2e47f4f1dcc5153a81cd6bca4e429803f4069a299e236e4dd510/protobuf-6.32.0-cp310-abi3-win32.whl", hash = "sha256:84f9e3c1ff6fb0308dbacb0950d8aa90694b0d0ee68e75719cb044b7078fe741", size = 424409, upload-time = "2025-08-14T21:21:12.366Z" }, - { url = "https://files.pythonhosted.org/packages/e1/59/0a820b7310f8139bd8d5a9388e6a38e1786d179d6f33998448609296c229/protobuf-6.32.0-cp310-abi3-win_amd64.whl", hash = "sha256:a8bdbb2f009cfc22a36d031f22a625a38b615b5e19e558a7b756b3279723e68e", size = 435735, upload-time = "2025-08-14T21:21:15.046Z" }, - { url = "https://files.pythonhosted.org/packages/cc/5b/0d421533c59c789e9c9894683efac582c06246bf24bb26b753b149bd88e4/protobuf-6.32.0-cp39-abi3-macosx_10_9_universal2.whl", hash = "sha256:d52691e5bee6c860fff9a1c86ad26a13afbeb4b168cd4445c922b7e2cf85aaf0", size = 426449, upload-time = "2025-08-14T21:21:16.687Z" }, - { url = "https://files.pythonhosted.org/packages/ec/7b/607764ebe6c7a23dcee06e054fd1de3d5841b7648a90fd6def9a3bb58c5e/protobuf-6.32.0-cp39-abi3-manylinux2014_aarch64.whl", hash = "sha256:501fe6372fd1c8ea2a30b4d9be8f87955a64d6be9c88a973996cef5ef6f0abf1", size = 322869, upload-time = "2025-08-14T21:21:18.282Z" }, - { url = "https://files.pythonhosted.org/packages/40/01/2e730bd1c25392fc32e3268e02446f0d77cb51a2c3a8486b1798e34d5805/protobuf-6.32.0-cp39-abi3-manylinux2014_x86_64.whl", hash = "sha256:75a2aab2bd1aeb1f5dc7c5f33bcb11d82ea8c055c9becbb41c26a8c43fd7092c", size = 322009, upload-time = "2025-08-14T21:21:19.893Z" }, - { url = "https://files.pythonhosted.org/packages/9c/f2/80ffc4677aac1bc3519b26bc7f7f5de7fce0ee2f7e36e59e27d8beb32dd1/protobuf-6.32.0-py3-none-any.whl", hash = "sha256:ba377e5b67b908c8f3072a57b63e2c6a4cbd18aea4ed98d2584350dbf46f2783", size = 169287, upload-time = "2025-08-14T21:21:23.515Z" }, + { url = "https://files.pythonhosted.org/packages/4e/b1/c79468184310de09d75095ed1314b839eb2f72df71097db9d1404a1b2717/protobuf-6.33.5-cp39-abi3-manylinux2014_aarch64.whl", hash = "sha256:9b71e0281f36f179d00cbcb119cb19dec4d14a81393e5ea220f64b286173e190", size = 324638, upload-time = "2026-01-29T21:51:26.423Z" }, + { url = "https://files.pythonhosted.org/packages/9b/53/a9443aa3ca9ba8724fdfa02dd1887c1bcd8e89556b715cfbacca6b63dbec/protobuf-6.33.5-cp39-abi3-manylinux2014_x86_64.whl", hash = "sha256:cbf16ba3350fb7b889fca858fb215967792dc125b35c7976ca4818bee3521cf0", size = 323465, upload-time = "2026-01-29T21:51:28.925Z" }, + { url = "https://files.pythonhosted.org/packages/57/bf/2086963c69bdac3d7cff1cc7ff79b8ce5ea0bec6797a017e1be338a46248/protobuf-6.33.5-py3-none-any.whl", hash = "sha256:69915a973dd0f60f31a08b8318b73eab2bd6a392c79184b3612226b0a3f8ec02", size = 170687, upload-time = "2026-01-29T21:51:32.557Z" }, ] [[package]] @@ -2620,13 +3039,8 @@ version = "7.0.0" source = { registry = "https://pypi.org/simple" } sdist = { url = "https://files.pythonhosted.org/packages/2a/80/336820c1ad9286a4ded7e845b2eccfcb27851ab8ac6abece774a6ff4d3de/psutil-7.0.0.tar.gz", hash = "sha256:7be9c3eba38beccb6495ea33afd982a44074b78f28c434a1f51cc07fd315c456", size = 497003, upload-time = "2025-02-13T21:54:07.946Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/ed/e6/2d26234410f8b8abdbf891c9da62bee396583f713fb9f3325a4760875d22/psutil-7.0.0-cp36-abi3-macosx_10_9_x86_64.whl", hash = "sha256:101d71dc322e3cffd7cea0650b09b3d08b8e7c4109dd6809fe452dfd00e58b25", size = 238051, upload-time = "2025-02-13T21:54:12.36Z" }, - { url = "https://files.pythonhosted.org/packages/04/8b/30f930733afe425e3cbfc0e1468a30a18942350c1a8816acfade80c005c4/psutil-7.0.0-cp36-abi3-macosx_11_0_arm64.whl", hash = "sha256:39db632f6bb862eeccf56660871433e111b6ea58f2caea825571951d4b6aa3da", size = 239535, upload-time = "2025-02-13T21:54:16.07Z" }, - { url = "https://files.pythonhosted.org/packages/2a/ed/d362e84620dd22876b55389248e522338ed1bf134a5edd3b8231d7207f6d/psutil-7.0.0-cp36-abi3-manylinux_2_12_i686.manylinux2010_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:1fcee592b4c6f146991ca55919ea3d1f8926497a713ed7faaf8225e174581e91", size = 275004, upload-time = "2025-02-13T21:54:18.662Z" }, { url = "https://files.pythonhosted.org/packages/bf/b9/b0eb3f3cbcb734d930fdf839431606844a825b23eaf9a6ab371edac8162c/psutil-7.0.0-cp36-abi3-manylinux_2_12_x86_64.manylinux2010_x86_64.manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:4b1388a4f6875d7e2aff5c4ca1cc16c545ed41dd8bb596cefea80111db353a34", size = 277986, upload-time = "2025-02-13T21:54:21.811Z" }, { url = "https://files.pythonhosted.org/packages/eb/a2/709e0fe2f093556c17fbafda93ac032257242cabcc7ff3369e2cb76a97aa/psutil-7.0.0-cp36-abi3-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:a5f098451abc2828f7dc6b58d44b532b22f2088f4999a937557b603ce72b1993", size = 279544, upload-time = "2025-02-13T21:54:24.68Z" }, - { url = "https://files.pythonhosted.org/packages/50/e6/eecf58810b9d12e6427369784efe814a1eec0f492084ce8eb8f4d89d6d61/psutil-7.0.0-cp37-abi3-win32.whl", hash = "sha256:ba3fcef7523064a6c9da440fc4d6bd07da93ac726b5733c29027d7dc95b39d99", size = 241053, upload-time = "2025-02-13T21:54:34.31Z" }, - { url = "https://files.pythonhosted.org/packages/50/1b/6921afe68c74868b4c9fa424dad3be35b095e16687989ebbb50ce4fceb7c/psutil-7.0.0-cp37-abi3-win_amd64.whl", hash = "sha256:4cf3d4eb1aa9b348dec30105c55cd9b7d4629285735a102beb4441e38db90553", size = 244885, upload-time = "2025-02-13T21:54:37.486Z" }, ] [[package]] @@ -2662,13 +3076,10 @@ version = "21.0.0" source = { registry = "https://pypi.org/simple" } sdist = { url = "https://files.pythonhosted.org/packages/ef/c2/ea068b8f00905c06329a3dfcd40d0fcc2b7d0f2e355bdb25b65e0a0e4cd4/pyarrow-21.0.0.tar.gz", hash = "sha256:5051f2dccf0e283ff56335760cbc8622cf52264d67e359d5569541ac11b6d5bc", size = 1133487, upload-time = "2025-07-18T00:57:31.761Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/ca/d4/d4f817b21aacc30195cf6a46ba041dd1be827efa4a623cc8bf39a1c2a0c0/pyarrow-21.0.0-cp312-cp312-macosx_12_0_arm64.whl", hash = "sha256:3a302f0e0963db37e0a24a70c56cf91a4faa0bca51c23812279ca2e23481fccd", size = 31160305, upload-time = "2025-07-18T00:55:35.373Z" }, - { url = "https://files.pythonhosted.org/packages/a2/9c/dcd38ce6e4b4d9a19e1d36914cb8e2b1da4e6003dd075474c4cfcdfe0601/pyarrow-21.0.0-cp312-cp312-macosx_12_0_x86_64.whl", hash = "sha256:b6b27cf01e243871390474a211a7922bfbe3bda21e39bc9160daf0da3fe48876", size = 32684264, upload-time = "2025-07-18T00:55:39.303Z" }, { url = "https://files.pythonhosted.org/packages/4f/74/2a2d9f8d7a59b639523454bec12dba35ae3d0a07d8ab529dc0809f74b23c/pyarrow-21.0.0-cp312-cp312-manylinux_2_28_aarch64.whl", hash = "sha256:e72a8ec6b868e258a2cd2672d91f2860ad532d590ce94cdf7d5e7ec674ccf03d", size = 41108099, upload-time = "2025-07-18T00:55:42.889Z" }, { url = "https://files.pythonhosted.org/packages/ad/90/2660332eeb31303c13b653ea566a9918484b6e4d6b9d2d46879a33ab0622/pyarrow-21.0.0-cp312-cp312-manylinux_2_28_x86_64.whl", hash = "sha256:b7ae0bbdc8c6674259b25bef5d2a1d6af5d39d7200c819cf99e07f7dfef1c51e", size = 42829529, upload-time = "2025-07-18T00:55:47.069Z" }, { url = "https://files.pythonhosted.org/packages/33/27/1a93a25c92717f6aa0fca06eb4700860577d016cd3ae51aad0e0488ac899/pyarrow-21.0.0-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:58c30a1729f82d201627c173d91bd431db88ea74dcaa3885855bc6203e433b82", size = 43367883, upload-time = "2025-07-18T00:55:53.069Z" }, { url = "https://files.pythonhosted.org/packages/05/d9/4d09d919f35d599bc05c6950095e358c3e15148ead26292dfca1fb659b0c/pyarrow-21.0.0-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:072116f65604b822a7f22945a7a6e581cfa28e3454fdcc6939d4ff6090126623", size = 45133802, upload-time = "2025-07-18T00:55:57.714Z" }, - { url = "https://files.pythonhosted.org/packages/71/30/f3795b6e192c3ab881325ffe172e526499eb3780e306a15103a2764916a2/pyarrow-21.0.0-cp312-cp312-win_amd64.whl", hash = "sha256:cf56ec8b0a5c8c9d7021d6fd754e688104f9ebebf1bf4449613c9531f5346a18", size = 26203175, upload-time = "2025-07-18T00:56:01.364Z" }, ] [[package]] @@ -2677,25 +3088,10 @@ version = "1.4.2" source = { registry = "https://pypi.org/simple" } sdist = { url = "https://files.pythonhosted.org/packages/04/14/43297a7b7f0c1bf0c00b596f754ee3ac946128c64d21047ccf9c9bbc5165/pybase64-1.4.2.tar.gz", hash = "sha256:46cdefd283ed9643315d952fe44de80dc9b9a811ce6e3ec97fd1827af97692d0", size = 137246, upload-time = "2025-07-27T13:08:57.808Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/28/6d/11ede991e800797b9f5ebd528013b34eee5652df93de61ffb24503393fa5/pybase64-1.4.2-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:db2c75d1388855b5a1015b65096d7dbcc708e7de3245dcbedeb872ec05a09326", size = 38326, upload-time = "2025-07-27T13:03:09.065Z" }, - { url = "https://files.pythonhosted.org/packages/fe/84/87f1f565f42e2397e2aaa2477c86419f5173c3699881c42325c090982f0a/pybase64-1.4.2-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:6b621a972a01841368fdb9dedc55fd3c6e0c7217d0505ba3b1ebe95e7ef1b493", size = 31661, upload-time = "2025-07-27T13:03:10.295Z" }, - { url = "https://files.pythonhosted.org/packages/cb/2a/a24c810e7a61d2cc6f73fe9ee4872a03030887fa8654150901b15f376f65/pybase64-1.4.2-cp312-cp312-manylinux1_i686.manylinux2014_i686.manylinux_2_17_i686.manylinux_2_5_i686.whl", hash = "sha256:f48c32ac6a16cbf57a5a96a073fef6ff7e3526f623cd49faa112b7f9980bafba", size = 68192, upload-time = "2025-07-27T13:03:11.467Z" }, { url = "https://files.pythonhosted.org/packages/ee/87/d9baf98cbfc37b8657290ad4421f3a3c36aa0eafe4872c5859cfb52f3448/pybase64-1.4.2-cp312-cp312-manylinux1_x86_64.manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:ace8b23093a6bb862477080d9059b784096ab2f97541e8bfc40d42f062875149", size = 71587, upload-time = "2025-07-27T13:03:12.719Z" }, { url = "https://files.pythonhosted.org/packages/0b/89/3df043cc56ef3b91b7aa0c26ae822a2d7ec8da0b0fd7c309c879b0eb5988/pybase64-1.4.2-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:1772c7532a7fb6301baea3dd3e010148dbf70cd1136a83c2f5f91bdc94822145", size = 59910, upload-time = "2025-07-27T13:03:14.266Z" }, - { url = "https://files.pythonhosted.org/packages/75/4f/6641e9edf37aeb4d4524dc7ba2168eff8d96c90e77f6283c2be3400ab380/pybase64-1.4.2-cp312-cp312-manylinux2014_armv7l.manylinux_2_17_armv7l.whl", hash = "sha256:f86f7faddcba5cbfea475f8ab96567834c28bf09ca6c7c3d66ee445adac80d8f", size = 56701, upload-time = "2025-07-27T13:03:15.6Z" }, - { url = "https://files.pythonhosted.org/packages/2d/7f/20d8ac1046f12420a0954a45a13033e75f98aade36eecd00c64e3549b071/pybase64-1.4.2-cp312-cp312-manylinux2014_ppc64le.manylinux_2_17_ppc64le.whl", hash = "sha256:0b8c8e275b5294089f314814b4a50174ab90af79d6a4850f6ae11261ff6a7372", size = 59288, upload-time = "2025-07-27T13:03:16.823Z" }, - { url = "https://files.pythonhosted.org/packages/17/ea/9c0ca570e3e50b3c6c3442e280c83b321a0464c86a9db1f982a4ff531550/pybase64-1.4.2-cp312-cp312-manylinux2014_s390x.manylinux_2_17_s390x.whl", hash = "sha256:864d85a0470c615807ae8b97d724d068b940a2d10ac13a5f1b9e75a3ce441758", size = 60267, upload-time = "2025-07-27T13:03:18.132Z" }, - { url = "https://files.pythonhosted.org/packages/f9/ac/46894929d71ccedebbfb0284173b0fea96bc029cd262654ba8451a7035d6/pybase64-1.4.2-cp312-cp312-manylinux_2_31_riscv64.whl", hash = "sha256:47254d97ed2d8351e30ecfdb9e2414547f66ba73f8a09f932c9378ff75cd10c5", size = 54801, upload-time = "2025-07-27T13:03:19.669Z" }, { url = "https://files.pythonhosted.org/packages/6a/1e/02c95218ea964f0b2469717c2c69b48e63f4ca9f18af01a5b2a29e4c1216/pybase64-1.4.2-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:264b65ecc4f0ee73f3298ab83bbd8008f7f9578361b8df5b448f985d8c63e02a", size = 58599, upload-time = "2025-07-27T13:03:20.951Z" }, - { url = "https://files.pythonhosted.org/packages/15/45/ccc21004930789b8fb439d43e3212a6c260ccddb2bf450c39a20db093f33/pybase64-1.4.2-cp312-cp312-musllinux_1_2_armv7l.whl", hash = "sha256:fbcc2b30cd740c16c9699f596f22c7a9e643591311ae72b1e776f2d539e9dd9d", size = 52388, upload-time = "2025-07-27T13:03:23.064Z" }, - { url = "https://files.pythonhosted.org/packages/c4/45/22e46e549710c4c237d77785b6fb1bc4c44c288a5c44237ba9daf5c34b82/pybase64-1.4.2-cp312-cp312-musllinux_1_2_i686.whl", hash = "sha256:cda9f79c22d51ee4508f5a43b673565f1d26af4330c99f114e37e3186fdd3607", size = 68802, upload-time = "2025-07-27T13:03:24.673Z" }, - { url = "https://files.pythonhosted.org/packages/55/0c/232c6261b81296e5593549b36e6e7884a5da008776d12665923446322c36/pybase64-1.4.2-cp312-cp312-musllinux_1_2_ppc64le.whl", hash = "sha256:0c91c6d2a7232e2a1cd10b3b75a8bb657defacd4295a1e5e80455df2dfc84d4f", size = 57841, upload-time = "2025-07-27T13:03:25.948Z" }, - { url = "https://files.pythonhosted.org/packages/20/8a/b35a615ae6f04550d696bb179c414538b3b477999435fdd4ad75b76139e4/pybase64-1.4.2-cp312-cp312-musllinux_1_2_riscv64.whl", hash = "sha256:a370dea7b1cee2a36a4d5445d4e09cc243816c5bc8def61f602db5a6f5438e52", size = 54320, upload-time = "2025-07-27T13:03:27.495Z" }, - { url = "https://files.pythonhosted.org/packages/d3/a9/8bd4f9bcc53689f1b457ecefed1eaa080e4949d65a62c31a38b7253d5226/pybase64-1.4.2-cp312-cp312-musllinux_1_2_s390x.whl", hash = "sha256:9aa4de83f02e462a6f4e066811c71d6af31b52d7484de635582d0e3ec3d6cc3e", size = 56482, upload-time = "2025-07-27T13:03:28.942Z" }, { url = "https://files.pythonhosted.org/packages/75/e5/4a7735b54a1191f61c3f5c2952212c85c2d6b06eb5fb3671c7603395f70c/pybase64-1.4.2-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:83a1c2f9ed00fee8f064d548c8654a480741131f280e5750bb32475b7ec8ee38", size = 70959, upload-time = "2025-07-27T13:03:30.171Z" }, - { url = "https://files.pythonhosted.org/packages/d3/67/e2b6cb32c782e12304d467418e70da0212567f42bd4d3b5eb1fdf64920ad/pybase64-1.4.2-cp312-cp312-win32.whl", hash = "sha256:a6e5688b18d558e8c6b8701cc8560836c4bbeba61d33c836b4dba56b19423716", size = 33683, upload-time = "2025-07-27T13:03:31.775Z" }, - { url = "https://files.pythonhosted.org/packages/4f/bc/d5c277496063a09707486180f17abbdbdebbf2f5c4441b20b11d3cb7dc7c/pybase64-1.4.2-cp312-cp312-win_amd64.whl", hash = "sha256:c995d21b8bd08aa179cd7dd4db0695c185486ecc72da1e8f6c37ec86cadb8182", size = 35817, upload-time = "2025-07-27T13:03:32.99Z" }, - { url = "https://files.pythonhosted.org/packages/e6/69/e4be18ae685acff0ae77f75d4586590f29d2cd187bf603290cf1d635cad4/pybase64-1.4.2-cp312-cp312-win_arm64.whl", hash = "sha256:e254b9258c40509c2ea063a7784f6994988f3f26099d6e08704e3c15dfed9a55", size = 30900, upload-time = "2025-07-27T13:03:34.499Z" }, ] [[package]] @@ -2716,34 +3112,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/a0/e3/59cd50310fc9b59512193629e1984c1f95e5c8ae6e5d8c69532ccc65a7fe/pycparser-2.23-py3-none-any.whl", hash = "sha256:e5c6e8d3fbad53479cab09ac03729e0a9faf2bee3db8208a550daf5af81a5934", size = 118140, upload-time = "2025-09-09T13:23:46.651Z" }, ] -[[package]] -name = "pycryptodomex" -version = "3.23.0" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/c9/85/e24bf90972a30b0fcd16c73009add1d7d7cd9140c2498a68252028899e41/pycryptodomex-3.23.0.tar.gz", hash = "sha256:71909758f010c82bc99b0abf4ea12012c98962fbf0583c2164f8b84533c2e4da", size = 4922157, upload-time = "2025-05-17T17:23:41.434Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/dd/9c/1a8f35daa39784ed8adf93a694e7e5dc15c23c741bbda06e1d45f8979e9e/pycryptodomex-3.23.0-cp37-abi3-macosx_10_9_universal2.whl", hash = "sha256:06698f957fe1ab229a99ba2defeeae1c09af185baa909a31a5d1f9d42b1aaed6", size = 2499240, upload-time = "2025-05-17T17:22:46.953Z" }, - { url = "https://files.pythonhosted.org/packages/7a/62/f5221a191a97157d240cf6643747558759126c76ee92f29a3f4aee3197a5/pycryptodomex-3.23.0-cp37-abi3-macosx_10_9_x86_64.whl", hash = "sha256:b2c2537863eccef2d41061e82a881dcabb04944c5c06c5aa7110b577cc487545", size = 1644042, upload-time = "2025-05-17T17:22:49.098Z" }, - { url = "https://files.pythonhosted.org/packages/8c/fd/5a054543c8988d4ed7b612721d7e78a4b9bf36bc3c5ad45ef45c22d0060e/pycryptodomex-3.23.0-cp37-abi3-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:43c446e2ba8df8889e0e16f02211c25b4934898384c1ec1ec04d7889c0333587", size = 2186227, upload-time = "2025-05-17T17:22:51.139Z" }, - { url = "https://files.pythonhosted.org/packages/c8/a9/8862616a85cf450d2822dbd4fff1fcaba90877907a6ff5bc2672cafe42f8/pycryptodomex-3.23.0-cp37-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:f489c4765093fb60e2edafdf223397bc716491b2b69fe74367b70d6999257a5c", size = 2272578, upload-time = "2025-05-17T17:22:53.676Z" }, - { url = "https://files.pythonhosted.org/packages/46/9f/bda9c49a7c1842820de674ab36c79f4fbeeee03f8ff0e4f3546c3889076b/pycryptodomex-3.23.0-cp37-abi3-manylinux_2_5_i686.manylinux1_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:bdc69d0d3d989a1029df0eed67cc5e8e5d968f3724f4519bd03e0ec68df7543c", size = 2312166, upload-time = "2025-05-17T17:22:56.585Z" }, - { url = "https://files.pythonhosted.org/packages/03/cc/870b9bf8ca92866ca0186534801cf8d20554ad2a76ca959538041b7a7cf4/pycryptodomex-3.23.0-cp37-abi3-musllinux_1_2_aarch64.whl", hash = "sha256:6bbcb1dd0f646484939e142462d9e532482bc74475cecf9c4903d4e1cd21f003", size = 2185467, upload-time = "2025-05-17T17:22:59.237Z" }, - { url = "https://files.pythonhosted.org/packages/96/e3/ce9348236d8e669fea5dd82a90e86be48b9c341210f44e25443162aba187/pycryptodomex-3.23.0-cp37-abi3-musllinux_1_2_i686.whl", hash = "sha256:8a4fcd42ccb04c31268d1efeecfccfd1249612b4de6374205376b8f280321744", size = 2346104, upload-time = "2025-05-17T17:23:02.112Z" }, - { url = "https://files.pythonhosted.org/packages/a5/e9/e869bcee87beb89040263c416a8a50204f7f7a83ac11897646c9e71e0daf/pycryptodomex-3.23.0-cp37-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:55ccbe27f049743a4caf4f4221b166560d3438d0b1e5ab929e07ae1702a4d6fd", size = 2271038, upload-time = "2025-05-17T17:23:04.872Z" }, - { url = "https://files.pythonhosted.org/packages/8d/67/09ee8500dd22614af5fbaa51a4aee6e342b5fa8aecf0a6cb9cbf52fa6d45/pycryptodomex-3.23.0-cp37-abi3-win32.whl", hash = "sha256:189afbc87f0b9f158386bf051f720e20fa6145975f1e76369303d0f31d1a8d7c", size = 1771969, upload-time = "2025-05-17T17:23:07.115Z" }, - { url = "https://files.pythonhosted.org/packages/69/96/11f36f71a865dd6df03716d33bd07a67e9d20f6b8d39820470b766af323c/pycryptodomex-3.23.0-cp37-abi3-win_amd64.whl", hash = "sha256:52e5ca58c3a0b0bd5e100a9fbc8015059b05cffc6c66ce9d98b4b45e023443b9", size = 1803124, upload-time = "2025-05-17T17:23:09.267Z" }, - { url = "https://files.pythonhosted.org/packages/f9/93/45c1cdcbeb182ccd2e144c693eaa097763b08b38cded279f0053ed53c553/pycryptodomex-3.23.0-cp37-abi3-win_arm64.whl", hash = "sha256:02d87b80778c171445d67e23d1caef279bf4b25c3597050ccd2e13970b57fd51", size = 1707161, upload-time = "2025-05-17T17:23:11.414Z" }, -] - [[package]] name = "pydantic" version = "2.12.5" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "annotated-types" }, - { name = "pydantic-core" }, - { name = "typing-extensions" }, - { name = "typing-inspection" }, + { name = "annotated-types", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "pydantic-core", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "typing-extensions", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "typing-inspection", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, ] sdist = { url = "https://files.pythonhosted.org/packages/69/44/36f1a6e523abc58ae5f928898e4aca2e0ea509b5aa6f6f392a5d882be928/pydantic-2.12.5.tar.gz", hash = "sha256:4d351024c75c0f085a9febbb665ce8c0c6ec5d30e903bdb6394b7ede26aebb49", size = 821591, upload-time = "2025-11-26T15:11:46.471Z" } wheels = [ @@ -2752,7 +3129,16 @@ wheels = [ [package.optional-dependencies] email = [ - { name = "email-validator" }, + { name = "email-validator", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, +] + +[[package]] +name = "pydantic-config" +version = "0.3.0" +source = { git = "https://github.com/samsja/pydantic_config.git?branch=main#1562882463aed6fca1964a1ae0aa4b5141f89d73" } +dependencies = [ + { name = "pydantic", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "tyro", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, ] [[package]] @@ -2760,26 +3146,14 @@ name = "pydantic-core" version = "2.41.5" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "typing-extensions" }, + { name = "typing-extensions", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, ] sdist = { url = "https://files.pythonhosted.org/packages/71/70/23b021c950c2addd24ec408e9ab05d59b035b39d97cdc1130e1bce647bb6/pydantic_core-2.41.5.tar.gz", hash = "sha256:08daa51ea16ad373ffd5e7606252cc32f07bc72b28284b6bc9c6df804816476e", size = 460952, upload-time = "2025-11-04T13:43:49.098Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/5f/5d/5f6c63eebb5afee93bcaae4ce9a898f3373ca23df3ccaef086d0233a35a7/pydantic_core-2.41.5-cp312-cp312-macosx_10_12_x86_64.whl", hash = "sha256:f41a7489d32336dbf2199c8c0a215390a751c5b014c2c1c5366e817202e9cdf7", size = 2110990, upload-time = "2025-11-04T13:39:58.079Z" }, - { url = "https://files.pythonhosted.org/packages/aa/32/9c2e8ccb57c01111e0fd091f236c7b371c1bccea0fa85247ac55b1e2b6b6/pydantic_core-2.41.5-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:070259a8818988b9a84a449a2a7337c7f430a22acc0859c6b110aa7212a6d9c0", size = 1896003, upload-time = "2025-11-04T13:39:59.956Z" }, { url = "https://files.pythonhosted.org/packages/68/b8/a01b53cb0e59139fbc9e4fda3e9724ede8de279097179be4ff31f1abb65a/pydantic_core-2.41.5-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:e96cea19e34778f8d59fe40775a7a574d95816eb150850a85a7a4c8f4b94ac69", size = 1919200, upload-time = "2025-11-04T13:40:02.241Z" }, - { url = "https://files.pythonhosted.org/packages/38/de/8c36b5198a29bdaade07b5985e80a233a5ac27137846f3bc2d3b40a47360/pydantic_core-2.41.5-cp312-cp312-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:ed2e99c456e3fadd05c991f8f437ef902e00eedf34320ba2b0842bd1c3ca3a75", size = 2052578, upload-time = "2025-11-04T13:40:04.401Z" }, - { url = "https://files.pythonhosted.org/packages/00/b5/0e8e4b5b081eac6cb3dbb7e60a65907549a1ce035a724368c330112adfdd/pydantic_core-2.41.5-cp312-cp312-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:65840751b72fbfd82c3c640cff9284545342a4f1eb1586ad0636955b261b0b05", size = 2208504, upload-time = "2025-11-04T13:40:06.072Z" }, - { url = "https://files.pythonhosted.org/packages/77/56/87a61aad59c7c5b9dc8caad5a41a5545cba3810c3e828708b3d7404f6cef/pydantic_core-2.41.5-cp312-cp312-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:e536c98a7626a98feb2d3eaf75944ef6f3dbee447e1f841eae16f2f0a72d8ddc", size = 2335816, upload-time = "2025-11-04T13:40:07.835Z" }, { url = "https://files.pythonhosted.org/packages/0d/76/941cc9f73529988688a665a5c0ecff1112b3d95ab48f81db5f7606f522d3/pydantic_core-2.41.5-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:eceb81a8d74f9267ef4081e246ffd6d129da5d87e37a77c9bde550cb04870c1c", size = 2075366, upload-time = "2025-11-04T13:40:09.804Z" }, - { url = "https://files.pythonhosted.org/packages/d3/43/ebef01f69baa07a482844faaa0a591bad1ef129253ffd0cdaa9d8a7f72d3/pydantic_core-2.41.5-cp312-cp312-manylinux_2_5_i686.manylinux1_i686.whl", hash = "sha256:d38548150c39b74aeeb0ce8ee1d8e82696f4a4e16ddc6de7b1d8823f7de4b9b5", size = 2171698, upload-time = "2025-11-04T13:40:12.004Z" }, { url = "https://files.pythonhosted.org/packages/b1/87/41f3202e4193e3bacfc2c065fab7706ebe81af46a83d3e27605029c1f5a6/pydantic_core-2.41.5-cp312-cp312-musllinux_1_1_aarch64.whl", hash = "sha256:c23e27686783f60290e36827f9c626e63154b82b116d7fe9adba1fda36da706c", size = 2132603, upload-time = "2025-11-04T13:40:13.868Z" }, - { url = "https://files.pythonhosted.org/packages/49/7d/4c00df99cb12070b6bccdef4a195255e6020a550d572768d92cc54dba91a/pydantic_core-2.41.5-cp312-cp312-musllinux_1_1_armv7l.whl", hash = "sha256:482c982f814460eabe1d3bb0adfdc583387bd4691ef00b90575ca0d2b6fe2294", size = 2329591, upload-time = "2025-11-04T13:40:15.672Z" }, { url = "https://files.pythonhosted.org/packages/cc/6a/ebf4b1d65d458f3cda6a7335d141305dfa19bdc61140a884d165a8a1bbc7/pydantic_core-2.41.5-cp312-cp312-musllinux_1_1_x86_64.whl", hash = "sha256:bfea2a5f0b4d8d43adf9d7b8bf019fb46fdd10a2e5cde477fbcb9d1fa08c68e1", size = 2319068, upload-time = "2025-11-04T13:40:17.532Z" }, - { url = "https://files.pythonhosted.org/packages/49/3b/774f2b5cd4192d5ab75870ce4381fd89cf218af999515baf07e7206753f0/pydantic_core-2.41.5-cp312-cp312-win32.whl", hash = "sha256:b74557b16e390ec12dca509bce9264c3bbd128f8a2c376eaa68003d7f327276d", size = 1985908, upload-time = "2025-11-04T13:40:19.309Z" }, - { url = "https://files.pythonhosted.org/packages/86/45/00173a033c801cacf67c190fef088789394feaf88a98a7035b0e40d53dc9/pydantic_core-2.41.5-cp312-cp312-win_amd64.whl", hash = "sha256:1962293292865bca8e54702b08a4f26da73adc83dd1fcf26fbc875b35d81c815", size = 2020145, upload-time = "2025-11-04T13:40:21.548Z" }, - { url = "https://files.pythonhosted.org/packages/f9/22/91fbc821fa6d261b376a3f73809f907cec5ca6025642c463d3488aad22fb/pydantic_core-2.41.5-cp312-cp312-win_arm64.whl", hash = "sha256:1746d4a3d9a794cacae06a5eaaccb4b8643a131d45fbc9af23e353dc0a5ba5c3", size = 1976179, upload-time = "2025-11-04T13:40:23.393Z" }, - { url = "https://files.pythonhosted.org/packages/09/32/59b0c7e63e277fa7911c2fc70ccfb45ce4b98991e7ef37110663437005af/pydantic_core-2.41.5-graalpy312-graalpy250_312_native-macosx_10_12_x86_64.whl", hash = "sha256:7da7087d756b19037bc2c06edc6c170eeef3c3bafcb8f532ff17d64dc427adfd", size = 2110495, upload-time = "2025-11-04T13:42:49.689Z" }, - { url = "https://files.pythonhosted.org/packages/aa/81/05e400037eaf55ad400bcd318c05bb345b57e708887f07ddb2d20e3f0e98/pydantic_core-2.41.5-graalpy312-graalpy250_312_native-macosx_11_0_arm64.whl", hash = "sha256:aabf5777b5c8ca26f7824cb4a120a740c9588ed58df9b2d196ce92fba42ff8dc", size = 1915388, upload-time = "2025-11-04T13:42:52.215Z" }, { url = "https://files.pythonhosted.org/packages/6e/0d/e3549b2399f71d56476b77dbf3cf8937cec5cd70536bdc0e374a421d0599/pydantic_core-2.41.5-graalpy312-graalpy250_312_native-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:c007fe8a43d43b3969e8469004e9845944f1a80e6acd47c150856bb87f230c56", size = 1942879, upload-time = "2025-11-04T13:42:56.483Z" }, { url = "https://files.pythonhosted.org/packages/f7/07/34573da085946b6a313d7c42f82f16e8920bfd730665de2d11c0c37a74b5/pydantic_core-2.41.5-graalpy312-graalpy250_312_native-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:76d0819de158cd855d1cbb8fcafdf6f5cf1eb8e470abe056d5d161106e38062b", size = 2139017, upload-time = "2025-11-04T13:42:59.471Z" }, ] @@ -2789,8 +3163,8 @@ name = "pydantic-extra-types" version = "2.10.5" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "pydantic" }, - { name = "typing-extensions" }, + { name = "pydantic", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "typing-extensions", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, ] sdist = { url = "https://files.pythonhosted.org/packages/7e/ba/4178111ec4116c54e1dc7ecd2a1ff8f54256cdbd250e576882911e8f710a/pydantic_extra_types-2.10.5.tar.gz", hash = "sha256:1dcfa2c0cf741a422f088e0dbb4690e7bfadaaf050da3d6f80d6c3cf58a2bad8", size = 138429, upload-time = "2025-06-02T09:31:52.713Z" } wheels = [ @@ -2799,21 +3173,21 @@ wheels = [ [package.optional-dependencies] pycountry = [ - { name = "pycountry" }, + { name = "pycountry", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, ] [[package]] name = "pydantic-settings" -version = "2.10.1" +version = "2.13.1" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "pydantic" }, - { name = "python-dotenv" }, - { name = "typing-inspection" }, + { name = "pydantic", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "python-dotenv", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "typing-inspection", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/68/85/1ea668bbab3c50071ca613c6ab30047fb36ab0da1b92fa8f17bbc38fd36c/pydantic_settings-2.10.1.tar.gz", hash = "sha256:06f0062169818d0f5524420a360d632d5857b83cffd4d42fe29597807a1614ee", size = 172583, upload-time = "2025-06-24T13:26:46.841Z" } +sdist = { url = "https://files.pythonhosted.org/packages/52/6d/fffca34caecc4a3f97bda81b2098da5e8ab7efc9a66e819074a11955d87e/pydantic_settings-2.13.1.tar.gz", hash = "sha256:b4c11847b15237fb0171e1462bf540e294affb9b86db4d9aa5c01730bdbe4025", size = 223826, upload-time = "2026-02-19T13:45:08.055Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/58/f0/427018098906416f580e3cf1366d3b1abfb408a0652e9f31600c24a1903c/pydantic_settings-2.10.1-py3-none-any.whl", hash = "sha256:a60952460b99cf661dc25c29c0ef171721f98bfcb52ef8d9ea4c943d7c8cc796", size = 45235, upload-time = "2025-06-24T13:26:45.485Z" }, + { url = "https://files.pythonhosted.org/packages/00/4b/ccc026168948fec4f7555b9164c724cf4125eac006e176541483d2c959be/pydantic_settings-2.13.1-py3-none-any.whl", hash = "sha256:d56fd801823dbeae7f0975e1f8c8e25c258eb75d278ea7abb5d9cebb01b56237", size = 58929, upload-time = "2026-02-19T13:45:06.034Z" }, ] [[package]] @@ -2836,34 +3210,37 @@ wheels = [ [package.optional-dependencies] crypto = [ - { name = "cryptography" }, + { name = "cryptography", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, ] [[package]] -name = "pylatexenc" -version = "3.0a33" +name = "pypika" +version = "0.51.1" source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/97/6c/37cf2bfa76f122885acec21d62018fd9503ece009a05ea5d32dd348133e2/pylatexenc-3.0a33.tar.gz", hash = "sha256:087c9d6d280ba1242e01caf0a6436b18e9790e07148eb9430aeaa51b68d2c114", size = 205228, upload-time = "2025-08-26T09:09:26.385Z" } +sdist = { url = "https://files.pythonhosted.org/packages/f8/78/cbaebba88e05e2dcda13ca203131b38d3640219f20ebb49676d26714861b/pypika-0.51.1.tar.gz", hash = "sha256:c30c7c1048fbf056fd3920c5a2b88b0c29dd190a9b2bee971fd17e4abe4d0ebe", size = 80919, upload-time = "2026-02-04T11:27:48.304Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/c6/43/0a8a0b4b42b9b38f01f255b777bd4697099eb75ea7201444b528794a2a4f/pylatexenc-3.0a33-py3-none-any.whl", hash = "sha256:71221905d8e731c0600cde10af08f16f539723b6ffe17f478e56f5c0208c507a", size = 267673, upload-time = "2025-08-26T09:09:24.832Z" }, + { url = "https://files.pythonhosted.org/packages/57/83/c77dfeed04022e8930b08eedca2b6e5efed256ab3321396fde90066efb65/pypika-0.51.1-py2.py3-none-any.whl", hash = "sha256:77985b4d7ce71b9905255bf12468cf598349e98837c037541cfc240e528aec46", size = 60585, upload-time = "2026-02-04T11:27:46.251Z" }, ] [[package]] -name = "pyparsing" -version = "3.2.3" +name = "pyproject-hooks" +version = "1.2.0" source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/bb/22/f1129e69d94ffff626bdb5c835506b3a5b4f3d070f17ea295e12c2c6f60f/pyparsing-3.2.3.tar.gz", hash = "sha256:b9c13f1ab8b3b542f72e28f634bad4de758ab3ce4546e4301970ad6fa77c38be", size = 1088608, upload-time = "2025-03-25T05:01:28.114Z" } +sdist = { url = "https://files.pythonhosted.org/packages/e7/82/28175b2414effca1cdac8dc99f76d660e7a4fb0ceefa4b4ab8f5f6742925/pyproject_hooks-1.2.0.tar.gz", hash = "sha256:1e859bd5c40fae9448642dd871adf459e5e2084186e8d2c2a79a824c970da1f8", size = 19228, upload-time = "2024-09-29T09:24:13.293Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/05/e7/df2285f3d08fee213f2d041540fa4fc9ca6c2d44cf36d3a035bf2a8d2bcc/pyparsing-3.2.3-py3-none-any.whl", hash = "sha256:a749938e02d6fd0b59b356ca504a24982314bb090c383e3cf201c95ef7e2bfcf", size = 111120, upload-time = "2025-03-25T05:01:24.908Z" }, + { url = "https://files.pythonhosted.org/packages/bd/24/12818598c362d7f300f18e74db45963dbcb85150324092410c8b49405e42/pyproject_hooks-1.2.0-py3-none-any.whl", hash = "sha256:9e5c6bfa8dcc30091c74b0cf803c81fdd29d94f01992a7707bc97babb1141913", size = 10216, upload-time = "2024-09-29T09:24:11.978Z" }, ] [[package]] -name = "pyproject-hooks" -version = "1.2.0" +name = "pyqwest" +version = "0.3.3" source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/e7/82/28175b2414effca1cdac8dc99f76d660e7a4fb0ceefa4b4ab8f5f6742925/pyproject_hooks-1.2.0.tar.gz", hash = "sha256:1e859bd5c40fae9448642dd871adf459e5e2084186e8d2c2a79a824c970da1f8", size = 19228, upload-time = "2024-09-29T09:24:13.293Z" } +sdist = { url = "https://files.pythonhosted.org/packages/37/d5/37c968f2ad6a1e6d443215063a9fcc109999f796aaf697806ea1ccf562af/pyqwest-0.3.3.tar.gz", hash = "sha256:473a9ede8a96d1d993a1628eb62edd056a3c57b74e65cfd30d72eb451be1a23d", size = 423176, upload-time = "2026-01-21T12:54:44.453Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/bd/24/12818598c362d7f300f18e74db45963dbcb85150324092410c8b49405e42/pyproject_hooks-1.2.0-py3-none-any.whl", hash = "sha256:9e5c6bfa8dcc30091c74b0cf803c81fdd29d94f01992a7707bc97babb1141913", size = 10216, upload-time = "2024-09-29T09:24:11.978Z" }, + { url = "https://files.pythonhosted.org/packages/90/05/35c546933d6765c9f38841ee703a299cd08d115e8f9df5a882669775b07d/pyqwest-0.3.3-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:b802af427c9e4ed33c5a6ed47612887dce1128ba3cf685c7a26a383d9412c811", size = 5267540, upload-time = "2026-01-21T12:53:46.707Z" }, + { url = "https://files.pythonhosted.org/packages/f9/a6/97990418194cd24640486a14f6cb390d9f7005bc1aa06fe2324ef0aa8050/pyqwest-0.3.3-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:81118f530429b10eb295ae5b9b1353593e17ad15e0166203d2dfca51b763db85", size = 5309338, upload-time = "2026-01-21T12:53:48.903Z" }, + { url = "https://files.pythonhosted.org/packages/70/12/1e503ba23b8cfe008178610ee649e4ec6536a4d98bce9786287ae82ebb33/pyqwest-0.3.3-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:420e26aed443f6beb8a10fe5eb06913084be48e549b537b0a8887088358ed354", size = 5455416, upload-time = "2026-01-21T12:53:50.466Z" }, + { url = "https://files.pythonhosted.org/packages/70/9f/d39e0ba8fa4396f9a68ad60e2544f76d79688e93f5bb6be777f6070435e8/pyqwest-0.3.3-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:f596adcd385d2125c8c715e6c107068139d36b5ebcb14709e1ed2f98423e624e", size = 5613290, upload-time = "2026-01-21T12:53:52.428Z" }, ] [[package]] @@ -2871,11 +3248,10 @@ name = "pytest" version = "8.4.2" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "colorama", marker = "sys_platform == 'win32'" }, - { name = "iniconfig" }, - { name = "packaging" }, - { name = "pluggy" }, - { name = "pygments" }, + { name = "iniconfig", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "packaging", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "pluggy", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "pygments", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, ] sdist = { url = "https://files.pythonhosted.org/packages/a3/5c/00a0e072241553e1a7496d638deababa67c5058571567b92a7eaa258397c/pytest-8.4.2.tar.gz", hash = "sha256:86c0d0b93306b961d58d62a4db4879f27fe25513d4b969df351abdddb3c30e01", size = 1519618, upload-time = "2025-09-04T14:34:22.711Z" } wheels = [ @@ -2887,7 +3263,7 @@ name = "python-dateutil" version = "2.9.0.post0" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "six" }, + { name = "six", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, ] sdist = { url = "https://files.pythonhosted.org/packages/66/c0/0c8b6ad9f17a802ee498c46e004a0eb49bc148f2fd230864601a86dcf6db/python-dateutil-2.9.0.post0.tar.gz", hash = "sha256:37dd54208da7e1cd875388217d5e00ebd4179249f90fb72437e91a35459a0ad3", size = 342432, upload-time = "2024-03-01T18:36:20.211Z" } wheels = [ @@ -2930,31 +3306,16 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/81/c4/34e93fe5f5429d7570ec1fa436f1986fb1f00c3e0f43a589fe2bbcd22c3f/pytz-2025.2-py2.py3-none-any.whl", hash = "sha256:5ddf76296dd8c44c26eb8f4b6f35488f3ccbf6fbbd7adee0b7262d43f0ec2f00", size = 509225, upload-time = "2025-03-25T02:24:58.468Z" }, ] -[[package]] -name = "pywin32" -version = "311" -source = { registry = "https://pypi.org/simple" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/e7/ab/01ea1943d4eba0f850c3c61e78e8dd59757ff815ff3ccd0a84de5f541f42/pywin32-311-cp312-cp312-win32.whl", hash = "sha256:750ec6e621af2b948540032557b10a2d43b0cee2ae9758c54154d711cc852d31", size = 8706543, upload-time = "2025-07-14T20:13:20.765Z" }, - { url = "https://files.pythonhosted.org/packages/d1/a8/a0e8d07d4d051ec7502cd58b291ec98dcc0c3fff027caad0470b72cfcc2f/pywin32-311-cp312-cp312-win_amd64.whl", hash = "sha256:b8c095edad5c211ff31c05223658e71bf7116daa0ecf3ad85f3201ea3190d067", size = 9495040, upload-time = "2025-07-14T20:13:22.543Z" }, - { url = "https://files.pythonhosted.org/packages/ba/3a/2ae996277b4b50f17d61f0603efd8253cb2d79cc7ae159468007b586396d/pywin32-311-cp312-cp312-win_arm64.whl", hash = "sha256:e286f46a9a39c4a18b319c28f59b61de793654af2f395c102b4f819e584b5852", size = 8710102, upload-time = "2025-07-14T20:13:24.682Z" }, -] - [[package]] name = "pyyaml" version = "6.0.2" source = { registry = "https://pypi.org/simple" } sdist = { url = "https://files.pythonhosted.org/packages/54/ed/79a089b6be93607fa5cdaedf301d7dfb23af5f25c398d5ead2525b063e17/pyyaml-6.0.2.tar.gz", hash = "sha256:d584d9ec91ad65861cc08d42e834324ef890a082e591037abe114850ff7bbc3e", size = 130631, upload-time = "2024-08-06T20:33:50.674Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/86/0c/c581167fc46d6d6d7ddcfb8c843a4de25bdd27e4466938109ca68492292c/PyYAML-6.0.2-cp312-cp312-macosx_10_9_x86_64.whl", hash = "sha256:c70c95198c015b85feafc136515252a261a84561b7b1d51e3384e0655ddf25ab", size = 183873, upload-time = "2024-08-06T20:32:25.131Z" }, - { url = "https://files.pythonhosted.org/packages/a8/0c/38374f5bb272c051e2a69281d71cba6fdb983413e6758b84482905e29a5d/PyYAML-6.0.2-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:ce826d6ef20b1bc864f0a68340c8b3287705cae2f8b4b1d932177dcc76721725", size = 173302, upload-time = "2024-08-06T20:32:26.511Z" }, { url = "https://files.pythonhosted.org/packages/c3/93/9916574aa8c00aa06bbac729972eb1071d002b8e158bd0e83a3b9a20a1f7/PyYAML-6.0.2-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:1f71ea527786de97d1a0cc0eacd1defc0985dcf6b3f17bb77dcfc8c34bec4dc5", size = 739154, upload-time = "2024-08-06T20:32:28.363Z" }, - { url = "https://files.pythonhosted.org/packages/95/0f/b8938f1cbd09739c6da569d172531567dbcc9789e0029aa070856f123984/PyYAML-6.0.2-cp312-cp312-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:9b22676e8097e9e22e36d6b7bda33190d0d400f345f23d4065d48f4ca7ae0425", size = 766223, upload-time = "2024-08-06T20:32:30.058Z" }, { url = "https://files.pythonhosted.org/packages/b9/2b/614b4752f2e127db5cc206abc23a8c19678e92b23c3db30fc86ab731d3bd/PyYAML-6.0.2-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:80bab7bfc629882493af4aa31a4cfa43a4c57c83813253626916b8c7ada83476", size = 767542, upload-time = "2024-08-06T20:32:31.881Z" }, { url = "https://files.pythonhosted.org/packages/d4/00/dd137d5bcc7efea1836d6264f049359861cf548469d18da90cd8216cf05f/PyYAML-6.0.2-cp312-cp312-musllinux_1_1_aarch64.whl", hash = "sha256:0833f8694549e586547b576dcfaba4a6b55b9e96098b36cdc7ebefe667dfed48", size = 731164, upload-time = "2024-08-06T20:32:37.083Z" }, { url = "https://files.pythonhosted.org/packages/c9/1f/4f998c900485e5c0ef43838363ba4a9723ac0ad73a9dc42068b12aaba4e4/PyYAML-6.0.2-cp312-cp312-musllinux_1_1_x86_64.whl", hash = "sha256:8b9c7197f7cb2738065c481a0461e50ad02f18c78cd75775628afb4d7137fb3b", size = 756611, upload-time = "2024-08-06T20:32:38.898Z" }, - { url = "https://files.pythonhosted.org/packages/df/d1/f5a275fdb252768b7a11ec63585bc38d0e87c9e05668a139fea92b80634c/PyYAML-6.0.2-cp312-cp312-win32.whl", hash = "sha256:ef6107725bd54b262d6dedcc2af448a266975032bc85ef0172c5f059da6325b4", size = 140591, upload-time = "2024-08-06T20:32:40.241Z" }, - { url = "https://files.pythonhosted.org/packages/0c/e8/4f648c598b17c3d06e8753d7d13d57542b30d56e6c2dedf9c331ae56312e/PyYAML-6.0.2-cp312-cp312-win_amd64.whl", hash = "sha256:7e7401d0de89a9a855c839bc697c079a4af81cf878373abd7dc625847d25cbd8", size = 156338, upload-time = "2024-08-06T20:32:41.93Z" }, ] [[package]] @@ -2962,47 +3323,29 @@ name = "pyzmq" version = "27.1.0" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "cffi", marker = "implementation_name == 'pypy'" }, + { name = "cffi", marker = "(implementation_name == 'pypy' and platform_machine == 'aarch64' and sys_platform == 'linux') or (implementation_name == 'pypy' and platform_machine == 'x86_64' and sys_platform == 'linux')" }, ] sdist = { url = "https://files.pythonhosted.org/packages/04/0b/3c9baedbdf613ecaa7aa07027780b8867f57b6293b6ee50de316c9f3222b/pyzmq-27.1.0.tar.gz", hash = "sha256:ac0765e3d44455adb6ddbf4417dcce460fc40a05978c08efdf2948072f6db540", size = 281750, upload-time = "2025-09-08T23:10:18.157Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/92/e7/038aab64a946d535901103da16b953c8c9cc9c961dadcbf3609ed6428d23/pyzmq-27.1.0-cp312-abi3-macosx_10_15_universal2.whl", hash = "sha256:452631b640340c928fa343801b0d07eb0c3789a5ffa843f6e1a9cee0ba4eb4fc", size = 1306279, upload-time = "2025-09-08T23:08:03.807Z" }, - { url = "https://files.pythonhosted.org/packages/e8/5e/c3c49fdd0f535ef45eefcc16934648e9e59dace4a37ee88fc53f6cd8e641/pyzmq-27.1.0-cp312-abi3-manylinux2014_i686.manylinux_2_17_i686.whl", hash = "sha256:1c179799b118e554b66da67d88ed66cd37a169f1f23b5d9f0a231b4e8d44a113", size = 895645, upload-time = "2025-09-08T23:08:05.301Z" }, { url = "https://files.pythonhosted.org/packages/f8/e5/b0b2504cb4e903a74dcf1ebae157f9e20ebb6ea76095f6cfffea28c42ecd/pyzmq-27.1.0-cp312-abi3-manylinux_2_26_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:3837439b7f99e60312f0c926a6ad437b067356dc2bc2ec96eb395fd0fe804233", size = 652574, upload-time = "2025-09-08T23:08:06.828Z" }, { url = "https://files.pythonhosted.org/packages/f8/9b/c108cdb55560eaf253f0cbdb61b29971e9fb34d9c3499b0e96e4e60ed8a5/pyzmq-27.1.0-cp312-abi3-manylinux_2_26_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:43ad9a73e3da1fab5b0e7e13402f0b2fb934ae1c876c51d0afff0e7c052eca31", size = 840995, upload-time = "2025-09-08T23:08:08.396Z" }, { url = "https://files.pythonhosted.org/packages/c2/bb/b79798ca177b9eb0825b4c9998c6af8cd2a7f15a6a1a4272c1d1a21d382f/pyzmq-27.1.0-cp312-abi3-musllinux_1_2_aarch64.whl", hash = "sha256:0de3028d69d4cdc475bfe47a6128eb38d8bc0e8f4d69646adfbcd840facbac28", size = 1642070, upload-time = "2025-09-08T23:08:09.989Z" }, - { url = "https://files.pythonhosted.org/packages/9c/80/2df2e7977c4ede24c79ae39dcef3899bfc5f34d1ca7a5b24f182c9b7a9ca/pyzmq-27.1.0-cp312-abi3-musllinux_1_2_i686.whl", hash = "sha256:cf44a7763aea9298c0aa7dbf859f87ed7012de8bda0f3977b6fb1d96745df856", size = 2021121, upload-time = "2025-09-08T23:08:11.907Z" }, { url = "https://files.pythonhosted.org/packages/46/bd/2d45ad24f5f5ae7e8d01525eb76786fa7557136555cac7d929880519e33a/pyzmq-27.1.0-cp312-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:f30f395a9e6fbca195400ce833c731e7b64c3919aa481af4d88c3759e0cb7496", size = 1878550, upload-time = "2025-09-08T23:08:13.513Z" }, - { url = "https://files.pythonhosted.org/packages/e6/2f/104c0a3c778d7c2ab8190e9db4f62f0b6957b53c9d87db77c284b69f33ea/pyzmq-27.1.0-cp312-abi3-win32.whl", hash = "sha256:250e5436a4ba13885494412b3da5d518cd0d3a278a1ae640e113c073a5f88edd", size = 559184, upload-time = "2025-09-08T23:08:15.163Z" }, - { url = "https://files.pythonhosted.org/packages/fc/7f/a21b20d577e4100c6a41795842028235998a643b1ad406a6d4163ea8f53e/pyzmq-27.1.0-cp312-abi3-win_amd64.whl", hash = "sha256:9ce490cf1d2ca2ad84733aa1d69ce6855372cb5ce9223802450c9b2a7cba0ccf", size = 619480, upload-time = "2025-09-08T23:08:17.192Z" }, - { url = "https://files.pythonhosted.org/packages/78/c2/c012beae5f76b72f007a9e91ee9401cb88c51d0f83c6257a03e785c81cc2/pyzmq-27.1.0-cp312-abi3-win_arm64.whl", hash = "sha256:75a2f36223f0d535a0c919e23615fc85a1e23b71f40c7eb43d7b1dedb4d8f15f", size = 552993, upload-time = "2025-09-08T23:08:18.926Z" }, ] [[package]] -name = "ray" -version = "2.49.1" +name = "quack-kernels" +version = "0.3.6" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "click" }, - { name = "filelock" }, - { name = "jsonschema" }, - { name = "msgpack" }, - { name = "packaging" }, - { name = "protobuf" }, - { name = "pyyaml" }, - { name = "requests" }, + { name = "apache-tvm-ffi", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "nvidia-cutlass-dsl", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "torch", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "torch-c-dlpack-ext", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, ] +sdist = { url = "https://files.pythonhosted.org/packages/75/af/5b77f565d4268905eb7e7fe3367f877ae98d0cae920c6d9f2d019a39f214/quack_kernels-0.3.6.tar.gz", hash = "sha256:43af5a1f31099682cb1c53aa835fcdfebeb88d917c08aa342b959af88c6e8749", size = 192593, upload-time = "2026-03-24T13:30:00.923Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/27/3d/bcfab305bbbb18b47071fc03688bfeb50460dd9b43fa0c18bbdb65712aa8/ray-2.49.1-cp312-cp312-macosx_12_0_arm64.whl", hash = "sha256:f7a715855d179c1dd6ae2e8b5f8919638cde379a5b157963a0bd74d1178b8b5a", size = 66857611, upload-time = "2025-09-03T00:25:37.096Z" }, - { url = "https://files.pythonhosted.org/packages/ed/76/8530d5fc58cf73aef433567db5df5b703436e4891bc16c6cd66bee7125c8/ray-2.49.1-cp312-cp312-macosx_12_0_x86_64.whl", hash = "sha256:f9365de3a9a661ccf089dfaac01c8b68ba00c98443330ef678e0c0248272c722", size = 69262881, upload-time = "2025-09-03T00:25:42.633Z" }, - { url = "https://files.pythonhosted.org/packages/c6/ba/77eae921fc2595087516df6efc0ca03caacc14d16592341985916f1aed13/ray-2.49.1-cp312-cp312-manylinux2014_aarch64.whl", hash = "sha256:cf63b916e399c2a4d484249611a96cee283cef32e544126115a4ad3e872c34eb", size = 69287149, upload-time = "2025-09-03T00:25:47.605Z" }, - { url = "https://files.pythonhosted.org/packages/00/02/c81260c0f94bd34a1442ea488bdd433dfc9e6ed6211c9a59bc4157b8e00e/ray-2.49.1-cp312-cp312-manylinux2014_x86_64.whl", hash = "sha256:484064fca02732e0b6f4a09dad0d1fb6abd4ca4b6d9bf7c26ab7a17a2887cd09", size = 70114899, upload-time = "2025-09-03T00:25:53.144Z" }, - { url = "https://files.pythonhosted.org/packages/f4/fe/049ec402bccf0fa38f648834b496fc9c707b7593273cf9bdaf085c72b071/ray-2.49.1-cp312-cp312-win_amd64.whl", hash = "sha256:96bfdc301f38ce626fd638396cd2f52c6e3c6ba751b475f54db17152c4bdf5ab", size = 26223630, upload-time = "2025-09-03T00:25:59.965Z" }, -] - -[package.optional-dependencies] -cgraph = [ - { name = "cupy-cuda12x", marker = "sys_platform != 'darwin'" }, + { url = "https://files.pythonhosted.org/packages/8e/28/e688da97806474963b2983a8231572733de11ed48f8899410c0afc8bdf15/quack_kernels-0.3.6-py3-none-any.whl", hash = "sha256:973d974ccca816014006af8f136294ba0f15b1c324197f8b54ffc6500666e0e5", size = 198865, upload-time = "2026-03-24T13:29:59.107Z" }, ] [[package]] @@ -3010,9 +3353,9 @@ name = "referencing" version = "0.36.2" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "attrs" }, - { name = "rpds-py" }, - { name = "typing-extensions" }, + { name = "attrs", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "rpds-py", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "typing-extensions", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, ] sdist = { url = "https://files.pythonhosted.org/packages/2f/db/98b5c277be99dd18bfd91dd04e1b759cad18d1a338188c936e92f921c7e2/referencing-0.36.2.tar.gz", hash = "sha256:df2e89862cd09deabbdba16944cc3f10feb6b3e6f18e902f7cc25609a34775aa", size = 74744, upload-time = "2025-01-25T08:48:16.138Z" } wheels = [ @@ -3021,24 +3364,31 @@ wheels = [ [[package]] name = "regex" -version = "2025.9.1" +version = "2026.3.32" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/81/93/5ab3e899c47fa7994e524447135a71cd121685a35c8fe35029005f8b236f/regex-2026.3.32.tar.gz", hash = "sha256:f1574566457161678297a116fa5d1556c5a4159d64c5ff7c760e7c564bf66f16", size = 415605, upload-time = "2026-03-28T21:49:22.012Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/4a/cf/1955bb5567bc491bd63068e17f75ab0c9ff5e9d08466beec7e347f5e768d/regex-2026.3.32-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:66a5083c3ffe5a5a95f8281ea47a88072d4f24001d562d1d9d28d4cdc005fec5", size = 796431, upload-time = "2026-03-28T21:46:33.101Z" }, + { url = "https://files.pythonhosted.org/packages/0a/fe/661043d1c263b0d9d10c6ff4e9c9745f3df9641c62b51f96a3473638e7ce/regex-2026.3.32-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:f54840bea73541652f1170dc63402a5b776fc851ad36a842da9e5163c1f504a0", size = 801512, upload-time = "2026-03-28T21:46:38.587Z" }, + { url = "https://files.pythonhosted.org/packages/b6/c8/d833397b70cd1bacfcdc0a611f0e2c1f5b91fee8eedd88affcee770cbbb6/regex-2026.3.32-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:66d3126afe7eac41759cd5f0b3b246598086e88e70527c0d68c9e615b81771c4", size = 785837, upload-time = "2026-03-28T21:46:42.926Z" }, + { url = "https://files.pythonhosted.org/packages/18/f4/04ed04ebf335a44083695c22772be6a42efa31900415555563acf02cb4de/regex-2026.3.32-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:b56993a7aeb4140c4770f4f7965c9e5af4f024457d06e23c01b0d47501cb18ed", size = 788332, upload-time = "2026-03-28T21:46:50.454Z" }, +] + +[[package]] +name = "renderers" +version = "0.1.6" source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/b2/5a/4c63457fbcaf19d138d72b2e9b39405954f98c0349b31c601bfcb151582c/regex-2025.9.1.tar.gz", hash = "sha256:88ac07b38d20b54d79e704e38aa3bd2c0f8027432164226bdee201a1c0c9c9ff", size = 400852, upload-time = "2025-09-01T22:10:10.479Z" } +dependencies = [ + { name = "jinja2", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "numpy", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "openai", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "openai-harmony", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "tiktoken", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "transformers", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/7c/a7/26162494dab2d7740ff02191cb87c30b68450fb154363c7f0a434e7f3ea9/renderers-0.1.6.tar.gz", hash = "sha256:b74bc3dc870bea3c37ff5b47826ace9b8dd608a4c1f56554c39be1b20b2c63dc", size = 163768, upload-time = "2026-05-07T14:12:36.634Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/39/ef/a0372febc5a1d44c1be75f35d7e5aff40c659ecde864d7fa10e138f75e74/regex-2025.9.1-cp312-cp312-macosx_10_13_universal2.whl", hash = "sha256:84a25164bd8dcfa9f11c53f561ae9766e506e580b70279d05a7946510bdd6f6a", size = 486317, upload-time = "2025-09-01T22:08:34.529Z" }, - { url = "https://files.pythonhosted.org/packages/b5/25/d64543fb7eb41a1024786d518cc57faf1ce64aa6e9ddba097675a0c2f1d2/regex-2025.9.1-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:645e88a73861c64c1af558dd12294fb4e67b5c1eae0096a60d7d8a2143a611c7", size = 289698, upload-time = "2025-09-01T22:08:36.162Z" }, - { url = "https://files.pythonhosted.org/packages/d8/dc/fbf31fc60be317bd9f6f87daa40a8a9669b3b392aa8fe4313df0a39d0722/regex-2025.9.1-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:10a450cba5cd5409526ee1d4449f42aad38dd83ac6948cbd6d7f71ca7018f7db", size = 287242, upload-time = "2025-09-01T22:08:37.794Z" }, - { url = "https://files.pythonhosted.org/packages/0f/74/f933a607a538f785da5021acf5323961b4620972e2c2f1f39b6af4b71db7/regex-2025.9.1-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:e9dc5991592933a4192c166eeb67b29d9234f9c86344481173d1bc52f73a7104", size = 797441, upload-time = "2025-09-01T22:08:39.108Z" }, - { url = "https://files.pythonhosted.org/packages/89/d0/71fc49b4f20e31e97f199348b8c4d6e613e7b6a54a90eb1b090c2b8496d7/regex-2025.9.1-cp312-cp312-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:a32291add816961aab472f4fad344c92871a2ee33c6c219b6598e98c1f0108f2", size = 862654, upload-time = "2025-09-01T22:08:40.586Z" }, - { url = "https://files.pythonhosted.org/packages/59/05/984edce1411a5685ba9abbe10d42cdd9450aab4a022271f9585539788150/regex-2025.9.1-cp312-cp312-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:588c161a68a383478e27442a678e3b197b13c5ba51dbba40c1ccb8c4c7bee9e9", size = 910862, upload-time = "2025-09-01T22:08:42.416Z" }, - { url = "https://files.pythonhosted.org/packages/b2/02/5c891bb5fe0691cc1bad336e3a94b9097fbcf9707ec8ddc1dce9f0397289/regex-2025.9.1-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:47829ffaf652f30d579534da9085fe30c171fa2a6744a93d52ef7195dc38218b", size = 801991, upload-time = "2025-09-01T22:08:44.072Z" }, - { url = "https://files.pythonhosted.org/packages/f1/ae/fd10d6ad179910f7a1b3e0a7fde1ef8bb65e738e8ac4fd6ecff3f52252e4/regex-2025.9.1-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:1e978e5a35b293ea43f140c92a3269b6ab13fe0a2bf8a881f7ac740f5a6ade85", size = 786651, upload-time = "2025-09-01T22:08:46.079Z" }, - { url = "https://files.pythonhosted.org/packages/30/cf/9d686b07bbc5bf94c879cc168db92542d6bc9fb67088d03479fef09ba9d3/regex-2025.9.1-cp312-cp312-musllinux_1_2_ppc64le.whl", hash = "sha256:4cf09903e72411f4bf3ac1eddd624ecfd423f14b2e4bf1c8b547b72f248b7bf7", size = 856556, upload-time = "2025-09-01T22:08:48.376Z" }, - { url = "https://files.pythonhosted.org/packages/91/9d/302f8a29bb8a49528abbab2d357a793e2a59b645c54deae0050f8474785b/regex-2025.9.1-cp312-cp312-musllinux_1_2_s390x.whl", hash = "sha256:d016b0f77be63e49613c9e26aaf4a242f196cd3d7a4f15898f5f0ab55c9b24d2", size = 849001, upload-time = "2025-09-01T22:08:50.067Z" }, - { url = "https://files.pythonhosted.org/packages/93/fa/b4c6dbdedc85ef4caec54c817cd5f4418dbfa2453214119f2538082bf666/regex-2025.9.1-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:656563e620de6908cd1c9d4f7b9e0777e3341ca7db9d4383bcaa44709c90281e", size = 788138, upload-time = "2025-09-01T22:08:51.933Z" }, - { url = "https://files.pythonhosted.org/packages/4a/1b/91ee17a3cbf87f81e8c110399279d0e57f33405468f6e70809100f2ff7d8/regex-2025.9.1-cp312-cp312-win32.whl", hash = "sha256:df33f4ef07b68f7ab637b1dbd70accbf42ef0021c201660656601e8a9835de45", size = 264524, upload-time = "2025-09-01T22:08:53.75Z" }, - { url = "https://files.pythonhosted.org/packages/92/28/6ba31cce05b0f1ec6b787921903f83bd0acf8efde55219435572af83c350/regex-2025.9.1-cp312-cp312-win_amd64.whl", hash = "sha256:5aba22dfbc60cda7c0853516104724dc904caa2db55f2c3e6e984eb858d3edf3", size = 275489, upload-time = "2025-09-01T22:08:55.037Z" }, - { url = "https://files.pythonhosted.org/packages/bd/ed/ea49f324db00196e9ef7fe00dd13c6164d5173dd0f1bbe495e61bb1fb09d/regex-2025.9.1-cp312-cp312-win_arm64.whl", hash = "sha256:ec1efb4c25e1849c2685fa95da44bfde1b28c62d356f9c8d861d4dad89ed56e9", size = 268589, upload-time = "2025-09-01T22:08:56.369Z" }, + { url = "https://files.pythonhosted.org/packages/5e/ad/2cf218b9fafe2333fb3e80e123e3e2022d4923d9a61fa73ee6d79f39b563/renderers-0.1.6-py3-none-any.whl", hash = "sha256:90c626713239ec108716b7c9d194ba81ffcebe94dc003324f14fbd70e6793e89", size = 83348, upload-time = "2026-05-07T14:12:35.218Z" }, ] [[package]] @@ -3046,26 +3396,38 @@ name = "requests" version = "2.32.5" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "certifi" }, - { name = "charset-normalizer" }, - { name = "idna" }, - { name = "urllib3" }, + { name = "certifi", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "charset-normalizer", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "idna", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "urllib3", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, ] sdist = { url = "https://files.pythonhosted.org/packages/c9/74/b3ff8e6c8446842c3f5c837e9c3dfcfe2018ea6ecef224c710c85ef728f4/requests-2.32.5.tar.gz", hash = "sha256:dbba0bac56e100853db0ea71b82b4dfd5fe2bf6d3754a8893c3af500cec7d7cf", size = 134517, upload-time = "2025-08-18T20:46:02.573Z" } wheels = [ { url = "https://files.pythonhosted.org/packages/1e/db/4254e3eabe8020b458f1a747140d32277ec7a271daf1d235b70dc0b4e6e3/requests-2.32.5-py3-none-any.whl", hash = "sha256:2462f94637a34fd532264295e186976db0f5d453d1cdd31473c85a6a161affb6", size = 64738, upload-time = "2025-08-18T20:46:00.542Z" }, ] +[[package]] +name = "requests-oauthlib" +version = "2.0.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "oauthlib", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "requests", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/42/f2/05f29bc3913aea15eb670be136045bf5c5bbf4b99ecb839da9b422bb2c85/requests-oauthlib-2.0.0.tar.gz", hash = "sha256:b3dffaebd884d8cd778494369603a9e7b58d29111bf6b41bdc2dcd87203af4e9", size = 55650, upload-time = "2024-03-22T20:32:29.939Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/3b/5d/63d4ae3b9daea098d5d6f5da83984853c1bbacd5dc826764b249fe119d24/requests_oauthlib-2.0.0-py2.py3-none-any.whl", hash = "sha256:7dd8a5c40426b779b0868c404bdef9768deccf22749cde15852df527e6269b36", size = 24179, upload-time = "2024-03-22T20:32:28.055Z" }, +] + [[package]] name = "reverse-text" -version = "0.1.4" +version = "0.1.5" source = { registry = "https://hub.primeintellect.ai/primeintellect/simple/" } dependencies = [ - { name = "datasets" }, - { name = "verifiers" }, + { name = "verifiers", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, ] wheels = [ - { url = "https://hub.primeintellect.ai/primeintellect/reverse-text/@efe79870/reverse_text-0.1.4-py2.py3-none-any.whl", hash = "sha256:0adc3abcb6ceba3fa82709258c732e95c90400bc899230dce47a2e0f31a74aec" }, + { url = "https://hub.primeintellect.ai/primeintellect/reverse-text/@e3cd3fe4/reverse_text-0.1.5-py2.py3-none-any.whl", hash = "sha256:f48c5a203603551219df30f734ff5f74b825ea2a29fec05c7a6b2cd4bb657a34" }, ] [[package]] @@ -3073,8 +3435,8 @@ name = "rich" version = "14.1.0" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "markdown-it-py" }, - { name = "pygments" }, + { name = "markdown-it-py", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "pygments", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, ] sdist = { url = "https://files.pythonhosted.org/packages/fe/75/af448d8e52bf1d8fa6a9d089ca6c07ff4453d86c65c145d0a300bb073b9b/rich-14.1.0.tar.gz", hash = "sha256:e497a48b844b0320d45007cdebfeaeed8db2a4f4bcf49f15e455cfc4af11eaa8", size = 224441, upload-time = "2025-07-25T07:32:58.125Z" } wheels = [ @@ -3086,9 +3448,9 @@ name = "rich-toolkit" version = "0.15.1" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "click" }, - { name = "rich" }, - { name = "typing-extensions" }, + { name = "click", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "rich", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "typing-extensions", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, ] sdist = { url = "https://files.pythonhosted.org/packages/67/33/1a18839aaa8feef7983590c05c22c9c09d245ada6017d118325bbfcc7651/rich_toolkit-0.15.1.tar.gz", hash = "sha256:6f9630eb29f3843d19d48c3bd5706a086d36d62016687f9d0efa027ddc2dd08a", size = 115322, upload-time = "2025-09-04T09:28:11.789Z" } wheels = [ @@ -3101,20 +3463,10 @@ version = "0.6.4" source = { registry = "https://pypi.org/simple" } sdist = { url = "https://files.pythonhosted.org/packages/73/46/05a94dc55ac03cf931d18e43b86ecee5ee054cb88b7853fffd741e35009c/rignore-0.6.4.tar.gz", hash = "sha256:e893fdd2d7fdcfa9407d0b7600ef2c2e2df97f55e1c45d4a8f54364829ddb0ab", size = 11633, upload-time = "2025-07-19T19:24:46.219Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/ec/6c/e5af4383cdd7829ef9aa63ac82a6507983e02dbc7c2e7b9aa64b7b8e2c7a/rignore-0.6.4-cp312-cp312-macosx_10_12_x86_64.whl", hash = "sha256:74720d074b79f32449d5d212ce732e0144a294a184246d1f1e7bcc1fc5c83b69", size = 885885, upload-time = "2025-07-19T19:23:53.236Z" }, - { url = "https://files.pythonhosted.org/packages/89/3e/1b02a868830e464769aa417ee195ac352fe71ff818df8ce50c4b998edb9c/rignore-0.6.4-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:0a8184fcf567bd6b6d7b85a0c138d98dd40f63054141c96b175844414c5530d7", size = 819736, upload-time = "2025-07-19T19:23:46.565Z" }, { url = "https://files.pythonhosted.org/packages/e0/75/b9be0c523d97c09f3c6508a67ce376aba4efe41c333c58903a0d7366439a/rignore-0.6.4-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:bcb0d7d7ecc3fbccf6477bb187c04a091579ea139f15f139abe0b3b48bdfef69", size = 892779, upload-time = "2025-07-19T19:22:35.167Z" }, - { url = "https://files.pythonhosted.org/packages/91/f4/3064b06233697f2993485d132f06fe95061fef71631485da75aed246c4fd/rignore-0.6.4-cp312-cp312-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:feac73377a156fb77b3df626c76f7e5893d9b4e9e886ac8c0f9d44f1206a2a91", size = 872116, upload-time = "2025-07-19T19:22:47.828Z" }, - { url = "https://files.pythonhosted.org/packages/99/94/cb8e7af9a3c0a665f10e2366144e0ebc66167cf846aca5f1ac31b3661598/rignore-0.6.4-cp312-cp312-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:465179bc30beb1f7a3439e428739a2b5777ed26660712b8c4e351b15a7c04483", size = 1163345, upload-time = "2025-07-19T19:23:00.557Z" }, - { url = "https://files.pythonhosted.org/packages/86/6b/49faa7ad85ceb6ccef265df40091d9992232d7f6055fa664fe0a8b13781c/rignore-0.6.4-cp312-cp312-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:4a4877b4dca9cf31a4d09845b300c677c86267657540d0b4d3e6d0ce3110e6e9", size = 939967, upload-time = "2025-07-19T19:23:13.494Z" }, { url = "https://files.pythonhosted.org/packages/80/c8/b91afda10bd5ca1e3a80463340b899c0dc26a7750a9f3c94f668585c7f40/rignore-0.6.4-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:456456802b1e77d1e2d149320ee32505b8183e309e228129950b807d204ddd17", size = 949717, upload-time = "2025-07-19T19:23:36.404Z" }, - { url = "https://files.pythonhosted.org/packages/3f/f1/88bfdde58ae3fb1c1a92bb801f492eea8eafcdaf05ab9b75130023a4670b/rignore-0.6.4-cp312-cp312-manylinux_2_5_i686.manylinux1_i686.whl", hash = "sha256:4c1ff2fc223f1d9473d36923160af37bf765548578eb9d47a2f52e90da8ae408", size = 975534, upload-time = "2025-07-19T19:23:25.988Z" }, { url = "https://files.pythonhosted.org/packages/aa/8f/a80b4a2e48ceba56ba19e096d41263d844757e10aa36ede212571b5d8117/rignore-0.6.4-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:e445fbc214ae18e0e644a78086ea5d0f579e210229a4fbe86367d11a4cd03c11", size = 1067837, upload-time = "2025-07-19T19:23:59.888Z" }, - { url = "https://files.pythonhosted.org/packages/7d/90/0905597af0e78748909ef58418442a480ddd93e9fc89b0ca9ab170c357c0/rignore-0.6.4-cp312-cp312-musllinux_1_2_armv7l.whl", hash = "sha256:e07d9c5270fc869bc431aadcfb6ed0447f89b8aafaa666914c077435dc76a123", size = 1134959, upload-time = "2025-07-19T19:24:12.396Z" }, - { url = "https://files.pythonhosted.org/packages/cc/7d/0fa29adf9183b61947ce6dc8a1a9779a8ea16573f557be28ec893f6ddbaa/rignore-0.6.4-cp312-cp312-musllinux_1_2_i686.whl", hash = "sha256:7a6ccc0ea83d2c0c6df6b166f2acacedcc220a516436490f41e99a5ae73b6019", size = 1109708, upload-time = "2025-07-19T19:24:24.176Z" }, { url = "https://files.pythonhosted.org/packages/4e/a7/92892ed86b2e36da403dd3a0187829f2d880414cef75bd612bfdf4dedebc/rignore-0.6.4-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:536392c5ec91755db48389546c833c4ab1426fe03e5a8522992b54ef8a244e7e", size = 1120546, upload-time = "2025-07-19T19:24:36.377Z" }, - { url = "https://files.pythonhosted.org/packages/31/1b/d29ae1fe901d523741d6d1d3ffe0d630734dd0ed6b047628a69c1e15ea44/rignore-0.6.4-cp312-cp312-win32.whl", hash = "sha256:f5f9dca46fc41c0a1e236767f68be9d63bdd2726db13a0ae3a30f68414472969", size = 642005, upload-time = "2025-07-19T19:24:56.671Z" }, - { url = "https://files.pythonhosted.org/packages/1a/41/a224944824688995374e4525115ce85fecd82442fc85edd5bcd81f4f256d/rignore-0.6.4-cp312-cp312-win_amd64.whl", hash = "sha256:e02eecb9e1b9f9bf7c9030ae73308a777bed3b2486204cc74dfcfbe699ab1497", size = 720358, upload-time = "2025-07-19T19:24:49.959Z" }, ] [[package]] @@ -3132,21 +3484,10 @@ version = "0.27.1" source = { registry = "https://pypi.org/simple" } sdist = { url = "https://files.pythonhosted.org/packages/e9/dd/2c0cbe774744272b0ae725f44032c77bdcab6e8bcf544bffa3b6e70c8dba/rpds_py-0.27.1.tar.gz", hash = "sha256:26a1c73171d10b7acccbded82bf6a586ab8203601e565badc74bbbf8bc5a10f8", size = 27479, upload-time = "2025-08-27T12:16:36.024Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/bd/fe/38de28dee5df58b8198c743fe2bea0c785c6d40941b9950bac4cdb71a014/rpds_py-0.27.1-cp312-cp312-macosx_10_12_x86_64.whl", hash = "sha256:ae2775c1973e3c30316892737b91f9283f9908e3cc7625b9331271eaaed7dc90", size = 361887, upload-time = "2025-08-27T12:13:10.233Z" }, - { url = "https://files.pythonhosted.org/packages/7c/9a/4b6c7eedc7dd90986bf0fab6ea2a091ec11c01b15f8ba0a14d3f80450468/rpds_py-0.27.1-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:2643400120f55c8a96f7c9d858f7be0c88d383cd4653ae2cf0d0c88f668073e5", size = 345795, upload-time = "2025-08-27T12:13:11.65Z" }, { url = "https://files.pythonhosted.org/packages/6f/0e/e650e1b81922847a09cca820237b0edee69416a01268b7754d506ade11ad/rpds_py-0.27.1-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:16323f674c089b0360674a4abd28d5042947d54ba620f72514d69be4ff64845e", size = 385121, upload-time = "2025-08-27T12:13:13.008Z" }, - { url = "https://files.pythonhosted.org/packages/1b/ea/b306067a712988e2bff00dcc7c8f31d26c29b6d5931b461aa4b60a013e33/rpds_py-0.27.1-cp312-cp312-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:9a1f4814b65eacac94a00fc9a526e3fdafd78e439469644032032d0d63de4881", size = 398976, upload-time = "2025-08-27T12:13:14.368Z" }, - { url = "https://files.pythonhosted.org/packages/2c/0a/26dc43c8840cb8fe239fe12dbc8d8de40f2365e838f3d395835dde72f0e5/rpds_py-0.27.1-cp312-cp312-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:7ba32c16b064267b22f1850a34051121d423b6f7338a12b9459550eb2096e7ec", size = 525953, upload-time = "2025-08-27T12:13:15.774Z" }, - { url = "https://files.pythonhosted.org/packages/22/14/c85e8127b573aaf3a0cbd7fbb8c9c99e735a4a02180c84da2a463b766e9e/rpds_py-0.27.1-cp312-cp312-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:e5c20f33fd10485b80f65e800bbe5f6785af510b9f4056c5a3c612ebc83ba6cb", size = 407915, upload-time = "2025-08-27T12:13:17.379Z" }, { url = "https://files.pythonhosted.org/packages/ed/7b/8f4fee9ba1fb5ec856eb22d725a4efa3deb47f769597c809e03578b0f9d9/rpds_py-0.27.1-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:466bfe65bd932da36ff279ddd92de56b042f2266d752719beb97b08526268ec5", size = 386883, upload-time = "2025-08-27T12:13:18.704Z" }, - { url = "https://files.pythonhosted.org/packages/86/47/28fa6d60f8b74fcdceba81b272f8d9836ac0340570f68f5df6b41838547b/rpds_py-0.27.1-cp312-cp312-manylinux_2_31_riscv64.whl", hash = "sha256:41e532bbdcb57c92ba3be62c42e9f096431b4cf478da9bc3bc6ce5c38ab7ba7a", size = 405699, upload-time = "2025-08-27T12:13:20.089Z" }, - { url = "https://files.pythonhosted.org/packages/d0/fd/c5987b5e054548df56953a21fe2ebed51fc1ec7c8f24fd41c067b68c4a0a/rpds_py-0.27.1-cp312-cp312-manylinux_2_5_i686.manylinux1_i686.whl", hash = "sha256:f149826d742b406579466283769a8ea448eed82a789af0ed17b0cd5770433444", size = 423713, upload-time = "2025-08-27T12:13:21.436Z" }, { url = "https://files.pythonhosted.org/packages/ac/ba/3c4978b54a73ed19a7d74531be37a8bcc542d917c770e14d372b8daea186/rpds_py-0.27.1-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:80c60cfb5310677bd67cb1e85a1e8eb52e12529545441b43e6f14d90b878775a", size = 562324, upload-time = "2025-08-27T12:13:22.789Z" }, - { url = "https://files.pythonhosted.org/packages/b5/6c/6943a91768fec16db09a42b08644b960cff540c66aab89b74be6d4a144ba/rpds_py-0.27.1-cp312-cp312-musllinux_1_2_i686.whl", hash = "sha256:7ee6521b9baf06085f62ba9c7a3e5becffbc32480d2f1b351559c001c38ce4c1", size = 593646, upload-time = "2025-08-27T12:13:24.122Z" }, { url = "https://files.pythonhosted.org/packages/11/73/9d7a8f4be5f4396f011a6bb7a19fe26303a0dac9064462f5651ced2f572f/rpds_py-0.27.1-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:a512c8263249a9d68cac08b05dd59d2b3f2061d99b322813cbcc14c3c7421998", size = 558137, upload-time = "2025-08-27T12:13:25.557Z" }, - { url = "https://files.pythonhosted.org/packages/6e/96/6772cbfa0e2485bcceef8071de7821f81aeac8bb45fbfd5542a3e8108165/rpds_py-0.27.1-cp312-cp312-win32.whl", hash = "sha256:819064fa048ba01b6dadc5116f3ac48610435ac9a0058bbde98e569f9e785c39", size = 221343, upload-time = "2025-08-27T12:13:26.967Z" }, - { url = "https://files.pythonhosted.org/packages/67/b6/c82f0faa9af1c6a64669f73a17ee0eeef25aff30bb9a1c318509efe45d84/rpds_py-0.27.1-cp312-cp312-win_amd64.whl", hash = "sha256:d9199717881f13c32c4046a15f024971a3b78ad4ea029e8da6b86e5aa9cf4594", size = 232497, upload-time = "2025-08-27T12:13:28.326Z" }, - { url = "https://files.pythonhosted.org/packages/e1/96/2817b44bd2ed11aebacc9251da03689d56109b9aba5e311297b6902136e2/rpds_py-0.27.1-cp312-cp312-win_arm64.whl", hash = "sha256:33aa65b97826a0e885ef6e278fbd934e98cdcfed80b63946025f01e2f5b29502", size = 222790, upload-time = "2025-08-27T12:13:29.71Z" }, ] [[package]] @@ -3155,24 +3496,10 @@ version = "0.13.0" source = { registry = "https://pypi.org/simple" } sdist = { url = "https://files.pythonhosted.org/packages/6e/1a/1f4b722862840295bcaba8c9e5261572347509548faaa99b2d57ee7bfe6a/ruff-0.13.0.tar.gz", hash = "sha256:5b4b1ee7eb35afae128ab94459b13b2baaed282b1fb0f472a73c82c996c8ae60", size = 5372863, upload-time = "2025-09-10T16:25:37.917Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/ac/fe/6f87b419dbe166fd30a991390221f14c5b68946f389ea07913e1719741e0/ruff-0.13.0-py3-none-linux_armv6l.whl", hash = "sha256:137f3d65d58ee828ae136a12d1dc33d992773d8f7644bc6b82714570f31b2004", size = 12187826, upload-time = "2025-09-10T16:24:39.5Z" }, - { url = "https://files.pythonhosted.org/packages/e4/25/c92296b1fc36d2499e12b74a3fdb230f77af7bdf048fad7b0a62e94ed56a/ruff-0.13.0-py3-none-macosx_10_12_x86_64.whl", hash = "sha256:21ae48151b66e71fd111b7d79f9ad358814ed58c339631450c66a4be33cc28b9", size = 12933428, upload-time = "2025-09-10T16:24:43.866Z" }, - { url = "https://files.pythonhosted.org/packages/44/cf/40bc7221a949470307d9c35b4ef5810c294e6cfa3caafb57d882731a9f42/ruff-0.13.0-py3-none-macosx_11_0_arm64.whl", hash = "sha256:64de45f4ca5441209e41742d527944635a05a6e7c05798904f39c85bafa819e3", size = 12095543, upload-time = "2025-09-10T16:24:46.638Z" }, { url = "https://files.pythonhosted.org/packages/f1/03/8b5ff2a211efb68c63a1d03d157e924997ada87d01bebffbd13a0f3fcdeb/ruff-0.13.0-py3-none-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:2b2c653ae9b9d46e0ef62fc6fbf5b979bda20a0b1d2b22f8f7eb0cde9f4963b8", size = 12312489, upload-time = "2025-09-10T16:24:49.556Z" }, - { url = "https://files.pythonhosted.org/packages/37/fc/2336ef6d5e9c8d8ea8305c5f91e767d795cd4fc171a6d97ef38a5302dadc/ruff-0.13.0-py3-none-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:4cec632534332062bc9eb5884a267b689085a1afea9801bf94e3ba7498a2d207", size = 11991631, upload-time = "2025-09-10T16:24:53.439Z" }, - { url = "https://files.pythonhosted.org/packages/39/7f/f6d574d100fca83d32637d7f5541bea2f5e473c40020bbc7fc4a4d5b7294/ruff-0.13.0-py3-none-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:dcd628101d9f7d122e120ac7c17e0a0f468b19bc925501dbe03c1cb7f5415b24", size = 13720602, upload-time = "2025-09-10T16:24:56.392Z" }, - { url = "https://files.pythonhosted.org/packages/fd/c8/a8a5b81d8729b5d1f663348d11e2a9d65a7a9bd3c399763b1a51c72be1ce/ruff-0.13.0-py3-none-manylinux_2_17_ppc64.manylinux2014_ppc64.whl", hash = "sha256:afe37db8e1466acb173bb2a39ca92df00570e0fd7c94c72d87b51b21bb63efea", size = 14697751, upload-time = "2025-09-10T16:24:59.89Z" }, - { url = "https://files.pythonhosted.org/packages/57/f5/183ec292272ce7ec5e882aea74937f7288e88ecb500198b832c24debc6d3/ruff-0.13.0-py3-none-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:0f96a8d90bb258d7d3358b372905fe7333aaacf6c39e2408b9f8ba181f4b6ef2", size = 14095317, upload-time = "2025-09-10T16:25:03.025Z" }, - { url = "https://files.pythonhosted.org/packages/9f/8d/7f9771c971724701af7926c14dab31754e7b303d127b0d3f01116faef456/ruff-0.13.0-py3-none-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:94b5e3d883e4f924c5298e3f2ee0f3085819c14f68d1e5b6715597681433f153", size = 13144418, upload-time = "2025-09-10T16:25:06.272Z" }, { url = "https://files.pythonhosted.org/packages/a8/a6/7985ad1778e60922d4bef546688cd8a25822c58873e9ff30189cfe5dc4ab/ruff-0.13.0-py3-none-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:03447f3d18479df3d24917a92d768a89f873a7181a064858ea90a804a7538991", size = 13370843, upload-time = "2025-09-10T16:25:09.965Z" }, - { url = "https://files.pythonhosted.org/packages/64/1c/bafdd5a7a05a50cc51d9f5711da704942d8dd62df3d8c70c311e98ce9f8a/ruff-0.13.0-py3-none-manylinux_2_31_riscv64.whl", hash = "sha256:fbc6b1934eb1c0033da427c805e27d164bb713f8e273a024a7e86176d7f462cf", size = 13321891, upload-time = "2025-09-10T16:25:12.969Z" }, { url = "https://files.pythonhosted.org/packages/bc/3e/7817f989cb9725ef7e8d2cee74186bf90555279e119de50c750c4b7a72fe/ruff-0.13.0-py3-none-musllinux_1_2_aarch64.whl", hash = "sha256:a8ab6a3e03665d39d4a25ee199d207a488724f022db0e1fe4002968abdb8001b", size = 12119119, upload-time = "2025-09-10T16:25:16.621Z" }, - { url = "https://files.pythonhosted.org/packages/58/07/9df080742e8d1080e60c426dce6e96a8faf9a371e2ce22eef662e3839c95/ruff-0.13.0-py3-none-musllinux_1_2_armv7l.whl", hash = "sha256:d2a5c62f8ccc6dd2fe259917482de7275cecc86141ee10432727c4816235bc41", size = 11961594, upload-time = "2025-09-10T16:25:19.49Z" }, - { url = "https://files.pythonhosted.org/packages/6a/f4/ae1185349197d26a2316840cb4d6c3fba61d4ac36ed728bf0228b222d71f/ruff-0.13.0-py3-none-musllinux_1_2_i686.whl", hash = "sha256:b7b85ca27aeeb1ab421bc787009831cffe6048faae08ad80867edab9f2760945", size = 12933377, upload-time = "2025-09-10T16:25:22.371Z" }, { url = "https://files.pythonhosted.org/packages/b6/39/e776c10a3b349fc8209a905bfb327831d7516f6058339a613a8d2aaecacd/ruff-0.13.0-py3-none-musllinux_1_2_x86_64.whl", hash = "sha256:79ea0c44a3032af768cabfd9616e44c24303af49d633b43e3a5096e009ebe823", size = 13418555, upload-time = "2025-09-10T16:25:25.681Z" }, - { url = "https://files.pythonhosted.org/packages/46/09/dca8df3d48e8b3f4202bf20b1658898e74b6442ac835bfe2c1816d926697/ruff-0.13.0-py3-none-win32.whl", hash = "sha256:4e473e8f0e6a04e4113f2e1de12a5039579892329ecc49958424e5568ef4f768", size = 12141613, upload-time = "2025-09-10T16:25:28.664Z" }, - { url = "https://files.pythonhosted.org/packages/61/21/0647eb71ed99b888ad50e44d8ec65d7148babc0e242d531a499a0bbcda5f/ruff-0.13.0-py3-none-win_amd64.whl", hash = "sha256:48e5c25c7a3713eea9ce755995767f4dcd1b0b9599b638b12946e892123d1efb", size = 13258250, upload-time = "2025-09-10T16:25:31.773Z" }, - { url = "https://files.pythonhosted.org/packages/e1/a3/03216a6a86c706df54422612981fb0f9041dbb452c3401501d4a22b942c9/ruff-0.13.0-py3-none-win_arm64.whl", hash = "sha256:ab80525317b1e1d38614addec8ac954f1b3e662de9d59114ecbf771d00cf613e", size = 12312357, upload-time = "2025-09-10T16:25:35.595Z" }, ] [[package]] @@ -3181,40 +3508,22 @@ version = "0.6.2" source = { registry = "https://pypi.org/simple" } sdist = { url = "https://files.pythonhosted.org/packages/ac/cc/738f3011628920e027a11754d9cae9abec1aed00f7ae860abbf843755233/safetensors-0.6.2.tar.gz", hash = "sha256:43ff2aa0e6fa2dc3ea5524ac7ad93a9839256b8703761e76e2d0b2a3fa4f15d9", size = 197968, upload-time = "2025-08-08T13:13:58.654Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/4d/b1/3f5fd73c039fc87dba3ff8b5d528bfc5a32b597fea8e7a6a4800343a17c7/safetensors-0.6.2-cp38-abi3-macosx_10_12_x86_64.whl", hash = "sha256:9c85ede8ec58f120bad982ec47746981e210492a6db876882aa021446af8ffba", size = 454797, upload-time = "2025-08-08T13:13:52.066Z" }, - { url = "https://files.pythonhosted.org/packages/8c/c9/bb114c158540ee17907ec470d01980957fdaf87b4aa07914c24eba87b9c6/safetensors-0.6.2-cp38-abi3-macosx_11_0_arm64.whl", hash = "sha256:d6675cf4b39c98dbd7d940598028f3742e0375a6b4d4277e76beb0c35f4b843b", size = 432206, upload-time = "2025-08-08T13:13:50.931Z" }, { url = "https://files.pythonhosted.org/packages/d3/8e/f70c34e47df3110e8e0bb268d90db8d4be8958a54ab0336c9be4fe86dac8/safetensors-0.6.2-cp38-abi3-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:1d2d2b3ce1e2509c68932ca03ab8f20570920cd9754b05063d4368ee52833ecd", size = 473261, upload-time = "2025-08-08T13:13:41.259Z" }, - { url = "https://files.pythonhosted.org/packages/2a/f5/be9c6a7c7ef773e1996dc214e73485286df1836dbd063e8085ee1976f9cb/safetensors-0.6.2-cp38-abi3-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:93de35a18f46b0f5a6a1f9e26d91b442094f2df02e9fd7acf224cfec4238821a", size = 485117, upload-time = "2025-08-08T13:13:43.506Z" }, - { url = "https://files.pythonhosted.org/packages/c9/55/23f2d0a2c96ed8665bf17a30ab4ce5270413f4d74b6d87dd663258b9af31/safetensors-0.6.2-cp38-abi3-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:89a89b505f335640f9120fac65ddeb83e40f1fd081cb8ed88b505bdccec8d0a1", size = 616154, upload-time = "2025-08-08T13:13:45.096Z" }, - { url = "https://files.pythonhosted.org/packages/98/c6/affb0bd9ce02aa46e7acddbe087912a04d953d7a4d74b708c91b5806ef3f/safetensors-0.6.2-cp38-abi3-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:fc4d0d0b937e04bdf2ae6f70cd3ad51328635fe0e6214aa1fc811f3b576b3bda", size = 520713, upload-time = "2025-08-08T13:13:46.25Z" }, { url = "https://files.pythonhosted.org/packages/fe/5d/5a514d7b88e310c8b146e2404e0dc161282e78634d9358975fd56dfd14be/safetensors-0.6.2-cp38-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:8045db2c872db8f4cbe3faa0495932d89c38c899c603f21e9b6486951a5ecb8f", size = 485835, upload-time = "2025-08-08T13:13:49.373Z" }, - { url = "https://files.pythonhosted.org/packages/7a/7b/4fc3b2ba62c352b2071bea9cfbad330fadda70579f617506ae1a2f129cab/safetensors-0.6.2-cp38-abi3-manylinux_2_5_i686.manylinux1_i686.whl", hash = "sha256:81e67e8bab9878bb568cffbc5f5e655adb38d2418351dc0859ccac158f753e19", size = 521503, upload-time = "2025-08-08T13:13:47.651Z" }, { url = "https://files.pythonhosted.org/packages/5a/50/0057e11fe1f3cead9254315a6c106a16dd4b1a19cd247f7cc6414f6b7866/safetensors-0.6.2-cp38-abi3-musllinux_1_2_aarch64.whl", hash = "sha256:b0e4d029ab0a0e0e4fdf142b194514695b1d7d3735503ba700cf36d0fc7136ce", size = 652256, upload-time = "2025-08-08T13:13:53.167Z" }, - { url = "https://files.pythonhosted.org/packages/e9/29/473f789e4ac242593ac1656fbece6e1ecd860bb289e635e963667807afe3/safetensors-0.6.2-cp38-abi3-musllinux_1_2_armv7l.whl", hash = "sha256:fa48268185c52bfe8771e46325a1e21d317207bcabcb72e65c6e28e9ffeb29c7", size = 747281, upload-time = "2025-08-08T13:13:54.656Z" }, - { url = "https://files.pythonhosted.org/packages/68/52/f7324aad7f2df99e05525c84d352dc217e0fa637a4f603e9f2eedfbe2c67/safetensors-0.6.2-cp38-abi3-musllinux_1_2_i686.whl", hash = "sha256:d83c20c12c2d2f465997c51b7ecb00e407e5f94d7dec3ea0cc11d86f60d3fde5", size = 692286, upload-time = "2025-08-08T13:13:55.884Z" }, { url = "https://files.pythonhosted.org/packages/ad/fe/cad1d9762868c7c5dc70c8620074df28ebb1a8e4c17d4c0cb031889c457e/safetensors-0.6.2-cp38-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:d944cea65fad0ead848b6ec2c37cc0b197194bec228f8020054742190e9312ac", size = 655957, upload-time = "2025-08-08T13:13:57.029Z" }, - { url = "https://files.pythonhosted.org/packages/59/a7/e2158e17bbe57d104f0abbd95dff60dda916cf277c9f9663b4bf9bad8b6e/safetensors-0.6.2-cp38-abi3-win32.whl", hash = "sha256:cab75ca7c064d3911411461151cb69380c9225798a20e712b102edda2542ddb1", size = 308926, upload-time = "2025-08-08T13:14:01.095Z" }, - { url = "https://files.pythonhosted.org/packages/2c/c3/c0be1135726618dc1e28d181b8c442403d8dbb9e273fd791de2d4384bcdd/safetensors-0.6.2-cp38-abi3-win_amd64.whl", hash = "sha256:c7b214870df923cbc1593c3faee16bec59ea462758699bd3fee399d00aac072c", size = 320192, upload-time = "2025-08-08T13:13:59.467Z" }, ] [[package]] -name = "scipy" -version = "1.16.1" -source = { registry = "https://pypi.org/simple" } +name = "science-env" +version = "0.1.2" +source = { registry = "https://hub.primeintellect.ai/primeintellect/simple/" } dependencies = [ - { name = "numpy" }, + { name = "math-verify", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "verifiers", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/f5/4a/b927028464795439faec8eaf0b03b011005c487bb2d07409f28bf30879c4/scipy-1.16.1.tar.gz", hash = "sha256:44c76f9e8b6e8e488a586190ab38016e4ed2f8a038af7cd3defa903c0a2238b3", size = 30580861, upload-time = "2025-07-27T16:33:30.834Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/f8/d9/ec4864f5896232133f51382b54a08de91a9d1af7a76dfa372894026dfee2/scipy-1.16.1-cp312-cp312-macosx_10_14_x86_64.whl", hash = "sha256:81b433bbeaf35728dad619afc002db9b189e45eebe2cd676effe1fb93fef2b9c", size = 36575194, upload-time = "2025-07-27T16:27:41.321Z" }, - { url = "https://files.pythonhosted.org/packages/5c/6d/40e81ecfb688e9d25d34a847dca361982a6addf8e31f0957b1a54fbfa994/scipy-1.16.1-cp312-cp312-macosx_12_0_arm64.whl", hash = "sha256:886cc81fdb4c6903a3bb0464047c25a6d1016fef77bb97949817d0c0d79f9e04", size = 28594590, upload-time = "2025-07-27T16:27:49.204Z" }, - { url = "https://files.pythonhosted.org/packages/0e/37/9f65178edfcc629377ce9a64fc09baebea18c80a9e57ae09a52edf84880b/scipy-1.16.1-cp312-cp312-macosx_14_0_arm64.whl", hash = "sha256:15240c3aac087a522b4eaedb09f0ad061753c5eebf1ea430859e5bf8640d5919", size = 20866458, upload-time = "2025-07-27T16:27:54.98Z" }, - { url = "https://files.pythonhosted.org/packages/2c/7b/749a66766871ea4cb1d1ea10f27004db63023074c22abed51f22f09770e0/scipy-1.16.1-cp312-cp312-macosx_14_0_x86_64.whl", hash = "sha256:65f81a25805f3659b48126b5053d9e823d3215e4a63730b5e1671852a1705921", size = 23539318, upload-time = "2025-07-27T16:28:01.604Z" }, - { url = "https://files.pythonhosted.org/packages/c4/db/8d4afec60eb833a666434d4541a3151eedbf2494ea6d4d468cbe877f00cd/scipy-1.16.1-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:6c62eea7f607f122069b9bad3f99489ddca1a5173bef8a0c75555d7488b6f725", size = 33292899, upload-time = "2025-07-27T16:28:09.147Z" }, - { url = "https://files.pythonhosted.org/packages/51/1e/79023ca3bbb13a015d7d2757ecca3b81293c663694c35d6541b4dca53e98/scipy-1.16.1-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:f965bbf3235b01c776115ab18f092a95aa74c271a52577bcb0563e85738fd618", size = 35162637, upload-time = "2025-07-27T16:28:17.535Z" }, - { url = "https://files.pythonhosted.org/packages/b6/49/0648665f9c29fdaca4c679182eb972935b3b4f5ace41d323c32352f29816/scipy-1.16.1-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:f006e323874ffd0b0b816d8c6a8e7f9a73d55ab3b8c3f72b752b226d0e3ac83d", size = 35490507, upload-time = "2025-07-27T16:28:25.705Z" }, - { url = "https://files.pythonhosted.org/packages/62/8f/66cbb9d6bbb18d8c658f774904f42a92078707a7c71e5347e8bf2f52bb89/scipy-1.16.1-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:e8fd15fc5085ab4cca74cb91fe0a4263b1f32e4420761ddae531ad60934c2119", size = 37923998, upload-time = "2025-07-27T16:28:34.339Z" }, - { url = "https://files.pythonhosted.org/packages/14/c3/61f273ae550fbf1667675701112e380881905e28448c080b23b5a181df7c/scipy-1.16.1-cp312-cp312-win_amd64.whl", hash = "sha256:f7b8013c6c066609577d910d1a2a077021727af07b6fab0ee22c2f901f22352a", size = 38508060, upload-time = "2025-07-27T16:28:43.242Z" }, + { url = "https://hub.primeintellect.ai/primeintellect/science-env/@4d326035/science_env-0.1.2-py3-none-any.whl", hash = "sha256:b601d370b32eb1bbee8f2ab7560fa5f3baa22c8ea3541a7639e355b5084d451e" }, ] [[package]] @@ -3223,14 +3532,8 @@ version = "0.2.1" source = { registry = "https://pypi.org/simple" } sdist = { url = "https://files.pythonhosted.org/packages/15/15/2e7a025fc62d764b151ae6d0f2a92f8081755ebe8d4a64099accc6f77ba6/sentencepiece-0.2.1.tar.gz", hash = "sha256:8138cec27c2f2282f4a34d9a016e3374cd40e5c6e9cb335063db66a0a3b71fad", size = 3228515, upload-time = "2025-08-12T07:00:51.718Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/4a/be/32ce495aa1d0e0c323dcb1ba87096037358edee539cac5baf8755a6bd396/sentencepiece-0.2.1-cp312-cp312-macosx_10_13_universal2.whl", hash = "sha256:57cae326c8727de58c85977b175af132a7138d84c764635d7e71bbee7e774133", size = 1943152, upload-time = "2025-08-12T06:59:40.048Z" }, - { url = "https://files.pythonhosted.org/packages/88/7e/ff23008899a58678e98c6ff592bf4d368eee5a71af96d0df6b38a039dd4f/sentencepiece-0.2.1-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:56dd39a3c4d6493db3cdca7e8cc68c6b633f0d4195495cbadfcf5af8a22d05a6", size = 1325651, upload-time = "2025-08-12T06:59:41.536Z" }, - { url = "https://files.pythonhosted.org/packages/19/84/42eb3ce4796777a1b5d3699dfd4dca85113e68b637f194a6c8d786f16a04/sentencepiece-0.2.1-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:d9381351182ff9888cc80e41c632e7e274b106f450de33d67a9e8f6043da6f76", size = 1253645, upload-time = "2025-08-12T06:59:42.903Z" }, { url = "https://files.pythonhosted.org/packages/89/fa/d3d5ebcba3cb9e6d3775a096251860c41a6bc53a1b9461151df83fe93255/sentencepiece-0.2.1-cp312-cp312-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:99f955df238021bf11f0fc37cdb54fd5e5b5f7fd30ecc3d93fb48b6815437167", size = 1316273, upload-time = "2025-08-12T06:59:44.476Z" }, { url = "https://files.pythonhosted.org/packages/04/88/14f2f4a2b922d8b39be45bf63d79e6cd3a9b2f248b2fcb98a69b12af12f5/sentencepiece-0.2.1-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:0cdfecef430d985f1c2bcbfff3defd1d95dae876fbd0173376012d2d7d24044b", size = 1387881, upload-time = "2025-08-12T06:59:46.09Z" }, - { url = "https://files.pythonhosted.org/packages/fd/b8/903e5ccb77b4ef140605d5d71b4f9e0ad95d456d6184688073ed11712809/sentencepiece-0.2.1-cp312-cp312-win32.whl", hash = "sha256:a483fd29a34c3e34c39ac5556b0a90942bec253d260235729e50976f5dba1068", size = 999540, upload-time = "2025-08-12T06:59:48.023Z" }, - { url = "https://files.pythonhosted.org/packages/2d/81/92df5673c067148c2545b1bfe49adfd775bcc3a169a047f5a0e6575ddaca/sentencepiece-0.2.1-cp312-cp312-win_amd64.whl", hash = "sha256:4cdc7c36234fda305e85c32949c5211faaf8dd886096c7cea289ddc12a2d02de", size = 1054671, upload-time = "2025-08-12T06:59:49.895Z" }, - { url = "https://files.pythonhosted.org/packages/fe/02/c5e3bc518655d714622bec87d83db9cdba1cd0619a4a04e2109751c4f47f/sentencepiece-0.2.1-cp312-cp312-win_arm64.whl", hash = "sha256:daeb5e9e9fcad012324807856113708614d534f596d5008638eb9b40112cd9e4", size = 1033923, upload-time = "2025-08-12T06:59:51.952Z" }, ] [[package]] @@ -3238,9 +3541,9 @@ name = "sentry-sdk" version = "3.0.0a6" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "certifi" }, - { name = "opentelemetry-sdk" }, - { name = "urllib3" }, + { name = "certifi", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "opentelemetry-sdk", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "urllib3", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, ] sdist = { url = "https://files.pythonhosted.org/packages/24/68/d2bebd64ef6f74f6925f8466fae251425836ea35317ecaf0eb17295c9e9c/sentry_sdk-3.0.0a6.tar.gz", hash = "sha256:cc1eb071dd581a490a335f9bfc2065880f4dc264b511987709d2c5ced5d377f5", size = 331068, upload-time = "2025-09-01T12:49:32.075Z" } wheels = [ @@ -3253,16 +3556,10 @@ version = "1.3.7" source = { registry = "https://pypi.org/simple" } sdist = { url = "https://files.pythonhosted.org/packages/8d/48/49393a96a2eef1ab418b17475fb92b8fcfad83d099e678751b05472e69de/setproctitle-1.3.7.tar.gz", hash = "sha256:bc2bc917691c1537d5b9bca1468437176809c7e11e5694ca79a9ca12345dcb9e", size = 27002, upload-time = "2025-09-05T12:51:25.278Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/fb/f0/2dc88e842077719d7384d86cc47403e5102810492b33680e7dadcee64cd8/setproctitle-1.3.7-cp312-cp312-macosx_10_13_universal2.whl", hash = "sha256:2dc99aec591ab6126e636b11035a70991bc1ab7a261da428491a40b84376654e", size = 18049, upload-time = "2025-09-05T12:49:36.241Z" }, - { url = "https://files.pythonhosted.org/packages/f0/b4/50940504466689cda65680c9e9a1e518e5750c10490639fa687489ac7013/setproctitle-1.3.7-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:cdd8aa571b7aa39840fdbea620e308a19691ff595c3a10231e9ee830339dd798", size = 13079, upload-time = "2025-09-05T12:49:38.088Z" }, { url = "https://files.pythonhosted.org/packages/d0/99/71630546b9395b095f4082be41165d1078204d1696c2d9baade3de3202d0/setproctitle-1.3.7-cp312-cp312-manylinux1_x86_64.manylinux_2_28_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:2906b6c7959cdb75f46159bf0acd8cc9906cf1361c9e1ded0d065fe8f9039629", size = 32932, upload-time = "2025-09-05T12:49:39.271Z" }, { url = "https://files.pythonhosted.org/packages/50/22/cee06af4ffcfb0e8aba047bd44f5262e644199ae7527ae2c1f672b86495c/setproctitle-1.3.7-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:6915964a6dda07920a1159321dcd6d94fc7fc526f815ca08a8063aeca3c204f1", size = 33736, upload-time = "2025-09-05T12:49:40.565Z" }, - { url = "https://files.pythonhosted.org/packages/5c/00/a5949a8bb06ef5e7df214fc393bb2fb6aedf0479b17214e57750dfdd0f24/setproctitle-1.3.7-cp312-cp312-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:cff72899861c765bd4021d1ff1c68d60edc129711a2fdba77f9cb69ef726a8b6", size = 35605, upload-time = "2025-09-05T12:49:42.362Z" }, { url = "https://files.pythonhosted.org/packages/b0/3a/50caca532a9343828e3bf5778c7a84d6c737a249b1796d50dd680290594d/setproctitle-1.3.7-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:b7cb05bd446687ff816a3aaaf831047fc4c364feff7ada94a66024f1367b448c", size = 33143, upload-time = "2025-09-05T12:49:43.515Z" }, - { url = "https://files.pythonhosted.org/packages/ca/14/b843a251296ce55e2e17c017d6b9f11ce0d3d070e9265de4ecad948b913d/setproctitle-1.3.7-cp312-cp312-musllinux_1_2_ppc64le.whl", hash = "sha256:3a57b9a00de8cae7e2a1f7b9f0c2ac7b69372159e16a7708aa2f38f9e5cc987a", size = 34434, upload-time = "2025-09-05T12:49:45.31Z" }, { url = "https://files.pythonhosted.org/packages/c8/b7/06145c238c0a6d2c4bc881f8be230bb9f36d2bf51aff7bddcb796d5eed67/setproctitle-1.3.7-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:d8828b356114f6b308b04afe398ed93803d7fca4a955dd3abe84430e28d33739", size = 32795, upload-time = "2025-09-05T12:49:46.419Z" }, - { url = "https://files.pythonhosted.org/packages/ef/dc/ef76a81fac9bf27b84ed23df19c1f67391a753eed6e3c2254ebcb5133f56/setproctitle-1.3.7-cp312-cp312-win32.whl", hash = "sha256:b0304f905efc845829ac2bc791ddebb976db2885f6171f4a3de678d7ee3f7c9f", size = 12552, upload-time = "2025-09-05T12:49:47.635Z" }, - { url = "https://files.pythonhosted.org/packages/e2/5b/a9fe517912cd6e28cf43a212b80cb679ff179a91b623138a99796d7d18a0/setproctitle-1.3.7-cp312-cp312-win_amd64.whl", hash = "sha256:9888ceb4faea3116cf02a920ff00bfbc8cc899743e4b4ac914b03625bdc3c300", size = 13247, upload-time = "2025-09-05T12:49:49.16Z" }, ] [[package]] @@ -3283,15 +3580,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/e0/f9/0595336914c5619e5f28a1fb793285925a8cd4b432c9da0a987836c7f822/shellingham-1.5.4-py2.py3-none-any.whl", hash = "sha256:7ecfff8f2fd72616f7481040475a65b2bf8af90a56c89140852d1120324e8686", size = 9755, upload-time = "2023-10-24T04:13:38.866Z" }, ] -[[package]] -name = "shtab" -version = "1.7.2" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/5a/3e/837067b970c1d2ffa936c72f384a63fdec4e186b74da781e921354a94024/shtab-1.7.2.tar.gz", hash = "sha256:8c16673ade76a2d42417f03e57acf239bfb5968e842204c17990cae357d07d6f", size = 45751, upload-time = "2025-04-12T20:28:03.271Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/74/03/3271b7bb470fbab4adf5bd30b0d32143909d96f3608d815b447357f47f2b/shtab-1.7.2-py3-none-any.whl", hash = "sha256:858a5805f6c137bb0cda4f282d27d08fd44ca487ab4a6a36d2a400263cd0b5c1", size = 14214, upload-time = "2025-04-12T20:28:01.82Z" }, -] - [[package]] name = "six" version = "1.17.0" @@ -3319,12 +3607,21 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/e9/44/75a9c9421471a6c4805dbf2356f7c181a29c1879239abab1ea2cc8f38b40/sniffio-1.3.1-py3-none-any.whl", hash = "sha256:2f6da418d1f1e0fddd844478f41680e794e6051915791a034ff65e5f100525a2", size = 10235, upload-time = "2024-02-25T23:20:01.196Z" }, ] +[[package]] +name = "soupsieve" +version = "2.8.3" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/7b/ae/2d9c981590ed9999a0d91755b47fc74f74de286b0f5cee14c9269041e6c4/soupsieve-2.8.3.tar.gz", hash = "sha256:3267f1eeea4251fb42728b6dfb746edc9acaffc4a45b27e19450b676586e8349", size = 118627, upload-time = "2026-01-20T04:27:02.457Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/46/2c/1462b1d0a634697ae9e55b3cecdcb64788e8b7d63f54d923fcd0bb140aed/soupsieve-2.8.3-py3-none-any.whl", hash = "sha256:ed64f2ba4eebeab06cc4962affce381647455978ffc1e36bb79a545b91f45a95", size = 37016, upload-time = "2026-01-20T04:27:01.012Z" }, +] + [[package]] name = "sse-starlette" version = "3.0.2" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "anyio" }, + { name = "anyio", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, ] sdist = { url = "https://files.pythonhosted.org/packages/42/6f/22ed6e33f8a9e76ca0a412405f31abb844b779d52c5f96660766edcd737c/sse_starlette-3.0.2.tar.gz", hash = "sha256:ccd60b5765ebb3584d0de2d7a6e4f745672581de4f5005ab31c3a25d10b52b3a", size = 20985, upload-time = "2025-07-27T09:07:44.565Z" } wheels = [ @@ -3336,9 +3633,9 @@ name = "stack-data" version = "0.6.3" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "asttokens" }, - { name = "executing" }, - { name = "pure-eval" }, + { name = "asttokens", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "executing", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "pure-eval", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, ] sdist = { url = "https://files.pythonhosted.org/packages/28/e3/55dcc2cfbc3ca9c29519eb6884dd1415ecb53b0e934862d3559ddcb7e20b/stack_data-0.6.3.tar.gz", hash = "sha256:836a778de4fec4dcd1dcd89ed8abff8a221f58308462e1c4aa2a3cf30148f0b9", size = 44707, upload-time = "2023-09-30T13:58:05.479Z" } wheels = [ @@ -3350,8 +3647,8 @@ name = "starlette" version = "0.50.0" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "anyio" }, - { name = "typing-extensions" }, + { name = "anyio", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "typing-extensions", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, ] sdist = { url = "https://files.pythonhosted.org/packages/ba/b8/73a0e6a6e079a9d9cfa64113d771e421640b6f679a52eeb9b32f72d871a1/starlette-0.50.0.tar.gz", hash = "sha256:a2a17b22203254bcbc2e1f926d2d55f3f9497f769416b3190768befe598fa3ca", size = 2646985, upload-time = "2025-11-01T15:25:27.516Z" } wheels = [ @@ -3367,18 +3664,55 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/0e/65/5e726c372da8a5e35022a94388b12252710aad0c2351699c3d76ae8dba78/supervisor-4.3.0-py2.py3-none-any.whl", hash = "sha256:0bcb763fddafba410f35cbde226aa7f8514b9fb82eb05a0c85f6588d1c13f8db", size = 320736, upload-time = "2025-08-23T18:25:00.767Z" }, ] +[[package]] +name = "swebench" +version = "4.1.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "beautifulsoup4", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "chardet", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "datasets", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "docker", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "ghapi", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "gitpython", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "modal", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "pre-commit", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "python-dotenv", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "requests", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "rich", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "tenacity", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "tqdm", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "unidiff", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/24/e1/c997299ad7bf088876d30398203aa1eed7dec897670dc1aa35b1d748ffcc/swebench-4.1.0.tar.gz", hash = "sha256:5aaa6a92c2db1aa64892d28a47483ca46a45a15cf1d2df673d7744f71811dc9a", size = 134341, upload-time = "2025-09-11T02:58:00.447Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/36/67/981d8b642ac3eac7c8a7b7832ff8b2fb74f96b28b5fcd9a8979879e5c46d/swebench-4.1.0-py3-none-any.whl", hash = "sha256:1243776f720047cc9e20a427f7a52b75c13a07abda6154fb60fe77f82ec8af57", size = 157231, upload-time = "2025-09-11T02:57:58.953Z" }, +] + [[package]] name = "sympy" version = "1.14.0" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "mpmath" }, + { name = "mpmath", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, ] sdist = { url = "https://files.pythonhosted.org/packages/83/d3/803453b36afefb7c2bb238361cd4ae6125a569b4db67cd9e79846ba2d68c/sympy-1.14.0.tar.gz", hash = "sha256:d3d3fe8df1e5a0b42f0e7bdf50541697dbe7d23746e894990c030e2b05e72517", size = 7793921, upload-time = "2025-04-27T18:05:01.611Z" } wheels = [ { url = "https://files.pythonhosted.org/packages/a2/09/77d55d46fd61b4a135c444fc97158ef34a095e5681d0a6c10b75bf356191/sympy-1.14.0-py3-none-any.whl", hash = "sha256:e091cc3e99d2141a0ba2847328f5479b05d94a6635cb96148ccb3f34671bd8f5", size = 6299353, upload-time = "2025-04-27T18:04:59.103Z" }, ] +[[package]] +name = "synchronicity" +version = "0.11.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "typing-extensions", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/a3/26/8874d34755691994266d4a844ba8d53d10c2690ec67f246ca4d6b6f34cbb/synchronicity-0.11.1.tar.gz", hash = "sha256:3628df9ab34bd7be89b729104114841c62612c5d5ec43b76f4b7b243185ec1a8", size = 58131, upload-time = "2025-12-19T18:28:42.291Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/3f/b9/71153db12f4ad029cfe9b7fbf9792ef3fc9ade4485d31a13470b52954e62/synchronicity-0.11.1-py3-none-any.whl", hash = "sha256:53959c7f8b9b852fb5ea4d3d290a47a04310ede483a4cf0f8452cb4b5fa09db2", size = 40399, upload-time = "2025-12-19T18:28:40.972Z" }, +] + [[package]] name = "tabulate" version = "0.9.0" @@ -3402,16 +3736,16 @@ name = "tensorboard" version = "2.20.0" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "absl-py" }, - { name = "grpcio" }, - { name = "markdown" }, - { name = "numpy" }, - { name = "packaging" }, - { name = "pillow" }, - { name = "protobuf" }, - { name = "setuptools" }, - { name = "tensorboard-data-server" }, - { name = "werkzeug" }, + { name = "absl-py", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "grpcio", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "markdown", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "numpy", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "packaging", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "pillow", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "protobuf", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "setuptools", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "tensorboard-data-server", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "werkzeug", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, ] wheels = [ { url = "https://files.pythonhosted.org/packages/9c/d9/a5db55f88f258ac669a92858b70a714bbbd5acd993820b41ec4a96a4d77f/tensorboard-2.20.0-py3-none-any.whl", hash = "sha256:9dc9f978cb84c0723acf9a345d96c184f0293d18f166bb8d59ee098e6cfaaba6", size = 5525680, upload-time = "2025-07-17T19:20:49.638Z" }, @@ -3423,38 +3757,19 @@ version = "0.7.2" source = { registry = "https://pypi.org/simple" } wheels = [ { url = "https://files.pythonhosted.org/packages/7a/13/e503968fefabd4c6b2650af21e110aa8466fe21432cd7c43a84577a89438/tensorboard_data_server-0.7.2-py3-none-any.whl", hash = "sha256:7e0610d205889588983836ec05dc098e80f97b7e7bbff7e994ebb78f578d0ddb", size = 2356, upload-time = "2023-10-23T21:23:32.16Z" }, - { url = "https://files.pythonhosted.org/packages/b7/85/dabeaf902892922777492e1d253bb7e1264cadce3cea932f7ff599e53fea/tensorboard_data_server-0.7.2-py3-none-macosx_10_9_x86_64.whl", hash = "sha256:9fe5d24221b29625dbc7328b0436ca7fc1c23de4acf4d272f1180856e32f9f60", size = 4823598, upload-time = "2023-10-23T21:23:33.714Z" }, { url = "https://files.pythonhosted.org/packages/73/c6/825dab04195756cf8ff2e12698f22513b3db2f64925bdd41671bfb33aaa5/tensorboard_data_server-0.7.2-py3-none-manylinux_2_31_x86_64.whl", hash = "sha256:ef687163c24185ae9754ed5650eb5bc4d84ff257aabdc33f0cc6f74d8ba54530", size = 6590363, upload-time = "2023-10-23T21:23:35.583Z" }, ] -[[package]] -name = "textarena" -version = "0.7.3" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "chess" }, - { name = "nltk" }, - { name = "openai" }, - { name = "python-dotenv" }, - { name = "requests" }, - { name = "rich" }, - { name = "websockets" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/5f/81/1a36dd198b837d4591fafe88c39aba1efcd482889ecb7dbd0e3083fcc163/textarena-0.7.3.tar.gz", hash = "sha256:ba3966926bf2b7e9a1f1ac18940a49c2370532d24348f6a13c26361b08d6a286", size = 902897, upload-time = "2025-07-31T05:57:04.13Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/df/e3/8d4d47ed0d59db74ac7f862e3b91e87a02b74882deaa68b573788afaaaf8/textarena-0.7.3-py3-none-any.whl", hash = "sha256:3c9d954b4422bca15b641b4defe5330d9fc847559c0aa7ddbad8dccde7dae807", size = 1007582, upload-time = "2025-07-31T05:57:00.253Z" }, -] - [[package]] name = "textual" version = "6.1.0" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "markdown-it-py", extra = ["linkify", "plugins"] }, - { name = "platformdirs" }, - { name = "pygments" }, - { name = "rich" }, - { name = "typing-extensions" }, + { name = "markdown-it-py", extra = ["linkify", "plugins"], marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "platformdirs", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "pygments", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "rich", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "typing-extensions", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, ] sdist = { url = "https://files.pythonhosted.org/packages/da/44/4b524b2f06e0fa6c4ede56a4e9af5edd5f3f83cf2eea5cb4fd0ce5bbe063/textual-6.1.0.tar.gz", hash = "sha256:cc89826ca2146c645563259320ca4ddc75d183c77afb7d58acdd46849df9144d", size = 1564786, upload-time = "2025-09-02T11:42:34.655Z" } wheels = [ @@ -3466,42 +3781,51 @@ name = "tiktoken" version = "0.11.0" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "regex" }, - { name = "requests" }, + { name = "regex", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "requests", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, ] sdist = { url = "https://files.pythonhosted.org/packages/a7/86/ad0155a37c4f310935d5ac0b1ccf9bdb635dcb906e0a9a26b616dd55825a/tiktoken-0.11.0.tar.gz", hash = "sha256:3c518641aee1c52247c2b97e74d8d07d780092af79d5911a6ab5e79359d9b06a", size = 37648, upload-time = "2025-08-08T23:58:08.495Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/e7/9e/eceddeffc169fc75fe0fd4f38471309f11cb1906f9b8aa39be4f5817df65/tiktoken-0.11.0-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:fd9e6b23e860973cf9526544e220b223c60badf5b62e80a33509d6d40e6c8f5d", size = 1055199, upload-time = "2025-08-08T23:57:45.076Z" }, - { url = "https://files.pythonhosted.org/packages/4f/cf/5f02bfefffdc6b54e5094d2897bc80efd43050e5b09b576fd85936ee54bf/tiktoken-0.11.0-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:6a76d53cee2da71ee2731c9caa747398762bda19d7f92665e882fef229cb0b5b", size = 996655, upload-time = "2025-08-08T23:57:46.304Z" }, { url = "https://files.pythonhosted.org/packages/65/8e/c769b45ef379bc360c9978c4f6914c79fd432400a6733a8afc7ed7b0726a/tiktoken-0.11.0-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:6ef72aab3ea240646e642413cb363b73869fed4e604dcfd69eec63dc54d603e8", size = 1128867, upload-time = "2025-08-08T23:57:47.438Z" }, { url = "https://files.pythonhosted.org/packages/d5/2d/4d77f6feb9292bfdd23d5813e442b3bba883f42d0ac78ef5fdc56873f756/tiktoken-0.11.0-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:7f929255c705efec7a28bf515e29dc74220b2f07544a8c81b8d69e8efc4578bd", size = 1183308, upload-time = "2025-08-08T23:57:48.566Z" }, { url = "https://files.pythonhosted.org/packages/7a/65/7ff0a65d3bb0fc5a1fb6cc71b03e0f6e71a68c5eea230d1ff1ba3fd6df49/tiktoken-0.11.0-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:61f1d15822e4404953d499fd1dcc62817a12ae9fb1e4898033ec8fe3915fdf8e", size = 1244301, upload-time = "2025-08-08T23:57:49.642Z" }, - { url = "https://files.pythonhosted.org/packages/f5/6e/5b71578799b72e5bdcef206a214c3ce860d999d579a3b56e74a6c8989ee2/tiktoken-0.11.0-cp312-cp312-win_amd64.whl", hash = "sha256:45927a71ab6643dfd3ef57d515a5db3d199137adf551f66453be098502838b0f", size = 884282, upload-time = "2025-08-08T23:57:50.759Z" }, +] + +[[package]] +name = "tilelang" +version = "0.1.9" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "apache-tvm-ffi", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "cloudpickle", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "ml-dtypes", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "numpy", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "psutil", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "torch", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "torch-c-dlpack-ext", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "tqdm", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "typing-extensions", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "z3-solver", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/56/70/5051f65821baa30a3d61fc48f8ba10c776490315e8c90f82559b92089756/tilelang-0.1.9.tar.gz", hash = "sha256:287f727c913bb648fcf6c1968809ba3390e55eeed257a5c6bb9a80bc05966af4", size = 93395292, upload-time = "2026-04-22T09:19:11.988Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/f7/8a/1cbeee79d62abaa02441c2d00621554e41aa62dbf3b94a4feb3867184b01/tilelang-0.1.9-cp38-abi3-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:4bbccfe9035aed775ffafb6dc25a5994504b24e2c5d95d0f39643edfafa7bf12", size = 45419374, upload-time = "2026-04-22T09:15:56.014Z" }, + { url = "https://files.pythonhosted.org/packages/c6/a7/f4bfb86f87e107703146e703204cec2c0eae2492b633e0052b0ace3febb6/tilelang-0.1.9-cp38-abi3-manylinux_2_34_aarch64.whl", hash = "sha256:77ab0ee2f40f66ea015b6b21426d482751e28cbc635ef9d1198cbd6502454a7c", size = 42110365, upload-time = "2026-04-22T09:17:18.292Z" }, ] [[package]] name = "tokenizers" -version = "0.22.0" +version = "0.22.2" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "huggingface-hub" }, + { name = "huggingface-hub", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/5e/b4/c1ce3699e81977da2ace8b16d2badfd42b060e7d33d75c4ccdbf9dc920fa/tokenizers-0.22.0.tar.gz", hash = "sha256:2e33b98525be8453f355927f3cab312c36cd3e44f4d7e9e97da2fa94d0a49dcb", size = 362771, upload-time = "2025-08-29T10:25:33.914Z" } +sdist = { url = "https://files.pythonhosted.org/packages/73/6f/f80cfef4a312e1fb34baf7d85c72d4411afde10978d4657f8cdd811d3ccc/tokenizers-0.22.2.tar.gz", hash = "sha256:473b83b915e547aa366d1eee11806deaf419e17be16310ac0a14077f1e28f917", size = 372115, upload-time = "2026-01-05T10:45:15.988Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/6d/b1/18c13648edabbe66baa85fe266a478a7931ddc0cd1ba618802eb7b8d9865/tokenizers-0.22.0-cp39-abi3-macosx_10_12_x86_64.whl", hash = "sha256:eaa9620122a3fb99b943f864af95ed14c8dfc0f47afa3b404ac8c16b3f2bb484", size = 3081954, upload-time = "2025-08-29T10:25:24.993Z" }, - { url = "https://files.pythonhosted.org/packages/c2/02/c3c454b641bd7c4f79e4464accfae9e7dfc913a777d2e561e168ae060362/tokenizers-0.22.0-cp39-abi3-macosx_11_0_arm64.whl", hash = "sha256:71784b9ab5bf0ff3075bceeb198149d2c5e068549c0d18fe32d06ba0deb63f79", size = 2945644, upload-time = "2025-08-29T10:25:23.405Z" }, - { url = "https://files.pythonhosted.org/packages/55/02/d10185ba2fd8c2d111e124c9d92de398aee0264b35ce433f79fb8472f5d0/tokenizers-0.22.0-cp39-abi3-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:ec5b71f668a8076802b0241a42387d48289f25435b86b769ae1837cad4172a17", size = 3254764, upload-time = "2025-08-29T10:25:12.445Z" }, - { url = "https://files.pythonhosted.org/packages/13/89/17514bd7ef4bf5bfff58e2b131cec0f8d5cea2b1c8ffe1050a2c8de88dbb/tokenizers-0.22.0-cp39-abi3-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:ea8562fa7498850d02a16178105b58803ea825b50dc9094d60549a7ed63654bb", size = 3161654, upload-time = "2025-08-29T10:25:15.493Z" }, - { url = "https://files.pythonhosted.org/packages/5a/d8/bac9f3a7ef6dcceec206e3857c3b61bb16c6b702ed7ae49585f5bd85c0ef/tokenizers-0.22.0-cp39-abi3-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:4136e1558a9ef2e2f1de1555dcd573e1cbc4a320c1a06c4107a3d46dc8ac6e4b", size = 3511484, upload-time = "2025-08-29T10:25:20.477Z" }, - { url = "https://files.pythonhosted.org/packages/aa/27/9c9800eb6763683010a4851db4d1802d8cab9cec114c17056eccb4d4a6e0/tokenizers-0.22.0-cp39-abi3-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:cdf5954de3962a5fd9781dc12048d24a1a6f1f5df038c6e95db328cd22964206", size = 3712829, upload-time = "2025-08-29T10:25:17.154Z" }, - { url = "https://files.pythonhosted.org/packages/10/e3/b1726dbc1f03f757260fa21752e1921445b5bc350389a8314dd3338836db/tokenizers-0.22.0-cp39-abi3-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:8337ca75d0731fc4860e6204cc24bb36a67d9736142aa06ed320943b50b1e7ed", size = 3408934, upload-time = "2025-08-29T10:25:18.76Z" }, - { url = "https://files.pythonhosted.org/packages/d4/61/aeab3402c26874b74bb67a7f2c4b569dde29b51032c5384db592e7b216f4/tokenizers-0.22.0-cp39-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:a89264e26f63c449d8cded9061adea7b5de53ba2346fc7e87311f7e4117c1cc8", size = 3345585, upload-time = "2025-08-29T10:25:22.08Z" }, - { url = "https://files.pythonhosted.org/packages/bc/d3/498b4a8a8764cce0900af1add0f176ff24f475d4413d55b760b8cdf00893/tokenizers-0.22.0-cp39-abi3-musllinux_1_2_aarch64.whl", hash = "sha256:790bad50a1b59d4c21592f9c3cf5e5cf9c3c7ce7e1a23a739f13e01fb1be377a", size = 9322986, upload-time = "2025-08-29T10:25:26.607Z" }, - { url = "https://files.pythonhosted.org/packages/a2/62/92378eb1c2c565837ca3cb5f9569860d132ab9d195d7950c1ea2681dffd0/tokenizers-0.22.0-cp39-abi3-musllinux_1_2_armv7l.whl", hash = "sha256:76cf6757c73a10ef10bf06fa937c0ec7393d90432f543f49adc8cab3fb6f26cb", size = 9276630, upload-time = "2025-08-29T10:25:28.349Z" }, - { url = "https://files.pythonhosted.org/packages/eb/f0/342d80457aa1cda7654327460f69db0d69405af1e4c453f4dc6ca7c4a76e/tokenizers-0.22.0-cp39-abi3-musllinux_1_2_i686.whl", hash = "sha256:1626cb186e143720c62c6c6b5371e62bbc10af60481388c0da89bc903f37ea0c", size = 9547175, upload-time = "2025-08-29T10:25:29.989Z" }, - { url = "https://files.pythonhosted.org/packages/14/84/8aa9b4adfc4fbd09381e20a5bc6aa27040c9c09caa89988c01544e008d18/tokenizers-0.22.0-cp39-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:da589a61cbfea18ae267723d6b029b84598dc8ca78db9951d8f5beff72d8507c", size = 9692735, upload-time = "2025-08-29T10:25:32.089Z" }, - { url = "https://files.pythonhosted.org/packages/bf/24/83ee2b1dc76bfe05c3142e7d0ccdfe69f0ad2f1ebf6c726cea7f0874c0d0/tokenizers-0.22.0-cp39-abi3-win32.whl", hash = "sha256:dbf9d6851bddae3e046fedfb166f47743c1c7bd11c640f0691dd35ef0bcad3be", size = 2471915, upload-time = "2025-08-29T10:25:36.411Z" }, - { url = "https://files.pythonhosted.org/packages/d1/9b/0e0bf82214ee20231845b127aa4a8015936ad5a46779f30865d10e404167/tokenizers-0.22.0-cp39-abi3-win_amd64.whl", hash = "sha256:c78174859eeaee96021f248a56c801e36bfb6bd5b067f2e95aa82445ca324f00", size = 2680494, upload-time = "2025-08-29T10:25:35.14Z" }, + { url = "https://files.pythonhosted.org/packages/d6/84/7990e799f1309a8b87af6b948f31edaa12a3ed22d11b352eaf4f4b2e5753/tokenizers-0.22.2-cp39-abi3-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:2249487018adec45d6e3554c71d46eb39fa8ea67156c640f7513eb26f318cec7", size = 3290736, upload-time = "2026-01-05T10:40:32.165Z" }, + { url = "https://files.pythonhosted.org/packages/2e/76/932be4b50ef6ccedf9d3c6639b056a967a86258c6d9200643f01269211ca/tokenizers-0.22.2-cp39-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:369cc9fc8cc10cb24143873a0d95438bb8ee257bb80c71989e3ee290e8d72c67", size = 3274982, upload-time = "2026-01-05T10:40:58.331Z" }, + { url = "https://files.pythonhosted.org/packages/1d/28/5f9f5a4cc211b69e89420980e483831bcc29dade307955cc9dc858a40f01/tokenizers-0.22.2-cp39-abi3-musllinux_1_2_aarch64.whl", hash = "sha256:29c30b83d8dcd061078b05ae0cb94d3c710555fbb44861139f9f83dcca3dc3e4", size = 9478245, upload-time = "2026-01-05T10:41:04.053Z" }, + { url = "https://files.pythonhosted.org/packages/05/a1/d62dfe7376beaaf1394917e0f8e93ee5f67fea8fcf4107501db35996586b/tokenizers-0.22.2-cp39-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:38337540fbbddff8e999d59970f3c6f35a82de10053206a7562f1ea02d046fa5", size = 10033429, upload-time = "2026-01-05T10:45:14.333Z" }, ] [[package]] @@ -3519,16 +3843,10 @@ version = "2.2.1" source = { registry = "https://pypi.org/simple" } sdist = { url = "https://files.pythonhosted.org/packages/18/87/302344fed471e44a87289cf4967697d07e532f2421fdaf868a303cbae4ff/tomli-2.2.1.tar.gz", hash = "sha256:cd45e1dc79c835ce60f7404ec8119f2eb06d38b1deba146f07ced3bbc44505ff", size = 17175, upload-time = "2024-11-27T22:38:36.873Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/52/e1/f8af4c2fcde17500422858155aeb0d7e93477a0d59a98e56cbfe75070fd0/tomli-2.2.1-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:4a8f6e44de52d5e6c657c9fe83b562f5f4256d8ebbfe4ff922c495620a7f6cea", size = 132762, upload-time = "2024-11-27T22:38:07.731Z" }, - { url = "https://files.pythonhosted.org/packages/03/b8/152c68bb84fc00396b83e7bbddd5ec0bd3dd409db4195e2a9b3e398ad2e3/tomli-2.2.1-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:8d57ca8095a641b8237d5b079147646153d22552f1c637fd3ba7f4b0b29167a8", size = 123453, upload-time = "2024-11-27T22:38:09.384Z" }, { url = "https://files.pythonhosted.org/packages/c8/d6/fc9267af9166f79ac528ff7e8c55c8181ded34eb4b0e93daa767b8841573/tomli-2.2.1-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:4e340144ad7ae1533cb897d406382b4b6fede8890a03738ff1683af800d54192", size = 233486, upload-time = "2024-11-27T22:38:10.329Z" }, { url = "https://files.pythonhosted.org/packages/5c/51/51c3f2884d7bab89af25f678447ea7d297b53b5a3b5730a7cb2ef6069f07/tomli-2.2.1-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:db2b95f9de79181805df90bedc5a5ab4c165e6ec3fe99f970d0e302f384ad222", size = 242349, upload-time = "2024-11-27T22:38:11.443Z" }, - { url = "https://files.pythonhosted.org/packages/ab/df/bfa89627d13a5cc22402e441e8a931ef2108403db390ff3345c05253935e/tomli-2.2.1-cp312-cp312-manylinux_2_5_i686.manylinux1_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:40741994320b232529c802f8bc86da4e1aa9f413db394617b9a256ae0f9a7f77", size = 252159, upload-time = "2024-11-27T22:38:13.099Z" }, { url = "https://files.pythonhosted.org/packages/9e/6e/fa2b916dced65763a5168c6ccb91066f7639bdc88b48adda990db10c8c0b/tomli-2.2.1-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:400e720fe168c0f8521520190686ef8ef033fb19fc493da09779e592861b78c6", size = 237243, upload-time = "2024-11-27T22:38:14.766Z" }, - { url = "https://files.pythonhosted.org/packages/b4/04/885d3b1f650e1153cbb93a6a9782c58a972b94ea4483ae4ac5cedd5e4a09/tomli-2.2.1-cp312-cp312-musllinux_1_2_i686.whl", hash = "sha256:02abe224de6ae62c19f090f68da4e27b10af2b93213d36cf44e6e1c5abd19fdd", size = 259645, upload-time = "2024-11-27T22:38:15.843Z" }, { url = "https://files.pythonhosted.org/packages/9c/de/6b432d66e986e501586da298e28ebeefd3edc2c780f3ad73d22566034239/tomli-2.2.1-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:b82ebccc8c8a36f2094e969560a1b836758481f3dc360ce9a3277c65f374285e", size = 244584, upload-time = "2024-11-27T22:38:17.645Z" }, - { url = "https://files.pythonhosted.org/packages/1c/9a/47c0449b98e6e7d1be6cbac02f93dd79003234ddc4aaab6ba07a9a7482e2/tomli-2.2.1-cp312-cp312-win32.whl", hash = "sha256:889f80ef92701b9dbb224e49ec87c645ce5df3fa2cc548664eb8a25e03127a98", size = 98875, upload-time = "2024-11-27T22:38:19.159Z" }, - { url = "https://files.pythonhosted.org/packages/ef/60/9b9638f081c6f1261e2688bd487625cd1e660d0a85bd469e91d8db969734/tomli-2.2.1-cp312-cp312-win_amd64.whl", hash = "sha256:7fc04e92e1d624a4a63c76474610238576942d6b8950a2d7f908a340494e67e4", size = 109418, upload-time = "2024-11-27T22:38:20.064Z" }, { url = "https://files.pythonhosted.org/packages/6e/c2/61d3e0f47e2b74ef40a68b9e6ad5984f6241a942f7cd3bbfbdbd03861ea9/tomli-2.2.1-py3-none-any.whl", hash = "sha256:cb55c73c5f4408779d0cf3eef9f762b9c9f147a77de7b258bef0a5628adc85cc", size = 14257, upload-time = "2024-11-27T22:38:35.385Z" }, ] @@ -3543,51 +3861,49 @@ wheels = [ [[package]] name = "torch" -version = "2.9.0+cu128" -source = { registry = "https://download.pytorch.org/whl/test/cu128" } +version = "2.11.0+cu128" +source = { registry = "https://download.pytorch.org/whl/cu128" } dependencies = [ - { name = "filelock" }, - { name = "fsspec" }, - { name = "jinja2" }, - { name = "networkx" }, - { name = "nvidia-cublas-cu12", marker = "sys_platform == 'linux'" }, - { name = "nvidia-cuda-cupti-cu12", marker = "sys_platform == 'linux'" }, - { name = "nvidia-cuda-nvrtc-cu12", marker = "sys_platform == 'linux'" }, - { name = "nvidia-cuda-runtime-cu12", marker = "sys_platform == 'linux'" }, - { name = "nvidia-cudnn-cu12", marker = "sys_platform == 'linux'" }, - { name = "nvidia-cufft-cu12", marker = "sys_platform == 'linux'" }, - { name = "nvidia-cufile-cu12", marker = "sys_platform == 'linux'" }, - { name = "nvidia-curand-cu12", marker = "sys_platform == 'linux'" }, - { name = "nvidia-cusolver-cu12", marker = "sys_platform == 'linux'" }, - { name = "nvidia-cusparse-cu12", marker = "sys_platform == 'linux'" }, - { name = "nvidia-cusparselt-cu12", marker = "sys_platform == 'linux'" }, - { name = "nvidia-nccl-cu12", marker = "sys_platform == 'linux'" }, - { name = "nvidia-nvjitlink-cu12", marker = "sys_platform == 'linux'" }, - { name = "nvidia-nvshmem-cu12", marker = "sys_platform == 'linux'" }, - { name = "nvidia-nvtx-cu12", marker = "sys_platform == 'linux'" }, - { name = "setuptools" }, - { name = "sympy" }, - { name = "triton", version = "3.5.0", source = { registry = "https://pypi.org/simple" }, marker = "sys_platform == 'linux'" }, - { name = "typing-extensions" }, -] -wheels = [ - { url = "https://download.pytorch.org/whl/test/cu128/torch-2.9.0%2Bcu128-cp312-cp312-manylinux_2_28_aarch64.whl", hash = "sha256:7f6f6d12241fe64f202a43cecfe13031724c543cef88eb2d41076b38085337db" }, - { url = "https://download.pytorch.org/whl/test/cu128/torch-2.9.0%2Bcu128-cp312-cp312-manylinux_2_28_x86_64.whl", hash = "sha256:357dd549f9d348c7dcdd97d2fdb3adf74dc7899a53255461e32330c37be3ebea" }, - { url = "https://download.pytorch.org/whl/test/cu128/torch-2.9.0%2Bcu128-cp312-cp312-win_amd64.whl", hash = "sha256:6c05c6244d5e5d33f0724aa5c9220bf8b93b4435f79f0b98a9ecc93a9a239e2e" }, + { name = "cuda-bindings", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "cuda-toolkit", extra = ["cublas", "cudart", "cufft", "cufile", "cupti", "curand", "cusolver", "cusparse", "nvjitlink", "nvrtc", "nvtx"], marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "filelock", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "fsspec", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "jinja2", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "networkx", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "nvidia-cudnn-cu12", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "nvidia-cusparselt-cu12", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "nvidia-nccl-cu12", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "nvidia-nvshmem-cu12", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "setuptools", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "sympy", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "triton", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "typing-extensions", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, +] +wheels = [ + { url = "https://download-r2.pytorch.org/whl/cu128/torch-2.11.0%2Bcu128-cp312-cp312-manylinux_2_28_aarch64.whl", hash = "sha256:9c8f38efee365cb9d334de8a83ce52fc7e5fc9e5a7b0853285efa1b69e00b0f2", upload-time = "2026-04-27T17:41:30Z" }, + { url = "https://download-r2.pytorch.org/whl/cu128/torch-2.11.0%2Bcu128-cp312-cp312-manylinux_2_28_x86_64.whl", hash = "sha256:d252cf975fb18c94a85336323ad425f473df56dab35a44b00399bd70c7a3b997", upload-time = "2026-04-27T17:42:06Z" }, ] [[package]] -name = "torchaudio" -version = "2.9.0" +name = "torch-c-dlpack-ext" +version = "0.1.5" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "torch" }, + { name = "torch", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/37/de/921b6491efce5c389a5ef9bbed3d2d6660005840dae488124173180859ab/torch_c_dlpack_ext-0.1.5.tar.gz", hash = "sha256:d06f0357d575d22a168cc77acb9020fc4bae30968ceb6718a055dcbe92bacabe", size = 12913, upload-time = "2026-01-12T11:25:08.484Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/87/06/8d760997307a5c3be4384424667bf31aae0a42060838c532c7d846516175/torch_c_dlpack_ext-0.1.5-cp312-cp312-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:3562ee411258676f9c38b8ad39306d1c8d027b6a86f6a87c920d2d009a9d1510", size = 443069, upload-time = "2026-01-12T11:24:45.451Z" }, + { url = "https://files.pythonhosted.org/packages/e2/79/a914539b4785f3e44f891aa012a886edb8bc10fe081c440981c57543ce21/torch_c_dlpack_ext-0.1.5-cp312-cp312-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:e6f9da4bb9af70e27facc777458be62e10dbbbddda7672d16138db0553c5a524", size = 897846, upload-time = "2026-01-12T11:24:48.168Z" }, ] + +[[package]] +name = "torchaudio" +version = "2.11.0+cu128" +source = { registry = "https://download.pytorch.org/whl/cu128" } wheels = [ - { url = "https://files.pythonhosted.org/packages/b7/63/3c0ede3aa3d19a8a6698ddd107fa88660549360b51bf8ce2717cd498d800/torchaudio-2.9.0-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:ab4cbcccfd873b0fb41fcb39c9869e59ef84bb95b093f6f58e2d05172a7500d2", size = 809116, upload-time = "2025-10-15T15:52:00.911Z" }, - { url = "https://files.pythonhosted.org/packages/be/d5/25e58745defe9d05893d3cba5c0e1a76aeaac503ac5ec4d9f83c871df71c/torchaudio-2.9.0-cp312-cp312-manylinux_2_28_aarch64.whl", hash = "sha256:7f93388b6e536c14d6015b6f75277a8b45efc532f61b35adc1ed06c98a86003e", size = 476020, upload-time = "2025-10-15T15:51:59.967Z" }, - { url = "https://files.pythonhosted.org/packages/f0/9c/58b8b49dfba2ae85e41ca86b0c52de45bbbea01987490de219c99c523a58/torchaudio-2.9.0-cp312-cp312-manylinux_2_28_x86_64.whl", hash = "sha256:508318a2130b40ad51378f90caf8727a4bd3ac2b296f2b90c900b44e6068a940", size = 2059901, upload-time = "2025-10-15T15:51:54.634Z" }, - { url = "https://files.pythonhosted.org/packages/d7/eb/58b05f75d12f69ccc460893a20c999da082e063082120ed06e05cca3a053/torchaudio-2.9.0-cp312-cp312-win_amd64.whl", hash = "sha256:82117e3a605f2959dc09b4cd8a11178d6e92727d5f85e5d4f9fe47502f84ee96", size = 665350, upload-time = "2025-10-15T15:52:08.384Z" }, + { url = "https://download-r2.pytorch.org/whl/cu128/torchaudio-2.11.0%2Bcu128-cp312-cp312-manylinux_2_28_aarch64.whl", hash = "sha256:d055c9ca9d4f0ffcbaa0fc22138bafd675256de392bdaadde00d797faa90ca56", upload-time = "2026-03-23T15:50:23Z" }, + { url = "https://download-r2.pytorch.org/whl/cu128/torchaudio-2.11.0%2Bcu128-cp312-cp312-manylinux_2_28_x86_64.whl", hash = "sha256:78b86a17f164bdaabdcee93fdfde2587fc43b9ebf15cd61dcf730b4f8615176b", upload-time = "2026-03-23T15:50:23Z" }, ] [[package]] @@ -3595,9 +3911,9 @@ name = "torchdata" version = "0.11.0" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "requests" }, - { name = "torch" }, - { name = "urllib3" }, + { name = "requests", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "torch", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "urllib3", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, ] wheels = [ { url = "https://files.pythonhosted.org/packages/95/d4/af694ef718aedbe95a72760ab9ff7a6a7a44ace2d7f70c27bfeb67c5c503/torchdata-0.11.0-py3-none-any.whl", hash = "sha256:52b940fbbe0e00fb21cabddf528449d1bec5bfb0d0823b7487b15f951658ee33", size = 61968, upload-time = "2025-02-20T22:26:30.666Z" }, @@ -3608,29 +3924,27 @@ name = "torchtitan" version = "0.1.0" source = { git = "https://github.com/pytorch/torchtitan?rev=a1fdd7e#a1fdd7e43694bbfeff5d6ad8ac738c067bb90d41" } dependencies = [ - { name = "datasets" }, - { name = "fsspec" }, - { name = "tensorboard" }, - { name = "tokenizers" }, - { name = "tomli" }, - { name = "torchdata" }, - { name = "tyro" }, + { name = "datasets", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "fsspec", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "tensorboard", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "tokenizers", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "tomli", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "torchdata", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "tyro", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, ] [[package]] name = "torchvision" -version = "0.24.0" -source = { registry = "https://pypi.org/simple" } +version = "0.26.0+cu128" +source = { registry = "https://download.pytorch.org/whl/cu128" } dependencies = [ - { name = "numpy" }, - { name = "pillow" }, - { name = "torch" }, + { name = "numpy", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "pillow", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "torch", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, ] wheels = [ - { url = "https://files.pythonhosted.org/packages/47/ef/81e4e69e02e2c4650b30e8c11c8974f946682a30e0ab7e9803a831beff76/torchvision-0.24.0-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:c61d40bcd2e2451e932902a702ad495ba1ec6f279e90b1e15cef2bb55dc911e2", size = 1891726, upload-time = "2025-10-15T15:51:16.977Z" }, - { url = "https://files.pythonhosted.org/packages/00/7b/e3809b3302caea9a12c13f3adebe4fef127188438e719fd6c8dc93db1da6/torchvision-0.24.0-cp312-cp312-manylinux_2_28_aarch64.whl", hash = "sha256:b0531d1483fc322d7da0d83be52f0df860a75114ab87dbeeb9de765feaeda843", size = 2419495, upload-time = "2025-10-15T15:51:11.885Z" }, - { url = "https://files.pythonhosted.org/packages/7e/e6/7324ead6793075a8c75c56abeed1236d1750de16a5613cfe2ddad164a92a/torchvision-0.24.0-cp312-cp312-manylinux_2_28_x86_64.whl", hash = "sha256:26b9dd9c083f8e5f7ac827de6d5b88c615d9c582dc87666770fbdf16887e4c25", size = 8050480, upload-time = "2025-10-15T15:51:24.012Z" }, - { url = "https://files.pythonhosted.org/packages/3e/ad/3c56fcd2a0d6e8afa80e115b5ade4302232ec99655220a51d05709819523/torchvision-0.24.0-cp312-cp312-win_amd64.whl", hash = "sha256:060b7c50ed4b3fb0316b08e2e31bfd874ec2f63ef5ae02f81e54341ca4e88703", size = 4292225, upload-time = "2025-10-15T15:51:27.699Z" }, + { url = "https://download-r2.pytorch.org/whl/cu128/torchvision-0.26.0%2Bcu128-cp312-cp312-manylinux_2_28_aarch64.whl", hash = "sha256:63e35234aed13b6edda37056f417b5c281249669db631e706811917af36b21d7", upload-time = "2026-04-09T23:21:35Z" }, + { url = "https://download-r2.pytorch.org/whl/cu128/torchvision-0.26.0%2Bcu128-cp312-cp312-manylinux_2_28_x86_64.whl", hash = "sha256:ccf26b4b659cfce6f2208cb8326071d51c70219a34856dfdf468d1e19af52c0d", upload-time = "2026-03-23T15:36:22Z" }, ] [[package]] @@ -3639,26 +3953,16 @@ version = "6.5.2" source = { registry = "https://pypi.org/simple" } sdist = { url = "https://files.pythonhosted.org/packages/09/ce/1eb500eae19f4648281bb2186927bb062d2438c2e5093d1360391afd2f90/tornado-6.5.2.tar.gz", hash = "sha256:ab53c8f9a0fa351e2c0741284e06c7a45da86afb544133201c5cc8578eb076a0", size = 510821, upload-time = "2025-08-08T18:27:00.78Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/f6/48/6a7529df2c9cc12efd2e8f5dd219516184d703b34c06786809670df5b3bd/tornado-6.5.2-cp39-abi3-macosx_10_9_universal2.whl", hash = "sha256:2436822940d37cde62771cff8774f4f00b3c8024fe482e16ca8387b8a2724db6", size = 442563, upload-time = "2025-08-08T18:26:42.945Z" }, - { url = "https://files.pythonhosted.org/packages/f2/b5/9b575a0ed3e50b00c40b08cbce82eb618229091d09f6d14bce80fc01cb0b/tornado-6.5.2-cp39-abi3-macosx_10_9_x86_64.whl", hash = "sha256:583a52c7aa94ee046854ba81d9ebb6c81ec0fd30386d96f7640c96dad45a03ef", size = 440729, upload-time = "2025-08-08T18:26:44.473Z" }, { url = "https://files.pythonhosted.org/packages/1b/4e/619174f52b120efcf23633c817fd3fed867c30bff785e2cd5a53a70e483c/tornado-6.5.2-cp39-abi3-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:b0fe179f28d597deab2842b86ed4060deec7388f1fd9c1b4a41adf8af058907e", size = 444295, upload-time = "2025-08-08T18:26:46.021Z" }, - { url = "https://files.pythonhosted.org/packages/95/fa/87b41709552bbd393c85dd18e4e3499dcd8983f66e7972926db8d96aa065/tornado-6.5.2-cp39-abi3-manylinux_2_5_i686.manylinux1_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:b186e85d1e3536d69583d2298423744740986018e393d0321df7340e71898882", size = 443644, upload-time = "2025-08-08T18:26:47.625Z" }, { url = "https://files.pythonhosted.org/packages/f9/41/fb15f06e33d7430ca89420283a8762a4e6b8025b800ea51796ab5e6d9559/tornado-6.5.2-cp39-abi3-manylinux_2_5_x86_64.manylinux1_x86_64.manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:e792706668c87709709c18b353da1f7662317b563ff69f00bab83595940c7108", size = 443878, upload-time = "2025-08-08T18:26:50.599Z" }, { url = "https://files.pythonhosted.org/packages/11/92/fe6d57da897776ad2e01e279170ea8ae726755b045fe5ac73b75357a5a3f/tornado-6.5.2-cp39-abi3-musllinux_1_2_aarch64.whl", hash = "sha256:06ceb1300fd70cb20e43b1ad8aaee0266e69e7ced38fa910ad2e03285009ce7c", size = 444549, upload-time = "2025-08-08T18:26:51.864Z" }, - { url = "https://files.pythonhosted.org/packages/9b/02/c8f4f6c9204526daf3d760f4aa555a7a33ad0e60843eac025ccfd6ff4a93/tornado-6.5.2-cp39-abi3-musllinux_1_2_i686.whl", hash = "sha256:74db443e0f5251be86cbf37929f84d8c20c27a355dd452a5cfa2aada0d001ec4", size = 443973, upload-time = "2025-08-08T18:26:53.625Z" }, { url = "https://files.pythonhosted.org/packages/ae/2d/f5f5707b655ce2317190183868cd0f6822a1121b4baeae509ceb9590d0bd/tornado-6.5.2-cp39-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:b5e735ab2889d7ed33b32a459cac490eda71a1ba6857b0118de476ab6c366c04", size = 443954, upload-time = "2025-08-08T18:26:55.072Z" }, - { url = "https://files.pythonhosted.org/packages/e8/59/593bd0f40f7355806bf6573b47b8c22f8e1374c9b6fd03114bd6b7a3dcfd/tornado-6.5.2-cp39-abi3-win32.whl", hash = "sha256:c6f29e94d9b37a95013bb669616352ddb82e3bfe8326fccee50583caebc8a5f0", size = 445023, upload-time = "2025-08-08T18:26:56.677Z" }, - { url = "https://files.pythonhosted.org/packages/c7/2a/f609b420c2f564a748a2d80ebfb2ee02a73ca80223af712fca591386cafb/tornado-6.5.2-cp39-abi3-win_amd64.whl", hash = "sha256:e56a5af51cc30dd2cae649429af65ca2f6571da29504a07995175df14c18f35f", size = 445427, upload-time = "2025-08-08T18:26:57.91Z" }, - { url = "https://files.pythonhosted.org/packages/5e/4f/e1f65e8f8c76d73658b33d33b81eed4322fb5085350e4328d5c956f0c8f9/tornado-6.5.2-cp39-abi3-win_arm64.whl", hash = "sha256:d6c33dc3672e3a1f3618eb63b7ef4683a7688e7b9e6e8f0d9aa5726360a004af", size = 444456, upload-time = "2025-08-08T18:26:59.207Z" }, ] [[package]] name = "tqdm" version = "4.67.1" source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "colorama", marker = "sys_platform == 'win32'" }, -] sdist = { url = "https://files.pythonhosted.org/packages/a8/4b/29b4ef32e036bb34e4ab51796dd745cdba7ed47ad142a9f4a1eb8e0c744d/tqdm-4.67.1.tar.gz", hash = "sha256:f8aef9c52c08c13a65f30ea34f4e5aac3fd1a34959879d7e59e63027286627f2", size = 169737, upload-time = "2024-11-24T20:12:22.481Z" } wheels = [ { url = "https://files.pythonhosted.org/packages/d0/30/dc54f88dd4a2b5dc8a0279bdd7270e735851848b762aeb1c1184ed1f6b14/tqdm-4.67.1-py3-none-any.whl", hash = "sha256:26445eca388f82e72884e0d580d5464cd801a3ea01e63e5601bdff9ba6a48de2", size = 78540, upload-time = "2024-11-24T20:12:19.698Z" }, @@ -3675,46 +3979,27 @@ wheels = [ [[package]] name = "transformers" -version = "4.56.1" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "filelock" }, - { name = "huggingface-hub" }, - { name = "numpy" }, - { name = "packaging" }, - { name = "pyyaml" }, - { name = "regex" }, - { name = "requests" }, - { name = "safetensors" }, - { name = "tokenizers" }, - { name = "tqdm" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/89/21/dc88ef3da1e49af07ed69386a11047a31dcf1aaf4ded3bc4b173fbf94116/transformers-4.56.1.tar.gz", hash = "sha256:0d88b1089a563996fc5f2c34502f10516cad3ea1aa89f179f522b54c8311fe74", size = 9855473, upload-time = "2025-09-04T20:47:13.14Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/71/7c/283c3dd35e00e22a7803a0b2a65251347b745474a82399be058bde1c9f15/transformers-4.56.1-py3-none-any.whl", hash = "sha256:1697af6addfb6ddbce9618b763f4b52d5a756f6da4899ffd1b4febf58b779248", size = 11608197, upload-time = "2025-09-04T20:47:04.895Z" }, -] - -[[package]] -name = "triton" -version = "3.3.1" -source = { registry = "https://pypi.org/simple" } -resolution-markers = [ - "sys_platform != 'linux'", -] +version = "5.5.0" +source = { git = "https://github.com/huggingface/transformers.git?rev=c1c3424#c1c34249fa27deefbd4a377dfbf883a39baf5c6d" } dependencies = [ - { name = "setuptools", marker = "sys_platform != 'linux'" }, + { name = "huggingface-hub", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "numpy", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "packaging", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "pyyaml", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "regex", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "safetensors", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "tokenizers", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "tqdm", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "typer", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, ] [[package]] name = "triton" -version = "3.5.0" +version = "3.6.0" source = { registry = "https://pypi.org/simple" } -resolution-markers = [ - "sys_platform == 'linux'", -] wheels = [ - { url = "https://files.pythonhosted.org/packages/87/9b/30988039e1e84df7554fba24e6a734d2d0e847af33cabdf9b532b3c51456/triton-3.5.0-cp312-cp312-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:7da21fccceafc163e3a5e857abe34351ef76345af06cabf9637a914742671f0b", size = 159946647, upload-time = "2025-10-15T19:15:56.325Z" }, - { url = "https://files.pythonhosted.org/packages/f5/3a/e991574f3102147b642e49637e0281e9bb7c4ba254edb2bab78247c85e01/triton-3.5.0-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:c9e71db82261c4ffa3921cd050cd5faa18322d2d405c30eb56084afaff3b0833", size = 170476535, upload-time = "2025-10-13T16:38:05.18Z" }, + { url = "https://files.pythonhosted.org/packages/17/5d/08201db32823bdf77a0e2b9039540080b2e5c23a20706ddba942924ebcd6/triton-3.6.0-cp312-cp312-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:374f52c11a711fd062b4bfbb201fd9ac0a5febd28a96fb41b4a0f51dde3157f4", size = 176128243, upload-time = "2026-01-20T16:16:07.857Z" }, + { url = "https://files.pythonhosted.org/packages/ab/a8/cdf8b3e4c98132f965f88c2313a4b493266832ad47fb52f23d14d4f86bb5/triton-3.6.0-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:74caf5e34b66d9f3a429af689c1c7128daba1d8208df60e81106b115c00d6fca", size = 188266850, upload-time = "2026-01-20T16:00:43.041Z" }, ] [[package]] @@ -3722,7 +4007,7 @@ name = "typeguard" version = "4.4.4" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "typing-extensions" }, + { name = "typing-extensions", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, ] sdist = { url = "https://files.pythonhosted.org/packages/c7/68/71c1a15b5f65f40e91b65da23b8224dad41349894535a97f63a52e462196/typeguard-4.4.4.tar.gz", hash = "sha256:3a7fd2dffb705d4d0efaed4306a704c89b9dee850b688f060a8b1615a79e5f74", size = 75203, upload-time = "2025-06-18T09:56:07.624Z" } wheels = [ @@ -3731,17 +4016,26 @@ wheels = [ [[package]] name = "typer" -version = "0.17.4" +version = "0.23.0" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "click" }, - { name = "rich" }, - { name = "shellingham" }, - { name = "typing-extensions" }, + { name = "annotated-doc", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "click", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "rich", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "shellingham", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/7e/e6/44e073787aa57cd71c151f44855232feb0f748428fd5242d7366e3c4ae8b/typer-0.23.0.tar.gz", hash = "sha256:d8378833e47ada5d3d093fa20c4c63427cc4e27127f6b349a6c359463087d8cc", size = 120181, upload-time = "2026-02-11T15:22:18.637Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/7a/ed/d6fca788b51d0d4640c4bc82d0e85bad4b49809bca36bf4af01b4dcb66a7/typer-0.23.0-py3-none-any.whl", hash = "sha256:79f4bc262b6c37872091072a3cb7cb6d7d79ee98c0c658b4364bdcde3c42c913", size = 56668, upload-time = "2026-02-11T15:22:21.075Z" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/92/e8/2a73ccf9874ec4c7638f172efc8972ceab13a0e3480b389d6ed822f7a822/typer-0.17.4.tar.gz", hash = "sha256:b77dc07d849312fd2bb5e7f20a7af8985c7ec360c45b051ed5412f64d8dc1580", size = 103734, upload-time = "2025-09-05T18:14:40.746Z" } + +[[package]] +name = "types-certifi" +version = "2021.10.8.3" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/52/68/943c3aeaf14624712a0357c4a67814dba5cea36d194f5c764dad7959a00c/types-certifi-2021.10.8.3.tar.gz", hash = "sha256:72cf7798d165bc0b76e1c10dd1ea3097c7063c42c21d664523b928e88b554a4f", size = 2095, upload-time = "2022-06-09T15:19:05.244Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/93/72/6b3e70d32e89a5cbb6a4513726c1ae8762165b027af569289e19ec08edd8/typer-0.17.4-py3-none-any.whl", hash = "sha256:015534a6edaa450e7007eba705d5c18c3349dcea50a6ad79a5ed530967575824", size = 46643, upload-time = "2025-09-05T18:14:39.166Z" }, + { url = "https://files.pythonhosted.org/packages/b5/63/2463d89481e811f007b0e1cd0a91e52e141b47f9de724d20db7b861dcfec/types_certifi-2021.10.8.3-py3-none-any.whl", hash = "sha256:b2d1e325e69f71f7c78e5943d410e650b4707bb0ef32e4ddf3da37f54176e88a", size = 2136, upload-time = "2022-06-09T15:19:03.127Z" }, ] [[package]] @@ -3749,13 +4043,22 @@ name = "types-requests" version = "2.32.4.20250809" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "urllib3" }, + { name = "urllib3", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, ] sdist = { url = "https://files.pythonhosted.org/packages/ed/b0/9355adb86ec84d057fea765e4c49cce592aaf3d5117ce5609a95a7fc3dac/types_requests-2.32.4.20250809.tar.gz", hash = "sha256:d8060de1c8ee599311f56ff58010fb4902f462a1470802cf9f6ed27bc46c4df3", size = 23027, upload-time = "2025-08-09T03:17:10.664Z" } wheels = [ { url = "https://files.pythonhosted.org/packages/2b/6f/ec0012be842b1d888d46884ac5558fd62aeae1f0ec4f7a581433d890d4b5/types_requests-2.32.4.20250809-py3-none-any.whl", hash = "sha256:f73d1832fb519ece02c85b1f09d5f0dd3108938e7d47e7f94bbfa18a6782b163", size = 20644, upload-time = "2025-08-09T03:17:09.716Z" }, ] +[[package]] +name = "types-toml" +version = "0.10.8.20240310" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/86/47/3e4c75042792bff8e90d7991aa5c51812cc668828cc6cce711e97f63a607/types-toml-0.10.8.20240310.tar.gz", hash = "sha256:3d41501302972436a6b8b239c850b26689657e25281b48ff0ec06345b8830331", size = 4392, upload-time = "2024-03-10T02:18:37.518Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/da/a2/d32ab58c0b216912638b140ab2170ee4b8644067c293b170e19fba340ccc/types_toml-0.10.8.20240310-py3-none-any.whl", hash = "sha256:627b47775d25fa29977d9c70dc0cbab3f314f32c8d8d0c012f2ef5de7aaec05d", size = 4777, upload-time = "2024-03-10T02:18:36.568Z" }, +] + [[package]] name = "typing-extensions" version = "4.15.0" @@ -3770,7 +4073,7 @@ name = "typing-inspection" version = "0.4.2" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "typing-extensions" }, + { name = "typing-extensions", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, ] sdist = { url = "https://files.pythonhosted.org/packages/55/e3/70399cb7dd41c10ac53367ae42139cf4b1ca5f36bb3dc6c9d33acdb43655/typing_inspection-0.4.2.tar.gz", hash = "sha256:ba561c48a67c5958007083d386c3295464928b01faa735ab8547c5692e87f464", size = 75949, upload-time = "2025-10-01T02:14:41.687Z" } wheels = [ @@ -3779,19 +4082,16 @@ wheels = [ [[package]] name = "tyro" -version = "0.9.31" +version = "1.0.8" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "colorama", marker = "sys_platform == 'win32'" }, - { name = "docstring-parser" }, - { name = "rich" }, - { name = "shtab" }, - { name = "typeguard" }, - { name = "typing-extensions" }, + { name = "docstring-parser", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "typeguard", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "typing-extensions", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/b0/68/807bb16fdf4fe5cdfd70f6e90b0fd78f0d11bee815996ce5c9ae864b3ee1/tyro-0.9.31.tar.gz", hash = "sha256:e858b21f4d179db9fb52b292d3a544aca0bb439399bd59f1943cf96382f15a68", size = 319632, upload-time = "2025-08-30T21:17:35.674Z" } +sdist = { url = "https://files.pythonhosted.org/packages/9e/89/f4d76ba5ac4ca627c3a0413fdc8d90f202929c45d9c25deec29c9a86a0ae/tyro-1.0.8.tar.gz", hash = "sha256:c22bf238ce029b8cc262759129fccfd968300c3761f00ce186570af3891e14bb", size = 460626, upload-time = "2026-02-25T00:18:45.388Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/9b/c6/c808aed7b41ddb7b25427d53185c12df3daea76a295d061a3c58ac61a313/tyro-0.9.31-py3-none-any.whl", hash = "sha256:565305246fbfe7c288b7c44d5af58782a582b885c662ff223e9933779c1c10a3", size = 131681, upload-time = "2025-08-30T21:17:34.534Z" }, + { url = "https://files.pythonhosted.org/packages/61/7c/9c5a4ab9e24f1428fd5ab2728b7f68878f54d549a11fa431ead76d3d127b/tyro-1.0.8-py3-none-any.whl", hash = "sha256:abb6a4054f20616d9ca5181ba0c6ae0e35cf8f35ae0c78d3d4dbd1f678e0643d", size = 181935, upload-time = "2026-02-25T00:18:43.689Z" }, ] [[package]] @@ -3812,6 +4112,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/37/87/1f677586e8ac487e29672e4b17455758fce261de06a0d086167bb760361a/uc_micro_py-1.0.3-py3-none-any.whl", hash = "sha256:db1dffff340817673d7b466ec86114a9dc0e9d4d9b5ba229d9d60e5c12600cd5", size = 6229, upload-time = "2024-02-09T16:52:00.371Z" }, ] +[[package]] +name = "unidiff" +version = "0.7.5" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/a3/48/81be0ac96e423a877754153699731ef439fd7b80b4c8b5425c94ed079ebd/unidiff-0.7.5.tar.gz", hash = "sha256:2e5f0162052248946b9f0970a40e9e124236bf86c82b70821143a6fc1dea2574", size = 20931, upload-time = "2023-03-10T01:05:39.185Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/8a/54/57c411a6e8f7bd7848c8b66e4dcaffa586bf4c02e63f2280db0327a4e6eb/unidiff-0.7.5-py2.py3-none-any.whl", hash = "sha256:c93bf2265cc1ba2a520e415ab05da587370bc2a3ae9e0414329f54f0c2fc09e8", size = 14386, upload-time = "2023-03-10T01:05:36.594Z" }, +] + [[package]] name = "urllib3" version = "2.5.0" @@ -3826,8 +4135,8 @@ name = "uvicorn" version = "0.35.0" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "click" }, - { name = "h11" }, + { name = "click", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "h11", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, ] sdist = { url = "https://files.pythonhosted.org/packages/5e/42/e0e305207bb88c6b8d3061399c6a961ffe5fbb7e2aa63c9234df7259e9cd/uvicorn-0.35.0.tar.gz", hash = "sha256:bc662f087f7cf2ce11a1d7fd70b90c9f98ef2e2831556dd078d131b96cc94a01", size = 78473, upload-time = "2025-06-28T16:15:46.058Z" } wheels = [ @@ -3836,13 +4145,12 @@ wheels = [ [package.optional-dependencies] standard = [ - { name = "colorama", marker = "sys_platform == 'win32'" }, - { name = "httptools" }, - { name = "python-dotenv" }, - { name = "pyyaml" }, - { name = "uvloop", marker = "platform_python_implementation != 'PyPy' and sys_platform != 'cygwin' and sys_platform != 'win32'" }, - { name = "watchfiles" }, - { name = "websockets" }, + { name = "httptools", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "python-dotenv", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "pyyaml", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "uvloop", marker = "(platform_machine == 'aarch64' and platform_python_implementation != 'PyPy' and sys_platform == 'linux') or (platform_machine == 'x86_64' and platform_python_implementation != 'PyPy' and sys_platform == 'linux')" }, + { name = "watchfiles", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "websockets", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, ] [[package]] @@ -3851,8 +4159,6 @@ version = "0.21.0" source = { registry = "https://pypi.org/simple" } sdist = { url = "https://files.pythonhosted.org/packages/af/c0/854216d09d33c543f12a44b393c402e89a920b1a0a7dc634c42de91b9cf6/uvloop-0.21.0.tar.gz", hash = "sha256:3bf12b0fda68447806a7ad847bfa591613177275d35b6724b1ee573faa3704e3", size = 2492741, upload-time = "2024-10-14T23:38:35.489Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/8c/4c/03f93178830dc7ce8b4cdee1d36770d2f5ebb6f3d37d354e061eefc73545/uvloop-0.21.0-cp312-cp312-macosx_10_13_universal2.whl", hash = "sha256:359ec2c888397b9e592a889c4d72ba3d6befba8b2bb01743f72fffbde663b59c", size = 1471284, upload-time = "2024-10-14T23:37:47.833Z" }, - { url = "https://files.pythonhosted.org/packages/43/3e/92c03f4d05e50f09251bd8b2b2b584a2a7f8fe600008bcc4523337abe676/uvloop-0.21.0-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:f7089d2dc73179ce5ac255bdf37c236a9f914b264825fdaacaded6990a7fb4c2", size = 821349, upload-time = "2024-10-14T23:37:50.149Z" }, { url = "https://files.pythonhosted.org/packages/a6/ef/a02ec5da49909dbbfb1fd205a9a1ac4e88ea92dcae885e7c961847cd51e2/uvloop-0.21.0-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:baa4dcdbd9ae0a372f2167a207cd98c9f9a1ea1188a8a526431eef2f8116cc8d", size = 4580089, upload-time = "2024-10-14T23:37:51.703Z" }, { url = "https://files.pythonhosted.org/packages/06/a7/b4e6a19925c900be9f98bec0a75e6e8f79bb53bdeb891916609ab3958967/uvloop-0.21.0-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:86975dca1c773a2c9864f4c52c5a55631038e387b47eaf56210f873887b6c8dc", size = 4693770, upload-time = "2024-10-14T23:37:54.122Z" }, { url = "https://files.pythonhosted.org/packages/ce/0c/f07435a18a4b94ce6bd0677d8319cd3de61f3a9eeb1e5f8ab4e8b5edfcb3/uvloop-0.21.0-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:461d9ae6660fbbafedd07559c6a2e57cd553b34b0065b6550685f6653a98c1cb", size = 4451321, upload-time = "2024-10-14T23:37:55.766Z" }, @@ -3861,23 +4167,33 @@ wheels = [ [[package]] name = "verifiers" -version = "0.1.9.post3" -source = { git = "https://github.com/PrimeIntellect-ai/verifiers.git?rev=209774a#209774a22b6915eeff2091a1ce2782ecbe580ae7" } -dependencies = [ - { name = "datasets" }, - { name = "jinja2" }, - { name = "math-verify" }, - { name = "mcp" }, - { name = "nest-asyncio" }, - { name = "openai" }, - { name = "openai-agents" }, - { name = "prime-sandboxes" }, - { name = "pydantic" }, - { name = "requests" }, - { name = "rich" }, - { name = "tenacity" }, - { name = "textual" }, - { name = "wget" }, +version = "0.1.14" +source = { git = "https://github.com/PrimeIntellect-ai/verifiers.git?rev=aa428f3#aa428f3941ae35a7cf7c0dad7e60c7eca525bac6" } +dependencies = [ + { name = "aiolimiter", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "anthropic", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "datasets", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "gepa", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "httpx", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "jinja2", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "math-verify", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "mcp", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "msgpack", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "nest-asyncio", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "numpy", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "openai", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "openai-agents", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "prime-sandboxes", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "prime-tunnel", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "pydantic", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "pyzmq", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "regex", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "requests", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "rich", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "setproctitle", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "tenacity", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "textual", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "wget", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, ] [[package]] @@ -3885,9 +4201,9 @@ name = "virtualenv" version = "20.34.0" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "distlib" }, - { name = "filelock" }, - { name = "platformdirs" }, + { name = "distlib", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "filelock", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "platformdirs", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, ] sdist = { url = "https://files.pythonhosted.org/packages/1c/14/37fcdba2808a6c615681cd216fecae00413c9dab44fb2e57805ecf3eaee3/virtualenv-20.34.0.tar.gz", hash = "sha256:44815b2c9dee7ed86e387b842a84f20b93f7f417f95886ca1996a72a4138eb1a", size = 6003808, upload-time = "2025-08-13T14:24:07.464Z" } wheels = [ @@ -3896,72 +4212,383 @@ wheels = [ [[package]] name = "vllm" -version = "0.12.0" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "aiohttp" }, - { name = "anthropic" }, +version = "0.20.2+cu129" +source = { url = "https://github.com/vllm-project/vllm/releases/download/v0.20.2/vllm-0.20.2+cu129-cp38-abi3-manylinux_2_31_aarch64.whl" } +resolution-markers = [ + "platform_machine == 'aarch64' and sys_platform == 'linux'", +] +dependencies = [ + { name = "aiohttp", marker = "platform_machine == 'aarch64' and sys_platform == 'linux'" }, + { name = "anthropic", marker = "platform_machine == 'aarch64' and sys_platform == 'linux'" }, + { name = "apache-tvm-ffi", marker = "platform_machine == 'aarch64' and sys_platform == 'linux'" }, + { name = "blake3", marker = "platform_machine == 'aarch64' and sys_platform == 'linux'" }, + { name = "cachetools", marker = "platform_machine == 'aarch64' and sys_platform == 'linux'" }, + { name = "cbor2", marker = "platform_machine == 'aarch64' and sys_platform == 'linux'" }, + { name = "cloudpickle", marker = "platform_machine == 'aarch64' and sys_platform == 'linux'" }, + { name = "compressed-tensors", marker = "platform_machine == 'aarch64' and sys_platform == 'linux'" }, + { name = "depyf", marker = "platform_machine == 'aarch64' and sys_platform == 'linux'" }, + { name = "diskcache", marker = "platform_machine == 'aarch64' and sys_platform == 'linux'" }, + { name = "einops", marker = "platform_machine == 'aarch64' and sys_platform == 'linux'" }, + { name = "fastapi", extra = ["standard"], marker = "platform_machine == 'aarch64' and sys_platform == 'linux'" }, + { name = "fastsafetensors", marker = "platform_machine == 'aarch64' and sys_platform == 'linux'" }, + { name = "filelock", marker = "platform_machine == 'aarch64' and sys_platform == 'linux'" }, + { name = "flashinfer-cubin", marker = "platform_machine == 'aarch64' and sys_platform == 'linux'" }, + { name = "flashinfer-python", marker = "platform_machine == 'aarch64' and sys_platform == 'linux'" }, + { name = "gguf", marker = "platform_machine == 'aarch64' and sys_platform == 'linux'" }, + { name = "ijson", marker = "platform_machine == 'aarch64' and sys_platform == 'linux'" }, + { name = "lark", marker = "platform_machine == 'aarch64' and sys_platform == 'linux'" }, + { name = "llguidance", marker = "platform_machine == 'aarch64' and sys_platform == 'linux'" }, + { name = "lm-format-enforcer", marker = "platform_machine == 'aarch64' and sys_platform == 'linux'" }, + { name = "mcp", marker = "platform_machine == 'aarch64' and sys_platform == 'linux'" }, + { name = "mistral-common", extra = ["image"], marker = "platform_machine == 'aarch64' and sys_platform == 'linux'" }, + { name = "model-hosting-container-standards", marker = "platform_machine == 'aarch64' and sys_platform == 'linux'" }, + { name = "msgspec", marker = "platform_machine == 'aarch64' and sys_platform == 'linux'" }, + { name = "ninja", marker = "platform_machine == 'aarch64' and sys_platform == 'linux'" }, + { name = "numba", marker = "platform_machine == 'aarch64' and sys_platform == 'linux'" }, + { name = "numpy", marker = "platform_machine == 'aarch64' and sys_platform == 'linux'" }, + { name = "nvidia-cudnn-frontend", marker = "platform_machine == 'aarch64' and sys_platform == 'linux'" }, + { name = "nvidia-cutlass-dsl", marker = "platform_machine == 'aarch64' and sys_platform == 'linux'" }, + { name = "openai", marker = "platform_machine == 'aarch64' and sys_platform == 'linux'" }, + { name = "openai-harmony", marker = "platform_machine == 'aarch64' and sys_platform == 'linux'" }, + { name = "opencv-python-headless", marker = "platform_machine == 'aarch64' and sys_platform == 'linux'" }, + { name = "opentelemetry-api", marker = "platform_machine == 'aarch64' and sys_platform == 'linux'" }, + { name = "opentelemetry-exporter-otlp", marker = "platform_machine == 'aarch64' and sys_platform == 'linux'" }, + { name = "opentelemetry-sdk", marker = "platform_machine == 'aarch64' and sys_platform == 'linux'" }, + { name = "opentelemetry-semantic-conventions-ai", marker = "platform_machine == 'aarch64' and sys_platform == 'linux'" }, + { name = "outlines-core", marker = "platform_machine == 'aarch64' and sys_platform == 'linux'" }, + { name = "partial-json-parser", marker = "platform_machine == 'aarch64' and sys_platform == 'linux'" }, + { name = "pillow", marker = "platform_machine == 'aarch64' and sys_platform == 'linux'" }, + { name = "prometheus-client", marker = "platform_machine == 'aarch64' and sys_platform == 'linux'" }, + { name = "prometheus-fastapi-instrumentator", marker = "platform_machine == 'aarch64' and sys_platform == 'linux'" }, + { name = "protobuf", marker = "platform_machine == 'aarch64' and sys_platform == 'linux'" }, + { name = "psutil", marker = "platform_machine == 'aarch64' and sys_platform == 'linux'" }, + { name = "py-cpuinfo", marker = "platform_machine == 'aarch64' and sys_platform == 'linux'" }, + { name = "pybase64", marker = "platform_machine == 'aarch64' and sys_platform == 'linux'" }, + { name = "pydantic", marker = "platform_machine == 'aarch64' and sys_platform == 'linux'" }, + { name = "python-json-logger", marker = "platform_machine == 'aarch64' and sys_platform == 'linux'" }, + { name = "pyyaml", marker = "platform_machine == 'aarch64' and sys_platform == 'linux'" }, + { name = "pyzmq", marker = "platform_machine == 'aarch64' and sys_platform == 'linux'" }, + { name = "quack-kernels", marker = "platform_machine == 'aarch64' and sys_platform == 'linux'" }, + { name = "regex", marker = "platform_machine == 'aarch64' and sys_platform == 'linux'" }, + { name = "requests", marker = "platform_machine == 'aarch64' and sys_platform == 'linux'" }, + { name = "sentencepiece", marker = "platform_machine == 'aarch64' and sys_platform == 'linux'" }, + { name = "setproctitle", marker = "platform_machine == 'aarch64' and sys_platform == 'linux'" }, + { name = "setuptools", marker = "platform_machine == 'aarch64' and sys_platform == 'linux'" }, + { name = "six", marker = "platform_machine == 'aarch64' and sys_platform == 'linux'" }, + { name = "tiktoken", marker = "platform_machine == 'aarch64' and sys_platform == 'linux'" }, + { name = "tilelang", marker = "platform_machine == 'aarch64' and sys_platform == 'linux'" }, + { name = "tokenizers", marker = "platform_machine == 'aarch64' and sys_platform == 'linux'" }, + { name = "torch", marker = "platform_machine == 'aarch64' and sys_platform == 'linux'" }, + { name = "torchaudio", marker = "platform_machine == 'aarch64' and sys_platform == 'linux'" }, + { name = "torchvision", marker = "platform_machine == 'aarch64' and sys_platform == 'linux'" }, + { name = "tqdm", marker = "platform_machine == 'aarch64' and sys_platform == 'linux'" }, + { name = "transformers", marker = "platform_machine == 'aarch64' and sys_platform == 'linux'" }, + { name = "typing-extensions", marker = "platform_machine == 'aarch64' and sys_platform == 'linux'" }, + { name = "watchfiles", marker = "platform_machine == 'aarch64' and sys_platform == 'linux'" }, + { name = "xgrammar", marker = "platform_machine == 'aarch64' and sys_platform == 'linux'" }, +] +wheels = [ + { url = "https://github.com/vllm-project/vllm/releases/download/v0.20.2/vllm-0.20.2+cu129-cp38-abi3-manylinux_2_31_aarch64.whl", hash = "sha256:8a58a086c5c4ed2883eee36aaaf6b79c83463d02da3015454acf92afcc8e150e" }, +] + +[package.metadata] +requires-dist = [ + { name = "aiohttp", specifier = ">=3.13.3" }, + { name = "anthropic", specifier = ">=0.71.0" }, + { name = "apache-tvm-ffi", specifier = "==0.1.9" }, + { name = "av", marker = "extra == 'audio'" }, { name = "blake3" }, { name = "cachetools" }, { name = "cbor2" }, { name = "cloudpickle" }, - { name = "compressed-tensors" }, - { name = "depyf" }, - { name = "diskcache" }, + { name = "compressed-tensors", specifier = "==0.15.0.1" }, + { name = "datasets", marker = "extra == 'bench'" }, + { name = "depyf", specifier = "==0.20.0" }, + { name = "diskcache", specifier = "==5.6.3" }, { name = "einops" }, - { name = "fastapi", extra = ["standard"] }, - { name = "filelock" }, - { name = "flashinfer-python" }, - { name = "gguf" }, - { name = "lark" }, - { name = "llguidance", marker = "platform_machine == 'aarch64' or platform_machine == 'arm64' or platform_machine == 'ppc64le' or platform_machine == 's390x' or platform_machine == 'x86_64'" }, - { name = "lm-format-enforcer" }, - { name = "mistral-common", extra = ["image"] }, - { name = "model-hosting-container-standards" }, + { name = "fastapi", extras = ["standard"], specifier = ">=0.115.0" }, + { name = "fastsafetensors", specifier = ">=0.2.2" }, + { name = "fastsafetensors", marker = "extra == 'fastsafetensors'", specifier = ">=0.2.2" }, + { name = "filelock", specifier = ">=3.16.1" }, + { name = "flashinfer-cubin", specifier = "==0.6.8.post1" }, + { name = "flashinfer-python", specifier = "==0.6.8.post1" }, + { name = "gguf", specifier = ">=0.17.0" }, + { name = "helion", marker = "extra == 'helion'", specifier = "==1.0.0" }, + { name = "ijson" }, + { name = "instanttensor", marker = "extra == 'instanttensor'", specifier = ">=0.1.5" }, + { name = "lark", specifier = "==1.2.2" }, + { name = "llguidance", marker = "platform_machine == 'aarch64' or platform_machine == 'arm64' or platform_machine == 'ppc64le' or platform_machine == 'x86_64'", specifier = ">=1.3.0,<1.4.0" }, + { name = "lm-format-enforcer", specifier = "==0.11.3" }, + { name = "matplotlib", marker = "extra == 'bench'" }, + { name = "mcp" }, + { name = "mistral-common", extras = ["audio"], marker = "extra == 'audio'" }, + { name = "mistral-common", extras = ["image"], specifier = ">=1.11.0" }, + { name = "model-hosting-container-standards", specifier = ">=0.1.13,<1.0.0" }, { name = "msgspec" }, { name = "ninja" }, - { name = "numba" }, + { name = "numba", specifier = "==0.65.0" }, { name = "numpy" }, - { name = "openai" }, - { name = "openai-harmony" }, - { name = "opencv-python-headless" }, - { name = "outlines-core" }, + { name = "nvidia-cudnn-frontend", specifier = ">=1.13.0,<1.19.0" }, + { name = "nvidia-cutlass-dsl", specifier = ">=4.4.2" }, + { name = "openai", specifier = ">=2.0.0" }, + { name = "openai-harmony", specifier = ">=0.0.3" }, + { name = "opencv-python-headless", specifier = ">=4.13.0" }, + { name = "opentelemetry-api", specifier = ">=1.27.0" }, + { name = "opentelemetry-api", marker = "extra == 'otel'", specifier = ">=1.26.0" }, + { name = "opentelemetry-exporter-otlp", specifier = ">=1.27.0" }, + { name = "opentelemetry-exporter-otlp", marker = "extra == 'otel'", specifier = ">=1.26.0" }, + { name = "opentelemetry-sdk", specifier = ">=1.27.0" }, + { name = "opentelemetry-sdk", marker = "extra == 'otel'", specifier = ">=1.26.0" }, + { name = "opentelemetry-semantic-conventions-ai", specifier = ">=0.4.1" }, + { name = "opentelemetry-semantic-conventions-ai", marker = "extra == 'otel'", specifier = ">=0.4.1" }, + { name = "outlines-core", specifier = "==0.2.14" }, + { name = "pandas", marker = "extra == 'bench'" }, { name = "partial-json-parser" }, { name = "pillow" }, - { name = "prometheus-client" }, - { name = "prometheus-fastapi-instrumentator" }, - { name = "protobuf" }, + { name = "plotly", marker = "extra == 'bench'" }, + { name = "prometheus-client", specifier = ">=0.18.0" }, + { name = "prometheus-fastapi-instrumentator", specifier = ">=7.0.0" }, + { name = "protobuf", specifier = ">=5.29.6,!=6.30.*,!=6.31.*,!=6.32.*,!=6.33.0.*,!=6.33.1.*,!=6.33.2.*,!=6.33.3.*,!=6.33.4.*" }, { name = "psutil" }, { name = "py-cpuinfo" }, { name = "pybase64" }, - { name = "pydantic" }, + { name = "pydantic", specifier = ">=2.12.0" }, { name = "python-json-logger" }, { name = "pyyaml" }, - { name = "pyzmq" }, - { name = "ray", extra = ["cgraph"] }, + { name = "pyzmq", specifier = ">=25.0.0" }, + { name = "quack-kernels", specifier = ">=0.3.3" }, { name = "regex" }, - { name = "requests" }, - { name = "scipy" }, + { name = "requests", specifier = ">=2.26.0" }, + { name = "runai-model-streamer", extras = ["azure", "gcs", "s3"], marker = "extra == 'runai'", specifier = ">=0.15.7" }, + { name = "scipy", marker = "extra == 'audio'" }, + { name = "scipy", marker = "extra == 'bench'" }, + { name = "seaborn", marker = "extra == 'bench'" }, { name = "sentencepiece" }, { name = "setproctitle" }, - { name = "setuptools" }, - { name = "six" }, - { name = "tiktoken" }, - { name = "tokenizers" }, - { name = "torch" }, - { name = "torchaudio" }, - { name = "torchvision" }, + { name = "setuptools", marker = "python_full_version >= '3.12'", specifier = ">=77.0.3,<81.0.0" }, + { name = "six", marker = "python_full_version >= '3.12'", specifier = ">=1.16.0" }, + { name = "smg-grpc-servicer", extras = ["vllm"], marker = "extra == 'grpc'", specifier = ">=0.5.2" }, + { name = "soundfile", marker = "extra == 'audio'" }, + { name = "tensorizer", marker = "extra == 'tensorizer'", specifier = "==2.10.1" }, + { name = "tiktoken", specifier = ">=0.6.0" }, + { name = "tilelang", specifier = "==0.1.9" }, + { name = "tokenizers", specifier = ">=0.21.1" }, + { name = "torch", specifier = "==2.11.0" }, + { name = "torchaudio", specifier = "==2.11.0" }, + { name = "torchvision", specifier = "==0.26.0" }, + { name = "tqdm" }, + { name = "transformers", specifier = ">=4.56.0,!=5.0.*,!=5.1.*,!=5.2.*,!=5.3.*,!=5.4.*,!=5.5.0" }, + { name = "typing-extensions", specifier = ">=4.10" }, + { name = "watchfiles" }, + { name = "xgrammar", marker = "platform_machine == 'aarch64' or platform_machine == 'arm64' or platform_machine == 'ppc64le' or platform_machine == 's390x' or platform_machine == 'x86_64'", specifier = ">=0.1.32,<1.0.0" }, + { name = "zentorch-weekly", marker = "extra == 'zen'", specifier = "==5.2.1.dev20260408" }, +] +provides-extras = ["zen", "bench", "tensorizer", "fastsafetensors", "instanttensor", "runai", "audio", "video", "flashinfer", "helion", "grpc", "otel"] + +[[package]] +name = "vllm" +version = "0.20.2+cu129" +source = { url = "https://github.com/vllm-project/vllm/releases/download/v0.20.2/vllm-0.20.2+cu129-cp38-abi3-manylinux_2_31_x86_64.whl" } +resolution-markers = [ + "platform_machine == 'x86_64' and sys_platform == 'linux'", +] +dependencies = [ + { name = "aiohttp", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "anthropic", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "apache-tvm-ffi", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "blake3", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "cachetools", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "cbor2", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "cloudpickle", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "compressed-tensors", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "depyf", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "diskcache", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "einops", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "fastapi", extra = ["standard"], marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "fastsafetensors", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "filelock", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "flashinfer-cubin", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "flashinfer-python", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "gguf", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "ijson", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "lark", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "llguidance", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "lm-format-enforcer", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "mcp", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "mistral-common", extra = ["image"], marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "model-hosting-container-standards", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "msgspec", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "ninja", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "numba", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "numpy", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "nvidia-cudnn-frontend", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "nvidia-cutlass-dsl", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "openai", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "openai-harmony", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "opencv-python-headless", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "opentelemetry-api", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "opentelemetry-exporter-otlp", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "opentelemetry-sdk", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "opentelemetry-semantic-conventions-ai", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "outlines-core", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "partial-json-parser", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "pillow", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "prometheus-client", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "prometheus-fastapi-instrumentator", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "protobuf", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "psutil", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "py-cpuinfo", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "pybase64", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "pydantic", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "python-json-logger", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "pyyaml", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "pyzmq", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "quack-kernels", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "regex", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "requests", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "sentencepiece", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "setproctitle", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "setuptools", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "six", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "tiktoken", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "tilelang", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "tokenizers", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "torch", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "torchaudio", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "torchvision", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "tqdm", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "transformers", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "typing-extensions", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "watchfiles", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "xgrammar", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, +] +wheels = [ + { url = "https://github.com/vllm-project/vllm/releases/download/v0.20.2/vllm-0.20.2+cu129-cp38-abi3-manylinux_2_31_x86_64.whl", hash = "sha256:2f8c2bf2ac6d3d16f930535e66822abd71065468521884eb5b910225b2abef4b" }, +] + +[package.metadata] +requires-dist = [ + { name = "aiohttp", specifier = ">=3.13.3" }, + { name = "anthropic", specifier = ">=0.71.0" }, + { name = "apache-tvm-ffi", specifier = "==0.1.9" }, + { name = "av", marker = "extra == 'audio'" }, + { name = "blake3" }, + { name = "cachetools" }, + { name = "cbor2" }, + { name = "cloudpickle" }, + { name = "compressed-tensors", specifier = "==0.15.0.1" }, + { name = "datasets", marker = "extra == 'bench'" }, + { name = "depyf", specifier = "==0.20.0" }, + { name = "diskcache", specifier = "==5.6.3" }, + { name = "einops" }, + { name = "fastapi", extras = ["standard"], specifier = ">=0.115.0" }, + { name = "fastsafetensors", specifier = ">=0.2.2" }, + { name = "fastsafetensors", marker = "extra == 'fastsafetensors'", specifier = ">=0.2.2" }, + { name = "filelock", specifier = ">=3.16.1" }, + { name = "flashinfer-cubin", specifier = "==0.6.8.post1" }, + { name = "flashinfer-python", specifier = "==0.6.8.post1" }, + { name = "gguf", specifier = ">=0.17.0" }, + { name = "helion", marker = "extra == 'helion'", specifier = "==1.0.0" }, + { name = "ijson" }, + { name = "instanttensor", marker = "extra == 'instanttensor'", specifier = ">=0.1.5" }, + { name = "lark", specifier = "==1.2.2" }, + { name = "llguidance", marker = "platform_machine == 'aarch64' or platform_machine == 'arm64' or platform_machine == 'ppc64le' or platform_machine == 'x86_64'", specifier = ">=1.3.0,<1.4.0" }, + { name = "lm-format-enforcer", specifier = "==0.11.3" }, + { name = "matplotlib", marker = "extra == 'bench'" }, + { name = "mcp" }, + { name = "mistral-common", extras = ["audio"], marker = "extra == 'audio'" }, + { name = "mistral-common", extras = ["image"], specifier = ">=1.11.0" }, + { name = "model-hosting-container-standards", specifier = ">=0.1.13,<1.0.0" }, + { name = "msgspec" }, + { name = "ninja" }, + { name = "numba", specifier = "==0.65.0" }, + { name = "numpy" }, + { name = "nvidia-cudnn-frontend", specifier = ">=1.13.0,<1.19.0" }, + { name = "nvidia-cutlass-dsl", specifier = ">=4.4.2" }, + { name = "openai", specifier = ">=2.0.0" }, + { name = "openai-harmony", specifier = ">=0.0.3" }, + { name = "opencv-python-headless", specifier = ">=4.13.0" }, + { name = "opentelemetry-api", specifier = ">=1.27.0" }, + { name = "opentelemetry-api", marker = "extra == 'otel'", specifier = ">=1.26.0" }, + { name = "opentelemetry-exporter-otlp", specifier = ">=1.27.0" }, + { name = "opentelemetry-exporter-otlp", marker = "extra == 'otel'", specifier = ">=1.26.0" }, + { name = "opentelemetry-sdk", specifier = ">=1.27.0" }, + { name = "opentelemetry-sdk", marker = "extra == 'otel'", specifier = ">=1.26.0" }, + { name = "opentelemetry-semantic-conventions-ai", specifier = ">=0.4.1" }, + { name = "opentelemetry-semantic-conventions-ai", marker = "extra == 'otel'", specifier = ">=0.4.1" }, + { name = "outlines-core", specifier = "==0.2.14" }, + { name = "pandas", marker = "extra == 'bench'" }, + { name = "partial-json-parser" }, + { name = "pillow" }, + { name = "plotly", marker = "extra == 'bench'" }, + { name = "prometheus-client", specifier = ">=0.18.0" }, + { name = "prometheus-fastapi-instrumentator", specifier = ">=7.0.0" }, + { name = "protobuf", specifier = ">=5.29.6,!=6.30.*,!=6.31.*,!=6.32.*,!=6.33.0.*,!=6.33.1.*,!=6.33.2.*,!=6.33.3.*,!=6.33.4.*" }, + { name = "psutil" }, + { name = "py-cpuinfo" }, + { name = "pybase64" }, + { name = "pydantic", specifier = ">=2.12.0" }, + { name = "python-json-logger" }, + { name = "pyyaml" }, + { name = "pyzmq", specifier = ">=25.0.0" }, + { name = "quack-kernels", specifier = ">=0.3.3" }, + { name = "regex" }, + { name = "requests", specifier = ">=2.26.0" }, + { name = "runai-model-streamer", extras = ["azure", "gcs", "s3"], marker = "extra == 'runai'", specifier = ">=0.15.7" }, + { name = "scipy", marker = "extra == 'audio'" }, + { name = "scipy", marker = "extra == 'bench'" }, + { name = "seaborn", marker = "extra == 'bench'" }, + { name = "sentencepiece" }, + { name = "setproctitle" }, + { name = "setuptools", marker = "python_full_version >= '3.12'", specifier = ">=77.0.3,<81.0.0" }, + { name = "six", marker = "python_full_version >= '3.12'", specifier = ">=1.16.0" }, + { name = "smg-grpc-servicer", extras = ["vllm"], marker = "extra == 'grpc'", specifier = ">=0.5.2" }, + { name = "soundfile", marker = "extra == 'audio'" }, + { name = "tensorizer", marker = "extra == 'tensorizer'", specifier = "==2.10.1" }, + { name = "tiktoken", specifier = ">=0.6.0" }, + { name = "tilelang", specifier = "==0.1.9" }, + { name = "tokenizers", specifier = ">=0.21.1" }, + { name = "torch", specifier = "==2.11.0" }, + { name = "torchaudio", specifier = "==2.11.0" }, + { name = "torchvision", specifier = "==0.26.0" }, { name = "tqdm" }, - { name = "transformers" }, - { name = "typing-extensions" }, + { name = "transformers", specifier = ">=4.56.0,!=5.0.*,!=5.1.*,!=5.2.*,!=5.3.*,!=5.4.*,!=5.5.0" }, + { name = "typing-extensions", specifier = ">=4.10" }, { name = "watchfiles" }, - { name = "xgrammar", marker = "platform_machine == 'aarch64' or platform_machine == 'arm64' or platform_machine == 'ppc64le' or platform_machine == 's390x' or platform_machine == 'x86_64'" }, + { name = "xgrammar", marker = "platform_machine == 'aarch64' or platform_machine == 'arm64' or platform_machine == 'ppc64le' or platform_machine == 's390x' or platform_machine == 'x86_64'", specifier = ">=0.1.32,<1.0.0" }, + { name = "zentorch-weekly", marker = "extra == 'zen'", specifier = "==5.2.1.dev20260408" }, +] +provides-extras = ["zen", "bench", "tensorizer", "fastsafetensors", "instanttensor", "runai", "audio", "video", "flashinfer", "helion", "grpc", "otel"] + +[[package]] +name = "vllm-router" +version = "0.1.22" +source = { url = "https://github.com/PrimeIntellect-ai/router/releases/download/v0.1.22/vllm_router-0.1.22-cp38-abi3-manylinux_2_28_x86_64.whl" } +dependencies = [ + { name = "aiohttp", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "fastapi", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "orjson", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "requests", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "setproctitle", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "uvicorn", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/3d/70/ce6104d49a8139d055fbe4db3ddbb4e670be2f5db2172e590623efb26537/vllm-0.12.0.tar.gz", hash = "sha256:dc2863be80260f378f1e72f9b82b4009610e434d05fdd7067383b231b89c4f4a", size = 17594403, upload-time = "2025-12-03T11:38:16.393Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/89/30/feb103208b5b55250c47e0582be72421e4d9c901ea81d9c95be0455fc3aa/vllm-0.12.0-cp38-abi3-manylinux_2_31_aarch64.whl", hash = "sha256:44684361b22b12898fb356074cc3e8786036f9f1e17ea1105bd27d1ffa63cbe9", size = 435774675, upload-time = "2025-12-03T11:37:57.693Z" }, - { url = "https://files.pythonhosted.org/packages/14/2e/cc616f226f389c4cbffcd9387581c1757f2d4a6cf0af7fb9a0cd4eddc1e4/vllm-0.12.0-cp38-abi3-manylinux_2_31_x86_64.whl", hash = "sha256:596673edcac5ea8ab8cb2d9b0200052978cd91d78a054926d9d0c6895d34237f", size = 466530617, upload-time = "2025-12-03T11:37:24.522Z" }, + { url = "https://github.com/PrimeIntellect-ai/router/releases/download/v0.1.22/vllm_router-0.1.22-cp38-abi3-manylinux_2_28_x86_64.whl", hash = "sha256:6361a0387241e56932f3ba2e51af27f58d11a462e3187e58286b2f96056e4d15" }, +] + +[package.metadata] +requires-dist = [ + { name = "aiohttp" }, + { name = "fastapi" }, + { name = "orjson" }, + { name = "pytest", marker = "extra == 'dev'", specifier = ">=7.0.0" }, + { name = "pytest-asyncio", marker = "extra == 'dev'", specifier = ">=0.21.0" }, + { name = "pytest-cov", marker = "extra == 'dev'", specifier = ">=4.0.0" }, + { name = "requests", specifier = ">=2.25.0" }, + { name = "setproctitle" }, + { name = "uvicorn" }, ] +provides-extras = ["dev"] [[package]] name = "wadler-lindig" @@ -3974,31 +4601,26 @@ wheels = [ [[package]] name = "wandb" -version = "0.21.4" +version = "0.26.1" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "click" }, - { name = "gitpython" }, - { name = "packaging" }, - { name = "platformdirs" }, - { name = "protobuf" }, - { name = "pydantic" }, - { name = "pyyaml" }, - { name = "requests" }, - { name = "sentry-sdk" }, - { name = "typing-extensions" }, + { name = "click", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "gitpython", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "packaging", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "platformdirs", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "protobuf", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "pydantic", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "pyyaml", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "requests", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "sentry-sdk", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "typing-extensions", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/59/a8/aaa3f3f8e410f34442466aac10b1891b3084d35b98aef59ebcb4c0efb941/wandb-0.21.4.tar.gz", hash = "sha256:b350d50973409658deb455010fafcfa81e6be3470232e316286319e839ffb67b", size = 40175929, upload-time = "2025-09-11T21:14:29.161Z" } +sdist = { url = "https://files.pythonhosted.org/packages/6a/a4/72a6640e1f566e81f184a426e3e45298d4c6672664de41adb7eb6f64370a/wandb-0.26.1.tar.gz", hash = "sha256:eef2dbaea06f0b1c0cdc5d76f544ae4c2b8848fc512442a00bd59f0502fc8aa1", size = 42159814, upload-time = "2026-04-23T16:27:34.033Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/d2/6b/3a8d9db18a4c4568599a8792c0c8b1f422d9864c7123e8301a9477fbf0ac/wandb-0.21.4-py3-none-macosx_10_14_x86_64.whl", hash = "sha256:c681ef7adb09925251d8d995c58aa76ae86a46dbf8de3b67353ad99fdef232d5", size = 18845369, upload-time = "2025-09-11T21:14:02.879Z" }, - { url = "https://files.pythonhosted.org/packages/60/e0/d7d6818938ec6958c93d979f9a90ea3d06bdc41e130b30f8cd89ae03c245/wandb-0.21.4-py3-none-macosx_12_0_arm64.whl", hash = "sha256:d35acc65c10bb7ac55d1331f7b1b8ab761f368f7b051131515f081a56ea5febc", size = 18339122, upload-time = "2025-09-11T21:14:06.455Z" }, - { url = "https://files.pythonhosted.org/packages/13/29/9bb8ed4adf32bed30e4d5df74d956dd1e93b6fd4bbc29dbe84167c84804b/wandb-0.21.4-py3-none-macosx_12_0_x86_64.whl", hash = "sha256:765e66b57b7be5f393ecebd9a9d2c382c9f979d19cdee4a3f118eaafed43fca1", size = 19081975, upload-time = "2025-09-11T21:14:09.317Z" }, - { url = "https://files.pythonhosted.org/packages/30/6e/4aa33bc2c56b70c0116e73687c72c7a674f4072442633b3b23270d2215e3/wandb-0.21.4-py3-none-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:06127ec49245d12fdb3922c1eca1ab611cefc94adabeaaaba7b069707c516cba", size = 18161358, upload-time = "2025-09-11T21:14:12.092Z" }, - { url = "https://files.pythonhosted.org/packages/f7/56/d9f845ecfd5e078cf637cb29d8abe3350b8a174924c54086168783454a8f/wandb-0.21.4-py3-none-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:48d4f65f1be5f5a25b868695e09cdbfe481678220df349a8c2cbed3992fb497f", size = 19602680, upload-time = "2025-09-11T21:14:14.987Z" }, - { url = "https://files.pythonhosted.org/packages/68/ea/237a3c2b679a35e02e577c5bf844d6a221a7d32925ab8d5230529e9f2841/wandb-0.21.4-py3-none-musllinux_1_2_aarch64.whl", hash = "sha256:ebd11f78351a3ca22caa1045146a6d2ad9e62fed6d0de2e67a0db5710d75103a", size = 18166392, upload-time = "2025-09-11T21:14:17.478Z" }, - { url = "https://files.pythonhosted.org/packages/12/e3/dbf2c575c79c99d94f16ce1a2cbbb2529d5029a76348c1ddac7e47f6873f/wandb-0.21.4-py3-none-musllinux_1_2_x86_64.whl", hash = "sha256:595b9e77591a805653e05db8b892805ee0a5317d147ef4976353e4f1cc16ebdc", size = 19678800, upload-time = "2025-09-11T21:14:20.264Z" }, - { url = "https://files.pythonhosted.org/packages/fa/eb/4ed04879d697772b8eb251c0e5af9a4ff7e2cc2b3fcd4b8eee91253ec2f1/wandb-0.21.4-py3-none-win32.whl", hash = "sha256:f9c86eb7eb7d40c6441533428188b1ae3205674e80c940792d850e2c1fe8d31e", size = 18738950, upload-time = "2025-09-11T21:14:23.08Z" }, - { url = "https://files.pythonhosted.org/packages/c3/4a/86c5e19600cb6a616a45f133c26826b46133499cd72d592772929d530ccd/wandb-0.21.4-py3-none-win_amd64.whl", hash = "sha256:2da3d5bb310a9f9fb7f680f4aef285348095a4cc6d1ce22b7343ba4e3fffcd84", size = 18738953, upload-time = "2025-09-11T21:14:25.539Z" }, + { url = "https://files.pythonhosted.org/packages/56/af/400d84a3bdce0b062b4baa70acb6becd2c8018697f4fbf5af9a9e1e406e5/wandb-0.26.1-py3-none-manylinux_2_28_aarch64.whl", hash = "sha256:7c78bc2454cfe1ffa1c3a256060a387356eed8a4488e024d9d2eba8f2b5bd51d", size = 25421317, upload-time = "2026-04-23T16:27:06.411Z" }, + { url = "https://files.pythonhosted.org/packages/7b/e9/b4bf8f3509dcea1cec52233a38991459654635b5a8e6a494eb912e1b9cfb/wandb-0.26.1-py3-none-manylinux_2_28_x86_64.whl", hash = "sha256:a2c8eeec8706dcd2872e69c3b4d20ec523082fdb4440295491556e219ad2aa67", size = 27192831, upload-time = "2026-04-23T16:27:10.308Z" }, + { url = "https://files.pythonhosted.org/packages/62/cf/4a6dce0c782223ef0eeea7139daee73418a7322befcf083512c31cebaa18/wandb-0.26.1-py3-none-musllinux_1_2_aarch64.whl", hash = "sha256:2fa768ee0636a569afb7541cf996e56309c47070566a38916823f94e02afe586", size = 25593326, upload-time = "2026-04-23T16:27:14.259Z" }, + { url = "https://files.pythonhosted.org/packages/df/99/58c3d8c36ae8e2b7d70bf6493eb5daa1cca0231a04b025717b4cd1a78f1e/wandb-0.26.1-py3-none-musllinux_1_2_x86_64.whl", hash = "sha256:5854928725cfeff1f284d5c043cd353f810e5da02eead2c120ef5056ad026fea", size = 27535542, upload-time = "2026-04-23T16:27:18.473Z" }, ] [[package]] @@ -4006,23 +4628,14 @@ name = "watchfiles" version = "1.1.0" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "anyio" }, + { name = "anyio", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, ] sdist = { url = "https://files.pythonhosted.org/packages/2a/9a/d451fcc97d029f5812e898fd30a53fd8c15c7bbd058fd75cfc6beb9bd761/watchfiles-1.1.0.tar.gz", hash = "sha256:693ed7ec72cbfcee399e92c895362b6e66d63dac6b91e2c11ae03d10d503e575", size = 94406, upload-time = "2025-06-15T19:06:59.42Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/f6/b8/858957045a38a4079203a33aaa7d23ea9269ca7761c8a074af3524fbb240/watchfiles-1.1.0-cp312-cp312-macosx_10_12_x86_64.whl", hash = "sha256:9dc001c3e10de4725c749d4c2f2bdc6ae24de5a88a339c4bce32300a31ede179", size = 402339, upload-time = "2025-06-15T19:05:24.516Z" }, - { url = "https://files.pythonhosted.org/packages/80/28/98b222cca751ba68e88521fabd79a4fab64005fc5976ea49b53fa205d1fa/watchfiles-1.1.0-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:d9ba68ec283153dead62cbe81872d28e053745f12335d037de9cbd14bd1877f5", size = 394409, upload-time = "2025-06-15T19:05:25.469Z" }, { url = "https://files.pythonhosted.org/packages/86/50/dee79968566c03190677c26f7f47960aff738d32087087bdf63a5473e7df/watchfiles-1.1.0-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:130fc497b8ee68dce163e4254d9b0356411d1490e868bd8790028bc46c5cc297", size = 450939, upload-time = "2025-06-15T19:05:26.494Z" }, - { url = "https://files.pythonhosted.org/packages/40/45/a7b56fb129700f3cfe2594a01aa38d033b92a33dddce86c8dfdfc1247b72/watchfiles-1.1.0-cp312-cp312-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:50a51a90610d0845a5931a780d8e51d7bd7f309ebc25132ba975aca016b576a0", size = 457270, upload-time = "2025-06-15T19:05:27.466Z" }, - { url = "https://files.pythonhosted.org/packages/b5/c8/fa5ef9476b1d02dc6b5e258f515fcaaecf559037edf8b6feffcbc097c4b8/watchfiles-1.1.0-cp312-cp312-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:dc44678a72ac0910bac46fa6a0de6af9ba1355669b3dfaf1ce5f05ca7a74364e", size = 483370, upload-time = "2025-06-15T19:05:28.548Z" }, - { url = "https://files.pythonhosted.org/packages/98/68/42cfcdd6533ec94f0a7aab83f759ec11280f70b11bfba0b0f885e298f9bd/watchfiles-1.1.0-cp312-cp312-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:a543492513a93b001975ae283a51f4b67973662a375a403ae82f420d2c7205ee", size = 598654, upload-time = "2025-06-15T19:05:29.997Z" }, - { url = "https://files.pythonhosted.org/packages/d3/74/b2a1544224118cc28df7e59008a929e711f9c68ce7d554e171b2dc531352/watchfiles-1.1.0-cp312-cp312-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:8ac164e20d17cc285f2b94dc31c384bc3aa3dd5e7490473b3db043dd70fbccfd", size = 478667, upload-time = "2025-06-15T19:05:31.172Z" }, { url = "https://files.pythonhosted.org/packages/8c/77/e3362fe308358dc9f8588102481e599c83e1b91c2ae843780a7ded939a35/watchfiles-1.1.0-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:f7590d5a455321e53857892ab8879dce62d1f4b04748769f5adf2e707afb9d4f", size = 452213, upload-time = "2025-06-15T19:05:32.299Z" }, { url = "https://files.pythonhosted.org/packages/6e/17/c8f1a36540c9a1558d4faf08e909399e8133599fa359bf52ec8fcee5be6f/watchfiles-1.1.0-cp312-cp312-musllinux_1_1_aarch64.whl", hash = "sha256:37d3d3f7defb13f62ece99e9be912afe9dd8a0077b7c45ee5a57c74811d581a4", size = 626718, upload-time = "2025-06-15T19:05:33.415Z" }, { url = "https://files.pythonhosted.org/packages/26/45/fb599be38b4bd38032643783d7496a26a6f9ae05dea1a42e58229a20ac13/watchfiles-1.1.0-cp312-cp312-musllinux_1_1_x86_64.whl", hash = "sha256:7080c4bb3efd70a07b1cc2df99a7aa51d98685be56be6038c3169199d0a1c69f", size = 623098, upload-time = "2025-06-15T19:05:34.534Z" }, - { url = "https://files.pythonhosted.org/packages/a1/e7/fdf40e038475498e160cd167333c946e45d8563ae4dd65caf757e9ffe6b4/watchfiles-1.1.0-cp312-cp312-win32.whl", hash = "sha256:cbcf8630ef4afb05dc30107bfa17f16c0896bb30ee48fc24bf64c1f970f3b1fd", size = 279209, upload-time = "2025-06-15T19:05:35.577Z" }, - { url = "https://files.pythonhosted.org/packages/3f/d3/3ae9d5124ec75143bdf088d436cba39812122edc47709cd2caafeac3266f/watchfiles-1.1.0-cp312-cp312-win_amd64.whl", hash = "sha256:cbd949bdd87567b0ad183d7676feb98136cde5bb9025403794a4c0db28ed3a47", size = 292786, upload-time = "2025-06-15T19:05:36.559Z" }, - { url = "https://files.pythonhosted.org/packages/26/2f/7dd4fc8b5f2b34b545e19629b4a018bfb1de23b3a496766a2c1165ca890d/watchfiles-1.1.0-cp312-cp312-win_arm64.whl", hash = "sha256:0a7d40b77f07be87c6faa93d0951a0fcd8cbca1ddff60a1b65d741bac6f3a9f6", size = 284343, upload-time = "2025-06-15T19:05:37.5Z" }, ] [[package]] @@ -4034,23 +4647,25 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/fd/84/fd2ba7aafacbad3c4201d395674fc6348826569da3c0937e75505ead3528/wcwidth-0.2.13-py2.py3-none-any.whl", hash = "sha256:3da69048e4540d84af32131829ff948f1e022c1c6bdb8d6102117aac784f6859", size = 34166, upload-time = "2024-01-06T02:10:55.763Z" }, ] +[[package]] +name = "websocket-client" +version = "1.9.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/2c/41/aa4bf9664e4cda14c3b39865b12251e8e7d239f4cd0e3cc1b6c2ccde25c1/websocket_client-1.9.0.tar.gz", hash = "sha256:9e813624b6eb619999a97dc7958469217c3176312b3a16a4bd1bc7e08a46ec98", size = 70576, upload-time = "2025-10-07T21:16:36.495Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/34/db/b10e48aa8fff7407e67470363eac595018441cf32d5e1001567a7aeba5d2/websocket_client-1.9.0-py3-none-any.whl", hash = "sha256:af248a825037ef591efbf6ed20cc5faa03d3b47b9e5a2230a529eeee1c1fc3ef", size = 82616, upload-time = "2025-10-07T21:16:34.951Z" }, +] + [[package]] name = "websockets" version = "15.0.1" source = { registry = "https://pypi.org/simple" } sdist = { url = "https://files.pythonhosted.org/packages/21/e6/26d09fab466b7ca9c7737474c52be4f76a40301b08362eb2dbc19dcc16c1/websockets-15.0.1.tar.gz", hash = "sha256:82544de02076bafba038ce055ee6412d68da13ab47f0c60cab827346de828dee", size = 177016, upload-time = "2025-03-05T20:03:41.606Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/51/6b/4545a0d843594f5d0771e86463606a3988b5a09ca5123136f8a76580dd63/websockets-15.0.1-cp312-cp312-macosx_10_13_universal2.whl", hash = "sha256:3e90baa811a5d73f3ca0bcbf32064d663ed81318ab225ee4f427ad4e26e5aff3", size = 175437, upload-time = "2025-03-05T20:02:16.706Z" }, - { url = "https://files.pythonhosted.org/packages/f4/71/809a0f5f6a06522af902e0f2ea2757f71ead94610010cf570ab5c98e99ed/websockets-15.0.1-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:592f1a9fe869c778694f0aa806ba0374e97648ab57936f092fd9d87f8bc03665", size = 173096, upload-time = "2025-03-05T20:02:18.832Z" }, - { url = "https://files.pythonhosted.org/packages/3d/69/1a681dd6f02180916f116894181eab8b2e25b31e484c5d0eae637ec01f7c/websockets-15.0.1-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:0701bc3cfcb9164d04a14b149fd74be7347a530ad3bbf15ab2c678a2cd3dd9a2", size = 173332, upload-time = "2025-03-05T20:02:20.187Z" }, { url = "https://files.pythonhosted.org/packages/a6/02/0073b3952f5bce97eafbb35757f8d0d54812b6174ed8dd952aa08429bcc3/websockets-15.0.1-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:e8b56bdcdb4505c8078cb6c7157d9811a85790f2f2b3632c7d1462ab5783d215", size = 183152, upload-time = "2025-03-05T20:02:22.286Z" }, - { url = "https://files.pythonhosted.org/packages/74/45/c205c8480eafd114b428284840da0b1be9ffd0e4f87338dc95dc6ff961a1/websockets-15.0.1-cp312-cp312-manylinux_2_5_i686.manylinux1_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:0af68c55afbd5f07986df82831c7bff04846928ea8d1fd7f30052638788bc9b5", size = 182096, upload-time = "2025-03-05T20:02:24.368Z" }, { url = "https://files.pythonhosted.org/packages/14/8f/aa61f528fba38578ec553c145857a181384c72b98156f858ca5c8e82d9d3/websockets-15.0.1-cp312-cp312-manylinux_2_5_x86_64.manylinux1_x86_64.manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:64dee438fed052b52e4f98f76c5790513235efaa1ef7f3f2192c392cd7c91b65", size = 182523, upload-time = "2025-03-05T20:02:25.669Z" }, { url = "https://files.pythonhosted.org/packages/ec/6d/0267396610add5bc0d0d3e77f546d4cd287200804fe02323797de77dbce9/websockets-15.0.1-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:d5f6b181bb38171a8ad1d6aa58a67a6aa9d4b38d0f8c5f496b9e42561dfc62fe", size = 182790, upload-time = "2025-03-05T20:02:26.99Z" }, - { url = "https://files.pythonhosted.org/packages/02/05/c68c5adbf679cf610ae2f74a9b871ae84564462955d991178f95a1ddb7dd/websockets-15.0.1-cp312-cp312-musllinux_1_2_i686.whl", hash = "sha256:5d54b09eba2bada6011aea5375542a157637b91029687eb4fdb2dab11059c1b4", size = 182165, upload-time = "2025-03-05T20:02:30.291Z" }, { url = "https://files.pythonhosted.org/packages/29/93/bb672df7b2f5faac89761cb5fa34f5cec45a4026c383a4b5761c6cea5c16/websockets-15.0.1-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:3be571a8b5afed347da347bfcf27ba12b069d9d7f42cb8c7028b5e98bbb12597", size = 182160, upload-time = "2025-03-05T20:02:31.634Z" }, - { url = "https://files.pythonhosted.org/packages/ff/83/de1f7709376dc3ca9b7eeb4b9a07b4526b14876b6d372a4dc62312bebee0/websockets-15.0.1-cp312-cp312-win32.whl", hash = "sha256:c338ffa0520bdb12fbc527265235639fb76e7bc7faafbb93f6ba80d9c06578a9", size = 176395, upload-time = "2025-03-05T20:02:33.017Z" }, - { url = "https://files.pythonhosted.org/packages/7d/71/abf2ebc3bbfa40f391ce1428c7168fb20582d0ff57019b69ea20fa698043/websockets-15.0.1-cp312-cp312-win_amd64.whl", hash = "sha256:fcd5cf9e305d7b8338754470cf69cf81f420459dbae8a3b40cee57417f4614a7", size = 176841, upload-time = "2025-03-05T20:02:34.498Z" }, { url = "https://files.pythonhosted.org/packages/fa/a8/5b41e0da817d64113292ab1f8247140aac61cbf6cfd085d6a0fa77f4984f/websockets-15.0.1-py3-none-any.whl", hash = "sha256:f7a866fbc1e97b5c617ee4116daaa09b722101d4a3c170c787450ba409f9736f", size = 169743, upload-time = "2025-03-05T20:03:39.41Z" }, ] @@ -4059,7 +4674,7 @@ name = "werkzeug" version = "3.1.3" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "markupsafe" }, + { name = "markupsafe", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, ] sdist = { url = "https://files.pythonhosted.org/packages/9f/69/83029f1f6300c5fb2471d621ab06f6ec6b3324685a2ce0f9777fd4a8b71e/werkzeug-3.1.3.tar.gz", hash = "sha256:60723ce945c19328679790e3282cc758aa4a6040e4bb330f53d30fa546d44746", size = 806925, upload-time = "2024-11-08T15:52:18.093Z" } wheels = [ @@ -4082,35 +4697,35 @@ wheels = [ ] [[package]] -name = "win32-setctime" -version = "1.2.0" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/b3/8f/705086c9d734d3b663af0e9bb3d4de6578d08f46b1b101c2442fd9aecaa2/win32_setctime-1.2.0.tar.gz", hash = "sha256:ae1fdf948f5640aae05c511ade119313fb6a30d7eabe25fef9764dca5873c4c0", size = 4867, upload-time = "2024-12-07T15:28:28.314Z" } +name = "wiki-search" +version = "0.1.23" +source = { registry = "https://hub.primeintellect.ai/primeintellect/simple/" } +dependencies = [ + { name = "chromadb", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "datasets", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "openai", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "verifiers", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, +] wheels = [ - { url = "https://files.pythonhosted.org/packages/e1/07/c6fe3ad3e685340704d314d765b7912993bcb8dc198f0e7a89382d37974b/win32_setctime-1.2.0-py3-none-any.whl", hash = "sha256:95d644c4e708aba81dc3704a116d8cbc974d70b3bdb8be1d150e36be6e9d1390", size = 4083, upload-time = "2024-12-07T15:28:26.465Z" }, + { url = "https://hub.primeintellect.ai/primeintellect/wiki-search/@10d58ffe/wiki_search-0.1.23-py3-none-any.whl", hash = "sha256:ffeff890f2d14d7b2910baf57c27f6939da0f669ae0c4545916762f3f4edd75b" }, ] [[package]] name = "xgrammar" -version = "0.1.27" +version = "0.1.33" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "mlx-lm", marker = "platform_machine == 'arm64' and sys_platform == 'darwin'" }, - { name = "ninja" }, - { name = "numpy" }, - { name = "pydantic" }, - { name = "torch" }, - { name = "transformers" }, - { name = "triton", version = "3.5.0", source = { registry = "https://pypi.org/simple" }, marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, - { name = "typing-extensions" }, + { name = "numpy", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "pydantic", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "torch", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "transformers", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "triton", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "typing-extensions", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/62/e1/b522b1e50fddd773d368c2945ef5ed628aa90c0c972027f9aa5a51d6d4f9/xgrammar-0.1.27.tar.gz", hash = "sha256:40af7bb2891f1633ec7f660723c74a92a963307d283aca9e3b4e53a0feaf1d46", size = 2303435, upload-time = "2025-11-04T03:11:53.512Z" } +sdist = { url = "https://files.pythonhosted.org/packages/db/43/e5dfddb1d2a4fccf3e3a88f103e88698cdefc3182f4e169a359ffe1c1794/xgrammar-0.1.33.tar.gz", hash = "sha256:8dbe5fc3d76651ab1fac7a68fc2a118b885fa0ec7189927fb6e0dce0081aea99", size = 2398956, upload-time = "2026-03-27T10:16:36.582Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/39/b6/09b43e2adff45d30ebcf9110d0ff753f4c96b368adaa2d166df3dee88d5f/xgrammar-0.1.27-cp312-cp312-macosx_10_14_x86_64.whl", hash = "sha256:6404a7714440eb86ab0379d749f33591274eeef04787dc00d61f22069f3ed51d", size = 663319, upload-time = "2025-11-04T03:11:28.682Z" }, - { url = "https://files.pythonhosted.org/packages/88/8b/53eb5c6d0df8df9f6350f182516a5b8c7b8b11d62650300d2c04af2bc4ea/xgrammar-0.1.27-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:d01fa9894bc44a7f6a70b0301b59f3e310c0e0e7b7ea4cf5ce190b12d8220dd8", size = 636168, upload-time = "2025-11-04T03:11:30.373Z" }, - { url = "https://files.pythonhosted.org/packages/08/1b/53d30395bb973f13255d3e3a72961f95fdfb4083877c3f93bb626e3d1522/xgrammar-0.1.27-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:906c0601bac9170e1bab77ca985259035ff9c386c347efcb191555eab86e984e", size = 8676340, upload-time = "2025-11-04T03:11:32.203Z" }, - { url = "https://files.pythonhosted.org/packages/48/74/70cfac0171d9f309cfe18c5384330e3edc9466c436b258495fd30ecf29a3/xgrammar-0.1.27-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:eb68988a122f544301c496f2cac8ee82960ca7f5b3a42a952b2a00c0a55e6ca5", size = 8870650, upload-time = "2025-11-04T03:11:34.322Z" }, - { url = "https://files.pythonhosted.org/packages/6a/a1/0392aa9c7669c56f7f88e4423b246476a74a72c3bb9db944e1bfc029985e/xgrammar-0.1.27-cp312-cp312-win_amd64.whl", hash = "sha256:3aac335ea052afc8f8dc34b9f2afcb9462a68189423aed9f60b0941db6cfc310", size = 708811, upload-time = "2025-11-04T03:11:36.214Z" }, + { url = "https://files.pythonhosted.org/packages/4e/04/43d4baca876f5ae1b45897ec30a59801a2da37f16da1fcd85f9555e4c125/xgrammar-0.1.33-cp312-cp312-manylinux_2_26_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:9c803e60d791854c5d1f271ece7e1f34d73c82dd4a8b2a06b7af5331482a78ac", size = 42133168, upload-time = "2026-03-27T10:15:16.994Z" }, + { url = "https://files.pythonhosted.org/packages/f0/a8/672833a3cff027253793aa999401d8364896ebf396967e475c7a878b895f/xgrammar-0.1.33-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:52b8eaa533282a0efb0835db6998ae72e7b3c7875d7a52e360ffebff9b78c30a", size = 42205803, upload-time = "2026-03-27T10:15:21.599Z" }, ] [[package]] @@ -4119,21 +4734,10 @@ version = "3.5.0" source = { registry = "https://pypi.org/simple" } sdist = { url = "https://files.pythonhosted.org/packages/00/5e/d6e5258d69df8b4ed8c83b6664f2b47d30d2dec551a29ad72a6c69eafd31/xxhash-3.5.0.tar.gz", hash = "sha256:84f2caddf951c9cbf8dc2e22a89d4ccf5d86391ac6418fe81e3c67d0cf60b45f", size = 84241, upload-time = "2024-08-17T09:20:38.972Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/07/0e/1bfce2502c57d7e2e787600b31c83535af83746885aa1a5f153d8c8059d6/xxhash-3.5.0-cp312-cp312-macosx_10_9_x86_64.whl", hash = "sha256:14470ace8bd3b5d51318782cd94e6f94431974f16cb3b8dc15d52f3b69df8e00", size = 31969, upload-time = "2024-08-17T09:18:24.025Z" }, - { url = "https://files.pythonhosted.org/packages/3f/d6/8ca450d6fe5b71ce521b4e5db69622383d039e2b253e9b2f24f93265b52c/xxhash-3.5.0-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:59aa1203de1cb96dbeab595ded0ad0c0056bb2245ae11fac11c0ceea861382b9", size = 30787, upload-time = "2024-08-17T09:18:25.318Z" }, { url = "https://files.pythonhosted.org/packages/5b/84/de7c89bc6ef63d750159086a6ada6416cc4349eab23f76ab870407178b93/xxhash-3.5.0-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:08424f6648526076e28fae6ea2806c0a7d504b9ef05ae61d196d571e5c879c84", size = 220959, upload-time = "2024-08-17T09:18:26.518Z" }, - { url = "https://files.pythonhosted.org/packages/fe/86/51258d3e8a8545ff26468c977101964c14d56a8a37f5835bc0082426c672/xxhash-3.5.0-cp312-cp312-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:61a1ff00674879725b194695e17f23d3248998b843eb5e933007ca743310f793", size = 200006, upload-time = "2024-08-17T09:18:27.905Z" }, - { url = "https://files.pythonhosted.org/packages/02/0a/96973bd325412feccf23cf3680fd2246aebf4b789122f938d5557c54a6b2/xxhash-3.5.0-cp312-cp312-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:f2f2c61bee5844d41c3eb015ac652a0229e901074951ae48581d58bfb2ba01be", size = 428326, upload-time = "2024-08-17T09:18:29.335Z" }, { url = "https://files.pythonhosted.org/packages/11/a7/81dba5010f7e733de88af9555725146fc133be97ce36533867f4c7e75066/xxhash-3.5.0-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:9d32a592cac88d18cc09a89172e1c32d7f2a6e516c3dfde1b9adb90ab5df54a6", size = 194380, upload-time = "2024-08-17T09:18:30.706Z" }, - { url = "https://files.pythonhosted.org/packages/fb/7d/f29006ab398a173f4501c0e4977ba288f1c621d878ec217b4ff516810c04/xxhash-3.5.0-cp312-cp312-manylinux_2_5_i686.manylinux1_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:70dabf941dede727cca579e8c205e61121afc9b28516752fd65724be1355cc90", size = 207934, upload-time = "2024-08-17T09:18:32.133Z" }, { url = "https://files.pythonhosted.org/packages/8a/6e/6e88b8f24612510e73d4d70d9b0c7dff62a2e78451b9f0d042a5462c8d03/xxhash-3.5.0-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:e5d0ddaca65ecca9c10dcf01730165fd858533d0be84c75c327487c37a906a27", size = 216301, upload-time = "2024-08-17T09:18:33.474Z" }, - { url = "https://files.pythonhosted.org/packages/af/51/7862f4fa4b75a25c3b4163c8a873f070532fe5f2d3f9b3fc869c8337a398/xxhash-3.5.0-cp312-cp312-musllinux_1_2_i686.whl", hash = "sha256:3e5b5e16c5a480fe5f59f56c30abdeba09ffd75da8d13f6b9b6fd224d0b4d0a2", size = 203351, upload-time = "2024-08-17T09:18:34.889Z" }, - { url = "https://files.pythonhosted.org/packages/22/61/8d6a40f288f791cf79ed5bb113159abf0c81d6efb86e734334f698eb4c59/xxhash-3.5.0-cp312-cp312-musllinux_1_2_ppc64le.whl", hash = "sha256:149b7914451eb154b3dfaa721315117ea1dac2cc55a01bfbd4df7c68c5dd683d", size = 210294, upload-time = "2024-08-17T09:18:36.355Z" }, - { url = "https://files.pythonhosted.org/packages/17/02/215c4698955762d45a8158117190261b2dbefe9ae7e5b906768c09d8bc74/xxhash-3.5.0-cp312-cp312-musllinux_1_2_s390x.whl", hash = "sha256:eade977f5c96c677035ff39c56ac74d851b1cca7d607ab3d8f23c6b859379cab", size = 414674, upload-time = "2024-08-17T09:18:38.536Z" }, { url = "https://files.pythonhosted.org/packages/31/5c/b7a8db8a3237cff3d535261325d95de509f6a8ae439a5a7a4ffcff478189/xxhash-3.5.0-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:fa9f547bd98f5553d03160967866a71056a60960be00356a15ecc44efb40ba8e", size = 192022, upload-time = "2024-08-17T09:18:40.138Z" }, - { url = "https://files.pythonhosted.org/packages/78/e3/dd76659b2811b3fd06892a8beb850e1996b63e9235af5a86ea348f053e9e/xxhash-3.5.0-cp312-cp312-win32.whl", hash = "sha256:f7b58d1fd3551b8c80a971199543379be1cee3d0d409e1f6d8b01c1a2eebf1f8", size = 30170, upload-time = "2024-08-17T09:18:42.163Z" }, - { url = "https://files.pythonhosted.org/packages/d9/6b/1c443fe6cfeb4ad1dcf231cdec96eb94fb43d6498b4469ed8b51f8b59a37/xxhash-3.5.0-cp312-cp312-win_amd64.whl", hash = "sha256:fa0cafd3a2af231b4e113fba24a65d7922af91aeb23774a8b78228e6cd785e3e", size = 30040, upload-time = "2024-08-17T09:18:43.699Z" }, - { url = "https://files.pythonhosted.org/packages/0f/eb/04405305f290173acc0350eba6d2f1a794b57925df0398861a20fbafa415/xxhash-3.5.0-cp312-cp312-win_arm64.whl", hash = "sha256:586886c7e89cb9828bcd8a5686b12e161368e0064d040e225e72607b43858ba2", size = 26796, upload-time = "2024-08-17T09:18:45.29Z" }, ] [[package]] @@ -4141,32 +4745,29 @@ name = "yarl" version = "1.20.1" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "idna" }, - { name = "multidict" }, - { name = "propcache" }, + { name = "idna", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "multidict", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "propcache", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, ] sdist = { url = "https://files.pythonhosted.org/packages/3c/fb/efaa23fa4e45537b827620f04cf8f3cd658b76642205162e072703a5b963/yarl-1.20.1.tar.gz", hash = "sha256:d017a4997ee50c91fd5466cef416231bb82177b93b029906cefc542ce14c35ac", size = 186428, upload-time = "2025-06-10T00:46:09.923Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/5f/9a/cb7fad7d73c69f296eda6815e4a2c7ed53fc70c2f136479a91c8e5fbdb6d/yarl-1.20.1-cp312-cp312-macosx_10_13_universal2.whl", hash = "sha256:bdcc4cd244e58593a4379fe60fdee5ac0331f8eb70320a24d591a3be197b94a9", size = 133667, upload-time = "2025-06-10T00:43:44.369Z" }, - { url = "https://files.pythonhosted.org/packages/67/38/688577a1cb1e656e3971fb66a3492501c5a5df56d99722e57c98249e5b8a/yarl-1.20.1-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:b29a2c385a5f5b9c7d9347e5812b6f7ab267193c62d282a540b4fc528c8a9d2a", size = 91025, upload-time = "2025-06-10T00:43:46.295Z" }, - { url = "https://files.pythonhosted.org/packages/50/ec/72991ae51febeb11a42813fc259f0d4c8e0507f2b74b5514618d8b640365/yarl-1.20.1-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:1112ae8154186dfe2de4732197f59c05a83dc814849a5ced892b708033f40dc2", size = 89709, upload-time = "2025-06-10T00:43:48.22Z" }, { url = "https://files.pythonhosted.org/packages/99/da/4d798025490e89426e9f976702e5f9482005c548c579bdae792a4c37769e/yarl-1.20.1-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:90bbd29c4fe234233f7fa2b9b121fb63c321830e5d05b45153a2ca68f7d310ee", size = 352287, upload-time = "2025-06-10T00:43:49.924Z" }, - { url = "https://files.pythonhosted.org/packages/1a/26/54a15c6a567aac1c61b18aa0f4b8aa2e285a52d547d1be8bf48abe2b3991/yarl-1.20.1-cp312-cp312-manylinux_2_17_armv7l.manylinux2014_armv7l.manylinux_2_31_armv7l.whl", hash = "sha256:680e19c7ce3710ac4cd964e90dad99bf9b5029372ba0c7cbfcd55e54d90ea819", size = 345429, upload-time = "2025-06-10T00:43:51.7Z" }, - { url = "https://files.pythonhosted.org/packages/d6/95/9dcf2386cb875b234353b93ec43e40219e14900e046bf6ac118f94b1e353/yarl-1.20.1-cp312-cp312-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:4a979218c1fdb4246a05efc2cc23859d47c89af463a90b99b7c56094daf25a16", size = 365429, upload-time = "2025-06-10T00:43:53.494Z" }, - { url = "https://files.pythonhosted.org/packages/91/b2/33a8750f6a4bc224242a635f5f2cff6d6ad5ba651f6edcccf721992c21a0/yarl-1.20.1-cp312-cp312-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:255b468adf57b4a7b65d8aad5b5138dce6a0752c139965711bdcb81bc370e1b6", size = 363862, upload-time = "2025-06-10T00:43:55.766Z" }, { url = "https://files.pythonhosted.org/packages/98/28/3ab7acc5b51f4434b181b0cee8f1f4b77a65919700a355fb3617f9488874/yarl-1.20.1-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:a97d67108e79cfe22e2b430d80d7571ae57d19f17cda8bb967057ca8a7bf5bfd", size = 355616, upload-time = "2025-06-10T00:43:58.056Z" }, - { url = "https://files.pythonhosted.org/packages/36/a3/f666894aa947a371724ec7cd2e5daa78ee8a777b21509b4252dd7bd15e29/yarl-1.20.1-cp312-cp312-manylinux_2_5_i686.manylinux1_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:8570d998db4ddbfb9a590b185a0a33dbf8aafb831d07a5257b4ec9948df9cb0a", size = 339954, upload-time = "2025-06-10T00:43:59.773Z" }, { url = "https://files.pythonhosted.org/packages/f1/81/5f466427e09773c04219d3450d7a1256138a010b6c9f0af2d48565e9ad13/yarl-1.20.1-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:97c75596019baae7c71ccf1d8cc4738bc08134060d0adfcbe5642f778d1dca38", size = 365575, upload-time = "2025-06-10T00:44:02.051Z" }, - { url = "https://files.pythonhosted.org/packages/2e/e3/e4b0ad8403e97e6c9972dd587388940a032f030ebec196ab81a3b8e94d31/yarl-1.20.1-cp312-cp312-musllinux_1_2_armv7l.whl", hash = "sha256:1c48912653e63aef91ff988c5432832692ac5a1d8f0fb8a33091520b5bbe19ef", size = 365061, upload-time = "2025-06-10T00:44:04.196Z" }, - { url = "https://files.pythonhosted.org/packages/ac/99/b8a142e79eb86c926f9f06452eb13ecb1bb5713bd01dc0038faf5452e544/yarl-1.20.1-cp312-cp312-musllinux_1_2_i686.whl", hash = "sha256:4c3ae28f3ae1563c50f3d37f064ddb1511ecc1d5584e88c6b7c63cf7702a6d5f", size = 364142, upload-time = "2025-06-10T00:44:06.527Z" }, - { url = "https://files.pythonhosted.org/packages/34/f2/08ed34a4a506d82a1a3e5bab99ccd930a040f9b6449e9fd050320e45845c/yarl-1.20.1-cp312-cp312-musllinux_1_2_ppc64le.whl", hash = "sha256:c5e9642f27036283550f5f57dc6156c51084b458570b9d0d96100c8bebb186a8", size = 381894, upload-time = "2025-06-10T00:44:08.379Z" }, - { url = "https://files.pythonhosted.org/packages/92/f8/9a3fbf0968eac704f681726eff595dce9b49c8a25cd92bf83df209668285/yarl-1.20.1-cp312-cp312-musllinux_1_2_s390x.whl", hash = "sha256:2c26b0c49220d5799f7b22c6838409ee9bc58ee5c95361a4d7831f03cc225b5a", size = 383378, upload-time = "2025-06-10T00:44:10.51Z" }, { url = "https://files.pythonhosted.org/packages/af/85/9363f77bdfa1e4d690957cd39d192c4cacd1c58965df0470a4905253b54f/yarl-1.20.1-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:564ab3d517e3d01c408c67f2e5247aad4019dcf1969982aba3974b4093279004", size = 374069, upload-time = "2025-06-10T00:44:12.834Z" }, - { url = "https://files.pythonhosted.org/packages/35/99/9918c8739ba271dcd935400cff8b32e3cd319eaf02fcd023d5dcd487a7c8/yarl-1.20.1-cp312-cp312-win32.whl", hash = "sha256:daea0d313868da1cf2fac6b2d3a25c6e3a9e879483244be38c8e6a41f1d876a5", size = 81249, upload-time = "2025-06-10T00:44:14.731Z" }, - { url = "https://files.pythonhosted.org/packages/eb/83/5d9092950565481b413b31a23e75dd3418ff0a277d6e0abf3729d4d1ce25/yarl-1.20.1-cp312-cp312-win_amd64.whl", hash = "sha256:48ea7d7f9be0487339828a4de0360d7ce0efc06524a48e1810f945c45b813698", size = 86710, upload-time = "2025-06-10T00:44:16.716Z" }, { url = "https://files.pythonhosted.org/packages/b4/2d/2345fce04cfd4bee161bf1e7d9cdc702e3e16109021035dbb24db654a622/yarl-1.20.1-py3-none-any.whl", hash = "sha256:83b8eb083fe4683c6115795d9fc1cfaf2cbbefb19b3a1cb68f6527460f483a77", size = 46542, upload-time = "2025-06-10T00:46:07.521Z" }, ] +[[package]] +name = "z3-solver" +version = "4.15.4.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/8a/8e/0c8f17309549d2e5cde9a3ccefa6365437f1e7bafe71878eaf9478e47b18/z3_solver-4.15.4.0.tar.gz", hash = "sha256:928c29b58c4eb62106da51c1914f6a4a55d0441f8f48a81b9da07950434a8946", size = 5018600, upload-time = "2025-10-29T18:12:03.062Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/21/c9/bb51a96af0091324c81b803f16c49f719f9f6ea0b0bb52200f5c97ec4892/z3_solver-4.15.4.0-py3-none-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:7e103a6f203f505b8b8b8e5c931cc407c95b61556512d4921c1ddc0b3f41b08e", size = 29268352, upload-time = "2025-10-29T18:11:53.032Z" }, + { url = "https://files.pythonhosted.org/packages/bf/2e/0b49f7e4e53817cfb09a0f6585012b782dfe0b666e8abefcb4fac0570606/z3_solver-4.15.4.0-py3-none-manylinux_2_34_aarch64.whl", hash = "sha256:62c7e9cbdd711932301f29919ad9158de9b2f58b4d281dd259bbcd0a2f408ba1", size = 27226534, upload-time = "2025-10-29T18:11:55.59Z" }, +] + [[package]] name = "zipp" version = "3.23.0"