From 3b0a84997e6051dd771f133618387aec9a80d9b2 Mon Sep 17 00:00:00 2001 From: Lorri Rao Date: Thu, 6 Aug 2026 23:33:45 +0000 Subject: [PATCH 1/3] Integrate MORI for FSDP all-gather Wire MORI into TorchTitan and Megatron FSDP2, add runtime preflight and multi-node validation, and provide a clear custom all-gather backend selector. --- docs/02-user-guide/cli-reference.md | 2 + .../node-smoke-test-instruction.md | 2 +- docs/02-user-guide/preflight.md | 58 ++- docs/04-technical-guides/README.md | 1 + docs/04-technical-guides/sdma-allgather.md | 257 +++++++++++ primus/backends/common/__init__.py | 7 + primus/backends/common/mori_allgather.py | 378 ++++++++++++++++ .../patches/mori_allgather_patches.py | 151 +++++++ .../sdma_symm_mem_collectives_patches.py | 11 +- .../backends/torchtitan/patches/__init__.py | 1 + .../torchtitan/patches/mori_allgather.py | 134 ++++++ .../patches/sdma_symm_mem_collectives.py | 15 +- primus/cli/subcommands/preflight.py | 7 + primus/tools/preflight/mori_preflight.py | 417 +++++++++++++++++ primus/tools/preflight/mori_preflight.sh | 419 ++++++++++++++++++ primus/tools/preflight/preflight_args.py | 78 ++++ primus/tools/preflight/preflight_perf_test.py | 12 +- pyproject.toml | 5 +- .../hooks/06_enable_sdma_all_gather.sh | 14 +- .../hooks/07_enable_mori_all_gather.sh | 60 +++ runner/helpers/mori/install_mori.sh | 57 +++ .../helpers/mori/multinode_allgather_smoke.py | 73 +++ runner/primus-cli-direct.sh | 12 + tests/runner/test_primus_cli_direct.sh | 15 +- .../cli/test_mori_preflight_helper.py | 96 ++++ .../cli/test_preflight_subcommand.py | 56 +++ 26 files changed, 2309 insertions(+), 29 deletions(-) create mode 100644 docs/04-technical-guides/sdma-allgather.md create mode 100644 primus/backends/common/__init__.py create mode 100644 primus/backends/common/mori_allgather.py create mode 100644 primus/backends/megatron/patches/mori_allgather_patches.py create mode 100644 primus/backends/torchtitan/patches/mori_allgather.py create mode 100644 primus/tools/preflight/mori_preflight.py create mode 100755 primus/tools/preflight/mori_preflight.sh create mode 100755 runner/helpers/hooks/07_enable_mori_all_gather.sh create mode 100755 runner/helpers/mori/install_mori.sh create mode 100644 runner/helpers/mori/multinode_allgather_smoke.py create mode 100644 tests/unit_tests/cli/test_mori_preflight_helper.py diff --git a/docs/02-user-guide/cli-reference.md b/docs/02-user-guide/cli-reference.md index e7928ba8f..530308182 100644 --- a/docs/02-user-guide/cli-reference.md +++ b/docs/02-user-guide/cli-reference.md @@ -165,6 +165,7 @@ These run under `primus/cli/main.py` unless you change `--script` in direct mode | `train posttrain --config ` | Post-training (SFT or LoRA-style workflows; same top-level flags as pretrain in the parser). | | `benchmark [args]` | Performance microbenchmarks (see table below). | | `preflight [--host] [--gpu] [--network] [--perf-test]` | Cluster and node diagnostics. | +| `preflight --mori [MORI options]` | Run NIC configuration checking, build MORI and do MORI single/cross node test. Recommended if you plan to enable MORI for training. | | `projection memory --config ` | Memory estimation from a merged config. | | `projection performance --config ` | Performance projection from a merged config. | | `projection both --config ` | Single benchmark → both performance and memory projections (cluster sizing). | @@ -208,6 +209,7 @@ Within a chosen file, nested keys follow normal YAML structure. Slurm and contai | Container pretrain | `./runner/primus-cli container --volume /data:/data -- train pretrain --config /data/exp.yaml` | | Slurm training | `./runner/primus-cli slurm srun -N 4 -- train pretrain --config exp.yaml` | | Preflight (fast) | `./runner/primus-cli slurm srun -N 4 -- preflight --host --gpu --network` | +| MORI preflight | `./runner/primus-cli direct -- preflight --mori` | | Inspect launch command | `./runner/primus-cli --dry-run direct -- train pretrain --config exp.yaml` | | Dry-run Slurm | `./runner/primus-cli --dry-run slurm srun -N 2 -- train pretrain --config exp.yaml` | diff --git a/docs/02-user-guide/node-smoke-test-instruction.md b/docs/02-user-guide/node-smoke-test-instruction.md index 41f163be1..661b280e2 100644 --- a/docs/02-user-guide/node-smoke-test-instruction.md +++ b/docs/02-user-guide/node-smoke-test-instruction.md @@ -5,7 +5,7 @@ A lightweight, distributed-rendezvous-free preflight check that runs on every no Use it to **screen a cluster fast and exclude bad nodes before launching a real training job**. A bad GPU, NIC, wedged driver, or leaked process on any node surfaces as a node FAIL — without a single global rendezvous, so a stuck node can't wedge its peers. - **Recommended launcher**: `runner/primus-cli slurm srun -- direct -- node_smoke ...` (auto-resolves the distributed env, applies `slurm.*` config defaults, same pattern as `train` / `benchmark`). The shorter `runner/primus-cli direct -- node_smoke ...` (bare `srun` + `direct`) is equivalent and handy for ad-hoc runs. -- **Companion tool**: [`preflight`](./preflight.md) — the heavier diagnostic with a global rendezvous and inter-node bandwidth tests. The recommended workflow is **node-smoke first, preflight second** (see [§10](#10-comparison-with-the-full-preflight)). +- **Companion tool**: [`preflight`](./preflight.md) — the heavier diagnostic with a global rendezvous and inter-node bandwidth tests. `preflight --mori` mode additionally performs NIC configuration checking and MORI single/cross node test. The recommended workflow is **node-smoke first, preflight second** (see [§10](#10-comparison-with-the-full-preflight)). --- diff --git a/docs/02-user-guide/preflight.md b/docs/02-user-guide/preflight.md index eb1b325b4..da7469b39 100644 --- a/docs/02-user-guide/preflight.md +++ b/docs/02-user-guide/preflight.md @@ -21,17 +21,19 @@ Preflight has two report types, controlled by a single precedence rule: | Mode | Triggered by | What it does | |---|---|---| +| **MORI runtime preflight** | `--mori` | Runs on every selected node, prints NIC/RDMA details, pulls the base image, builds pinned MORI with live NIC detection, runs an 8-GPU correctness smoke per node, verifies matching fingerprints, and optionally runs one `8 × N`-rank all-gather. Exclusive with the standard selectors below. | | **Info-only** | `--host`, `--gpu`, `--network` (in any combination) | Lightweight host / GPU / network introspection. Emits a per-node report **without requiring a rendezvous**; multi-node aggregation then uses a **timeout-bounded** rendezvous (`--dist-timeout-sec`), so it never hangs indefinitely on network misconfig. | | **Perf-only** | `--perf-test`, `--tests ...`, or `--quick` | Runs the configured perf tests under a global rendezvous. **Implied** by `--tests` and `--quick`. | | **Default (info + perf)** | No flags at all | Runs the info report first, then every perf test. | ### Mode precedence -1. **Any of `--perf-test` / `--tests` / `--quick` is set → perf-only mode.** +1. **`--mori` is set → MORI runtime mode.** +2. **Any of `--perf-test` / `--tests` / `--quick` is set → perf-only mode.** If info selectors (`--host`/`--gpu`/`--network`) are also present, they are dropped and a `WARN` is emitted (also written as a `> Note:` at the top of the perf report). To get both reports, run two invocations. -2. **Otherwise, any of `--host`/`--gpu`/`--network` is set → info-only mode.** +3. **Otherwise, any of `--host`/`--gpu`/`--network` is set → info-only mode.** Perf-only tuning knobs (e.g. `--comm-sizes-mb`) are inert in this mode and trigger a single `WARN` listing them. -3. **Otherwise (no flags) → default**: info report **first** (no rendezvous), then perf tests. +4. **Otherwise (no flags) → default**: info report **first** (no rendezvous), then perf tests. The default order ensures you always get a report even if `torch.distributed` initialization later hangs. @@ -63,6 +65,51 @@ primus-cli direct -- preflight --perf-test primus-cli direct -- preflight --quick ``` +### MORI runtime build and local correctness smoke + +```bash +primus-cli direct -- preflight --mori +``` + +MORI mode must use the host/direct launcher because it starts its own +privileged temporary container. Under Slurm, include the explicit `direct` +entry: + +```bash +primus-cli slurm srun -N 1 --ntasks-per-node=1 \ + -- direct -- preflight --mori +``` + +General multi-node preflight builds/tests every listed node, verifies that their +NIC-stack fingerprints match, then runs one all-gather across all GPUs: + +```bash +primus-cli direct -- preflight --mori \ + --mori-nodes node1,node2,node3,node4 \ + --mori-socket-ifname fenic \ + --mori-gid-index 1 +``` + +Detailed phase behavior, timing, and validated MI355X commands are documented +in [`docs/04-technical-guides/sdma-allgather.md`](../04-technical-guides/sdma-allgather.md#mori-for-primus-fsdp). + +MORI mode options: + +| Flag | Default | Purpose | +|---|---|---| +| `--mori-base-image` | ROCm 7.15 Primus nightly | Base image pulled on every run. | +| `--mori-repo` | `https://github.com/ROCm/mori.git` | MORI source repository. | +| `--mori-ref` | pinned validated commit | Revision built by preflight. | +| `--mori-max-jobs` | `32` | Parallel source-build jobs. | +| `--mori-smoke-numel` | `67108864` | BF16 elements/rank in local and N-node smokes (128 MiB/rank). | +| `--mori-keep-container` | off | Keep the temporary build container for debugging. | +| `--mori-log-dir DIR` | under `--dump-path` | Override timed phase-log directory. | +| `--mori-nodes NODES` | current node | Comma-separated hosts, Slurm hostlist, or `@file`. Each node runs full local preflight before the N-node smoke. | +| `--mori-master-addr IP` | auto | Override master bootstrap address. | +| `--mori-master-port PORT` | `29610` | N-node torchrun port. | +| `--mori-socket-ifname IFACE` | auto | Override bootstrap interface. | +| `--mori-gid-index N` | auto | Override RoCEv2 GID index. | + Equivalent on SLURM via `primus-cli slurm`: ```bash @@ -314,6 +361,11 @@ sudo sysctl --system | `--report-file-name NAME` | auto-generated `preflight-${NNODES}N-YYYYMMDD-HHMMSS` | Base name for report files. Omit to let preflight auto-generate a unique timestamped name (prevents stale leftovers from prior runs being mistaken for fresh output). Pass an explicit value when you want a stable / well-known filename. | | `--disable-pdf` | enabled | Skip PDF generation (Markdown only). Useful when `weasyprint`/`markdown2` aren't installed. | +MORI mode writes timed phase logs and container diagnostics under +`/mori-preflight--/`, or the directory supplied by +`--mori-log-dir`. It does not generate the standard Markdown/PDF performance +report. + Output files: | File | Produced when | Notes | diff --git a/docs/04-technical-guides/README.md b/docs/04-technical-guides/README.md index 2f45c99c3..51e898be8 100644 --- a/docs/04-technical-guides/README.md +++ b/docs/04-technical-guides/README.md @@ -5,6 +5,7 @@ Deep technical topics for advanced users. - [Parallelism strategies](parallelism-strategies.md): DP, TP, PP, SP, CP, EP, FSDP explained - [Parallelism configuration](parallelism-configuration.md): per-backend parallelism setup and batch size relationships - [Collective operations](collective-operations.md): NCCL/RCCL operations and their role in each parallelism strategy +- [SDMA and MORI AllGather for FSDP](sdma-allgather.md): RCCL symmetric-memory SDMA and MORI hierarchical FSDP2 communication paths - [Performance tuning](performance-tuning.md): HipBLASLt, Primus-Turbo, FP8, MoE optimization - [MoE training deep-dive](moe-training.md): bottlenecks and Primus-Turbo optimizations for Mixture-of-Experts models - [MegaMoE fused MoE layer](mega-moe.md): FlyDSL-based fused MoE layer for EP-only bf16 training, setup and reproduction diff --git a/docs/04-technical-guides/sdma-allgather.md b/docs/04-technical-guides/sdma-allgather.md new file mode 100644 index 000000000..02a16c36a --- /dev/null +++ b/docs/04-technical-guides/sdma-allgather.md @@ -0,0 +1,257 @@ +# SDMA and MORI AllGather for FSDP + +Primus provides two custom PyTorch FSDP2 communication paths: + +| `FSDP_ALL_GATHER_BACKEND` | FSDP all-gather implementation | Intra-node path | Cross-node path | +|---|---|---|---| +| `rccl_sdma` | PyTorch symmetric memory over RCCL | RCCL with SDMA | RDMA without SDMA | +| `mori` | MORI `HierAllGather` | MORI SDMA | RDMA & SDMA | + +Both paths move all-gather traffic away from CU-resident RCCL kernels so that +FSDP communication can overlap with GEMM-heavy forward compute. Leave +`FSDP_ALL_GATHER_BACKEND` unset to use the framework default. + +## RCCL symmetric-memory SDMA + +### Enablement + +Set one user-facing switch before launching Primus: + +```bash +export FSDP_ALL_GATHER_BACKEND=rccl_sdma + +runner/primus-cli direct -- train pretrain --config +``` + +When `FSDP_ALL_GATHER_BACKEND=rccl_sdma` is set, the Primus hook +`runner/helpers/hooks/06_enable_sdma_all_gather.sh` emits the runtime +environment needed by the training container and torchrun children. + +The hook sets: + +```bash +NCCL_CTA_POLICY=2 +NCCL_CUMEM_ENABLE=1 +NCCL_LOCAL_REGISTER=0 +TORCH_NCCL_USE_TENSOR_REGISTER_ALLOCATOR_HOOK=true +FSDP_ALL_GATHER_BACKEND=rccl_sdma +LD_PRELOAD=/tmp/libhip_attr_drain.so +``` + +It also rebuilds `runner/helpers/hooks/sdma/hip_attr_drain_preload.c` into +`/tmp/libhip_attr_drain.so`. The interposer drains a stale HIP TLS error from +RCCL's cuMem capability probe on ROCm builds that do not have the upstream +fix. It does not change RCCL return values. + +### What the Primus patch does + +The Python backend patch is gated by +`FSDP_ALL_GATHER_BACKEND=rccl_sdma`. It wires PyTorch FSDP2 modules to use +symmetric-memory collectives: + +```python +from torch.distributed.fsdp._fully_shard._fsdp_collectives import ( + SymmMemAllGather, + SymmMemReduceScatter, +) + +module.set_custom_all_gather(SymmMemAllGather(group)) +module.set_custom_reduce_scatter(SymmMemReduceScatter(group)) +``` + +`SymmMemAllGather` allocates all-gather buffers from PyTorch symmetric memory. +Those buffers are cuMem-backed and rendezvoused across ranks. With zero-CTA +policy enabled, RCCL can dispatch the all-gather through the ROCm copy-engine +path (`__amd_rocclr_batchMemOp.kd` / `hsa_amd_memory_async_batch_copy`) instead +of running the data movement inside `ncclDevKernel_Generic_2` on CUs. + +The important discriminator is the buffer provenance: + +| Buffer source | Expected data path | +|---|---| +| `symm_mem.empty` / FSDP `SymmMemAllGather` | SDMA / copy engine | +| regular `torch.empty` / default FSDP all-gather | CU-resident RCCL kernel | + +The environment variables make the copy-engine path legal and observable, but +they do not by themselves turn regular `torch.empty` FSDP buffers into SDMA +buffers. The FSDP custom all-gather hook is the key. + +### Validation + +For low-level validation, use a symmetric-memory probe and verify that +`symm_mem.rendezvous()` completes and that a symmetric-memory all-gather runs. +For profiling validation, count HSA API calls or inspect traces for: + +```text +hsa_amd_memory_async_batch_copy +__amd_rocclr_batchMemOp.kd +``` + +Non-zero counts during all-gather indicate the SDMA copy-engine path. A trace +showing only `ncclDevKernel_Generic_2` for the data movement is not the SDMA +path, even if the communicator reports cuMem transport setup. + +### Driver and runtime compatibility + +The SDMA/FSDP path depends on PyTorch symmetric-memory rendezvous, which uses +ROCr virtual-memory APIs under the hood. The ROCm runtime and loaded amdgpu +driver must be compatible. + +One known failure mode was observed with a ROCm 7.15 nightly image where +`symm_mem.rendezvous()` hung inside: + +```text +hsa_amd_vmem_set_access + -> hsaKmtMemoryVaMap + -> driver ioctl / timeline wait +``` + +The userspace change involved: + +```text +b58362f60ff4f0b2b31a32a2a368db6bffdd5883 +ROCM-21775 Use DRM_IOCTL_SYNCOBJ_TIMELINE_WAIT ioctl in hsaKmt map/unmap ops +``` + +With an older loaded amdgpu driver, the relevant ioctl did not return, causing +`torch.distributed._symmetric_memory.rendezvous()` to hang. Updating and +reloading the amdgpu driver fixed the hang on the affected MI300X system: + +```text +$ sudo dkms status +amdgpu/7.1.3-2377367.22.04, 6.5.0-45-generic, x86_64: installed +$ uname -r +uname -r: 6.5.0-45-generic +``` + +If an SDMA run hangs before any FSDP forward progress, please try to dump stack and see where it hangs. + +## MORI hierarchical all-gather + +MORI replaces FSDP2 all-gather with `mori.ccl.HierAllGather`. It uses SDMA for +the intra-node comm and vendor direct verbs for the cross-node RDMA comm. + +### Enablement + +Set the single user-facing switch before launching Primus: + +```bash +export FSDP_ALL_GATHER_BACKEND=mori + +runner/primus-cli direct -- train pretrain --config +``` + +TorchTitan applies MORI to each compatible FSDP2 module. Megatron applies it to +FSDP2 transformer layers and additionally requires: + +```bash +--use_torch_fsdp2 true +``` + + +### What the Primus patch does + +The TorchTitan and Megatron patches wrap `fully_shard()` and attach one shared +adapter to compatible modules: + +```python +from primus.backends.common.mori_allgather import MoriAllGather + +mori_all_gather = MoriAllGather() +module.set_custom_all_gather(mori_all_gather) +``` + +The adapter: + +1. Initializes MORI SHMEM once from torchrun's default c10d `TCPStore`. This + avoids creating an eager cross-node RCCL transport solely for MORI + bootstrap. +2. Derives ranks per node from `LOCAL_WORLD_SIZE`. +3. Builds and caches `HierAllGather` for the FSDP process group and largest + observed per-rank input. +4. Launches MORI on the current CUDA stream and returns a Work-like object when + FSDP requests asynchronous completion. + + +### Runtime preflight + +MORI is sensitive to the live NIC driver, firmware, direct-verbs library, GID, +and capabilities such as Ionic CCQE. Any slight misalignment / misconfig will likly cause MORI to fail. To mitigate this issue, we provide an unified Primus preflight command, which can detect all the known critical configs and do a test all-gather on every target node: + +```bash +runner/primus-cli direct -- preflight --mori +``` + +For multi-node validation: + +```bash +runner/primus-cli direct -- preflight --mori \ + --mori-nodes node1,node2 \ + --mori-socket-ifname \ + --mori-gid-index +``` + +The CLI invokes `primus/tools/preflight/mori_preflight.py`, which runs +`mori_preflight.sh` on every selected node. The shell worker: + +1. Prints host identity, GPU, IP, RDMA links, valid GIDs, NIC + driver/firmware, vendor-library hash, and required DV symbols. +2. Starts a privileged temporary container from the pinned Primus CI image. +3. Mounts the detected host vendor library into that container. +4. Calls `runner/helpers/mori/install_mori.sh` to install dependencies, clone + the pinned source/submodules, and build MORI with live RDMA visibility. +5. Calls MORI's runtime detector and records `ccqe_runtime` in the node + fingerprint. +6. Runs an 8-GPU bit-exact all-gather smoke. +7. When `--mori-nodes` is set, keeps the temporary containers for this same + information/build/local smoke on + every node, verifies matching node fingerprints, then launches one + all-gather over all `8 × N` ranks before removing them. + +The mori source version pinning is required for now, till this PR is stablized in our base rocm docker: https://github.com/ROCm/mori/pull/441 + +Logs and phase timing are written under +`/tmp/primus-mori-preflight--/`. + +### Vendor library names + +The library names used by preflight come directly from MORI: + +| NIC | MORI runtime loader names | +|---|---| +| Ionic / AINIC | `libionic.so` | +| Broadcom BNXT | `libbnxt_re.so`, then `libbnxt_re-rdmav59.so`, then `libbnxt_re-rdmav34.so` | +| Mellanox mlx5 | `libmlx5.so` | + +MORI's +[`dv_loader.hpp`](https://github.com/ROCm/mori/blob/dc4bc75a8ae63cb79a3ce17e55f2be3d8aa692c2/include/mori/application/transport/rdma/providers/dv_loader.hpp#L133) +uses these exact `dlopen()` names. Its +[`MoriDetectDevice.cmake`](https://github.com/ROCm/mori/blob/dc4bc75a8ae63cb79a3ce17e55f2be3d8aa692c2/cmake/MoriDetectDevice.cmake#L140) +uses the same names for build-time `find_library()` detection. Preflight mounts +the host's detected vendor library under these aliases so build-time detection +and runtime loading use the same library. + +### Automatic MORI installation + +When `FSDP_ALL_GATHER_BACKEND=mori` is set, the launcher hook checks whether +`mori.ccl.HierAllGather` is available. If it is missing, the hook calls +`runner/helpers/mori/install_mori.sh` before torchrun starts, so no separate +user installation command is needed. Automatic installation requires root +inside the training container. + +Useful overrides are `MORI_REPO`, `MORI_REF`, `MORI_SOURCE_DIR`, `MAX_JOBS`, +and `ROCM_PATH`. The installer clears `MORI_DEVICE_NIC` so MORI detects the +live NIC and mounted vendor library. + + +### Troubleshooting + +- `ccqe_candidate=true` on one node and `mixed` or `false` on another: choose + nodes with matching Ionic firmware and vendor-library capabilities. +- `ccqe_runtime=true` on one node and `false` on another: the training image + sees different effective vendor-library or NIC support; do not launch the + cross-node collective. +- `local GID N/A`: inspect `/sys/class/infiniband/ionic_*/ports/1/gids/`; + this pair uses `NCCL_IB_GID_INDEX=1`, not `3`. +- BNXT `231.x`: unsupported for MORI IBGDA; use supported firmware/userspace or + a validated mlx5/ionic pair. diff --git a/primus/backends/common/__init__.py b/primus/backends/common/__init__.py new file mode 100644 index 000000000..14fbae675 --- /dev/null +++ b/primus/backends/common/__init__.py @@ -0,0 +1,7 @@ +############################################################################### +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +"""Shared backend integration helpers.""" diff --git a/primus/backends/common/mori_allgather.py b/primus/backends/common/mori_allgather.py new file mode 100644 index 000000000..730f54e56 --- /dev/null +++ b/primus/backends/common/mori_allgather.py @@ -0,0 +1,378 @@ +############################################################################### +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +"""MORI-backed FSDP2 all-gather adapter. + +The implementation mirrors the public FSDP2 adapter shape from +ROCm/mori's ``examples/fsdp_sdma/mori_allgather.py`` while keeping the +Primus backend patches small. It is intentionally all-gather only: +FSDP reduce-scatter stays on the framework default path. +""" + +from __future__ import annotations + +import importlib +import os +from collections.abc import Sequence +from typing import Any + +import torch +import torch.distributed as dist + +from primus.core.utils.module_utils import log_rank_0 + +try: + from torch.distributed.fsdp._fully_shard._fsdp_api import ( + AllGather as _FSDPAllGather, + ) +except Exception as e: # pragma: no cover - depends on torch internal version + _FSDPAllGather = object + _FSDP_ALL_GATHER_IMPORT_ERROR = e +else: + _FSDP_ALL_GATHER_IMPORT_ERROR = None + +_MORI_SHMEM_INITIALIZED = False + + +def _safe_log_rank_0(message: str) -> None: + """Log through Primus when initialized; otherwise fall back to print.""" + try: + log_rank_0(message) + except Exception: + if not dist.is_available() or not dist.is_initialized() or dist.get_rank() == 0: + print(message, flush=True) + + +def mori_all_gather_enabled() -> bool: + """Return whether Primus should install the MORI FSDP all-gather backend.""" + return os.environ.get("FSDP_ALL_GATHER_BACKEND", "") == "mori" + + +def ensure_mori_shmem_initialized(pg_name: str = "default") -> None: + """Initialize MORI SHMEM from torchrun's c10d store once. + + MORI's convenience process-group initializer broadcasts its unique ID via + the process-group backend. For NCCL groups that eagerly creates a second + cross-node transport before MORI is ready. Use the rendezvous TCPStore + instead so MORI remains the only RDMA data path. + """ + global _MORI_SHMEM_INITIALIZED + + if _MORI_SHMEM_INITIALIZED: + return + if not dist.is_available() or not dist.is_initialized(): + raise RuntimeError("MORI FSDP all-gather requires torch.distributed to be initialized") + + shmem = importlib.import_module("mori.shmem") + rank = dist.get_rank() + world_size = dist.get_world_size() + store = dist.distributed_c10d._get_default_store() + uid_key = f"primus_mori_shmem_uid_{pg_name}_{world_size}" + if rank == 0: + store.set(uid_key, shmem.shmem_get_unique_id()) + uid = store.get(uid_key) + shmem.shmem_init_attr( + shmem.MORI_SHMEM_INIT_WITH_UNIQUEID, + rank, + world_size, + uid, + ) + + if shmem.shmem_mype() != rank or shmem.shmem_npes() != world_size: + raise RuntimeError( + "MORI SHMEM PE mapping must match the FSDP process group: " + f"rank/world_size={rank}/{world_size}, " + f"mype/npes={shmem.shmem_mype()}/{shmem.shmem_npes()}" + ) + + _MORI_SHMEM_INITIALIZED = True + _safe_log_rank_0("[MORI:FSDP] initialized MORI SHMEM from torch process group") + + +class _CudaEventWork: + """Small Work-like object for async FSDP all-gather calls.""" + + def __init__(self, event: torch.cuda.Event, device: torch.device) -> None: + self._event = event + self._device = device + self._waited = False + + def wait(self) -> bool: + if not self._waited: + torch.cuda.current_stream(self._device).wait_event(self._event) + self._waited = True + return True + + +class _DeviceDeferredHostSyncWork(dist.distributed_c10d.Work): + """Defer MORI's reliable host landing fence until FSDP consumes the AG.""" + + def __init__(self, stream: torch.cuda.Stream, event: torch.cuda.Event | None = None) -> None: + super().__init__() + self._stream = stream + self._event = event + self._done = False + + def wait(self, timeout=None) -> bool: # noqa: ARG002 + if not self._done: + if self._event is not None: + self._event.synchronize() + else: + self._stream.synchronize() + self._done = True + return True + + def is_completed(self) -> bool: + return self._done + + +class _HostProxyDeferredWork(dist.distributed_c10d.Work): + """Defer host-proxy completion to FSDP's wait/copy-out point.""" + + def __init__(self, collective: Any, handle: Any, drain: bool = False) -> None: + super().__init__() + self._collective = collective + self._handle = handle + self._drain = drain + self._done = False + + def wait(self, timeout=None) -> bool: # noqa: ARG002 + if not self._done: + self._collective._complete(self._handle) + if self._drain: + self._handle["stream"].synchronize() + self._collective._pending = None + self._done = True + return True + + def is_completed(self) -> bool: + return self._done + + +class MoriAllGather(_FSDPAllGather): + """FSDP2 custom all-gather backed by ``mori.ccl.HierAllGather``.""" + + def __init__(self, ranks_per_node: int | None = None) -> None: + if _FSDP_ALL_GATHER_IMPORT_ERROR is not None: + raise ImportError( + "MORI FSDP all-gather requires PyTorch FSDP2's internal " "AllGather API" + ) from _FSDP_ALL_GATHER_IMPORT_ERROR + + os.environ.setdefault("MORI_ENABLE_SDMA", "1") + os.environ.setdefault("MORI_SHMEM_HEAP_SIZE", "8G") + os.environ.setdefault("MORI_HIER_CUDA_GRAPH", "0") + if "MORI_SOCKET_IFNAME" not in os.environ and "NCCL_SOCKET_IFNAME" in os.environ: + os.environ["MORI_SOCKET_IFNAME"] = os.environ["NCCL_SOCKET_IFNAME"].lstrip("=") + + self._ranks_per_node = ranks_per_node + self._collective: Any | None = None + self._rank: int | None = None + self._world_size: int | None = None + self._cap_bytes = 0 + self._output_buffer: torch.Tensor | None = None + + world = int(os.environ.get("WORLD_SIZE", "0") or "0") + if world > 0: + rpn = self._ranks_per_node_value(world) + num_nodes = world // rpn if rpn else 1 + if num_nodes >= 2: + # These defaults are copied from MORI's FSDP example and only + # apply if the user did not explicitly tune the same variables. + setdefault = os.environ.setdefault + setdefault("MORI_HIER_FUSE_LOCAL", "1") + setdefault("MORI_HIER_FUSE_REMOTE", "1") + setdefault("MORI_HIER_LOCAL_PUSHONLY", "1") + if rpn < 8: + setdefault("MORI_HIER_DEEP_PIPE", "auto") + setdefault("MORI_SDMA_NUM_CHANNELS", "8") + else: + setdefault("MORI_HIER_DEBUG_SYNC", "1") + setdefault("MORI_HIER_CUDA_GRAPH", "0") + setdefault("MORI_FSDP_DEFER_HOSTSYNC", "1") + setdefault("MORI_FSDP_EVENT_FENCE", "1") + setdefault("MORI_FSDP_FWD_PREFETCH", "1") + + self._host_proxy = os.environ.get("MORI_FSDP_HOST_PROXY", "") not in ( + "", + "0", + "false", + "False", + ) + self._hostproxy_async = os.environ.get("MORI_HOSTPROXY_ASYNC", "") not in ( + "", + "0", + "false", + "False", + ) + if self._hostproxy_async: + os.environ.setdefault("MORI_HOSTPROXY_ASYNC_DRAIN", "1") + os.environ.setdefault("MORI_HOSTPROXY_ASYNC_RING", "2") + self._hostproxy_async_drain = os.environ.get("MORI_HOSTPROXY_ASYNC_DRAIN", "") not in ( + "", + "0", + "false", + "False", + ) + self._defer_hostsync = os.environ.get("MORI_FSDP_DEFER_HOSTSYNC", "") not in ( + "", + "0", + "false", + "False", + ) + self._event_fence = os.environ.get("MORI_FSDP_EVENT_FENCE", "") not in ( + "", + "0", + "false", + "False", + ) + + def allocate( + self, + size: Sequence[int | torch.SymInt], + *, + dtype: torch.dtype, + device: torch.device, + ) -> torch.Tensor: + numel = 1 + for dim in size: + numel *= int(dim) + if ( + self._output_buffer is not None + and self._output_buffer.dtype == dtype + and self._output_buffer.device == device + and self._output_buffer.numel() >= numel + ): + return self._output_buffer.narrow(0, 0, numel) + self._output_buffer = torch.empty(numel, dtype=dtype, device=device) + return self._output_buffer + + def _ranks_per_node_value(self, world_size: int) -> int: + if self._ranks_per_node is not None: + return self._ranks_per_node + env_value = os.environ.get("LOCAL_WORLD_SIZE") + if env_value: + return int(env_value) + return min(torch.cuda.device_count(), world_size) + + def _get_collective(self, group: dist.ProcessGroup, per_rank_bytes: int) -> Any: + rank, world_size = group.rank(), group.size() + if ( + self._collective is not None + and self._rank == rank + and self._world_size == world_size + and self._cap_bytes >= per_rank_bytes + ): + return self._collective + + ensure_mori_shmem_initialized("default") + + shmem = importlib.import_module("mori.shmem") + ccl = importlib.import_module("mori.ccl") + my_pe = shmem.shmem_mype() + npes = shmem.shmem_npes() + if my_pe != rank or npes != world_size: + raise RuntimeError( + "MORI FSDP HierAllGather requires the FSDP process group to match " + f"SHMEM PEs, got rank/world_size={rank}/{world_size} and " + f"my_pe/npes={my_pe}/{npes}" + ) + + cap = max(per_rank_bytes, self._cap_bytes) + ranks_per_node = self._ranks_per_node_value(world_size) + if self._host_proxy: + cap_floor = int(os.environ.get("MORI_FSDP_HOSTPROXY_CAP_MB", "160")) * (1 << 20) + cap = max(cap, cap_floor) + if self._collective is not None: + raise RuntimeError( + "HostProxyHierAllGather built with cap " + f"{self._cap_bytes} B but a {per_rank_bytes} B AG arrived; " + "raise MORI_FSDP_HOSTPROXY_CAP_MB" + ) + collective = ccl.HostProxyHierAllGather( + rank, + world_size, + ranks_per_node, + output_buffer_size=cap * world_size, + ) + else: + collective = ccl.HierAllGather( + my_pe, + npes, + input_buffer_size=cap, + output_buffer_size=cap * world_size, + copy_output_to_user=True, + ranks_per_node=ranks_per_node, + ) + + self._collective = collective + self._rank = rank + self._world_size = world_size + self._cap_bytes = cap + return collective + + def _validate( + self, + output_tensor: torch.Tensor, + input_tensor: torch.Tensor, + group: dist.ProcessGroup, + ) -> None: + if not input_tensor.is_cuda or not output_tensor.is_cuda: + raise RuntimeError("MORI FSDP HierAllGather requires CUDA tensors") + if input_tensor.device != output_tensor.device: + raise RuntimeError("MORI FSDP HierAllGather requires tensors on the same device") + if input_tensor.dtype != output_tensor.dtype: + raise RuntimeError("MORI FSDP HierAllGather requires matching dtypes") + expected = input_tensor.numel() * group.size() + if output_tensor.numel() != expected: + raise RuntimeError( + f"MORI FSDP HierAllGather expected output numel {expected}, " f"got {output_tensor.numel()}" + ) + if (input_tensor.numel() * input_tensor.element_size()) % 4 != 0: + raise RuntimeError("MORI FSDP HierAllGather requires 4-byte-aligned input bytes") + + def __call__( + self, + output_tensor: torch.Tensor, + input_tensor: torch.Tensor, + group: dist.ProcessGroup, + async_op: bool = False, + ) -> Any | None: + self._validate(output_tensor, input_tensor, group) + per_rank_bytes = input_tensor.numel() * input_tensor.element_size() + collective = self._get_collective(group, per_rank_bytes) + device = input_tensor.device + stream = torch.cuda.current_stream(device) + + input_tensor.record_stream(stream) + output_tensor.record_stream(stream) + + if self._host_proxy and self._hostproxy_async: + pending = getattr(collective, "_pending", None) + if pending is not None: + pending.wait() + handle = collective.call_async(input_tensor, output_tensor, input_tensor.numel(), stream=stream) + if handle is None: + return None + work = _HostProxyDeferredWork(collective, handle, drain=self._hostproxy_async_drain) + collective._pending = work + return work + + ok = collective(input_tensor, output_tensor, input_tensor.numel(), stream=stream) + if not ok: + raise RuntimeError("MORI HierAllGather call failed") + + if self._defer_hostsync: + event = None + if self._event_fence: + event = torch.cuda.Event() + event.record(stream) + return _DeviceDeferredHostSyncWork(stream, event) + + if async_op: + event = torch.cuda.Event() + event.record(stream) + return _CudaEventWork(event, device) + return None diff --git a/primus/backends/megatron/patches/mori_allgather_patches.py b/primus/backends/megatron/patches/mori_allgather_patches.py new file mode 100644 index 000000000..4fcf8e971 --- /dev/null +++ b/primus/backends/megatron/patches/mori_allgather_patches.py @@ -0,0 +1,151 @@ +############################################################################### +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +"""Megatron FSDP2 MORI all-gather patch. + +Activation: + Export ``FSDP_ALL_GATHER_BACKEND=mori`` and run Megatron with + ``use_torch_fsdp2: true``. This patch wraps the ``fully_shard`` symbol + used by Megatron's Torch FSDP2 wrapper and attaches MORI's FSDP2 + ``AllGather`` backend to transformer-layer units. + +Scope: + Transformer layers only for the first integration pass. Embedding, + lm_head, rotary, and reduce-scatter stay on framework defaults. +""" + +import functools +import os +import sys + +from primus.core.patches import PatchContext, get_args, register_patch +from primus.core.utils.module_utils import log_rank_0, warning_rank_0 + + +def _mori_all_gather_enabled(ctx: PatchContext) -> bool: + """Gate on the MORI backend and the Megatron FSDP2 path.""" + if os.environ.get("FSDP_ALL_GATHER_BACKEND", "") != "mori": + return False + return getattr(get_args(ctx), "use_torch_fsdp2", False) + + +@register_patch( + "megatron.fsdp.mori_allgather", + backend="megatron", + phase="before_train", + description=( + "Attach MORI HierAllGather to Megatron FSDP2 transformer-layer " + "modules so all-gather uses intra-node SDMA and cross-node RDMA. " + "Gated on FSDP_ALL_GATHER_BACKEND=mori and use_torch_fsdp2." + ), + condition=_mori_all_gather_enabled, + priority=40, +) +def patch_megatron_fsdp_mori_allgather(ctx: PatchContext) -> None: + """Install MORI's FSDP2 all-gather backend for Megatron transformer layers.""" + import megatron.core.distributed.torch_fully_sharded_data_parallel as _mfsdp_mod + import torch.distributed.fsdp as _fsdp_pkg + from torch.distributed.fsdp._fully_shard import _fully_shard as _ffs_mod + + from primus.backends.common.mori_allgather import MoriAllGather + + if not getattr(_mfsdp_mod, "HAVE_FSDP", False) or not hasattr(_mfsdp_mod, "fully_shard"): + warning_rank_0( + "[Patch:megatron.fsdp.mori_allgather] " + "torch_fully_sharded_data_parallel.fully_shard not found; skipping." + ) + return + + try: + from megatron.core.transformer.transformer_layer import TransformerLayer + except Exception as e: + warning_rank_0( + "[Patch:megatron.fsdp.mori_allgather] could not import " f"TransformerLayer; skipping: {e}" + ) + return + + orig_fully_shard = _fsdp_pkg.fully_shard + mori_all_gather = MoriAllGather() + log_all = os.environ.get("MORI_LOG_ATTACH", "0") == "1" + + def _attach_mori_all_gather(fsdp_module) -> None: + """Attach MoriAllGather to transformer-layer units only.""" + if not isinstance(fsdp_module, TransformerLayer): + if log_all: + log_rank_0( + "[Patch:megatron.fsdp.mori_allgather] skip non-transformer " + f"module: {type(fsdp_module).__name__}" + ) + return + try: + state = fsdp_module._get_fsdp_state() + except Exception as e: + if log_all: + log_rank_0( + "[Patch:megatron.fsdp.mori_allgather] skip " + f"{type(fsdp_module).__name__} without FSDP state: {e}" + ) + return + + groups = getattr(state, "_fsdp_param_groups", None) or [] + if len(groups) != 1: + if log_all: + log_rank_0( + "[Patch:megatron.fsdp.mori_allgather] skip " + f"{type(fsdp_module).__name__} with {len(groups)} param groups" + ) + return + + try: + fsdp_module.set_custom_all_gather(mori_all_gather) + except (AttributeError, ValueError, AssertionError) as e: + warning_rank_0( + f"[Patch:megatron.fsdp.mori_allgather] WARN: failed to " + f"attach MoriAllGather to {type(fsdp_module).__name__}: {e}" + ) + return + if log_all: + pg = groups[0]._all_gather_process_group + log_rank_0( + "[Patch:megatron.fsdp.mori_allgather] attached MoriAllGather " + f"to {type(fsdp_module).__name__} (group={pg.group_name})" + ) + + @functools.wraps(orig_fully_shard) + def wrapped_fully_shard(module, *args, **kwargs): + result = orig_fully_shard(module, *args, **kwargs) + _attach_mori_all_gather(result if result is not None else module) + return result + + for attr in dir(orig_fully_shard): + if attr.startswith("__"): + continue + try: + setattr(wrapped_fully_shard, attr, getattr(orig_fully_shard, attr)) + except (AttributeError, TypeError): + pass + + _fsdp_pkg.fully_shard = wrapped_fully_shard + if hasattr(_ffs_mod, "fully_shard"): + _ffs_mod.fully_shard = wrapped_fully_shard + _mfsdp_mod.fully_shard = wrapped_fully_shard + + patched_aliases = 0 + for module in tuple(sys.modules.values()): + if module is None: + continue + try: + if getattr(module, "fully_shard", None) is orig_fully_shard: + setattr(module, "fully_shard", wrapped_fully_shard) + patched_aliases += 1 + except (AttributeError, TypeError): + continue + + log_rank_0( + "[Patch:megatron.fsdp.mori_allgather] installed: FSDP2 " + "transformer-layer modules in Megatron's TorchFullyShardedDataParallel " + f"will use MoriAllGather; patched {patched_aliases} loaded aliases." + ) diff --git a/primus/backends/megatron/patches/sdma_symm_mem_collectives_patches.py b/primus/backends/megatron/patches/sdma_symm_mem_collectives_patches.py index 8f6bf5708..bc07aca4e 100644 --- a/primus/backends/megatron/patches/sdma_symm_mem_collectives_patches.py +++ b/primus/backends/megatron/patches/sdma_symm_mem_collectives_patches.py @@ -25,8 +25,8 @@ FSDP default. Activation: - Export ``SDMA_ALL_GATHER=1`` AND run with ``use_torch_fsdp2: true``. - No-op otherwise. The companion hook + Export ``FSDP_ALL_GATHER_BACKEND=rccl_sdma`` AND run with + ``use_torch_fsdp2: true``. No-op otherwise. The companion hook ``runner/helpers/hooks/06_enable_sdma_all_gather.sh`` exports the zero-CTA env (``NCCL_CTA_POLICY=2``, ...) and the LD_PRELOAD interposer so no YAML changes are required to opt in. @@ -45,8 +45,8 @@ def _sdma_all_gather_enabled(ctx: PatchContext) -> bool: - """Gate on SDMA_ALL_GATHER=1 AND the use_torch_fsdp2 path.""" - if os.environ.get("SDMA_ALL_GATHER", "0") != "1": + """Gate on the RCCL SDMA backend and the use_torch_fsdp2 path.""" + if os.environ.get("FSDP_ALL_GATHER_BACKEND", "") != "rccl_sdma": return False return getattr(get_args(ctx), "use_torch_fsdp2", False) @@ -59,7 +59,8 @@ def _sdma_all_gather_enabled(ctx: PatchContext) -> bool: "Attach SymmMemAllGather to FSDP2 transformer-layer modules in " "Megatron's TorchFullyShardedDataParallel so their all-gather uses " "the SDMA (copy-engine) dispatch path; other units stay on FSDP's " - "default all-gather. Gated on SDMA_ALL_GATHER=1 and use_torch_fsdp2." + "default all-gather. Gated on FSDP_ALL_GATHER_BACKEND=rccl_sdma " + "and use_torch_fsdp2." ), condition=_sdma_all_gather_enabled, ) diff --git a/primus/backends/torchtitan/patches/__init__.py b/primus/backends/torchtitan/patches/__init__.py index b4c6e2d9a..a0556f14a 100644 --- a/primus/backends/torchtitan/patches/__init__.py +++ b/primus/backends/torchtitan/patches/__init__.py @@ -34,6 +34,7 @@ metrics_output_format, mock_dataset_patches, model_override_patches, + mori_allgather, peak_flops_patches, pipelining_schedule_patches, sdma_symm_mem_collectives, diff --git a/primus/backends/torchtitan/patches/mori_allgather.py b/primus/backends/torchtitan/patches/mori_allgather.py new file mode 100644 index 000000000..00aaf64ed --- /dev/null +++ b/primus/backends/torchtitan/patches/mori_allgather.py @@ -0,0 +1,134 @@ +############################################################################### +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +"""TorchTitan / FSDP2 MORI all-gather patch. + +Activation: + Export ``FSDP_ALL_GATHER_BACKEND=mori`` before launching a TorchTitan + pretrain. + The patch wraps ``torch.distributed.fsdp.fully_shard`` and attaches + MORI's FSDP2 ``AllGather`` backend to each fully-sharded module via + ``set_custom_all_gather``. + +MORI routes intra-node traffic over SDMA and cross-node traffic over RDMA. +Reduce-scatter stays on the framework default path. +""" + +import functools +import os +import sys + +from primus.core.patches import PatchContext, register_patch +from primus.core.utils.module_utils import log_rank_0 + + +def _mori_all_gather_enabled(ctx: PatchContext) -> bool: + """Gate on the MORI all-gather backend.""" + return os.environ.get("FSDP_ALL_GATHER_BACKEND", "") == "mori" + + +@register_patch( + "torchtitan.fsdp.mori_allgather", + backend="torchtitan", + phase="setup", + description=( + "Auto-attach MORI HierAllGather to every fully_shard'd module so " + "FSDP all-gather uses intra-node SDMA and cross-node RDMA. " + "Gated on FSDP_ALL_GATHER_BACKEND=mori." + ), + condition=_mori_all_gather_enabled, + priority=40, +) +def patch_torchtitan_fsdp_mori_allgather(ctx: PatchContext) -> None: + """Install MORI's FSDP2 all-gather backend for TorchTitan modules.""" + import torch.distributed.fsdp as _fsdp_pkg + from torch.distributed.fsdp._fully_shard import _fsdp_collectives as _ffsc + from torch.distributed.fsdp._fully_shard import _fully_shard as _ffs_mod + + from primus.backends.common.mori_allgather import MoriAllGather + + orig_fully_shard = _fsdp_pkg.fully_shard + mori_all_gather = MoriAllGather() + log_all = os.environ.get("MORI_LOG_ATTACH", "0") == "1" + + def _attach_mori_all_gather(fsdp_module) -> None: + try: + state = fsdp_module._get_fsdp_state() + except Exception as e: + if log_all: + log_rank_0( + f"[Patch:torchtitan.fsdp.mori_allgather] " + f"skip (no _fsdp_state): {type(fsdp_module).__name__}: {e}" + ) + return + + groups = getattr(state, "_fsdp_param_groups", None) or [] + if len(groups) != 1: + if log_all: + log_rank_0( + f"[Patch:torchtitan.fsdp.mori_allgather] " + f"skip multi-group module ({len(groups)} groups): " + f"{type(fsdp_module).__name__}" + ) + return + + try: + fsdp_module.set_custom_all_gather(mori_all_gather) + except (AttributeError, ValueError, AssertionError) as e: + log_rank_0( + f"[Patch:torchtitan.fsdp.mori_allgather] WARN: failed to " + f"attach MoriAllGather to {type(fsdp_module).__name__}: {e}" + ) + return + if log_all: + pg = groups[0]._all_gather_process_group + log_rank_0( + f"[Patch:torchtitan.fsdp.mori_allgather] attached " + f"MoriAllGather to {type(fsdp_module).__name__} " + f"(group={pg.group_name})" + ) + + @functools.wraps(orig_fully_shard) + def wrapped_fully_shard(module, *args, **kwargs): + result = orig_fully_shard(module, *args, **kwargs) + _attach_mori_all_gather(result if result is not None else module) + return result + + for _attr in dir(orig_fully_shard): + if _attr.startswith("__"): + continue + try: + setattr(wrapped_fully_shard, _attr, getattr(orig_fully_shard, _attr)) + except (AttributeError, TypeError): + pass + + _fsdp_pkg.fully_shard = wrapped_fully_shard + if hasattr(_ffs_mod, "fully_shard"): + _ffs_mod.fully_shard = wrapped_fully_shard + + # TorchTitan model modules may already hold a local + # ``from torch.distributed.fsdp import fully_shard`` alias by the time + # Primus runs setup patches. Replace every loaded alias that still points + # at the exact original function so those call sites also attach MORI. + patched_aliases = 0 + for module in tuple(sys.modules.values()): + if module is None: + continue + try: + if getattr(module, "fully_shard", None) is orig_fully_shard: + setattr(module, "fully_shard", wrapped_fully_shard) + patched_aliases += 1 + except (AttributeError, TypeError): + continue + + log_rank_0( + "[Patch:torchtitan.fsdp.mori_allgather] installed: every " + "fully_shard()-d module will use MoriAllGather for FSDP all-gather; " + f"patched {patched_aliases} loaded aliases." + ) + + _ffsc._primus_mori_orig_fully_shard = orig_fully_shard + _ffsc._primus_mori_attach = _attach_mori_all_gather diff --git a/primus/backends/torchtitan/patches/sdma_symm_mem_collectives.py b/primus/backends/torchtitan/patches/sdma_symm_mem_collectives.py index 60aa0a639..8af9907f7 100644 --- a/primus/backends/torchtitan/patches/sdma_symm_mem_collectives.py +++ b/primus/backends/torchtitan/patches/sdma_symm_mem_collectives.py @@ -14,18 +14,19 @@ immediately after construction. Activation: - Export ``SDMA_ALL_GATHER=1`` in the shell before launching any + Export ``FSDP_ALL_GATHER_BACKEND=rccl_sdma`` before launching any torchtitan pretrain (e.g. ``primus-cli direct -- train pretrain --config ``). The patch is a no-op otherwise. The companion hook ``runner/helpers/hooks/06_enable_sdma_all_gather.sh`` runs at - ``primus-cli`` startup and, when ``SDMA_ALL_GATHER=1``, exports + ``primus-cli`` startup and, when the backend is ``rccl_sdma``, exports the standard zero-CTA env (``NCCL_CTA_POLICY=2``, ``NCCL_CUMEM_ENABLE=1``, ...) and the LD_PRELOAD interposer so no YAML or script changes are required to opt in. - ``SDMA_ALL_GATHER`` is the only knob; there are no sub-options. + ``FSDP_ALL_GATHER_BACKEND`` is the only selector; there are no + backend-specific sub-options. Requirements: - PyTorch >= 2.12 (introduces ``SymmMemAllGather``). @@ -43,8 +44,8 @@ def _sdma_all_gather_enabled(ctx: PatchContext) -> bool: - """Single env-driven gate. Triggered only by ``SDMA_ALL_GATHER=1``.""" - return os.environ.get("SDMA_ALL_GATHER", "0") == "1" + """Gate on the RCCL symmetric-memory SDMA backend.""" + return os.environ.get("FSDP_ALL_GATHER_BACKEND", "") == "rccl_sdma" @register_patch( @@ -54,7 +55,7 @@ def _sdma_all_gather_enabled(ctx: PatchContext) -> bool: description=( "Auto-attach SymmMemAllGather to every fully_shard'd module " "so FSDP's all-gather uses the SDMA (copy-engine) dispatch " - "path. Gated on SDMA_ALL_GATHER=1." + "path. Gated on FSDP_ALL_GATHER_BACKEND=rccl_sdma." ), condition=_sdma_all_gather_enabled, ) @@ -73,7 +74,7 @@ def patch_torchtitan_fsdp_sdma_symm_mem(ctx: PatchContext) -> None: from torch.distributed.fsdp._fully_shard import _fully_shard as _ffs_mod from torch.distributed.fsdp._fully_shard._fsdp_collectives import SymmMemAllGather - # Hardcoded sensible defaults. SDMA_ALL_GATHER is the only knob. + # Hardcoded sensible defaults. FSDP_ALL_GATHER_BACKEND is the only selector. backend = "NCCL" log_all = False diff --git a/primus/cli/subcommands/preflight.py b/primus/cli/subcommands/preflight.py index fc3c581d8..a43c1b093 100644 --- a/primus/cli/subcommands/preflight.py +++ b/primus/cli/subcommands/preflight.py @@ -13,6 +13,7 @@ primus-cli preflight --gpu # GPU info only primus-cli preflight --network # Network info only primus-cli preflight --perf-test # Perf tests only (skip info) + primus-cli preflight --mori # MORI runtime build + smoke """ from __future__ import annotations @@ -24,6 +25,11 @@ def run(args: Any, extra_args: List[str]) -> None: """ Entry point for the 'preflight' subcommand. """ + if getattr(args, "mori", False): + from primus.tools.preflight.mori_preflight import run_mori_preflight + + raise SystemExit(run_mori_preflight(args, extra_args)) + from primus.tools.preflight.preflight_perf_test import run_preflight if extra_args: @@ -43,6 +49,7 @@ def register_subcommand(subparsers): primus-cli preflight --gpu # GPU info only primus-cli preflight --network # Network info only primus-cli preflight --perf-test # Perf only + primus-cli preflight --mori # MORI runtime build + smoke """ from primus.tools.preflight.preflight_args import add_preflight_parser diff --git a/primus/tools/preflight/mori_preflight.py b/primus/tools/preflight/mori_preflight.py new file mode 100644 index 000000000..9f7255257 --- /dev/null +++ b/primus/tools/preflight/mori_preflight.py @@ -0,0 +1,417 @@ +############################################################################### +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +"""MORI per-node preflight and general N-node correctness orchestrator.""" + +from __future__ import annotations + +import concurrent.futures +import os +import shlex +import shutil +import socket +import subprocess +import sys +import time +from dataclasses import dataclass +from datetime import datetime +from pathlib import Path +from typing import Any, Sequence + +SCRIPT_DIR = Path(__file__).resolve().parent +DEFAULT_PRIMUS_ROOT = SCRIPT_DIR.parents[2] +FINGERPRINT_PREFIX = "[preflight] NODE_FINGERPRINT " + + +@dataclass +class NodeResult: + node: str + returncode: int + fingerprint: str | None + log_file: Path + + +def _dedupe(items: list[str]) -> list[str]: + seen = set() + result = [] + for item in items: + item = item.strip() + if item and item not in seen: + seen.add(item) + result.append(item) + return result + + +def resolve_nodes(spec: str | None) -> list[str]: + if not spec: + return [socket.gethostname().split(".")[0]] + if spec.startswith("@"): + return _dedupe(Path(spec[1:]).read_text().splitlines()) + if "," in spec: + return _dedupe(spec.split(",")) + if shutil.which("scontrol"): + result = subprocess.run( + ["scontrol", "show", "hostnames", spec], + text=True, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + check=False, + ) + nodes = _dedupe(result.stdout.splitlines()) + if result.returncode == 0 and nodes: + return nodes + return [spec] + + +def is_local(node: str) -> bool: + names = { + "localhost", + "127.0.0.1", + socket.gethostname(), + socket.gethostname().split(".")[0], + socket.getfqdn(), + } + return node in names + + +def preflight_container_name(node: str) -> str: + short_node = node.split(".")[0] + return f"primus_mori_preflight_{os.environ.get('USER', 'user')}_{short_node}" + + +def shell_join_env(env: dict[str, str], command: list[str]) -> str: + assignments = [f"{name}={value}" for name, value in env.items()] + return shlex.join(["env", *assignments, *command]) + + +def local_preflight_command( + args: Any, + node: str, + node_log_dir: Path, + repo_root: Path, +) -> list[str]: + env = { + "BASE_IMAGE": args.mori_base_image, + "MORI_REPO": args.mori_repo, + "MORI_REF": args.mori_ref, + "MAX_JOBS": str(args.mori_max_jobs), + "SMOKE_NUMEL": str(args.mori_smoke_numel), + # Reuse these containers for the N-node smoke and remove them after. + "KEEP_CONTAINER": "1", + "CONTAINER_NAME": preflight_container_name(node), + "LOG_DIR": str(node_log_dir), + "PRIMUS_ROOT": str(repo_root), + } + helper = str(repo_root / "primus" / "tools" / "preflight" / "mori_preflight.sh") + command = ["bash", helper] + if is_local(node): + return ["env", *[f"{k}={v}" for k, v in env.items()], *command] + + remote_command = f"cd {shlex.quote(str(repo_root))} && {shell_join_env(env, command)}" + return ["ssh", "-o", "BatchMode=yes", node, remote_command] + + +def run_node_preflight( + args: Any, + node: str, + root_log_dir: Path, + repo_root: Path, +) -> NodeResult: + short_node = node.split(".")[0] + node_log_dir = root_log_dir / "nodes" / short_node + node_log_dir.mkdir(parents=True, exist_ok=True) + launcher_log = node_log_dir / "launcher.log" + command = local_preflight_command(args, node, node_log_dir, repo_root) + fingerprint = None + + with launcher_log.open("w", encoding="utf-8") as log: + process = subprocess.Popen( + command, + text=True, + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + bufsize=1, + ) + assert process.stdout is not None + for line in process.stdout: + log.write(line) + log.flush() + print(f"[{short_node}] {line}", end="", flush=True) + if line.startswith(FINGERPRINT_PREFIX): + fingerprint = line[len(FINGERPRINT_PREFIX) :].strip() + returncode = process.wait() + + return NodeResult(node, returncode, fingerprint, launcher_log) + + +def remote_output(node: str, command: str) -> str: + if is_local(node): + argv = ["bash", "-lc", command] + else: + argv = ["ssh", "-o", "BatchMode=yes", node, command] + result = subprocess.run( + argv, + text=True, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + check=False, + ) + if result.returncode != 0: + raise RuntimeError( + f"{node}: command failed ({result.returncode}): " f"{result.stderr.strip() or command}" + ) + return result.stdout.strip() + + +def probe_network( + node: str, + interface_override: str | None, + gid_override: int | None, + address_override: str | None, +) -> tuple[str, int, str]: + if interface_override and gid_override is not None and address_override: + return interface_override, gid_override, address_override + + command = f""" +interface={shlex.quote(interface_override or "")} +gid_index={shlex.quote("" if gid_override is None else str(gid_override))} +master_addr={shlex.quote(address_override or "")} + +if [ -z "$interface" ]; then + if ip -o -4 addr show dev fenic scope global >/dev/null 2>&1; then + interface=fenic + else + for path in /sys/class/infiniband/*/device/net/*; do + [ -e "$path" ] || continue + candidate=${{path##*/}} + if ip -o -4 addr show dev "$candidate" scope global >/dev/null 2>&1; then + interface=$candidate + break + fi + done + fi +fi +[ -n "$interface" ] || {{ echo "no RDMA interface with IPv4" >&2; exit 1; }} + +if [ -z "$gid_index" ]; then + gid_index=0 + for dev in /sys/class/infiniband/*; do + [ -d "$dev" ] || continue + for path in "$dev"/ports/1/gids/*; do + [ -f "$path" ] || continue + idx=${{path##*/}} + gid=$(cat "$path") + type=$(cat "$dev/ports/1/gid_attrs/types/$idx" 2>/dev/null || true) + case "$type:$gid" in + "RoCE v2:"*":ffff:"*) gid_index=$idx; break 2 ;; + esac + done + done +fi + +if [ -z "$master_addr" ]; then + master_addr=$(ip -o -4 addr show dev "$interface" scope global | + awk 'NR==1{{split($4,a,"/"); print a[1]}}') +fi +[ -n "$master_addr" ] || {{ echo "no IPv4 address on $interface" >&2; exit 1; }} +printf '%s\\t%s\\t%s\\n' "$interface" "$gid_index" "$master_addr" +""" + fields = remote_output(node, command).split("\t") + if len(fields) != 3: + raise RuntimeError(f"{node}: malformed network probe output: {fields}") + return fields[0], int(fields[1]), fields[2] + + +def validate_fingerprints(results: list[NodeResult]) -> None: + missing = [result.node for result in results if not result.fingerprint] + if missing: + raise RuntimeError(f"Missing MORI fingerprint from nodes: {', '.join(missing)}") + grouped: dict[str, list[str]] = {} + for result in results: + assert result.fingerprint is not None + grouped.setdefault(result.fingerprint, []).append(result.node) + if len(grouped) != 1: + detail = "; ".join(f"{nodes}: {fingerprint}" for fingerprint, nodes in grouped.items()) + raise RuntimeError(f"MORI node-stack mismatch: {detail}") + + +def run_multinode(args: Any, nodes: list[str], log_dir: Path) -> int: + master_node = nodes[0] + interface, gid_index, master_addr = probe_network( + master_node, + args.mori_socket_ifname, + args.mori_gid_index, + args.mori_master_addr, + ) + multinode_log_dir = log_dir / "multinode" + multinode_log_dir.mkdir(parents=True, exist_ok=True) + + def launch(node: str, rank: int) -> tuple[int, Path]: + torchrun = shlex.join( + [ + "torchrun", + f"--nnodes={len(nodes)}", + "--nproc_per_node=8", + f"--node_rank={rank}", + f"--master_addr={master_addr}", + f"--master_port={args.mori_master_port}", + "/src/primus/runner/helpers/mori/multinode_allgather_smoke.py", + "--numel", + str(args.mori_smoke_numel), + ] + ) + docker_command = [ + "docker", + "exec", + "-e", + "PYTHONPATH=/src/primus", + "-e", + f"NCCL_SOCKET_IFNAME={interface}", + "-e", + f"NCCL_IB_GID_INDEX={gid_index}", + preflight_container_name(node), + "bash", + "-lc", + "export LD_LIBRARY_PATH=/opt/mori-host-libs:${LD_LIBRARY_PATH}; " f"exec {torchrun}", + ] + command = ( + docker_command + if is_local(node) + else ["ssh", "-o", "BatchMode=yes", node, shlex.join(docker_command)] + ) + rank_log = multinode_log_dir / f"node-{rank}-{node}.log" + with rank_log.open("w", encoding="utf-8") as output: + result = subprocess.run( + command, + stdout=output, + stderr=subprocess.STDOUT, + check=False, + ) + return result.returncode, rank_log + + print( + f"[MORI:Preflight] launching {len(nodes)} nodes via {master_node} " + f"({master_addr}, {interface}, gid={gid_index})", + flush=True, + ) + failures = [] + with concurrent.futures.ThreadPoolExecutor(max_workers=len(nodes)) as executor: + futures = {} + for rank, node in enumerate(nodes): + futures[executor.submit(launch, node, rank)] = node + if rank == 0: + time.sleep(1) + for future, node in futures.items(): + returncode, rank_log = future.result() + if returncode: + failures.append((node, rank_log)) + for node, rank_log in failures: + print(f"[MORI:Preflight] FAIL {node}: {rank_log}", file=sys.stderr) + return int(bool(failures)) + + +def remove_preflight_containers(nodes: list[str]) -> None: + def remove(node: str) -> None: + name = preflight_container_name(node) + remote_output( + node, + f"docker rm -f {shlex.quote(name)} >/dev/null 2>&1 || true", + ) + + with concurrent.futures.ThreadPoolExecutor(max_workers=len(nodes)) as executor: + list(executor.map(remove, nodes)) + + +def run_orchestrator(args: Any, repo_root: Path, log_dir: Path) -> int: + nodes = resolve_nodes(args.mori_nodes) + log_dir.mkdir(parents=True, exist_ok=True) + print(f"[MORI:Preflight] nodes ({len(nodes)}): {', '.join(nodes)}") + print(f"[MORI:Preflight] logs: {log_dir}") + + try: + results = [] + with concurrent.futures.ThreadPoolExecutor(max_workers=len(nodes)) as executor: + future_map = { + executor.submit(run_node_preflight, args, node, log_dir, repo_root): node for node in nodes + } + for future in concurrent.futures.as_completed(future_map): + results.append(future.result()) + + results.sort(key=lambda result: nodes.index(result.node)) + failed = [result for result in results if result.returncode != 0] + if failed: + for result in failed: + print( + f"[MORI:Preflight] FAIL {result.node}: {result.log_file}", + file=sys.stderr, + ) + return 1 + + validate_fingerprints(results) + print("[MORI:Preflight] all node fingerprints match", flush=True) + if len(nodes) > 1: + return run_multinode(args, nodes, log_dir) + return 0 + finally: + if not args.mori_keep_container: + remove_preflight_containers(nodes) + + +def _validate_mori_mode(args: Any, extra_args: Sequence[str]) -> int | None: + if extra_args: + print( + f"[Primus:Preflight] ERROR: unknown MORI arguments: {list(extra_args)}", + file=sys.stderr, + ) + return 2 + + incompatible = [] + for attr, flag in ( + ("check_host", "--host"), + ("check_gpu", "--gpu"), + ("check_network", "--network"), + ("perf_test", "--perf-test"), + ("plot", "--plot"), + ("tests", "--tests"), + ("comm_sizes_mb", "--comm-sizes-mb"), + ("intra_comm_sizes_mb", "--intra-comm-sizes-mb"), + ("inter_comm_sizes_mb", "--inter-comm-sizes-mb"), + ("intra_group_sizes", "--intra-group-sizes"), + ("inter_group_sizes", "--inter-group-sizes"), + ("ring_p2p_sizes_mb", "--ring-p2p-sizes-mb"), + ("quick", "--quick"), + ): + if getattr(args, attr, None): + incompatible.append(flag) + if getattr(args, "split_nodes_subgroup", True) is False: + incompatible.append("--no-split-nodes-subgroup") + if incompatible: + print( + "[Primus:Preflight] ERROR: --mori cannot be combined with " + ", ".join(incompatible), + file=sys.stderr, + ) + return 2 + return None + + +def run_mori_preflight(args: Any, extra_args: Sequence[str] = ()) -> int: + validation_rc = _validate_mori_mode(args, extra_args) + if validation_rc is not None: + return validation_rc + + repo_root = Path( + os.environ.get( + "PRIMUS_PATH", + str(DEFAULT_PRIMUS_ROOT), + ) + ).resolve() + timestamp = datetime.now().strftime("%Y%m%d-%H%M%S") + short_host = socket.gethostname().split(".")[0] + report_name = args.report_file_name or f"mori-preflight-{short_host}-{timestamp}" + log_dir = Path(args.mori_log_dir or (Path(args.dump_path) / report_name)).resolve() + + print(f"[Primus:Preflight] MORI logs: {log_dir}", flush=True) + return run_orchestrator(args, repo_root, log_dir) diff --git a/primus/tools/preflight/mori_preflight.sh b/primus/tools/preflight/mori_preflight.sh new file mode 100755 index 000000000..9220aaa80 --- /dev/null +++ b/primus/tools/preflight/mori_preflight.sh @@ -0,0 +1,419 @@ +#!/usr/bin/env bash +############################################################################### +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +set -Eeuo pipefail + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +PRIMUS_ROOT="${PRIMUS_ROOT:-$(cd "${SCRIPT_DIR}/../../.." && pwd)}" + +BASE_IMAGE="${BASE_IMAGE:-unifiedtrainingdockers.azurecr.io/utd/nightly:primus_the_rock_rocm7.15_20260728}" +MORI_REPO="${MORI_REPO:-https://github.com/ROCm/mori.git}" +MORI_REF="${MORI_REF:-12d1bc32d0c93dcd5062e74f4e0f772e36e1aac4}" +MAX_JOBS="${MAX_JOBS:-32}" +SMOKE_NUMEL="${SMOKE_NUMEL:-67108864}" +KEEP_CONTAINER="${KEEP_CONTAINER:-0}" +LOG_DIR="${LOG_DIR:-/tmp/primus-mori-preflight-$(hostname -s)-$(date +%Y%m%d-%H%M%S)}" +CONTAINER_NAME="${CONTAINER_NAME:-primus_mori_preflight_${USER}_$(hostname -s)}" + +mkdir -p "${LOG_DIR}" + +CURRENT_PHASE="initialization" +PHASE_SUMMARY=() + +section() { + echo + echo "================================================================================" + echo "$*" + echo "================================================================================" +} + +format_seconds() { + local seconds="$1" + printf "%dm%02ds" "$((seconds / 60))" "$((seconds % 60))" +} + +run_phase() { + local name="$1" + shift + local start end elapsed rc + CURRENT_PHASE="${name}" + start="$(date +%s)" + section "Phase: ${name}" + + set +e + "$@" 2>&1 | tee "${LOG_DIR}/${name// /_}.log" + rc="${PIPESTATUS[0]}" + set -e + + end="$(date +%s)" + elapsed="$((end - start))" + PHASE_SUMMARY+=("${name}|${elapsed}|${rc}") + echo "[preflight] ${name}: $(format_seconds "${elapsed}") (rc=${rc})" + if [[ "${rc}" -ne 0 ]]; then + return "${rc}" + fi +} + +cleanup() { + local rc=$? + if [[ "${rc}" -ne 0 ]]; then + echo + echo "[preflight] FAILED during phase: ${CURRENT_PHASE}" >&2 + echo "[preflight] Logs: ${LOG_DIR}" >&2 + docker inspect "${CONTAINER_NAME}" >"${LOG_DIR}/container-inspect.json" 2>/dev/null || true + fi + + if [[ "${KEEP_CONTAINER}" != "1" ]]; then + docker rm -f "${CONTAINER_NAME}" >/dev/null 2>&1 || true + else + echo "[preflight] Keeping container: ${CONTAINER_NAME}" + fi +} +trap cleanup EXIT INT TERM + +command -v docker >/dev/null || { + echo "[preflight] docker is required" >&2 + exit 1 +} + +section "Host identity" +echo "hostname : $(hostname -f)" +echo "kernel : $(uname -r)" +echo "user : $(id)" +echo "base image : ${BASE_IMAGE}" +echo "MORI revision : ${MORI_REF}" +echo "Primus root : ${PRIMUS_ROOT}" +echo "log directory : ${LOG_DIR}" + +section "GPU information" +if command -v rocm-smi >/dev/null; then + rocm-smi --showproductname --showuse --showmemuse --csv 2>&1 || true +else + echo "rocm-smi not found" +fi +if command -v rocminfo >/dev/null; then + rocminfo 2>/dev/null | + awk '/Name: +gfx/{print "GPU architecture : " $2}' | + sort -u || true +fi + +section "IP interfaces" +ip -o -4 addr show scope global 2>&1 || true + +section "RDMA devices and links" +if command -v ibv_devices >/dev/null; then + ibv_devices 2>&1 || true +fi +if command -v ibdev2netdev >/dev/null; then + ibdev2netdev 2>&1 || true +fi +if command -v rdma >/dev/null; then + rdma link show 2>&1 || true +fi + +section "Valid GIDs" +for dev_path in /sys/class/infiniband/*; do + [[ -d "${dev_path}" ]] || continue + dev="${dev_path##*/}" + for gid_path in "${dev_path}"/ports/1/gids/*; do + [[ -f "${gid_path}" ]] || continue + idx="${gid_path##*/}" + gid="$(<"${gid_path}")" + [[ "${gid}" != "0000:0000:0000:0000:0000:0000:0000:0000" ]] || continue + type="unknown" + ndev="unknown" + if [[ -r "${dev_path}/ports/1/gid_attrs/types/${idx}" ]]; then + type="$(<"${dev_path}/ports/1/gid_attrs/types/${idx}")" + fi + if [[ -r "${dev_path}/ports/1/gid_attrs/ndevs/${idx}" ]]; then + ndev="$(<"${dev_path}/ports/1/gid_attrs/ndevs/${idx}")" + fi + printf "%-12s index=%-3s type=%-8s netdev=%-12s gid=%s\n" \ + "${dev}" "${idx}" "${type}" "${ndev}" "${gid}" + done +done + +section "NIC driver and firmware" +mapfile -t NETDEVS < <( + for dev_path in /sys/class/infiniband/*; do + [[ -d "${dev_path}" ]] || continue + for ndev_path in "${dev_path}"/device/net/*; do + [[ -e "${ndev_path}" ]] && basename "${ndev_path}" + done + done | sort -u +) +STACK_DATA="" +for netdev in "${NETDEVS[@]:-}"; do + driver_info="$(ethtool -i "${netdev}" 2>&1 || true)" + echo "--- ${netdev} ---" + echo "${driver_info}" + driver_version="$(awk -F': ' '$1=="version"{print $2}' <<<"${driver_info}")" + firmware_version="$(awk -F': ' '$1=="firmware-version"{print $2}' <<<"${driver_info}")" + STACK_DATA+="${netdev}|${driver_version}|${firmware_version};" +done +STACK_SHA="$(printf "%s" "${STACK_DATA}" | sha256sum | awk '{print $1}')" + +detect_nic() { + if compgen -G "/sys/class/infiniband/ionic*" >/dev/null; then + echo ionic + elif compgen -G "/sys/class/infiniband/bnxt_re*" >/dev/null; then + echo bnxt + elif compgen -G "/sys/class/infiniband/mlx5*" >/dev/null; then + echo mlx5 + else + echo unknown + fi +} + +DETECTED_NIC="$(detect_nic)" + +section "Vendor library checks (detected NIC: ${DETECTED_NIC})" + +find_library() { + local name="$1" + local path + path="$(ldconfig -p 2>/dev/null | awk -v n="${name}" '$1==n{print $NF; exit}')" + if [[ -z "${path}" ]]; then + for candidate in \ + "/usr/local/lib/${name}" \ + "/usr/lib/x86_64-linux-gnu/${name}" \ + "/lib/x86_64-linux-gnu/${name}"; do + if [[ -e "${candidate}" ]]; then + path="${candidate}" + break + fi + done + fi + echo "${path}" +} + +has_symbol() { + local path="$1" + local symbol="$2" + nm -D "${path}" 2>/dev/null | + awk -v s="${symbol}" '$3 == s || index($3, s "@") == 1 {found=1} END {exit !found}' +} + +check_symbols() { + local path="$1" + shift + local symbol + for symbol in "$@"; do + if has_symbol "${path}" "${symbol}"; then + echo " ${symbol}: present" + else + echo " ${symbol}: MISSING" + fi + done +} + +ionic_ccqe_fw_supported() { + local value="$1" + if [[ ! "${value}" =~ ^([0-9]+)\.([0-9]+)\.([0-9]+)-[[:alpha:]]-?([0-9]+)$ ]]; then + return 1 + fi + local major="${BASH_REMATCH[1]}" + local minor="${BASH_REMATCH[2]}" + local patch="${BASH_REMATCH[3]}" + local build="${BASH_REMATCH[4]}" + ((major > 1 || + (major == 1 && minor > 117) || + (major == 1 && minor == 117 && patch > 5) || + (major == 1 && minor == 117 && patch == 5 && build >= 58))) +} + +VENDOR_LIB="" +VENDOR_NAMES=() +VENDOR_SHA="missing" +CCQE_CANDIDATE="n/a" +case "${DETECTED_NIC}" in + ionic) + VENDOR_LIB="$(find_library libionic.so)" + VENDOR_NAMES=(libionic.so) + ;; + bnxt) + VENDOR_LIB="$(find_library libbnxt_re.so)" + VENDOR_NAMES=(libbnxt_re.so libbnxt_re-rdmav34.so) + ;; + mlx5) + VENDOR_LIB="$(find_library libmlx5.so)" + VENDOR_NAMES=(libmlx5.so) + ;; +esac + +if [[ -n "${VENDOR_LIB}" ]]; then + VENDOR_LIB="$(readlink -f "${VENDOR_LIB}")" + VENDOR_SHA="$(sha256sum "${VENDOR_LIB}" | awk '{print $1}')" + echo "library : ${VENDOR_LIB}" + echo "sha256 : ${VENDOR_SHA}" + case "${DETECTED_NIC}" in + ionic) + check_symbols "${VENDOR_LIB}" \ + ionic_dv_get_ctx \ + ionic_dv_get_cq \ + ionic_dv_get_qp \ + ionic_dv_create_cq_ex + ;; + bnxt) + check_symbols "${VENDOR_LIB}" \ + bnxt_re_dv_umem_reg \ + bnxt_re_dv_umem_dereg \ + bnxt_re_dv_create_cq \ + bnxt_re_dv_destroy_cq \ + bnxt_re_dv_init_obj \ + bnxt_re_dv_create_qp \ + bnxt_re_dv_destroy_qp \ + bnxt_re_dv_modify_qp + ;; + esac +else + echo "vendor library : not found" +fi + +if [[ "${DETECTED_NIC}" == "ionic" ]]; then + section "Ionic CCQE prerequisites" + ccqe_total=0 + ccqe_fw_eligible=0 + for device_path in /sys/class/infiniband/ionic*; do + [[ -d "${device_path}" ]] || continue + device="${device_path##*/}" + firmware="unknown" + [[ -r "${device_path}/fw_ver" ]] && firmware="$(<"${device_path}/fw_ver")" + ((ccqe_total += 1)) + if ionic_ccqe_fw_supported "${firmware}"; then + ((ccqe_fw_eligible += 1)) + echo " ${device}: firmware=${firmware} eligible=true" + else + echo " ${device}: firmware=${firmware} eligible=false" + fi + done + + ccqe_symbol="false" + if [[ -n "${VENDOR_LIB}" ]] && has_symbol "${VENDOR_LIB}" ionic_dv_create_cq_ex; then + ccqe_symbol="true" + fi + if [[ "${ccqe_symbol}" == "true" && "${ccqe_total}" -gt 0 && "${ccqe_fw_eligible}" -eq "${ccqe_total}" ]]; then + CCQE_CANDIDATE="true" + elif [[ "${ccqe_symbol}" == "true" && "${ccqe_fw_eligible}" -gt 0 ]]; then + CCQE_CANDIDATE="mixed" + else + CCQE_CANDIDATE="false" + fi + echo " firmware eligible : ${ccqe_fw_eligible}/${ccqe_total}" + echo " create_cq_ex symbol: ${ccqe_symbol}" + echo " host candidate : ${CCQE_CANDIDATE}" + echo " effective runtime : not tested (requires RDMA CQ creation)" +fi + +run_phase "pull base image" docker pull "${BASE_IMAGE}" + +docker rm -f "${CONTAINER_NAME}" >/dev/null 2>&1 || true + +DOCKER_ARGS=( + docker run -d + --name "${CONTAINER_NAME}" + --device=/dev/kfd + --device=/dev/dri + --group-add video + --cap-add SYS_PTRACE + --security-opt seccomp=unconfined + --privileged + --ipc=host + --network=host + -e PYTHONDONTWRITEBYTECODE=1 + -v "${PRIMUS_ROOT}:/src/primus:ro" +) + +if [[ -n "${VENDOR_LIB}" ]]; then + for name in "${VENDOR_NAMES[@]}"; do + DOCKER_ARGS+=(-v "${VENDOR_LIB}:/opt/mori-host-libs/${name}:ro") + done +fi + +DOCKER_ARGS+=("${BASE_IMAGE}" sleep infinity) +run_phase "start container" "${DOCKER_ARGS[@]}" + +# shellcheck disable=SC2016 +run_phase "container network check" \ + docker exec "${CONTAINER_NAME}" bash -lc ' + set -e + echo "torch=$(python3 -c "import torch; print(torch.__version__)")" + echo "ROCM_PATH=${ROCM_PATH}" + ip -o -4 addr show scope global || true + ibv_devices || true + export LD_LIBRARY_PATH=/opt/mori-host-libs:${LD_LIBRARY_PATH} + python3 - <<'"'"'PY'"'"' +import ctypes +for name in ("libionic.so", "libbnxt_re.so", "libmlx5.so"): + try: + ctypes.CDLL(name) + print(f"{name}: loadable") + except OSError as exc: + print(f"{name}: unavailable ({exc})") +PY + ' + +run_phase "install MORI" \ + docker exec \ + -e MORI_REPO="${MORI_REPO}" \ + -e MORI_REF="${MORI_REF}" \ + -e MAX_JOBS="${MAX_JOBS}" \ + "${CONTAINER_NAME}" \ + bash /src/primus/runner/helpers/mori/install_mori.sh + +# Use MORI's own runtime detector so fingerprints catch cases where identical +# host firmware produces different CCQE decisions inside the training image. +# shellcheck disable=SC2016 +run_phase "MORI runtime capability check" \ + docker exec "${CONTAINER_NAME}" bash -lc ' + export LD_LIBRARY_PATH="/opt/mori-host-libs:${LD_LIBRARY_PATH}" + python3 - <<'"'"'PY'"'"' +from mori.jit.core import detect_nic_type, is_ccqe_enabled + +nic = detect_nic_type() +ccqe = str(is_ccqe_enabled()).lower() if nic == "ionic" else "n/a" +print(f"[preflight] MORI_CAPABILITY nic={nic} ccqe_runtime={ccqe}") +PY + ' + +CCQE_RUNTIME="n/a" +if [[ "${DETECTED_NIC}" == "ionic" ]]; then + CCQE_RUNTIME="$( + awk -F'ccqe_runtime=' '/MORI_CAPABILITY/ {print $2; exit}' \ + "${LOG_DIR}/MORI_runtime_capability_check.log" + )" + CCQE_RUNTIME="${CCQE_RUNTIME:-unknown}" +fi + +# shellcheck disable=SC2016 +run_phase "8-GPU MORI all-gather smoke" \ + docker exec \ + -e PYTHONPATH=/src/primus \ + -e NCCL_SOCKET_IFNAME=lo \ + -e SMOKE_NUMEL="${SMOKE_NUMEL}" \ + "${CONTAINER_NAME}" bash -lc ' + set -e + export LD_LIBRARY_PATH="/opt/mori-host-libs:${LD_LIBRARY_PATH}" + torchrun --standalone --nproc_per_node=8 \ + /src/primus/runner/helpers/mori/multinode_allgather_smoke.py \ + --numel "${SMOKE_NUMEL}" + ' + +echo "[preflight] NODE_FINGERPRINT nic=${DETECTED_NIC} stack_sha=${STACK_SHA} vendor_sha=${VENDOR_SHA} ccqe_candidate=${CCQE_CANDIDATE} ccqe_runtime=${CCQE_RUNTIME}" + +section "Timing summary" +total=0 +for entry in "${PHASE_SUMMARY[@]}"; do + IFS="|" read -r name elapsed rc <<<"${entry}" + printf "%-32s %8s rc=%s\n" "${name}" "$(format_seconds "${elapsed}")" "${rc}" + total="$((total + elapsed))" +done +printf "%-32s %8s\n" "TOTAL" "$(format_seconds "${total}")" + +echo +echo "[preflight] PASS" +echo "[preflight] Logs: ${LOG_DIR}" diff --git a/primus/tools/preflight/preflight_args.py b/primus/tools/preflight/preflight_args.py index 1710e14b4..9db5c47f7 100644 --- a/primus/tools/preflight/preflight_args.py +++ b/primus/tools/preflight/preflight_args.py @@ -49,6 +49,7 @@ def add_preflight_parser(parser: argparse.ArgumentParser) -> argparse.ArgumentPa primus-cli preflight --gpu # GPU info only primus-cli preflight --network # Network info only primus-cli preflight --gpu --network # GPU + Network info + primus-cli preflight --mori # MORI runtime build + smoke primus-cli preflight --perf-test # Perf only, all tests primus-cli preflight --quick # Perf only, fast preset primus-cli preflight --tests gemm # Perf only, GEMM only @@ -229,4 +230,81 @@ def add_preflight_parser(parser: argparse.ArgumentParser) -> argparse.ArgumentPa action="store_false", help="Disable PDF report generation.", ) + + # MORI preflight is an exclusive orchestration mode. It runs once per node + # (not once per GPU), builds MORI inside a privileged runtime container, + # and executes local plus optional cross-node all-gather correctness tests. + mori = parser.add_argument_group("MORI runtime preflight") + mori.add_argument( + "--mori", + action="store_true", + help="Run MORI NIC diagnostics, runtime source build, and all-gather smoke. " + "Cannot be combined with standard info/perf selectors.", + ) + mori.add_argument( + "--mori-base-image", + default=("unifiedtrainingdockers.azurecr.io/utd/nightly:" "primus_the_rock_rocm7.15_20260728"), + help="Base image pulled by MORI preflight.", + ) + mori.add_argument( + "--mori-repo", + default="https://github.com/ROCm/mori.git", + help="MORI git repository.", + ) + mori.add_argument( + "--mori-ref", + default="12d1bc32d0c93dcd5062e74f4e0f772e36e1aac4", + help="Pinned MORI git revision to build.", + ) + mori.add_argument( + "--mori-max-jobs", + type=int, + default=32, + help="Maximum parallel MORI build jobs.", + ) + mori.add_argument( + "--mori-smoke-numel", + type=int, + default=67108864, + help="BF16 elements per rank in MORI all-gather smokes " "(default: 67108864 = 128 MiB).", + ) + mori.add_argument( + "--mori-keep-container", + action="store_true", + help="Keep the temporary build container after preflight.", + ) + mori.add_argument( + "--mori-log-dir", + default=None, + help="MORI phase-log directory. Defaults under --dump-path.", + ) + mori.add_argument( + "--mori-nodes", + default=None, + help="General multi-node target: comma-separated hosts, Slurm hostlist, " + "or @file. Each node runs full diagnostics/build/local smoke before " + "the N-node all-gather.", + ) + mori.add_argument( + "--mori-master-addr", + default=None, + help="Master IP for the N-node smoke (auto-detected when unset).", + ) + mori.add_argument( + "--mori-master-port", + type=int, + default=29610, + help="torchrun rendezvous port for the N-node MORI smoke.", + ) + mori.add_argument( + "--mori-socket-ifname", + default=None, + help="Bootstrap interface for the N-node smoke.", + ) + mori.add_argument( + "--mori-gid-index", + type=int, + default=None, + help="RCCL RoCEv2 GID index for the N-node smoke.", + ) return parser diff --git a/primus/tools/preflight/preflight_perf_test.py b/primus/tools/preflight/preflight_perf_test.py index 636dee728..fd97ad639 100644 --- a/primus/tools/preflight/preflight_perf_test.py +++ b/primus/tools/preflight/preflight_perf_test.py @@ -415,14 +415,20 @@ def run_preflight(args): Mode precedence (single rule): - 1. Any of --perf-test / --tests / --quick is set -> perf-only mode. + 1. --mori is set -> exclusive MORI runtime preflight mode. + 2. Any of --perf-test / --tests / --quick is set -> perf-only mode. If info selectors (--host/--gpu/--network) are also present, they are dropped with a WARN. - 2. Otherwise, any of --host/--gpu/--network is set -> info-only mode. + 3. Otherwise, any of --host/--gpu/--network is set -> info-only mode. Perf tuning knobs (e.g. --comm-sizes-mb), if set, are inert and a WARN is emitted. - 3. Otherwise (no flags) -> default: info AND all perf tests. + 4. Otherwise (no flags) -> default: info AND all perf tests. """ + if getattr(args, "mori", False): + from primus.tools.preflight.mori_preflight import run_mori_preflight + + return run_mori_preflight(args) + # R4: canonical normalization point for the report file name. Done here # before any downstream code reads args.report_file_name, so every code # path (info-only, perf-only, info+perf, dist-init failure) sees the same diff --git a/pyproject.toml b/pyproject.toml index 4851a3253..5d772f12e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -94,7 +94,10 @@ path = "primus/__init__.py" packages = ["primus"] # _thirdparty.lock is generated at build time (and git-ignored); force-include it # so the wheel ships it for `primus-cli deps sync`. -artifacts = ["primus/_thirdparty.lock"] +artifacts = [ + "primus/_thirdparty.lock", + "primus/tools/preflight/mori_preflight.sh", +] [tool.hatch.build.targets.wheel.force-include] # Ship the whole primus-cli bash toolkit as package data under primus/runner/. diff --git a/runner/helpers/hooks/06_enable_sdma_all_gather.sh b/runner/helpers/hooks/06_enable_sdma_all_gather.sh index 7364fdbf4..8f0ab15e6 100755 --- a/runner/helpers/hooks/06_enable_sdma_all_gather.sh +++ b/runner/helpers/hooks/06_enable_sdma_all_gather.sh @@ -8,15 +8,15 @@ # Global hook: opt into the SDMA/RCCL dispatch path for FSDP # all-gather. # -# Single trigger -- the only knob: +# Backend selector: # -# export SDMA_ALL_GATHER=1 +# export FSDP_ALL_GATHER_BACKEND=rccl_sdma # primus-cli direct -- train pretrain --config # -# When SDMA_ALL_GATHER=1, this hook: +# When the backend is rccl_sdma, this hook: # 1. Exports the zero-CTA env that RCCL needs to actually take the # copy-engine path (NCCL_CTA_POLICY=2, NCCL_CUMEM_ENABLE=1, ...). -# 2. Propagates SDMA_ALL_GATHER=1 into the launched torchrun children +# 2. Propagates FSDP_ALL_GATHER_BACKEND=rccl_sdma into torchrun children # so the companion Python patch's gate fires there too. See # primus/backends/torchtitan/patches/sdma_symm_mem_collectives.py. # 3. Rebuilds the bundled LD_PRELOAD interposer @@ -33,7 +33,7 @@ set -euo pipefail -if [[ "${SDMA_ALL_GATHER:-0}" != "1" ]]; then +if [[ "${FSDP_ALL_GATHER_BACKEND:-}" != "rccl_sdma" ]]; then exit 0 fi @@ -48,11 +48,11 @@ echo "env.NCCL_CUMEM_ENABLE=1" echo "env.NCCL_LOCAL_REGISTER=0" echo "env.TORCH_NCCL_USE_TENSOR_REGISTER_ALLOCATOR_HOOK=true" -# 2) Make the trigger visible to torchrun children so the Python patch +# 2) Make the selector visible to torchrun children so the Python patch # fires there. primus-cli direct doesn't inherit the host env into # the child unless it's either CLI-passed via --env or hook-emitted # via env.*. -echo "env.SDMA_ALL_GATHER=1" +echo "env.FSDP_ALL_GATHER_BACKEND=rccl_sdma" # 3) Always (re)build the interposer. The source is tiny and gcc is # typically <1s; we don't bother with a staleness check so the .so diff --git a/runner/helpers/hooks/07_enable_mori_all_gather.sh b/runner/helpers/hooks/07_enable_mori_all_gather.sh new file mode 100755 index 000000000..d16891eb5 --- /dev/null +++ b/runner/helpers/hooks/07_enable_mori_all_gather.sh @@ -0,0 +1,60 @@ +#!/bin/bash +############################################################################### +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### +# +# Global hook: opt into MORI FSDP all-gather. +# +# Trigger: +# +# export FSDP_ALL_GATHER_BACKEND=mori +# primus-cli direct -- train pretrain --config +# +# When enabled, this hook installs MORI when needed and propagates MORI_* env +# vars into torchrun children. The Python patches then attach MoriAllGather to +# FSDP2 modules. + +set -euo pipefail + +if [[ "${FSDP_ALL_GATHER_BACKEND:-}" != "mori" ]]; then + exit 0 +fi + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +mori_installer="${SCRIPT_DIR}/../mori/install_mori.sh" + +if ! python3 -c "from mori.ccl import HierAllGather" >/dev/null 2>&1; then + echo "[MORI] HierAllGather is unavailable; installing MORI before launch." >&2 + if [[ ! -x "${mori_installer}" ]]; then + echo "[ERROR] MORI installer is not executable: ${mori_installer}" >&2 + exit 1 + fi + "${mori_installer}" >&2 + if ! python3 -c "from mori.ccl import HierAllGather" >/dev/null 2>&1; then + echo "[ERROR] MORI installation completed but HierAllGather is unavailable." >&2 + exit 1 + fi +fi + +# MORI all-gather uses SDMA for the intra-node leg. Allow an explicit caller +# value to win, but default it on for this feature. +export MORI_ENABLE_SDMA="${MORI_ENABLE_SDMA:-1}" +export MORI_SHMEM_HEAP_SIZE="${MORI_SHMEM_HEAP_SIZE:-8G}" + +# MORI's single-node eager path is the correctness-safe default on the ROCm +# versions used by Primus v26.4. Explicit user settings still win. +export MORI_HIER_CUDA_GRAPH="${MORI_HIER_CUDA_GRAPH:-0}" + +if [[ -z "${MORI_SOCKET_IFNAME:-}" && -n "${NCCL_SOCKET_IFNAME:-}" ]]; then + export MORI_SOCKET_IFNAME="${NCCL_SOCKET_IFNAME#=}" +fi + +# primus-cli direct does not implicitly propagate host env into torchrun +# children. Emit every MORI_* variable as env.* so user-selected MORI tuning +# knobs (host-proxy, RDMA devices, async, graph/debug flags, etc.) survive. +echo "env.FSDP_ALL_GATHER_BACKEND=mori" +for name in "${!MORI_@}"; do + echo "env.${name}=${!name}" +done diff --git a/runner/helpers/mori/install_mori.sh b/runner/helpers/mori/install_mori.sh new file mode 100755 index 000000000..4f32e5c8b --- /dev/null +++ b/runner/helpers/mori/install_mori.sh @@ -0,0 +1,57 @@ +#!/usr/bin/env bash +############################################################################### +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +set -euo pipefail + +MORI_REPO="${MORI_REPO:-https://github.com/ROCm/mori.git}" +MORI_REF="${MORI_REF:-12d1bc32d0c93dcd5062e74f4e0f772e36e1aac4}" +MORI_SOURCE_DIR="${MORI_SOURCE_DIR:-/opt/mori}" +MAX_JOBS="${MAX_JOBS:-32}" +ROCM_PATH="${ROCM_PATH:-/opt/rocm}" + +if [[ "$(id -u)" -ne 0 ]]; then + echo "[MORI:Install] run as root inside the training container" >&2 + exit 1 +fi +if [[ -z "${MORI_SOURCE_DIR}" || "${MORI_SOURCE_DIR}" == "/" ]]; then + echo "[MORI:Install] unsafe MORI_SOURCE_DIR=${MORI_SOURCE_DIR@Q}" >&2 + exit 1 +fi +if [[ ! -d "${ROCM_PATH}" ]]; then + echo "[MORI:Install] ROCM_PATH does not exist: ${ROCM_PATH}" >&2 + exit 1 +fi + +echo "[MORI:Install] install dependencies" +apt-get update +apt-get install -y --no-install-recommends \ + git ibverbs-utils libibverbs-dev libnuma-dev libpci-dev +python3 -m pip install --no-cache-dir \ + "packaging<26" "setuptools==81.0.0" setuptools_scm Cython pybind11 ninja + +echo "[MORI:Install] clone ${MORI_REPO}@${MORI_REF}" +rm -rf "${MORI_SOURCE_DIR}" +git clone --filter=blob:none "${MORI_REPO}" "${MORI_SOURCE_DIR}" +cd "${MORI_SOURCE_DIR}" +git checkout "${MORI_REF}" +git submodule update --init --depth 1 3rdparty/msgpack-c 3rdparty/spdlog + +if [[ ! -e /usr/lib64/libc.so ]]; then + mkdir -p /usr/lib64 + ln -s /usr/lib/x86_64-linux-gnu/libc.so /usr/lib64/libc.so +fi + +echo "[MORI:Install] build and verify" +export ROCM_PATH +export CMAKE_PREFIX_PATH="${ROCM_PATH}:${ROCM_PATH}/lib/rocm_sysdeps${CMAKE_PREFIX_PATH:+:${CMAKE_PREFIX_PATH}}" +export CMAKE_LIBRARY_PATH="/opt/mori-host-libs${CMAKE_LIBRARY_PATH:+:${CMAKE_LIBRARY_PATH}}" +export LD_LIBRARY_PATH="/opt/mori-host-libs:${ROCM_PATH}/lib${LD_LIBRARY_PATH:+:${LD_LIBRARY_PATH}}" +unset MORI_DEVICE_NIC +MAX_JOBS="${MAX_JOBS}" python3 -m pip install --no-build-isolation . +python3 -c \ + "import mori; from mori.ccl import HierAllGather; print(f'MORI {mori.__version__}: {HierAllGather}')" +echo "[MORI:Install] PASS" diff --git a/runner/helpers/mori/multinode_allgather_smoke.py b/runner/helpers/mori/multinode_allgather_smoke.py new file mode 100644 index 000000000..a8804e2b8 --- /dev/null +++ b/runner/helpers/mori/multinode_allgather_smoke.py @@ -0,0 +1,73 @@ +#!/usr/bin/env python3 +"""Two-node correctness smoke for Primus's MORI FSDP all-gather adapter.""" + +import argparse +import os + +import torch +import torch.distributed as dist + +from primus.backends.common.mori_allgather import MoriAllGather + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--numel", type=int, default=64 * 1024 * 1024) + args = parser.parse_args() + + local_rank = int(os.environ["LOCAL_RANK"]) + torch.cuda.set_device(local_rank) + device = torch.device("cuda", local_rank) + # Keep RCCL lazy: this smoke validates MORI's RDMA path, and no torch + # collective is needed before MORI initializes from the default c10d store. + dist.init_process_group("nccl") + + rank = dist.get_rank() + world_size = dist.get_world_size() + ranks_per_node = int(os.environ["LOCAL_WORLD_SIZE"]) + + input_tensor = torch.full( + (args.numel,), + rank + 1, + dtype=torch.bfloat16, + device=device, + ) + output_tensor = torch.empty( + args.numel * world_size, + dtype=input_tensor.dtype, + device=device, + ) + + work = MoriAllGather(ranks_per_node=ranks_per_node)( + output_tensor, + input_tensor, + dist.group.WORLD, + async_op=True, + ) + if work is not None: + work.wait() + torch.cuda.synchronize(device) + + expected = torch.repeat_interleave( + torch.arange(1, world_size + 1, dtype=input_tensor.dtype, device=device), + args.numel, + ) + if not torch.equal(output_tensor, expected): + raise RuntimeError(f"MORI all-gather mismatch on rank {rank}") + + if rank == 0: + size_mb = args.numel * input_tensor.element_size() / (1 << 20) + print( + f"MORI_MULTINODE_ALLGATHER_PASS world={world_size} " + f"ranks_per_node={ranks_per_node} per_rank_mb={size_mb:.1f}", + flush=True, + ) + + import mori.shmem as shmem + + shmem.shmem_finalize() + dist.destroy_process_group() + + +if __name__ == "__main__": + main() diff --git a/runner/primus-cli-direct.sh b/runner/primus-cli-direct.sh index c90a0d86d..4db1ec93a 100755 --- a/runner/primus-cli-direct.sh +++ b/runner/primus-cli-direct.sh @@ -100,6 +100,9 @@ Examples: # Run silently (back-pocket option; launcher errors and log file preserved) primus-cli direct --silent -- preflight --quick + # Runtime NIC inspection, MORI build, and local all-gather correctness smoke + primus-cli direct -- preflight --mori + Notes: - If --single is specified, Primus skips torchrun and uses python3 directly. - run_mode auto-detection: when the primus subcommand is 'node_smoke', run_mode @@ -415,6 +418,15 @@ if [[ -z "${direct_config[run_mode]:-}" ]]; then break fi done + if [[ "$_detected_subcmd" == "preflight" ]]; then + for _arg in "${primus_args[@]}"; do + if [[ "$_arg" == "--mori" ]]; then + _default_run_mode="single" + LOG_INFO_RANK0 "[direct] Auto-selected run_mode=single for 'preflight --mori'" + break + fi + done + fi direct_config[run_mode]="$_default_run_mode" fi diff --git a/tests/runner/test_primus_cli_direct.sh b/tests/runner/test_primus_cli_direct.sh index 86ea1a542..bc3381751 100755 --- a/tests/runner/test_primus_cli_direct.sh +++ b/tests/runner/test_primus_cli_direct.sh @@ -441,7 +441,18 @@ test_auto_single_for_node_smoke() { assert_contains "$out_preflight" "Run Mode : torchrun" "preflight defaults to torchrun" assert_contains "$out_preflight" "torchrun --nproc_per_node" "preflight uses torchrun command" - # Sub-test 13b: node_smoke auto-selects single mode. + # Sub-test 13b: MORI preflight auto-selects single mode because it + # launches its own temporary-container torchrun workload. + local out_mori + out_mori=$(timeout 30 bash "$RUNNER_DIR/primus-cli-direct.sh" --dry-run -- preflight --mori 2>&1 || true) + assert_contains "$out_mori" "Auto-selected run_mode=single for 'preflight --mori'" \ + "MORI preflight auto-detect fires" + assert_contains "$out_mori" "Run Mode : single" "MORI preflight uses single mode" + assert_contains "$out_mori" "python3" "MORI preflight uses python3 launcher" + assert_not_contains "$out_mori" "torchrun --nproc_per_node" \ + "MORI preflight does NOT use launcher torchrun" + + # Sub-test 13c: node_smoke auto-selects single mode. local out_smoke out_smoke=$(timeout 30 bash "$RUNNER_DIR/primus-cli-direct.sh" --dry-run -- node_smoke --tier2-perf 2>&1 || true) assert_contains "$out_smoke" "Auto-selected run_mode=single for subcommand 'node_smoke'" \ @@ -450,7 +461,7 @@ test_auto_single_for_node_smoke() { assert_contains "$out_smoke" "python3" "node_smoke uses python3 launcher" assert_not_contains "$out_smoke" "torchrun --nproc_per_node" "node_smoke does NOT use torchrun" - # Sub-test 13c: explicit --single on a non-node_smoke subcommand still + # Sub-test 13d: explicit --single on a non-node_smoke subcommand still # works (regression guard). local out_explicit out_explicit=$(timeout 30 bash "$RUNNER_DIR/primus-cli-direct.sh" --single --dry-run -- benchmark gemm 2>&1 || true) diff --git a/tests/unit_tests/cli/test_mori_preflight_helper.py b/tests/unit_tests/cli/test_mori_preflight_helper.py new file mode 100644 index 000000000..f278c7dea --- /dev/null +++ b/tests/unit_tests/cli/test_mori_preflight_helper.py @@ -0,0 +1,96 @@ +"""Unit tests for the native MORI multi-node preflight orchestrator.""" + +from __future__ import annotations + +from types import SimpleNamespace + +import pytest + +from primus.tools.preflight import mori_preflight as helper + + +def test_resolve_nodes_csv_deduplicates(): + assert helper.resolve_nodes("node1,node2,node1,node3") == [ + "node1", + "node2", + "node3", + ] + + +def test_resolve_nodes_file(tmp_path): + node_file = tmp_path / "nodes.txt" + node_file.write_text("node1\nnode2\nnode1\n") + assert helper.resolve_nodes(f"@{node_file}") == ["node1", "node2"] + + +def test_network_probe_uses_one_remote_command(monkeypatch): + calls = [] + monkeypatch.setattr( + helper, + "remote_output", + lambda *args: calls.append(args) or "fenic\t1\t10.0.0.1", + ) + assert helper.probe_network("node1", None, None, None) == ("fenic", 1, "10.0.0.1") + assert len(calls) == 1 + + +def test_multinode_smoke_launches_every_rank(monkeypatch, tmp_path): + args = SimpleNamespace( + mori_socket_ifname="fenic", + mori_gid_index=1, + mori_master_addr="10.0.0.1", + mori_master_port=29610, + mori_smoke_numel=1024, + ) + commands = [] + + def fake_run(command, **_kwargs): + commands.append(command) + return SimpleNamespace(returncode=0) + + monkeypatch.setattr(helper.subprocess, "run", fake_run) + monkeypatch.setattr(helper, "is_local", lambda _node: True) + monkeypatch.setattr(helper.time, "sleep", lambda _seconds: None) + + assert helper.run_multinode(args, ["node1", "node2"], tmp_path) == 0 + launches = [command[-1] for command in commands] + assert len(launches) == 2 + assert any("--node_rank=0" in command for command in launches) + assert any("--node_rank=1" in command for command in launches) + + +def test_matching_fingerprints_pass(tmp_path): + results = [ + helper.NodeResult( + "node1", + 0, + "nic=ionic hash=abc ccqe_candidate=true ccqe_runtime=true", + tmp_path / "1.log", + ), + helper.NodeResult( + "node2", + 0, + "nic=ionic hash=abc ccqe_candidate=true ccqe_runtime=true", + tmp_path / "2.log", + ), + ] + helper.validate_fingerprints(results) + + +def test_mismatched_fingerprints_fail(tmp_path): + results = [ + helper.NodeResult( + "node1", + 0, + "nic=ionic hash=abc ccqe_candidate=true ccqe_runtime=true", + tmp_path / "1.log", + ), + helper.NodeResult( + "node2", + 0, + "nic=ionic hash=abc ccqe_candidate=true ccqe_runtime=false", + tmp_path / "2.log", + ), + ] + with pytest.raises(RuntimeError, match="node-stack mismatch"): + helper.validate_fingerprints(results) diff --git a/tests/unit_tests/cli/test_preflight_subcommand.py b/tests/unit_tests/cli/test_preflight_subcommand.py index 47a997b5d..b9887b6e1 100644 --- a/tests/unit_tests/cli/test_preflight_subcommand.py +++ b/tests/unit_tests/cli/test_preflight_subcommand.py @@ -14,8 +14,10 @@ from __future__ import annotations import argparse +from pathlib import Path from primus.cli.subcommands import preflight +from primus.tools.preflight import mori_preflight def _build_parser(): @@ -43,6 +45,9 @@ def test_defaults(): # (preflight-{NNODES}N-{YYYYMMDD-HHMMSS}) at run time. assert args.report_file_name is None assert args.save_pdf is True + assert args.mori is False + assert args.mori_nodes is None + assert args.mori_smoke_numel == 67108864 def test_selection_flags(): @@ -71,3 +76,54 @@ def test_perf_test_flag(): args = parser.parse_args(["preflight", "--perf-test", "--plot"]) assert args.perf_test is True assert args.plot is True + + +def test_mori_mode_flags(): + parser, _ = _build_parser() + args = parser.parse_args( + [ + "preflight", + "--mori", + "--mori-nodes", + "node-1,node-2,node-3", + "--mori-socket-ifname", + "fenic", + "--mori-gid-index", + "1", + ] + ) + assert args.mori is True + assert args.mori_nodes == "node-1,node-2,node-3" + assert args.mori_socket_ifname == "fenic" + assert args.mori_gid_index == 1 + + +def test_mori_mode_maps_cli_to_helper_args(monkeypatch, tmp_path): + parser, _ = _build_parser() + args = parser.parse_args( + [ + "preflight", + "--mori", + "--dump-path", + str(tmp_path), + "--mori-nodes", + "node-1,node-2", + "--mori-max-jobs", + "12", + ] + ) + repo_root = Path(mori_preflight.__file__).resolve().parents[3] + monkeypatch.setenv("PRIMUS_PATH", str(repo_root)) + captured = [] + monkeypatch.setattr( + mori_preflight, + "run_orchestrator", + lambda *values: captured.extend(values) or 0, + ) + assert mori_preflight.run_mori_preflight(args, []) == 0 + native_args, native_repo_root, native_log_dir = captured + assert native_args is args + assert native_args.mori_nodes == "node-1,node-2" + assert native_args.mori_max_jobs == 12 + assert native_log_dir.parent == tmp_path + assert native_repo_root == repo_root From c7b6277f2461ac6d70c4d1ba9e1d52fdbafa934c Mon Sep 17 00:00:00 2001 From: Lorri Rao Date: Wed, 12 Aug 2026 23:56:54 +0000 Subject: [PATCH 2/3] Optimize MORI FSDP memory and overlap Size device-driven workspaces from observed shards and default dense-node collectives to async completion so large models avoid excess HBM while overlapping compute. Document execution modes and add comparison tooling and tests. --- docs/04-technical-guides/sdma-allgather.md | 60 ++++- primus/backends/common/mori_allgather.py | 110 ++++++++- .../patches/mori_allgather_patches.py | 1 + .../torchtitan/patches/mori_allgather.py | 1 + runner/helpers/mori/run_llama405b_compare.sh | 232 ++++++++++++++++++ .../backends/test_mori_allgather.py | 121 +++++++++ 6 files changed, 513 insertions(+), 12 deletions(-) create mode 100755 runner/helpers/mori/run_llama405b_compare.sh create mode 100644 tests/unit_tests/backends/test_mori_allgather.py diff --git a/docs/04-technical-guides/sdma-allgather.md b/docs/04-technical-guides/sdma-allgather.md index 02a16c36a..c0958928e 100644 --- a/docs/04-technical-guides/sdma-allgather.md +++ b/docs/04-technical-guides/sdma-allgather.md @@ -128,8 +128,9 @@ If an SDMA run hangs before any FSDP forward progress, please try to dump stack ## MORI hierarchical all-gather -MORI replaces FSDP2 all-gather with `mori.ccl.HierAllGather`. It uses SDMA for -the intra-node comm and vendor direct verbs for the cross-node RDMA comm. +MORI replaces FSDP2 all-gather with a hierarchical intra-node and cross-node +collective. The default device-driven mode uses SDMA inside each node and +vendor direct verbs for cross-node RDMA. A host-proxy mode is also available. ### Enablement @@ -167,12 +168,63 @@ The adapter: avoids creating an eager cross-node RCCL transport solely for MORI bootstrap. 2. Derives ranks per node from `LOCAL_WORLD_SIZE`. -3. Builds and caches `HierAllGather` for the FSDP process group and largest - observed per-rank input. +3. Inspects every compatible FSDP parameter group before training, computes the + largest padded per-rank shard using its effective communication dtype, and + builds `HierAllGather` once at that capacity. 4. Launches MORI on the current CUDA stream and returns a Work-like object when FSDP requests asynchronous completion. +### Inter-node execution modes + +Primus supports two MORI implementations for the cross-node leg. + +#### Device-driven RDMA (default) + +When `MORI_FSDP_HOST_PROXY` is unset or false, Primus constructs +`mori.ccl.HierAllGather`. GPU kernels post cross-node RDMA operations directly +through MORI's device-verbs/IBGDA path, while intra-node traffic uses MORI SDMA. +The CPU is not in the per-collective data path. + +This mode provides direct GPU/NIC overlap, but requires a compatible live NIC +stack and working device-side queue support. For Ionic devices, that includes +the effective CCQE capability checked by MORI preflight. Driver, firmware, +vendor-library, GID, or CCQE mismatches can prevent initialization or cause a +device-side collective failure. + +#### CPU host proxy + +Set the following to use MORI's persistent host-proxy implementation: + +```bash +export MORI_FSDP_HOST_PROXY=1 +``` + +Primus then constructs `mori.ccl.HostProxyHierAllGather`. A CPU proxy posts +RDMA work requests and polls completion queues on behalf of the GPU. Collective +payloads remain in registered GPU memory; host proxy does not bounce the data +through CPU memory. This path can maintain deeper NIC send queues and avoids +depending on GPU-posted RDMA, but it adds CPU progress and synchronization to +the hot path. + +By default, host proxy uses PyTorch/RCCL for its intra-node gather legs. Enable +MORI SDMA for those legs with: + +```bash +export MORI_HOSTPROXY_SDMA_INTRA=1 +``` + +The +[currently pinned host-proxy implementation](https://github.com/ROCm/mori/blob/12d1bc32d0c93dcd5062e74f4e0f772e36e1aac4/python/mori/ccl/host_proxy_ag.py#L177-L185) +has a single-node degenerate path and a two-node cross-node path; more than two +nodes raises `NotImplementedError`. This limitation applies only to host proxy, +not the default device-driven `HierAllGather`. Host proxy allocates a +persistent full-output GPU staging buffer. Primus automatically derives the +maximum per-rank shard from the attached FSDP groups; no manual size calculation +is required. The host-proxy-specific `MORI_FSDP_HOSTPROXY_CAP_MB` remains +available as an additional minimum. + + ### Runtime preflight MORI is sensitive to the live NIC driver, firmware, direct-verbs library, GID, diff --git a/primus/backends/common/mori_allgather.py b/primus/backends/common/mori_allgather.py index 730f54e56..d1bf6655d 100644 --- a/primus/backends/common/mori_allgather.py +++ b/primus/backends/common/mori_allgather.py @@ -35,6 +35,54 @@ _FSDP_ALL_GATHER_IMPORT_ERROR = None _MORI_SHMEM_INITIALIZED = False +_MIB = 1 << 20 +_DEFAULT_SLICE_MIN_BYTES = 8 * _MIB + + +def _env_flag(name: str, default: bool) -> bool: + value = os.environ.get(name) + if value is None: + return default + return value.strip().lower() not in ("", "0", "false", "no", "off") + + +def _env_nonnegative_int(name: str, default: int = 0) -> int: + value = int(os.environ.get(name, str(default))) + if value < 0: + raise ValueError(f"{name} must be non-negative, got {value}") + return value + + +def _compact_workspace_sizes( + cap_bytes: int, + world_size: int, + ranks_per_node: int, + slice_min_bytes: int, +) -> tuple[int, int]: + """Return MORI's input/output capacities for its compact direct path. + + Large messages use MORI's sliced path: the inter-node ring holds one shard + per node, and the intra-node phase writes directly to the FSDP output. + Messages below ``slice_min_bytes`` may use the non-sliced fallback, whose + full-world output must also fit. The pinned MORI implementation uses the + larger capacity for both fused transits, so this is conservative while + avoiding full-layer allocations. + """ + if cap_bytes <= 0: + raise ValueError(f"cap_bytes must be positive, got {cap_bytes}") + if ranks_per_node <= 0 or world_size % ranks_per_node != 0: + raise ValueError( + f"world_size ({world_size}) must be divisible by ranks_per_node ({ranks_per_node})" + ) + if slice_min_bytes < 0: + raise ValueError(f"slice_min_bytes must be non-negative, got {slice_min_bytes}") + + num_nodes = world_size // ranks_per_node + fallback_per_rank = min(cap_bytes, slice_min_bytes) + sliced_ring_bytes = num_nodes * cap_bytes + fallback_output_bytes = world_size * fallback_per_rank + workspace_bytes = max(sliced_ring_bytes, fallback_output_bytes) + return fallback_per_rank, workspace_bytes def _safe_log_rank_0(message: str) -> None: @@ -172,6 +220,7 @@ def __init__(self, ranks_per_node: int | None = None) -> None: self._rank: int | None = None self._world_size: int | None = None self._cap_bytes = 0 + self._observed_max_shard_bytes = 0 self._output_buffer: torch.Tensor | None = None world = int(os.environ.get("WORLD_SIZE", "0") or "0") @@ -189,7 +238,7 @@ def __init__(self, ranks_per_node: int | None = None) -> None: setdefault("MORI_HIER_DEEP_PIPE", "auto") setdefault("MORI_SDMA_NUM_CHANNELS", "8") else: - setdefault("MORI_HIER_DEBUG_SYNC", "1") + setdefault("MORI_HIER_DEBUG_SYNC", "0") setdefault("MORI_HIER_CUDA_GRAPH", "0") setdefault("MORI_FSDP_DEFER_HOSTSYNC", "1") setdefault("MORI_FSDP_EVENT_FENCE", "1") @@ -229,6 +278,24 @@ def __init__(self, ranks_per_node: int | None = None) -> None: "False", ) + def observe_fsdp_param_group(self, param_group: Any) -> int: + """Record a conservative all-gather shard size before training starts.""" + param_dtype = getattr(getattr(param_group, "mp_policy", None), "param_dtype", None) + shard_bytes = 0 + for fsdp_param in getattr(param_group, "fsdp_params", ()): + tensor = getattr(fsdp_param, "_sharded_param_data", None) + if tensor is None: + continue + dtype = tensor.dtype + if param_dtype is not None and dtype.is_floating_point: + dtype = param_dtype + shard_bytes += tensor.numel() * torch.empty((), dtype=dtype).element_size() + + if shard_bytes > 0: + shard_bytes = ((shard_bytes + _MIB - 1) // _MIB) * _MIB + self._observed_max_shard_bytes = max(self._observed_max_shard_bytes, shard_bytes) + return shard_bytes + def allocate( self, size: Sequence[int | torch.SymInt], @@ -259,11 +326,16 @@ def _ranks_per_node_value(self, world_size: int) -> int: def _get_collective(self, group: dist.ProcessGroup, per_rank_bytes: int) -> Any: rank, world_size = group.rank(), group.size() + cap_floor = self._observed_max_shard_bytes + if self._host_proxy: + hostproxy_floor = _env_nonnegative_int("MORI_FSDP_HOSTPROXY_CAP_MB", 160) * _MIB + cap_floor = max(cap_floor, hostproxy_floor) + required_cap = max(per_rank_bytes, cap_floor) if ( self._collective is not None and self._rank == rank and self._world_size == world_size - and self._cap_bytes >= per_rank_bytes + and self._cap_bytes >= required_cap ): return self._collective @@ -280,15 +352,13 @@ def _get_collective(self, group: dist.ProcessGroup, per_rank_bytes: int) -> Any: f"my_pe/npes={my_pe}/{npes}" ) - cap = max(per_rank_bytes, self._cap_bytes) + cap = max(required_cap, self._cap_bytes) ranks_per_node = self._ranks_per_node_value(world_size) if self._host_proxy: - cap_floor = int(os.environ.get("MORI_FSDP_HOSTPROXY_CAP_MB", "160")) * (1 << 20) - cap = max(cap, cap_floor) if self._collective is not None: raise RuntimeError( "HostProxyHierAllGather built with cap " - f"{self._cap_bytes} B but a {per_rank_bytes} B AG arrived; " + f"{self._cap_bytes} B but requires {required_cap} B; " "raise MORI_FSDP_HOSTPROXY_CAP_MB" ) collective = ccl.HostProxyHierAllGather( @@ -298,13 +368,37 @@ def _get_collective(self, group: dist.ProcessGroup, per_rank_bytes: int) -> Any: output_buffer_size=cap * world_size, ) else: + num_nodes = world_size // ranks_per_node + compact = num_nodes >= 2 and _env_flag("MORI_FSDP_COMPACT_WORKSPACE", True) + if compact: + slice_min_bytes = _env_nonnegative_int( + "MORI_FSDP_SLICE_MIN_MB", _DEFAULT_SLICE_MIN_BYTES // _MIB + ) * _MIB + input_buffer_size, output_buffer_size = _compact_workspace_sizes( + cap, + world_size, + ranks_per_node, + slice_min_bytes, + ) + else: + slice_min_bytes = _DEFAULT_SLICE_MIN_BYTES + input_buffer_size = cap + output_buffer_size = cap * world_size + + _safe_log_rank_0( + "[MORI:FSDP] building HierAllGather " + f"cap={cap} B input_workspace={input_buffer_size} B " + f"output_workspace={output_buffer_size} B compact={compact}" + ) collective = ccl.HierAllGather( my_pe, npes, - input_buffer_size=cap, - output_buffer_size=cap * world_size, + input_buffer_size=input_buffer_size, + output_buffer_size=output_buffer_size, copy_output_to_user=True, ranks_per_node=ranks_per_node, + slice_min_bytes=slice_min_bytes, + slice_direct=True if compact else None, ) self._collective = collective diff --git a/primus/backends/megatron/patches/mori_allgather_patches.py b/primus/backends/megatron/patches/mori_allgather_patches.py index 4fcf8e971..e90d2a15c 100644 --- a/primus/backends/megatron/patches/mori_allgather_patches.py +++ b/primus/backends/megatron/patches/mori_allgather_patches.py @@ -99,6 +99,7 @@ def _attach_mori_all_gather(fsdp_module) -> None: ) return + mori_all_gather.observe_fsdp_param_group(groups[0]) try: fsdp_module.set_custom_all_gather(mori_all_gather) except (AttributeError, ValueError, AssertionError) as e: diff --git a/primus/backends/torchtitan/patches/mori_allgather.py b/primus/backends/torchtitan/patches/mori_allgather.py index 00aaf64ed..86d4bd9c9 100644 --- a/primus/backends/torchtitan/patches/mori_allgather.py +++ b/primus/backends/torchtitan/patches/mori_allgather.py @@ -75,6 +75,7 @@ def _attach_mori_all_gather(fsdp_module) -> None: ) return + mori_all_gather.observe_fsdp_param_group(groups[0]) try: fsdp_module.set_custom_all_gather(mori_all_gather) except (AttributeError, ValueError, AssertionError) as e: diff --git a/runner/helpers/mori/run_llama405b_compare.sh b/runner/helpers/mori/run_llama405b_compare.sh new file mode 100755 index 000000000..c2c4f0008 --- /dev/null +++ b/runner/helpers/mori/run_llama405b_compare.sh @@ -0,0 +1,232 @@ +#!/usr/bin/env bash +############################################################################### +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +set -euo pipefail + +BACKEND="${1:-}" +case "${BACKEND}" in + rccl | mori) ;; + *) + echo "Usage: $0 {rccl|mori}" >&2 + exit 2 + ;; +esac + +# Edit these values for the experiment you want to run. +NODES="${NODES:-smci355-ccs-aus-n04-33,smci355-ccs-aus-n05-21}" +IMAGE="${IMAGE:-unifiedtrainingdockers.azurecr.io/utd/nightly:primus_the_rock_rocm7.15_20260728}" +CONFIG="${CONFIG:-examples/torchtitan/configs/MI355X/llama3.1_405B-BF16-pretrain.yaml}" +MODEL_LAYERS="${MODEL_LAYERS:-64}" # Set 126 for the full 405B model. +TRAINING_STEPS="${TRAINING_STEPS:-21}" +# SEQ_LEN="${SEQ_LEN:-128}" +LOCAL_BATCH_SIZE="${LOCAL_BATCH_SIZE:-1}" +MOCK_DATA="${MOCK_DATA:-True}" +# DISABLE_COMPILE="${DISABLE_COMPILE:-1}" +# DISABLE_PRIMUS_TURBO="${DISABLE_PRIMUS_TURBO:-1}" + +# profile_freq=20, warmup=1, active=1 records iteration 19 and saves at 20. +PROFILE_STEP="${PROFILE_STEP:-19}" +PROFILER_WARMUP="${PROFILER_WARMUP:-1}" +PROFILER_ACTIVE="${PROFILER_ACTIVE:-1}" +PROFILE_FREQ="$((PROFILE_STEP + 1))" + +PRIMUS_ROOT="${PRIMUS_ROOT:-/apps/tas/lorrirao/sdma_rccl_pytorch/primus}" +SHARED_DATA="${SHARED_DATA:-/apps/tas/lorrirao/sdma_rccl_pytorch/mori_multinode_data}" +HF_TOKEN_PATH="${HF_TOKEN_PATH:-/apps/tas/lorrirao/.cache/huggingface/token}" +OUTPUT_ROOT="${OUTPUT_ROOT:-/apps/tas/lorrirao/sdma_rccl_pytorch/mori_perf_compare_405b}" + +# Optional local tokenizer/assets override. Leave empty to use SHARED_DATA. +HF_ASSETS_HOST_DIR="${HF_ASSETS_HOST_DIR:-}" +HF_ASSETS_CONTAINER_DIR="${HF_ASSETS_CONTAINER_DIR:-/workspace/Primus/data/torchtitan/Llama-3.1-405B}" + +MASTER_PORT="${MASTER_PORT:-29660}" +SOCKET_IFNAME="${SOCKET_IFNAME:-fenic}" +NCCL_IB_GID_INDEX="${NCCL_IB_GID_INDEX:-1}" +# MORI_SHMEM_HEAP_SIZE="${MORI_SHMEM_HEAP_SIZE:-8G}" +EXTRA_TRAIN_ARGS="${EXTRA_TRAIN_ARGS:-}" + +IFS="," read -r -a NODE_ARRAY <<<"${NODES}" +NNODES="${#NODE_ARRAY[@]}" +if ((NNODES != 2)); then + echo "This comparison script expects exactly two nodes, got ${NNODES}." >&2 + exit 2 +fi + +quote_cmd() { + printf "%q " "$@" +} + +remote() { + local node="$1" + shift + if [[ "${node}" == "$(hostname -s)" || "${node}" == "$(hostname -f)" ]]; then + "$@" + else + ssh -o BatchMode=yes "${node}" "$(quote_cmd "$@")" + fi +} + +MASTER_ADDR="${MASTER_ADDR:-}" +if [[ -z "${MASTER_ADDR}" ]]; then + MASTER_ADDR="$( + remote "${NODE_ARRAY[0]}" bash -lc \ + "ip -o -4 addr show dev $(printf '%q' "${SOCKET_IFNAME}") scope global | + awk 'NR==1{split(\$4,a,\"/\"); print a[1]}'" + )" +fi +if [[ -z "${MASTER_ADDR}" ]]; then + echo "Unable to determine MASTER_ADDR on ${NODE_ARRAY[0]}/${SOCKET_IFNAME}." >&2 + exit 1 +fi + +mkdir -p "${OUTPUT_ROOT}/${BACKEND}" + +TRAIN_ARGS=( + -- + train + pretrain + --config + "${CONFIG}" + --training.steps + "${TRAINING_STEPS}" + --training.mock_data + "${MOCK_DATA}" + --training.local_batch_size + "${LOCAL_BATCH_SIZE}" + --metrics.log_freq + 1 + --metrics.enable_tensorboard + True + --metrics.save_tb_folder + tb + --metrics.disable_color_printing + True + --profiling.enable_profiling + True + --profiling.save_traces_folder + profile_traces + --profiling.profile_freq + "${PROFILE_FREQ}" + --profiling.profiler_warmup + "${PROFILER_WARMUP}" + --profiling.profiler_active + "${PROFILER_ACTIVE}" + --job.dump_folder + "/workspace/results/${BACKEND}" +) + +[[ -n "${MODEL_LAYERS}" ]] && TRAIN_ARGS+=(--model.n_layers "${MODEL_LAYERS}") +[[ -n "${SEQ_LEN:-}" ]] && TRAIN_ARGS+=(--training.seq_len "${SEQ_LEN}") +[[ "${DISABLE_COMPILE:-0}" == "1" ]] && TRAIN_ARGS+=(--compile.enable False) +[[ "${DISABLE_PRIMUS_TURBO:-0}" == "1" ]] && + TRAIN_ARGS+=(--primus_turbo.enable_primus_turbo False) +if [[ -n "${EXTRA_TRAIN_ARGS}" ]]; then + read -r -a extra_args <<<"${EXTRA_TRAIN_ARGS}" + TRAIN_ARGS+=("${extra_args[@]}") +fi +printf -v TRAIN_ARGS_QUOTED "%q " "${TRAIN_ARGS[@]}" + +container_name() { + local rank="$1" + echo "llama405_${BACKEND}_${USER}_${rank}" +} + +cleanup() { + local rank + for rank in "${!NODE_ARRAY[@]}"; do + remote "${NODE_ARRAY[rank]}" docker rm -f "$(container_name "${rank}")" \ + >/dev/null 2>&1 || true + done +} +trap cleanup EXIT INT TERM + +launch_rank() { + local node="$1" + local rank="$2" + local name + name="$(container_name "${rank}")" + + local docker_args=( + docker run --rm + --name "${name}" + --device=/dev/kfd + --device=/dev/dri + --group-add video + --cap-add SYS_PTRACE + --security-opt seccomp=unconfined + --privileged + --ipc=host + --network=host + -v "${PRIMUS_ROOT}/primus:/workspace/Primus/primus:ro" + -v "${PRIMUS_ROOT}/runner:/workspace/Primus/runner:ro" + -v "${PRIMUS_ROOT}/examples:/workspace/Primus/examples:ro" + -v "${SHARED_DATA}:/workspace/Primus/data" + -v "${OUTPUT_ROOT}:/workspace/results" + -v "${HF_TOKEN_PATH}:/run/hf_token:ro" + -e HF_TOKEN_FILE=/run/hf_token + -e NNODES="${NNODES}" + -e NODE_RANK="${rank}" + -e GPUS_PER_NODE=8 + -e MASTER_ADDR="${MASTER_ADDR}" + -e MASTER_PORT="${MASTER_PORT}" + -e NCCL_SOCKET_IFNAME="${SOCKET_IFNAME}" + -e GLOO_SOCKET_IFNAME="${SOCKET_IFNAME}" + -e NCCL_IB_GID_INDEX="${NCCL_IB_GID_INDEX}" + ) + if [[ -n "${HF_ASSETS_HOST_DIR}" ]]; then + docker_args+=(-v "${HF_ASSETS_HOST_DIR}:${HF_ASSETS_CONTAINER_DIR}:ro") + fi + if [[ "${BACKEND}" == "mori" ]]; then + docker_args+=( + -e FSDP_ALL_GATHER_BACKEND=mori + -e MORI_SOCKET_IFNAME="${SOCKET_IFNAME}" + -e MORI_HIER_CUDA_GRAPH=0 + -e MORI_SHMEM_HEAP_SIZE="${MORI_SHMEM_HEAP_SIZE:-8G}" + -e MORI_FSDP_COMPACT_WORKSPACE="${MORI_FSDP_COMPACT_WORKSPACE:-1}" + ) + fi + + remote "${node}" docker rm -f "${name}" >/dev/null 2>&1 || true + remote "${node}" \ + "${docker_args[@]}" \ + "${IMAGE}" \ + bash -lc \ + "export HF_TOKEN=\"\$(< /run/hf_token)\"; \ + cd /workspace/Primus; \ + bash runner/primus-cli direct \ + --log_file /workspace/results/${BACKEND}_node${rank}.log \ + ${TRAIN_ARGS_QUOTED}" +} + +echo "Backend : ${BACKEND}" +echo "Nodes : ${NODE_ARRAY[*]}" +echo "Master : ${MASTER_ADDR}:${MASTER_PORT}" +echo "Config : ${CONFIG}" +echo "Layers : ${MODEL_LAYERS:-config default}" +echo "Steps : ${TRAINING_STEPS}" +echo "Profile step : ${PROFILE_STEP}" +echo "Output : ${OUTPUT_ROOT}/${BACKEND}" + +pids=() +for rank in "${!NODE_ARRAY[@]}"; do + node="${NODE_ARRAY[rank]}" + launch_rank "${node}" "${rank}" \ + >"${OUTPUT_ROOT}/${BACKEND}/launcher-node${rank}.log" 2>&1 & + pids+=("$!") + [[ "${rank}" == "0" ]] && sleep 1 +done + +status=0 +for pid in "${pids[@]}"; do + wait "${pid}" || status=$? +done +if [[ "${status}" -ne 0 ]]; then + echo "${BACKEND} run failed. Logs: ${OUTPUT_ROOT}/${BACKEND}" >&2 + exit "${status}" +fi + +echo "${BACKEND} run passed. Results: ${OUTPUT_ROOT}/${BACKEND}" diff --git a/tests/unit_tests/backends/test_mori_allgather.py b/tests/unit_tests/backends/test_mori_allgather.py new file mode 100644 index 000000000..6869784d9 --- /dev/null +++ b/tests/unit_tests/backends/test_mori_allgather.py @@ -0,0 +1,121 @@ +############################################################################### +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +from types import SimpleNamespace + +import pytest +import torch + +from primus.backends.common import mori_allgather + + +def test_dense_node_defaults_to_async_completion(monkeypatch): + monkeypatch.setenv("WORLD_SIZE", "16") + monkeypatch.setenv("LOCAL_WORLD_SIZE", "8") + monkeypatch.delenv("MORI_HIER_DEBUG_SYNC", raising=False) + + mori_allgather.MoriAllGather() + + assert mori_allgather.os.environ["MORI_HIER_DEBUG_SYNC"] == "0" + + +def test_compact_workspace_sizes_cover_sliced_and_fallback_paths(): + mib = 1 << 20 + + input_bytes, output_bytes = mori_allgather._compact_workspace_sizes( + cap_bytes=381 * mib, + world_size=16, + ranks_per_node=8, + slice_min_bytes=8 * mib, + ) + + assert input_bytes == 8 * mib + assert output_bytes == 2 * 381 * mib + + +def test_compact_workspace_sizes_keep_full_output_for_small_messages(): + mib = 1 << 20 + + input_bytes, output_bytes = mori_allgather._compact_workspace_sizes( + cap_bytes=4 * mib, + world_size=16, + ranks_per_node=8, + slice_min_bytes=8 * mib, + ) + + assert input_bytes == 4 * mib + assert output_bytes == 16 * 4 * mib + + +def test_compact_workspace_sizes_reject_invalid_topology(): + with pytest.raises(ValueError, match="must be divisible"): + mori_allgather._compact_workspace_sizes( + cap_bytes=1024, + world_size=16, + ranks_per_node=6, + slice_min_bytes=1024, + ) + + +def test_observe_fsdp_param_group_uses_effective_dtype_and_rounds_up(): + adapter = mori_allgather.MoriAllGather.__new__(mori_allgather.MoriAllGather) + adapter._observed_max_shard_bytes = 0 + group = SimpleNamespace( + mp_policy=SimpleNamespace(param_dtype=torch.bfloat16), + fsdp_params=[ + SimpleNamespace(_sharded_param_data=torch.empty(100, dtype=torch.float32)), + SimpleNamespace(_sharded_param_data=torch.empty(20, dtype=torch.int32)), + ], + ) + + assert adapter.observe_fsdp_param_group(group) == 1 << 20 + assert adapter._observed_max_shard_bytes == 1 << 20 + + +def test_observed_capacity_builds_compact_collective_once(monkeypatch): + mib = 1 << 20 + calls = [] + collective = object() + fake_shmem = SimpleNamespace(shmem_mype=lambda: 0, shmem_npes=lambda: 16) + + def hier_all_gather(*args, **kwargs): + calls.append((args, kwargs)) + return collective + + fake_ccl = SimpleNamespace(HierAllGather=hier_all_gather) + + def import_module(name): + if name == "mori.shmem": + return fake_shmem + if name == "mori.ccl": + return fake_ccl + raise AssertionError(f"unexpected import: {name}") + + monkeypatch.setenv("MORI_FSDP_COMPACT_WORKSPACE", "1") + monkeypatch.setattr(mori_allgather, "ensure_mori_shmem_initialized", lambda _: None) + monkeypatch.setattr(mori_allgather.importlib, "import_module", import_module) + monkeypatch.setattr(mori_allgather, "_safe_log_rank_0", lambda _: None) + + adapter = mori_allgather.MoriAllGather.__new__(mori_allgather.MoriAllGather) + adapter._ranks_per_node = 8 + adapter._collective = None + adapter._rank = None + adapter._world_size = None + adapter._cap_bytes = 0 + adapter._observed_max_shard_bytes = 381 * mib + adapter._host_proxy = False + group = SimpleNamespace(rank=lambda: 0, size=lambda: 16) + + assert adapter._get_collective(group, 380 * mib) is collective + assert adapter._get_collective(group, 128 * mib) is collective + assert len(calls) == 1 + + args, kwargs = calls[0] + assert args == (0, 16) + assert kwargs["input_buffer_size"] == 8 * mib + assert kwargs["output_buffer_size"] == 2 * 381 * mib + assert kwargs["slice_min_bytes"] == 8 * mib + assert kwargs["slice_direct"] is True From a8b4dc0cad191bd2bbe5488f96460395a976dc48 Mon Sep 17 00:00:00 2001 From: Lorri Rao Date: Thu, 13 Aug 2026 22:30:56 +0000 Subject: [PATCH 3/3] Reduce MORI FSDP synchronization overhead Use device events, automatic SHMEM sizing, and persistent backing-buffer registration to eliminate host waits and per-step IPC churn. Add standalone bandwidth and GEMM-interference reproducers with focused tests and documentation. --- docs/04-technical-guides/sdma-allgather.md | 5 +- primus/backends/common/mori_allgather.py | 125 ++++++-- .../hooks/07_enable_mori_all_gather.sh | 1 - .../mori/gemm_allgather_overlap_repro.py | 302 ++++++++++++++++++ .../helpers/mori/multinode_allgather_bench.py | 157 +++++++++ runner/helpers/mori/run_llama405b_compare.sh | 2 - .../backends/test_mori_allgather.py | 90 +++++- 7 files changed, 649 insertions(+), 33 deletions(-) create mode 100644 runner/helpers/mori/gemm_allgather_overlap_repro.py create mode 100644 runner/helpers/mori/multinode_allgather_bench.py diff --git a/docs/04-technical-guides/sdma-allgather.md b/docs/04-technical-guides/sdma-allgather.md index c0958928e..cdd99ce0c 100644 --- a/docs/04-technical-guides/sdma-allgather.md +++ b/docs/04-technical-guides/sdma-allgather.md @@ -171,7 +171,10 @@ The adapter: 3. Inspects every compatible FSDP parameter group before training, computes the largest padded per-rank shard using its effective communication dtype, and builds `HierAllGather` once at that capacity. -4. Launches MORI on the current CUDA stream and returns a Work-like object when +4. Derives the symmetric-heap size from those same workspaces before MORI SHMEM + initializes. The pinned MORI static-heap default is 4 GiB; Primus uses a + calculated 2 GiB minimum with at least 512 MiB of workspace headroom. +5. Launches MORI on the current CUDA stream and returns a Work-like object when FSDP requests asynchronous completion. diff --git a/primus/backends/common/mori_allgather.py b/primus/backends/common/mori_allgather.py index d1bf6655d..df09762dd 100644 --- a/primus/backends/common/mori_allgather.py +++ b/primus/backends/common/mori_allgather.py @@ -36,6 +36,7 @@ _MORI_SHMEM_INITIALIZED = False _MIB = 1 << 20 +_GIB = 1 << 30 _DEFAULT_SLICE_MIN_BYTES = 8 * _MIB @@ -85,6 +86,31 @@ def _compact_workspace_sizes( return fallback_per_rank, workspace_bytes +def _auto_shmem_heap_bytes( + input_buffer_size: int, + output_buffer_size: int, + world_size: int, + ranks_per_node: int, + *, + host_proxy: bool = False, + host_proxy_sdma: bool = False, +) -> int: + """Size MORI's static heap from the collective workspaces it will own.""" + num_nodes = world_size // ranks_per_node + if host_proxy: + data_bytes = output_buffer_size + output_buffer_size // max(num_nodes, 1) if host_proxy_sdma else 0 + elif num_nodes >= 2: + intra_bytes = max(ranks_per_node * input_buffer_size, output_buffer_size) + data_bytes = intra_bytes + output_buffer_size + else: + data_bytes = input_buffer_size + output_buffer_size + + margin_bytes = max(512 * _MIB, data_bytes // 4) + required_bytes = data_bytes + margin_bytes + rounded_bytes = ((required_bytes + _GIB - 1) // _GIB) * _GIB + return max(2 * _GIB, rounded_bytes) + + def _safe_log_rank_0(message: str) -> None: """Log through Primus when initialized; otherwise fall back to print.""" try: @@ -155,21 +181,28 @@ def wait(self) -> bool: return True -class _DeviceDeferredHostSyncWork(dist.distributed_c10d.Work): - """Defer MORI's reliable host landing fence until FSDP consumes the AG.""" +class _DeviceDeferredEventWork(dist.distributed_c10d.Work): + """Insert a device-side dependency when FSDP consumes the MORI result.""" - def __init__(self, stream: torch.cuda.Stream, event: torch.cuda.Event | None = None) -> None: + def __init__( + self, + stream: torch.cuda.Stream, + device: torch.device, + event: torch.cuda.Event | None = None, + ) -> None: super().__init__() self._stream = stream + self._device = device self._event = event self._done = False def wait(self, timeout=None) -> bool: # noqa: ARG002 if not self._done: + consumer_stream = torch.cuda.current_stream(self._device) if self._event is not None: - self._event.synchronize() + consumer_stream.wait_event(self._event) else: - self._stream.synchronize() + consumer_stream.wait_stream(self._stream) self._done = True return True @@ -210,7 +243,6 @@ def __init__(self, ranks_per_node: int | None = None) -> None: ) from _FSDP_ALL_GATHER_IMPORT_ERROR os.environ.setdefault("MORI_ENABLE_SDMA", "1") - os.environ.setdefault("MORI_SHMEM_HEAP_SIZE", "8G") os.environ.setdefault("MORI_HIER_CUDA_GRAPH", "0") if "MORI_SOCKET_IFNAME" not in os.environ and "NCCL_SOCKET_IFNAME" in os.environ: os.environ["MORI_SOCKET_IFNAME"] = os.environ["NCCL_SOCKET_IFNAME"].lstrip("=") @@ -316,6 +348,20 @@ def allocate( self._output_buffer = torch.empty(numel, dtype=dtype, device=device) return self._output_buffer + def _registration_output(self, output_tensor: torch.Tensor) -> torch.Tensor: + """Use the persistent backing extent for device-path IPC registration.""" + backing = self._output_buffer + if ( + self._host_proxy + or backing is None + or backing.dtype != output_tensor.dtype + or backing.device != output_tensor.device + or backing.data_ptr() != output_tensor.data_ptr() + or backing.numel() < output_tensor.numel() + ): + return output_tensor + return backing + def _ranks_per_node_value(self, world_size: int) -> int: if self._ranks_per_node is not None: return self._ranks_per_node @@ -339,6 +385,45 @@ def _get_collective(self, group: dist.ProcessGroup, per_rank_bytes: int) -> Any: ): return self._collective + cap = max(required_cap, self._cap_bytes) + ranks_per_node = self._ranks_per_node_value(world_size) + num_nodes = world_size // ranks_per_node + if self._host_proxy: + input_buffer_size = cap + output_buffer_size = cap * world_size + slice_min_bytes = _DEFAULT_SLICE_MIN_BYTES + compact = False + else: + compact = num_nodes >= 2 and _env_flag("MORI_FSDP_COMPACT_WORKSPACE", True) + if compact: + slice_min_bytes = _env_nonnegative_int( + "MORI_FSDP_SLICE_MIN_MB", _DEFAULT_SLICE_MIN_BYTES // _MIB + ) * _MIB + input_buffer_size, output_buffer_size = _compact_workspace_sizes( + cap, + world_size, + ranks_per_node, + slice_min_bytes, + ) + else: + slice_min_bytes = _DEFAULT_SLICE_MIN_BYTES + input_buffer_size = cap + output_buffer_size = cap * world_size + + heap_bytes = _auto_shmem_heap_bytes( + input_buffer_size, + output_buffer_size, + world_size, + ranks_per_node, + host_proxy=self._host_proxy, + host_proxy_sdma=_env_flag("MORI_HOSTPROXY_SDMA_INTRA", False), + ) + if "MORI_SHMEM_HEAP_SIZE" not in os.environ: + os.environ["MORI_SHMEM_HEAP_SIZE"] = f"{heap_bytes // _GIB}G" + _safe_log_rank_0( + f"[MORI:FSDP] auto-sized MORI SHMEM heap to {heap_bytes // _GIB} GiB" + ) + ensure_mori_shmem_initialized("default") shmem = importlib.import_module("mori.shmem") @@ -352,8 +437,6 @@ def _get_collective(self, group: dist.ProcessGroup, per_rank_bytes: int) -> Any: f"my_pe/npes={my_pe}/{npes}" ) - cap = max(required_cap, self._cap_bytes) - ranks_per_node = self._ranks_per_node_value(world_size) if self._host_proxy: if self._collective is not None: raise RuntimeError( @@ -365,26 +448,9 @@ def _get_collective(self, group: dist.ProcessGroup, per_rank_bytes: int) -> Any: rank, world_size, ranks_per_node, - output_buffer_size=cap * world_size, + output_buffer_size=output_buffer_size, ) else: - num_nodes = world_size // ranks_per_node - compact = num_nodes >= 2 and _env_flag("MORI_FSDP_COMPACT_WORKSPACE", True) - if compact: - slice_min_bytes = _env_nonnegative_int( - "MORI_FSDP_SLICE_MIN_MB", _DEFAULT_SLICE_MIN_BYTES // _MIB - ) * _MIB - input_buffer_size, output_buffer_size = _compact_workspace_sizes( - cap, - world_size, - ranks_per_node, - slice_min_bytes, - ) - else: - slice_min_bytes = _DEFAULT_SLICE_MIN_BYTES - input_buffer_size = cap - output_buffer_size = cap * world_size - _safe_log_rank_0( "[MORI:FSDP] building HierAllGather " f"cap={cap} B input_workspace={input_buffer_size} B " @@ -454,7 +520,10 @@ def __call__( collective._pending = work return work - ok = collective(input_tensor, output_tensor, input_tensor.numel(), stream=stream) + registration_output = self._registration_output(output_tensor) + if registration_output is not output_tensor: + registration_output.record_stream(stream) + ok = collective(input_tensor, registration_output, input_tensor.numel(), stream=stream) if not ok: raise RuntimeError("MORI HierAllGather call failed") @@ -463,7 +532,7 @@ def __call__( if self._event_fence: event = torch.cuda.Event() event.record(stream) - return _DeviceDeferredHostSyncWork(stream, event) + return _DeviceDeferredEventWork(stream, device, event) if async_op: event = torch.cuda.Event() diff --git a/runner/helpers/hooks/07_enable_mori_all_gather.sh b/runner/helpers/hooks/07_enable_mori_all_gather.sh index d16891eb5..9b1499f7a 100755 --- a/runner/helpers/hooks/07_enable_mori_all_gather.sh +++ b/runner/helpers/hooks/07_enable_mori_all_gather.sh @@ -41,7 +41,6 @@ fi # MORI all-gather uses SDMA for the intra-node leg. Allow an explicit caller # value to win, but default it on for this feature. export MORI_ENABLE_SDMA="${MORI_ENABLE_SDMA:-1}" -export MORI_SHMEM_HEAP_SIZE="${MORI_SHMEM_HEAP_SIZE:-8G}" # MORI's single-node eager path is the correctness-safe default on the ROCm # versions used by Primus v26.4. Explicit user settings still win. diff --git a/runner/helpers/mori/gemm_allgather_overlap_repro.py b/runner/helpers/mori/gemm_allgather_overlap_repro.py new file mode 100644 index 000000000..4fd659562 --- /dev/null +++ b/runner/helpers/mori/gemm_allgather_overlap_repro.py @@ -0,0 +1,302 @@ +#!/usr/bin/env python3 +"""Standalone two-node reproducer for GEMM interference from all-gather. + +Dependencies are limited to PyTorch and MORI. Launch with 8 ranks per node. +The default shapes are taken from the Llama 405B-width batch-size-2 trace. +""" + +import argparse +import json +import os +import statistics +import time +from pathlib import Path + +import torch +import torch.distributed as dist +from torch.profiler import ProfilerActivity, profile, record_function + +import mori.shmem as shmem +from mori.ccl import HierAllGather + + +SHAPES = { + # name: (M, N, K, transpose A storage, transpose B storage) + "mlp_up": (4096, 53248, 16384, False, True), + "mlp_down": (4096, 16384, 53248, False, False), + "mlp_wgrad": (53248, 16384, 4096, True, False), + "attention_proj": (4096, 16384, 16384, False, True), +} + + +def _matrix(rows, cols, transposed, device): + if transposed: + base = torch.empty((cols, rows), dtype=torch.bfloat16, device=device) + base.normal_(mean=0.0, std=0.01) + return base.t(), base + tensor = torch.empty((rows, cols), dtype=torch.bfloat16, device=device) + tensor.normal_(mean=0.0, std=0.01) + return tensor, tensor + + +def _check_all_gather(output, numel, world_size): + indices = torch.tensor([0, numel // 2, numel - 1], device=output.device) + for source_rank in range(world_size): + values = output[source_rank * numel + indices] + if not torch.equal(values, torch.full_like(values, source_rank + 1)): + raise RuntimeError(f"all-gather sample mismatch for source rank {source_rank}") + + +def _percentile(values, percentile): + ordered = sorted(values) + index = round((len(ordered) - 1) * percentile / 100) + return ordered[index] + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--shapes", nargs="+", choices=["all", *SHAPES], default=["all"]) + parser.add_argument("--modes", nargs="+", choices=["baseline", "copy", "mori", "rccl"], default=None) + parser.add_argument("--reps", type=int, default=10) + parser.add_argument("--warmup", type=int, default=3) + parser.add_argument("--fsdp-numel", type=int, default=199_231_488) + parser.add_argument("--output-json") + parser.add_argument("--trace-dir") + parser.add_argument("--trace-shape", choices=SHAPES, default="mlp_down") + parser.add_argument("--trace-mode", choices=["copy", "mori", "rccl"], default="mori") + args = parser.parse_args() + + shape_names = list(SHAPES) if "all" in args.shapes else args.shapes + modes = args.modes or ["baseline", "copy", "mori", "rccl"] + if "baseline" not in modes: + modes = ["baseline", *modes] + + os.environ.setdefault("MORI_ENABLE_SDMA", "1") + os.environ.setdefault("MORI_SHMEM_HEAP_SIZE", "2G") + os.environ.setdefault("MORI_HIER_CUDA_GRAPH", "0") + os.environ.setdefault("MORI_HIER_DEBUG_SYNC", "0") + os.environ.setdefault("MORI_HIER_FUSE_LOCAL", "1") + os.environ.setdefault("MORI_HIER_FUSE_REMOTE", "1") + os.environ.setdefault("MORI_HIER_LOCAL_PUSHONLY", "1") + + local_rank = int(os.environ["LOCAL_RANK"]) + torch.cuda.set_device(local_rank) + device = torch.device("cuda", local_rank) + dist.init_process_group("cpu:gloo,cuda:nccl") + sync_group = dist.new_group(backend="gloo") + torch._C._distributed_c10d._register_process_group("default", dist.group.WORLD) + + rank = dist.get_rank() + world_size = dist.get_world_size() + ranks_per_node = int(os.environ["LOCAL_WORLD_SIZE"]) + num_nodes = world_size // ranks_per_node + shmem.shmem_torch_process_group_init("default") + + per_rank_bytes = args.fsdp_numel * 2 + slice_min_bytes = 8 << 20 + output_workspace_bytes = max( + num_nodes * per_rank_bytes, + world_size * min(per_rank_bytes, slice_min_bytes), + ) + mori_all_gather = HierAllGather( + my_pe=rank, + npes=world_size, + input_buffer_size=min(per_rank_bytes, slice_min_bytes), + output_buffer_size=output_workspace_bytes, + copy_output_to_user=True, + ranks_per_node=ranks_per_node, + slice_min_bytes=slice_min_bytes, + slice_direct=True, + ) + + ag_input = torch.full( + (args.fsdp_numel,), + rank + 1, + dtype=torch.bfloat16, + device=device, + ) + ag_output = torch.empty( + args.fsdp_numel * world_size, + dtype=torch.bfloat16, + device=device, + ) + copy_output = torch.empty_like(ag_input) + compute_stream = torch.cuda.current_stream(device) + comm_stream = torch.cuda.Stream(device=device, priority=-1) + + def launch_comm(mode): + if mode == "baseline": + return None + with torch.cuda.stream(comm_stream): + if mode == "copy": + copy_output.copy_(ag_input, non_blocking=True) + return None + if mode == "mori": + if not mori_all_gather( + ag_input, + ag_output, + args.fsdp_numel, + stream=comm_stream, + ): + raise RuntimeError("MORI all-gather failed") + return None + return dist.all_gather_into_tensor( + ag_output, + ag_input, + group=dist.group.WORLD, + async_op=True, + ) + + # Eagerly initialize both communication paths outside measured regions. + launch_comm("mori") + comm_stream.synchronize() + _check_all_gather(ag_output, args.fsdp_numel, world_size) + launch_comm("rccl") + comm_stream.synchronize() + _check_all_gather(ag_output, args.fsdp_numel, world_size) + + local_results = {} + for shape_name in shape_names: + m, n, k, transpose_a, transpose_b = SHAPES[shape_name] + a, a_storage = _matrix(m, k, transpose_a, device) + b, b_storage = _matrix(k, n, transpose_b, device) + output = torch.empty((m, n), dtype=torch.bfloat16, device=device) + + def gemm(): + torch.mm(a, b, out=output) + + for _ in range(args.warmup): + gemm() + torch.cuda.synchronize() + + for mode in modes: + for _ in range(args.warmup): + launch_comm(mode) + gemm() + torch.cuda.synchronize() + + gemm_times = [] + comm_times = [] + wall_times = [] + for _ in range(args.reps): + dist.barrier(group=sync_group) + torch.cuda.synchronize() + gemm_start = torch.cuda.Event(enable_timing=True) + gemm_end = torch.cuda.Event(enable_timing=True) + comm_start = torch.cuda.Event(enable_timing=True) + comm_end = torch.cuda.Event(enable_timing=True) + start = time.perf_counter() + with torch.cuda.stream(comm_stream): + comm_start.record() + work = launch_comm(mode) + if work is not None: + work.wait() + comm_end.record() + with torch.cuda.stream(compute_stream): + gemm_start.record() + gemm() + gemm_end.record() + gemm_end.synchronize() + comm_end.synchronize() + torch.cuda.synchronize() + wall_times.append((time.perf_counter() - start) * 1e3) + gemm_times.append(gemm_start.elapsed_time(gemm_end)) + comm_times.append(comm_start.elapsed_time(comm_end)) + + local_results[(shape_name, mode)] = { + "gemm_ms": gemm_times, + "comm_ms": comm_times, + "wall_ms": wall_times, + } + + del a, b, a_storage, b_storage, output + torch.cuda.empty_cache() + + gathered_results = [None] * world_size + dist.all_gather_object(gathered_results, local_results, group=sync_group) + rows = [] + if rank == 0: + for shape_name in shape_names: + baseline_values = [ + value + for rank_results in gathered_results + for value in rank_results[(shape_name, "baseline")]["gemm_ms"] + ] + baseline_median = statistics.median(baseline_values) + for mode in modes: + gemm_values = [ + value + for rank_results in gathered_results + for value in rank_results[(shape_name, mode)]["gemm_ms"] + ] + comm_values = [ + value + for rank_results in gathered_results + for value in rank_results[(shape_name, mode)]["comm_ms"] + ] + wall_values = [ + value + for rank_results in gathered_results + for value in rank_results[(shape_name, mode)]["wall_ms"] + ] + row = { + "shape": shape_name, + "mode": mode, + "m": SHAPES[shape_name][0], + "n": SHAPES[shape_name][1], + "k": SHAPES[shape_name][2], + "gemm_median_ms": statistics.median(gemm_values), + "gemm_p95_ms": _percentile(gemm_values, 95), + "gemm_slowdown_pct": ( + 100 * (statistics.median(gemm_values) / baseline_median - 1) + ), + "comm_median_ms": statistics.median(comm_values), + "wall_median_ms": statistics.median(wall_values), + } + rows.append(row) + print( + f"GEMM_AG_REPRO shape={shape_name} mode={mode} " + f"gemm_ms={row['gemm_median_ms']:.3f} " + f"slowdown_pct={row['gemm_slowdown_pct']:.2f} " + f"comm_ms={row['comm_median_ms']:.3f} " + f"wall_ms={row['wall_median_ms']:.3f}", + flush=True, + ) + + if args.output_json: + output_path = Path(args.output_json) + output_path.parent.mkdir(parents=True, exist_ok=True) + output_path.write_text(json.dumps(rows, indent=2) + "\n") + + if args.trace_dir: + shape_name = args.trace_shape + if shape_name not in shape_names: + raise ValueError("--trace-shape must also be included in --shapes") + m, n, k, transpose_a, transpose_b = SHAPES[shape_name] + a, a_storage = _matrix(m, k, transpose_a, device) + b, b_storage = _matrix(k, n, transpose_b, device) + output = torch.empty((m, n), dtype=torch.bfloat16, device=device) + dist.barrier(group=sync_group) + with profile( + activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA], + record_shapes=True, + ) as profiler: + with record_function(f"gemm_overlap_{args.trace_mode}_{shape_name}"): + launch_comm(args.trace_mode) + torch.mm(a, b, out=output) + torch.cuda.synchronize() + trace_dir = Path(args.trace_dir) + trace_dir.mkdir(parents=True, exist_ok=True) + profiler.export_chrome_trace( + str(trace_dir / f"rank{rank}_{args.trace_mode}_{shape_name}.json") + ) + del a, b, a_storage, b_storage, output + + dist.barrier(group=sync_group) + del mori_all_gather + shmem.shmem_finalize() + dist.destroy_process_group() + + +if __name__ == "__main__": + main() diff --git a/runner/helpers/mori/multinode_allgather_bench.py b/runner/helpers/mori/multinode_allgather_bench.py new file mode 100644 index 000000000..245842c4d --- /dev/null +++ b/runner/helpers/mori/multinode_allgather_bench.py @@ -0,0 +1,157 @@ +#!/usr/bin/env python3 +"""Cross-node MORI/RCCL all-gather bandwidth at FSDP-scale message sizes.""" + +import argparse +import json +import os +import statistics +import time + +import torch +import torch.distributed as dist + +from primus.backends.common.mori_allgather import MoriAllGather + + +def _time_collective(fn, sync_group, reps, warmup): + for _ in range(warmup): + fn() + torch.cuda.synchronize() + + local_times = [] + for _ in range(reps): + dist.barrier(group=sync_group) + start = time.perf_counter() + fn() + torch.cuda.synchronize() + local_times.append((time.perf_counter() - start) * 1e3) + + per_rank_times = [None] * dist.get_world_size() + dist.all_gather_object(per_rank_times, local_times, group=sync_group) + return [max(rank_times[index] for rank_times in per_rank_times) for index in range(reps)] + + +def _check_samples(output, numel, world_size): + sample_indices = sorted({0, numel // 2, numel - 1}) + for rank in range(world_size): + values = output[rank * numel + torch.tensor(sample_indices, device=output.device)] + expected = torch.full_like(values, rank + 1) + if not torch.equal(values, expected): + raise RuntimeError(f"all-gather sample mismatch for source rank {rank}") + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument( + "--sizes-mib", + type=float, + nargs="+", + default=[8, 32, 64, 128, 256, 380.00390625], + ) + parser.add_argument("--reps", type=int, default=10) + parser.add_argument("--warmup", type=int, default=3) + parser.add_argument("--output-json") + args = parser.parse_args() + + local_rank = int(os.environ["LOCAL_RANK"]) + torch.cuda.set_device(local_rank) + device = torch.device("cuda", local_rank) + dist.init_process_group("cpu:gloo,cuda:nccl") + sync_group = dist.new_group(backend="gloo") + + rank = dist.get_rank() + world_size = dist.get_world_size() + ranks_per_node = int(os.environ["LOCAL_WORLD_SIZE"]) + sizes_numel = [int(size * (1 << 20) / 2) for size in args.sizes_mib] + max_numel = max(sizes_numel) + + input_tensor = torch.full( + (max_numel,), + rank + 1, + dtype=torch.bfloat16, + device=device, + ) + output_tensor = torch.empty( + max_numel * world_size, + dtype=torch.bfloat16, + device=device, + ) + + mori = MoriAllGather(ranks_per_node=ranks_per_node) + + def run_mori(numel): + work = mori( + output_tensor[: numel * world_size], + input_tensor[:numel], + dist.group.WORLD, + async_op=True, + ) + if work is not None: + work.wait() + + def run_rccl(numel): + work = dist.all_gather_into_tensor( + output_tensor[: numel * world_size], + input_tensor[:numel], + group=dist.group.WORLD, + async_op=True, + ) + work.wait() + + # Construct MORI once at the maximum capacity so every measured size reuses + # the same buffers and does not include setup or resize time. + run_mori(max_numel) + torch.cuda.synchronize() + _check_samples(output_tensor, max_numel, world_size) + + rows = [] + for mode, fn in (("mori", run_mori), ("rccl", run_rccl)): + for requested_mib, numel in zip(args.sizes_mib, sizes_numel): + times_ms = _time_collective( + lambda numel=numel: fn(numel), + sync_group, + args.reps, + args.warmup, + ) + fn(numel) + torch.cuda.synchronize() + _check_samples(output_tensor, numel, world_size) + + median_ms = statistics.median(times_ms) + per_rank_bytes = numel * 2 + algorithm_bytes = per_rank_bytes * (world_size - 1) + output_bytes = per_rank_bytes * world_size + row = { + "mode": mode, + "requested_mib": requested_mib, + "per_rank_mib": per_rank_bytes / (1 << 20), + "median_ms": median_ms, + "min_ms": min(times_ms), + "max_ms": max(times_ms), + "algorithm_gbps": algorithm_bytes / (median_ms * 1e6), + "output_gbps": output_bytes / (median_ms * 1e6), + } + rows.append(row) + if rank == 0: + print( + f"AG_BENCH mode={mode} per_rank_mib={row['per_rank_mib']:.6f} " + f"median_ms={median_ms:.3f} min_ms={row['min_ms']:.3f} " + f"algo_GBps={row['algorithm_gbps']:.2f} " + f"output_GBps={row['output_gbps']:.2f}", + flush=True, + ) + + if rank == 0 and args.output_json: + with open(args.output_json, "w") as output_file: + json.dump(rows, output_file, indent=2) + output_file.write("\n") + + import mori.shmem as shmem + + dist.barrier(group=sync_group) + shmem.shmem_finalize() + dist.destroy_process_group() + + +if __name__ == "__main__": + main() diff --git a/runner/helpers/mori/run_llama405b_compare.sh b/runner/helpers/mori/run_llama405b_compare.sh index c2c4f0008..a5bb5afa1 100755 --- a/runner/helpers/mori/run_llama405b_compare.sh +++ b/runner/helpers/mori/run_llama405b_compare.sh @@ -46,7 +46,6 @@ HF_ASSETS_CONTAINER_DIR="${HF_ASSETS_CONTAINER_DIR:-/workspace/Primus/data/torch MASTER_PORT="${MASTER_PORT:-29660}" SOCKET_IFNAME="${SOCKET_IFNAME:-fenic}" NCCL_IB_GID_INDEX="${NCCL_IB_GID_INDEX:-1}" -# MORI_SHMEM_HEAP_SIZE="${MORI_SHMEM_HEAP_SIZE:-8G}" EXTRA_TRAIN_ARGS="${EXTRA_TRAIN_ARGS:-}" IFS="," read -r -a NODE_ARRAY <<<"${NODES}" @@ -185,7 +184,6 @@ launch_rank() { -e FSDP_ALL_GATHER_BACKEND=mori -e MORI_SOCKET_IFNAME="${SOCKET_IFNAME}" -e MORI_HIER_CUDA_GRAPH=0 - -e MORI_SHMEM_HEAP_SIZE="${MORI_SHMEM_HEAP_SIZE:-8G}" -e MORI_FSDP_COMPACT_WORKSPACE="${MORI_FSDP_COMPACT_WORKSPACE:-1}" ) fi diff --git a/tests/unit_tests/backends/test_mori_allgather.py b/tests/unit_tests/backends/test_mori_allgather.py index 6869784d9..e81bf9052 100644 --- a/tests/unit_tests/backends/test_mori_allgather.py +++ b/tests/unit_tests/backends/test_mori_allgather.py @@ -12,6 +12,34 @@ from primus.backends.common import mori_allgather +def test_deferred_work_waits_on_device_event_without_host_sync(monkeypatch): + waited_events = [] + + class ConsumerStream: + def wait_event(self, event): + waited_events.append(event) + + def wait_stream(self, stream): + raise AssertionError("event path should not wait on the producer stream") + + class Event: + def synchronize(self): + raise AssertionError("wait must not synchronize the host") + + event = Event() + monkeypatch.setattr(mori_allgather.torch.cuda, "current_stream", lambda _: ConsumerStream()) + work = mori_allgather._DeviceDeferredEventWork( + stream=SimpleNamespace(), + device=torch.device("cuda", 0), + event=event, + ) + + assert work.wait() + assert work.wait() + assert waited_events == [event] + assert work.is_completed() + + def test_dense_node_defaults_to_async_completion(monkeypatch): monkeypatch.setenv("WORLD_SIZE", "16") monkeypatch.setenv("LOCAL_WORLD_SIZE", "8") @@ -60,6 +88,34 @@ def test_compact_workspace_sizes_reject_invalid_topology(): ) +def test_auto_shmem_heap_uses_two_gib_for_compact_workspace(): + mib = 1 << 20 + gib = 1 << 30 + + heap_bytes = mori_allgather._auto_shmem_heap_bytes( + input_buffer_size=8 * mib, + output_buffer_size=762 * mib, + world_size=16, + ranks_per_node=8, + ) + + assert heap_bytes == 2 * gib + + +def test_auto_shmem_heap_grows_for_large_workspace(): + mib = 1 << 20 + gib = 1 << 30 + + heap_bytes = mori_allgather._auto_shmem_heap_bytes( + input_buffer_size=8 * mib, + output_buffer_size=6 * gib, + world_size=16, + ranks_per_node=8, + ) + + assert heap_bytes == 15 * gib + + def test_observe_fsdp_param_group_uses_effective_dtype_and_rounds_up(): adapter = mori_allgather.MoriAllGather.__new__(mori_allgather.MoriAllGather) adapter._observed_max_shard_bytes = 0 @@ -75,6 +131,31 @@ def test_observe_fsdp_param_group_uses_effective_dtype_and_rounds_up(): assert adapter._observed_max_shard_bytes == 1 << 20 +def test_registration_output_uses_persistent_backing_extent(): + adapter = mori_allgather.MoriAllGather.__new__(mori_allgather.MoriAllGather) + adapter._host_proxy = False + adapter._output_buffer = torch.empty(1024) + output_view = adapter._output_buffer.narrow(0, 0, 512) + + registration_output = adapter._registration_output(output_view) + + assert registration_output is adapter._output_buffer + assert registration_output.numel() == 1024 + + +def test_registration_output_rejects_offset_view_and_host_proxy(): + adapter = mori_allgather.MoriAllGather.__new__(mori_allgather.MoriAllGather) + adapter._host_proxy = False + adapter._output_buffer = torch.empty(1024) + offset_view = adapter._output_buffer.narrow(0, 1, 512) + + assert adapter._registration_output(offset_view) is offset_view + + adapter._host_proxy = True + prefix_view = adapter._output_buffer.narrow(0, 0, 512) + assert adapter._registration_output(prefix_view) is prefix_view + + def test_observed_capacity_builds_compact_collective_once(monkeypatch): mib = 1 << 20 calls = [] @@ -94,8 +175,14 @@ def import_module(name): return fake_ccl raise AssertionError(f"unexpected import: {name}") + initialized_heaps = [] monkeypatch.setenv("MORI_FSDP_COMPACT_WORKSPACE", "1") - monkeypatch.setattr(mori_allgather, "ensure_mori_shmem_initialized", lambda _: None) + monkeypatch.delenv("MORI_SHMEM_HEAP_SIZE", raising=False) + monkeypatch.setattr( + mori_allgather, + "ensure_mori_shmem_initialized", + lambda _: initialized_heaps.append(mori_allgather.os.environ["MORI_SHMEM_HEAP_SIZE"]), + ) monkeypatch.setattr(mori_allgather.importlib, "import_module", import_module) monkeypatch.setattr(mori_allgather, "_safe_log_rank_0", lambda _: None) @@ -112,6 +199,7 @@ def import_module(name): assert adapter._get_collective(group, 380 * mib) is collective assert adapter._get_collective(group, 128 * mib) is collective assert len(calls) == 1 + assert initialized_heaps == ["2G"] args, kwargs = calls[0] assert args == (0, 16)