[Release RFC] Release process and add staged RC GPU validation

.github/scripts/validate_release_gpu.sh

@@ -0,0 +1,246 @@
+#!/bin/bash
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+set -euo pipefail
+
+required_env_vars=(
+  VALIDATION_SUITE
+  TORCH_VERSION
+  TORCHVISION_VERSION
+  TORCHAO_VERSION
+  PYTORCH_INDEX_URL
+)
+for var in "${required_env_vars[@]}"; do
+  if [[ -z "${!var:-}" ]]; then
+    echo "Missing required environment variable: ${var}"
+    exit 1
+  fi
+done
+
+case "${VALIDATION_SUITE}" in
+  core|graph-trainer|graph-trainer-h100|h100) ;;
+  *)
+    echo "Unknown validation suite: ${VALIDATION_SUITE}" >&2
+    exit 1
+    ;;
+esac
+
+REPO_ROOT="$(git rev-parse --show-toplevel)"
+RC_VERSION="$(tr -d '[:space:]' < "${REPO_ROOT}/assets/version.txt")"
+if [[ "${RC_VERSION}" != *rc* ]]; then
+  echo "Expected an RC version, got: ${RC_VERSION}"
+  exit 1
+fi
+export RC_VERSION TORCH_VERSION TORCHVISION_VERSION TORCHAO_VERSION
+
+eval "$(conda shell.bash hook)"
+conda activate "$(conda env list --json | jq -r '.envs[-1]')"
+
+export HF_HOME="${RUNNER_TEMP}/hf_home"
+export PYTHONUNBUFFERED=1
+
+DRIVER_VERSION=$(nvidia-smi --query-gpu=driver_version --format=csv,noheader | head -n 1 || true)
+echo "CUDA driver version: ${DRIVER_VERSION}"
+nvidia-smi
+
+python -m pip uninstall -y torch torchvision torchao torchtitan
+
+USE_CPP=0 PIP_EXTRA_INDEX_URL='' python -m pip install --pre \
+  --index-url "${PYTORCH_INDEX_URL}" \
+  "torch==${TORCH_VERSION}" \
+  "torchvision==${TORCHVISION_VERSION}" \
+  "torchao==${TORCHAO_VERSION}"
+
+install_release_candidate() {
+  for attempt in {1..5}; do
+    echo "Installing torchtitan==${RC_VERSION} from TestPyPI (attempt ${attempt}/5)"
+    if python -m pip install \
+      --index-url https://test.pypi.org/simple/ \
+      --extra-index-url https://pypi.org/simple/ \
+      "torchtitan==${RC_VERSION}"; then
+      return 0
+    fi
+    if [[ "${attempt}" -lt 5 ]]; then
+      sleep 15
+    fi
+  done
+
+  echo "Failed to install torchtitan==${RC_VERSION} after 5 attempts" >&2
+  return 1
+}
+install_release_candidate
+
+ARTIFACTS="${RUNNER_TEMP}/artifacts-to-be-uploaded"
+VALIDATION_ROOT="${RUNNER_TEMP}/validate-release-gpu"
+mkdir -p \
+  "${ARTIFACTS}" \
+  "${VALIDATION_ROOT}/scripts" \
+  "${VALIDATION_ROOT}/torchtitan/experiments/graph_trainer" \
+  "${VALIDATION_ROOT}/torchtitan/models/flux"
+
+# Copy test support files without copying TorchTitan Python packages. Training
+# subprocesses must import the installed TestPyPI RC rather than the checkout.
+cp "${REPO_ROOT}/run_train.sh" "${VALIDATION_ROOT}/run_train.sh"
+cp "${REPO_ROOT}/scripts/loss_compare.py" "${VALIDATION_ROOT}/scripts/loss_compare.py"
+cp -R "${REPO_ROOT}/tests" "${VALIDATION_ROOT}/tests"
+cp \
+  "${REPO_ROOT}/torchtitan/models/flux/run_infer.sh" \
+  "${VALIDATION_ROOT}/torchtitan/models/flux/run_infer.sh"
+cp \
+  "${REPO_ROOT}/torchtitan/experiments/graph_trainer/run_train_precompile.sh" \
+  "${VALIDATION_ROOT}/torchtitan/experiments/graph_trainer/run_train_precompile.sh"
+
+cd "${VALIDATION_ROOT}"
+unset PYTHONPATH
+
+python - <<'PY'
+import os
+
+import torch
+import torchao
+import torchtitan
+import torchvision
+import triton
+
+assert torchtitan.__version__ == os.environ["RC_VERSION"]
+assert "site-packages" in torchtitan.__file__
+assert torch.__version__.split("+")[0] == os.environ["TORCH_VERSION"]
+assert torchvision.__version__.split("+")[0] == os.environ["TORCHVISION_VERSION"]
+assert torchao.__version__.split("+")[0] == os.environ["TORCHAO_VERSION"]
+assert torch.version.cuda == "13.0"
+assert torch.cuda.is_available()
+assert torch.cuda.device_count() >= 8
+print(f"torchtitan={torchtitan.__version__} ({torchtitan.__file__})")
+print(f"torch={torch.__version__} ({torch.version.git_version})")
+print(f"torchvision={torchvision.__version__}")
+print(f"torchao={torchao.__version__}")
+print(f"triton={triton.__version__}")
+print(f"num_cuda_devices={torch.cuda.device_count()}")
+PY
+
+TORCHTITAN_PACKAGE_ROOT="$(python - <<'PY'
+from pathlib import Path
+
+import torchtitan
+
+print(Path(torchtitan.__file__).parent)
+PY
+)"
+
+run_core_tests() {
+  python -m pytest tests/unit_tests/flex_shard/test_dist_muon.py \
+    --durations=20 -vv
+
+  python scripts/loss_compare.py . . \
+    --baseline-options="--parallelism.data_parallel_replicate_degree=1" \
+    --baseline-ngpus=8 \
+    --steps=100 \
+    --job-dump-folder="${ARTIFACTS}/llama3_golden" \
+    --import-result=tests/assets/losses/llama3_cuda.txt \
+    --assert-equal
+
+  python scripts/loss_compare.py . . \
+    --baseline-options="--parallelism.data_parallel_replicate_degree=1" \
+    --test-options="--parallelism.data_parallel_replicate_degree=4" \
+    --baseline-ngpus=8 \
+    --test-ngpus=8 \
+    --steps=1 \
+    --job-dump-folder="${ARTIFACTS}/llama3_fsdp_hsdp_parity" \
+    --assert-equal
+
+  python scripts/loss_compare.py . . \
+    --baseline-module=qwen3 \
+    --baseline-config=qwen3_moe_debug \
+    --baseline-options="--parallelism.tensor_parallel_degree=2 --parallelism.expert_parallel_degree=4 --parallelism.spmd_backend=spmd_types --training.disable_cuda_graphs" \
+    --baseline-ngpus=8 \
+    --steps=100 \
+    --job-dump-folder="${ARTIFACTS}/qwen3_moe_golden" \
+    --import-result=tests/assets/losses/qwen3_moe_cuda.txt \
+    --assert-equal
+
+  python -m tests.integration_tests.run_tests \
+    --gpu_arch_type cuda \
+    --test_suite features \
+    "${ARTIFACTS}/feature_tests" \
+    --ngpu 8
+
+  python -m tests.integration_tests.run_tests \
+    --gpu_arch_type cuda \
+    --test_suite models \
+    "${ARTIFACTS}/model_tests" \
+    --ngpu 8
+
+  python -m tests.integration_tests.flux \
+    "${ARTIFACTS}/flux_tests" \
+    --ngpu 8
+}
+
+run_graph_trainer_tests() {
+  local tests_root="${TORCHTITAN_PACKAGE_ROOT}/experiments/graph_trainer/tests"
+
+  python -m torchtitan.experiments.graph_trainer.tests.integration_tests \
+    --test_suite graph_trainer_default \
+    --gpu_arch_type cuda \
+    "${ARTIFACTS}/graph_trainer" \
+    --ngpu 8
+  python -m pytest "${tests_root}/test_numerics.py::TestSimpleFSDP" -v
+  python -m pytest "${tests_root}/test_numerics.py::TestGraphTrainerNumerics" \
+    -v -k dense
+  python -m pytest "${tests_root}/test_passes.py" -v
+  python -m pytest "${tests_root}/test_profiler.py" -v
+  python -m pytest "${tests_root}/test_trace_module.py" -v
+  python -m pytest "${tests_root}/test_precompile.py" -v
+  python -m torchtitan.experiments.graph_trainer.tests.run_precompile_tests \
+    "${ARTIFACTS}/graph_trainer_precompile" \
+    --ngpu 8 \
+    --test_name aot_fx_trace_llama3_precompile_fsdp_tp
+  python -m pytest "${tests_root}/test_bitwise_deterministic.py" -v
+  python -m pytest "${tests_root}/test_sac_peak_memory.py" -v
+}
+
+install_deepep() {
+  export CUDA_HOME=/usr/local/cuda
+  export NCCL_NVLS_ENABLE=0
+  export TORCH_SHOW_CPP_STACKTRACES=1
+  bash /install_deepep.sh
+}
+
+run_h100_tests() {
+  install_deepep
+  python -m tests.integration_tests.run_tests \
+    --test_suite h100 \
+    --gpu_arch_type cuda \
+    "${ARTIFACTS}/h100" \
+    --ngpu 8
+}
+
+run_graph_trainer_h100_tests() {
+  local tests_root="${TORCHTITAN_PACKAGE_ROOT}/experiments/graph_trainer/tests"
+
+  install_deepep
+  python -m torchtitan.experiments.graph_trainer.tests.integration_tests \
+    --test_suite graph_trainer_h100 \
+    --gpu_arch_type cuda \
+    "${ARTIFACTS}/graph_trainer_h100" \
+    --ngpu 8
+  python -m pytest "${tests_root}/test_numerics.py::TestGraphTrainerNumerics" \
+    -v -k moe
+  python -m torchtitan.experiments.graph_trainer.tests.run_precompile_tests \
+    "${ARTIFACTS}/graph_trainer_h100_precompile" \
+    --ngpu 8 \
+    --test_name aot_fx_trace_deepseek_v3_precompile_fsdp_tp_ep
+  python -m pytest "${tests_root}/test_bitwise_deterministic.py" -v
+}
+
+case "${VALIDATION_SUITE}" in
+  core) run_core_tests ;;
+  graph-trainer) run_graph_trainer_tests ;;
+  graph-trainer-h100) run_graph_trainer_h100_tests ;;
+  h100) run_h100_tests ;;
+esac
+
+find "${ARTIFACTS}" -type d -name checkpoint -prune -exec rm -rf {} +

