Skip to content

GraphTrainer 8 GPU Integration Tests #41

GraphTrainer 8 GPU Integration Tests

GraphTrainer 8 GPU Integration Tests #41

name: GraphTrainer 8 GPU Integration Tests
on:
push:
branches: [ main ]
paths:
- 'torchtitan/experiments/graph_trainer/**'
- '.github/workflows/integration_test_8gpu_graph_trainer.yaml'
pull_request:
types: [opened, synchronize, reopened, ready_for_review]
paths:
- 'torchtitan/experiments/graph_trainer/**'
- '.github/workflows/integration_test_8gpu_graph_trainer.yaml'
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:
build-test:
# 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' || github.event.pull_request.draft == false)
uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main
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: 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 GPU info / 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
# torch nightly now depends on cuda-toolkit[...] extras (e.g.
# nvidia-curand==10.4.0.35) whose wheels live on PyPI, not on the pytorch
# nightly index. --index-url *replaces* the default PyPI index, so add it
# back via --extra-index-url or pip cannot resolve those cuda deps.
python -m pip install --force-reinstall --pre \
torch \
--index-url https://download.pytorch.org/whl/nightly/cu130 \
--extra-index-url https://pypi.org/simple
python -m pip install git+https://github.com/meta-pytorch/autoparallel.git
sudo mkdir -p "$RUNNER_TEMP/artifacts-to-be-uploaded"
sudo chown -R $(id -u):$(id -g) "$RUNNER_TEMP/artifacts-to-be-uploaded"
python -m torchtitan.experiments.graph_trainer.tests.integration_tests --test_suite graph_trainer_default --gpu_arch_type cuda $RUNNER_TEMP/artifacts-to-be-uploaded --ngpu 8
python -m torchtitan.experiments.graph_trainer.tests.integration_tests --test_suite graph_trainer_autoparallel --gpu_arch_type cuda $RUNNER_TEMP/artifacts-to-be-uploaded/autoparallel --ngpu 4
# Run the numerics unit tests. Standard MoE and DeepSeek V3 AutoParallel
# tests run in the H100 workflow.
pytest torchtitan/experiments/graph_trainer/tests/test_numerics.py::TestSimpleFSDP -v
pytest torchtitan/experiments/graph_trainer/tests/test_numerics.py::TestGraphTrainerNumerics -v -k "dense"
pytest torchtitan/experiments/graph_trainer/tests/test_numerics.py::TestGraphTrainerAutoParallelNumerics::test_llama3_aot_fx_trace_autoparallel_vs_eager -v
# Run the graph passes unit tests
pytest torchtitan/experiments/graph_trainer/tests/test_passes.py -v
# Run kernel annotation profiler tests (skips if cuda-compat < 13.1)
pytest torchtitan/experiments/graph_trainer/tests/test_profiler.py -v
# Run the make_fx tracer unit tests
pytest torchtitan/experiments/graph_trainer/tests/test_trace_module.py -v
# Run precompile unit tests
pytest torchtitan/experiments/graph_trainer/tests/test_precompile.py -v
# Run precompile integration tests (Llama3 only; DSv3 runs in H100 workflow)
python -m torchtitan.experiments.graph_trainer.tests.run_precompile_tests $RUNNER_TEMP/artifacts-to-be-uploaded/precompile --ngpu 8 --test_name aot_fx_trace_llama3_precompile_fsdp_tp
# Run bitwise deterministic and SAC peak-memory guardrail tests
pytest torchtitan/experiments/graph_trainer/tests/test_bitwise_deterministic.py -v
pytest torchtitan/experiments/graph_trainer/tests/test_sac_peak_memory.py -v
rm -rf $RUNNER_TEMP/artifacts-to-be-uploaded/*/checkpoint