diff --git a/.github/workflows/collectivex-sweep.yml b/.github/workflows/collectivex-sweep.yml index b56bc285aa..f3688998c0 100644 --- a/.github/workflows/collectivex-sweep.yml +++ b/.github/workflows/collectivex-sweep.yml @@ -1,44 +1,34 @@ -# CollectiveX Sweep — one structured run instead of thousands of dispatches. -# -# Shape (mirrors the InferenceX CI tracker): setup -> sweep (a MATRIX job = "a job with other jobs -# in it") -> aggregate (the collector "at the end"). The matrix unit is a SHARD = one allocation that -# sweeps many cases sharing (sku, backend, mode, resource) — generate_matrix's own grouping, chunked -# so no cell exceeds the job budget. Each cell emits a handful of per-case JSONs; the aggregate job -# collects every shard into ONE line-delimited file (results/aggregate/*.ndjson) so there aren't -# thousands of individual result files. Run once per backend (deepep / uccl / flashinfer / -# deepep-hybrid / nccl-ep, + deepep_v2) for full parity. +# Generate a shard matrix, run each shard on its GPU pool, and upload raw results. name: CollectiveX Sweep +permissions: + actions: read + contents: read on: workflow_dispatch: inputs: backend: - description: EP library to sweep (deepep matrix is remapped onto the others, capability-filtered) + description: "EP library to sweep — 'all' runs every EP backend in one matrix" type: choice - default: deepep - options: [deepep, uccl, flashinfer, deepep-hybrid, nccl-ep] - deepep_v2: - description: DeepEP V2 from-source kernels (kernel_gen=v2; deepep backend only) - type: boolean - default: false - suites: - description: "'all' or comma-list of suite names" - type: string default: all + options: [all, deepep-v2, mori] only_sku: - description: Restrict to one SKU (h100-dgxc|h200|b300|b200-dgxc|gb200|gb300|mi355x); blank = all + description: Restrict to one GHA runner pool; blank = all type: string default: '' - max_cases: - description: Max cases per shard cell (chunk larger shards) + exclude_skus: + description: Comma-list of runner pools to drop from the matrix (partial run, e.g. b300); disjoint from only_sku; blank = none type: string - default: '14' - + default: '' + ep_sizes: + description: Keep only shards whose expert-parallel degree is in this comma-list (8 keeps EP8 only, dropping EP16 so GB SKUs co-schedule with 8-GPU SKUs; blank = all) + type: string + default: '' concurrency: - group: cx-sweep-${{ github.ref }}-${{ inputs.backend }}-${{ inputs.deepep_v2 }}-${{ inputs.only_sku }} + group: cx-${{ github.ref }}-${{ inputs.backend }}-${{ inputs.only_sku }} cancel-in-progress: false jobs: - # ---- setup: resolve the suites into the shard matrix (the "pending jobs" node) ---- + # ---- setup: generate the shard matrix and upload its requested coverage ---- setup: runs-on: ubuntu-latest outputs: @@ -46,26 +36,41 @@ jobs: n: ${{ steps.gen.outputs.n }} steps: - uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v5.0.0 - with: { clean: true } - - run: pip install --quiet pyyaml + with: { clean: true, persist-credentials: false } - id: gen working-directory: experimental/CollectiveX + env: + INPUT_BACKEND: ${{ inputs.backend }} + INPUT_ONLY_SKU: ${{ inputs.only_sku }} + INPUT_EXCLUDE_SKUS: ${{ inputs.exclude_skus }} + INPUT_EP_SIZES: ${{ inputs.ep_sizes }} run: | set -euo pipefail - ov=""; [ "${{ inputs.backend }}" != "deepep" ] && ov="--backend ${{ inputs.backend }}" - v2=""; [ "${{ inputs.deepep_v2 }}" = "true" ] && v2="--deepep-v2" - os=""; [ -n "${{ inputs.only_sku }}" ] && os="--only-sku ${{ inputs.only_sku }}" - # full matrix (with cases) -> artifact for the cells; slim (no cases) -> the strategy output. - python3 sweep_matrix.py --suites "${{ inputs.suites }}" --max-cases "${{ inputs.max_cases }}" $ov $v2 $os --out matrix_full.json >/dev/null - SLIM=$(python3 -c "import json;m=json.load(open('matrix_full.json'));print(json.dumps({'include':[{k:v for k,v in x.items() if k!='cases'} for x in m['include']]}))") - echo "matrix=$SLIM" >> "$GITHUB_OUTPUT" - echo "n=$(python3 -c "import json;print(len(json.load(open('matrix_full.json'))['include']))")" >> "$GITHUB_OUTPUT" - python3 -c "import json;m=json.load(open('matrix_full.json'));print('shard-cells:',len(m['include']),'cases:',sum(x['n'] for x in m['include']))" + args=(--backend "$INPUT_BACKEND") + [ -n "$INPUT_ONLY_SKU" ] && args+=(--only-sku "$INPUT_ONLY_SKU") + [ -n "$INPUT_EXCLUDE_SKUS" ] && args+=(--exclude-skus "$INPUT_EXCLUDE_SKUS") + [ -n "$INPUT_EP_SIZES" ] && args+=(--ep-sizes "$INPUT_EP_SIZES") + python3 sweep_matrix.py "${args[@]}" --out matrix_full.json >/dev/null + python3 - "$GITHUB_OUTPUT" <<'PY' + import json + import pathlib + import sys + + matrix = json.loads(pathlib.Path("matrix_full.json").read_text()) + cells = matrix["include"] + slim = {"include": cells} + with open(sys.argv[1], "a", encoding="utf-8") as output: + output.write(f"matrix={json.dumps(slim, separators=(',', ':'))}\n") + output.write(f"n={len(cells)}\n") + print(f"execution-cells: {len(cells)}") + PY - uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1 with: name: cxsweep-matrix-${{ github.run_id }} path: experimental/CollectiveX/matrix_full.json if-no-files-found: error + overwrite: true + retention-days: 14 # ---- sweep: ONE matrix cell per shard (the parent job with child jobs) ---- sweep: @@ -73,82 +78,148 @@ jobs: if: ${{ fromJSON(needs.setup.outputs.n) > 0 }} strategy: fail-fast: false - max-parallel: 10 # don't saturate the ~20-runner fleet; cells queue as slots free + # Fixed global cap; real throttling is each SKU pool's runner count and its + # cluster's Slurm partition. setup interleaves shards across SKUs, so the + # first jobs under this cap spread over pools instead of queuing on one. + max-parallel: 10 matrix: ${{ fromJSON(needs.setup.outputs.matrix) }} - # h200 label spans two clusters; pin to the validated dgxc pool (mirrors collectivex-experimental). - runs-on: ${{ matrix.sku == 'h200' && 'h200-dgxc' || matrix.sku }} + runs-on: ${{ matrix.sku }} timeout-minutes: 350 env: - CX_BENCH: ${{ matrix.backend }} - CX_DEEPEP_V2: ${{ matrix.deepep_v2 && '1' || '' }} - CX_NODES: ${{ matrix.nodes }} - CX_SHARD_FILE: results/.shard_${{ matrix.id }}.json + COLLX_BENCH: ${{ matrix.backend }} + COLLX_NODES: ${{ matrix.nodes }} + COLLX_GPUS_PER_NODE: ${{ matrix.gpus_per_node }} + COLLX_SCALE_UP_DOMAIN: ${{ matrix.scale_up_domain }} + COLLX_SHARD_FILE: .shards/${{ matrix.id }}.json + COLLX_SHARD_SKU: ${{ matrix.sku }} + COLLECTIVEX_CANONICAL_GHA: '1' COLLECTIVEX_SOURCE_SHA: ${{ github.sha }} - CX_NODELIST: ${{ matrix.sku == 'mi355x' && 'mia1-p01-g10,mia1-p01-g15' || '' }} - CX_STAGE_DIR: ${{ matrix.sku == 'gb200' && '/mnt/lustre01/users-public/sa-shared/cx-stage' || '' }} + # Consolidated shards run one bounded build-group in one Slurm allocation. + # MI300X's compute partition has a 180-minute ceiling; the other production + # pools accept 300 minutes. Allocations release as soon as the shard finishes. + COLLX_TIME: ${{ matrix.sku == 'mi300x' && '180' || '300' }} + COLLECTIVEX_EXECUTION_ID: ${{ github.run_id }}_${{ github.run_attempt }}_${{ matrix.id }} + COLLX_JOB_PARENT: ${{ matrix.launcher == 'mi-amds' && format('/tmp/inferencex-collectivex-parent-{0}-{1}-{2}', github.run_id, github.run_attempt, matrix.id) || '/tmp' }} + COLLX_JOB_ROOT: ${{ matrix.launcher == 'mi-amds' && format('/tmp/inferencex-collectivex-parent-{0}-{1}-{2}/inferencex-collectivex-{0}-{1}-{2}', github.run_id, github.run_attempt, matrix.id) || format('/tmp/inferencex-collectivex-{0}-{1}-{2}', github.run_id, github.run_attempt, matrix.id) }} + COLLX_SOURCE_ROOT: ${{ matrix.launcher == 'mi-amds' && format('/tmp/inferencex-collectivex-parent-{0}-{1}-{2}/inferencex-collectivex-{0}-{1}-{2}/source', github.run_id, github.run_attempt, matrix.id) || format('/tmp/inferencex-collectivex-{0}-{1}-{2}/source', github.run_id, github.run_attempt, matrix.id) }} + HOME: ${{ matrix.launcher == 'mi-amds' && format('/tmp/inferencex-collectivex-parent-{0}-{1}-{2}/inferencex-collectivex-{0}-{1}-{2}/home', github.run_id, github.run_attempt, matrix.id) || format('/tmp/inferencex-collectivex-{0}-{1}-{2}/home', github.run_id, github.run_attempt, matrix.id) }} steps: - - uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v5.0.0 - with: { clean: true } + # Create the private mode-0700 /tmp job root and git-fetch the exact COLLECTIVEX_SOURCE_SHA into it; shards never execute from the runner's shared checkout. + - name: Prepare isolated source + id: source + env: + COLLECTIVEX_REPOSITORY: ${{ github.repository }} + run: | + set -euo pipefail + [ "${COLLX_JOB_ROOT%/*}" = "$COLLX_JOB_PARENT" ] \ + && [[ "${COLLX_JOB_ROOT##*/}" =~ ^inferencex-collectivex-[0-9]+-[0-9]+-[A-Za-z0-9._-]+$ ]] \ + || { echo "CollectiveX isolated root is invalid" >&2; exit 1; } + [ "$COLLX_SOURCE_ROOT" = "$COLLX_JOB_ROOT/source" ] \ + || { echo "CollectiveX source root is invalid" >&2; exit 1; } + if [ "$COLLX_JOB_PARENT" != /tmp ]; then + [[ "$COLLX_JOB_PARENT" =~ ^/tmp/inferencex-collectivex-parent-[0-9]+-[0-9]+-[A-Za-z0-9._-]+$ ]] \ + || { echo "CollectiveX isolated parent is invalid" >&2; exit 1; } + shared_parent="$GITHUB_WORKSPACE/.collectivex-jobs" + install -d -m 700 "$shared_parent" + [ "$(stat -c '%u:%a' "$shared_parent")" = "$(id -u):700" ] \ + || { echo "CollectiveX shared parent is unsafe" >&2; exit 1; } + [ ! -e "$COLLX_JOB_PARENT" ] && [ ! -L "$COLLX_JOB_PARENT" ] \ + || { echo "CollectiveX isolated parent already exists" >&2; exit 1; } + ln -s "$shared_parent" "$COLLX_JOB_PARENT" + fi + umask 077 + [ ! -e "$COLLX_JOB_ROOT" ] && [ ! -L "$COLLX_JOB_ROOT" ] \ + || { echo "CollectiveX isolated root already exists" >&2; exit 1; } + mkdir -m 700 "$COLLX_JOB_ROOT" + trap 'rc=$?; if [ "$rc" != 0 ]; then rm -rf -- "$COLLX_JOB_ROOT"; [ "$COLLX_JOB_PARENT" = /tmp ] || rm -f -- "$COLLX_JOB_PARENT"; fi; exit "$rc"' EXIT + mkdir -m 700 "$HOME" "$COLLX_JOB_ROOT/control" "$COLLX_JOB_ROOT/artifact" "$COLLX_SOURCE_ROOT" + export GIT_CONFIG_NOSYSTEM=1 GIT_CONFIG_GLOBAL=/dev/null + if ! git init -q "$COLLX_SOURCE_ROOT" \ + || ! git -C "$COLLX_SOURCE_ROOT" remote add origin \ + "https://github.com/${COLLECTIVEX_REPOSITORY}.git" \ + || ! git -C "$COLLX_SOURCE_ROOT" -c credential.helper= -c protocol.version=2 \ + fetch -q --no-tags --depth=1 origin "$COLLECTIVEX_SOURCE_SHA" \ + || ! git -C "$COLLX_SOURCE_ROOT" -c advice.detachedHead=false \ + checkout -q --detach FETCH_HEAD \ + || [ "$(git -C "$COLLX_SOURCE_ROOT" rev-parse HEAD)" != "$COLLECTIVEX_SOURCE_SHA" ]; then + echo "CollectiveX source preparation failed" >&2 + exit 1 + fi + echo 'prepared=true' >> "$GITHUB_OUTPUT" + trap - EXIT - uses: actions/download-artifact@3e5f45b2cfb9172054b4087a40e8e0b5a5461e7c # v8.0.1 with: name: cxsweep-matrix-${{ github.run_id }} - path: experimental/CollectiveX - - name: Extract this shard's cases (stdlib only — no runner deps) - working-directory: experimental/CollectiveX + path: ${{ env.COLLX_JOB_ROOT }}/control + # Materialize this shard's .shards/.json case control from the run-scoped matrix artifact downloaded above. + - name: Extract execution control run: | set -euo pipefail - python3 -c " - import json - m=json.load(open('matrix_full.json')) - s=[x for x in m['include'] if x['id']=='${{ matrix.id }}'] - assert s, 'shard ${{ matrix.id }} not in matrix' - s=s[0] - json.dump({'id':s['id'],'sku':s['sku'],'backend':s['backend'],'nodes':s['nodes'],'deepep_v2':s['deepep_v2'],'cases':s['cases']}, open('results/.shard_${{ matrix.id }}.json','w')) - print('shard ${{ matrix.id }}:', len(s['cases']), 'cases') - " - - name: Sweep shard ${{ matrix.id }} (${{ matrix.n }} cases, one allocation) - env: - RUNNER_NAME: ${{ runner.name }} - run: bash "experimental/CollectiveX/launchers/launch_${RUNNER_NAME%%_*}.sh" - - name: Shard summary - if: always() - run: python3 experimental/CollectiveX/summarize.py --results-dir experimental/CollectiveX/results --markdown >> "$GITHUB_STEP_SUMMARY" || true + cd "$COLLX_SOURCE_ROOT/experimental/CollectiveX" 2>/dev/null \ + || { echo "CollectiveX source is unavailable" >&2; exit 1; } + python3 sweep_matrix.py \ + --extract-from "$COLLX_JOB_ROOT/control/matrix_full.json" \ + --shard-id '${{ matrix.id }}' \ + --out '${{ env.COLLX_SHARD_FILE }}' >/dev/null + # Merge the base + per-SKU operator-config secrets into one validated mode-0600 file, then run the SKU launcher (allocate -> stage -> build -> run cases). + - name: Execute sweep cell ${{ matrix.id }} + id: sweep_shard + run: | + set -euo pipefail + umask 077 + # Runner-local Slurm/storage config comes from the tracked + # configs/platform_config.json baseline (no operator secret); the + # launcher's collx_load_operator_config emits it per SKU. + cd "$COLLX_SOURCE_ROOT" 2>/dev/null \ + || { echo "CollectiveX source is unavailable" >&2; exit 1; } + bash "experimental/CollectiveX/launchers/launch_${{ matrix.launcher }}.sh" + # always(): cancel any Slurm allocation a killed launcher left recorded, append the summary table, and stage result JSONs so a red or partial leg still uploads. + - name: Summarize and stage shard results + id: stage_artifact + if: ${{ always() && steps.source.outputs.prepared == 'true' }} + run: | + set -euo pipefail + cd "$COLLX_SOURCE_ROOT" 2>/dev/null \ + || { echo "CollectiveX source is unavailable" >&2; exit 1; } + source experimental/CollectiveX/runtime/common.sh + collx_cleanup_allocation "$COLLX_JOB_ROOT" \ + || { echo "CollectiveX allocation cleanup failed; results may still be changing" >&2; exit 1; } + python3 experimental/CollectiveX/summarize.py \ + --results-dir experimental/CollectiveX/results >> "$GITHUB_STEP_SUMMARY" || true + shopt -s nullglob + results=(experimental/CollectiveX/results/*.json) + if [ "${#results[@]}" -eq 0 ]; then + echo "staged=false" >> "$GITHUB_OUTPUT" + echo "No result JSON to stage; leg produced none." + exit 0 + fi + cp -- "${results[@]}" "$COLLX_JOB_ROOT/artifact/" + echo "staged=true" >> "$GITHUB_OUTPUT" + # The neutral per-case JSONs are the workflow's only output; a consumer downloads them and decides what to display. - name: Upload shard results - if: always() + id: upload_artifact + if: always() && steps.stage_artifact.outputs.staged == 'true' uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1 with: - name: cxshard-${{ matrix.id }}-${{ github.run_id }} - path: experimental/CollectiveX/results/*.json # glob skips the hidden .shard_*.json - if-no-files-found: warn - - # ---- aggregate: collect every shard into ONE ndjson (the "result aggregator at the end") ---- - aggregate: - needs: sweep - if: always() - runs-on: ubuntu-latest - steps: - - uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v5.0.0 - with: { clean: true } - - uses: actions/download-artifact@3e5f45b2cfb9172054b4087a40e8e0b5a5461e7c # v8.0.1 - with: - pattern: cxshard-*-${{ github.run_id }} - path: _shards - merge-multiple: true - - name: Aggregate shards -> one ndjson - working-directory: experimental/CollectiveX + name: ${{ format('cxshard-{0}-{1}-{2}', matrix.id, github.run_id, github.run_attempt) }} + path: | + ${{ env.COLLX_JOB_ROOT }}/artifact/*.json + if-no-files-found: error + retention-days: 14 + # Delete the private /tmp job root after the preceding step releases any allocation. + - name: Cleanup isolated workspace + if: ${{ always() && steps.source.outputs.prepared == 'true' }} run: | set -euo pipefail - tag="${{ inputs.backend }}${{ inputs.deepep_v2 && '-v2' || '' }}" - python3 aggregate_results.py --in-dir ../../_shards --out "results/aggregate/collectivex_${tag}_${{ github.run_id }}.ndjson" - { - echo "## CollectiveX sweep aggregate (${tag})" - echo '```' - wc -l results/aggregate/*.ndjson 2>/dev/null || echo "no ndjson" - echo '```' - } >> "$GITHUB_STEP_SUMMARY" - - name: Upload aggregate - uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1 - with: - name: cxsweep-aggregate-${{ inputs.backend }}${{ inputs.deepep_v2 && '-v2' || '' }}-${{ github.run_id }} - path: experimental/CollectiveX/results/aggregate/*.ndjson - if-no-files-found: warn + [ "$COLLX_JOB_PARENT" = /tmp ] \ + || [[ "$COLLX_JOB_PARENT" =~ ^/tmp/inferencex-collectivex-parent-[0-9]+-[0-9]+-[A-Za-z0-9._-]+$ ]] \ + || { echo "CollectiveX cleanup parent is invalid" >&2; exit 1; } + [ "${COLLX_JOB_ROOT%/*}" = "$COLLX_JOB_PARENT" ] \ + && [[ "${COLLX_JOB_ROOT##*/}" =~ ^inferencex-collectivex-[0-9]+-[0-9]+-[A-Za-z0-9._-]+$ ]] \ + || { echo "CollectiveX cleanup root is invalid" >&2; exit 1; } + [ "$COLLX_SOURCE_ROOT" = "$COLLX_JOB_ROOT/source" ] \ + || { echo "CollectiveX cleanup source is invalid" >&2; exit 1; } + cd "$COLLX_SOURCE_ROOT" 2>/dev/null \ + || { echo "CollectiveX source is unavailable" >&2; exit 1; } + rm -rf -- "$COLLX_JOB_ROOT" + [ "$COLLX_JOB_PARENT" = /tmp ] || rm -f -- "$COLLX_JOB_PARENT" diff --git a/experimental/CollectiveX/.gitignore b/experimental/CollectiveX/.gitignore new file mode 100644 index 0000000000..f34ac59bad --- /dev/null +++ b/experimental/CollectiveX/.gitignore @@ -0,0 +1,14 @@ +__pycache__/ +*.pyc +results/ +unsupported/ +.shards/ +.collx_workloads/ +.collx_backend/ +/matrix_full.json +gpucore.* + +# Local plans and infrastructure inventory. +goal.md +notes.md +private-infra.md diff --git a/experimental/CollectiveX/README.md b/experimental/CollectiveX/README.md new file mode 100644 index 0000000000..71f3cde0ea --- /dev/null +++ b/experimental/CollectiveX/README.md @@ -0,0 +1,152 @@ +# CollectiveX + +CollectiveX is an experimental MoE expert-parallel communication benchmark. It measures dispatch, +combine, and paired roundtrip latency across EP libraries and accelerator systems, then uploads +neutral result artifacts. + +CollectiveX schedules benchmarks, executes them on real allocations, and uploads the neutral +artifacts each run emits. It does not validate those artifacts, promote, rank, recommend, select, or +decide what a consumer displays. Any downstream display or comparison is the consumer's +responsibility. The full measurement methodology is in [docs/methodology.md](docs/methodology.md). + +## Execution Profile + +The workload uses packed placement and one pinned `fixed-profile` resource configuration per +backend/topology; there is no tuning sweep. Dispatch and combine are fixed BF16 on every backend; +precision is not a swept dimension. Every case runs the single normal-mode contract: + +- Normal mode uses `layout-and-dispatch-v1`, rank-deduplicated token payloads, and activation-only + combine. Coverage is uniform routing only. + +Cases use a fixed timing profile from `configs/sweep.json`: 256 trials x 8 timed iterations (2048 +samples per component) with 32 synchronized full roundtrip warmups before each measured component at +every trial/point. Component measurement order rotates each trial so every timed component occupies +every position in the sequence; each iteration takes the cross-rank maximum before nearest-rank +p50/p90/p95/p99, and roundtrip p99 is the headline latency. A keyed BLAKE2b counter produces +byte-identical routing and gate weights on every runtime. + +Correctness is checked against an implementation-independent oracle that reproduces the backend's +two-level reduction — intra-scale-up-domain FP32, then a BF16 cast of each domain's partial for the +scale-out send. The combine gate is a tight max elementwise relative error below `8 * 2^-8` +(denominator clamped at 0.02), which holds across scale-up and multi-node scale-out topologies +alike. Any failed rank or point makes the case ineligible in the result it writes. + +The matrix covers H100, H200, B200, B300, GB200, GB300, MI300X, MI325X, and MI355X. `sweep_matrix.py` materializes +the requested SKUs, backends, EP sizes, and token ladders, then extracts strict per-shard controls +and rejects missing, stale, malformed, or altered shard controls. `--only-sku`, `--exclude-skus`, and +`--ep-sizes` select a subset; the matrix is generated per dispatch, with no frozen digest or locked +case count. + +| Systems | EP8 | EP16 | +|---|---|---| +| H100/H200/B200/B300 | 1x8 NVLink, scale-up | 2x8 NVLink + RDMA, scale-out | +| MI300X/MI325X/MI355X | 1x8 XGMI, scale-up | 2x8 XGMI + RDMA, scale-out | +| GB200/GB300 | 2x4 MNNVL, scale-up | 4x4 MNNVL, scale-up | + +Physical host count does not determine scope: both GB topologies stay inside one 72-GPU MNNVL +scale-up domain. + +| Backend | Current scope | +|---|---| +| DeepEP V2 | PR #605 `ElasticBuffer` plus exact upstream #630 and #640 fixes: LSA for scale-up and GIN for x86 EP16 scale-out | +| MoRI | AMD EP8 uses IntraNode-family kernels (MI355X IntraNode, MI300X/MI325X asyncLL); EP16 pins InterNodeV1 over 2x8 XGMI + RDMA | + +DeepEP V2 means the `ElasticBuffer` implementation introduced by +[DeepEP PR #605](https://github.com/deepseek-ai/DeepEP/pull/605), not a newer legacy `Buffer` build. +The pinned source is the [PR #630](https://github.com/deepseek-ai/DeepEP/pull/630) head, whose parent +is the #605 merge tree, plus the exact one-line library matcher from upstream +[PR #640](https://github.com/deepseek-ai/DeepEP/pull/640). The first fixes pure scale-up +initialization when GIN is unavailable; the second prevents NCCL shared-memory mappings from being +misclassified as duplicate NCCL libraries. Scale-up cases request NCCL Device API LSA and fail closed +unless the realized LSA team covers the full EP world. x86 EP16 scale-out cases instead require the +hybrid path with GIN, two logical scale-out domains represented by two physical RDMA ranks, and eight +scale-up ranks per domain; GB EP16 remains MNNVL scale-up and therefore uses LSA. Whether a given +SKU/backend/EP cell is attempted is a capability fact; whether it succeeded is decided by the +benchmark's return code. + +## Workflow And Artifacts + +`.github/workflows/collectivex-sweep.yml` has two jobs. `setup` generates a public-SKU matrix +(`backend`, `only_sku`, `exclude_skus`, `ep_sizes` inputs) and uploads the matrix. +`sweep` extracts a strict ignored `.shards/.json` control per matrix entry, executes one +allocation per shard, fetches pinned DeepEP source before allocation when required, and uploads the +result artifacts with `always()` so a red or partial run still uploads. + +Each shard emits per-case result JSON and a small mechanical summary. A case counts as successful on +the benchmark's own return code; there is no completeness or privacy validation step, and failed or +unsupported cells produce no synthetic record. No step promotes a run, +builds a dataset, or advances a channel; the neutral artifacts are the output. A consumer downloads +them and decides what to display. + +No operator credentials are passed to the workflow or uploaded; runner-local overrides and any +selectors stay on the runner. Per-step runner logs are kept on the runner for postmortem, and +result artifacts carry only the fields listed in the methodology. + +## Runner Configuration + +Each SKU's Slurm and storage values come from its tracked baseline in the registry. An optional +runner-local JSON document at `$XDG_CONFIG_HOME/inferencex/collectivex.json` or +`COLLECTIVEX_OPERATOR_CONFIG` overlays that baseline per field; unknown runners, fields, duplicate +keys, and non-JSON input fail closed, and configuration is never evaluated as shell. GHA passes no +operator secret, so a SKU runs entirely from its tracked baseline unless a runner-local document is +present. + +All public per-SKU platform data lives in the tracked `configs/platform_config.json` registry: +architecture/product, container image and platform, fixed placement, launcher, runnable backend/EP +pairs, the scale-out `fabric` identity (NIC and switch — so same-GPU clusters on different fabrics +are distinct entries, e.g. a second b200 cluster), tracked operator defaults, and scale-out RDMA +selectors. Operator documents can override the defaults. Launchers +declare and check the fields they actually require. `sweep_matrix.py` derives EP topology from the +placement fields; the sweep includes every registered SKU by default. + +Every selected non-MNNVL EP16 placement additionally requires `socket_ifname` and `rdma_devices` for +its operator-approved fabric; optional `ib_gid_index`, `rdma_service_level`, `rdma_traffic_class`, +and `rail_isolated` are also allowlisted. Service level and traffic class are mapped into MoRI's +RDMA/IO QoS environment. +CollectiveX does not heuristically select a management route or HCA. After allocation, every +non-MNNVL scale-out node must prove that all configured interfaces and active HCA ports exist before +backend setup. Scale-up and MNNVL jobs clear these overrides. Scale-out NCCL/RCCL is pinned to `IB` +with exact-match HCA selectors so a socket fallback fails instead of being mislabeled as RDMA. +Scale-out also disables NCCL dual-port NIC fusion (`NCCL_IB_MERGE_NICS=0`): a fused device disables +NCCL GIN, which the DeepEP V2 EP16 hybrid path requires, and a rail-isolated fabric +(`rail_isolated=1`, e.g. B300's multi-plane RoCE) additionally sets `NCCL_CROSS_NIC=0`. + +`ib_gid_index` is applied only when every selected HCA port reports an Ethernet link layer, where it +selects the operator-approved RoCE GID. Native InfiniBand profiles retain explicit HCA and service +level pinning but leave the RoCE-only GID override unset so NVSHMEM/NCCL can use the native LID path. +Mixed Ethernet and InfiniBand HCA lists are rejected. + +`stage_dir` is a pre-existing, runner-owned, non-symlinked base outside the checkout and workflow +workspace. It is not group- or world-writable and is visible at the same path on the runner and every +allocated node. Jobs create only a marked mode-0700 execution child, prove cross-node read/write +visibility, and remove that exact child after allocation teardown; they never mount the runner +checkout or create a stage beneath image storage on AMD. When an AMD operator row omits `stage_dir`, +the runner derives a private base beside its standard `_work` directory on the shared runner +filesystem; the root-owned squash cache is never used as a repository stage. + +H200, B200, and B300 runners may omit `stage_dir`; their isolated execution child is created under a +runner-owned mode-0700 base in the validated operating-system account home, independent of the +workflow's temporary `HOME`. H100 may also omit `stage_dir`; its private base is created beside, never +beneath, the configured shared container directory so it is compute-visible. Canonical B300 execution +ignores any legacy configured `stage_dir` and always uses the validated compute-visible account-home +base; an execution-ID suffix isolates parallel B300 workers. Canonical GB300 execution likewise +ignores its legacy group-writable `stage_dir` and derives an execution-specific private base beneath the +validated compute-visible account home. Backend preparation runs from that staged tree on every node. + +Enroot imports the configured image tag into a per-run-scoped squash keyed by image tag and image +platform, so one run never reuses another run's imported filesystem. The image tag and platform are +per-SKU registry fields; the DeepEP V2 source pin lives in `runtime/common.sh` and its build is +fetched and verified at the pinned commit, checked for `ElasticBuffer`, and cached in a +cluster-local build cache keyed by architecture, image, and commit. Only the fixed `/cx-cache` mount +reaches the container. + +## Local Checks + +```bash +python3 -m unittest discover experimental/CollectiveX/tests -p 'test_*.py' +python3 experimental/CollectiveX/sweep_matrix.py --backend all --out /tmp/cx-matrix.json >/dev/null +bash -n experimental/CollectiveX/runtime/*.sh experimental/CollectiveX/launchers/*.sh +``` + +Core paths are `configs/`, `sweep_matrix.py`, `summarize.py`, `bench/`, `runtime/`, `launchers/`, +and `tests/`. diff --git a/experimental/CollectiveX/bench/ep_backend.py b/experimental/CollectiveX/bench/ep_backend.py new file mode 100644 index 0000000000..ccf0c74934 --- /dev/null +++ b/experimental/CollectiveX/bench/ep_backend.py @@ -0,0 +1,313 @@ +#!/usr/bin/env python3 +"""Shared lifecycle and input generation for EP backends.""" +from __future__ import annotations + +import abc +import types +from dataclasses import dataclass, field +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + import torch + +from ep_harness import ( + time_us, + token_ladder, +) + + +@dataclass +class RankInputs: + """Inputs for one token-ladder shape at tokens_per_rank tokens on this rank. + + topk_idx/topk_weights are this rank's contiguous slice of the global routing trace + (host tensors; moved to device at make_problem time); activations are the rank's token + activations (already on device). The global trace is retained so Pass 1 can compute + routing statistics and input snapshots. + """ + + tokens_per_rank: int + topk_idx: "torch.Tensor" + topk_weights: "torch.Tensor" + activations: "torch.Tensor" + global_idx: "torch.Tensor | None" = None + global_weights: "torch.Tensor | None" = None + + +@dataclass +class WorkloadSpec: + """Numeric shape + materialised inputs for one fully-specified sweep line. + + Fully default-constructible so make_inputs can early-return a tensor-free + spec (ok=False + rc) on an empty ladder; the driver prints message and + returns rc + """ + + ok: bool = True + rc: int = 0 + message: str = "" + ep_size: int = 0 + experts_per_rank: int = 0 + cap: "int | None" = None + dropped: list = field(default_factory=list) + max_tokens_per_rank: int = 0 + ladder: list = field(default_factory=list) + points: dict = field(default_factory=dict) + + +class EPBackend(abc.ABC): + """One expert-parallel dispatch/combine transport under a fixed benchmark contract. + + Subclasses implement the transport (create_buffer, dispatch, stage, + combine, recv_tokens, inspect_dispatch, combine_transformed); + everything the driver and the oracles need beyond that is provided here. + Dispatch and combine are fixed BF16, so no adapter selects a precision codec. + """ + + name: str = "" + SUPPORTED_MODES: tuple = ("normal",) + stage_device_work = False + combine_needs_redispatch = False + dispatch_needs_combine_cleanup = False + # Adapters that reduce activations and top-k weights independently must carry + # the complete local weighted expert sum in the activation tensor. + combine_weight_semantics = "unweighted-rank-sum" + roundtrip_only = False + + def __init_subclass__(cls, **kwargs): + super().__init_subclass__(**kwargs) + if not getattr(cls, "name", ""): + raise TypeError( + f"{cls.__name__} must declare a non-empty class-level `name`" + ) + + def __init__(self, args, rank, world_size, local_rank, device): + self.args = args + self.rank = rank + self.world_size = world_size + self.local_rank = local_rank + self.device = device + self.mode = args.mode + if self.mode not in self.SUPPORTED_MODES: + raise ValueError(f"{self.name} does not support mode {self.mode!r}") + + # ---- Abstract transport contract ------------------------------------------------- + + @abc.abstractmethod + def create_buffer(self, spec: WorkloadSpec): + """Size the communicator from spec before the first dispatch.""" + + @abc.abstractmethod + def dispatch(self, problem): + """Scatter tokens to their experts; return an opaque per-call handle.""" + + @abc.abstractmethod + def stage(self, problem, handle): + """Prepare the combine input on handle (copy into place).""" + + @abc.abstractmethod + def combine(self, problem, handle): + """Gather the staged tokens back to their source rank; return combined activations.""" + + @abc.abstractmethod + def recv_tokens(self, handle): + """Number of tokens this rank received in dispatch (stable for a fixed trace).""" + + @abc.abstractmethod + def inspect_dispatch(self, problem, handle): + """Normalized post-dispatch view for the token-rank correctness oracle.""" + + @abc.abstractmethod + def combine_transformed(self, problem, handle, transformed): + """Combine an oracle-transformed payload in place of the staged input.""" + + # ---- Input generation (shared) --------------------------------------------------- + + def buffer_cap(self, args): + """Max tokens/rank the communicator can serve, or None when unbounded.""" + return None + + def make_inputs(self, args) -> WorkloadSpec: + """Resolve the token ladder and materialise per-rank inputs for the sweep. + + Buffer sizing needs the ladder *numbers* (not the input tensors), so this + runs before create_buffer. Returns a tensor-free spec with ok=False + when the ladder is empty. + """ + ep_size = self.world_size + experts_per_rank = args.experts // ep_size + cap = self.buffer_cap(args) + ladder, dropped = token_ladder(args.tokens_ladder, cap) + if not ladder: + return WorkloadSpec( + ok=False, rc=2, + message=f"empty token ladder (phase={args.phase}, cap={cap})", + ) + spec = WorkloadSpec( + ep_size=ep_size, + experts_per_rank=experts_per_rank, + cap=cap, + dropped=list(dropped), + max_tokens_per_rank=max(ladder), + ladder=list(ladder), + ) + for tokens_per_rank in ladder: + spec.points[tokens_per_rank] = self._build_rank_inputs(args, tokens_per_rank) + return spec + + def _build_rank_inputs(self, args, tokens_per_rank) -> RankInputs: + """Build one rank's deterministic inputs for a tokens-per-rank shape.""" + import torch + import routing + + ep_size = self.world_size + num_logical = getattr(args, "num_logical_experts", args.experts) + global_tokens = tokens_per_rank * ep_size + idx_g, w_g = routing.build_global_routing( + global_tokens, num_logical, args.topk, args.routing, args.seed + ) + idx_s, w_s = routing.rank_slice(idx_g, w_g, self.rank, tokens_per_rank) + activations = routing.rank_activations( + tokens_per_rank, args.hidden, args.seed, self.rank, self.device, torch.bfloat16 + ) + return RankInputs( + tokens_per_rank=tokens_per_rank, + topk_idx=idx_s.contiguous(), + topk_weights=w_s.contiguous(), + activations=activations, + global_idx=idx_g, + global_weights=w_g, + ) + + def make_problem(self, T, idx, weights, x): + """Assemble the per-shape problem namespace (BF16 dispatch sends x directly).""" + import torch + + return types.SimpleNamespace( + T=T, + x=x, + dispatch_x=x, + topk_idx=idx.to(self._topk_idx_dtype()), + topk_weights=weights.to(torch.float32), + ) + + def _topk_idx_dtype(self): + """Integer dtype the backend's kernels expect for top-k routing indices.""" + import torch + return torch.int64 + + # ---- Timing template methods ----------------------------------------------------- + + def timed_components(self): + """Components measured for this backend: roundtrip always; the rest unless + the backend exposes only a stateful paired round trip.""" + components = ["roundtrip"] + if not self.roundtrip_only: + components.extend(["dispatch", "combine"]) + if self.stage_device_work: + components.append("stage") + return components + + def warm(self, problem, count): + """Untimed synchronized full round trips (fabric/clock warm-up; cold-jump-safe). + + Caches the dynamic receive cardinality once so adapters never read a device + scalar during a timed trial (the count is stable for a fixed routing trace). + """ + import torch + + for _ in range(count): + handle = self.dispatch(problem) + if not hasattr(problem, "recv_tokens"): + problem.recv_tokens = self.recv_tokens(handle) + self.stage(problem, handle) + self.combine(problem, handle) + torch.cuda.synchronize() + + def run_roundtrip(self, problem): + """One full dispatch -> stage -> combine round trip; returns combined activations.""" + handle = self.dispatch(problem) + self.stage(problem, handle) + return self.combine(problem, handle) + + def benchmark_component(self, component, problem, warmup, iters): + """Measure one named component; every component gets the same warm-up first.""" + if component == "roundtrip": + return self.benchmark_roundtrip(problem, warmup, iters) + if component == "dispatch": + return self.benchmark_dispatch(problem, warmup, iters) + if component == "stage": + return self.benchmark_stage(problem, warmup, iters) + if component == "combine": + return self.benchmark_combine(problem, warmup, iters) + raise RuntimeError(f"unknown timed component {component!r}") + + def benchmark_roundtrip(self, problem, warmup, iters): + import torch + + self.warm(problem, warmup) + return time_us(torch, lambda p=problem: self.run_roundtrip(p), 0, iters) + + def benchmark_dispatch(self, problem, warmup, iters): + import torch + + self.warm(problem, warmup) + + def finish_dispatch(hh, p=problem): + self.stage(p, hh) + self.combine(p, hh) + + dispatch_needs_cleanup = self.dispatch_needs_combine_cleanup + return time_us( + torch, lambda p=problem: self.dispatch(p), 0, iters, + post=finish_dispatch if dispatch_needs_cleanup else None, + ) + + def benchmark_stage(self, problem, warmup, iters): + import torch + + self.warm(problem, warmup) + + def prep_stage(p=problem): + return self.dispatch(p) + + def stage_op(hh, p=problem): + self.stage(p, hh) + return hh + + # Drain each timed stage's dispatch with an untimed combine where the + # backend requires the pair (same rule as benchmark_dispatch). + return time_us( + torch, stage_op, 0, iters, pre=prep_stage, + post=(lambda hh, p=problem: self.combine(p, hh)) + if self.dispatch_needs_combine_cleanup else None, + ) + + def benchmark_combine(self, problem, warmup, iters): + import torch + + self.warm(problem, warmup) + + def prep_combine(p=problem): + hh = self.dispatch(p) + self.stage(p, hh) + return hh + + if self.combine_needs_redispatch: + return time_us( + torch, lambda hh, p=problem: self.combine(p, hh), 0, iters, pre=prep_combine, + ) + hh = prep_combine() + torch.cuda.synchronize() + return time_us(torch, lambda p=problem, hx=hh: self.combine(p, hx), 0, iters) + + def finalize(self, rc): + """Barrier and tear down the process group; returns rc.""" + import torch.distributed as dist + + try: + dist.barrier() + dist.destroy_process_group() + except Exception: + pass + return rc diff --git a/experimental/CollectiveX/bench/ep_deepep_v2.py b/experimental/CollectiveX/bench/ep_deepep_v2.py new file mode 100644 index 0000000000..e6b1f05e78 --- /dev/null +++ b/experimental/CollectiveX/bench/ep_deepep_v2.py @@ -0,0 +1,222 @@ +#!/usr/bin/env python3 +"""DeepEP PR #605 adapter with the exact upstream PR #630 and #640 fixes.""" + +from __future__ import annotations + +import inspect +import os +import sys +import types +from pathlib import Path + +import torch +import torch.distributed as dist +from ep_backend import EPBackend + +try: + import deep_ep + from deep_ep import ElasticBuffer # type: ignore +except Exception as exc: # pragma: no cover - requires the benchmark image + print(f"ERROR: DeepEP V2 import failed: {exc!r}", file=sys.stderr) + raise + + +# Source pins (PR #605 head + #630/#640 fixes) live in runtime/common.sh; +# the launcher fetches and builds them from that checkout. This adapter no longer +# verifies the wheel's commit tag against the pin — it checks only that the loaded +# deep_ep exposes ElasticBuffer (the from-source PR #605 capability). + + +def _jit_cache_directory( + args, + world_size: int, + max_tokens: int, + allow_hybrid_mode: bool, + realized: dict[str, int | bool], +) -> str: + values = ( + args.runner, world_size, args.hidden, args.topk, args.experts, + getattr(args, "num_logical_experts", args.experts), max_tokens, + int(allow_hybrid_mode), realized["allocated_qps"], realized["num_sms"], + ) + return "jit-" + "-".join(str(value) for value in values) + + +# GIN/GDAKI allocates num_allocated_qps device QPs per peer rank on the local NIC +# (contexts x world_size QPs, before NCCL's own connection QPs). Upstream's hybrid +# default (129, or 65 with fast RDMA atomics) exhausts the per-NIC QP budget at +# EP16: construction dies in ncclDevCommCreate with ibv_create_qp ENOMEM once +# NCCL's regular QPs land on top (identical on H200 bare-metal and B200 pods; on +# CX-7 the budget sits between 784 and 1040 QPs — 49x16 initializes, 65x16 does +# not). Spending a fixed ~512-QP budget keeps every EP size inside that limit +# with headroom: EP8 resolves to 65 (the allocation CX-8 racks already run +# successfully), EP16 to 33 and EP32 to 17 (33 and 49 verified on the failing +# H200 pair). An explicit value also skips upstream's rank-local ibstat probe, +# which is not guaranteed to resolve identically across ranks. +_GIN_QP_BUDGET = 512 + + +def _hybrid_num_allocated_qps(world_size: int) -> int: + return max(9, 1 + _GIN_QP_BUDGET // world_size) + + +def _configure_gin_mode(args, world_size: int) -> bool: + scale_up_domain = int(args.scale_up_domain) + allow_hybrid_mode = world_size > scale_up_domain + if allow_hybrid_mode: + os.environ.pop("EP_DISABLE_GIN", None) + else: + os.environ["EP_DISABLE_GIN"] = "1" + return allow_hybrid_mode + + +def _require_runtime() -> None: + """Capability check only: the loaded deep_ep must expose ElasticBuffer (still + catches the b300 image-bundled deep_ep 1.2.1 shadowing the from-source build, + which lacks the class).""" + if not inspect.isclass(ElasticBuffer) or ElasticBuffer.__name__ != "ElasticBuffer": + raise RuntimeError("invalid DeepEP V2 runtime: deep_ep.ElasticBuffer is absent") + + +class DeepEPV2Backend(EPBackend): + name = "deepep-v2" + # Invariant by identity contract: this backend IS the PR #605 ElasticBuffer + # implementation; LSA vs hybrid GIN are transport paths, not kernel families. + kernel_generation = "v2-elastic-buffer" + stage_device_work = False + combine_needs_redispatch = False + combine_weight_semantics = "unweighted-rank-sum" + + def __init__(self, args, rank, world_size, local_rank, device): + # deepep-v2 is normal-mode only; base SUPPORTED_MODES=("normal",) enforces it. + super().__init__(args, rank, world_size, local_rank, device) + self.group = dist.group.WORLD + + def create_buffer(self, spec): + # max_tokens is the measured-ladder maximum; the historical values (which + # also folded in the conditioning ramp) are identical because the ramp + # never exceeded the measured maximum, so the JIT directory stays stable. + args, world_size = self.args, self.world_size + self.max_tokens = spec.max_tokens_per_rank + _require_runtime() + jit_root = Path(os.environ["EP_JIT_CACHE_DIR"]) + allow_hybrid_mode = _configure_gin_mode(args, world_size) + self.buffer = ElasticBuffer( + self.group, + num_max_tokens_per_rank=self.max_tokens, + hidden=args.hidden, + num_topk=args.topk, + use_fp8_dispatch=False, # BF16 communication path. + deterministic=False, + allow_hybrid_mode=allow_hybrid_mode, + allow_multiple_reduction=True, + prefer_overlap_with_compute=True, + num_gpu_timeout_secs=100, + explicitly_destroy=True, + # 0 is upstream's use-the-default sentinel; only hybrid (GIN) mode + # needs the explicit budget-derived allocation. + num_allocated_qps=( + _hybrid_num_allocated_qps(world_size) if allow_hybrid_mode else 0 + ), + ) + tuning_num_experts = int(getattr(args, "num_logical_experts", args.experts)) + self.num_sms = int( + self.buffer.get_theoretical_num_sms(tuning_num_experts, args.topk) + ) + self.num_qps = int(self.buffer.get_theoretical_num_qps(self.num_sms)) + realized = { + "num_sms": self.num_sms, + "allocated_qps": int(self.buffer.num_allocated_qps), + } + jit_cache_directory = _jit_cache_directory( + args, + world_size, + self.max_tokens, + allow_hybrid_mode, + realized, + ) + os.environ["EP_JIT_CACHE_DIR"] = str(jit_root / jit_cache_directory) + + def _topk_idx_dtype(self): + # DeepEP V2's kernels key routing indices on deep_ep.topk_idx_t, not int64. + return deep_ep.topk_idx_t + + def dispatch(self, p): + recv_x, recv_topk_idx, recv_topk_weights, handle, _ = self.buffer.dispatch( + p.dispatch_x, + topk_idx=p.topk_idx, + topk_weights=p.topk_weights, + num_experts=self.args.experts, + num_max_tokens_per_rank=self.max_tokens, + expert_alignment=1, + num_sms=self.num_sms, + num_qps=self.num_qps, + async_with_compute_stream=False, + do_handle_copy=True, + do_cpu_sync=True, + do_expand=False, + ) + return types.SimpleNamespace( + recv_x=recv_x, + recv_topk_idx=recv_topk_idx, + recv_topk_weights=recv_topk_weights, + handle=handle, + ) + + def stage(self, p, h): + # BF16: the received buffer is already the semantic payload to combine. + h.combine_input = h.recv_x + + def combine(self, p, h): + combined_x, _, _ = self.buffer.combine( + h.combine_input, + handle=h.handle, + num_sms=self.num_sms, + num_qps=self.num_qps, + async_with_compute_stream=False, + ) + return combined_x + + def inspect_dispatch(self, p, h): + count = self.recv_tokens(h) + local_idx = h.recv_topk_idx[:count] + valid = local_idx >= 0 + expert_ids = torch.where( + valid, + local_idx + self.rank * (self.args.experts // self.world_size), + local_idx, + ) + local = local_idx[valid].to(torch.int64) + return types.SimpleNamespace( + payload=h.recv_x[:count], + expert_ids=expert_ids, + weights=h.recv_topk_weights[:count].masked_fill(~valid, 0), + local_expert_counts=torch.bincount( + local, minlength=self.args.experts // self.world_size + ), + ) + + def combine_transformed(self, p, h, transformed): + combine_input = torch.zeros_like(h.recv_x) + combine_input[: transformed.shape[0]].copy_(transformed.to(combine_input.dtype)) + combined, _, _ = self.buffer.combine( + combine_input, + handle=h.handle, + num_sms=self.num_sms, + num_qps=self.num_qps, + async_with_compute_stream=False, + ) + return combined + + def recv_tokens(self, h): + return int(h.handle.psum_num_recv_tokens_per_scaleup_rank[-1].item()) + + def finalize(self, rc): + try: + dist.barrier() + self.buffer.destroy() + dist.barrier() + dist.destroy_process_group() + except Exception: + return 1 + return rc diff --git a/experimental/CollectiveX/bench/ep_harness.py b/experimental/CollectiveX/bench/ep_harness.py new file mode 100644 index 0000000000..d01250fa25 --- /dev/null +++ b/experimental/CollectiveX/bench/ep_harness.py @@ -0,0 +1,868 @@ +#!/usr/bin/env python3 +"""Shared EP timing, correctness, and result generation.""" +from __future__ import annotations + +import argparse +import datetime as _dt +import json +import math +import os +import re + + +_CASE_ID = re.compile(r"^[a-z0-9][a-z0-9.-]*$") +_NON_SLUG = re.compile(r"[^a-z0-9]+") + + +def is_case_id(value) -> bool: + return bool(isinstance(value, str) and _CASE_ID.fullmatch(value)) + + +def case_id(sku: str, case: dict) -> str: + parts = ( + sku, + case["backend"], + case["workload"], + case["mode"], + case["phase"], + f"ep{int(case['ep'])}", + case["routing"], + ) + values = [_NON_SLUG.sub("-", str(part).lower()).strip("-") for part in parts] + if not all(values): + raise ValueError("case ID contains an empty factor") + return "-".join(values) + + +# Workload and timing values arrive from configs/sweep.json through the matrix. +CONDITIONING_ROUNDS_PER_SHAPE = 8 +# Dispatch and combine are fixed BF16, so the combine oracle uses one frozen gate. +# _expected_transformed_combine reproduces the two-level (intra-domain FP32, +# per-domain BF16 scale-out partial) reduction, so a correct backend's only +# residual is the accumulation-order ambiguity the model cannot pin down: at most +# topk (8) BF16 stores at one ulp (2^-8) each. Below the magnitude floor the gate +# is effectively absolute (cancellation makes relative error meaningless there). +COMBINE_REL_TOL = 8 * 2.0 ** -8 +COMBINE_MAG_FLOOR = 2e-2 + +def logical_byte_provenance(logical_copies: int, hidden: int) -> dict[str, int]: + """Return comparable logical BF16 activation bytes for one direction. + + Dispatch and combine both move BF16 (2 bytes/value) with no separate scale + payload, so ``scale_bytes`` is always zero. + """ + if logical_copies < 0 or hidden < 0: + raise ValueError("logical byte dimensions must be non-negative") + activation_data_bytes = logical_copies * hidden * 2 + return { + "activation_data_bytes": activation_data_bytes, + "scale_bytes": 0, + "total_logical_bytes": activation_data_bytes, + } + +def format_collective_version(raw) -> str: + """Normalize PyTorch's tuple or packed NCCL/RCCL version representation.""" + if isinstance(raw, int): + if raw < 10_000: + return f"{raw // 1000}.{raw // 100 % 10}.{raw % 100}" + return f"{raw // 10_000}.{raw // 100 % 100}.{raw % 100}" + if isinstance(raw, (tuple, list)): + return ".".join(map(str, raw)) + return str(raw) if raw not in (None, "") else "unknown" + + +def add_common_args(ap: argparse.ArgumentParser) -> None: + """Add the varying v1 inputs; fixed profile values are not CLI axes.""" + ap.add_argument("--mode", required=True, choices=["normal"]) + ap.add_argument("--phase", required=True, choices=["decode", "prefill"], + help="token-size regime label: decode (small T) / prefill (large T)") + ap.add_argument("--tokens-ladder", required=True, + help="space/comma-separated source-tokens-per-rank sweep; the matrix " + "supplies the workload's phase ladder from configs/sweep.json") + ap.add_argument("--hidden", type=int, required=True) + ap.add_argument("--topk", type=int, required=True) + ap.add_argument("--experts", type=int, required=True, + help="TOTAL experts (fixed across EP degrees)") + ap.add_argument("--routing", required=True, choices=["uniform"]) + ap.add_argument("--case-id", required=True) + ap.add_argument("--suite", required=True) + ap.add_argument("--workload-name", required=True) + ap.add_argument("--seed", type=int, required=True, + help="routing-trace seed; part of the workload identity in configs/sweep.json") + ap.add_argument( + "--version", + type=int, + required=True, + help="iterable benchmark version copied verbatim into the emitted result", + ) + # The single cross-SKU profile lives in configs/sweep.json + # `timing:`; the matrix bakes it into every scheduled case. + ap.add_argument("--warmup", type=int, required=True, + help="untimed full roundtrips before each trial/point") + ap.add_argument("--iters", type=int, required=True, + help="timed iterations per trial") + ap.add_argument("--trials", type=int, required=True, + help="timed trials") + # provenance / output + ap.add_argument("--runner", required=True) + ap.add_argument("--topology-class", required=True) + ap.add_argument("--transport", required=True) + ap.add_argument("--scope", required=True, choices=["scale-up", "scale-out"]) + ap.add_argument("--scale-up-transport", required=True) + ap.add_argument("--scale-out-transport", required=True) + ap.add_argument("--gpus-per-node", type=int, required=True) + ap.add_argument("--scale-up-domain", type=int, required=True) + ap.add_argument("--out", required=True) + + +def token_ladder(spec: str, cap: int | None) -> tuple[list[int], list[int]]: + """Return (ladder, dropped) from an explicit spec (there is no default — the + model-specific ladders live in configs/sweep.json); positive ints; clamped to + `cap` with dropped points reported (never silently truncated).""" + want = sorted({t for t in (int(t) for t in spec.replace(",", " ").split() if t) if t > 0}) + if cap is not None: + return [t for t in want if t <= cap], [t for t in want if t > cap] + return want, [] + + +def trial_order(values: list, trial_index: int) -> list: + """Rotate and reverse values so each occupies every timing position.""" + if not values or len(values) != len(set(values)): + raise ValueError("trial order requires non-empty unique values") + if type(trial_index) is not int or trial_index < 0: + raise ValueError("trial_index must be a non-negative integer") + cycle, offset = divmod(trial_index, len(values)) + base = list(values) if cycle % 2 == 0 else list(reversed(values)) + return base[offset:] + base[:offset] + + +def percentile(xs: list[float], q: float) -> float: + if not xs: + return float("nan") + s = sorted(xs) + i = max(0, min(len(s) - 1, math.ceil(q / 100.0 * len(s)) - 1)) + return s[i] + + +def _pcts(xs): + return ({"p50": percentile(xs, 50), "p90": percentile(xs, 90), + "p95": percentile(xs, 95), "p99": percentile(xs, 99)} if xs else None) + + +def _component(percentiles, count, *, derived=False): + if percentiles is None: + return {"availability": "unavailable", "origin": None, + "percentiles_us": None, "sample_count": 0} + return { + "availability": "derived" if derived else "measured", + "origin": "derived-percentile-sum" if derived else "measured", + "percentiles_us": percentiles, + "sample_count": 0 if derived else count, + } + + +# The exact routing fields each row publishes — a whitelist so a new stat in +# routing.routing_stats never leaks into the artifact unreviewed. +_ROUTING_FIELDS = ( + "empty_expert_count", "empty_rank_count", "expert_assignment_rank_cv", + "expert_assignments_per_rank", "expert_load_cv", "expert_load_max", + "expert_load_mean", "expert_load_min", "fanout_histogram", "fanout_max", + "fanout_mean", "fanout_min", "hotspot_ratio", "locality", + "payload_copies_per_rank", "payload_rank_cv", "routed_copies", +) + + +def _write_bytes_atomic(path: str, payload: bytes) -> None: + os.makedirs(os.path.dirname(os.path.abspath(path)), exist_ok=True) + temporary = f"{path}.tmp-{os.getpid()}" + try: + with open(temporary, "wb") as handle: + handle.write(payload) + handle.flush() + os.fsync(handle.fileno()) + os.replace(temporary, path) + finally: + try: + os.unlink(temporary) + except FileNotFoundError: + pass + + +def _write_json_atomic(path: str, value) -> None: + payload = json.dumps( + value, allow_nan=False, ensure_ascii=False, separators=(",", ":") + ).encode() + b"\n" + _write_bytes_atomic(path, payload) + + +def time_us(torch, fn, warmup: int, iters: int, pre=None, post=None) -> list[float]: + """Per-iteration CUDA-event latencies (µs) for THIS rank. + + Without `pre`: times `fn()`. With `pre`: runs `pre()` UNTIMED each iteration (sync + before the start event so its GPU work can't bleed in), then times `fn(pre_result)`. + `post(result)` runs after the end event and synchronization, so stateful backends can + consume/reset a timed operation without charging that cleanup to its latency. Returns + the raw per-iteration series; the caller reduces across ranks per iteration before + percentiling. + """ + def sample(): + arg = pre() if pre is not None else None + if pre is not None: + torch.cuda.synchronize() + s = torch.cuda.Event(enable_timing=True) + e = torch.cuda.Event(enable_timing=True) + s.record() + result = fn(arg) if pre is not None else fn() + e.record() + torch.cuda.synchronize() + elapsed = s.elapsed_time(e) * 1000.0 # ms -> us + if post is not None: + post(result) + torch.cuda.synchronize() + return elapsed + + for _ in range(max(0, warmup)): + if pre is not None: + a = pre() + torch.cuda.synchronize() + fn(a) + else: + fn() + # sync EACH warmup iteration, not just once after the loop: the measured-roundtrip fn + # interleaves dispatch+combine on a backend's persistent comm buffer, so back-to-back + # un-synced warmup iterations let iter N+1's dispatch race iter N's combine (CUDA abort + # on a rank -> NCCL-watchdog SIGABRT). Cheap (warmup is small); timed samples already sync. + torch.cuda.synchronize() + return [sample() for _ in range(iters)] + + +def kernel_generation(backend) -> str: + """Return the adapter's declared kernel family.""" + return getattr(backend, "kernel_generation", None) or "n-a" + + +def _reduce_vec(torch, dist, device, vals, op): + t = torch.tensor(vals, device=device, dtype=torch.float64) + dist.all_reduce(t, op=op) + return [float(x) for x in t.tolist()] + + +def _reduce_int(torch, dist, device, v: int, op) -> int: + t = torch.tensor([int(v)], device=device, dtype=torch.int64) + dist.all_reduce(t, op=op) + return int(t.item()) + + +def _same_tensors_across_ranks(torch, dist, device, *tensors) -> bool: + matches = True + for tensor in tensors: + observed = tensor.to(device=device, non_blocking=False) + reference = observed.clone() if dist.get_rank() == 0 else torch.empty_like(observed) + dist.broadcast(reference, src=0) + matches = matches and bool(torch.equal(observed, reference)) + result = torch.tensor([int(matches)], device=device, dtype=torch.int64) + dist.all_reduce(result, op=dist.ReduceOp.MIN) + return bool(result.item()) + + +def _normalized_expert_metadata(torch, expert_ids, weights): + """Sort each row by global expert ID while keeping -1 sentinels last.""" + valid = expert_ids >= 0 + keys = torch.where(valid, expert_ids.to(torch.int64), torch.full_like(expert_ids, 1 << 30)) + order = torch.argsort(keys, dim=1, stable=True) + sorted_ids = torch.gather(expert_ids.to(torch.int64), 1, order) + sorted_weights = torch.gather(weights.to(torch.float32), 1, order) + sorted_valid = sorted_ids >= 0 + return ( + torch.where(sorted_valid, sorted_ids, torch.full_like(sorted_ids, -1)), + sorted_weights.masked_fill(~sorted_valid, 0), + ) + + +def _expert_coefficients(torch, expert): + """Per-expert affine coefficients — the transform and its independently derived + expectation must use these exact formulas, so they exist only here.""" + scale = ((expert * 17 + 5) % 31 + 1).to(torch.float32) / 32 + offset_a = (((expert * 29 + 7) % 37) - 18).to(torch.float32) / 64 + offset_b = (((expert * 43 + 11) % 41) - 20).to(torch.float32) / 128 + return scale, offset_a, offset_b + + +def _column_pattern(torch, ncols, device): + columns = torch.arange(ncols, device=device, dtype=torch.int64) + return (((columns * 13) % 17) - 8).to(torch.float32) / 8 + + +def _expert_transform(torch, payload, expert_ids, weights, combine_weight_semantics): + """Build one local expert aggregate for the v1 unweighted combine contract.""" + if combine_weight_semantics != "unweighted-rank-sum": + raise ValueError("benchmark requires unweighted rank-sum combine") + valid = expert_ids >= 0 + expert = expert_ids.clamp(min=0).to(torch.int64) + gate = weights.to(torch.float32).masked_fill(~valid, 0) + scale, offset_a, offset_b = _expert_coefficients(torch, expert) + scale_sum = (gate * scale).sum(dim=1, keepdim=True) + offset_a_sum = (gate * offset_a).sum(dim=1, keepdim=True) + offset_b_sum = (gate * offset_b).sum(dim=1, keepdim=True) + pattern = _column_pattern(torch, payload.shape[1], payload.device) + transformed = ( + payload.float() * scale_sum + offset_a_sum + offset_b_sum * pattern.unsqueeze(0) + ) + return transformed.to(payload.dtype) + + +def _expected_transformed_combine(torch, problem, experts_per_rank, scale_up_domain): + """Reproduce the two-level reduction combine actually performs, so the + expectation carries the same BF16 rounding a correct backend does rather than + hiding it in a wide tolerance. Each destination rank casts its FP32 local + aggregate to the payload dtype (as _expert_transform does). Ranks sharing a + scale-up domain (NVLink/MNNVL) then reduce in FP32, and each domain casts its + aggregate to the payload dtype for the scale-out send before those + communicated BF16 partials are summed. When the whole EP group fits in one + scale-up domain (ep_size <= scale_up_domain — every EP8 case and the MNNVL + EP16 cases) there is a single domain and no scale-out rounding; a multi-node + RoCE EP16 group has one BF16 partial per node, and omitting that cast is what + left the scale-out combine ~0.048 off a single-domain reference.""" + semantic_x = getattr(problem, "oracle_x", problem.x) + expert_ids = problem.topk_idx.to(torch.int64) + weights = problem.topk_weights.to(torch.float32) + pattern = _column_pattern(torch, semantic_x.shape[1], semantic_x.device) + dtype = semantic_x.dtype + destination = expert_ids // experts_per_rank + ranks_per_domain = max(1, scale_up_domain) + domains: dict[int, object] = {} + scale, offset_a, offset_b = _expert_coefficients(torch, expert_ids) + for rank_id in destination.unique().tolist(): + gate = weights * (destination == rank_id) + # Per-rank BF16 output, FP32-accumulated within its scale-up domain. + contribution = ( + semantic_x.float() * (gate * scale).sum(dim=1, keepdim=True) + + (gate * offset_a).sum(dim=1, keepdim=True) + + (gate * offset_b).sum(dim=1, keepdim=True) * pattern.unsqueeze(0) + ).to(dtype).float() + domain = rank_id // ranks_per_domain + if domain in domains: + domains[domain] += contribution + else: + domains[domain] = contribution + # Each domain's aggregate is cast to the communicated payload dtype (the + # scale-out send) before the partials are summed. Unrouted tokens carry an + # exact zero through every level (all gates zero) — no mask needed. + expected = torch.zeros_like(semantic_x, dtype=torch.float32) + for domain in sorted(domains): + expected += domains[domain].to(dtype).float() + return expected + + +_ORACLE_CHECKS = ( + "combine_values", "counts", "metadata", "multiplicity", "payload", + "source_set", "weights", +) + + +def _oracle_report(**fields): + """One report shape for both the fail-soft and full oracle paths, so every + emitted correctness dict carries identical keys.""" + report = { + "passed": False, + "rel_tol": COMBINE_REL_TOL, + "mag_floor": COMBINE_MAG_FLOOR, + "combine_weight_semantics": "undeclared", + "receive_count": 0, + "max_absolute_error": None, + "max_elementwise_relative_error": None, + "max_weight_error": None, + "checks": dict.fromkeys(_ORACLE_CHECKS, False), + } + assert set(fields) <= set(report), sorted(set(fields) - set(report)) + report.update(fields) + assert set(report["checks"]) == set(_ORACLE_CHECKS) + return report + + +def _run_expert_oracle( + torch, + routing, + backend, + problem, + global_idx, + global_weights, + rank: int, + experts_per_rank: int, + scale_up_domain: int, + seed: int, +): + """Verify one real dispatch/transform/combine without entering a timed region.""" + handle = backend.dispatch(problem) + torch.cuda.synchronize() + try: + view = backend.inspect_dispatch(problem, handle) + source_ids = routing.decode_source_ids(view.payload, seed) + except Exception as inspection_error: + # Drain the in-flight dispatch before reporting: an abandoned handle + # would deadlock the other ranks. + try: + problem.recv_tokens = backend.recv_tokens(handle) + backend.stage(problem, handle) + backend.combine(problem, handle) + torch.cuda.synchronize() + except Exception as cleanup_error: + raise inspection_error from cleanup_error + return _oracle_report( + combine_weight_semantics=getattr( + backend, "combine_weight_semantics", "undeclared" + ), + ) + + receive_count = int(view.payload.shape[0]) + shape_ok = ( + view.payload.ndim == 2 + and view.expert_ids.shape == (receive_count, problem.topk_idx.shape[1]) + and view.weights.shape == view.expert_ids.shape + ) + source_range = bool( + receive_count == 0 + or ((source_ids >= 0) & (source_ids < global_idx.shape[0])).all().item() + ) + if source_range: + expected_idx = global_idx.to(problem.x.device).index_select(0, source_ids) + expected_weights = global_weights.to(problem.x.device).index_select(0, source_ids) + local = (expected_idx // experts_per_rank) == rank + expected_ids = torch.where(local, expected_idx, torch.full_like(expected_idx, -1)) + expected_weights = expected_weights.masked_fill(~local, 0) + expected_payload = routing.activations_for_source_ids( + source_ids, problem.x.shape[1], seed, problem.x.dtype + ) + else: + expected_ids = torch.full_like(view.expert_ids, -1) + expected_weights = torch.zeros_like(view.weights) + expected_payload = torch.empty_like(view.payload) + actual_ids, actual_weights = _normalized_expert_metadata( + torch, view.expert_ids, view.weights + ) + expected_ids, expected_weights = _normalized_expert_metadata( + torch, expected_ids, expected_weights + ) + expected_sources = ( + ((global_idx // experts_per_rank) == rank).any(dim=1).nonzero(as_tuple=True)[0] + ).to(problem.x.device) + source_set_ok = ( + source_range + and source_ids.numel() == torch.unique(source_ids).numel() + and torch.equal(torch.sort(source_ids).values, expected_sources) + ) + payload_ok = source_range and torch.equal(view.payload, expected_payload) + metadata_ok = shape_ok and torch.equal(actual_ids, expected_ids) + max_weight_error = ( + float((actual_weights - expected_weights).abs().max().item()) + if actual_weights.numel() + else 0.0 + ) + weights_ok = max_weight_error == 0.0 + valid_expected = expected_ids >= 0 + expected_local = expected_ids[valid_expected] - rank * experts_per_rank + expected_counts = torch.bincount(expected_local, minlength=experts_per_rank) + counts_ok = torch.equal( + view.local_expert_counts.to(torch.int64), expected_counts.to(torch.int64) + ) + multiplicity_ok = torch.equal( + (actual_ids >= 0).sum(dim=1), (expected_ids >= 0).sum(dim=1) + ) + problem.recv_tokens = receive_count + combine_weight_semantics = backend.combine_weight_semantics + transformed = _expert_transform( + torch, view.payload, actual_ids, actual_weights, combine_weight_semantics + ) + view.combine_input = transformed + combined = backend.combine_transformed(problem, handle, transformed) + torch.cuda.synchronize() + expected_combined = _expected_transformed_combine( + torch, problem, experts_per_rank, scale_up_domain + ) + if combined.shape == expected_combined.shape: + # Zero errors stand when the rank legitimately combined nothing. + max_absolute_error = max_elementwise_relative_error = 0.0 + combine_values_ok = True + if combined.numel(): + absolute_error = (combined.float() - expected_combined).abs() + max_absolute_error = float(absolute_error.max().item()) + max_elementwise_relative_error = float( + (absolute_error / expected_combined.abs().clamp_min(COMBINE_MAG_FLOOR)) + .max().item() + ) + combine_values_ok = max_elementwise_relative_error < COMBINE_REL_TOL + else: + max_absolute_error = max_elementwise_relative_error = None + combine_values_ok = False + checks = { + "combine_values": combine_values_ok, + "counts": counts_ok, + "metadata": metadata_ok, + "multiplicity": multiplicity_ok, + "payload": payload_ok, + "source_set": source_set_ok, + "weights": weights_ok, + } + return _oracle_report( + passed=all(checks.values()), + combine_weight_semantics=combine_weight_semantics, + receive_count=receive_count, + max_absolute_error=max_absolute_error, + max_elementwise_relative_error=max_elementwise_relative_error, + max_weight_error=max_weight_error, + checks=checks, + ) + + +def run_sweep(args, backend, torch, dist, device, rank: int, world_size: int) -> int: + """Drive the source-tokens-per-rank sweep for one fully-specified line.""" + mode = args.mode + if mode != "normal": + if rank == 0: + print(f"ERROR: unknown CollectiveX case mode {mode!r}") + return 2 + if min(args.iters, args.trials, args.warmup) <= 0: + if rank == 0: + print(f"ERROR: iters/trials/warmup must be positive; got " + f"{args.iters}:{args.trials}:{args.warmup}") + return 2 + import routing # torch-based; imported lazily so the module byte-compiles without torch + + ep_size = world_size + num_logical = getattr(args, "num_logical_experts", args.experts) + if args.experts % ep_size != 0: + if rank == 0: + print(f"ERROR: experts ({args.experts}) must divide ep_size ({ep_size})") + return 2 + experts_per_rank = args.experts // ep_size + gpn = args.gpus_per_node + scale_up_domain = args.scale_up_domain + suite = args.suite + workload_name = args.workload_name + if getattr(backend, "mode", None) != mode: + if rank == 0: + print(f"ERROR: backend mode {getattr(backend, 'mode', None)!r} != {mode!r}") + return 2 + expected_weight_semantics = "unweighted-rank-sum" + if getattr(backend, "combine_weight_semantics", None) != expected_weight_semantics: + if rank == 0: + print( + f"ERROR: {mode} requires combine semantics {expected_weight_semantics}" + ) + return 2 + + spec = backend.make_inputs(args) + if not spec.ok: + if rank == 0: + print(f"ERROR: {spec.message}") + return spec.rc + cap = spec.cap + ladder, dropped = spec.ladder, spec.dropped + if rank == 0 and dropped: + print(f"NOTE: dropped tokens/rank {dropped} — exceed {backend.name} buffer cap {cap} " + f"(hidden={args.hidden}); not silently truncated.") + MAX, MIN, SUM = dist.ReduceOp.MAX, dist.ReduceOp.MIN, dist.ReduceOp.SUM + + # Inputs determine the communicator capacity. + backend.create_buffer(spec) + + # ---- Pass 1: per shape, ascending (a cold-jump-safe ramp): warm untimed, + # then prove workload identity and run the expert oracle. The untimed warm + # rounds settle clocks/fabric BEFORE anything gate-bearing runs at that shape + # and are never measured or emitted. ---- + problems, gate, gts, global_traces, input_snapshots = {}, {}, {}, {}, {} + routing_consistent = True + for T in ladder: + gt = T * ep_size + gts[T] = gt + point = spec.points[T] + problem = backend.make_problem( + T, point.topk_idx.to(device), point.topk_weights.to(device), point.activations + ) + backend.warm(problem, CONDITIONING_ROUNDS_PER_SHAPE) + torch.cuda.synchronize() + problems[T] = problem + idx_g, w_g = point.global_idx, point.global_weights + rstats = routing.routing_stats(idx_g, args.experts, experts_per_rank) + rstats["locality"] = routing.routing_locality( + idx_g, experts_per_rank, ep_size, max(1, T), gpn, scale_up_domain + ) + point_routing_consistent = _same_tensors_across_ranks( + torch, dist, device, idx_g, w_g + ) + routing_consistent = routing_consistent and point_routing_consistent + input_snapshots[T] = ( + problem.x.clone(), problem.topk_idx.clone(), problem.topk_weights.clone() + ) + oracle = _run_expert_oracle( + torch, routing, backend, problem, idx_g, w_g, rank, experts_per_rank, + scale_up_domain, args.seed, + ) + before_x, before_idx, before_weights = input_snapshots[T] + pre_input_unchanged = ( + torch.equal(problem.x, before_x) + and torch.equal(problem.topk_idx, before_idx) + and torch.equal(problem.topk_weights, before_weights) + ) + global_traces[T] = (idx_g, w_g) + gate[T] = { + "rstats": rstats, + "recv_local": oracle["receive_count"], + "max_rel": oracle["max_elementwise_relative_error"] or 0.0, + "local_ok": int(oracle["passed"]), + "oracle_pre": oracle, + "pre_input_unchanged": pre_input_unchanged, + } + + # ---- Pass 2: every backend uses the same rotated point order. + # Per-iteration cross-rank MAX samples are pooled across trials. ---- + disp_pool = {T: [] for T in ladder} # pooled per-iteration cross-rank MAX (dispatch) + stage_pool = {T: [] for T in ladder} # measured only when stage launches device work + comb_pool = {T: [] for T in ladder} # ... combine + rt_pool = {T: [] for T in ladder} # independently measured round trip + for trial_index in range(args.trials): + order = trial_order(list(ladder), trial_index) + for T in order: + problem = problems[T] + # timed_components() encodes the roundtrip-only vs full-component contract + # (and whether stage launches device work) once, in the base class. + component_order = trial_order(backend.timed_components(), trial_index) + measured = {name: [] for name in ("dispatch", "stage", "combine", "roundtrip")} + for component_name in component_order: + # The base template gives every component the same synchronized + # full-roundtrip warm-up before its timed trial and encodes the two + # branch rules (dispatch cleanup, combine re-dispatch) internally. + measured[component_name] = backend.benchmark_component( + component_name, problem, args.warmup, args.iters + ) + # per-iteration cross-rank MAX (the distributed-op latency per iter), pooled. + if measured["dispatch"]: + disp_pool[T] += _reduce_vec(torch, dist, device, measured["dispatch"], MAX) + comb_pool[T] += _reduce_vec(torch, dist, device, measured["combine"], MAX) + if measured["stage"]: + stage_pool[T] += _reduce_vec(torch, dist, device, measured["stage"], MAX) + rt_pool[T] += _reduce_vec(torch, dist, device, measured["roundtrip"], MAX) + + # ---- Pass 3: prove timed inputs were immutable and repeat the full oracle. ---- + for T in ladder: + problem = problems[T] + before_x, before_idx, before_weights = input_snapshots[T] + input_unchanged = gate[T]["pre_input_unchanged"] and ( + torch.equal(problem.x, before_x) + and torch.equal(problem.topk_idx, before_idx) + and torch.equal(problem.topk_weights, before_weights) + ) + idx_g, w_g = global_traces[T] + post = _run_expert_oracle( + torch, routing, backend, problem, idx_g, w_g, rank, experts_per_rank, + scale_up_domain, args.seed, + ) + pre = gate[T]["oracle_pre"] + gate[T].update({ + "input_unchanged": input_unchanged, + "local_ok": int(pre["passed"] and post["passed"] and input_unchanged), + "max_rel": max( + pre["max_elementwise_relative_error"] or 0.0, + post["max_elementwise_relative_error"] or 0.0, + ), + "oracle_post": post, + }) + + # ---- Pass 4: percentiles (p50/p90/p95/p99, nearest-rank) from pooled samples + bytes + row ---- + rows = [] + for T in ladder: + gt = gts[T] + g = gate[T] + rstats = g["rstats"] + d, s, c, rt = disp_pool[T], stage_pool[T], comb_pool[T], rt_pool[T] + dp, sp, cp, rtp = _pcts(d), _pcts(s), _pcts(c), _pcts(rt) + # isolated_sum = SUM of the isolated dispatch+stage+combine percentiles. Stage contributes + # zero when it is explicitly not applicable. This is NOT a measured chained operation + # (can't reveal shared sync / launch amortization / overlap) — do NOT use for throughput + # or SLO capacity. The MEASURED round trip (rtp) is the real chained latency. + isum = ( + {key: dp[key] + (sp[key] if sp is not None else 0.0) + cp[key] for key in dp} + if dp and cp else None + ) + recv_total = _reduce_int(torch, dist, device, g["recv_local"], SUM) + recv_max = _reduce_int(torch, dist, device, g["recv_local"], MAX) + recv_min = _reduce_int(torch, dist, device, g["recv_local"], MIN) + global_ok = _reduce_int(torch, dist, device, g["local_ok"], MIN) + max_rel = _reduce_vec(torch, dist, device, [g["max_rel"]], MAX)[0] + point_ok = bool(global_ok) and recv_total > 0 + throughput = { + percentile_name: gt / (latency_us * 1e-6) + for percentile_name, latency_us in rtp.items() + } + # Canonical LOGICAL payload bytes come from the routing trace (NOT backend recv + # tensors): one copy per unique (token, dest-rank) pair. Dispatch and combine + # move the same logical payload; the roundtrip is their sum and stage moves nothing. + one_way_bytes = logical_byte_provenance(rstats["routed_copies"], args.hidden) + roundtrip_bytes = {field: 2 * value for field, value in one_way_bytes.items()} + stage_bytes = dict.fromkeys(one_way_bytes, 0) + rows.append({ + "components": { + "combine": _component(cp, len(c)), + "dispatch": _component(dp, len(d)), + "isolated_sum": _component(isum, 0, derived=True), + "roundtrip": _component(rtp, len(rt)), + "stage": _component(sp, len(s)), + }, + "correctness": { + # Max elementwise relative error (COMBINE_MAG_FLOOR-clamped) + # against the BF16-faithful expected combine. + "max_relative_error": max_rel, + "passed": point_ok, + }, + "global_tokens": gt, + "byte_provenance": { + "combine": one_way_bytes, + "dispatch": one_way_bytes, + "roundtrip": roundtrip_bytes, + "stage": stage_bytes, + }, + "receive": { + "max": recv_max, + "mean": recv_total / world_size, + "min": recv_min, + "total": recv_total, + }, + "routing": {key: rstats[key] for key in _ROUTING_FIELDS}, + "token_rate_at_latency_percentile": throughput, + "tokens_per_rank": T, + }) + if rank == 0: + component_log = (f"disp p50/p99={dp['p50']:7.1f}/{dp['p99']:7.1f} " + f"comb {cp['p50']:6.1f}/{cp['p99']:6.1f} " if dp and cp + else "components=unavailable ") + print(f" T={T:<5} {component_log}" + f"RT p50/p99={rtp['p50']:7.1f}/{rtp['p99']:7.1f}us n={len(rt)} fanout={rstats['fanout_mean']:.2f} " + f"recv[min/mean/max]={recv_min}/{recv_total // world_size}/{recv_max} " + f"correct={point_ok}") + + # status=valid requires correctness AND a proven-identical routing trace across ranks. + all_ok = bool(rows) and all(r["correctness"]["passed"] for r in rows) and routing_consistent + + generated_at = _dt.datetime.now().astimezone().isoformat() + nodes = int(os.environ.get("SLURM_NNODES", "1")) + scheduled_case = { + "backend": backend.name, + "ep": ep_size, + "experts": num_logical, + "gpus_per_node": gpn, + "hidden": args.hidden, + "ladder": " ".join(map(str, ladder)), + "mode": mode, + "nodes": nodes, + "phase": args.phase, + "routing": args.routing, + "scale_up_domain": scale_up_domain, + "scale_up_transport": args.scale_up_transport, + "scale_out_transport": args.scale_out_transport or None, + "scope": args.scope, + "suite": suite, + "topk": args.topk, + "topology_class": args.topology_class, + "transport": args.transport, + "workload": workload_name, + } + case_factors = {"case": scheduled_case, "sku": args.runner} + computed_case_id = case_id(args.runner, scheduled_case) + if args.case_id != computed_case_id: + raise ValueError( + f"scheduled case ID does not match realized factors: {args.case_id} != {computed_case_id}" + ) + git_run = getattr(args, "git_run", None) or {} + allocation_factors = { + "run_attempt": git_run.get("run_attempt"), + "run_id": git_run.get("run_id"), + "source_sha": git_run.get("source_sha"), + } + try: + attempt_ordinal = int(os.environ.get("COLLX_ATTEMPT_ID", "1")) + except ValueError: + attempt_ordinal = 0 + if attempt_ordinal <= 0: + raise ValueError("COLLX_ATTEMPT_ID must be a positive integer") + doc = { + "version": args.version, + "record_type": "case-attempt", + "generated_at": generated_at, + "identity": { + "allocation_factors": allocation_factors, + "attempt_ordinal": attempt_ordinal, + "case_factors": case_factors, + "case_id": args.case_id, + }, + "workload": { + "cross_rank_consistent": routing_consistent, + }, + "measurement": { + "combine_dtype": "bf16", + "combine_semantics": "activation-only", + "dispatch_dtype": "bf16", + "payload_unit": "token-rank", + "rows": rows, + "sampling": { + "iterations_per_trial": args.iters, + "samples_per_component": args.iters * args.trials, + "trials": args.trials, + "warmup_iterations": args.warmup, + }, + }, + "implementation": { + "kernel_generation": kernel_generation(backend), + "name": backend.name, + }, + "topology": { + "device_product": getattr(args, "runtime_device_product", None), + "gpus_per_node": gpn, + "nodes": nodes, + "placement": "packed", + "scale_up_domain": scale_up_domain, + "scale_up_transport": args.scale_up_transport, + "scale_out_transport": args.scale_out_transport or None, + "scope": args.scope, + "topology_class": args.topology_class, + "transport": args.transport, + "world_size": world_size, + }, + "runtime": getattr(args, "runtime", {}), + "provenance": { + "image": getattr(args, "image", "") or None, + "source_sha": git_run.get("source_sha"), + }, + "outcome": { + "reasons": [] if all_ok else ["semantic correctness or routing identity failed"], + "status": "success" if all_ok else "invalid", + }, + } + if rank == 0: + _write_json_atomic(args.out, doc) + # Ladder ends + two interior points — one mid-ladder headline hides the + # low-token (startup-dominated) behavior. + summary_rows = [] + for tokens in (ladder[0], 8, 64, ladder[-1]): + row = next((r for r in rows if r["tokens_per_rank"] == tokens), None) + if row is not None and row not in summary_rows: + summary_rows.append(row) + + def _point_summary(row): + percentiles = row["components"]["dispatch"]["percentiles_us"] + if not percentiles: + return f"T={row['tokens_per_rank']}:n/a" + return f"T={row['tokens_per_rank']}:disp_p99={percentiles['p99']:.1f}us" + + component_summary = " ".join(_point_summary(row) for row in summary_rows) + print(f"{backend.name} ep-dispatch-combine [{args.phase}/{mode}]: " + f"status={doc['outcome']['status']} {len(rows)} pts, routing_consistent={routing_consistent}, " + f"{component_summary} " + f"-> {args.out}") + # CI honesty: run_sweep's return code is the only success signal collx_run_shard (and thus CI) + # reads — the doc is uploaded regardless, via the launcher's always() stage step. A captured + # `invalid` outcome (semantic correctness or cross-rank routing identity failed) must therefore + # fail the leg, not ride as a green success; otherwise a persistent oracle failure is invisible + # in CI and could autopublish an invalid doc. Agree the verdict across ranks (MIN) so every + # rank exits identically and the distributed case fails as one. + outcome_ok = bool(_reduce_int(torch, dist, device, int(all_ok), dist.ReduceOp.MIN)) + return 0 if outcome_ok else 3 diff --git a/experimental/CollectiveX/bench/ep_mori.py b/experimental/CollectiveX/bench/ep_mori.py new file mode 100644 index 0000000000..c0daa62783 --- /dev/null +++ b/experimental/CollectiveX/bench/ep_mori.py @@ -0,0 +1,264 @@ +#!/usr/bin/env python3 +"""CollectiveX MoRI adapter: native BF16 dispatch/combine over mori.ops.""" +from __future__ import annotations + +import os +import sys +import types + +# MoRI reads the symmetric-heap size when the heap is created (at shmem init, +# once a process group exists — see create_buffer). The pinned upstream +# inter-node benchmark uses 6 GiB for its InterNodeV1 staging and signal +# buffers; under VMM_HEAP this is a virtual reservation backed on demand. +os.environ["MORI_SHMEM_HEAP_SIZE"] = "6G" + +import torch +import torch.distributed as dist + +from ep_backend import EPBackend + +try: + import mori # type: ignore +except Exception as exc: # pragma: no cover - requires the benchmark image + print(f"ERROR: mori import failed: {exc!r}", file=sys.stderr) + raise + + +def _project_local_metadata(torch_module, raw_expert_ids, raw_weights, rank, experts_per_rank): + local_start = rank * experts_per_rank + local = (raw_expert_ids >= local_start) & ( + raw_expert_ids < local_start + experts_per_rank + ) + expert_ids = torch_module.where( + local, raw_expert_ids, torch_module.full_like(raw_expert_ids, -1) + ) + weights = torch_module.where(local, raw_weights, torch_module.zeros_like(raw_weights)) + return expert_ids, weights, raw_expert_ids[local] - local_start + + +class MoRIBackend(EPBackend): + name = "mori" + combine_needs_redispatch = True + dispatch_needs_combine_cleanup = True + + def __init__(self, args, rank, world_size, local_rank, device): + super().__init__(args, rank, world_size, local_rank, device) + self.ep_size = world_size + self.experts_per_rank = args.experts // self.ep_size + gpus_per_node = int(args.gpus_per_node) + scale_out = args.scope == "scale-out" + + # The kernel is a pinned function of the cell, not an operator choice, + # and no cell runs a low-latency-family kernel (mirroring the + # normal-mode-only NVIDIA backends): scale-up uses the direct IntraNode + # kernel on every CDNA SKU (mori's default; `kernel_type` kwarg + # omitted); scale-out EP16 uses InterNodeV1, whose required enum member + # is an image-lineage check. + # (kernel, generation label, (block_num, rdma_block_num, dispatch_warps, combine_warps)) + kernel_name, self.kernel_generation, blocks = ( + ("InterNodeV1", "inter-node-v1", (96, 64, 8, 8)) if scale_out + else ("IntraNode", "intranode", (80, 0, 16, 8)) + ) + self._kernel_type = None + if kernel_name != "IntraNode": + kernel_enum = getattr(mori.ops, "EpDispatchCombineKernelType", None) + if kernel_enum is None or not hasattr(kernel_enum, kernel_name): + raise RuntimeError( + f"this MoRI image lacks EpDispatchCombineKernelType.{kernel_name}" + ) + self._kernel_type = getattr(kernel_enum, kernel_name) + self._inter_node = kernel_name == "InterNodeV1" + self.num_qps = 1 + self.block_num, self.rdma_block_num, self.dispatch_warps, self.combine_warps = blocks + self._external_input = self._inter_node + # Registered-input MoRI copies expert output into a device-side symmetric buffer. External + # input kernels consume the dispatch output directly, so their stage is not applicable. + self.stage_device_work = not self._external_input + # Stash the __init__-only locals the moved create_buffer body reads back. + self._gpus_per_node = gpus_per_node + + def create_buffer(self, spec): + args, world_size, rank = self.args, self.world_size, self.rank + gpus_per_node = self._gpus_per_node + + world_group = torch.distributed.group.WORLD + torch._C._distributed_c10d._register_process_group("default", world_group) + # Scale-out EP16 registers the symmetric heap over the AMD AI NIC. The + # default STATIC_HEAP registers it as one contiguous MR; on the Ionic + # stack that registration fails during InterNodeV1 init (an EINVAL that + # is a firmware command failure, not an MR-size violation — the NIC + # advertises multi-GiB max_mr_size). VMM_HEAP backs the same reservation + # with on-demand 64 MiB DMA-BUF chunks, the supported inter-node path + # (MoRI PR #155, validated on MI355X + AI NIC). Read at heap init below, + # so it must precede shmem_torch_process_group_init. Scale-up EP8 is + # intranode (no RDMA registration) and keeps the default heap. + if self._inter_node: + os.environ["MORI_SHMEM_MODE"] = "VMM_HEAP" + mori.shmem.shmem_torch_process_group_init("default") + realized_qps = int(mori.shmem.shmem_num_qp_per_pe()) + if realized_qps < self.num_qps: + raise RuntimeError( + f"MoRI realized {realized_qps} QPs per PE; {self.num_qps} required" + ) + + # MoRI preallocates one communicator buffer for the case's entire ladder. + self._cap = max(512, spec.max_tokens_per_rank) + # BF16 communication path: no FP8 dispatch dtype, no scale payload, no combine quantization. + dispatch_dtype = torch.bfloat16 + config_kwargs = { + "data_type": dispatch_dtype, + "rank": rank, + "world_size": world_size, + "hidden_dim": args.hidden, + "scale_dim": 0, + "scale_type_size": 1, + "max_token_type_size": torch.tensor([], dtype=dispatch_dtype).element_size(), + "max_num_inp_token_per_rank": self._cap, + "num_experts_per_rank": self.experts_per_rank, + "num_experts_per_token": args.topk, + "use_external_inp_buf": self._external_input, + "quant_type": "none", + } + if self._kernel_type is not None: + config_kwargs["kernel_type"] = self._kernel_type + if self._inter_node: + config_kwargs.update({ + "block_num": self.block_num, + "warp_num_per_block": self.dispatch_warps, + "gpu_per_node": gpus_per_node, + "rdma_block_num": self.rdma_block_num, + "num_qp_per_pe": self.num_qps, + }) + self.config = mori.ops.EpDispatchCombineConfig(**config_kwargs) + expected_config = { + "data_type": dispatch_dtype, + "scale_dim": 0, + "scale_type_size": 1, + "use_external_inp_buf": self._external_input, + "quant_type": config_kwargs["quant_type"], + } + if self._inter_node: + expected_config.update({ + "block_num": self.block_num, + "warp_num_per_block": self.dispatch_warps, + "gpu_per_node": 8, + "rdma_block_num": 64, + "num_qp_per_pe": 1, + }) + if any(getattr(self.config, key, None) != value for key, value in expected_config.items()): + raise RuntimeError("MoRI requested launch/topology configuration was not realized") + # The newer pinned MoRI revision can otherwise replace explicit values + # with token-dependent tuning rules from the image. + os.environ["MORI_EP_LAUNCH_CONFIG_MODE"] = "MANUAL" + self.op = mori.ops.EpDispatchCombineOp(self.config) + if getattr(self.op, "launch_config_mode", None) != "MANUAL": + raise RuntimeError("MoRI explicit launch configuration was not applied") + + def make_problem(self, T, idx, weights, x): + indices = idx.to(torch.int32) + gate_weights = weights.to(torch.float32) + return types.SimpleNamespace( + T=T, + x=x, + dispatch_x=x, + topk_idx=indices, + topk_weights=gate_weights, + indices=indices, + weights=gate_weights, + scales=torch.empty((T, 0), dtype=torch.uint8, device=self.device), + ) + + def dispatch(self, p): + dispatch_output, dispatch_weights, _scales, dispatch_indices, recv_num = ( + self.op.dispatch( + p.dispatch_x, + p.weights, + p.scales, + p.indices, + block_num=self.block_num, + rdma_block_num=self.rdma_block_num, + warp_per_block=self.dispatch_warps, + ) + ) + return types.SimpleNamespace( + dispatch_output=dispatch_output, + dispatch_weights=dispatch_weights, + dispatch_indices=dispatch_indices, + recv_num=recv_num[0], + combine_input=None, + ) + + def stage(self, p, h): + rows = getattr(p, "recv_tokens", None) + if not isinstance(rows, int) or rows < 0 or rows > h.dispatch_output.size(0): + raise RuntimeError("MoRI receive count was not validated before staging") + h.combine_input = h.dispatch_output + if self._external_input: + return None + buffer = self.op.get_registered_combine_input_buffer( + torch.bfloat16, hidden_dim=h.combine_input.size(1) + ) + buffer[:rows, :].copy_(h.combine_input[:rows, :]) + h.combine_input = buffer + + def combine(self, p, h): + combined, _weights = self.op.combine( + h.combine_input, + None, + h.dispatch_indices, + block_num=self.block_num, + rdma_block_num=self.rdma_block_num, + warp_per_block=self.combine_warps, + ) + return combined[:p.T] + + def inspect_dispatch(self, p, h): + count = self.recv_tokens(h) + if h.dispatch_weights is None: + raise RuntimeError("MoRI dispatch did not expose gate weights") + if count < 0 or any( + tensor.ndim == 0 or count > tensor.size(0) + for tensor in (h.dispatch_output, h.dispatch_indices, h.dispatch_weights) + ): + raise RuntimeError("MoRI receive count exceeds dispatch metadata") + raw_expert_ids = h.dispatch_indices[:count].to(torch.int64) + expert_ids, weights, local_expert_ids = _project_local_metadata( + torch, + raw_expert_ids, + h.dispatch_weights[:count].to(torch.float32), + self.rank, + self.experts_per_rank, + ) + return types.SimpleNamespace( + payload=h.dispatch_output[:count], + expert_ids=expert_ids, + weights=weights, + local_expert_counts=torch.bincount( + local_expert_ids, minlength=self.experts_per_rank + ), + ) + + def combine_transformed(self, p, h, transformed): + h.combine_input = transformed.to(torch.bfloat16) + rows = getattr(p, "recv_tokens", None) + if not isinstance(rows, int) or rows < 0 or rows > h.combine_input.size(0): + raise RuntimeError("MoRI receive count was not validated before transformed combine") + if not self._external_input: + buffer = self.op.get_registered_combine_input_buffer( + torch.bfloat16, hidden_dim=h.combine_input.size(1) + ) + buffer[:rows, :].copy_(h.combine_input[:rows, :]) + h.combine_input = buffer + return self.combine(p, h) + + def recv_tokens(self, h): + return int(h.recv_num.item()) + + def finalize(self, rc): + try: + dist.barrier() + except Exception: + pass + sys.stdout.flush() + sys.stderr.flush() + os._exit(rc if 0 <= rc <= 255 else 1) diff --git a/experimental/CollectiveX/bench/routing.py b/experimental/CollectiveX/bench/routing.py new file mode 100644 index 0000000000..b61d2be160 --- /dev/null +++ b/experimental/CollectiveX/bench/routing.py @@ -0,0 +1,173 @@ +#!/usr/bin/env python3 +"""Deterministic platform-independent MoE routing and activations.""" +from __future__ import annotations + +import hashlib +import struct + +import torch + +_MASK64 = (1 << 64) - 1 + +SOURCE_ID_BITS = 32 +SOURCE_ID_COLUMNS = SOURCE_ID_BITS + + +def build_global_routing(global_tokens: int, experts: int, topk: int, routing: str, seed: int): + """Return one byte-stable counter-generated routing window on CPU.""" + if routing != "uniform": + raise ValueError(f"unknown routing {routing!r} (uniform)") + if global_tokens <= 0 or experts <= 0 or topk <= 0 or topk > experts: + raise ValueError("global_tokens/experts/topk must be positive and topk <= experts") + + # Keyed BLAKE2b as a deterministic random oracle over the coordinate tuple: + # stdlib, byte-stable on every runtime, and uniform by construction. + key = (int(seed) & _MASK64).to_bytes(8, "little") + + def counter(token: int, slot: int, attempt: int, stream: int) -> int: + message = struct.pack("<4Q", token, slot, attempt, stream) + return int.from_bytes( + hashlib.blake2b(message, key=key, digest_size=8).digest(), "little" + ) + + indices, weights = [], [] + for token in range(int(global_tokens)): + selected, used = [], set() + for slot in range(int(topk)): + attempt = 0 + while True: + expert = counter(token, slot, attempt, 0) % int(experts) + if expert not in used: + used.add(expert) + selected.append(expert) + break + attempt += 1 + raw = [1 + counter(token, slot, 0, 1) % 65535 for slot in range(int(topk))] + denominator = float(sum(raw)) + indices.append(selected) + weights.append([value / denominator for value in raw]) + return ( + torch.tensor(indices, dtype=torch.int64), + torch.tensor(weights, dtype=torch.float32), + ) + + +def rank_slice(idx, weights, rank: int, tokens_per_rank: int): + lo = rank * tokens_per_rank + return idx[lo:lo + tokens_per_rank].contiguous(), weights[lo:lo + tokens_per_rank].contiguous() + + +def rank_activations(tokens: int, hidden: int, seed: int, rank: int, device, + dtype=torch.bfloat16): + """Exact counter-derived inputs with a quantization-safe source-token prefix.""" + source = torch.arange(tokens, device=device, dtype=torch.int64) + rank * tokens + return activations_for_source_ids(source, hidden, seed, dtype) + + +def activations_for_source_ids(source, hidden: int, seed: int, dtype=torch.bfloat16): + """Materialize canonical activations for arbitrary global source-token IDs.""" + if hidden < SOURCE_ID_COLUMNS: + raise ValueError(f"hidden must be at least {SOURCE_ID_COLUMNS}") + source = source.to(torch.int64) + column = torch.arange(hidden, device=source.device, dtype=torch.int64) + # Integer lattice, frozen (the oracle regenerates these exact bytes): + # 131/17/19 are odd multipliers coprime to the prime modulus 257, so distinct (source, column, seed) + # land on distinct residues and corruption cannot alias into a correct-looking pattern; %257-128 yields k in + # [-128, 128] and k/64 is exactly representable in bfloat16, so the oracle's bit-exact payload compare sees + # transport corruption, never representation error. + values = (source[:, None] * 131 + column[None, :] * 17 + int(seed) * 19) % 257 - 128 + output = values.to(dtype).mul_(1 / 64) + if bool((source < 0).any().item()) or bool((source >= (1 << SOURCE_ID_BITS)).any().item()): + raise ValueError("source token ID is outside the bounded identity contract") + source_columns = torch.arange(SOURCE_ID_BITS, device=source.device, dtype=torch.int64) + source_bits = ((source[:, None] >> source_columns[None, :]) & 1) * 2 - 1 + # Magnitude one sits inside the ordinary [-2, 2] activation range, so the identity cannot set + # an FP8 block scale. Decode depends only on sign and remains stable after dequantization. + output[:, :SOURCE_ID_BITS] = source_bits.to(dtype) + return output + + +def decode_source_ids(payload, seed: int): + """Decode and validate source IDs carried by rank_activations.""" + if payload.ndim != 2 or payload.shape[1] < SOURCE_ID_COLUMNS: + raise ValueError("received payload cannot carry the source-token prefix") + prefix = payload[:, :SOURCE_ID_COLUMNS].float() + if not bool(torch.isfinite(prefix).all().item()) or bool((prefix.abs() < 0.25).any().item()): + raise ValueError("received source-token prefix is not quantization-stable") + bits = prefix >= 0 + powers = 1 << torch.arange(SOURCE_ID_BITS, device=payload.device, dtype=torch.int64) + source = (bits[:, :SOURCE_ID_BITS].to(torch.int64) * powers).sum(dim=1) + return source + + +def routing_locality(idx, experts_per_rank: int, ep_size: int, tokens_per_rank: int, + gpus_per_node: int, scale_up_domain: int = None) -> dict: + """Locality of rank-deduplicated payload copies under packed placement.""" + gt = idx.shape[0] + assignments = (idx // experts_per_rank).clamp(max=ep_size - 1) + destinations = torch.zeros((gt, ep_size), dtype=torch.bool) + destinations.scatter_(1, assignments, True) + token, dest = destinations.nonzero(as_tuple=True) + src = (token // max(1, tokens_per_rank)).clamp(max=ep_size - 1) + sud = scale_up_domain or (gpus_per_node * ep_size) # default: all one domain + phys = torch.arange(ep_size, dtype=torch.int64) + pd, ps = phys[dest], phys[src] + local = (dest == src) + same_node = (pd // gpus_per_node) == (ps // gpus_per_node) + same_dom = (pd // sud) == (ps // sud) + n = dest.numel() + return { + "placement": "packed", + "local_rank_fraction": float(local.float().mean()), + "same_node_fraction": float(same_node.float().mean()), + "same_scaleup_domain_fraction": float(same_dom.float().mean()), + "cross_node_fraction": float((~same_node).float().mean()), + "cross_domain_fraction": float((~same_dom).float().mean()), + "gpus_per_node": gpus_per_node, "scale_up_domain": sud, "copies": int(n), + } + + +def routing_stats(idx, experts: int, experts_per_rank: int) -> dict: + """Realized routing properties for the GLOBAL trace — published per point so the + fan-out / load can never be silently misread. idx is the global [gt, topk] tensor; + """ + ep = max(1, experts // max(1, experts_per_rank)) + ranks = (idx // experts_per_rank) # [gt, topk] destination rank per assignment + # unique destination ranks per token (fan-out) + onehot = torch.zeros(idx.shape[0], ep, dtype=torch.bool) + onehot.scatter_(1, ranks.clamp(max=ep - 1), True) + fanout = onehot.sum(dim=1) # [gt] + hist = torch.bincount(fanout, minlength=ep + 1)[1:ep + 1].tolist() # counts for fan-out 1..ep + load = torch.bincount(idx.reshape(-1), minlength=experts).float() + # Keep expert assignments (compute load) separate from rank-deduplicated payload copies + # (network load). Conflating them overstates traffic when two experts share a rank. + assignment_load = torch.bincount( + ranks.reshape(-1).clamp(max=ep - 1), minlength=ep + ).float() + payload_load = onehot.sum(dim=0).float() + # One-number imbalance summaries so a row is self-describing for the distribution-sensitivity + # suite (no need to read the full histograms): CV = std/mean of the load; hotspot_ratio = + # worst expert load over the mean. + def _cv(t): + m = float(t.mean()) + return float(t.std(unbiased=False) / m) if m > 0 else 0.0 + expert_load_cv = _cv(load) + assignment_rank_cv = _cv(assignment_load) + payload_rank_cv = _cv(payload_load) + hotspot_ratio = float(load.max() / load.mean()) if float(load.mean()) > 0 else 0.0 + # Empty experts capture compute skew; empty destination ranks capture network skew. + empty_expert_count = int((load == 0).sum()) + empty_rank_count = int((payload_load == 0).sum()) + return { + "fanout_mean": float(fanout.float().mean()), + "fanout_min": int(fanout.min()), "fanout_max": int(fanout.max()), + "fanout_histogram": hist, # index k-1 = #tokens with fan-out k + "expert_assignments_per_rank": [int(x) for x in assignment_load.tolist()], + "payload_copies_per_rank": [int(x) for x in payload_load.tolist()], + "routed_copies": int(fanout.sum()), # total (token, dest-rank) pairs + "expert_load_min": int(load.min()), "expert_load_max": int(load.max()), + "expert_load_mean": float(load.mean()), "expert_load_cv": expert_load_cv, + "expert_assignment_rank_cv": assignment_rank_cv, + "payload_rank_cv": payload_rank_cv, "hotspot_ratio": hotspot_ratio, + "empty_expert_count": empty_expert_count, "empty_rank_count": empty_rank_count, + } diff --git a/experimental/CollectiveX/bench/run_ep.py b/experimental/CollectiveX/bench/run_ep.py new file mode 100644 index 0000000000..ef96436ec9 --- /dev/null +++ b/experimental/CollectiveX/bench/run_ep.py @@ -0,0 +1,141 @@ +#!/usr/bin/env python3 +"""CollectiveX v1 EP benchmark entrypoint for torchrun or rank environments.""" + +from __future__ import annotations + +import argparse +import ctypes +import os +import sys + +# Make the sibling bench/ modules importable when run as `bench/run_ep.py` under +# torchrun (it executes the file as __main__, not as a package). +HERE = os.path.dirname(os.path.abspath(__file__)) +sys.path[:0] = [HERE, os.path.dirname(HERE)] + +import ep_harness # noqa: E402 (stdlib-only; safe before torch) + + +def _loaded_collective_version() -> str | None: + try: + with open("/proc/self/maps", encoding="utf-8") as handle: + paths = { + os.path.realpath(line.rstrip().split()[-1]) + for line in handle + if any(name in line for name in ("libnccl.so", "librccl.so")) + and os.path.isfile(line.rstrip().split()[-1]) + } + if len(paths) != 1: + return None + version = ctypes.c_int() + library = ctypes.CDLL(paths.pop()) + if library.ncclGetVersion(ctypes.byref(version)) != 0: + return None + return ep_harness.format_collective_version(version.value) + except (AttributeError, OSError): + return None + + +def _runtime_info(torch, *, vendor: str) -> dict: + """Return the runtime versions needed to compare and debug results.""" + runtime_kind = "cuda" if vendor == "nvidia" else "hip" + collective_kind = "nccl" if vendor == "nvidia" else "rccl" + return { + "accelerator_runtime": getattr(torch.version, runtime_kind), + "collective_library": {"kind": collective_kind, "version": _loaded_collective_version()}, + "framework": str(torch.__version__), + "vendor": vendor, + } + + +def main() -> int: + ap = argparse.ArgumentParser(description="CollectiveX EP dispatch/combine sweep") + ap.add_argument("--backend", required=True, choices=["deepep-v2", "mori"]) + ep_harness.add_common_args(ap) + args = ap.parse_args() + + if not ep_harness.is_case_id(args.case_id): + print(f"ERROR: invalid native case ID {args.case_id!r}", file=sys.stderr) + return 2 + # Seed and timing arrive baked into the case argv from the single + # configs/sweep.json source; there is no separate canonical constant to + # cross-check against. + + try: + import torch + import torch.distributed as dist + except Exception as exc: # pragma: no cover + print(f"ERROR: torch unavailable: {exc!r}", file=sys.stderr) + return 3 + + rank = int(os.environ.get("RANK", "0")) + world_size = int(os.environ.get("WORLD_SIZE", "1")) + local_rank = int(os.environ.get("LOCAL_RANK", "0")) + torch.cuda.set_device(local_rank) + device = torch.device(f"cuda:{local_rank}") + + vendor = "amd" if torch.version.hip else "nvidia" + device_name = torch.cuda.get_device_name(device) + args.runtime_device_product = device_name + args.image = os.environ.get("COLLECTIVEX_IMAGE", "") + _run = { + "run_id": os.environ.get("GITHUB_RUN_ID"), + "run_attempt": os.environ.get("GITHUB_RUN_ATTEMPT"), + "source_sha": os.environ.get("COLLECTIVEX_SOURCE_SHA") + or os.environ.get("GITHUB_SHA"), + } + args.git_run = _run if any(_run.values()) else None + + # Import the backend class only after torch initializes. The selected mode is an + # explicit case dimension; adapters do not infer it from the token ladder. + if args.backend == "mori": + from ep_mori import MoRIBackend as Backend + else: + from ep_deepep_v2 import DeepEPV2Backend as Backend + + # MoRI registers the default GPU process group with its SHMEM runtime. Keep that + # group device-only so scale-out does not also depend on a host Gloo fabric. + if not dist.is_initialized(): + if args.backend == "mori": + dist.init_process_group( + backend="nccl", + rank=rank, + world_size=world_size, + device_id=device, + ) + else: + # PR #605 reuses PyTorch's NCCL communicator through ``_comm_ptr``. Supplying + # device_id eagerly forms it before ElasticBuffer construction. + dist.init_process_group("nccl", device_id=device) + + args.runtime = _runtime_info(torch, vendor=vendor) + + # Construct + run inside a try so a backend exception (esp. a new adapter on GPU) prints its + # FULL traceback to STDOUT — torchrun captures per-rank stdout but only summarizes stderr, so an + # uncaught exception is otherwise invisible in CI. Print on every rank (prefixed) then re-raise. + try: + backend = Backend(args, rank, world_size, local_rank, device) + if rank == 0: + print( + f"[run_ep] backend={args.backend} phase={args.phase} mode={args.mode} " + f"world={world_size} ep_size={world_size} hidden={args.hidden} " + f"topk={args.topk} experts={args.experts} dtype=bf16 " + f"routing={args.routing} seed={args.seed}" + ) + rc = ep_harness.run_sweep(args, backend, torch, dist, device, rank, world_size) + except Exception: + import traceback + + print( + f"[run_ep][rank{rank}] backend={args.backend} FAILED:\n" + + traceback.format_exc(), + flush=True, + ) + raise + # finalize() handles backend-specific teardown: DeepEP returns rc cleanly; + # MoRI hard-exits past its post-shmem_finalize teardown assertion. + return backend.finalize(rc) + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/experimental/CollectiveX/configs/platform_config.json b/experimental/CollectiveX/configs/platform_config.json new file mode 100644 index 0000000000..b89487ca1a --- /dev/null +++ b/experimental/CollectiveX/configs/platform_config.json @@ -0,0 +1,175 @@ +{ + "platforms": { + "h100-dgxc": { + "arch": "sm90", + "product": "h100", + "image": "lmsysorg/sglang:v0.5.11-cu130", + "image_platform": "linux/amd64", + "gpus_per_node": 8, + "scale_up_domain": 8, + "scale_up_transport": "nvlink", + "launcher": "single-slurm", + "backends": {"deepep-v2": [8, 16]}, + "fabric": {"nic": "ConnectX-7 2x200GbE", "switch": "Arista 7060DX5-64S (Tomahawk4, 25.6T)"}, + "operator": { + "partition": "hpc-gpu-1", + "account": "customer", + "squash_dir": "/mnt/nfs/sa-shared/cx-squash", + "exclude_nodes": "hpc-gpu-1-0,hpc-gpu-1-1,hpc-gpu-1-4,hpc-gpu-1-5,hpc-gpu-1-7,hpc-gpu-1-8,hpc-gpu-1-13,hpc-gpu-1-16,hpc-gpu-1-19" + }, + "network": { + "socket_ifname": "eth0", + "rdma_devices": "mlx5_18,mlx5_19,mlx5_20,mlx5_21,mlx5_22,mlx5_23,mlx5_24,mlx5_25,mlx5_26,mlx5_27,mlx5_28,mlx5_29,mlx5_30,mlx5_31,mlx5_32,mlx5_33", + "ib_gid_index": "3" + } + }, + "h200-dgxc": { + "arch": "sm90", + "product": "h200", + "image": "lmsysorg/sglang:v0.5.11-cu130", + "image_platform": "linux/amd64", + "gpus_per_node": 8, + "scale_up_domain": 8, + "scale_up_transport": "nvlink", + "launcher": "single-slurm", + "backends": {"deepep-v2": [8, 16]}, + "fabric": {"nic": "ConnectX-7 400G", "switch": "NVIDIA Quantum-2 QM9790 (25.6T, InfiniBand)"}, + "operator": { + "partition": "main", + "squash_dir": "/home/sa-shared/containers" + }, + "network": { + "rdma_devices": "mlx5_0,mlx5_1,mlx5_2,mlx5_3,mlx5_4,mlx5_5,mlx5_6,mlx5_7" + } + }, + "b200-dgxc": { + "arch": "sm100", + "product": "b200", + "image": "lmsysorg/sglang:v0.5.11-cu130", + "image_platform": "linux/amd64", + "gpus_per_node": 8, + "scale_up_domain": 8, + "scale_up_transport": "nvlink", + "launcher": "single-slurm", + "backends": {"deepep-v2": [8, 16]}, + "fabric": {"nic": "ConnectX-7 400GbE", "switch": "Whitebox Tomahawk3 leaf + Tomahawk4 (RoCE)"}, + "operator": { + "partition": "gpu-2", + "account": "benchmark", + "qos": "gpu-2_qos", + "squash_dir": "/home/sa-shared/containers" + }, + "network": { + "rdma_devices": "mlx5_0,mlx5_1,mlx5_2,mlx5_3,mlx5_4,mlx5_5,mlx5_6,mlx5_7", + "ib_gid_index": "3" + } + }, + "b300": { + "arch": "sm103", + "product": "b300", + "image": "lmsysorg/sglang:v0.5.11-cu130", + "image_platform": "linux/amd64", + "gpus_per_node": 8, + "scale_up_domain": 8, + "scale_up_transport": "nvlink", + "launcher": "single-slurm", + "backends": {"deepep-v2": [8, 16]}, + "fabric": {"nic": "ConnectX-8 2x400GbE", "switch": "NVIDIA Spectrum-X SN5600 (51.2T)"}, + "operator": { + "partition": "batch_1", + "account": "benchmark", + "qos": "batch_1_qos", + "squash_dir": "/data/home/sa-shared/sqsh" + }, + "network": { + "socket_ifname": "bond0", + "rdma_devices": "mlx5_0,mlx5_1,mlx5_2,mlx5_3,mlx5_4,mlx5_5,mlx5_8,mlx5_9,mlx5_10,mlx5_11,mlx5_16,mlx5_17,mlx5_20,mlx5_21,mlx5_22,mlx5_23", + "ib_gid_index": "3", + "rail_isolated": "1" + } + }, + "gb200": { + "arch": "sm100", + "product": "gb200", + "image": "lmsysorg/sglang:v0.5.11-cu130", + "image_platform": "linux/arm64", + "gpus_per_node": 4, + "scale_up_domain": 72, + "scale_up_transport": "mnnvl", + "launcher": "gb-nv", + "backends": {"deepep-v2": [8, 16]}, + "fabric": {"nic": "MNNVL (scale-out not used)", "switch": "NVLink NVL72"}, + "operator": { + "partition": "batch", + "account": "benchmark", + "storage_roots": [ + "/mnt/lustre01/users-public/sa-shared" + ] + } + }, + "gb300": { + "arch": "sm103", + "product": "gb300", + "image": "lmsysorg/sglang:v0.5.11-cu130", + "image_platform": "linux/arm64", + "gpus_per_node": 4, + "scale_up_domain": 72, + "scale_up_transport": "mnnvl", + "launcher": "gb-nv", + "backends": {"deepep-v2": [8, 16]}, + "fabric": {"nic": "MNNVL (scale-out not used)", "switch": "NVLink NVL72"}, + "operator": { + "partition": "batch_1", + "account": "benchmark", + "qos": "batch_1_qos", + "squash_dir": "/data/home/sa-shared/collectivex/containers", + "enroot_cache_path": "/data/home/sa-shared/collectivex/enroot-cache" + } + }, + "mi300x": { + "arch": "gfx942", + "product": "mi300x", + "image": "rocm/sgl-dev:sglang-0.5.14-rocm720-mi35x-mori-0701", + "image_platform": "linux/amd64", + "gpus_per_node": 8, + "scale_up_domain": 8, + "scale_up_transport": "xgmi", + "launcher": "mi-amds", + "backends": {"mori": [8]}, + "fabric": {"nic": "Pollara 400GbE", "switch": "Tomahawk5 (51.2T)"} + }, + "mi325x": { + "arch": "gfx942", + "product": "mi325x", + "image": "rocm/sgl-dev:sglang-0.5.14-rocm720-mi35x-mori-0701", + "image_platform": "linux/amd64", + "gpus_per_node": 8, + "scale_up_domain": 8, + "scale_up_transport": "xgmi", + "launcher": "mi-amds", + "backends": {"mori": [8]}, + "fabric": {"nic": "Pollara 400GbE", "switch": "Tomahawk5 (51.2T)"} + }, + "mi355x": { + "arch": "gfx950", + "product": "mi355x", + "image": "rocm/sgl-dev:sglang-0.5.14-rocm720-mi35x-mori-0701", + "image_platform": "linux/amd64", + "gpus_per_node": 8, + "scale_up_domain": 8, + "scale_up_transport": "xgmi", + "launcher": "mi-amds", + "backends": {"mori": [8, 16]}, + "fabric": {"nic": "Pollara 400GbE", "switch": "Arista 7060X6-64PE (Tomahawk5, 51.2T)"}, + "network": { + "socket_ifname": "eno0", + "rdma_devices": "rdma0,rdma1,rdma2,rdma3,rdma4,rdma5,rdma6,rdma7", + "ib_gid_index": 1 + }, + "operator": { + "partition": "compute", + "squash_dir": "/var/lib/squash" + } + } + } +} diff --git a/experimental/CollectiveX/configs/sweep.json b/experimental/CollectiveX/configs/sweep.json new file mode 100644 index 0000000000..a222466cb0 --- /dev/null +++ b/experimental/CollectiveX/configs/sweep.json @@ -0,0 +1,23 @@ +{ + "version": 1, + "suite": "ep-core", + "mode": "normal", + "routing": "uniform", + "ep_degrees": [8, 16], + "timing": { + "iters_per_trial": 8, + "trials_per_point": 256, + "warmup_iters_per_trial": 32 + }, + "workload": { + "name": "deepseek-v3", + "hidden": 7168, + "topk": 8, + "routed_experts": 256, + "seed": 67, + "token_ladders": { + "decode": [1, 2, 4, 8, 16, 32, 64, 128, 256, 512], + "prefill": [1024, 2048, 4096, 8192] + } + } +} diff --git a/experimental/CollectiveX/docs/methodology.md b/experimental/CollectiveX/docs/methodology.md new file mode 100644 index 0000000000..9619d29df1 --- /dev/null +++ b/experimental/CollectiveX/docs/methodology.md @@ -0,0 +1,229 @@ +# CollectiveX EP Benchmark Methodology + +CollectiveX schedules expert-parallel (EP) communication benchmarks, executes them on real +accelerator allocations, and uploads the neutral artifacts each run emits. It does **not** validate +those artifacts, promote, rank, recommend, select, hide, or decide what any consumer displays. The +frontend reads the neutral matrix, result, and summary artifacts and makes its own coverage +and display decisions. This document describes how a case is scheduled, measured, checked, and +recorded — not a publication or qualification contract. + +## Product Boundary + +CollectiveX is a communication microbenchmark for: + +- comparing EP libraries on one chip/topology; +- comparing EP latency and logical payload bandwidth across systems under the same workload; and +- surfacing unsupported, failed, invalid, and unstable cases rather than hiding them. + +It does not predict serving throughput without a separate correlation study. + +## Matrix + +The implemented workload is `deepseek-v3`: hidden 7168, top-k 8, 256 routed experts, packed +placement, and one pinned fixed resource profile per backend/topology. Dispatch and combine are +fixed BF16 on every backend; precision is not a swept dimension. Every case uses the normal +`layout-and-dispatch-v1` semantics. + +- `ep-core`: uniform routing over the workload's token ladders — for `deepseek-v3`, decode + T=1..512 powers of two and prefill T=1024..8192 powers of two. Ladders are model-specific and + live with the workload in `configs/sweep.json`. + +`sweep_matrix.py` materializes the requested SKUs, backends, EP sizes, and token ladders into a +matrix document, then extracts strict per-shard controls. `--only-sku`, `--exclude-skus`, and +`--ep-sizes` select a subset; a subset produces a smaller matrix, not a different contract. The +matrix is generated per dispatch; there is no frozen matrix digest or locked case count. + +| Systems | EP8 | EP16 | +|---|---|---| +| H100/H200/B200/B300 | 1x8 NVLink, scale-up | 2x8 NVLink + RDMA, scale-out | +| MI300X/MI325X/MI355X | 1x8 XGMI, scale-up | 2x8 XGMI + RDMA, scale-out | +| GB200/GB300 | 2x4 MNNVL, scale-up | 4x4 MNNVL, scale-up | + +Physical host count does not define scope. Both GB cells remain inside one 72-GPU MNNVL scale-up +domain. + +Unsupported combinations are explicitly classified in the matrix, not silently skipped coverage. DeepEP V2 is the +`ElasticBuffer` introduced by PR #605, pinned with upstream PR #630's minimal pure-scale-up fix and +the exact upstream PR #640 library matcher that excludes NCCL shared-memory mappings. Scale-up cases +request NCCL Device API LSA and fail closed unless the realized LSA team covers the full EP world. +x86 EP16 scale-out uses the hybrid path with GIN and requires two logical scale-out domains +represented by two physical RDMA ranks, with eight scale-up ranks per domain. GB EP16 remains MNNVL +scale-up and uses LSA. MoRI EP8 uses the direct IntraNode kernel on every CDNA SKU; EP16 uses pinned +InterNodeV1 over 2x8 XGMI + RDMA with 96 blocks, 64 RDMA blocks, 8 warps, one QP per PE, and external +input. No cell runs a low-latency-family kernel: throughput-oriented kernels are measured across the +full token ladder on both vendors. +Whether a given SKU/backend/EP cell is attempted is a capability fact; whether it succeeded is +decided only by the emitted artifact. + +## Workload Identity + +One deterministic workload is generated over the global token batch from the workload's seed in +`configs/sweep.json` (part of the workload identity, baked into every scheduled case) and sliced by +source rank; a keyed BLAKE2b counter over the (token, slot, attempt, stream) coordinates produces +byte-identical expert indices and gate weights on every runtime, and the harness proves the +realized routing trace identical across ranks before a case can succeed. + +Routing traffic distinguishes: + +- token-expert assignments, which determine expert compute load; and +- rank-deduplicated token payload copies, which determine EP activation traffic. + +Adapters may not generate routing or reinterpret one quantity as the other. + +## Measurement + +Normal mode uses `layout-and-dispatch-v1`: dispatch timing includes layout plus communication, and +combine returns activation payload through an unweighted rank-sum path. Expert-output staging is +outside isolated combine timing and inside the measured paired roundtrip. Each component declares +availability, origin, and sample count. A paired-only API reports null isolated components. +`isolated_sum` is derived. The artifact records the mode so a reader can keep distinct measurement +contracts separate. + +Every measured component uses one fixed timing profile, defined once in `configs/sweep.json` +and baked into every scheduled case: + +- 256 trials x 8 timed iterations = 2048 observations; +- 32 synchronized full dispatch-stage-combine warmups before each available measured component at + every trial/point; +- component measurement order rotates each trial (`trial_order`) so every timed component occupies + every position in the sequence, over a per-trial-rotated token ladder; and +- per-iteration maximum latency across ranks before nearest-rank p50/p90/p95/p99. + +Measured roundtrip p99 is the headline latency. Decode and prefill identify the serving regime +represented by one MoE-layer collective; they do not change the timed primitive at an otherwise +identical shape. Ascending through the ladder, each measured shape is conditioned with 8 untimed +full roundtrips — settling clocks, fabric, and buffer state — before it is correctness-checked; +all timing happens after every shape is warmed and checked. Conditioning rounds are never +measured or emitted. + +Logical payload bandwidth is: + +`logical_payload_bytes / measured_latency_seconds` + +Normal-mode payload bytes use rank-deduplicated token-rank BF16 activations (2 bytes per value, no +scale payload) and exclude expert metadata, padding, and backend buffer capacity. Algorithm bandwidth, bus bandwidth, wire +utilization, and physical-link utilization are not emitted without a defined primitive model or +transport counters. Logical bandwidth must never be labeled physical bandwidth. Payload and token +rates are named `rate_at_latency_percentile`: bytes or tokens divided by the matching latency +percentile. They are lower-tail service rates at p99 latency, not p99 percentiles of an inverted +rate distribution. + +## Correctness + +An implementation-independent oracle uses an expert-specific deterministic transform so wrong expert +routing cannot pass an identity roundtrip. For every rank and point it verifies: + +1. destination rank/expert, source token, multiplicity, gate weight, and receive counts; +2. dispatched payload and metadata before timing; +3. combined output before timing; +4. unchanged semantic inputs through all timed samples; and +5. dispatched payload/metadata and combined output again after timing. + +Normal-mode adapters use activation-only, unweighted rank-sum combine. The oracle builds each rank's +gate-weighted expert aggregate before combine and derives the expected combine from the values +actually communicated, reproducing the two-level reduction: each destination rank casts its FP32 +aggregate to the payload dtype (BF16) exactly as the adapter does; ranks sharing a scale-up domain +(NVLink/MNNVL) reduce in FP32, and each domain casts its aggregate to BF16 for the scale-out send +before those partials are summed. A group that fits in one scale-up domain (`ep_size <= +scale_up_domain` — every EP8 case and the MNNVL EP16 cases) has a single domain and no scale-out +rounding; a multi-node RoCE EP16 group carries one BF16 partial per node. Modelling that per-domain +cast is what lets the gate stay tight — max elementwise relative error (denominator clamped at 0.02) +below `8 * 2^-8`, the residual accumulation-order ambiguity — across scale-up and scale-out topologies +alike (omitting it left multi-node EP16 ~0.048 off, above the gate). It is a correctness gate, not an +estimate of transport error. Any failed rank or point makes the case ineligible in the result it writes. +Pre/post dispatch behavior is checked against canonical source-token metadata and expected output. +Native receive slots may be assigned nondeterministically, so physical receive order is not treated +as a correctness property. + +## Result Artifact + +One raw case document carries `record_type: "case-attempt"` and the single `version`, and contains: + +- `identity`: `case_id`, `attempt_ordinal`, `case_factors` (SKU and the scheduled case — backend, + EP size, mode, phase, suite, workload, and the topology coordinate), and `allocation_factors` + (run id, run attempt, source SHA); +- `workload`: `cross_rank_consistent`, whether the routing trace was proven identical across ranks; +- `measurement`: dispatch/combine dtype and semantics, `sampling`, and the per-point `rows`; +- `implementation`: backend name and kernel generation; +- `topology`: requested SKU/product, placement, nodes, scale-up domain, transport, and world size; +- `provenance`: the mounted image tag and source SHA; and +- `outcome`: `status` (`success` or `invalid`) and `reasons`. + +Each `rows` entry carries point latency, byte accounting, token rate, correctness, load, and fanout; +per-point statistics are summarized in place, not emitted as separate documents. Each dispatched +case writes exactly this one raw result document; unsupported or never-run cells produce no +synthetic record. + +## Identity + +Identifiers are readable factor strings: + +- `case_id`: `{sku}-{backend}-{workload}-{mode}-{phase}-ep{ep}-{routing}`, each factor + slug-normalized; and +- `attempt_ordinal`: a positive integer distinguishing repeat executions of one `case_id`. + +Backend source pins live in `runtime/common.sh` and are enforced by exact fetched-commit comparison; +the loaded DeepEP V2 build is checked for the required `ElasticBuffer` API. + +These IDs let a consumer group matched configurations and separate distinct ones. The backend does +not itself compute cohorts, controlled comparisons, sensitivity pairs, eligibility, or +recommendations — a reader decides which cases to surface and how to compare them. + +## Execution Isolation + +Every non-MNNVL scale-out case uses operator-pinned socket and RDMA selectors. The launcher rejects +missing or partial profiles, then probes every allocated node for the configured interface, active +HCA port, and configured GID before backend initialization. It never substitutes a default route, +inherited runner environment, or transport fallback. Scale-up and MNNVL cases clear the profile; +scale-out NVIDIA forces `NCCL_NET=IB`, while AMD leaves plugin selection to RCCL. Both use exact HCA +matching. Scale-out also pins `NCCL_IB_MERGE_NICS=0` so dual-port NIC fusion cannot disable NCCL GIN +— which the DeepEP V2 EP16 hybrid path requires — and a rail-isolated fabric (`rail_isolated`) adds +`NCCL_CROSS_NIC=0`. Selectors come from the tracked platform registry, optionally overlaid by an +operator config, and appear only in mode-0600 private logs. + +Repository staging uses a pre-existing, runner-owned, group/world non-writable shared base outside +the checkout and workflow workspace. The parent process resolves the exact execution child before +copying; backend preparation then runs from that tree on every allocated node. Cleanup waits for +confirmed allocation teardown and removes only that child. DeepEP V2 source is fetched before allocation at an +exact pinned revision, initializes its pinned `fmt` submodule, and applies the required local patch. + +H200, B200, and B300 may derive that private base beneath the validated operating-system account home +when it is compute-visible. H100 instead derives a sibling of its shared container directory, never a +child of image storage. +Canonical B300 execution ignores the legacy operator `stage_dir` field and always derives the base +from the validated shared account home. Its UID-mapped Actions shell may accept that exact base when +its owner matches the private parent owner; explicit stages and all other runners retain the strict +effective-UID ownership rule. An execution-ID suffix isolates parallel B300 workers. The current +NFS export may realize a newly created base as +UID 0; only that creation path is accepted, while a pre-existing root-owned base is rejected. +Canonical GB300 execution likewise ignores its legacy group-writable `stage_dir` and derives an +execution-specific private base beneath the validated compute-visible account home. + +## Image Pinning And Build Isolation + +Enroot imports configured container tags into a per-run-scoped squash keyed by the image tag and +image platform, so one run never reuses another run's imported filesystem. Image-provided DeepEP is +also checked against exact package versions and its expected API. Source-built DeepEP V2 uses +a separate mode-0700 cluster-local cache mounted only as `/cx-cache`. Its path binds CPU/GPU +architecture, image, and upstream commit. The cache is never an artifact; per-execution +source/results stages remain isolated and disposable, and runtime probes fail closed before reuse. The runner UID is +inside the trusted cluster boundary: this cache guards against stale or accidental mutation, not +hostile same-UID jobs. Only an unpublished partial build may be reset automatically; a cache that +fails integrity or runtime checks is left intact and rejected so a concurrent allocation cannot lose +files it is using. + +## Neutral Artifact Delivery + +There is no results server, attached store, or managed object store. Each shard runs one allocation, +emits per-case result JSON and a small mechanical summary, and uploads them as GitHub artifacts with +`always()` so a red or partial run still uploads. A case counts as successful on the benchmark's own +return code; there is no completeness or privacy validation step before upload, and failed or +unsupported cells produce no synthetic record. + +No step promotes a run, builds a dataset, or advances a channel; the artifacts are the output. Any +downstream display or comparison is the consumer's responsibility. + +## Legacy Data + +Historical numeric schemas 3-5 are outside this benchmark's artifacts. They remain historical +diagnostic evidence and are not produced or consumed by the current sweep. diff --git a/experimental/CollectiveX/launchers/launch_gb-nv.sh b/experimental/CollectiveX/launchers/launch_gb-nv.sh new file mode 100644 index 0000000000..a2d4b7b5ac --- /dev/null +++ b/experimental/CollectiveX/launchers/launch_gb-nv.sh @@ -0,0 +1,79 @@ +#!/usr/bin/env bash +# CollectiveX shared GB200/GB300 NVL72 (aarch64) launcher. +# shellcheck disable=SC2034 +# +# EP8/EP16 use one Slurm task per GPU across two or four trays in the same +# MNNVL scale-up domain. +# +# Flow: +# identity -> setup -> repository-stage -> backend-setup -> scheduler-allocation +# -> container-import -> container-launch -> artifact-collection +set -euo pipefail +HERE="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +COLLX_DIR="$(cd "$HERE/.." && pwd)"; REPO_ROOT="$(cd "$COLLX_DIR/../.." && pwd)" +# shellcheck source=../runtime/common.sh +source "$HERE/../runtime/common.sh" + +# ---- identity: resolve SKU, backend, platform ------------------------------- +PRODUCT="${COLLX_SHARD_SKU:-}" +case "$PRODUCT" in + gb200|gb300) ;; + *) collx_die "COLLX_SHARD_SKU must be gb200 or gb300" ;; +esac +RUNNER="$PRODUCT" +export COLLX_RUNNER="$RUNNER" COLLX_BENCH="${COLLX_BENCH:-deepep-v2}" +export COLLX_VENDOR=nvidia +# ---- setup: operator config, canonical env, topology, network profile ------- +collx_launcher_prologue "$RUNNER" +NODES="${COLLX_NODES:-2}"; GPN="${COLLX_GPUS_PER_NODE:-4}" +SCALE_UP_DOMAIN="${COLLX_SCALE_UP_DOMAIN:-72}" +NGPUS="${COLLX_NGPUS:-$((NODES * GPN))}" +if [ "$PRODUCT" = gb200 ]; then default_time=30; else default_time=90; fi +TIME_MIN="${COLLX_TIME:-$default_time}" +IMAGE="$COLLX_IMAGE" +TS="$(date -u +%Y-%m-%dT%H-%M-%SZ)" +export COLLX_TRANSPORT=mnnvl +export COLLX_NODES="$NODES" COLLX_GPUS_PER_NODE="$GPN" COLLX_SCALE_UP_DOMAIN="$SCALE_UP_DOMAIN" +export COLLX_NGPUS="$NGPUS" +case "$COLLX_BENCH" in + deepep-v2) ;; + *) collx_die "unsupported $PRODUCT EP backend: $COLLX_BENCH" ;; +esac +collx_require_vars COLLX_IMAGE COLLX_IMAGE_PLATFORM COLLX_PARTITION COLLX_ACCOUNT COLLX_SQUASH_DIR COLLX_STAGE_DIR +[ "$PRODUCT" != gb300 ] || collx_require_vars COLLX_ENROOT_CACHE_PATH +PARTITION="$COLLX_PARTITION"; ACCOUNT="$COLLX_ACCOUNT"; SQUASH_DIR="$COLLX_SQUASH_DIR" +[ -z "${COLLX_ENROOT_CACHE_PATH:-}" ] || export ENROOT_CACHE_PATH="$COLLX_ENROOT_CACHE_PATH" +export NCCL_CUMEM_ENABLE=1 NCCL_MNNVL_ENABLE=1 MC_FORCE_MNNVL=1 +collx_apply_network_profile "$NODES" "$COLLX_TRANSPORT" + +collx_log "$PRODUCT nodes=$NODES x ${GPN}gpu world=$NGPUS bench=$COLLX_BENCH" +collx_select_image "$IMAGE" + +# ---- repository-stage: compute-visible copy of the checkout ----------------- +MOUNT_SRC="$(collx_stage_path "$REPO_ROOT" "$COLLX_STAGE_DIR")" +collx_stage_repo "$REPO_ROOT" "$MOUNT_SRC" +CONTAINER_MOUNTS="$MOUNT_SRC:/ix" +# ---- backend-setup: pinned DeepEP source + isolated build cache ------------- +# The backend case above admits only deepep-v2, so its staging is unconditional. +collx_prepare_deepep_source "$MOUNT_SRC" \ + || collx_die "cannot stage the pinned backend source" +export COLLX_BACKEND_SOURCE_ROOT=/ix/experimental/CollectiveX/.collx_sources +collx_prepare_backend_cache "$COLLX_SQUASH_DIR" \ + || collx_die "cannot prepare the isolated backend cache" +CONTAINER_MOUNTS="$CONTAINER_MOUNTS,$COLLX_PREPARED_BACKEND_CACHE:/cx-cache" +export COLLX_BACKEND_CACHE_ROOT=/cx-cache + +# ---- scheduler-allocation: salloc the trays --------------------------------- +command -v salloc >/dev/null || collx_die "salloc not found" +collx_salloc_jobid --partition="$PARTITION" --account="$ACCOUNT" --nodes="$NODES" \ + --gres=gpu:"$GPN" --ntasks-per-node="$GPN" --exclusive --mem=0 --cpus-per-task=35 \ + --time="$TIME_MIN" +[ -n "$JOB_ID" ] || collx_die "no JOB_ID from salloc" + +# ---- container-import: squash file resolved on the allocation --------------- +SQUASH_FILE="$(collx_ensure_squash_on_job "$JOB_ID" "$SQUASH_DIR" "$IMAGE")" + +# ---- container-launch -> artifact-collection (shared tail) ------------------ +COLLX_DISTRIBUTED_CONTAINER_ARGS=(--container-writable --container-remap-root) +collx_execute_and_collect "$MOUNT_SRC" "$REPO_ROOT" +exit "$FINAL_RC" diff --git a/experimental/CollectiveX/launchers/launch_mi-amds.sh b/experimental/CollectiveX/launchers/launch_mi-amds.sh new file mode 100644 index 0000000000..70b3263db2 --- /dev/null +++ b/experimental/CollectiveX/launchers/launch_mi-amds.sh @@ -0,0 +1,127 @@ +#!/usr/bin/env bash +# CollectiveX shared AMD Slurm launcher (one or two nodes). +# shellcheck disable=SC2034 +# +# Flow (container import runs inside the allocation retry loop): +# identity -> setup -> repository-stage -> scheduler-allocation + container-import +# -> container-launch -> artifact-collection +set -euo pipefail + +HERE="$(cd -P -- "$(dirname "${BASH_SOURCE[0]}")" && pwd -P)" +COLLX_DIR="$(cd "$HERE/.." && pwd)" +REPO_ROOT="$(cd "$COLLX_DIR/../.." && pwd)" +# shellcheck source=../runtime/common.sh +source "$HERE/../runtime/common.sh" + +# ---- identity: resolve SKU, backend, platform ------------------------------- +RUNNER="${COLLX_SHARD_SKU:-}" +case "$RUNNER" in + mi300x|mi325x) CPUS_PER_NODE=256; DEVICE_MOUNTS=",/dev/kfd:/dev/kfd,/dev/dri:/dev/dri" ;; + mi355x) CPUS_PER_NODE=128; DEVICE_MOUNTS="" ;; + *) collx_die "COLLX_SHARD_SKU is not a registered AMD SKU" ;; +esac +export COLLX_RUNNER="$RUNNER" COLLX_BENCH="${COLLX_BENCH:-mori}" +export COLLX_VENDOR=amd +# ---- setup: operator config, canonical env, topology, network profile ------- +collx_launcher_prologue "$RUNNER" + +NODES="${COLLX_NODES:-1}"; GPN="${COLLX_GPUS_PER_NODE:-8}" +SCALE_UP_DOMAIN="${COLLX_SCALE_UP_DOMAIN:-8}" +NGPUS="${COLLX_NGPUS:-$((NODES * GPN))}" +TIME_MIN="${COLLX_TIME:-60}" +EXCLUDE_NODES="${COLLX_EXCLUDE_NODES:-}" +NODELIST="${COLLX_NODELIST:-}" +MOUNT_DIR=/ix +TS="$(date -u +%Y-%m-%dT%H-%M-%SZ)" +case "$COLLX_BENCH" in + mori) ;; + *) collx_die "unsupported AMD EP backend: $COLLX_BENCH" ;; +esac + +export MORI_DISABLE_AUTO_XGMI="${MORI_DISABLE_AUTO_XGMI:-0}" +export MORI_ENABLE_SDMA="${MORI_ENABLE_SDMA:-1}" +export MORI_APP_LOG_LEVEL="${MORI_APP_LOG_LEVEL:-info}" +export MORI_SHMEM_LOG_LEVEL="${MORI_SHMEM_LOG_LEVEL:-info}" +export MORI_IO_LOG_LEVEL="${MORI_IO_LOG_LEVEL:-info}" +IMAGE="$COLLX_IMAGE" +export COLLX_NGPUS="$NGPUS" COLLX_NODES="$NODES" +export COLLX_GPUS_PER_NODE="$GPN" COLLX_SCALE_UP_DOMAIN="$SCALE_UP_DOMAIN" +if [ "$NODES" -gt 1 ]; then + export COLLX_TRANSPORT=xgmi-rdma +else + export COLLX_TRANSPORT=xgmi +fi +export COLLX_RUN_TIMEOUT="${COLLX_RUN_TIMEOUT:-1800}" +collx_apply_network_profile "$NODES" "$COLLX_TRANSPORT" +collx_require_vars COLLX_IMAGE COLLX_IMAGE_PLATFORM COLLX_PARTITION COLLX_SQUASH_DIR COLLX_STAGE_DIR +PARTITION="$COLLX_PARTITION"; SQUASH_DIR="$COLLX_SQUASH_DIR" + +collx_log "runner=$RUNNER nodes=$NODES x ${GPN}gpu world=$NGPUS bench=$COLLX_BENCH" + +# ---- repository-stage: compute-visible copy of the checkout ----------------- +MOUNT_SRC="$(collx_stage_path "$REPO_ROOT" "$COLLX_STAGE_DIR")" +collx_stage_repo "$REPO_ROOT" "$MOUNT_SRC" +collx_select_image "$IMAGE" + +# ---- scheduler-allocation + container-import: retry until nodes validate ---- +# Each attempt must pass the network profile AND import the squash; a rejected +# allocation is cancelled and its nodes excluded from the next attempt. +command -v salloc >/dev/null || collx_die "salloc not found on this runner" + +allocation=(--partition="$PARTITION" --nodes="$NODES" --gres=gpu:"$GPN" + --time="$TIME_MIN" --ntasks-per-node="$GPN" + --cpus-per-task="$((CPUS_PER_NODE / GPN))") +if [ "$RUNNER" = mi355x ]; then + allocation+=(--exclusive) +fi +excluded_nodes="$EXCLUDE_NODES" +for allocation_attempt in 1 2 3; do + attempt_allocation=("${allocation[@]}") + if [ -n "$NODELIST" ]; then + collx_log "using configured node pin" + attempt_allocation+=(--nodelist="$NODELIST") + elif [ -n "$excluded_nodes" ]; then + attempt_allocation+=(--exclude="$excluded_nodes") + fi + export COLLX_SALLOC_ATTEMPT="$allocation_attempt" + export COLLX_NETWORK_VALIDATION_ATTEMPT="$allocation_attempt" + collx_salloc_jobid "${attempt_allocation[@]}" + [ -n "$JOB_ID" ] || collx_die "could not resolve allocated JOB_ID from salloc" + reject_reason="" + if ! collx_validate_network_profile_on_job "$JOB_ID" "$NODES" "$COLLX_TRANSPORT"; then + # A node whose RoCE devices do not match the pinned selector (e.g. an + # outlier still using default rocepXXXs0 names instead of the rdmaN udev + # names the rest of the fleet exposes) must be rejected and retried + # elsewhere, not treated as a hard failure. + reject_reason=network + else + if SQUASH_FILE="$(collx_ensure_squash_on_job \ + "$JOB_ID" "$SQUASH_DIR" "$IMAGE" "${COLLX_LOCK_DIR:-}")"; then + break + fi + reject_reason=container-import + fi + if [ -n "$NODELIST" ] || [ "$allocation_attempt" = 3 ]; then + if [ "$reject_reason" = network ]; then + collx_log_tail "${COLLX_NETWORK_PROFILE_LOG:-}" + collx_die "allocated nodes failed the network profile" + fi + collx_die "allocated nodes failed container import" + fi + rejected_nodes="$(collx_allocation_nodes_csv "$JOB_ID")" \ + || collx_die "cannot identify nodes from a rejected allocation" + collx_log "allocated nodes failed $reject_reason validation; retrying elsewhere" + collx_cleanup_allocation || collx_die "cannot release a rejected allocation" + JOB_ID="" + [ -z "$excluded_nodes" ] || excluded_nodes+=, + excluded_nodes+="$rejected_nodes" +done +unset COLLX_SALLOC_ATTEMPT COLLX_NETWORK_VALIDATION_ATTEMPT +CONTAINER_MOUNTS="$MOUNT_SRC:$MOUNT_DIR$DEVICE_MOUNTS" + +# ---- container-launch -> artifact-collection (shared tail) ------------------ +COLLX_DISTRIBUTED_CONTAINER_ARGS=(--container-writable --container-remap-root) +collx_execute_and_collect "$MOUNT_SRC" "$REPO_ROOT" +rm -f "$MOUNT_SRC"/experimental/CollectiveX/gpucore.* 2>/dev/null || true +collx_log "done - result artifacts collected" +exit "$FINAL_RC" diff --git a/experimental/CollectiveX/launchers/launch_single-slurm.sh b/experimental/CollectiveX/launchers/launch_single-slurm.sh new file mode 100644 index 0000000000..c776eef250 --- /dev/null +++ b/experimental/CollectiveX/launchers/launch_single-slurm.sh @@ -0,0 +1,147 @@ +#!/usr/bin/env bash +# CollectiveX shared standard NVIDIA Slurm launcher (one or two nodes). +# shellcheck disable=SC2034 +# +# Flow: +# identity -> setup -> repository-stage -> backend-setup -> scheduler-allocation +# -> container-import -> container-launch -> artifact-collection +set -euo pipefail + +HERE="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +COLLX_DIR="$(cd "$HERE/.." && pwd)" +REPO_ROOT="$(cd "$COLLX_DIR/../.." && pwd)" +# shellcheck source=../runtime/common.sh +source "$HERE/../runtime/common.sh" + +# ---- identity: resolve SKU, backend, platform ------------------------------- +RUNNER="${COLLX_SHARD_SKU:-}" +ALLOC_EXTRA=(); SRUN_EXTRA=(); LOCAL_IMPORT=0 +case "$RUNNER" in + h100-dgxc) PRODUCT=h100; DEFAULT_TIME=45; REQUIRE_ACCOUNT=1 ;; + h200-dgxc) + PRODUCT=h200; DEFAULT_TIME=45; REQUIRE_ACCOUNT=0 + SRUN_EXTRA=(--container-remap-root) + ;; + b200-dgxc) + PRODUCT=b200; DEFAULT_TIME=30; REQUIRE_ACCOUNT=1 + ALLOC_EXTRA=(--mem=0) + ;; + b300) + PRODUCT=b300; DEFAULT_TIME=45; REQUIRE_ACCOUNT=1 + # Do not restore ALLOC_EXTRA=(-N 1 --mem=0); it blocks two-node B300 jobs. + ALLOC_EXTRA=(--mem=0) + SRUN_EXTRA=(--mpi=none --container-remap-root) + LOCAL_IMPORT=1 + ;; + *) collx_die "COLLX_SHARD_SKU is not a registered NVIDIA SKU" ;; +esac +export COLLX_RUNNER="$RUNNER" COLLX_BENCH="${COLLX_BENCH:-deepep-v2}" +export COLLX_VENDOR=nvidia +# ---- setup: operator config, canonical env, topology, network profile ------- +collx_launcher_prologue "$RUNNER" + +NODES="${COLLX_NODES:-1}"; GPN="${COLLX_GPUS_PER_NODE:-8}" +SCALE_UP_DOMAIN="${COLLX_SCALE_UP_DOMAIN:-8}" +NGPUS="${COLLX_NGPUS:-$((NODES * GPN))}" +TIME_MIN="${COLLX_TIME:-$DEFAULT_TIME}" +IMAGE="$COLLX_IMAGE" +TS="$(date -u +%Y-%m-%dT%H-%M-%SZ)" +case "$COLLX_BENCH" in + deepep-v2) ;; + *) collx_die "unsupported $RUNNER EP backend: $COLLX_BENCH" ;; +esac + +export COLLX_NGPUS="$NGPUS" COLLX_NODES="$NODES" +export COLLX_GPUS_PER_NODE="$GPN" COLLX_SCALE_UP_DOMAIN="$SCALE_UP_DOMAIN" +if [ "$NODES" -gt 1 ]; then + export COLLX_TRANSPORT=nvlink-rdma +else + export COLLX_TRANSPORT=nvlink +fi +export NCCL_CUMEM_ENABLE=1 +collx_apply_network_profile "$NODES" "$COLLX_TRANSPORT" +collx_require_vars COLLX_IMAGE COLLX_IMAGE_PLATFORM COLLX_PARTITION COLLX_SQUASH_DIR +[ "$REQUIRE_ACCOUNT" = 0 ] || collx_require_vars COLLX_ACCOUNT +# b300 /home and /scratch are node-local; h100-dgxc /home is login-local (absent on +# compute nodes) — on both, the implicit passwd-home stage base is compute-invisible, +# so the operator config must pin an explicit compute-visible stage_dir. +case "$RUNNER" in + b300|h100-dgxc) collx_require_vars COLLX_STAGE_DIR ;; +esac + +collx_log "runner=$RUNNER nodes=$NODES x ${GPN}gpu world=$NGPUS bench=$COLLX_BENCH" +collx_select_image "$IMAGE" + +# ---- repository-stage: compute-visible copy of the checkout ----------------- +MOUNT_SRC="$(collx_stage_path "$REPO_ROOT" "${COLLX_STAGE_DIR:-}")" +collx_stage_repo "$REPO_ROOT" "$MOUNT_SRC" +CONTAINER_MOUNTS="$MOUNT_SRC:/ix" +# ---- backend-setup: pinned DeepEP source + isolated build cache ------------- +# The backend case above admits only deepep-v2, so its staging is unconditional. +collx_prepare_deepep_source "$MOUNT_SRC" \ + || collx_die "cannot stage the pinned backend source" +export COLLX_BACKEND_SOURCE_ROOT=/ix/experimental/CollectiveX/.collx_sources +collx_prepare_backend_cache "$COLLX_SQUASH_DIR" \ + || collx_die "cannot prepare the isolated backend cache" +CONTAINER_MOUNTS="$CONTAINER_MOUNTS,$COLLX_PREPARED_BACKEND_CACHE:/cx-cache" +export COLLX_BACKEND_CACHE_ROOT=/cx-cache + +# ---- scheduler-allocation: salloc, retry until nodes validate --------------- +# Each attempt must pass the network profile (and accelerator-context on b300); +# a rejected allocation is cancelled and its nodes excluded from the next attempt. +command -v salloc >/dev/null || collx_die "salloc not found on this runner" +allocation=(--partition="$COLLX_PARTITION" --nodes="$NODES" --gres=gpu:"$GPN" + --ntasks-per-node="$GPN" --exclusive --time="$TIME_MIN" "${ALLOC_EXTRA[@]}") +[ -z "${COLLX_ACCOUNT:-}" ] || allocation+=(--account="$COLLX_ACCOUNT") +[ -z "${COLLX_QOS:-}" ] || allocation+=(--qos="$COLLX_QOS") +[ -z "${COLLX_NODELIST:-}" ] || allocation+=(--nodelist="$COLLX_NODELIST") +excluded_nodes="${COLLX_EXCLUDE_NODES:-}" +for allocation_attempt in 1 2 3; do + validation_failure="" + attempt_allocation=("${allocation[@]}") + [ -z "$excluded_nodes" ] || attempt_allocation+=(--exclude="$excluded_nodes") + export COLLX_SALLOC_ATTEMPT="$allocation_attempt" + export COLLX_NETWORK_VALIDATION_ATTEMPT="$allocation_attempt" + collx_salloc_jobid "${attempt_allocation[@]}" + [ -n "$JOB_ID" ] || collx_die "could not resolve allocated JOB_ID from salloc" + if ! collx_validate_network_profile_on_job "$JOB_ID" "$NODES" "$COLLX_TRANSPORT"; then + validation_failure=network + elif [ "$RUNNER" = b300 ] \ + && ! collx_validate_cuda_context_on_job "$JOB_ID" "$NODES" "$GPN"; then + validation_failure=cuda-context + else + break + fi + retryable=0 + [ "$RUNNER:$validation_failure" != h100-dgxc:network ] || retryable=1 + [ "$RUNNER:$validation_failure" != b300:cuda-context ] || retryable=1 + if [ "$retryable" = 0 ] || [ "$allocation_attempt" = 3 ]; then + if [ "$validation_failure" = network ]; then + collx_log_tail "${COLLX_NETWORK_PROFILE_LOG:-}" + collx_die "allocated nodes failed the network profile" + fi + collx_log_tail "$COLLX_CUDA_CONTEXT_LOG" + collx_die "allocated nodes failed accelerator context validation" + fi + rejected_nodes="$(collx_allocation_nodes_csv "$JOB_ID")" \ + || collx_die "cannot identify nodes from a rejected allocation" + collx_log "allocated nodes failed $validation_failure validation; retrying elsewhere" + collx_cleanup_allocation || collx_die "cannot release a rejected allocation" + JOB_ID="" + [ -z "$excluded_nodes" ] || excluded_nodes+=, + excluded_nodes+="$rejected_nodes" +done +unset COLLX_SALLOC_ATTEMPT COLLX_NETWORK_VALIDATION_ATTEMPT + +# ---- container-import: squash file (login-local on b300, else on the job) ---- +if [ "$LOCAL_IMPORT" = 1 ]; then + SQUASH_FILE="$(COLLX_ENROOT_LOCAL_IMPORT=1 collx_ensure_squash "$COLLX_SQUASH_DIR" "$IMAGE")" +else + SQUASH_FILE="$(collx_ensure_squash_on_job "$JOB_ID" "$COLLX_SQUASH_DIR" "$IMAGE")" +fi + +# ---- container-launch -> artifact-collection (shared tail) ------------------ +COLLX_DISTRIBUTED_CONTAINER_ARGS=(--container-writable "${SRUN_EXTRA[@]}") +collx_execute_and_collect "$MOUNT_SRC" "$REPO_ROOT" +collx_log "done - result artifacts collected" +exit "$FINAL_RC" diff --git a/experimental/CollectiveX/runtime/common.sh b/experimental/CollectiveX/runtime/common.sh new file mode 100644 index 0000000000..bb0bc38387 --- /dev/null +++ b/experimental/CollectiveX/runtime/common.sh @@ -0,0 +1,898 @@ +# shellcheck shell=bash +# CollectiveX — shared launcher helpers (sourced, not executed). +# +# Cluster-generic scaffolding only (Slurm/container/build/staging); no +# model-serving. Logging goes to stderr so functions can `echo` a single +# result on stdout. + +unset COLLECTIVEX_OPERATOR_CONFIG_LOADED +COLLX_RUNTIME_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd -P)" + +collx_log() { printf '[collectivex] %s\n' "$*" >&2; } +collx_die() { printf '[collectivex] FATAL: %s\n' "$*" >&2; exit 1; } + +COLLX_DEEPEP_V2_REPO="https://github.com/deepseek-ai/DeepEP" +COLLX_DEEPEP_V2_COMMIT="fa8a9b16898204afd347c663b89e65ef87dc6ce6" + +# Print bounded command output without maintaining a parallel failure taxonomy. +collx_log_tail() { + local log_path="$1" + if [ -s "$log_path" ]; then + collx_log "--- command log tail ---" + tail -n 100 -- "$log_path" >&2 || true + collx_log "--- end command log tail ---" + fi +} + +# Shared launcher skeleton: the identity-stage boilerplate every launcher runs +# before any SKU-specific work. +collx_launcher_prologue() { + JOB_ID="" + collx_install_launcher_fail_safe + [ -n "${COLLX_SHARD_FILE:-}" ] || collx_die "COLLX_SHARD_FILE is required" + collx_load_operator_config + collx_prepare_stage_dir "$1" +} + +# Shared launcher tail: run the shard, collect artifacts, and fold both return +# codes into FINAL_RC (run failures win; collection failures surface otherwise). +collx_execute_and_collect() { + local mount_src="$1" repo_root="$2" run_rc=0 collect_rc=0 + collx_run_shard || run_rc=$? + collx_collect_results "$mount_src" "$repo_root" || collect_rc=$? + FINAL_RC="$run_rc" + [ "$FINAL_RC" != 0 ] || FINAL_RC="$collect_rc" +} + +collx_job_root_is_safe() { + local root="$1" + if [[ "$root" =~ ^/tmp/inferencex-collectivex-[0-9]+-[0-9]+-[A-Za-z0-9._-]+$ ]]; then + : + elif [[ "$root" =~ ^/tmp/inferencex-collectivex-parent-([0-9]+)-([0-9]+)-([A-Za-z0-9._-]+)/inferencex-collectivex-([0-9]+)-([0-9]+)-([A-Za-z0-9._-]+)$ ]]; then + [ "${BASH_REMATCH[1]}" = "${BASH_REMATCH[4]}" ] \ + && [ "${BASH_REMATCH[2]}" = "${BASH_REMATCH[5]}" ] \ + && [ "${BASH_REMATCH[3]}" = "${BASH_REMATCH[6]}" ] || return 1 + else + return 1 + fi + [ -d "$root" ] && [ ! -L "$root" ] \ + && [ "$(stat -c '%u:%a' "$root" 2>/dev/null)" = "$(id -u):700" ] +} + +# Load the selected SKU's public platform settings plus any allowlisted local +# operator overrides; JSON values are never sourced or evaluated as shell. +collx_load_operator_config() { + [ -n "${COLLECTIVEX_OPERATOR_CONFIG_LOADED:-}" ] \ + && [ "$COLLECTIVEX_OPERATOR_CONFIG_LOADED" = "$$" ] && return 0 + local config_path parsed_path key value + unset COLLX_IMAGE COLLX_IMAGE_PLATFORM + unset COLLX_PARTITION COLLX_ACCOUNT COLLX_QOS COLLX_SQUASH_DIR COLLX_STAGE_DIR COLLX_ENROOT_CACHE_PATH + unset ENROOT_CACHE_PATH + unset COLLX_EXCLUDE_NODES COLLX_NODELIST COLLX_LOCK_DIR COLLX_MASTER_PORT + unset COLLX_SOCKET_IFNAME COLLX_RDMA_DEVICES COLLX_IB_GID_INDEX COLLX_RDMA_SERVICE_LEVEL + unset COLLX_RDMA_TRAFFIC_CLASS COLLX_RAIL_ISOLATED + unset MASTER_ADDR MASTER_PORT RANK WORLD_SIZE LOCAL_RANK LOCAL_WORLD_SIZE + config_path="${COLLECTIVEX_OPERATOR_CONFIG:-${XDG_CONFIG_HOME:-${HOME}/.config}/inferencex/collectivex.json}" + if [ ! -e "$config_path" ]; then + # No operator document: a host-utility step (no SKU) is a no-op; a known SKU + # emits the tracked platform_config.json operator baseline ("-" sentinel). + # An optional local file at COLLECTIVEX_OPERATOR_CONFIG/XDG still overlays it. + if [ -z "${COLLX_RUNNER:-${COLLX_SHARD_SKU:-}}" ]; then + COLLECTIVEX_OPERATOR_CONFIG_LOADED="$$" + return 0 + fi + config_path="-" + fi + umask 077 + parsed_path="$(mktemp /tmp/inferencex-collectivex-parsed.XXXXXX)" \ + || collx_die "cannot parse runner configuration" + if ! python3 "$COLLX_RUNTIME_DIR/config.py" operator-config "$config_path" \ + "${COLLX_RUNNER:-${COLLX_SHARD_SKU:-}}" \ + > "$parsed_path" + then + rm -f -- "$parsed_path" + unset COLLECTIVEX_OPERATOR_CONFIG + collx_die "runner-local configuration failed" + fi + while IFS= read -r -d '' key && IFS= read -r -d '' value; do + printf -v "$key" '%s' "$value" + export "${key?}" + done < "$parsed_path" + rm -f -- "$parsed_path" + unset COLLECTIVEX_OPERATOR_CONFIG + COLLECTIVEX_OPERATOR_CONFIG_LOADED="$$" +} + +# Per-step log files: several callers parse these for markers (salloc grant, +# stage copy-error, per-node network selectors), so they are a data channel, +# not just failure display. Logs persist after the run for postmortem. +collx_private_log_path() { + local path="${COLLX_JOB_ROOT:-/tmp/inferencex-collectivex-$(id -u)}/logs/$1.log" + mkdir -p "${path%/*}" || collx_die "cannot create log directory" + : > "$path" || collx_die "cannot create runtime log" + printf '%s' "$path" +} + +# Host-side utility steps need only the basic login paths. They never receive +# the complete Actions or runner environment. +collx_host_exports() { + printf '%s' 'HOME,PATH,USER,XDG_CACHE_HOME,ENROOT_CACHE_PATH' +} + +collx_require_vars() { + local name + local -a missing=() + for name in "$@"; do + [ -n "${!name:-}" ] || missing+=("$name") + done + [ "${#missing[@]}" -eq 0 ] || collx_die \ + "missing platform or runner configuration: ${missing[*]}" +} + +collx_export_gid_index_for_link_layer() { + local link_layer="$1" + unset NVSHMEM_IB_GID_INDEX NCCL_IB_GID_INDEX + [ -n "${COLLX_IB_GID_INDEX:-}" ] || return 0 + case "$link_layer" in + roce) + export NVSHMEM_IB_GID_INDEX="$COLLX_IB_GID_INDEX" + export NCCL_IB_GID_INDEX="$COLLX_IB_GID_INDEX" + ;; + infiniband) ;; + *) collx_die "unsupported RDMA link layer" ;; + esac +} + +# Convert private, runner-local network selectors into the public library +# variables needed inside the container. Values are interface/HCA identifiers, +# never addresses; the rendezvous hostname is derived from the allocation. +collx_apply_network_profile() { + local nodes="$1" transport="$2" + local selector rdma_name rdma_names="" ep_nic="" + local -a selectors + [[ "$nodes" =~ ^[1-9][0-9]*$ ]] || collx_die "invalid network placement" + unset NCCL_NET NCCL_SOCKET_IFNAME GLOO_SOCKET_IFNAME NCCL_IB_HCA + unset NCCL_IB_GID_INDEX NCCL_IB_SL NCCL_IB_MERGE_NICS NCCL_CROSS_NIC + unset NVSHMEM_ENABLE_NIC_PE_MAPPING + unset NVSHMEM_HCA_LIST NVSHMEM_IB_GID_INDEX NVSHMEM_IB_SL + unset NVSHMEM_IB_ENABLE_IBGDA NVSHMEM_IBGDA_NIC_HANDLER + unset EP_NIC_NAME EP_OVERRIDE_RDMA_SL + unset MORI_RDMA_DEVICES + unset MORI_RDMA_TC MORI_IO_TC MORI_RDMA_SL MORI_IO_SL + # Single-node and MNNVL runs need only the scrub above; everything past this + # point is the scale-out path, so no per-branch scale-out guards remain. + { [ "$nodes" -gt 1 ] && [ "$transport" != mnnvl ]; } || return 0 + [ -n "${COLLX_RDMA_DEVICES:-}" ] \ + || collx_die "RDMA execution requires a private device selector" + if [ -n "${COLLX_SOCKET_IFNAME:-}" ]; then + [[ "$COLLX_SOCKET_IFNAME" =~ ^[A-Za-z][A-Za-z0-9_.-]{0,31}$ ]] \ + || collx_die "invalid private socket interface selector" + export NCCL_SOCKET_IFNAME="$COLLX_SOCKET_IFNAME" GLOO_SOCKET_IFNAME="$COLLX_SOCKET_IFNAME" + fi + [[ "$COLLX_RDMA_DEVICES" =~ ^[A-Za-z][A-Za-z0-9_.-]{0,31}(:[1-9][0-9]*)?(,[A-Za-z][A-Za-z0-9_.-]{0,31}(:[1-9][0-9]*)?)*$ ]] \ + || collx_die "invalid private RDMA device selector" + IFS=, read -r -a selectors <<< "$COLLX_RDMA_DEVICES" + for selector in "${selectors[@]}"; do + rdma_name="${selector%%:*}" + rdma_names="${rdma_names}${rdma_names:+,}${rdma_name}" + [ -n "$ep_nic" ] || ep_nic="$rdma_name" + done + export NVSHMEM_HCA_LIST="$COLLX_RDMA_DEVICES" + export NVSHMEM_ENABLE_NIC_PE_MAPPING=1 + # RCCL selects its own net plugin; NCCL_NET=IB breaks AMD SKUs. + if [ "${COLLX_VENDOR:-nvidia}" = amd ]; then + unset NCCL_NET + else + export NCCL_NET=IB + fi + export NCCL_IB_HCA="=$COLLX_RDMA_DEVICES" + export MORI_RDMA_DEVICES="$rdma_names" EP_NIC_NAME="$ep_nic" + # The selector enumerates individual ports. NCCL's default dual-port fusion + # would collapse each card into one "fused" device, and any fused device + # disables NCCL GIN (init.cc nicFused gate) — the deep_ep EP16 hybrid path + # then asserts railedGinType == NCCL_GIN_TYPE_NONE. Single-port selectors are + # unaffected, so pin unmerged operation for every scale-out run. + export NCCL_IB_MERGE_NICS=0 + if [ -n "${COLLX_RAIL_ISOLATED:-}" ]; then + [[ "$COLLX_RAIL_ISOLATED" =~ ^[01]$ ]] \ + || collx_die "invalid private rail isolation flag" + # Rail-isolated multi-plane fabrics (per-port rail subnets, no cross-rail + # routing): cross-NIC pairs black-hole at QP RTR, and NCCL's RAIL GIN + # connection type is the one the fabric supports. + [ "$COLLX_RAIL_ISOLATED" != 1 ] || export NCCL_CROSS_NIC=0 + fi + if [ -n "${COLLX_IB_GID_INDEX:-}" ]; then + [[ "$COLLX_IB_GID_INDEX" =~ ^[0-9]+$ ]] && [ "$COLLX_IB_GID_INDEX" -le 255 ] \ + || collx_die "invalid private IB GID index" + fi + if [ -n "${COLLX_RDMA_SERVICE_LEVEL:-}" ]; then + [[ "$COLLX_RDMA_SERVICE_LEVEL" =~ ^[0-9]+$ ]] && [ "$COLLX_RDMA_SERVICE_LEVEL" -le 15 ] \ + || collx_die "invalid private RDMA service level" + export NVSHMEM_IB_SL="$COLLX_RDMA_SERVICE_LEVEL" + export NCCL_IB_SL="$COLLX_RDMA_SERVICE_LEVEL" + export EP_OVERRIDE_RDMA_SL="$COLLX_RDMA_SERVICE_LEVEL" + export MORI_RDMA_SL="$COLLX_RDMA_SERVICE_LEVEL" MORI_IO_SL="$COLLX_RDMA_SERVICE_LEVEL" + fi + if [ -n "${COLLX_RDMA_TRAFFIC_CLASS:-}" ]; then + [[ "$COLLX_RDMA_TRAFFIC_CLASS" =~ ^[0-9]+$ ]] && [ "$COLLX_RDMA_TRAFFIC_CLASS" -le 255 ] \ + || collx_die "invalid private RDMA traffic class" + export MORI_RDMA_TC="$COLLX_RDMA_TRAFFIC_CLASS" MORI_IO_TC="$COLLX_RDMA_TRAFFIC_CLASS" + fi + local nic_handler=gpu + export NVSHMEM_IB_ENABLE_IBGDA=1 NVSHMEM_IBGDA_NIC_HANDLER="$nic_handler" + if [ -n "${COLLX_RDMA_LINK_LAYER:-}" ]; then + case "$COLLX_RDMA_LINK_LAYER" in + roce|infiniband) ;; + *) collx_die "invalid validated RDMA link layer" ;; + esac + collx_export_gid_index_for_link_layer "$COLLX_RDMA_LINK_LAYER" + fi +} + +# Slurm may remove NCCL's leading exact-match marker while propagating an +# inherited environment. Reconstruct it from the validated private selector at +# the container boundary instead of accepting a prefix-matched HCA list. +collx_restore_exact_hca_selector() { + if [ "${COLLX_NODES:-1}" -le 1 ] || [ "${COLLX_TRANSPORT:-}" = mnnvl ]; then + return 0 + fi + [ -n "${COLLX_RDMA_DEVICES:-}" ] \ + || { collx_log "ERROR: scale-out RDMA selector is unavailable"; return 1; } + [[ "$COLLX_RDMA_DEVICES" =~ ^[A-Za-z][A-Za-z0-9_.-]{0,31}(:[1-9][0-9]*)?(,[A-Za-z][A-Za-z0-9_.-]{0,31}(:[1-9][0-9]*)?)*$ ]] \ + || { collx_log "ERROR: invalid scale-out RDMA selector"; return 1; } + export NCCL_IB_HCA="=$COLLX_RDMA_DEVICES" +} + +collx_default_route_interface() { + python3 "$COLLX_RUNTIME_DIR/probe.py" default-route-interface +} + +# Prove that the operator-pinned scale-out fabric exists on every allocated +# node before image import or backend initialization. Selector values and node +# diagnostics stay in the runner-private log. +collx_validate_network_profile_on_job() { + local job_id="$1" nodes="$2" transport="$3" + local log_label=network-profile log rc=0 marker_count link_layer + { [ "$nodes" -gt 1 ] && [ "$transport" != mnnvl ]; } || return 0 + [[ "$job_id" =~ ^[1-9][0-9]*$ && "$nodes" =~ ^[1-9][0-9]*$ ]] \ + || return 1 + [ -n "${COLLX_RDMA_DEVICES:-}" ] || return 1 + case "${COLLX_NETWORK_VALIDATION_ATTEMPT:-1}" in + 1) ;; + 2|3) log_label+="-a${COLLX_NETWORK_VALIDATION_ATTEMPT}" ;; + *) return 1 ;; + esac + log="$(collx_private_log_path "$log_label")" || return 1 + export COLLX_NETWORK_PROFILE_LOG="$log" + srun --jobid="$job_id" --nodes="$nodes" --ntasks="$nodes" --ntasks-per-node=1 \ + --chdir=/tmp --input=all --export="$(collx_host_exports)" \ + python3 /dev/stdin network-profile "${COLLX_SOCKET_IFNAME:-}" \ + "$COLLX_RDMA_DEVICES" "${COLLX_IB_GID_INDEX:-}" \ + < "$COLLX_RUNTIME_DIR/probe.py" > "$log" 2>&1 || rc=$? + if [ "$rc" != 0 ]; then + marker="$(grep -aoE '(socket-interface|rdma-(device|port))-[0-9]+=(missing|down|inactive|default-route-missing|gid-missing|gid-empty|link-layer-missing|link-layer-invalid|link-layer-mixed)' "$log" \ + | tail -n 1 || true)" + [ -z "$marker" ] || collx_log "ERROR: network-profile-$marker" + return "$rc" + fi + socket_ifname="$( + sed -nE 's/^\[collectivex-private\] socket-interface-selected=([A-Za-z][A-Za-z0-9_.-]{0,31})$/\1/p' "$log" \ + | sort -u + )" + marker_count="$(grep -Ec '^\[collectivex-private\] socket-interface-selected=' "$log")" + socket_unique_count="$(printf '%s\n' "$socket_ifname" | sed '/^$/d' | wc -l | tr -d ' ')" + if [ "$socket_unique_count" -lt 1 ] || [ "$marker_count" != "$nodes" ]; then + collx_log "ERROR: network-profile-socket-markers=$marker_count/$nodes unique=$socket_unique_count" + return 1 + fi + if [ "$socket_unique_count" = 1 ]; then + export COLLX_SOCKET_IFNAME="$socket_ifname" + else + unset COLLX_SOCKET_IFNAME + fi + link_layer="$( + sed -nE 's/^\[collectivex-private\] rdma-link-layer=(roce|infiniband)$/\1/p' "$log" \ + | sort -u + )" + marker_count="$(grep -Ec '^\[collectivex-private\] rdma-link-layer=(roce|infiniband)$' "$log")" + case "$marker_count:$link_layer" in + "$nodes":roce|"$nodes":infiniband) ;; + *) return 1 ;; + esac + export COLLX_RDMA_LINK_LAYER="$link_layer" + collx_export_gid_index_for_link_layer "$link_layer" +} + +collx_allocation_nodes_csv() { + local job_id="$1" nodelist node output="" + [[ "$job_id" =~ ^[1-9][0-9]*$ ]] || return 1 + nodelist="$(squeue -h -j "$job_id" -o %N 2>/dev/null)" || return 1 + [[ "$nodelist" =~ ^[][A-Za-z0-9._,-]+$ ]] || return 1 + while IFS= read -r node; do + [ -n "$node" ] || continue + [[ "$node" =~ ^[A-Za-z0-9][A-Za-z0-9._-]*$ ]] || return 1 + [ -z "$output" ] || output+=, + output+="$node" + done < <(scontrol show hostnames "$nodelist" 2>/dev/null) + [ -n "$output" ] || return 1 + printf '%s' "$output" +} + +collx_resolve_slurm_rendezvous() { + local job_id="$1" master_addr master_port socket_ifname="${COLLX_SOCKET_IFNAME:-}" + [[ "$job_id" =~ ^[1-9][0-9]*$ ]] || collx_die "invalid rendezvous allocation" + # Query relative node zero directly so MASTER_ADDR always hosts global rank 0. + # Prefer the address on the already validated cross-node socket interface; + # a short hostname may resolve onto a management network that ranks cannot use. + if [[ "$socket_ifname" =~ ^[A-Za-z][A-Za-z0-9_.-]{0,31}$ ]]; then + master_addr="$(srun --jobid="$job_id" --nodes=1 --ntasks=1 --relative=0 \ + --chdir=/tmp --export="$(collx_host_exports)" bash -s -- "$socket_ifname" \ + 2>/dev/null <<'BASH' | head -n1 +set -euo pipefail +ip -o -4 address show dev "$1" scope global \ + | awk 'NR == 1 {split($4, address, "/"); print address[1]}' +BASH +)" + [[ "$master_addr" =~ ^([0-9]{1,3}\.){3}[0-9]{1,3}$ ]] \ + || collx_die "could not resolve the allocated primary interface" + else + master_addr="$(srun --jobid="$job_id" --nodes=1 --ntasks=1 --relative=0 \ + --chdir=/tmp --export="$(collx_host_exports)" hostname -s 2>/dev/null | head -n1)" + [[ "$master_addr" =~ ^[A-Za-z0-9][A-Za-z0-9._-]*$ ]] \ + || collx_die "could not resolve the allocated primary node" + fi + master_port="${COLLX_MASTER_PORT:-29551}" + [[ "$master_port" =~ ^[1-9][0-9]*$ ]] && [ "$master_port" -le 65535 ] \ + || collx_die "invalid distributed rendezvous port" + export MASTER_ADDR="$master_addr" MASTER_PORT="$master_port" +} + +# Printed into `bash -c` ahead of the rank wrapper. Sources the per-node backend +# environment written during preparation. +collx_source_backend_env() { + cat <<'BASH' +case "${SLURM_NODEID:-}" in ""|*[!0-9]*) exit 66;; esac +. "/ix/experimental/CollectiveX/.collx_backend/env/node-${SLURM_NODEID}.sh" || exit 66 +BASH +} + +# Printed into `bash -c` for one Slurm task per GPU. Every rank derives its +# identity from Slurm rather than accepting caller-supplied rank values. +collx_slurm_rank_wrapper() { + cat <<'BASH' +case "${SLURM_PROCID:-}:${SLURM_NTASKS:-}:${SLURM_LOCALID:-}:${SLURM_NODEID:-}" in + *[!0-9:]*|:*|*::*|*:) exit 67 ;; +esac +[ "$SLURM_NTASKS" = "$COLLX_NGPUS" ] || exit 67 +[ "$SLURM_LOCALID" -lt "$COLLX_GPUS_PER_NODE" ] || exit 67 +. /ix/experimental/CollectiveX/runtime/common.sh || exit 68 +if [ "${COLLX_NODES:-1}" -gt 1 ] && [ "${COLLX_TRANSPORT:-}" != mnnvl ]; then + if [ -z "${COLLX_SOCKET_IFNAME:-}" ]; then + COLLX_SOCKET_IFNAME="$(collx_default_route_interface)" || exit 68 + [[ "$COLLX_SOCKET_IFNAME" =~ ^[A-Za-z][A-Za-z0-9_.-]{0,31}$ ]] || exit 68 + export COLLX_SOCKET_IFNAME + fi + collx_apply_network_profile "$COLLX_NODES" "$COLLX_TRANSPORT" || exit 68 +fi +export RANK="$SLURM_PROCID" WORLD_SIZE="$SLURM_NTASKS" +export LOCAL_RANK="$SLURM_LOCALID" LOCAL_WORLD_SIZE="$COLLX_GPUS_PER_NODE" +exec python3 bench/run_ep.py "$@" +BASH +} + +# Allocate via salloc's stable grant message and assign JOB_ID in this shell. +# Record it so workflow cleanup can release a launcher interrupted by Actions. +collx_salloc_jobid() { + local log_label=scheduler-allocation log job_id root="${COLLX_JOB_ROOT:-}" + case "${COLLX_SALLOC_ATTEMPT:-1}" in + 1) ;; + 2|3) log_label+="-a${COLLX_SALLOC_ATTEMPT}" ;; + *) return 1 ;; + esac + if ! log="$(collx_private_log_path "$log_label")"; then + collx_log "ERROR: scheduler log is unavailable" + return 1 + fi + collx_log "scheduler-request=submit" + if ! (salloc "$@" --no-shell) > "$log" 2>&1; then + collx_log "ERROR: scheduler allocation failed" + collx_log_tail "$log" + return 1 + fi + job_id="$(sed -nE \ + 's/.*Granted job allocation ([1-9][0-9]*).*/\1/p' "$log" | head -n1)" + [[ "$job_id" =~ ^[1-9][0-9]*$ ]] || return 1 + JOB_ID="$job_id" + if [ -n "$root" ]; then + collx_job_root_is_safe "$root" || return 1 + (umask 077; printf '%s\n' "$JOB_ID" > "$root/jobid") || return 1 + fi +} + +# Idempotent cleanup for launcher traps, allocation retries, and workflow recovery. +collx_cleanup_allocation() { + local root="${1:-${COLLX_JOB_ROOT:-}}" path="" job_id="${JOB_ID:-}" active + if [ -n "$root" ]; then + collx_job_root_is_safe "$root" || return 1 + path="$root/jobid" + if [ -z "$job_id" ] && [ -f "$path" ]; then + IFS= read -r job_id < "$path" || return 1 + fi + fi + [ -n "$job_id" ] || return 0 + [[ "$job_id" =~ ^[1-9][0-9]*$ ]] || return 1 + scancel "$job_id" >/dev/null 2>&1 || true + for _ in {1..30}; do + active="$(squeue -h -j "$job_id" -o %A 2>/dev/null)" || active=unknown + if [ -z "$active" ]; then + [ -z "$path" ] || rm -f -- "$path" + return + fi + sleep 1 + done + collx_log "ERROR: scheduled allocation did not terminate during cleanup" + return 1 +} + +# Import uses the configured tag because Enroot cannot reliably import a +# digest-qualified Docker Hub reference non-interactively. +collx_select_image() { + local image="$1" + [[ "$image" =~ ^[A-Za-z0-9._/-]+:[A-Za-z0-9._-]+$ ]] \ + || collx_die "configured image reference is malformed" + export COLLECTIVEX_IMAGE="$image" +} + +# Create a per-UID cache under validated cluster-local storage. Only the fixed +# /cx-cache mount enters the container; the operator host path does not. +collx_prepare_backend_cache() { + local cache + unset COLLX_PREPARED_BACKEND_CACHE + cache="$(python3 "$COLLX_RUNTIME_DIR/probe.py" prepare-cache "$1")" || return 1 + [[ "$cache" = /* ]] || return 1 + export COLLX_PREPARED_BACKEND_CACHE="$cache" +} + +# Fetch the pinned DeepEP tree before allocating GPUs. +collx_prepare_deepep_source() { + local mount_src="$1" root source temporary log + root="$mount_src/experimental/CollectiveX/.collx_sources" + source="$root/deepep-v2-$COLLX_DEEPEP_V2_COMMIT" + [ ! -d "$source" ] || return 0 + mkdir -p -- "$root" && chmod 700 "$root" || return 1 + temporary="$(mktemp -d "$root/.deepep-v2.XXXXXX")" || return 1 + log="$(collx_private_log_path backend-source-deepep-v2)" || return 1 + # On b300 the NFS export can realize a newly created stage dir as UID 0 while + # git runs as the UID-mapped Actions user, tripping git's "dubious ownership" + # guard on the source tree and its fmt submodule. HOME is this job's ephemeral + # dir and the runner UID is inside the trusted cluster boundary, so scope the + # exemption globally (also reaches the submodule child git). + git config --global --add safe.directory '*' >> "$log" 2>&1 || true + if GIT_TERMINAL_PROMPT=0 git init -q "$temporary" > "$log" 2>&1 \ + && git -C "$temporary" remote add origin "$COLLX_DEEPEP_V2_REPO" >> "$log" 2>&1 \ + && GIT_TERMINAL_PROMPT=0 git -C "$temporary" fetch -q --no-tags --depth 1 \ + origin "$COLLX_DEEPEP_V2_COMMIT" >> "$log" 2>&1 \ + && git -C "$temporary" -c advice.detachedHead=false checkout -q --detach FETCH_HEAD \ + >> "$log" 2>&1 \ + && [ "$(git -C "$temporary" rev-parse HEAD)" = "$COLLX_DEEPEP_V2_COMMIT" ] \ + && GIT_TERMINAL_PROMPT=0 git -C "$temporary" submodule update -q --init --depth 1 \ + third-party/fmt >> "$log" 2>&1 \ + && python3 "$COLLX_RUNTIME_DIR/stage.py" rewrite-deepep-v2 \ + "$temporary/deep_ep/__init__.py" >> "$log" 2>&1 \ + && mv -- "$temporary" "$source" >> "$log" 2>&1; then + return 0 + fi + rm -rf -- "$temporary" + collx_log "ERROR: DeepEP source preparation failed" + collx_log_tail "$log" + return 1 +} + +collx_materialize_deepep_source() { + local destination="$1" source + [ -n "${COLLX_BACKEND_SOURCE_ROOT:-}" ] || return 1 + source="$COLLX_BACKEND_SOURCE_ROOT/deepep-v2-$COLLX_DEEPEP_V2_COMMIT" + [ -d "$source" ] || return 1 + rm -rf -- "$destination" && cp -R -- "$source" "$destination" +} + +collx_prepare_implicit_stage_base() { + python3 "$COLLX_RUNTIME_DIR/stage.py" implicit-stage-base "${1:-}" "${2:-}" +} + +collx_prepare_runner_shared_stage_base() { + local runner_temp="${RUNNER_TEMP:-}" runner_root + case "$runner_temp" in + /*/_work/_temp) runner_root="${runner_temp%/_work/_temp}" ;; + *) collx_die "canonical AMD execution requires a standard shared runner temp" ;; + esac + [ -n "$runner_root" ] && [ "$runner_root" != "$runner_temp" ] \ + || collx_die "canonical AMD execution requires a shared runner root" + collx_prepare_implicit_stage_base "$runner_root" +} + +collx_prepare_stage_dir() { + local runner="$1" + [ "${COLLECTIVEX_CANONICAL_GHA:-0}" = 1 ] || return 0 + [ -n "${COLLX_SQUASH_DIR:-}" ] \ + || collx_die "canonical CollectiveX execution requires shared container storage" + case "$runner" in b300|gb300) COLLX_STAGE_DIR="" ;; esac + if [ -z "${COLLX_STAGE_DIR:-}" ]; then + case "$runner" in + h100-dgxc) + COLLX_STAGE_DIR="$(collx_prepare_implicit_stage_base "${COLLX_SQUASH_DIR%/*}")" \ + || collx_die "canonical CollectiveX execution cannot create an isolated shared stage directory" + ;; + b300|gb300) + COLLX_STAGE_DIR="$(collx_prepare_implicit_stage_base "" \ + "${COLLECTIVEX_EXECUTION_ID:-${GITHUB_RUN_ID:-}}")" \ + || collx_die "canonical CollectiveX execution cannot create an isolated stage directory" + ;; + h200-dgxc|b200-dgxc) + COLLX_STAGE_DIR="$(collx_prepare_implicit_stage_base)" \ + || collx_die "canonical CollectiveX execution cannot create an isolated stage directory" + ;; + mi300x|mi325x|mi355x) + COLLX_STAGE_DIR="$(collx_prepare_runner_shared_stage_base)" \ + || collx_die "canonical AMD execution cannot create an isolated shared stage directory" + ;; + *) collx_die "canonical CollectiveX execution requires a configured shared stage directory" ;; + esac + elif [ "$runner" = mi300x ]; then + COLLX_STAGE_DIR="$(python3 "$COLLX_RUNTIME_DIR/stage.py" resolve-directory \ + "$COLLX_STAGE_DIR")" \ + || collx_die "canonical MI300X execution cannot resolve the shared stage directory" + fi + export COLLX_STAGE_DIR +} + +collx_squash_path() { + local squash_dir="$1" image="$2" key platform run_scope + case "${COLLX_IMAGE_PLATFORM:-}" in + linux/amd64) platform="" ;; + linux/arm64) platform="_linux_arm64" ;; + *) return 1 ;; + esac + run_scope="${GITHUB_RUN_ID:-${COLLECTIVEX_EXECUTION_ID:-manual}}-${GITHUB_RUN_ATTEMPT:-1}" + run_scope="$(printf '%s' "$run_scope" | tr -cs 'A-Za-z0-9_.-' '-')" || return 1 + run_scope="${run_scope#-}"; run_scope="${run_scope%-}" + [ -n "$run_scope" ] || return 1 + key="${platform}_${run_scope}_$( + printf '%s' "$image" | sed 's#[/:@#]#_#g' + )" + printf '%s' "$squash_dir/${key}.sqsh" +} + +# collx_ensure_squash -> echoes the squash file path. +# Imports via Enroot only if a valid squash is not already present, under a lock. +collx_ensure_squash() { + local squash_dir="$1" image="$2" key sq locks lock_fd log + local enroot_local="" import_rc=0 machine + log="$(collx_private_log_path container-import)" + machine="$(uname -m)" + case "${COLLX_IMAGE_PLATFORM:-}:$machine" in + linux/amd64:x86_64|linux/amd64:amd64|linux/arm64:aarch64|linux/arm64:arm64) ;; + *) collx_log_tail "$log"; return 1 ;; + esac + mkdir -p "$squash_dir" 2>> "$log" \ + || { collx_log_tail "$log"; return 1; } + sq="$(collx_squash_path "$squash_dir" "$image")" \ + || { collx_log_tail "$log"; return 1; } + key="${sq##*/}" + key="${key%.sqsh}" + locks="$squash_dir/.locks" + mkdir -p "$locks" 2>> "$log" \ + || { collx_log_tail "$log"; return 1; } + { exec {lock_fd}>"$locks/${key}.lock"; } 2>> "$log" \ + || { collx_log_tail "$log"; return 1; } + # A concurrent leg of the same run holds this lock for its full import + # (measured ~18 minutes for the 32 GB sglang squash on b300), so the wait + # must outlast an import and a timeout must say so — the empty import log + # would otherwise make this the only silent launcher death. + flock -w 2700 "$lock_fd" 2>> "$log" \ + || { collx_log "ERROR: timed out waiting for the container import lock" + collx_log_tail "$log"; return 1; } + if unsquashfs -l "$sq" >/dev/null 2>&1; then + collx_log "container squash ready" + else + collx_log "importing configured container image" + rm -f "$sq" 2>> "$log" \ + || { collx_log_tail "$log"; return 1; } + # > "$log" 2>&1 || import_rc=$? + rm -rf -- "$enroot_local" >/dev/null 2>&1 || true + [ "$import_rc" = 0 ] \ + || { collx_log_tail "$log"; return 1; } + else + enroot import -o "$sq" "docker://$image" > "$log" 2>&1 \ + || { collx_log_tail "$log"; return 1; } + fi + unsquashfs -l "$sq" >> "$log" 2>&1 \ + || { collx_log_tail "$log"; return 1; } + fi + flock -u "$lock_fd" + exec {lock_fd}>&- + echo "$sq" +} + +# Import on an allocated compute node so multiarch tags resolve for the target +# architecture. The squash directory must be shared with the submit host. +collx_ensure_squash_on_job() { + local job_id="$1" squash_dir="$2" image="$3" lock_dir="${4:-}" sq key lock + local log_label=container-import log + [[ "$job_id" =~ ^[0-9]+$ ]] || return 1 + case "${COLLX_SALLOC_ATTEMPT:-1}" in + 1) ;; + 2|3) log_label+="-a${COLLX_SALLOC_ATTEMPT}" ;; + *) return 1 ;; + esac + sq="$(collx_squash_path "$squash_dir" "$image")" || return 1 + key="${sq##*/}" + key="${key%.sqsh}" + [ -n "$lock_dir" ] || lock_dir="$squash_dir/.locks" + lock="$lock_dir/${key}.lock" + log="$(collx_private_log_path "$log_label")" + # Run once per node because some clusters use node-local squash storage. + if ! srun --jobid="$job_id" --nodes="${COLLX_NODES:-1}" --ntasks="${COLLX_NODES:-1}" \ + --ntasks-per-node=1 --chdir=/tmp \ + --export="$(collx_host_exports)" \ + bash -s -- "$sq" "$lock" "$image" "$COLLX_IMAGE_PLATFORM" \ + > "$log" 2>&1 <<'BASH' +set -euo pipefail +sq="$1"; lock="$2"; image="$3"; platform="$4" +machine="$(uname -m)" +case "$platform:$machine" in + linux/amd64:x86_64|linux/amd64:amd64|linux/arm64:aarch64|linux/arm64:arm64) ;; + *) exit 13 ;; +esac +compute_home="$(mktemp -d /tmp/inferencex-collectivex-home.XXXXXX)" +trap 'rm -rf -- "$compute_home"' EXIT +export HOME="$compute_home" XDG_CACHE_HOME="$compute_home/.cache" +export ENROOT_TEMP_PATH="$compute_home/enroot-tmp" +export ENROOT_CACHE_PATH="$compute_home/enroot-cache" +export ENROOT_DATA_PATH="$compute_home/enroot-data" +export ENROOT_RUNTIME_PATH="$compute_home/enroot-run" +mkdir -p "$(dirname "$sq")" "$(dirname "$lock")" \ + "$ENROOT_TEMP_PATH" "$ENROOT_CACHE_PATH" "$ENROOT_DATA_PATH" "$ENROOT_RUNTIME_PATH" +exec 9>"$lock" +# Shared storage serializes the import; node-local storage imports in parallel. +flock 9 +if unsquashfs -l "$sq" >/dev/null 2>&1; then + echo 'container squash ready' +else + rm -f -- "$sq" + enroot import -o "$sq" "docker://$image" /dev/null 2>&1 +fi +BASH + then + collx_log "ERROR: container import failed" + collx_log_tail "$log" + return 1 + fi + printf '%s' "$sq" +} + +# A clean nvidia-smi inventory does not prove that a prior cancelled workload +# released every CUDA context. Retaining each primary context catches poisoned +# allocations before a full shard spends time failing every case. +collx_validate_cuda_context_on_job() { + local job_id="$1" nodes="$2" gpus_per_node="$3" log_label=cuda-context log + case "${COLLX_SALLOC_ATTEMPT:-1}" in + 1) ;; + 2|3) log_label+="-a${COLLX_SALLOC_ATTEMPT}" ;; + *) return 1 ;; + esac + log="$(collx_private_log_path "$log_label")" + export COLLX_CUDA_CONTEXT_LOG="$log" + srun --jobid="$job_id" --nodes="$nodes" --ntasks="$nodes" --ntasks-per-node=1 \ + --gres=gpu:"$gpus_per_node" --chdir=/tmp --input=all \ + --export="$(collx_host_exports)" python3 /dev/stdin cuda-context "$gpus_per_node" \ + < "$COLLX_RUNTIME_DIR/probe.py" >"$log" 2>&1 +} + +# Resolve the exact per-execution child before any copy starts, so the parent +# EXIT trap can remove an interrupted partial stage. The configured base must +# already exist on compute-visible storage and must not traverse symlinks. +collx_stage_path() { + local repo_root="$1" stage_base="${2:-}" tag stage_path + tag="${COLLECTIVEX_EXECUTION_ID:-${GITHUB_RUN_ID:-manual-$$}}" + [[ "$tag" =~ ^[A-Za-z0-9][A-Za-z0-9._-]*$ ]] \ + || collx_die "invalid staging execution identity" + if [ -z "$stage_base" ] || [ "$stage_base" = "$repo_root" ]; then + [ -n "${COLLX_SQUASH_DIR:-}" ] \ + || collx_die "CollectiveX staging requires COLLX_STAGE_DIR or COLLX_SQUASH_DIR" + stage_base="$COLLX_SQUASH_DIR" + stage_path="${stage_base%/}/.collectivex-stage-$tag" + else + stage_path="${stage_base%/}/job_$tag" + fi + python3 "$COLLX_RUNTIME_DIR/stage.py" validate-stage-path "$repo_root" "$stage_base" \ + "$stage_path" "${COLLX_JOB_ROOT:-}" "${GITHUB_WORKSPACE:-}" +} + +# Stage only the public benchmark tree into the private execution child. +collx_stage_repo() { + local repo_root="$1" stage_dir="$2" log + python3 "$COLLX_RUNTIME_DIR/stage.py" create-stage "$stage_dir" \ + || collx_die "cannot create the configured stage directory" + collx_log "staging CollectiveX on compute-visible storage" + log="$(collx_private_log_path repository-stage)" + if ! python3 "$COLLX_RUNTIME_DIR/stage.py" copy-repository \ + "$repo_root/experimental/CollectiveX" \ + "$stage_dir/experimental/CollectiveX" > "$log" 2>&1; then + rm -rf -- "$stage_dir" >/dev/null 2>&1 \ + || collx_log "ERROR: cannot remove the incomplete execution stage" + collx_log "ERROR: repository staging failed" + collx_log_tail "$log" + return 1 + fi +} + +# collx_collect_results +# When the run used a staged (compute-visible) mount, copy result JSONs back to +# the original checkout's results/ so the workflow's upload-artifact (which reads +# the checkout, not the stage dir) finds them. No-op when no staging was used. +collx_collect_results() { + local mount_src="$1" repo_root="$2" dst log + local -a files + [ "$mount_src" = "$repo_root" ] && return 0 + log="$(collx_private_log_path "artifact-collection-$$-${RANDOM}")" + dst="$repo_root/experimental/CollectiveX/results" + mkdir -p "$dst" 2>> "$log" \ + || { collx_log "ERROR: cannot create checkout result directory"; return 1; } + shopt -s nullglob + files=("$mount_src/experimental/CollectiveX/results/"*.json) + shopt -u nullglob + [ "${#files[@]}" -gt 0 ] || { collx_log "ERROR: staged run produced no result JSON"; return 1; } + cp -- "${files[@]}" "$dst/" >> "$log" 2>&1 \ + || { collx_log "ERROR: staged result collection failed"; return 1; } + collx_log "collected staged results for artifact validation" +} + +collx_cleanup_stage() { + local mount_src="$1" repo_root="$2" + [ "$mount_src" != "$repo_root" ] || return 0 + if ! python3 "$COLLX_RUNTIME_DIR/stage.py" validate-cleanup "$mount_src"; then + collx_log "ERROR: refusing to remove an invalid stage directory" + return 1 + fi + rm -rf -- "$mount_src" >/dev/null 2>&1 || { + collx_log "ERROR: cannot remove generated stage directory" + return 1 + } + collx_log "removed generated per-execution stage directory" +} + +# Run one shard with one Slurm task per GPU on one or more nodes. +# Launchers provide only allocation/container policy through globals and +# COLLX_DISTRIBUTED_CONTAINER_ARGS; per-case benchmark inputs travel as run_ep.py +# argv decoded from the shard control (config.py case-args), never as env. +# shellcheck disable=SC2153 +collx_run_shard() { + local build_log expected_cases ci=0 failed_cases=0 + local runtime_log argv_file shard wrap + local -a container_args ep_args + [ "${NODES:-0}" -ge 1 ] && [ "${NGPUS:-0}" = "$((NODES * GPN))" ] \ + || collx_die "invalid shard launcher placement" + [ -n "${JOB_ID:-}" ] && [ -n "${SQUASH_FILE:-}" ] \ + && [ -n "${CONTAINER_MOUNTS:-}" ] || collx_die "shard launcher is incomplete" + wrap="$(collx_source_backend_env)"$'\n'"$(collx_slurm_rank_wrapper)" + + collx_resolve_slurm_rendezvous "$JOB_ID" + collx_apply_network_profile "$NODES" "${COLLX_TRANSPORT:-}" + mkdir -p "$MOUNT_SRC/experimental/CollectiveX/results" + container_args=(--container-mounts="$CONTAINER_MOUNTS" --no-container-mount-home + --container-workdir=/ix/experimental/CollectiveX --no-container-entrypoint) + if declare -p COLLX_DISTRIBUTED_CONTAINER_ARGS >/dev/null 2>&1; then + container_args+=("${COLLX_DISTRIBUTED_CONTAINER_ARGS[@]}") + fi + local container_name="cxep_${JOB_ID}" + + shard="${COLLX_SHARD_FILE:-}" + [ -f "$shard" ] || shard="$COLLX_DIR/$shard" + [ -f "$shard" ] || collx_die "shard control is unavailable" + expected_cases="$(python3 "$COLLX_RUNTIME_DIR/config.py" case-count "$shard")" \ + && [[ "$expected_cases" =~ ^[1-9][0-9]*$ ]] \ + || collx_die "could not enumerate shard cases" + + collx_log "shard backend preparation: bench=$COLLX_BENCH nodes=$NODES" + build_log="$(collx_private_log_path backend-prepare)" + if ! srun --jobid="$JOB_ID" --nodes="$NODES" --ntasks-per-node=1 --chdir=/tmp \ + --container-name="$container_name" --container-image="$SQUASH_FILE" \ + "${container_args[@]}" --export=ALL \ + bash /ix/experimental/CollectiveX/runtime/prepare_backend.sh \ + "$build_log" 2>&1; then + collx_log "ERROR: backend preparation failed" + collx_log_tail "$build_log" + return 1 + fi + + argv_file="$(mktemp)" || return 1 + while [ "$ci" -lt "$expected_cases" ]; do + python3 "$COLLX_RUNTIME_DIR/config.py" case-args "$shard" "$ci" \ + "$RUNNER" "$TS" \ + "$NGPUS" "$NODES" "$GPN" "$SCALE_UP_DOMAIN" > "$argv_file" \ + || { rm -f "$argv_file"; collx_die "shard case $ci does not decode against this allocation"; } + mapfile -d '' -t ep_args < "$argv_file" + [ "${#ep_args[@]}" -gt 0 ] \ + || { rm -f "$argv_file"; collx_die "case $ci produced no benchmark arguments"; } + collx_log "EP${NGPUS}[$((ci + 1))/$expected_cases] $COLLX_BENCH" + runtime_log="$(collx_private_log_path "runtime-c$(printf '%03d' "$ci")")" + if ! timeout -k 30 "${COLLX_RUN_TIMEOUT:-900}" \ + srun --jobid="$JOB_ID" --nodes="$NODES" \ + --ntasks="$NGPUS" --ntasks-per-node="$GPN" --chdir=/tmp \ + --container-name="$container_name" --container-image="$SQUASH_FILE" \ + "${container_args[@]}" \ + --export=ALL \ + bash -c "$wrap" _ "${ep_args[@]}" \ + "$runtime_log" 2>&1; then + collx_log "ERROR: case $ci failed" + collx_log_tail "$runtime_log" + failed_cases=$((failed_cases + 1)) + fi + ci=$((ci + 1)) + done + rm -f "$argv_file" + [ "$failed_cases" = 0 ] || { + collx_log "ERROR: $failed_cases/$expected_cases case(s) failed" + return 1 + } +} + +# Remove this allocation's persistent pyxis container before the allocation is +# released. Clusters may run pyxis with container_scope=global, where the named +# --container-writable container every shard uses (cxep_) survives job +# teardown and its unpacked rootfs — tens of GB per node — would otherwise +# accumulate on every allocated node's local image store until it fills and the +# next writable extraction fails with ENOSPC. Best-effort and bounded: teardown +# must never hang or fail on this. +collx_remove_distributed_container() { + local job_id="$1" nodes="${2:-1}" + [ -n "$job_id" ] || return 0 + [ "$nodes" -ge 1 ] 2>/dev/null || return 0 + timeout 120 srun --jobid="$job_id" --nodes="$nodes" --ntasks-per-node=1 \ + --chdir=/tmp enroot remove -f "pyxis_cxep_${job_id}" \ + /dev/null 2>&1 || true +} + +collx_launcher_cleanup() { + local rc="$1" stage_root="${MOUNT_SRC:-}" + trap - EXIT HUP INT TERM + if [ -n "${JOB_ID:-}" ]; then + collx_remove_distributed_container "$JOB_ID" "${NODES:-1}" + if ! collx_cleanup_allocation; then + [ "$rc" != 0 ] || rc=1 + exit "$rc" + fi + fi + if [ -n "${REPO_ROOT:-}" ] && [ -n "$stage_root" ] \ + && [ "$stage_root" != "$REPO_ROOT" ]; then + if [ "$rc" != 0 ] && [ -d "$stage_root/experimental/CollectiveX" ]; then + collx_collect_results "$stage_root" "$REPO_ROOT" || true + fi + if ! collx_cleanup_stage "$stage_root" "$REPO_ROOT"; then + [ "$rc" != 0 ] || rc=1 + fi + fi + exit "$rc" +} + +collx_install_launcher_fail_safe() { + trap 'collx_launcher_cleanup "$?"' EXIT + trap 'collx_launcher_cleanup 129' HUP + trap 'collx_launcher_cleanup 130' INT + trap 'collx_launcher_cleanup 143' TERM +} diff --git a/experimental/CollectiveX/runtime/config.py b/experimental/CollectiveX/runtime/config.py new file mode 100644 index 0000000000..dc3c847fd1 --- /dev/null +++ b/experimental/CollectiveX/runtime/config.py @@ -0,0 +1,174 @@ +#!/usr/bin/env python3 +"""Load private runner settings, the public backend registry, and shard controls.""" + +from __future__ import annotations + +import argparse +import json +import os +import sys + + +OPERATOR_FIELDS = { + "partition", "account", "qos", "squash_dir", "stage_dir", + "enroot_cache_path", "exclude_nodes", "nodelist", "lock_dir", +} +NETWORK_FIELDS = { + "socket_ifname", "rdma_devices", "ib_gid_index", "rdma_service_level", + "rdma_traffic_class", "rail_isolated", +} + + +def _platforms() -> dict: + """The per-SKU platform registry (configs/platform_config.json). Callers + fail closed on a missing file, unknown SKU, or missing field.""" + path = os.path.join( + os.path.dirname(os.path.dirname(os.path.abspath(__file__))), + "configs", "platform_config.json", + ) + with open(path, encoding="utf-8") as stream: + return json.load(stream)["platforms"] + + +def emit(values: dict[str, object]) -> None: + for field, value in values.items(): + name = f"COLLX_{field.upper()}" + sys.stdout.buffer.write(name.encode() + b"\0" + str(value).encode() + b"\0") + + +def _network_overlay(runner: str) -> dict[str, object]: + """Repo-tracked per-SKU scale-out RDMA selectors — the `network` block of the + SKU's configs/platform_config.json entry — overlaid onto the base operator + config. Only NETWORK_FIELDS are taken, so identity keys and notes are ignored; + a missing/invalid file is a no-op fallback to the base/secret network fields.""" + try: + block = _platforms().get(runner, {}).get("network", {}) + except (KeyError, OSError, TypeError, json.JSONDecodeError): + return {} + return {key: value for key, value in block.items() if key in NETWORK_FIELDS} + + +def operator_config(path: str, runner: str) -> None: + try: + platform = _platforms()[runner] + # The registry's tracked per-SKU `operator` block is the baseline + # (de-secreted by operator decision); an operator config document, when + # provided, overrides it per field. Path "-" means registry-only. + selected = dict(platform.get("operator", {})) + if path != "-": + with open(path, encoding="utf-8") as stream: + document = json.load(stream) + selected.update(document["runners"].get(runner, {})) + # Overlay repo-tracked scale-out RDMA selectors onto the base runner config; + # SKUs without a platform_config.json network block keep their base/secret + # network fields. + selected.update(_network_overlay(runner)) + allowed = OPERATOR_FIELDS | NETWORK_FIELDS | {"storage_roots"} + if set(selected) - allowed: + raise ValueError + roots = selected.pop("storage_roots", None) + if roots: + for root in roots: + squash = os.path.join(root, "collectivex", "containers") + stage = os.path.join(root, "collectivex", "stage") + try: + os.makedirs(squash, mode=0o700, exist_ok=True) + os.makedirs(stage, mode=0o700, exist_ok=True) + selected.update(squash_dir=squash, stage_dir=stage) + break + except OSError: + continue + else: + raise ValueError + if any(not isinstance(value, (str, int)) or "\0" in str(value) for value in selected.values()): + raise ValueError + selected.update(image=platform["image"], image_platform=platform["image_platform"]) + emit(selected) + except (KeyError, OSError, TypeError, ValueError, json.JSONDecodeError): + print("validation-invalid-config", file=sys.stderr) + raise SystemExit(1) + + +def load(path: str) -> dict: + with open(path, encoding="utf-8") as stream: + return json.load(stream) + + +def case_count(path: str) -> None: + print(len(load(path)["cases"]), end="") + + +def _emit_argv(case: dict, version: object, runner: str, ts: str, index: int) -> None: + """Emit one null-delimited run_ep.py argv — the only case-to-invocation codec.""" + get = lambda key, default="": str(case.get(key) or default) + argv = [ + "--backend", str(case["backend"]), + "--mode", str(case["mode"]), + "--phase", str(case["phase"]), + "--routing", str(case["routing"]), + "--gpus-per-node", str(case["gpus_per_node"]), + "--scale-up-domain", str(case["scale_up_domain"]), + "--scope", str(case["scope"]), + "--scale-up-transport", str(case["scale_up_transport"]), + "--scale-out-transport", get("scale_out_transport"), + "--tokens-ladder", str(case["ladder"]), + "--hidden", str(case["hidden"]), + "--topk", str(case["topk"]), + "--experts", str(case["experts"]), + "--seed", str(case["seed"]), + "--runner", runner, + "--topology-class", str(case["topology_class"]), + "--transport", str(case["transport"]), + "--case-id", str(case["case_id"]), + "--suite", str(case["suite"]), + "--workload-name", str(case["workload"]), + "--version", str(version), + ] + iters, trials, warmup = str(case["timing"]).split(":") + for flag, value in (("--iters", iters), ("--trials", trials), ("--warmup", warmup)): + argv += [flag, value] + out = f"results/{runner}_{case['backend']}_{case['phase']}_{ts}-c{index:03d}.json" + argv += ["--out", out] + sys.stdout.buffer.write(b"\0".join(part.encode() for part in argv) + b"\0") + + +def case_args( + path: str, index: int, runner: str, ts: str, + ngpus: str, nodes: str, gpus_per_node: str, scale_up_domain: str, +) -> None: + document = load(path) + cases = document["cases"] + if not 0 <= index < len(cases): + raise SystemExit(1) + case = cases[index] + placement = tuple( + str(case.get(field, "")) + for field in ("ep", "nodes", "gpus_per_node", "scale_up_domain") + ) + if placement != (ngpus, nodes, gpus_per_node, scale_up_domain): + print(f"case placement {placement} differs from the allocation", file=sys.stderr) + raise SystemExit(1) + _emit_argv(case, document["version"], runner, ts, index) + + +def main() -> None: + parser = argparse.ArgumentParser() + commands = parser.add_subparsers(dest="command", required=True) + for name, names in { + "operator-config": ("path", "runner"), + "case-count": ("path",), + "case-args": ("path", "index", "runner", "ts", + "ngpus", "nodes", "gpus_per_node", "scale_up_domain"), + }.items(): + command = commands.add_parser(name) + for arg in names: command.add_argument(arg) + args = parser.parse_args() + if args.command == "operator-config": operator_config(args.path, args.runner) + elif args.command == "case-count": case_count(args.path) + elif args.command == "case-args": + case_args(args.path, int(args.index), args.runner, args.ts, + args.ngpus, args.nodes, args.gpus_per_node, args.scale_up_domain) + + +if __name__ == "__main__": + main() diff --git a/experimental/CollectiveX/runtime/prepare_backend.sh b/experimental/CollectiveX/runtime/prepare_backend.sh new file mode 100644 index 0000000000..132203fec1 --- /dev/null +++ b/experimental/CollectiveX/runtime/prepare_backend.sh @@ -0,0 +1,326 @@ +#!/usr/bin/env bash +# Prepare one backend per allocated node and persist its rank environment. +set -euo pipefail + +cd /ix/experimental/CollectiveX +# shellcheck source=../runtime/common.sh +source runtime/common.sh + +: "${COLLX_RUNNER:?COLLX_RUNNER not set}" +: "${COLLX_BENCH:?COLLX_BENCH not set}" + +collx_log "backend preparation: runner=$COLLX_RUNNER bench=$COLLX_BENCH nodes=${COLLX_NODES:-1}" + +# Fresh rank tasks source only these backend-created values. Network variables +# are reapplied by the rank wrapper from the platform profile. +readonly -a RANK_ENV_VARS=( + PATH VIRTUAL_ENV LD_LIBRARY_PATH PYTHONPATH CUDA_HOME CPATH NVCC_PREPEND_FLAGS + NVSHMEM_DIR EP_NCCL_ROOT_DIR EP_NVSHMEM_ROOT_DIR EP_JIT_CACHE_DIR + EP_REUSE_NCCL_COMM NCCL_CUMEM_ENABLE +) +readonly -a DEEPEP_RANK_UNSETS=(EP_SUPPRESS_NCCL_CHECK) + +# ---- discovery -------------------------------------------------------------- + +cuda_arch() { + local expected detected + expected="$(python3 - "$COLLX_RUNNER" <<'PY' +import json, sys +arch = json.load(open("configs/platform_config.json"))["platforms"][sys.argv[1]]["arch"] +digits = arch.removeprefix("sm") +print(f"{digits[:-1]}.{digits[-1]}" if arch.startswith("sm") and digits.isdigit() else "") +PY +)" || { collx_log "ERROR: no platform registry entry for $COLLX_RUNNER"; return 1; } + [ -n "$expected" ] || { + collx_log "ERROR: no CUDA target registered for $COLLX_RUNNER"; return 1 + } + detected="$(python3 - <<'PY' +import torch + +major, minor = torch.cuda.get_device_capability() +print(f"{major}.{minor}") +PY +)" || return 1 + [ "$detected" = "$expected" ] || { + collx_log "ERROR: $COLLX_RUNNER expected CUDA target $expected, detected $detected" + return 1 + } + printf '%s' "$detected" +} + +nvidia_package_root() { + local python="$1" package="$2" component="$3" + "$python" - "$package" "$component" <<'PY' +from importlib import metadata +from pathlib import Path, PurePosixPath +import sys + +package, component = sys.argv[1:] +try: + distribution = metadata.distribution(package) + prefix = f"nvidia/{component}/" + entries = [str(entry).replace("\\", "/") for entry in distribution.files or ()] + if not any(entry.startswith(prefix) for entry in entries): + raise ValueError + root = Path(distribution.locate_file(PurePosixPath("nvidia") / component)).resolve() + if not root.is_dir(): + raise ValueError +except (metadata.PackageNotFoundError, OSError, TypeError, ValueError): + raise SystemExit(1) +print(root, end="") +PY +} + +cuda_toolchain_paths() { + local cccl="" candidate cuda_home nvcc + nvcc="$(command -v nvcc)" || { collx_log "ERROR: CUDA nvcc is unavailable"; return 1; } + nvcc="$(readlink -f -- "$nvcc")" || { collx_log "ERROR: CUDA nvcc cannot be resolved"; return 1; } + case "$nvcc" in + */bin/nvcc) cuda_home="${nvcc%/bin/nvcc}" ;; + *) collx_log "ERROR: CUDA nvcc has an unexpected path"; return 1 ;; + esac + [ -x "$cuda_home/bin/nvcc" ] && [ -d "$cuda_home/include" ] && [ -d "$cuda_home/lib64" ] \ + || { collx_log "ERROR: CUDA toolkit root is incomplete"; return 1; } + for candidate in "$cuda_home"/targets/*/include/cccl; do + if [ -d "$candidate" ]; then + cccl="$candidate" + break + fi + done + [ -n "$cccl" ] || { collx_log "ERROR: CUDA CCCL headers are unavailable"; return 1; } + printf '%s\t%s' "$cuda_home" "$cccl" +} + +deepep_nvshmem_overlay() { + local root="$1" packaged="$2" overlay path temporary + overlay="$root/nvshmem-overlay" + if ! ( + umask 077 + exec 8>"$root/nvshmem-overlay.lock" || exit 1 + flock 8 || exit 1 + if [ ! -d "$overlay" ]; then + temporary="$root/.nvshmem-overlay.$$" + rm -rf "$temporary" || exit 1 + mkdir -p "$temporary/lib" || exit 1 + ln -s "$packaged/include" "$temporary/include" || exit 1 + for path in "$packaged"/lib/*; do + ln -s "$path" "$temporary/lib/${path##*/}" || exit 1 + done + [ ! -e "$packaged/lib/libnvshmem_host.so.3" ] \ + || ln -sf "$packaged/lib/libnvshmem_host.so.3" \ + "$temporary/lib/libnvshmem_host.so" || exit 1 + mv "$temporary" "$overlay" || exit 1 + fi + [ ! -L "$overlay" ] \ + && [ "$(readlink -f "$overlay/include")" = "$(readlink -f "$packaged/include")" ] \ + && [ -e "$overlay/lib/libnvshmem_host.so" ] \ + && [ -e "$overlay/lib/libnvshmem_device.a" ] + ); then + collx_log "ERROR: DeepEP V2 NVSHMEM overlay is invalid" + return 1 + fi + printf '%s' "$overlay" +} + +deepep_cache_root() { + local arch="$1" cpu base image + cpu="$(uname -m)" + [[ "$cpu" =~ ^[A-Za-z0-9._-]+$ ]] || return 1 + base="${COLLX_BACKEND_CACHE_ROOT:-}" + [[ "$base" = /* ]] || return 1 + image="$(printf '%s' "${COLLECTIVEX_IMAGE:-manual}" | tr -cs 'A-Za-z0-9_.-' '-')" + printf '%s/deepep-v2-%s-sm%s-%s-%s' \ + "$base" "$cpu" "${arch/./}" "${image#-}" "${COLLX_DEEPEP_V2_COMMIT:0:12}" +} + +deepep_activate() { + local root="$1" venv venv_site nccl_root nvshmem_package overlay + local toolchain cuda_home cccl execution_id + venv="$root/venv" + [ -x "$venv/bin/python" ] \ + || { collx_log "ERROR: DeepEP V2 venv interpreter is unavailable"; return 1; } + for venv_site in "$venv"/lib/python*/site-packages; do break; done + [ -d "$venv_site" ] \ + || { collx_log "ERROR: DeepEP V2 venv site-packages is unavailable"; return 1; } + nccl_root="$(nvidia_package_root "$venv/bin/python" nvidia-nccl-cu13 nccl)" \ + || { collx_log "ERROR: DeepEP V2 NCCL package root is unavailable"; return 1; } + nvshmem_package="$(nvidia_package_root \ + "$venv/bin/python" nvidia-nvshmem-cu12 nvshmem)" \ + || { collx_log "ERROR: DeepEP V2 NVSHMEM package root is unavailable"; return 1; } + overlay="$(deepep_nvshmem_overlay "$root" "$nvshmem_package")" || return 1 + toolchain="$(cuda_toolchain_paths)" || return 1 + IFS=$'\t' read -r cuda_home cccl <<< "$toolchain" + [ -n "$cuda_home" ] && [ -n "$cccl" ] || return 1 + execution_id="${COLLECTIVEX_EXECUTION_ID:-manual}" + [[ "$execution_id" =~ ^[A-Za-z0-9._-]+$ ]] \ + || { collx_log "ERROR: DeepEP V2 execution identity is invalid"; return 1; } + + export \ + VIRTUAL_ENV="$venv" \ + PATH="$venv/bin:${PATH#"$venv/bin:"}" \ + PYTHONPATH="$venv_site${PYTHONPATH:+:$PYTHONPATH}" \ + CUDA_HOME="$cuda_home" \ + CPATH="$cccl:${CPATH:-}" \ + NVCC_PREPEND_FLAGS="-I$cccl ${NVCC_PREPEND_FLAGS:-}" \ + NVSHMEM_DIR="$overlay" \ + EP_NCCL_ROOT_DIR="$nccl_root" \ + EP_NVSHMEM_ROOT_DIR="$overlay" \ + EP_JIT_CACHE_DIR="/tmp/collectivex-deepep-v2-jit-$execution_id" \ + EP_REUSE_NCCL_COMM=1 \ + NCCL_CUMEM_ENABLE=1 \ + LD_LIBRARY_PATH="$overlay/lib:$nccl_root/lib:$nvshmem_package/lib:${LD_LIBRARY_PATH:-}" + unset "${DEEPEP_RANK_UNSETS[@]}" + + # Shared JIT caches race across nodes; keep this cache node-local. CUMEM is + # persisted here too because image environment overrides launcher exports. + [ ! -L "$EP_JIT_CACHE_DIR" ] \ + || { collx_log "ERROR: DeepEP V2 JIT cache path is unsafe"; return 1; } + if ! mkdir -p "$EP_JIT_CACHE_DIR" || ! chmod 700 "$EP_JIT_CACHE_DIR"; then + collx_log "ERROR: DeepEP V2 JIT cache is unavailable" + return 1 + fi +} + +deepep_probe() { + "$VIRTUAL_ENV/bin/python" - <<'PY' +import inspect +import deep_ep +assert inspect.isclass(deep_ep.ElasticBuffer) +PY +} + +deepep_install() { + local root="$1" arch="$2" venv="$1/venv" source_dir="$1/source" + local -a pip + if [ -e "$root" ] || [ -L "$root" ]; then + rm -rf "$root" \ + || { collx_log "ERROR: incomplete DeepEP V2 cache-reset failed"; return 1; } + fi + mkdir -m 700 "$root" \ + || { collx_log "ERROR: DeepEP V2 cache-create failed"; return 1; } + python3 -m venv "$venv" \ + || { collx_log "ERROR: DeepEP V2 venv creation failed"; return 1; } + pip=("$venv/bin/python" -m pip install -q --disable-pip-version-check --no-input) + "${pip[@]}" \ + "pip==26.1.2" "setuptools==82.0.1" "wheel==0.47.0" "ninja==1.13.0" \ + "numpy==2.2.6" "nvidia-nvshmem-cu12==3.3.9" >&2 2>&1 \ + || { collx_log "ERROR: DeepEP V2 build-tool installation failed"; return 1; } + "${pip[@]}" --index-url https://download.pytorch.org/whl/cu130 \ + --extra-index-url https://pypi.org/simple "torch==2.10.0" >&2 2>&1 \ + || { collx_log "ERROR: torch 2.10.0+cu130 installation failed"; return 1; } + # Torch pins NCCL 2.28.9; ElasticBuffer requires 2.30.4. + "${pip[@]}" --force-reinstall --no-deps "nvidia-nccl-cu13==2.30.4" >&2 2>&1 \ + || { collx_log "ERROR: NCCL 2.30.4 installation failed"; return 1; } + deepep_activate "$root" \ + || { collx_log "ERROR: DeepEP V2 environment activation failed"; return 1; } + collx_materialize_deepep_source "$source_dir" \ + || { collx_log "ERROR: DeepEP V2 staged source is invalid"; return 1; } + (cd "$source_dir" && TORCH_CUDA_ARCH_LIST="$arch" MAX_JOBS=16 \ + "$venv/bin/python" -m pip install -q --no-build-isolation --no-deps \ + --force-reinstall .) >&2 2>&1 \ + || { collx_log "ERROR: DeepEP V2 build failed"; return 1; } + deepep_probe \ + || { collx_log "ERROR: DeepEP V2 import probe failed"; return 1; } + : > "$root/.ready" +} + +# ---- DeepEP lifecycle ------------------------------------------------------- + +deepep_prepare() { + local arch root venv source_dir ready lock_path + arch="$(cuda_arch)" || return 1 + root="$(deepep_cache_root "$arch")" || return 1 + venv="$root/venv"; source_dir="$root/source"; ready="$root/.ready" + lock_path="${root}.lock" + command -v flock >/dev/null || { collx_log "ERROR: flock is required for DeepEP V2"; return 1; } + mkdir -p "${root%/*}" || return 1 + collx_log "DeepEP V2: preparing PR #605 with upstream PR #630 and #640 fixes ($COLLX_DEEPEP_V2_COMMIT)" + if ! ( + [ ! -L "$lock_path" ] \ + || { collx_log "ERROR: DeepEP V2 cache lock is unsafe"; exit 1; } + (umask 077; : >> "$lock_path") && chmod 600 "$lock_path" \ + || { collx_log "ERROR: DeepEP V2 cache-lock-create failed"; exit 1; } + exec 9<>"$lock_path" \ + || { collx_log "ERROR: DeepEP V2 cache-lock-open failed"; exit 1; } + flock 9 \ + || { collx_log "ERROR: DeepEP V2 cache-lock-acquire failed"; exit 1; } + if [ ! -f "$ready" ] || [ ! -x "$venv/bin/python" ] || [ ! -d "$source_dir" ]; then + deepep_install "$root" "$arch" || exit 1 + fi + ); then + collx_log "ERROR: shared DeepEP V2 environment is incomplete" + return 1 + fi + deepep_activate "$root" || return 1 + deepep_probe || { collx_log "ERROR: DeepEP V2 shared runtime probe failed"; return 1; } + collx_log "DeepEP V2 ready ($COLLX_DEEPEP_V2_COMMIT, ElasticBuffer, NCCL Device API; LSA/Gin selected by adapter)" +} + +# ---- container boundary ---------------------------------------------------- + +write_rank_env() { + local root="$PWD/.collx_backend/env" node_id="${SLURM_NODEID:-0}" path temporary name + [[ "$node_id" =~ ^[0-9]+$ ]] || return 1 + mkdir -p "$root" || return 1 + chmod 700 "$root" || return 1 + temporary="$(mktemp "$root/.node-${node_id}.XXXXXX")" || return 1 + chmod 600 "$temporary" || { rm -f "$temporary"; return 1; } + for name in "${RANK_ENV_VARS[@]}"; do + if declare -p "$name" >/dev/null 2>&1; then + printf 'export %s=%q\n' "$name" "${!name}" >> "$temporary" \ + || { rm -f "$temporary"; return 1; } + fi + done + if [ "$COLLX_BENCH" = deepep-v2 ]; then + for name in "${DEEPEP_RANK_UNSETS[@]}"; do + printf 'unset %s\n' "$name" >> "$temporary" \ + || { rm -f "$temporary"; return 1; } + done + fi + path="$root/node-${node_id}.sh" + mv -f -- "$temporary" "$path" || { rm -f "$temporary"; return 1; } +} + +validate_container_network() { + local interface device rdma_name + local -a interfaces devices + if [ "${COLLX_NODES:-1}" -le 1 ] || [ "${COLLX_TRANSPORT:-}" = mnnvl ]; then + return 0 + fi + collx_restore_exact_hca_selector || return 1 + [ -n "${GLOO_SOCKET_IFNAME:-}" ] && [ -n "${NCCL_IB_HCA:-}" ] \ + || { collx_log "ERROR: scale-out network selectors are unavailable"; return 1; } + IFS=, read -r -a interfaces <<< "$GLOO_SOCKET_IFNAME" + for interface in "${interfaces[@]}"; do + [ -d "/sys/class/net/$interface" ] \ + || { collx_log "ERROR: configured scale-out socket interface is absent"; return 1; } + done + IFS=, read -r -a devices <<< "$NCCL_IB_HCA" + for device in "${devices[@]}"; do + device="${device#=}" + rdma_name="${device%%:*}" + [ -d "/sys/class/infiniband/$rdma_name" ] \ + || { collx_log "ERROR: configured scale-out RDMA device is absent"; return 1; } + done +} + +main() { + collx_apply_network_profile "${COLLX_NODES:-1}" "${COLLX_TRANSPORT:-}" || return 1 + validate_container_network || return 1 + case "$COLLX_BENCH" in + deepep-v2) deepep_prepare || return 1 ;; + mori) + python3 -c "import mori" \ + || { collx_log "ERROR: MoRI backend import failed"; return 1; } + ;; + *) + collx_log "ERROR: unknown backend preparation request" + return 1 + ;; + esac + write_rank_env +} + +rc=0; main || rc=$? +collx_log "backend preparation: bench=$COLLX_BENCH rc=$rc" +exit "$rc" diff --git a/experimental/CollectiveX/runtime/probe.py b/experimental/CollectiveX/runtime/probe.py new file mode 100644 index 0000000000..922a99682e --- /dev/null +++ b/experimental/CollectiveX/runtime/probe.py @@ -0,0 +1,124 @@ +#!/usr/bin/env python3 +"""Allocation and network checks used by CollectiveX launchers.""" + +from __future__ import annotations + +import argparse +import ctypes +import os +from pathlib import Path + + +def default_route_interface(route_path: Path = Path("/proc/net/route")) -> str: + for line in route_path.read_text().splitlines()[1:]: + fields = line.split() + if len(fields) >= 4 and fields[1] == "00000000" and int(fields[3], 16) & 1: + return fields[0] + return "" + + +def prepare_cache(parent_path: str) -> str: + path = Path(parent_path).resolve() / f".collectivex-backend-cache-{os.getuid()}" + path.mkdir(mode=0o700, exist_ok=True) + os.chmod(path, 0o700) + return str(path) + + +def validate_cuda_context(expected: int) -> None: + cuda = ctypes.CDLL("libcuda.so.1") + count = ctypes.c_int() + if cuda.cuInit(0) != 0 or cuda.cuDeviceGetCount(ctypes.byref(count)) != 0 or count.value != expected: + raise SystemExit(1) + + +def _emit(marker: str) -> None: + # collx_validate_network_profile_on_job (runtime/common.sh) greps these exact strings + # out of the per-node probe log to derive COLLX_SOCKET_IFNAME / COLLX_RDMA_LINK_LAYER and to + # diagnose failures. The marker vocabulary is a string contract with that function — + # keep the two halves in lockstep (see tests/test_runtime.py::NetworkProfileContract). + print(f"[collectivex-private] {marker}") + + +def _check_port(port_path: Path, ordinal: int, gid_index: str, profile: str): + # Return the port's link layer ("roce"/"infiniband") when it is active, carries a + # non-empty GID at the pinned index, and agrees with any already-seen link layer; + # otherwise emit the matching rdma-port-= marker and return None. + if not port_path.is_dir(): + _emit(f"rdma-port-{ordinal}=missing"); return None + state = port_path / "state" + if not state.is_file() or state.read_text().split()[:1] != ["4:"]: + _emit(f"rdma-port-{ordinal}=inactive"); return None + if gid_index: + gid = port_path / "gids" / gid_index + if not gid.is_file(): + _emit(f"rdma-port-{ordinal}=gid-missing"); return None + if not "".join(c for c in gid.read_text() if c not in ":0" and not c.isspace()): + _emit(f"rdma-port-{ordinal}=gid-empty"); return None + link = port_path / "link_layer" + if not link.is_file(): + _emit(f"rdma-port-{ordinal}=link-layer-missing"); return None + layer = {"Ethernet": "roce", "InfiniBand": "infiniband"}.get(link.read_text().strip()) + if layer is None: + _emit(f"rdma-port-{ordinal}=link-layer-invalid"); return None + if profile and profile != layer: + _emit(f"rdma-port-{ordinal}=link-layer-mixed"); return None + return layer + + +def validate_network_profile(socket_names: str, rdma_devices: str, gid_index: str, + sys_root: Path = Path("/sys"), + route_path: Path = Path("/proc/net/route")) -> None: + # Prove the operator-pinned scale-out fabric on this node: resolve the cross-node socket + # interface (operator selector, else this node's default route), confirm it is live, and + # confirm every pinned RDMA port is active with a consistent link layer. On success emit + # the socket-interface-selected and rdma-link-layer markers the launcher consumes; on any + # failure emit the diagnostic marker and exit non-zero. The seam carries exactly ONE socket + # interface: the launcher's marker-extraction regex has no comma, so a multi-interface + # selector could never survive past this probe — fail it loudly here instead. + interface = socket_names or default_route_interface(route_path) + if not interface: + _emit("socket-interface-1=default-route-missing") + raise SystemExit(1) + _emit(f"socket-interface-selected={interface}") + net = sys_root / "class" / "net" / interface + if not net.is_dir(): + _emit("socket-interface-1=missing"); raise SystemExit(1) + operstate = net / "operstate" + state = operstate.read_text().strip() if operstate.is_file() else "" + if state not in ("up", "unknown"): + _emit("socket-interface-1=down"); raise SystemExit(1) + profile = "" + for ordinal, selector in enumerate((s for s in rdma_devices.split(",") if s), start=1): + device, _, configured_port = selector.partition(":") + ports = sys_root / "class" / "infiniband" / device / "ports" + if not ports.is_dir(): + _emit(f"rdma-device-{ordinal}=missing"); raise SystemExit(1) + if configured_port: + layer = _check_port(ports / configured_port, ordinal, gid_index, profile) + if layer is None: raise SystemExit(1) + profile = layer + else: + active = False + for port_path in sorted(p for p in ports.iterdir() if p.is_dir()): + layer = _check_port(port_path, ordinal, gid_index, profile) + if layer is not None: + profile, active = layer, True + if not active: raise SystemExit(1) + if not profile: raise SystemExit(1) + _emit(f"rdma-link-layer={profile}") + + +def main() -> None: + parser = argparse.ArgumentParser(); commands = parser.add_subparsers(dest="command", required=True) + commands.add_parser("default-route-interface") + command = commands.add_parser("prepare-cache"); command.add_argument("parent") + command = commands.add_parser("cuda-context"); command.add_argument("expected", type=int) + command = commands.add_parser("network-profile"); command.add_argument("socket_names"); command.add_argument("rdma_devices"); command.add_argument("gid_index") + args = parser.parse_args() + if args.command == "default-route-interface": print(default_route_interface(), end="") + elif args.command == "prepare-cache": print(prepare_cache(args.parent), end="") + elif args.command == "cuda-context": validate_cuda_context(args.expected) + else: validate_network_profile(args.socket_names, args.rdma_devices, args.gid_index) + + +if __name__ == "__main__": main() diff --git a/experimental/CollectiveX/runtime/stage.py b/experimental/CollectiveX/runtime/stage.py new file mode 100644 index 0000000000..119b77496f --- /dev/null +++ b/experimental/CollectiveX/runtime/stage.py @@ -0,0 +1,103 @@ +#!/usr/bin/env python3 +"""Create, copy, and clean isolated CollectiveX workspaces.""" + +from __future__ import annotations + +import argparse +import os +import pwd +from pathlib import Path +import shutil + + +EXCLUDES = {"__pycache__", "results", ".shards", ".collx_workloads", ".collx_backend", + ".collx_sources", ".venv", ".pytest_cache", "private-infra.md", "goal.md", + "notes.md"} + + +def implicit_stage_base(args) -> None: + # Resolve the account home from /etc/passwd, not $HOME. The GHA launcher deliberately + # points $HOME at a runner-local /tmp sandbox. The passwd home is compute-visible. + base = args.home or pwd.getpwuid(os.getuid()).pw_dir + home = Path(base).resolve() + suffix = "" + if args.isolation_key: + if not all(char.isalnum() or char in "._-" for char in args.isolation_key): + raise SystemExit(1) + suffix = "-" + args.isolation_key + path = home / f".inferencex-collectivex-stage{suffix}" + path.mkdir(mode=0o700, exist_ok=True) + print(path, end="") + + +def resolve_directory(args) -> None: + path = Path(args.path).resolve() + if not path.is_dir(): raise SystemExit(1) + print(path, end="") + + +def validate_stage_path(args) -> None: + base, child = Path(args.base).resolve(), Path(args.child) + if child.parent.resolve() != base or child.exists() or base == Path("/"): + raise SystemExit(1) + for excluded in (args.repo, args.job_root, args.workspace): + if excluded and base == Path(excluded).resolve(): raise SystemExit(1) + print(child, end="") + + +def create_stage(args) -> None: + Path(args.stage).mkdir(mode=0o700) + + +def copy_repository(args) -> None: + source, target = Path(args.source), Path(args.target) + shutil.copytree(source, target, ignore=shutil.ignore_patterns(*EXCLUDES), dirs_exist_ok=False) + + +def validate_cleanup(args) -> None: + root = Path(args.root) + if not root.is_dir() or root.is_symlink() or root == Path("/"): + raise SystemExit(1) + + +def rewrite_deepep_v2(args) -> None: + path = Path(args.path) + old = "for so in [line.strip().split(' ')[-1] for line in f if 'nccl' in line]:" + new = "for so in [line.strip().split(' ')[-1] for line in f if 'libnccl' in line]:" + text = path.read_text() + if text.count(old) != 1: raise SystemExit(1) + path.write_text(text.replace(old, new)) + + +# The runtime/common.sh launcher shells out to these subcommands by literal name and +# positional argv; there are no optional flags. That argv shape is a string contract with +# common.sh — a subcommand or flag common.sh passes but this parser does not declare fails +# with "unrecognized arguments" and aborts the leg at repository-stage. Keep the two halves in +# lockstep (see tests/test_runtime.py::StageContract), which is exactly the contract that broke +# when the --allow-* flags were dropped here but left on the common.sh callers. +SPECS = { + "implicit-stage-base": (("home", "?"), ("isolation_key", "?")), + "resolve-directory": (("path",),), + "validate-stage-path": (("repo",), ("base",), ("child",), ("job_root", "?"), ("workspace", "?")), + "create-stage": (("stage",),), "copy-repository": (("source",), ("target",)), + "validate-cleanup": (("root",),), "rewrite-deepep-v2": (("path",),), +} + + +def build_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser() + commands = parser.add_subparsers(dest="command", required=True) + handlers = globals() + for name, arguments in SPECS.items(): + command = commands.add_parser(name) + for item in arguments: + command.add_argument(item[0], nargs=item[1] if len(item) > 1 else None, default="") + command.set_defaults(handler=handlers[name.replace("-", "_")]) + return parser + + +def main() -> None: + args = build_parser().parse_args(); args.handler(args) + + +if __name__ == "__main__": main() diff --git a/experimental/CollectiveX/summarize.py b/experimental/CollectiveX/summarize.py new file mode 100644 index 0000000000..85c7048e29 --- /dev/null +++ b/experimental/CollectiveX/summarize.py @@ -0,0 +1,92 @@ +#!/usr/bin/env python3 +"""Render a small shard summary; benchmark gating remains in the harness.""" +from __future__ import annotations + +import argparse +import json +from pathlib import Path + +# Emitted case-attempt documents this summary reads, discriminated by record_type. +# This is a best-effort renderer over whatever raw attempts a shard produced; it +# validates nothing. +CASE_RECORD_TYPE = "case-attempt" + + +def load_results(directory: str, runner: str | None, timestamp: str | None) -> list[dict]: + documents: list[dict] = [] + for path in sorted(Path(directory).glob("*.json")): + if runner and not path.name.startswith(f"{runner}_"): + continue + if timestamp and timestamp not in path.name: + continue + try: + with path.open() as handle: + document = json.load(handle) + except (OSError, ValueError): + continue + if isinstance(document, dict) and document.get("record_type") == CASE_RECORD_TYPE: + documents.append(document) + return documents + + +def _identity(document: dict) -> tuple[str, str, str, str, int]: + factors = document["identity"]["case_factors"] + case = factors["case"] + return ( + factors["sku"], case["suite"], case["routing"], case["phase"], case["ep"], + ) + + +def _headline(document: dict) -> tuple[int | str, float | str, float | str]: + rows = document["measurement"]["rows"] + row = next((item for item in rows if item["tokens_per_rank"] == 64), rows[len(rows) // 2]) + latency = row["components"]["roundtrip"]["percentiles_us"] + return row["tokens_per_rank"], latency["p50"], latency["p99"] + + +def render(documents: list[dict]) -> str: + documents = sorted(documents, key=_identity) + invalid = [d for d in documents if d["outcome"]["status"] != "success"] + lines = ["## CollectiveX EP results", ""] + if invalid: + # The leg is already red (ep_harness.run_sweep returns nonzero on a non-success + # outcome); call the count out loudly so it is not lost in the per-row table. + lines.append( + f"> **{len(invalid)} of {len(documents)} outcome(s) INVALID** — " + "the leg fails; see the outcome column below." + ) + lines.append("") + lines += [ + "| ver | sku | backend | suite | phase | routing | ep | outcome | T* | p50 us | p99 us |", + "|--:|---|---|---|---|---|--:|---|--:|--:|--:|", + ] + for document in documents: + sku, suite, routing, phase, ep = _identity(document) + backend = document["identity"]["case_factors"]["case"]["backend"] + token, p50, p99 = _headline(document) + lines.append( + f"| {document['version']} | {sku} | `{backend}` | {suite} | {phase} | " + f"{routing} | {ep} | " + f"{document['outcome']['status']} | {token} | {p50} | {p99} |" + ) + if not documents: + lines.append("\n> No valid native outcome documents found.") + return "\n".join(lines) + + +def main() -> int: + parser = argparse.ArgumentParser(description="Summarize CollectiveX native v1 outcomes") + parser.add_argument("--results-dir", default="results") + parser.add_argument("--runner") + parser.add_argument("--ts") + args = parser.parse_args() + documents = load_results(args.results_dir, args.runner, args.ts) + print(render(documents)) + # Pure renderer — never gates CI. The per-case leg gate lives in ep_harness.run_sweep: a + # non-success outcome returns nonzero and fails the shard (see collx_run_shard). A dead + # "exit 1 when no success doc" gate here lost its only caller in 41caeaa0. + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/experimental/CollectiveX/sweep_matrix.py b/experimental/CollectiveX/sweep_matrix.py new file mode 100644 index 0000000000..72c9dcdb98 --- /dev/null +++ b/experimental/CollectiveX/sweep_matrix.py @@ -0,0 +1,223 @@ +#!/usr/bin/env python3 +"""Build the CollectiveX sweep matrix and extract execution shards.""" +from __future__ import annotations + +import argparse +import json +from pathlib import Path +import sys +from typing import Any + +HERE = Path(__file__).resolve().parent +sys.path.insert(0, str(HERE / "bench")) + +import ep_harness # noqa: E402 + + +TOPOLOGY_FIELDS = ( + "nodes", "gpus_per_node", "scale_up_domain", "scope", "scale_up_transport", + "scale_out_transport", "transport", "topology_class", +) + + +def _load_config(name: str) -> dict[str, Any]: + return json.loads((HERE / "configs" / name).read_text(encoding="utf-8")) + + +SWEEP = _load_config("sweep.json") +PLATFORMS = _load_config("platform_config.json")["platforms"] +SWEEP_BACKENDS = tuple(dict.fromkeys( + backend for platform in PLATFORMS.values() for backend in platform["backends"] +)) + + +def _topology(platform: dict[str, Any], ep: int) -> dict[str, Any]: + gpus_per_node = platform["gpus_per_node"] + if ep % gpus_per_node: + raise SystemExit(f"EP{ep} is not divisible by {gpus_per_node} GPUs per node") + product = platform["product"] + domain = platform["scale_up_domain"] + scale_up = platform["scale_up_transport"] + scale_out = ep > domain + if scale_up == "mnnvl": + scale_up_class = f"{product}-nvl{domain}-mnnvl" + elif scale_up == "xgmi": + scale_up_class = f"{product}-xgmi" + else: + scale_up_class = f"{product}-{scale_up}-island" + return { + "nodes": ep // gpus_per_node, + "gpus_per_node": gpus_per_node, + "scale_up_domain": domain, + "scope": "scale-out" if scale_out else "scale-up", + "scale_up_transport": scale_up, + "scale_out_transport": "rdma" if scale_out else None, + "transport": f"{scale_up}-rdma" if scale_out else scale_up, + "topology_class": f"{product}-{scale_up}-rdma" if scale_out else scale_up_class, + } + + +def _selected_backends(backend: str) -> list[str]: + if backend == "all": + return list(SWEEP_BACKENDS) + if backend not in SWEEP_BACKENDS: + raise SystemExit(f"unknown --backend {backend!r}; have {list(SWEEP_BACKENDS)}") + return [backend] + + +def resolve_matrix( + backend: str = "all", + only_sku: str = "", + exclude_skus: str = "", + ep_sizes: str = "", +) -> dict[str, Any]: + """Resolve the fixed sweep into allocation-sized workflow shards.""" + selected_eps: set[int] = set() + for value in filter(None, (part.strip() for part in ep_sizes.split(","))): + if not value.isdigit() or int(value) <= 0: + raise SystemExit(f"invalid --ep-sizes {ep_sizes!r}; expected positive integers") + selected_eps.add(int(value)) + + if only_sku and only_sku not in PLATFORMS: + raise SystemExit(f"unknown --only-sku {only_sku!r}; have {sorted(PLATFORMS)}") + excluded = {value.strip() for value in exclude_skus.split(",") if value.strip()} + unknown = sorted(excluded - set(PLATFORMS)) + if unknown: + raise SystemExit(f"unknown --exclude-skus {unknown}; have {sorted(PLATFORMS)}") + if only_sku in excluded: + raise SystemExit("--only-sku and --exclude-skus select disjoint pools") + + timing = SWEEP["timing"] + timing_profile = ":".join(str(timing[key]) for key in ( + "iters_per_trial", "trials_per_point", "warmup_iters_per_trial", + )) + workload = SWEEP["workload"] + targets = _selected_backends(backend) + requested_cases: list[dict[str, Any]] = [] + shards: dict[tuple[str, str, int], list[dict[str, Any]]] = {} + + for sku in sorted(PLATFORMS): + if (only_sku and sku != only_sku) or sku in excluded: + continue + platform = PLATFORMS[sku] + for ep in SWEEP["ep_degrees"]: + if selected_eps and ep not in selected_eps: + continue + topology = _topology(platform, ep) + for phase, ladder in workload["token_ladders"].items(): + for target in targets: + runnable_eps = platform["backends"].get(target) + if runnable_eps is None: + continue + runnable = ep in runnable_eps + case = { + "suite": SWEEP["suite"], + "workload": workload["name"], + "backend": target, + "routing": SWEEP["routing"], + "phase": phase, + "ep": ep, + "hidden": workload["hidden"], + "topk": workload["topk"], + "experts": workload["routed_experts"], + "seed": workload["seed"], + "ladder": " ".join(map(str, ladder)), + "mode": SWEEP["mode"], + "timing": timing_profile, + **{field: topology[field] for field in TOPOLOGY_FIELDS}, + } + case["case_id"] = ep_harness.case_id(sku, case) + requested_cases.append({ + "sku": sku, + "case": case, + "disposition": "runnable" if runnable else "unsupported", + "reason": None if runnable else "backend-platform-unsupported", + "detail": None, + }) + if runnable: + shards.setdefault((sku, target, topology["nodes"]), []).append(case) + + shards_by_sku: dict[str, list[dict[str, Any]]] = {} + for (sku, target, nodes), cases in sorted(shards.items()): + first = cases[0] + shards_by_sku.setdefault(sku, []).append({ + "id": f"{sku}-{target}-n{nodes}", + "sku": sku, + "backend": target, + "launcher": PLATFORMS[sku]["launcher"], + "nodes": nodes, + "gpus_per_node": first["gpus_per_node"], + "scale_up_domain": first["scale_up_domain"], + "cases": cases, + }) + include = [ + shards_by_sku[sku][index] + for index in range(max(map(len, shards_by_sku.values()), default=0)) + for sku in sorted(shards_by_sku) + if index < len(shards_by_sku[sku]) + ] + return { + "version": SWEEP["version"], + "requested_cases": requested_cases, + "include": include, + } + + +def extract_shard(matrix_path: str, shard_id: str, output_path: str) -> dict[str, Any]: + """Write one generator-produced shard as a runner control document.""" + document = json.loads(Path(matrix_path).read_text(encoding="utf-8")) + matches = [item for item in document["include"] if item["id"] == shard_id] + if len(matches) != 1: + raise SystemExit(f"expected one shard {shard_id!r}, found {len(matches)}") + source = matches[0] + control = {key: source[key] for key in ("id", "sku", "backend", "nodes", "cases")} + control["version"] = document["version"] + output = Path(output_path) + output.parent.mkdir(parents=True, exist_ok=True) + output.write_text(json.dumps(control, sort_keys=True, separators=(",", ":")) + "\n") + return control + + +def main() -> int: + parser = argparse.ArgumentParser(description="CollectiveX matrix resolver") + parser.add_argument("--backend", default="all") + parser.add_argument("--only-sku", default="") + parser.add_argument("--exclude-skus", default="") + parser.add_argument("--ep-sizes", default="") + parser.add_argument("--extract-from", default="", metavar="MATRIX") + parser.add_argument("--shard-id", default="") + parser.add_argument("--out", default="") + args = parser.parse_args() + + if args.extract_from: + if not all((args.shard_id, args.out)): + parser.error("shard extraction requires --shard-id and --out") + control = extract_shard(args.extract_from, args.shard_id, args.out) + print(f"extracted {control['id']}: {len(control['cases'])} cases", file=sys.stderr) + print(json.dumps(control, separators=(",", ":"))) + return 0 + + matrix = resolve_matrix( + backend=args.backend, + only_sku=args.only_sku, + exclude_skus=args.exclude_skus, + ep_sizes=args.ep_sizes, + ) + if args.out: + Path(args.out).write_text( + json.dumps(matrix, sort_keys=True, separators=(",", ":")) + "\n", + encoding="utf-8", + ) + runnable = sum(item["disposition"] == "runnable" for item in matrix["requested_cases"]) + unsupported = len(matrix["requested_cases"]) - runnable + print( + f"resolved {len(matrix['include'])} shard-cells, " + f"{runnable} runnable and {unsupported} unsupported cases", + file=sys.stderr, + ) + print(json.dumps(matrix)) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/experimental/CollectiveX/tests/test_ep_backend.py b/experimental/CollectiveX/tests/test_ep_backend.py new file mode 100644 index 0000000000..2f2077c9f2 --- /dev/null +++ b/experimental/CollectiveX/tests/test_ep_backend.py @@ -0,0 +1,143 @@ +#!/usr/bin/env python3 +"""Small torch-free smoke tests for the shared EP backend lifecycle.""" +from __future__ import annotations + +import sys +import types +import unittest +from pathlib import Path +from unittest import mock + +ROOT = Path(__file__).resolve().parents[1] +sys.path[:0] = [str(ROOT), str(ROOT / "bench")] + +import ep_backend # noqa: E402 +from ep_backend import EPBackend, RankInputs # noqa: E402 + + +def args(**updates): + values = dict( + experts=8, phase="decode", tokens_ladder="", routing="uniform", seed=0, + hidden=16, topk=2, mode="normal", + ) + values.update(updates) + return types.SimpleNamespace(**values) + + +class FakeBackend(EPBackend): + name = "fake" + + def __init__(self, options, *, cap=None, world_size=1): + super().__init__(options, 0, world_size, 0, "cpu") + self.cap = cap + self.calls: list[str] = [] + + def create_buffer(self, spec): + return None + + def dispatch(self, problem): + self.calls.append("dispatch") + return object() + + def stage(self, problem, handle): + self.calls.append("stage") + + def combine(self, problem, handle): + self.calls.append("combine") + + def recv_tokens(self, handle): + return 0 + + def inspect_dispatch(self, problem, handle): + return None + + def combine_transformed(self, problem, handle, transformed): + return None + + def buffer_cap(self, options): + return self.cap + + def _build_rank_inputs(self, options, tokens): + return RankInputs( + tokens_per_rank=tokens, topk_idx=None, topk_weights=None, + activations=None, + ) + + +class BackendTests(unittest.TestCase): + def test_input_plan_sizes_for_the_measured_ladder(self): + backend = FakeBackend(args(tokens_ladder="8 16"), world_size=2) + spec = backend.make_inputs(backend.args) + self.assertTrue(spec.ok) + self.assertEqual(spec.ladder, [8, 16]) + self.assertEqual(spec.max_tokens_per_rank, 16) + self.assertEqual((spec.ep_size, spec.experts_per_rank), (2, 4)) + self.assertEqual(sorted(spec.points), [8, 16]) + + def test_invalid_or_fully_clamped_ladder_fails_before_execution(self): + for backend, message in ( + (FakeBackend(args(tokens_ladder="0")), "empty token ladder"), + (FakeBackend(args(tokens_ladder="128"), cap=64), "cap=64"), + ): + with self.subTest(message=message): + spec = backend.make_inputs(backend.args) + self.assertEqual(spec.rc, 2) + self.assertIn(message, spec.message) + + def test_timed_components_follow_backend_contract(self): + backend = FakeBackend(args()) + self.assertEqual(backend.timed_components(), ["roundtrip", "dispatch", "combine"]) + backend.stage_device_work = True + self.assertEqual( + backend.timed_components(), ["roundtrip", "dispatch", "combine", "stage"] + ) + backend.roundtrip_only = True + self.assertEqual(backend.timed_components(), ["roundtrip"]) + + def test_dispatch_cleanup_is_outside_timed_call(self): + backend = FakeBackend(args()) + backend.dispatch_needs_combine_cleanup = True + captured = {} + + def fake_time(_torch, operation, _warmup, _iters, **kwargs): + handle = operation() + kwargs["post"](handle) + captured.update(kwargs) + return [1.0] + + with mock.patch.dict(sys.modules, {"torch": types.SimpleNamespace()}), mock.patch.object( + ep_backend, "time_us", side_effect=fake_time + ): + backend.benchmark_dispatch(object(), 0, 1) + self.assertIn("post", captured) + self.assertEqual(backend.calls, ["dispatch", "stage", "combine"]) + + def test_stage_cleanup_matches_the_dispatch_contract(self): + # MoRI-shaped backends (dispatch_needs_combine_cleanup) must not leak an + # un-combined dispatch out of an isolated-stage iteration. + for needs_cleanup, calls in ( + (True, ["dispatch", "stage", "combine"]), (False, ["dispatch", "stage"]), + ): + backend = FakeBackend(args()) + backend.dispatch_needs_combine_cleanup = needs_cleanup + + def fake_time(_torch, operation, _warmup, _iters, **kwargs): + result = operation(kwargs["pre"]()) + if kwargs["post"] is not None: + kwargs["post"](result) + return [1.0] + + with mock.patch.dict(sys.modules, {"torch": types.SimpleNamespace()}), mock.patch.object( + ep_backend, "time_us", side_effect=fake_time + ): + backend.benchmark_stage(object(), 0, 1) + with self.subTest(needs_cleanup=needs_cleanup): + self.assertEqual(backend.calls, calls) + + def test_mode_is_fail_closed(self): + with self.assertRaises(ValueError): + FakeBackend(args(mode="unsupported")) + + +if __name__ == "__main__": + unittest.main() diff --git a/experimental/CollectiveX/tests/test_matrix.py b/experimental/CollectiveX/tests/test_matrix.py new file mode 100644 index 0000000000..4166993007 --- /dev/null +++ b/experimental/CollectiveX/tests/test_matrix.py @@ -0,0 +1,81 @@ +#!/usr/bin/env python3 +"""Matrix, subset, and shard-extraction tests.""" +from __future__ import annotations + +import json +import sys +import tempfile +import unittest +from pathlib import Path + +ROOT = Path(__file__).resolve().parents[1] +sys.path.insert(0, str(ROOT)) + +import sweep_matrix # noqa: E402 + + +def matrix(**options): + return sweep_matrix.resolve_matrix(**options) + + +class MatrixTests(unittest.TestCase): + def test_shard_extraction_is_deterministic_and_preserves_cases(self): + document = matrix(backend="deepep-v2", only_sku="h200-dgxc") + cell = document["include"][0] + with tempfile.TemporaryDirectory() as temporary: + root = Path(temporary) + source = root / "matrix.json" + source.write_text(json.dumps(document, sort_keys=True)) + outputs = [ + sweep_matrix.extract_shard( + source, cell["id"], root / f"shard-{index}.json", + ) + for index in range(2) + ] + self.assertEqual(outputs[0], outputs[1]) + self.assertEqual(outputs[0]["cases"], cell["cases"]) + + def test_sku_and_ep_filters_only_remove_cases(self): + full = matrix(backend="all") + for options, keep in ( + ({"exclude_skus": "b300"}, lambda item: item["sku"] != "b300"), + ({"ep_sizes": "8"}, lambda item: item["case"]["ep"] == 8), + ): + partial = matrix(backend="all", **options) + expected = { + item["case"]["case_id"]: item for item in full["requested_cases"] if keep(item) + } + actual = {item["case"]["case_id"]: item for item in partial["requested_cases"]} + self.assertEqual(actual, expected) + + def test_only_real_platform_cells_are_unsupported(self): + document = matrix(backend="all") + unsupported = { + (item["sku"], item["case"]["backend"], item["case"]["ep"]) + for item in document["requested_cases"] if item["disposition"] == "unsupported" + } + expected = { + (sku, backend, ep) + for sku, platform in sweep_matrix.PLATFORMS.items() + for backend, runnable_eps in platform["backends"].items() + for ep in sweep_matrix.SWEEP["ep_degrees"] + if ep not in runnable_eps + } + self.assertEqual(unsupported, expected) + for item in document["requested_cases"]: + self.assertIn(item["case"]["backend"], sweep_matrix.PLATFORMS[item["sku"]]["backends"]) + + def test_invalid_filters_fail_closed(self): + for options in ( + {"exclude_skus": "unknown"}, + {"only_sku": "b300", "exclude_skus": "b300"}, + {"ep_sizes": "0"}, + {"ep_sizes": "eight"}, + {"backend": "unknown"}, + ): + with self.subTest(options=options), self.assertRaises(SystemExit): + sweep_matrix.resolve_matrix(**options) + + +if __name__ == "__main__": + unittest.main() diff --git a/experimental/CollectiveX/tests/test_runtime.py b/experimental/CollectiveX/tests/test_runtime.py new file mode 100644 index 0000000000..0347e16ca7 --- /dev/null +++ b/experimental/CollectiveX/tests/test_runtime.py @@ -0,0 +1,384 @@ +#!/usr/bin/env python3 +"""Focused tests for the standalone runtime helpers.""" + +from __future__ import annotations + +import argparse +import contextlib +import io +import json +import os +from pathlib import Path +import re +import subprocess +import sys +import tempfile +import unittest + + +RUNTIME = Path(__file__).resolve().parents[1] / "runtime" +BENCH = Path(__file__).resolve().parents[1] / "bench" +sys.path.insert(0, str(RUNTIME)) +sys.path.insert(0, str(BENCH)) + +import probe # noqa: E402 +import config # noqa: E402 +import stage # noqa: E402 +import ep_harness # noqa: E402 (stdlib-only at module top) + + +# configs/platform_config.json is shared by matrix scheduling, operator/network +# loading, and backend builds. +class PlatformRegistryTests(unittest.TestCase): + REGISTRY = RUNTIME.parent / "configs" / "platform_config.json" + NETWORK_FIELDS = { + "socket_ifname", "rdma_devices", "ib_gid_index", + "rdma_service_level", "rdma_traffic_class", "rail_isolated", + } + + def test_every_platform_entry_is_complete_and_typed(self) -> None: + platforms = json.loads(self.REGISTRY.read_text())["platforms"] + self.assertTrue(platforms) + for name, entry in platforms.items(): + with self.subTest(sku=name): + for field in ( + "arch", "product", "image", "image_platform", + "scale_up_transport", "launcher", + ): + self.assertIsInstance(entry[field], str) + self.assertTrue(entry[field]) + for field in ("gpus_per_node", "scale_up_domain"): + self.assertIsInstance(entry[field], int) + self.assertGreater(entry[field], 0) + self.assertTrue(entry["backends"]) + for degrees in entry["backends"].values(): + self.assertTrue(degrees) + self.assertLessEqual(set(degrees), {8, 16}) + self.assertLessEqual( + set(entry.get("network", {})), self.NETWORK_FIELDS + ) + # Fabric provenance: each cluster records its scale-out NIC and + # switch so same-GPU clusters on different fabrics stay distinct. + fabric = entry["fabric"] + self.assertEqual(set(fabric), {"nic", "switch"}) + for value in fabric.values(): + self.assertIsInstance(value, str) + self.assertTrue(value) + self.assertRegex(entry["arch"], r"^(sm|gfx)\d+$") + self.assertRegex(entry["image"], r"^[A-Za-z0-9._/-]+:[A-Za-z0-9._-]+$") + self.assertIn(entry["image_platform"], {"linux/amd64", "linux/arm64"}) + + +class ProbeTests(unittest.TestCase): + def test_default_route_interface(self) -> None: + with tempfile.TemporaryDirectory() as directory: + route = Path(directory) / "route" + route.write_text( + "Iface Destination Gateway Flags RefCnt Use Metric Mask MTU Window IRTT\n" + "eth9 00000000 00000000 0003 0 0 0 00000000 0 0 0\n" + ) + self.assertEqual(probe.default_route_interface(route), "eth9") + + def test_prepare_cache_is_private_and_reusable(self) -> None: + with tempfile.TemporaryDirectory() as directory: + first = Path(probe.prepare_cache(directory)) + second = Path(probe.prepare_cache(directory)) + self.assertEqual(first, second) + self.assertEqual(first.stat().st_mode & 0o777, 0o700) + + +class ConfigTests(unittest.TestCase): + def test_operator_config_emits_allowlisted_values(self) -> None: + with tempfile.TemporaryDirectory() as directory: + path = Path(directory) / "operator.json" + path.write_text(json.dumps({ + "runners": { + "h100-dgxc": { + "partition": "gpu", + "account": "bench", + "squash_dir": directory, + } + }, + })) + path.chmod(0o600) + read_fd, write_fd = os.pipe() + stdout = sys.stdout + try: + sys.stdout = os.fdopen(write_fd, "w") + config.operator_config(str(path), "h100-dgxc") + sys.stdout.flush() + finally: + sys.stdout.close() + sys.stdout = stdout + payload = os.read(read_fd, 4096) + os.close(read_fd) + self.assertIn(b"COLLX_PARTITION\0gpu\0", payload) + self.assertIn(b"COLLX_SQUASH_DIR\0" + directory.encode() + b"\0", payload) + self.assertIn(b"COLLX_IMAGE\0lmsysorg/sglang:v0.5.11-cu130\0", payload) + self.assertIn(b"COLLX_IMAGE_PLATFORM\0linux/amd64\0", payload) + + def _emit_registry_only(self, runner: str) -> bytes: + read_fd, write_fd = os.pipe() + stdout = sys.stdout + try: + sys.stdout = os.fdopen(write_fd, "w") + config.operator_config("-", runner) + sys.stdout.flush() + finally: + sys.stdout.close() + sys.stdout = stdout + payload = os.read(read_fd, 4096) + os.close(read_fd) + return payload + + def test_operator_config_registry_only_emits_tracked_baseline(self) -> None: + # "-" = no operator document: the registry's per-SKU operator block is + # the tracked baseline (plus its network overlay where present). + payload = self._emit_registry_only("h200-dgxc") + self.assertIn(b"COLLX_PARTITION\0main\0", payload) + self.assertIn(b"COLLX_SQUASH_DIR\0/home/sa-shared/containers\0", payload) + self.assertIn(b"COLLX_RDMA_DEVICES\0", payload) + + def test_operator_config_registry_only_emits_image_for_secret_fed_sku(self) -> None: + # A SKU without tracked operator settings still gets its public image + # configuration; private scheduler values can arrive through the overlay. + payload = self._emit_registry_only("mi325x") + self.assertIn(b"COLLX_IMAGE\0rocm/sgl-dev:sglang-0.5.14-rocm720-mi35x-mori-0701\0", payload) + self.assertIn(b"COLLX_IMAGE_PLATFORM\0linux/amd64\0", payload) + + +class StageTests(unittest.TestCase): + def test_create_copy_and_validate_cleanup(self) -> None: + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + source = root / "source" + target = root / "stage" + (source / "runtime").mkdir(parents=True) + (source / "runtime" / "common.sh").write_text("test") + (source / "goal.md").write_text("private") + (source / ".shards").mkdir() + (source / ".shards" / "leg.json").write_text("{}") + args = type("Args", (), {"stage": str(target)}) + stage.create_stage(args) + copy_args = type( + "Args", (), {"source": str(source), "target": str(target / "experimental" / "CollectiveX")} + ) + stage.copy_repository(copy_args) + staged = target / "experimental" / "CollectiveX" + self.assertTrue((staged / "runtime" / "common.sh").is_file()) + self.assertFalse((staged / ".shards").exists()) + self.assertFalse((staged / "goal.md").exists()) + cleanup_args = type("Args", (), {"root": str(target)}) + stage.validate_cleanup(cleanup_args) + +# The per-node probe (runtime/probe.py) and the launcher gate +# (runtime/common.sh: collx_validate_network_profile_on_job) share an implicit string contract: +# the probe prints these markers, the launcher greps them back out to derive COLLX_SOCKET_IFNAME +# and COLLX_RDMA_LINK_LAYER. The patterns are duplicated here on purpose — the test fails if +# either side drifts, which is exactly the failure that slipped through when 5506c623 moved the +# probe into Python but left the emit statements behind, silently zeroing the marker count for +# every non-MNNVL multi-node leg. +SOCKET_MARKER = r"^\[collectivex-private\] socket-interface-selected=([A-Za-z][A-Za-z0-9_.-]{0,31})$" +LINK_MARKER = r"^\[collectivex-private\] rdma-link-layer=(roce|infiniband)$" +FAILURE_MARKER = ( + r"(socket-interface|rdma-(device|port))-[0-9]+=" + r"(missing|down|inactive|default-route-missing|gid-missing|gid-empty|" + r"link-layer-missing|link-layer-invalid|link-layer-mixed)" +) + + +class NetworkProfileContract(unittest.TestCase): + def _fabric(self, root: Path, *, state: str = "4: ACTIVE", + link_layer: str = "Ethernet", gid: str = "fe80::1") -> None: + net = root / "class" / "net" / "eth0" + net.mkdir(parents=True) + (net / "operstate").write_text("up\n") + port = root / "class" / "infiniband" / "mlx5_0" / "ports" / "1" + (port / "gids").mkdir(parents=True) + (port / "state").write_text(state + "\n") + (port / "link_layer").write_text(link_layer + "\n") + (port / "gids" / "3").write_text(gid + "\n") + + def _run(self, root: Path, route: Path, socket_names: str = "eth0"): + buffer = io.StringIO() + rc = 0 + try: + with contextlib.redirect_stdout(buffer): + probe.validate_network_profile(socket_names, "mlx5_0:1", "3", + sys_root=root, route_path=route) + except SystemExit: + rc = 1 + return rc, buffer.getvalue().splitlines() + + @staticmethod + def _captures(pattern: str, lines: list) -> list: + return [match.group(1) for line in lines + for match in [re.match(pattern, line)] if match] + + def test_launcher_still_declares_the_marker_patterns(self) -> None: + common = (RUNTIME / "common.sh").read_text() + self.assertIn(SOCKET_MARKER, common) + self.assertIn(LINK_MARKER, common) + self.assertIn(FAILURE_MARKER, common) + + def test_healthy_fabric_emits_the_success_markers_the_launcher_extracts(self) -> None: + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + self._fabric(root) + rc, lines = self._run(root, root / "route") + self.assertEqual(rc, 0) + self.assertEqual(self._captures(SOCKET_MARKER, lines), ["eth0"]) + self.assertEqual(self._captures(LINK_MARKER, lines), ["roce"]) + + def test_infiniband_link_layer_maps_to_the_launcher_token(self) -> None: + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + self._fabric(root, link_layer="InfiniBand") + rc, lines = self._run(root, root / "route") + self.assertEqual(rc, 0) + self.assertEqual(self._captures(LINK_MARKER, lines), ["infiniband"]) + + def test_socket_interface_resolves_from_default_route(self) -> None: + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + self._fabric(root) + route = root / "route" + route.write_text( + "Iface Destination Gateway Flags RefCnt Use Metric Mask MTU Window IRTT\n" + "eth0 00000000 00000000 0003 0 0 0 00000000 0 0 0\n" + ) + rc, lines = self._run(root, route, socket_names="") + self.assertEqual(rc, 0) + self.assertEqual(self._captures(SOCKET_MARKER, lines), ["eth0"]) + + def test_inactive_port_emits_a_launcher_recognized_failure_marker(self) -> None: + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + self._fabric(root, state="1: DOWN") + rc, lines = self._run(root, root / "route") + self.assertEqual(rc, 1) + failures = [line for line in lines if re.search(FAILURE_MARKER, line)] + self.assertTrue(any("rdma-port-1=inactive" in line for line in failures), failures) + + def test_all_zero_gid_emits_gid_empty(self) -> None: + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + self._fabric(root, gid="0000:0000:0000:0000:0000:0000:0000:0000") + rc, lines = self._run(root, root / "route") + self.assertEqual(rc, 1) + self.assertTrue(any("rdma-port-1=gid-empty" in line for line in lines), lines) + + +class StageContract(unittest.TestCase): + # runtime/common.sh drives runtime/stage.py purely by literal subcommand name and positional + # argv — there are no optional flags. That argv shape is a string contract: a subcommand or + # flag the launcher passes but stage.py does not declare fails with "unrecognized arguments" + # and aborts the leg at repository-stage. This extracts every stage.py call out of common.sh + # and proves stage.py's parser accepts it — the guard that would have caught the --allow-* + # flags surviving on the callers after they were dropped from stage.py's argparse. + @staticmethod + def _invocations(text: str) -> list: + calls = [] + for line in text.splitlines(): + if "stage.py" not in line or line.lstrip().startswith("#"): + continue + subcommand, flags = None, [] + for raw in line.split("stage.py", 1)[1].split(): + token = raw.strip('"').strip("'") + if token.startswith("--"): + flags.append(token.split("=", 1)[0]) + elif subcommand is None and token and not token.startswith(("$", "${")): + subcommand = token + if subcommand: + calls.append((subcommand, flags)) + return calls + + def test_launcher_only_invokes_declared_subcommands_and_flags(self) -> None: + invocations = self._invocations((RUNTIME / "common.sh").read_text()) + self.assertGreaterEqual(len(invocations), len(stage.SPECS), invocations) + parser = stage.build_parser() + for subcommand, flags in invocations: + self.assertIn(subcommand, stage.SPECS, subcommand) + argv = [subcommand] + ["x"] * len(stage.SPECS[subcommand]) + flags + with contextlib.redirect_stderr(io.StringIO()): + try: + parser.parse_args(argv) + except SystemExit: + self.fail(f"common.sh invokes stage.py with an argv shape it rejects: {argv}") + + def test_contract_test_has_teeth(self) -> None: + # A flag common.sh must never pass has to be rejected by the parser — this is the exact + # failure (unrecognized arguments: --allow-parent-owner) the reconcile removed. + parser = stage.build_parser() + with contextlib.redirect_stderr(io.StringIO()): + with self.assertRaises(SystemExit): + parser.parse_args(["validate-stage-path", "x", "x", "x", "--allow-parent-owner"]) + + +# config.py case-args is the single case→invocation codec: collx_run_shard decodes one +# null-delimited argv per case and hands it verbatim to bench/run_ep.py. Parse the +# emitted argv with the same parser shape run_ep builds so the two sides cannot +# drift — a flag the codec emits but run_ep does not declare (or vice versa) fails +# here instead of on a GPU allocation. +class CaseArgvContract(unittest.TestCase): + CASE = { + "backend": "deepep-v2", "mode": "normal", "phase": "decode", + "routing": "uniform", "ep": 16, "nodes": 2, "gpus_per_node": 8, + "scale_up_domain": 8, "scope": "scale-out", + "scale_up_transport": "nvlink", "scale_out_transport": "rdma", + "transport": "nvlink-rdma", "topology_class": "h200-nvlink-rdma", + "hidden": 7168, "topk": 8, "experts": 256, "seed": 67, + "ladder": "1 2 4", "timing": "8:256:32", + "case_id": "h200-dgxc-deepep-v2-deepseek-v3-normal-decode-ep16-uniform", + "suite": "ep-core", "workload": "deepseek-v3", + } + + @staticmethod + def _run_ep_parser() -> argparse.ArgumentParser: + # Mirror of the parser bench/run_ep.py builds in main(). + parser = argparse.ArgumentParser() + parser.add_argument( + "--backend", required=True, choices=["deepep-v2", "mori"] + ) + ep_harness.add_common_args(parser) + return parser + + def _decode(self, stdout: bytes) -> list: + parts = stdout.split(b"\0") + self.assertEqual(parts[-1], b"") + return [part.decode() for part in parts[:-1]] + + def _case_argv(self, placement: list) -> list: + with tempfile.TemporaryDirectory() as directory: + path = Path(directory) / "shard.json" + path.write_text(json.dumps({"version": 1, "cases": [self.CASE]})) + result = subprocess.run( + [sys.executable, str(RUNTIME / "config.py"), "case-args", + str(path), "0", "h200-dgxc", "TS", *placement], + capture_output=True, check=True, + ) + return self._decode(result.stdout) + + def test_case_args_round_trips_through_the_run_ep_parser(self) -> None: + argv = self._case_argv(["16", "2", "8", "8"]) + args = self._run_ep_parser().parse_args(argv) + self.assertEqual( + (args.backend, args.mode, args.phase, args.routing, args.scope), + ("deepep-v2", "normal", "decode", "uniform", "scale-out"), + ) + self.assertEqual((args.hidden, args.topk, args.experts), (7168, 8, 256)) + self.assertEqual((args.gpus_per_node, args.scale_up_domain), (8, 8)) + self.assertEqual(args.tokens_ladder, "1 2 4") + self.assertEqual(args.scale_out_transport, "rdma") + self.assertEqual(args.case_id, self.CASE["case_id"]) + self.assertEqual(args.version, 1) + self.assertEqual(args.seed, self.CASE["seed"]) + self.assertEqual((args.iters, args.trials, args.warmup), (8, 256, 32)) + self.assertEqual(args.out, "results/h200-dgxc_deepep-v2_decode_TS-c000.json") + + def test_case_args_fails_closed_on_placement_mismatch(self) -> None: + with self.assertRaises(subprocess.CalledProcessError): + self._case_argv(["8", "1", "8", "8"]) + +if __name__ == "__main__": + unittest.main()