Add blockwise FP8 training recipe for linears and MoE grouped experts #1024
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 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 | |
| # Run E2E RL integration tests (up to 8 GPUs). | |
| # Includes the MoE TP=4 EP=4 random_init test which needs 8 GPUs. | |
| python -m torchtitan.experiments.rl.tests.integration_tests \ | |
| $RUNNER_TEMP/artifacts-to-be-uploaded --ngpu 8 \ | |
| --hf_assets_path "$MODEL_PATH" |