Skip to content

feat(trainer): support VLM training with dynamo #14419

feat(trainer): support VLM training with dynamo

feat(trainer): support VLM training with dynamo #14419

Workflow file for this run

name: CPU Tests
on:
push:
branches: [main]
pull_request:
types: [opened, synchronize, reopened, ready_for_review]
jobs:
unit-tests:
name: Unit tests
runs-on: ubuntu-latest
steps:
- name: Checkout repository
uses: actions/checkout@v7
with:
submodules: false
- name: Init public submodules
run: |
git config --global url."https://github.com/".insteadOf "git@github.com:"
git submodule update --init --recursive -- deps/verifiers deps/renderers deps/prime-envs deps/pydantic-config
- name: Install uv
uses: astral-sh/setup-uv@v10.0.1
with:
enable-cache: true
cache-dependency-glob: "uv.lock"
- name: Install dependencies
run: uv sync --all-extras --all-packages --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/unit -m "not gpu"
- name: Cleanup output_dir
run: rm -rf /tmp/outputs
slim-configs-install:
name: Slim install (prime-rl-configs only)
runs-on: ubuntu-latest
steps:
- name: Checkout repository
uses: actions/checkout@v7
- name: Install uv
uses: astral-sh/setup-uv@v10.0.1
with:
enable-cache: true
cache-dependency-glob: "uv.lock"
- 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
# `verifiers` (+ its `datasets` dep) is a declared slim dep: the v1 config types
# (EnvConfig, Task, ...) extend verifiers.v1, which is pure-pydantic and pulls no
# GPU/ML deps. We still forbid the actual heavy training deps below.
forbidden = ["torch", "transformers", "vllm", "wandb", "ring_flash_attn",
"prime", "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 real SFT config (with [slurm]) parses on slim install
run: |
/tmp/slim/bin/python - <<'PY'
from pydantic_config import cli
from prime_rl.configs.sft import SFTConfig
# The slim wheel intentionally excludes taskset packages, so validate a real
# SFT config here. RL recipes are covered by the full config suite.
for cfg_cls, args in [
(SFTConfig, ["@", "configs/ci/slim_sft.toml"]),
]:
cfg = cli(cfg_cls, args=args)
assert cfg.slurm is not None, f"expected [slurm] in {args}"
# On slim, template_path can't be auto-resolved (heavy wheel absent) — must stay None.
assert cfg.slurm.template_path is None, f"unexpected template_path on slim: {cfg.slurm.template_path}"
print("slim parse OK")
PY