AMD-AGI · lorri-rao · Aug 6, 2026
@@ -165,6 +165,7 @@ These run under `primus/cli/main.py` unless you change `--script` in direct mode
 | `train posttrain --config <yaml>` | Post-training (SFT or LoRA-style workflows; same top-level flags as pretrain in the parser). |
 | `benchmark <suite> [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 <yaml>` | Memory estimation from a merged config. |
 | `projection performance --config <yaml>` | Performance projection from a merged config. |
 | `projection both --config <yaml>` | 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` |
 

@@ -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)).
 
 ---
 

@@ -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
+`<dump-path>/mori-preflight-<host>-<timestamp>/`, or the directory supplied by
+`--mori-log-dir`. It does not generate the standard Markdown/PDF performance
+report.
+
 Output files:
 
 | File | Produced when | Notes |

@@ -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

@@ -0,0 +1,252 @@
+# 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 <experiment.yaml>
+```
+
+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 <experiment.yaml>
+```
+
+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 <bootstrap-interface> \
+  --mori-gid-index <rocev2-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. Runs an 8-GPU bit-exact all-gather smoke.
+6. 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-<node>-<timestamp>/`.
+
+### 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=True` on one node and `ccqe=False` on the other: choose nodes with
+  matching ionic stacks.
+- `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.
@@ -0,0 +1,7 @@
+###############################################################################
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""Shared backend integration helpers."""