Add NVFP4 quantization-aware distillation (QAD) #2655
Workflow file for this run
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 H100 Integration Tests | |
| on: | |
| push: | |
| branches: [ main ] | |
| paths: | |
| - 'torchtitan/experiments/rl/**' | |
| - '.github/workflows/integration_test_8gpu_rl_h100.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: | |
| runner-cuda: mt-l-bx86iamx-176-1800-h100-8 | |
| 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_v3.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 }} | |
| # Private-ECR image, tagged with the .ci/docker hash from set-matrix.yaml. | |
| docker-image: 308535385114.dkr.ecr.us-east-1.amazonaws.com/torchtitan/${{ matrix.docker-image }}:${{ needs.set-matrix.outputs.docker-hash }} | |
| repository: pytorch/torchtitan | |
| upload-artifact: outputs | |
| timeout: 45 | |
| 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 Flash Attention 3 (required for H100/H200) | |
| uv pip install flash-attn-3 --extra-index-url=https://download.pytorch.org/whl/test/cu130 | |
| # 4. 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. | |
| # Clear PIP_EXTRA_INDEX_URL so the default cpu index can't supply a +cpu torch. | |
| PIP_EXTRA_INDEX_URL= uv pip install torch torchvision vllm xformers --pre \ | |
| --extra-index-url ${{ matrix.index-url }} \ | |
| --index-strategy unsafe-best-match | |
| # 5. Make the checkout importable for subprocesses spawned by the test. | |
| export PYTHONPATH="$PWD:${PYTHONPATH:-}" | |
| # 6. Download HF model checkpoint for tests | |
| MODEL_PATH=$(python -c "from huggingface_hub import snapshot_download; print(snapshot_download('Qwen/Qwen3-0.6B'))") | |
| mkdir -p "$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 | |
| # Run bitwise parity tests (TP=2, batch-invariant mode). Varlen and | |
| # Flex are intentionally separate torchrun invocations so each backend | |
| # gets a fresh process group and vLLM engine lifecycle. | |
| HF_ASSETS_PATH="$MODEL_PATH" torchrun --nproc-per-node=2 -m pytest \ | |
| torchtitan/experiments/rl/tests/test_bitwise_parity.py::TestBitwiseParityVarlen -v | |
| HF_ASSETS_PATH="$MODEL_PATH" torchrun --nproc-per-node=2 -m pytest \ | |
| torchtitan/experiments/rl/tests/test_bitwise_parity.py::TestBitwiseParityFlex -v | |
| # Run batch-invariant E2E RL integration tests (up to 8 GPUs). | |
| # Includes the MoE TP=4 EP=4 batch-invariant bitwise-identity test. | |
| python -m torchtitan.experiments.rl.tests.integration_tests \ | |
| $RUNNER_TEMP/artifacts-to-be-uploaded --test_suite h100 --ngpu 8 \ | |
| --hf_assets_path "$MODEL_PATH" | |
| # RL Loss Guard: deterministic loss regression check | |
| LOSS_FILE="torchtitan/experiments/rl/tests/assets/losses/rl_grpo_cuda.txt" | |
| python torchtitan/experiments/rl/scripts/loss_compare.py \ | |
| --hf-assets-path "$MODEL_PATH" \ | |
| --dump-folder "$RUNNER_TEMP/rl_loss_guard" \ | |
| --config rl_grpo_qwen3_0_6b_varlen_batch_invariant \ | |
| --import-result "$LOSS_FILE" \ | |
| --assert-equal |