RL Integration Tests #1077
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| name: RL Integration Tests | |
| on: | |
| push: | |
| branches: [ main ] | |
| paths: | |
| - 'torchtitan/experiments/rl/**' | |
| - '.github/workflows/integration_test_8gpu_rl.yaml' | |
| pull_request: | |
| types: [labeled, synchronize] | |
| branches: [ main ] | |
| schedule: | |
| # Runs every 12 hours | |
| - cron: '0 */12 * * *' | |
| concurrency: | |
| group: unit-test${{ github.workflow }}-${{ github.ref == 'refs/heads/main' && github.run_number || github.ref }} | |
| cancel-in-progress: true | |
| defaults: | |
| run: | |
| shell: bash -l -eo pipefail {0} | |
| permissions: | |
| id-token: write | |
| contents: read | |
| jobs: | |
| # Step 1: Dynamically compute the matrix based on conditions | |
| set-matrix: | |
| # Skip scheduled runs on forks, where they would only fail and email the fork owner | |
| if: >- | |
| (github.repository_owner == 'pytorch' || github.event_name != 'schedule') && | |
| (github.event_name != 'pull_request' || contains(github.event.pull_request.labels.*.name, 'ciflow/rl')) | |
| uses: ./.github/workflows/set-matrix.yaml | |
| with: | |
| is-experimental: true | |
| # Step 2: Use the dynamic matrix in the build-test job | |
| build-test: | |
| needs: set-matrix | |
| uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main | |
| strategy: | |
| fail-fast: false | |
| matrix: ${{ fromJSON(needs.set-matrix.outputs.matrix) }} | |
| with: | |
| runner: ${{ matrix.runner }} | |
| gpu-arch-type: ${{ matrix.gpu-arch-type }} | |
| gpu-arch-version: ${{ matrix.gpu-arch-version }} | |
| docker-image: ${{ matrix.docker-image }} | |
| repository: pytorch/torchtitan | |
| upload-artifact: outputs | |
| timeout: 60 | |
| script: | | |
| set -eux | |
| # The generic Linux job chooses to use base env, not the one setup by the image | |
| CONDA_ENV=$(conda env list --json | jq -r ".envs | .[-1]") | |
| conda activate "${CONDA_ENV}" | |
| # Log CUDA driver version for debugging. | |
| DRIVER_VERSION=$(nvidia-smi --query-gpu=driver_version --format=csv,noheader | head -n 1 || true) | |
| echo "CUDA driver version: ${DRIVER_VERSION}" | |
| pip config --user set global.progress_bar off | |
| # Install uv for faster dependency resolution | |
| pip install uv | |
| # 1. Install Monarch, TorchStore, and Renderers | |
| uv pip install torchmonarch | |
| uv pip install --no-deps "git+https://github.com/meta-pytorch/torchstore.git@main" | |
| uv pip install pygtrie portpicker | |
| uv pip install "git+https://github.com/PrimeIntellect-ai/renderers.git@main" | |
| # 2. Install batch-invariant ops | |
| uv pip install --no-deps "git+https://github.com/thinking-machines-lab/batch_invariant_ops.git@main" | |
| # 3. Install PyTorch nightly, vllm, and xformers | |
| # torchvision must be installed from the nightly channel alongside torch | |
| # so its C++ extensions (e.g. torchvision::nms) match the torch ABI. | |
| uv pip install torch torchvision vllm xformers --pre \ | |
| --extra-index-url https://download.pytorch.org/whl/nightly/cu130 \ | |
| --index-strategy unsafe-best-match | |
| # 4. Make the checkout importable for subprocesses spawned by the test. | |
| export PYTHONPATH="$PWD:${PYTHONPATH:-}" | |
| # 5. Download HF model checkpoint for tests | |
| MODEL_PATH=$(python -c "from huggingface_hub import snapshot_download; print(snapshot_download('Qwen/Qwen3-0.6B'))") | |
| sudo mkdir -p "$RUNNER_TEMP/artifacts-to-be-uploaded" | |
| sudo chown -R $(id -u):$(id -g) "$RUNNER_TEMP/artifacts-to-be-uploaded" | |
| # Install nvcc so FlashInfer can JIT compile its CUDA kernels. | |
| # vLLM uses FlashInfer sampler by default (vllm-project/vllm#40376) | |
| # but the CI docker image only has CUDA runtime, not the toolkit. | |
| uv pip install nvidia-cuda-nvcc | |
| # RL Loss Guard runs FIRST, in a clean container, so a lingering port from | |
| # the E2E suite can't collide with the generator's distributed init | |
| # (EADDRINUSE). Deterministic loss regression check on the REAL Qwen3-0.6B | |
| # (--hf-assets-path=$MODEL_PATH): real weights -> varied rewards -> non-zero | |
| # advantage -> meaningful, full-precision losses (random init gives all | |
| # zeros). The golden is arch-specific; regenerate on A10G if it changes. | |
| LOSS_FILE="torchtitan/experiments/rl/tests/assets/losses/rl_grpo_cuda.txt" | |
| # Override to trainer TP=4 + 1 generator TP=4: the config's TP=2 OOMs with | |
| # batch-invariant on A10G; the golden is from TP=4. | |
| python torchtitan/experiments/rl/scripts/loss_compare.py \ | |
| --hf-assets-path "$MODEL_PATH" \ | |
| --dump-folder "$RUNNER_TEMP/artifacts-to-be-uploaded/rl_loss_guard" \ | |
| --config rl_grpo_qwen3_0_6b_varlen_batch_invariant \ | |
| --import-result "$LOSS_FILE" \ | |
| --assert-equal \ | |
| --options='--trainer.parallelism.tensor_parallel_degree 4 --generator.parallelism.tensor_parallel_degree 4 --num_generators 1' | |
| # Bitwise trainer/generator parity tests (batch-invariant). Varlen and Flex | |
| # (both TP=2) run as separate torchrun invocations so each backend gets a | |
| # fresh process group and vLLM engine; nproc matches each config's | |
| # tensor_parallel_degree (the test builds ParallelDims from it). | |
| RL_TEST_DUMP_FOLDER="$RUNNER_TEMP/artifacts-to-be-uploaded" \ | |
| HF_ASSETS_PATH="$MODEL_PATH" torchrun --nproc-per-node=2 -m pytest \ | |
| torchtitan/experiments/rl/tests/test_bitwise_parity.py::TestBitwiseParityVarlen -v | |
| RL_TEST_DUMP_FOLDER="$RUNNER_TEMP/artifacts-to-be-uploaded" \ | |
| HF_ASSETS_PATH="$MODEL_PATH" torchrun --nproc-per-node=2 -m pytest \ | |
| torchtitan/experiments/rl/tests/test_bitwise_parity.py::TestBitwiseParityFlex -v | |
| # Run E2E RL integration tests (up to 8 GPUs): includes the MoE TP=4 EP=4 | |
| # random-init test and the batch-invariant debug tests (both need 8 GPUs). | |
| python -m torchtitan.experiments.rl.tests.integration_tests \ | |
| $RUNNER_TEMP/artifacts-to-be-uploaded --ngpu 8 \ | |
| --hf_assets_path "$MODEL_PATH" |