diff --git a/.github/CODEOWNERS b/.github/CODEOWNERS new file mode 100644 index 000000000..53b3ea467 --- /dev/null +++ b/.github/CODEOWNERS @@ -0,0 +1 @@ +/src/art/trainer_rank/** @bradhilton diff --git a/.github/workflows/trainer-rank-gpu.yml b/.github/workflows/trainer-rank-gpu.yml new file mode 100644 index 000000000..5fe8be13c --- /dev/null +++ b/.github/workflows/trainer-rank-gpu.yml @@ -0,0 +1,137 @@ +name: TrainerRank GPU validation + +on: + pull_request: + types: [opened, reopened, synchronize, ready_for_review] + +permissions: + contents: read + +concurrency: + group: trainer-rank-gpu-${{ github.event.pull_request.number }} + cancel-in-progress: true + +jobs: + classify: + runs-on: ubuntu-latest + outputs: + required: ${{ steps.changes.outputs.required }} + steps: + - uses: actions/checkout@v4 + with: + fetch-depth: 0 + + - id: changes + name: Classify changed files + env: + BASE_SHA: ${{ github.event.pull_request.base.sha }} + HEAD_SHA: ${{ github.event.pull_request.head.sha }} + run: | + set -euo pipefail + pattern='^src/art/megatron/(prefix_tree(_packing|_state)?\.py|context_parallel/|flex_attn/|gdn/|lora\.py|megatron_patches\.py|training/(finalize_grads|microbatches)\.py)' + changed="$(git diff --name-only "${BASE_SHA}...${HEAD_SHA}")" + critical="$(printf '%s\n' "${changed}" | grep -E "${pattern}" || true)" + if [ -n "${critical}" ]; then + echo "required=true" >> "${GITHUB_OUTPUT}" + { + echo "### TrainerRank GPU validation required" + echo + echo '```text' + printf '%s\n' "${critical}" + echo '```' + } >> "${GITHUB_STEP_SUMMARY}" + else + echo "required=false" >> "${GITHUB_OUTPUT}" + echo "### TrainerRank GPU validation not required" >> "${GITHUB_STEP_SUMMARY}" + fi + + validate: + name: Run on 2x H200 + needs: classify + if: needs.classify.outputs.required == 'true' + runs-on: ubuntu-latest + timeout-minutes: 45 + environment: trainer-rank-gpu-validation + env: + SKY_INFRA: k8s/cks-wb3 + steps: + - uses: actions/checkout@v4 + with: + ref: ${{ github.event.pull_request.head.sha }} + + - name: Install launch dependencies + run: | + curl -LsSf https://astral.sh/uv/install.sh | sh + echo "${HOME}/.local/bin" >> "${GITHUB_PATH}" + sudo apt-get update + sudo apt-get install -y --no-install-recommends socat + kubectl_version="$(curl -LsSf https://dl.k8s.io/release/stable.txt)" + kubectl_path="${RUNNER_TEMP}/kubectl" + curl -LsSf -o "${kubectl_path}" \ + "https://dl.k8s.io/release/${kubectl_version}/bin/linux/amd64/kubectl" + sudo install -m 0755 "${kubectl_path}" /usr/local/bin/kubectl + + - name: Configure Kubernetes + env: + KUBECONFIG_VALUE: ${{ secrets.CKS_WB3_KUBECONFIG }} + run: | + set -euo pipefail + test -n "${KUBECONFIG_VALUE}" + mkdir -p "${HOME}/.kube" + KUBECONFIG_PATH="${HOME}/.kube/config" python3 - <<'PY' + import base64 + import binascii + import os + from pathlib import Path + + value = os.environ["KUBECONFIG_VALUE"].encode() + try: + decoded = base64.b64decode(b"".join(value.split()), validate=True) + data = decoded if decoded.lstrip().startswith(b"apiVersion:") else value + except binascii.Error: + data = value + Path(os.environ["KUBECONFIG_PATH"]).write_bytes(data) + PY + chmod 600 "${HOME}/.kube/config" + kubectl --context cks-wb3 get nodes >/dev/null + + - name: Run TrainerRank GPU validation + run: | + set -euo pipefail + sky_venv="${RUNNER_TEMP}/skypilot" + uv venv --python 3.10 "${sky_venv}" + uv pip install --python "${sky_venv}/bin/python" \ + 'skypilot[kubernetes,remote]==0.12.0' 'kubernetes==34.1.0' + sky="${sky_venv}/bin/sky" + cluster="trainer-rank-gpu-${GITHUB_RUN_ID}-${GITHUB_RUN_ATTEMPT}" + cleanup() { "${sky}" down -y "${cluster}" || true; } + trap cleanup EXIT + "${sky}" check kubernetes + "${sky}" launch -y -c "${cluster}" \ + --infra "${SKY_INFRA}" \ + scripts/ci/trainer-rank-gpu.sky.yaml + + gate: + name: trainer-rank-gpu-validation + needs: [classify, validate] + if: always() + runs-on: ubuntu-latest + steps: + - env: + CLASSIFY_RESULT: ${{ needs.classify.result }} + REQUIRED: ${{ needs.classify.outputs.required }} + RESULT: ${{ needs.validate.result }} + run: | + if [ "${CLASSIFY_RESULT}" != "success" ]; then + echo "TrainerRank GPU classification failed." >&2 + exit 1 + fi + if [ "${REQUIRED}" = "true" ] && [ "${RESULT}" != "success" ]; then + echo "TrainerRank GPU validation required but result was ${RESULT}." >&2 + exit 1 + fi + if [ "${REQUIRED}" = "true" ]; then + echo "TrainerRank GPU validation passed." + else + echo "TrainerRank GPU validation was not required." + fi diff --git a/scripts/ci/trainer-rank-gpu-tests.sh b/scripts/ci/trainer-rank-gpu-tests.sh new file mode 100755 index 000000000..d5ed0e2b9 --- /dev/null +++ b/scripts/ci/trainer-rank-gpu-tests.sh @@ -0,0 +1,23 @@ +#!/usr/bin/env bash +set -euo pipefail + +export CUDA_VISIBLE_DEVICES=0,1 +export PYTHONUNBUFFERED=1 + +uv run --no-sync pytest --tb=short \ + tests/integration/megatron/cp_attn/test_attention_packed_vs_flattened.py \ + 'tests/integration/megatron/gdn_shared_prefix/test_gdn_cp_packed_correctness.py::test_gdn_cp_packed_sibling_order_matches_cp1_oracle[2]' \ + 'tests/integration/megatron/gdn_shared_prefix/test_gdn_cp_packed_correctness.py::test_gdn_cp_tree_chain_matches_cp1_oracle[2]' \ + 'tests/integration/megatron/gdn_shared_prefix/test_gdn_cp_packed_correctness.py::test_gdn_cp_tree_trainability_updates_parameters[2]' \ + tests/integration/megatron/gdn_shared_prefix/test_real_gdn_tp_lora.py::test_real_qwen35_gdn_tp2_gradients_match_flattened \ + tests/integration/megatron/lora/test_dynamic_lora_slots.py::test_dynamic_lora_slots_capture_recompute_context_and_step_independently \ + 'tests/integration/megatron/lora/test_dynamic_lora_slots.py::test_trainer_rank_tp_head_backward_matches_unsharded_oracle[2]' + +ART_MEGATRON_CONTEXT_PARALLEL_SIZE=2 \ + uv run --no-sync torchrun --standalone --nproc-per-node=2 \ + dev/trainer_rank_check.py \ + --model Qwen/Qwen3-0.6B \ + --layers 1 \ + --depths 0,4 \ + --chunks 17,8192 \ + --slots 0 diff --git a/scripts/ci/trainer-rank-gpu.sky.yaml b/scripts/ci/trainer-rank-gpu.sky.yaml new file mode 100644 index 000000000..e5eeb45c4 --- /dev/null +++ b/scripts/ci/trainer-rank-gpu.sky.yaml @@ -0,0 +1,27 @@ +name: trainer-rank-gpu-validation + +workdir: . + +resources: + accelerators: H200:2 + image_id: docker:docker.io/bradhiltonnw/art-gpu:latest + +setup: | + uv sync --frozen --extra megatron --group dev + +run: | + timeout --signal=TERM --kill-after=30s 20m \ + bash scripts/ci/trainer-rank-gpu-tests.sh + +config: + kubernetes: + pod_config: + spec: + schedulerName: binpack-scheduler + activeDeadlineSeconds: 3600 + containers: + - name: ray-node + imagePullPolicy: Always + env: + - name: UV_LINK_MODE + value: copy diff --git a/src/art/megatron/prefix_tree.py b/src/art/megatron/prefix_tree.py index 4683ca79f..82ac8aa21 100644 --- a/src/art/megatron/prefix_tree.py +++ b/src/art/megatron/prefix_tree.py @@ -1,5 +1,6 @@ from __future__ import annotations +# Temporary no-op used to validate the manual TrainerRank GPU workflow. from dataclasses import dataclass import torch