.github/workflows/validate_release_gpu.yml

@@ -0,0 +1,113 @@
+name: Validate a TestPyPI Release Candidate on GPUs
+
+on:
+  workflow_dispatch:
+
+permissions:
+  contents: read
+  id-token: write
+
+# RC GPU validation coverage:
+# - validate-standard-gpu: 8x A10G; core golden losses, models, features,
+#   Flux, and the standard GraphTrainer suite.
+# - validate-h100: 8x H100; core H100 tests, including DeepSeek HybridEP/DeepEP.
+# - validate-graph-trainer-h100: 8x H100; GraphTrainer DeepSeek and Qwen3 MoE
+#   tests that require H100-class hardware.
+jobs:
+  validate-version:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v7
+
+      - name: Verify RC version
+        run: |
+          version=$(tr -d '\n' < assets/version.txt)
+          case "$version" in
+            *rc*) echo "Validating RC version: $version" ;;
+            *) echo "Expected an RC version, got: $version"; exit 1 ;;
+          esac
+
+  validate-standard-gpu:
+    name: Validate ${{ matrix.validation-suite }} on 8x A10G
+    needs:
+      - validate-version
+    uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main
+    strategy:
+      fail-fast: false
+      matrix:
+        validation-suite:
+          - core
+          - graph-trainer
+    with:
+      runner: linux.g5.48xlarge.nvidia.gpu
+      gpu-arch-type: cuda
+      gpu-arch-version: "13.0"
+      docker-image: torchtitan-ubuntu-22.04-clang12
+      repository: pytorch/torchtitan
+      upload-artifact: outputs
+      timeout: 180
+      script: |
+        VALIDATION_SUITE="${{ matrix.validation-suite }}" \
+        TORCH_VERSION="2.14.0" \
+        TORCHVISION_VERSION="0.29.0" \
+        TORCHAO_VERSION="0.18.0" \
+        PYTORCH_INDEX_URL="https://download.pytorch.org/whl/test/cu130" \
+        bash .github/scripts/validate_release_gpu.sh
+
+  set-h100-matrix:
+    needs:
+      - validate-version
+    uses: ./.github/workflows/set-matrix.yaml
+    with:
+      runner-cuda: mt-l-bx86iamx-176-1800-h100-8
+      gpu-arch: cuda
+
+  validate-h100:
+    name: Validate core H100 tests
+    needs:
+      - set-h100-matrix
+    if: ${{ needs.set-h100-matrix.outputs.matrix != '' && fromJSON(needs.set-h100-matrix.outputs.matrix).include[0] != null }}
+    uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main
+    strategy:
+      fail-fast: false
+      matrix: ${{ fromJSON(needs.set-h100-matrix.outputs.matrix) }}
+    with:
+      runner: ${{ matrix.runner }}
+      gpu-arch-type: ${{ matrix.gpu-arch-type }}
+      gpu-arch-version: ${{ matrix.gpu-arch-version }}
+      docker-image: 308535385114.dkr.ecr.us-east-1.amazonaws.com/torchtitan/${{ matrix.docker-image }}:${{ needs.set-h100-matrix.outputs.docker-hash }}
+      repository: pytorch/torchtitan
+      upload-artifact: outputs
+      timeout: 180
+      script: |
+        VALIDATION_SUITE="h100" \
+        TORCH_VERSION="2.14.0" \
+        TORCHVISION_VERSION="0.29.0" \
+        TORCHAO_VERSION="0.18.0" \
+        PYTORCH_INDEX_URL="https://download.pytorch.org/whl/test/cu130" \
+        bash .github/scripts/validate_release_gpu.sh
+
+  validate-graph-trainer-h100:
+    name: Validate GraphTrainer H100 tests
+    needs:
+      - set-h100-matrix
+    if: ${{ needs.set-h100-matrix.outputs.matrix != '' && fromJSON(needs.set-h100-matrix.outputs.matrix).include[0] != null }}
+    uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main
+    strategy:
+      fail-fast: false
+      matrix: ${{ fromJSON(needs.set-h100-matrix.outputs.matrix) }}
+    with:
+      runner: ${{ matrix.runner }}
+      gpu-arch-type: ${{ matrix.gpu-arch-type }}
+      gpu-arch-version: ${{ matrix.gpu-arch-version }}
+      docker-image: 308535385114.dkr.ecr.us-east-1.amazonaws.com/torchtitan/${{ matrix.docker-image }}:${{ needs.set-h100-matrix.outputs.docker-hash }}
+      repository: pytorch/torchtitan
+      upload-artifact: outputs
+      timeout: 180
+      script: |
+        VALIDATION_SUITE="graph-trainer-h100" \
+        TORCH_VERSION="2.14.0" \
+        TORCHVISION_VERSION="0.29.0" \
+        TORCHAO_VERSION="0.18.0" \
+        PYTORCH_INDEX_URL="https://download.pytorch.org/whl/test/cu130" \
+        bash .github/scripts/validate_release_gpu.sh