From bc0191d7755cbf676278dc31ea5a1287393f2d2f Mon Sep 17 00:00:00 2001 From: Shao Duan <46639847+shaoxiongduan@users.noreply.github.com> Date: Wed, 7 Oct 2026 04:01:53 -0700 Subject: [PATCH] [perf] Harden Ulysses agreement and tune bounded H3 transfers (#1807) Co-authored-by: SolitaryThinker Co-authored-by: Claude Opus 5.5 (1M context) --- docs/contributing/env_vars.md | 1 + docs/contributing/profiling.md | 3 + docs/contributing/ulysses_performance.md | 67 +++++++ .../csrc/comm/ulysses_all_to_all.cu | 20 ++- .../include/comm/ulysses_all_to_all.cuh | 10 +- .../python/fastvideo_kernel/comm_ops.py | 14 +- .../device_communicators/ulysses_a2a.py | 167 ++++++++++++++---- fastvideo/envs.py | 9 + .../distributed/benchmark_ulysses_a2a.py | 91 ++++++++++ .../distributed/check_ulysses_h3_native.py | 81 +++++++++ .../test_ulysses_fault_injection.py | 120 ++++++++++++- .../distributed/test_ulysses_h3_policy.py | 142 +++++++++++++++ .../ulysses_native_test_support.py | 36 ++++ 13 files changed, 718 insertions(+), 43 deletions(-) create mode 100644 docs/contributing/ulysses_performance.md create mode 100644 fastvideo/tests/distributed/benchmark_ulysses_a2a.py create mode 100644 fastvideo/tests/distributed/check_ulysses_h3_native.py create mode 100644 fastvideo/tests/distributed/test_ulysses_h3_policy.py create mode 100644 fastvideo/tests/distributed/ulysses_native_test_support.py diff --git a/docs/contributing/env_vars.md b/docs/contributing/env_vars.md index 3e3756c90..036d4e6ad 100644 --- a/docs/contributing/env_vars.md +++ b/docs/contributing/env_vars.md @@ -183,6 +183,7 @@ longer exists also fails the test, so the fixing pull request deletes its entry. | `FASTVIDEO_HCCL_SO_PATH` | str | unset | distributed | Path to the HCCL library file on Ascend NPUs. Deprecated names: `HCCL_SO_PATH`. | | `FASTVIDEO_WORKER_MULTIPROC_METHOD` | one of spawn, fork, forkserver | `spawn` | distributed | Multiprocessing start method for worker processes. | | `FASTVIDEO_ULYSSES_A2A` | one of off, auto | `off` | distributed | Sequence-parallel all-to-all backend. off uses the NCCL path in DistributedAutograd.AllToAll4D. auto uses the fused NVLink kernel when the group is a load-store accessible mesh of 2, 4, 6, or 8 ranks in eager execution, and the NCCL path otherwise. | +| `FASTVIDEO_ULYSSES_A2A_LONG_TRAINING` | one of auto, chunked | `auto` | distributed | Fused Ulysses all-to-all policy for grad-enabled calls whose window exceeds the long-training plane limit. auto keeps the original unchunked transport. chunked uses bounded chunks; validate the recipe's activation memory first. | | `FASTVIDEO_CONFIGURE_LOGGING` | bool | `1` | logging | Configure logging at import. When true, FastVideo uses its default logging configuration or the file in FASTVIDEO_LOGGING_CONFIG_PATH. | | `FASTVIDEO_LOGGING_CONFIG_PATH` | str | unset | logging | Path to a JSON logging configuration file. | | `FASTVIDEO_LOGGING_LEVEL` | str | `INFO` | logging | Default logging level. | diff --git a/docs/contributing/profiling.md b/docs/contributing/profiling.md index 144b9975c..67d3dfa5d 100644 --- a/docs/contributing/profiling.md +++ b/docs/contributing/profiling.md @@ -1,5 +1,8 @@ # Profiling FastVideo +For GB200 sequence-parallel transport tuning and its long-training memory +policy, see [Ulysses performance](ulysses_performance.md). + !!! warning Profiling is only intended for FastVideo developers and maintainers to understand the proportion of time spent in different parts of the codebase. **FastVideo end-users should never turn on profiling** as it will significantly slow down inference. diff --git a/docs/contributing/ulysses_performance.md b/docs/contributing/ulysses_performance.md new file mode 100644 index 000000000..0eb736ece --- /dev/null +++ b/docs/contributing/ulysses_performance.md @@ -0,0 +1,67 @@ +# Ulysses performance on GB200 + +Enable the fused transport with `FASTVIDEO_ULYSSES_A2A=auto` and rebuild +`fastvideo-kernel` from the same checkout. The tuned path applies to contiguous +BF16 operands with 56 global heads, head dimension 128, and sequence parallelism +of four on GB200. It uses 144 CTAs and exchanges one batch plane at a time. +Older kernel builds retain the original 36-CTA path. + +Each rank keeps at most 1 GiB of registered window storage. A 250,000-token +sequence needs 896,000,000 bytes per plane with this geometry, so packed QKV or +QKVG can use the fused path without a window large enough for the entire pack. +The full output still needs its own GPU allocation. Chunking reduces registered +storage; it does not remove activation or optimizer memory costs. + +Every call collectively agrees on geometry, window capacity, chunking and CTA +count before entering the kernel. Outputs own their storage, and backward uses +the plan saved by its forward. Contiguity alone does not replace these checks. +Existing unsupported-layout, capture, topology and lifecycle fallbacks remain. +The fused path also requires every rank to be on one host: LSA covers +addressability, not locality, so a multi-node NVLink mesh declines the fused +path and keeps the NCCL fallback. + +## Long training sequences + +By default, grad-tracked operands whose per-plane size exceeds 512 MiB keep +the original transport policy, including NCCL fallback for oversized packs. +This avoids enabling a measured allocation-pressure regression in a 250k-token +FSDP4 training recipe. No-grad inference can use tuned chunks directly. + +After configuring sufficient activation memory headroom, enable long training +chunks explicitly: + +```bash +export FASTVIDEO_ULYSSES_A2A=auto +export FASTVIDEO_ULYSSES_A2A_LONG_TRAINING=chunked +``` + +An activation offload policy can provide that headroom. Measure the complete +forward, backward, clipping and optimizer step with resident optimizer state; +transport microbenchmarks alone do not predict long-sequence training speed. +Compare allocator retries alongside step latency. Keep the allocator, FSDP +mesh, attention backend, checkpointing and precision identical across routes. + +`FASTVIDEO_ULYSSES_A2A_LONG_TRAINING=auto` restores the conservative policy. +The 512 MiB boundary describes a transport operand; it is not a model-wide +activation-memory estimate. Revalidate other training recipes and hardware. + +## Validation + +CPU policy regressions live in +`fastvideo/tests/distributed/test_ulysses_h3_policy.py`. On an exclusive group +of four GB200 GPUs, run the native transport gate after rebuilding the kernel: + +```bash +FASTVIDEO_ULYSSES_A2A=auto torchrun --standalone --nproc_per_node=4 \ + fastvideo/tests/distributed/check_ulysses_h3_native.py +``` + +The gate covers exact transport gradients at 32k, 128k and 250k, the long +training opt-in, retained-output ownership, and recovery after rank capability +disagreement. Full-model throughput and absolute model FLOP utilization (MFU) +are separate measurements; this transport gate reports neither. + +`benchmark_ulysses_a2a.py` compares revisions with identical argv and reports +host-to-completion latency including agreement; set `FASTVIDEO_ULYSSES_A2A=off` +for the NCCL baseline. `ulysses_native_test_support.py` loads a separately built +communication extension for development validation. diff --git a/fastvideo-kernel/csrc/comm/ulysses_all_to_all.cu b/fastvideo-kernel/csrc/comm/ulysses_all_to_all.cu index 10193826d..724537e51 100644 --- a/fastvideo-kernel/csrc/comm/ulysses_all_to_all.cu +++ b/fastvideo-kernel/csrc/comm/ulysses_all_to_all.cu @@ -23,6 +23,7 @@ // The per-group context is an ncclDevComm plus a registered symmetric window, // both created here from the caller's ncclComm_t. +#include #include #include #include @@ -162,13 +163,19 @@ bool ulysses_lsa_covers_group(int64_t comm_ptr, int64_t world_size) { // mode == 0: inp [B, S_local, H, D] -> out [B, S_global, H_local, D] // mode == 1: inp [B, S_global, H_local, D] -> out [B, S_local, H, D] // where H is the *global* head count and H_local = H / world_size. -void ulysses_a2a(int64_t handle, torch::Tensor inp, torch::Tensor out, int64_t B, int64_t S_local, - int64_t H, int64_t D, int64_t mode) { +void ulysses_a2a_tuned(int64_t handle, torch::Tensor inp, torch::Tensor out, int64_t B, int64_t S_local, + int64_t H, int64_t D, int64_t mode, int64_t launch_blocks) { auto* ctx = reinterpret_cast(handle); TORCH_CHECK(ctx != nullptr, "handle must come from allocate_ulysses_a2a"); const at::cuda::CUDAGuard device_guard(inp.device()); auto stream = at::cuda::getCurrentCUDAStream(); + TORCH_CHECK(launch_blocks == 36 || launch_blocks == 144, "Ulysses launch must use 36 or 144 CTAs"); + if (launch_blocks == 144) { + const auto* props = at::cuda::getCurrentDeviceProperties(); + TORCH_CHECK(props->major == 10 && props->minor == 0 && props->multiProcessorCount >= 144, + "144-CTA Ulysses launch requires an sm100 device with at least 144 SMs"); + } TORCH_CHECK(inp.is_cuda() && out.is_cuda(), "inp and out must be CUDA tensors"); TORCH_CHECK(inp.is_contiguous() && out.is_contiguous(), "inp and out must be contiguous"); @@ -204,7 +211,7 @@ void ulysses_a2a(int64_t handle, torch::Tensor inp, torch::Tensor out, int64_t B const int64_t num_rows = B * static_cast(W) * S_local; const int blocks = - static_cast(std::max(1, std::min(fi::kMaxBlocks, num_rows))); + static_cast(std::max(1, std::min(launch_blocks, num_rows))); const int threads = fi::kUlyssesThreads; #define LAUNCH_ULYSSES_A2A(T, NG, MODE) \ @@ -266,6 +273,12 @@ void ulysses_a2a(int64_t handle, torch::Tensor inp, torch::Tensor out, int64_t B TORCH_CHECK(status == cudaSuccess, "ulysses_a2a copy-out failed: ", cudaGetErrorString(status)); } +// Preserve the original wheel API and its 36-CTA launch policy. +void ulysses_a2a(int64_t handle, torch::Tensor inp, torch::Tensor out, int64_t B, int64_t S_local, + int64_t H, int64_t D, int64_t mode) { + ulysses_a2a_tuned(handle, inp, out, B, S_local, H, D, mode, 36); +} + void register_ulysses_a2a(pybind11::module_& m) { m.def("allocate_ulysses_a2a", &allocate_ulysses_a2a, "allocate a local ulysses a2a window"); m.def("register_ulysses_a2a_window", ®ister_ulysses_a2a_window, @@ -276,4 +289,5 @@ void register_ulysses_a2a(pybind11::module_& m) { m.def("ulysses_lsa_covers_group", &ulysses_lsa_covers_group, "whether the whole group is load-store accessible"); m.def("ulysses_a2a", &ulysses_a2a, "fused-transpose Ulysses all-to-all over NVLink"); + m.def("ulysses_a2a_tuned", &ulysses_a2a_tuned, "Ulysses all-to-all with a collectively agreed CTA count"); } diff --git a/fastvideo-kernel/include/comm/ulysses_all_to_all.cuh b/fastvideo-kernel/include/comm/ulysses_all_to_all.cuh index f69dc6035..f2281ebe2 100644 --- a/fastvideo-kernel/include/comm/ulysses_all_to_all.cuh +++ b/fastvideo-kernel/include/comm/ulysses_all_to_all.cuh @@ -48,9 +48,13 @@ namespace comm { namespace ulysses { constexpr int kUlyssesThreads = 512; -// Deliberately modest: this is link-bandwidth bound, so a small grid leaves the -// rest of the GPU free without costing throughput. -constexpr int kMaxBlocks = 36; +// Upper bound on the CTAs one launch may use, and therefore the number of NCCL +// LSA barrier slots the host reserves when it creates devComm: the kernel +// indexes its barrier by blockIdx.x, so this must cover the largest launch +// (36 CTAs for the original entrypoint, 144 for the tuned GB200 plan). The +// original 36-CTA grid is link-bandwidth bound and leaves the rest of the GPU +// free without costing throughput. +constexpr int kMaxBlocks = 144; // Shared movement body for the fused-transpose all-to-all (no barriers). // diff --git a/fastvideo-kernel/python/fastvideo_kernel/comm_ops.py b/fastvideo-kernel/python/fastvideo_kernel/comm_ops.py index cc5067303..168376e70 100644 --- a/fastvideo-kernel/python/fastvideo_kernel/comm_ops.py +++ b/fastvideo-kernel/python/fastvideo_kernel/comm_ops.py @@ -29,6 +29,11 @@ def is_available() -> bool: return _ops is not None and all(hasattr(_ops, name) for name in _REQUIRED_OPS) +def supports_tuned_launch() -> bool: + """Whether this build has the expanded barrier capacity and tuned entrypoint.""" + return is_available() and hasattr(_ops, "ulysses_a2a_tuned") + + def _require() -> None: if not is_available(): raise RuntimeError( @@ -73,7 +78,7 @@ def dispose(handle: int) -> None: def all_to_all(handle: int, inp: torch.Tensor, out: torch.Tensor, B: int, S_local: int, H: int, - D: int, mode: int) -> None: + D: int, mode: int, *, blocks: int = 36) -> None: """Run one fused all-to-all on the current stream, writing into ``out``. ``mode == 0``: ``[B, S_local, H, D] -> [B, S_global, H_local, D]`` @@ -83,4 +88,9 @@ def all_to_all(handle: int, inp: torch.Tensor, out: torch.Tensor, B: int, S_loca geometry in the same order. """ _require() - _ops.ulysses_a2a(int(handle), inp, out, int(B), int(S_local), int(H), int(D), int(mode)) + if blocks == 36: + _ops.ulysses_a2a(int(handle), inp, out, int(B), int(S_local), int(H), int(D), int(mode)) + elif blocks == 144 and supports_tuned_launch(): + _ops.ulysses_a2a_tuned(int(handle), inp, out, int(B), int(S_local), int(H), int(D), int(mode), int(blocks)) + else: + raise ValueError(f"unsupported Ulysses launch: blocks={blocks}; rebuild the kernel for tuned launches") diff --git a/fastvideo/distributed/device_communicators/ulysses_a2a.py b/fastvideo/distributed/device_communicators/ulysses_a2a.py index ffbd94b92..eb23a9e98 100644 --- a/fastvideo/distributed/device_communicators/ulysses_a2a.py +++ b/fastvideo/distributed/device_communicators/ulysses_a2a.py @@ -6,6 +6,9 @@ load-store accessible NVLink mesh: same layout, byte-identical results, fewer passes over local memory. Anything else falls back to the NCCL path. """ +import socket +from array import array + import torch import torch.distributed as dist from torch.distributed import ProcessGroup @@ -27,6 +30,8 @@ _DTYPE_CODES = { # Bound persistent registered memory per rank. Larger operands use NCCL instead # of growing the window without limit. MAX_WINDOW_BYTES = 1024**3 +H3_TRAINING_PLANE_LIMIT_BYTES = 512 * 1024**2 +_CONTRACT_SIZE = 12 # (scatter_dim, gather_dim) -> kernel mode. # 0: [B, S_local, H, D] -> [B, S_global, H_local, D] @@ -39,6 +44,28 @@ def is_enabled() -> bool: return envs.FASTVIDEO_ULYSSES_A2A.get() == "auto" +def _host_identity() -> str: + """Namespace-independent identity of this process's physical host. + + ``socket.gethostname()`` is the UTS-namespace name - the container or pod + name under Docker and Kubernetes - so ranks that share one host can report + different names and ranks on different hosts can share one. The kernel boot + id is not namespaced: every container on a host sees the same value and no + two concurrently running hosts share one. The hostname is only a fallback + for images where /proc is not readable. Ranks where only some can read + /proc mix the two schemes and the group declines: the conservative + direction, matching the pre-boot_id behaviour. + """ + try: + with open("/proc/sys/kernel/random/boot_id") as boot_id: + identity = boot_id.read().strip() + if identity: + return identity + except OSError: + pass + return socket.gethostname() + + class _FusedUlyssesA2A(torch.autograd.Function): """Differentiable fused all-to-all. @@ -47,17 +74,22 @@ class _FusedUlyssesA2A(torch.autograd.Function): """ @staticmethod - def forward(ctx, helper: "UlyssesA2AHelper", x: torch.Tensor, mode: int) -> torch.Tensor: # type: ignore[override] + def forward(ctx, helper: "UlyssesA2AHelper", x: torch.Tensor, mode: int, chunked: bool, + blocks: int) -> torch.Tensor: # type: ignore[override] ctx.helper = helper ctx.mode = mode - return helper.run_armed(x, mode) + ctx.chunked = chunked + ctx.blocks = blocks + return helper.run_armed(x, mode, chunked, blocks) @staticmethod def backward(ctx, grad_output: torch.Tensor): # type: ignore[override] # Same numel and dtype as the forward output, so the window is already # sized for it; only contiguity needs restoring. - grad_input = ctx.helper.run_armed(grad_output.contiguous(), 1 - ctx.mode) - return None, grad_input, None + # Reuse the forward plan, even if later calls chose another plan or + # backward executes with grad tracking disabled. + grad_input = ctx.helper.run_armed(grad_output.contiguous(), 1 - ctx.mode, ctx.chunked, ctx.blocks) + return None, grad_input, None, None, None class UlyssesA2AHelper: @@ -76,8 +108,15 @@ class UlyssesA2AHelper: self.pynccl_comm = pynccl_comm self._handle: int | None = None + # Reuse storage, but exchange the current contract on every call. A + # rank-local cache hit cannot establish what peers are doing now. + self._local_contract = array("q", [0] * _CONTRACT_SIZE) + self._local_tensor = torch.frombuffer(self._local_contract, dtype=torch.int64) + self._gathered_tensor = torch.empty(world_size * _CONTRACT_SIZE, dtype=torch.int64, device="cpu") self._nbytes = 0 self._disabled_reason: str | None = None + self._h3_tuning_available: bool | None = None + self._contract_mismatch_logged = False if world_size not in SUPPORTED_WORLD_SIZES: self._disabled_reason = (f"world size {world_size} is not one of " @@ -95,12 +134,12 @@ class UlyssesA2AHelper: return int(getattr(comm, "value", comm)) def _can_attempt(self) -> tuple[bool, str]: - """Whether this rank could use the fused path, without allocating anything.""" + """Check local capability only; the caller exchanges every rank's result.""" try: from fastvideo_kernel import comm_ops if not comm_ops.is_available(): return False, "fastvideo-kernel was built without the Ulysses a2a kernel" - if not comm_ops.lsa_covers_group(self._comm_ptr(), self.world_size): + elif not comm_ops.lsa_covers_group(self._comm_ptr(), self.world_size): return False, "the group is not a load-store-accessible (NVLink) mesh" except Exception as e: # noqa: BLE001 return False, f"backend unavailable ({type(e).__name__}: {e})" @@ -108,7 +147,7 @@ class UlyssesA2AHelper: def _agree(self, ok: bool) -> bool: """Reduce a local yes/no to a group-wide verdict: True only if all agree.""" - vote = torch.tensor([1 if ok else 0], dtype=torch.int32) + vote = torch.tensor([1 if ok else 0], dtype=torch.int32, device="cpu") dist.all_reduce(vote, op=dist.ReduceOp.MIN, group=self.cpu_group) return bool(vote.item()) @@ -158,6 +197,20 @@ class UlyssesA2AHelper: logger.warning("Ulysses partial-context cleanup failed", exc_info=True) return self._agree(cleanup_ok) + def _execution_plan(self, x: torch.Tensor, mode: int) -> tuple[bool, int]: + """Use the measured GB200/SP4 launch only for H3's bf16 head geometry.""" + global_heads = x.shape[2] if mode == 0 else x.shape[2] * self.world_size + if self.world_size != 4 or x.dtype != torch.bfloat16 or global_heads != 56 or x.shape[3] != 128: + return False, 36 + if self._h3_tuning_available is None: + from fastvideo_kernel import comm_ops + + props = torch.cuda.get_device_properties(self.device) + self._h3_tuning_available = ("GB200" in props.name and props.major == 10 and props.minor == 0 + and props.multi_processor_count >= 144 + and getattr(comm_ops, "supports_tuned_launch", lambda: False)()) + return (True, 144) if self._h3_tuning_available else (False, 36) + def _call_signature(self, x: torch.Tensor, scatter_dim: int, gather_dim: int) -> tuple[tuple[int, ...], str]: """Return a rank-comparable call contract and any local decline reason.""" mode = _MODE_FROM_DIMS.get((scatter_dim, gather_dim)) @@ -188,17 +241,33 @@ class UlyssesA2AHelper: status, reason = 0, "sequence length is not divisible by the group" nbytes = int(x.numel() * x.element_size()) + chunked, blocks = False, 36 + if status == 1 and nbytes: + assert mode is not None + chunked, blocks = self._execution_plan(x, mode) + window_bytes = nbytes // shape[0] if chunked else nbytes + if (status == 1 and chunked and torch.is_grad_enabled() and x.requires_grad + and window_bytes > H3_TRAINING_PLANE_LIMIT_BYTES): + policy = envs.FASTVIDEO_ULYSSES_A2A_LONG_TRAINING.get() + if policy == "auto": + # Preserve the complete original long-training path, including + # its gather launch. Faster gathers alone did not avoid the + # allocation-pressure regression in the measured FSDP4 recipe. + chunked, blocks, window_bytes = False, 36, nbytes + elif policy != "chunked": + status, reason = 0, "unsupported FASTVIDEO_ULYSSES_A2A_LONG_TRAINING policy" if status == 1 and nbytes == 0: status, reason = 0, "input is empty" - elif status == 1 and nbytes > MAX_WINDOW_BYTES: - status, reason = 0, f"operand exceeds the {MAX_WINDOW_BYTES}-byte window cap" + elif status == 1 and window_bytes > MAX_WINDOW_BYTES: + status, reason = 0, f"operand window exceeds the {MAX_WINDOW_BYTES}-byte cap" - # status, armed, mode, dtype, B, S, H, D, bytes, capacity. Comparing the + # status, armed, mode, dtype, B, S, H, D, window bytes, capacity, + # chunked, blocks. Comparing the # whole vector prevents equal-size but differently-shaped ranks from # entering the fused kernel with incompatible address math. CUDA device # ordinals are deliberately absent: rank-local ordinals normally differ. - signature = (status, int(self._handle is not None), -1 if mode is None else mode, dtype_code, *shape, nbytes, - self._nbytes) + signature = (status, int(self._handle is not None), -1 if mode is None else mode, dtype_code, *shape, + window_bytes, self._nbytes, int(chunked), blocks) return signature, reason def _agree_call(self, signature: tuple[int, ...]) -> tuple[bool, bool, bool]: @@ -208,18 +277,32 @@ class UlyssesA2AHelper: collective is intentionally eager-only; compiled regions use the NCCL implementation before reaching here. """ - # Host-side Gloo control keeps this agreement outside CUDA graph capture - # and avoids inserting a second NCCL collective ahead of the data path. - local = torch.tensor(signature, dtype=torch.int64) - gathered = torch.empty(self.world_size * local.numel(), dtype=local.dtype) - dist.all_gather_into_tensor(gathered, local, group=self.cpu_group) - contracts = gathered.view(self.world_size, local.numel()) - identical = bool(torch.all(contracts == contracts[0]).item()) - statuses = contracts[:, 0] - use_fused = identical and bool(torch.all(statuses == 1).item()) - permanently_unavailable = bool(torch.any(statuses < 0).item()) - lifecycle_consistent = (bool(torch.all(contracts[:, 1] == contracts[0, 1]).item()) - and bool(torch.all(contracts[:, -1] == contracts[0, -1]).item())) + # Host-side Gloo control avoids inserting a second NCCL collective ahead + # of the data path. It is CPU-only and allocation-free, so it also stays + # valid when the caller is inside a CUDA graph capture. + if len(signature) != _CONTRACT_SIZE: + raise RuntimeError(f"Ulysses call contract has {len(signature)} fields, expected {_CONTRACT_SIZE}") + self._local_contract[:] = array("q", signature) + dist.all_gather_into_tensor(self._gathered_tensor, self._local_tensor, group=self.cpu_group) + values = self._gathered_tensor.tolist() + contracts = [values[start:start + _CONTRACT_SIZE] for start in range(0, len(values), _CONTRACT_SIZE)] + first = contracts[0] + use_fused = first[0] == 1 and all(contract == first for contract in contracts) + permanently_unavailable = any(contract[0] < 0 for contract in contracts) + lifecycle_consistent = all(contract[1] == first[1] and contract[9] == first[9] for contract in contracts) + if (not use_fused and not permanently_unavailable and lifecycle_consistent + and any(contract != first for contract in contracts) and not self._contract_mismatch_logged): + # A transient disagreement (a peer declined this call) recovers on a + # later call, but a persistent one - mixed kernel builds or launch + # capabilities - would silently keep every later call on NCCL. + self._contract_mismatch_logged = True + differing = [ + index for index in range(_CONTRACT_SIZE) if any(contract[index] != first[index] for contract in contracts) + ] + details = ", ".join( + f"field {index}: {sorted({contract[index] for contract in contracts})}" for index in differing) + logger.info("Ulysses fused all-to-all declined: ranks disagree on the call contract (%s); " + "this and later disagreeing calls use the NCCL path", details) return use_fused, permanently_unavailable, lifecycle_consistent def _build(self, nbytes: int) -> bool: @@ -327,8 +410,8 @@ class UlyssesA2AHelper: # -- collective ---------------------------------------------------------- - def run_armed(self, x: torch.Tensor, mode: int) -> torch.Tensor: - """Run one collective on an already-armed context.""" + def run_armed(self, x: torch.Tensor, mode: int, chunked: bool, blocks: int) -> torch.Tensor: + """Run one call's collective(s) on an already-armed context.""" assert self._handle is not None, "run_armed called on an unarmed helper" from fastvideo_kernel import comm_ops @@ -340,7 +423,22 @@ class UlyssesA2AHelper: B, S_global, H_local, D = x.shape S_local, H = S_global // w, H_local * w out = torch.empty(B, S_local, H, D, dtype=x.dtype, device=x.device) - comm_ops.all_to_all(self._handle, x, out, B, S_local, H, D, mode) + if chunked: + # Each copy completes on this stream before the next plane reuses + # the registered window. The full result owns its storage; saved + # activations never alias the reusable window. + for plane in range(B): + comm_ops.all_to_all(self._handle, + x[plane:plane + 1], + out[plane:plane + 1], + 1, + S_local, + H, + D, + mode, + blocks=blocks) + else: + comm_ops.all_to_all(self._handle, x, out, B, S_local, H, D, mode, blocks=blocks) return out def try_all_to_all_4D(self, x: torch.Tensor, scatter_dim: int, gather_dim: int) -> torch.Tensor | None: @@ -365,7 +463,7 @@ class UlyssesA2AHelper: return None mode = signature[2] - nbytes = signature[-2] + nbytes = signature[8] if self._handle is None: if not self._build(nbytes): return None @@ -376,7 +474,7 @@ class UlyssesA2AHelper: if not self._build(nbytes): return None - return _FusedUlyssesA2A.apply(self, x, mode) + return _FusedUlyssesA2A.apply(self, x, mode, bool(signature[10]), signature[11]) def maybe_create_helper(cpu_group: ProcessGroup | None, device_group: ProcessGroup | None, world_size: int, @@ -404,9 +502,14 @@ def maybe_create_helper(cpu_group: ProcessGroup | None, device_group: ProcessGro except Exception as e: # noqa: BLE001 - converted to a group verdict below reason = f"helper construction failed ({type(e).__name__}: {e})" - vote = torch.tensor([int(helper is not None)], dtype=torch.int32) - dist.all_reduce(vote, op=dist.ReduceOp.MIN, group=cpu_group) - if not bool(vote.item()): + # Every rank reaches the same exchange, including configuration, constructor, + # and backend failures. LSA covers addressability, not single-host locality. + gathered: list[tuple[str, bool]] = [("", False)] * world_size + dist.all_gather_object(gathered, (_host_identity(), helper is not None), group=cpu_group) + host_ids = {identity for identity, _ in gathered} + if len(host_ids) != 1: + reason = f"ranks do not share one host (host identities differ: {sorted(host_ids)})" + if len(host_ids) != 1 or not all(ok for _, ok in gathered): if dist.get_rank(cpu_group) == 0: logger.info("Ulysses fused all-to-all unavailable: %s", reason or "a peer rank declined") return None diff --git a/fastvideo/envs.py b/fastvideo/envs.py index 1adb704c0..b52d7866b 100644 --- a/fastvideo/envs.py +++ b/fastvideo/envs.py @@ -295,6 +295,15 @@ FASTVIDEO_ULYSSES_A2A = EnvChoice( doc="Sequence-parallel all-to-all backend. off uses the NCCL path in DistributedAutograd.AllToAll4D. auto uses " "the fused NVLink kernel when the group is a load-store accessible mesh of 2, 4, 6, or 8 ranks in eager " "execution, and the NCCL path otherwise.") +# Keep the original long-training transport by default. Opt into bounded +# chunks after validating the training recipe's activation memory budget. +FASTVIDEO_ULYSSES_A2A_LONG_TRAINING = EnvChoice( + "auto", + choices=("auto", "chunked"), + category="distributed", + doc="Fused Ulysses all-to-all policy for grad-enabled calls whose window exceeds the long-training plane limit. " + "auto keeps the original unchunked transport. chunked uses bounded chunks; validate the recipe's activation " + "memory first.") # ================== Logging ================== diff --git a/fastvideo/tests/distributed/benchmark_ulysses_a2a.py b/fastvideo/tests/distributed/benchmark_ulysses_a2a.py new file mode 100644 index 000000000..982e003c0 --- /dev/null +++ b/fastvideo/tests/distributed/benchmark_ulysses_a2a.py @@ -0,0 +1,91 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Run with torchrun on one NVLink host; compare revisions using identical argv. + +Reports host-to-completion latency (including agreement), p50/p95 of rank-max +samples, for bf16 Wan2.1-14B and MiniMax-H3 5s/720p attention operands. No model +weights are required. Set FASTVIDEO_ULYSSES_A2A=off for the NCCL baseline. +""" + +import argparse +import json +import os +import statistics +import time +from pathlib import Path + +import torch +import torch.distributed as dist + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--iters", type=int, default=60) + parser.add_argument("--warmup", type=int, default=15) + parser.add_argument("--rounds", type=int, default=3) + parser.add_argument("--tag", required=True) + parser.add_argument("--output", type=Path, required=True) + parser.add_argument("--expect-fused", action="store_true") + args = parser.parse_args() + + from fastvideo.distributed import cleanup_dist_env_and_memory, maybe_init_distributed_environment_and_model_parallel + from fastvideo.distributed.device_communicators.base_device_communicator import DeviceCommunicatorBase + from fastvideo.distributed.parallel_state import get_sp_group + + world = int(os.environ["WORLD_SIZE"]) + rank = int(os.environ["RANK"]) + torch.cuda.set_device(int(os.environ["LOCAL_RANK"])) + device = torch.device("cuda", torch.cuda.current_device()) + torch.manual_seed(2026 + rank) + maybe_init_distributed_environment_and_model_parallel(1, world) + comm = get_sp_group().device_communicator + records = [] + try: + for name, sequence, heads in [("Wan2.1-14B", 75600, 40), ("MiniMax-H3", 37296, 56)]: + assert sequence % world == 0 + scatter = torch.randn(3, sequence // world, heads, 128, dtype=torch.bfloat16, device=device) + gather = torch.randn(1, sequence, heads // world, 128, dtype=torch.bfloat16, device=device) + for operand, dims in [(scatter, (2, 1)), (gather, (1, 2))]: + actual = comm.all_to_all_4D(operand, *dims) + expected = DeviceCommunicatorBase.all_to_all_4D(comm, operand, *dims) + assert torch.equal(actual, expected), f"{name} {dims}: parity failed" + del actual, expected + armed = comm.ulysses_a2a is not None and comm.ulysses_a2a._handle is not None + if args.expect_fused: + assert armed, "benchmark requires all ranks to engage the fused kernel" + for repeat in range(args.rounds): + for operation in ("scatter", "gather", "layer"): + def run(): + if operation in ("scatter", "layer"): + comm.all_to_all_4D(scatter, 2, 1) + if operation in ("gather", "layer"): + comm.all_to_all_4D(gather, 1, 2) + + for _ in range(args.warmup): + run() + torch.cuda.synchronize() + samples = [] + for _ in range(args.iters): + dist.barrier(group=get_sp_group().cpu_group) + start = time.perf_counter_ns() + run() + torch.cuda.synchronize() + samples.append((time.perf_counter_ns() - start) / 1000) + samples_tensor = torch.tensor(samples, dtype=torch.float64, device=device) + dist.all_reduce(samples_tensor, op=dist.ReduceOp.MAX) + maximums = samples_tensor.cpu().tolist() + record = dict(tag=args.tag, model=name, operation=operation, round=repeat, + world_size=world, fused=armed, p50_us=statistics.median(maximums), + p95_us=sorted(maximums)[int(0.95 * (len(maximums) - 1))], + rank_max_samples_us=maximums) + records.append(record) + if rank == 0: + print(json.dumps({k: v for k, v in record.items() if k != "rank_max_samples_us"}), flush=True) + del scatter, gather + if rank == 0: + args.output.write_text(json.dumps(records, indent=2) + "\n") + finally: + cleanup_dist_env_and_memory() + + +if __name__ == "__main__": + main() diff --git a/fastvideo/tests/distributed/check_ulysses_h3_native.py b/fastvideo/tests/distributed/check_ulysses_h3_native.py new file mode 100644 index 000000000..46b8c1104 --- /dev/null +++ b/fastvideo/tests/distributed/check_ulysses_h3_native.py @@ -0,0 +1,81 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Explicit torchrun GPU gate for the native H3 transport policy (4x GB200).""" +import os +from unittest.mock import patch + +import torch +import torch.distributed as dist + +from fastvideo import envs +from fastvideo.distributed import cleanup_dist_env_and_memory, maybe_init_distributed_environment_and_model_parallel +from fastvideo.distributed.device_communicators.base_device_communicator import DeviceCommunicatorBase +from fastvideo.distributed.parallel_state import get_sp_group +from fastvideo.tests.distributed.test_ulysses_a2a_parity import _check_shape + + +def _check_native_shape(helper, batch, local_sequence, dtype, device, expected_plans): + # Numerical parity also passes when both sides use NCCL. Assert the actual + # helper executions, including the two backward collectives, separately. + with patch.object(helper, 'run_armed', wraps=helper.run_armed) as run: + _check_shape(batch, local_sequence, 56, 128, dtype, 4, device) + plans = [call.args[1:] for call in run.call_args_list] + assert plans == expected_plans, (batch, local_sequence, dtype, plans, expected_plans) + + +def main(): + # A caller may run this from a shell configured for long training. + with envs.FASTVIDEO_ULYSSES_A2A_LONG_TRAINING.override('auto'): + _run() + + +def _run(): + rank = int(os.environ['RANK']) + assert int(os.environ['WORLD_SIZE']) == 4 + torch.cuda.set_device(int(os.environ['LOCAL_RANK'])) + device = torch.device('cuda', int(os.environ['LOCAL_RANK'])) + torch.manual_seed(20260906 + rank) + maybe_init_distributed_environment_and_model_parallel(1, 4) + comm = get_sp_group().device_communicator + helper = comm.ulysses_a2a + assert helper is not None + for dtype in [torch.bfloat16, torch.float16, torch.float32]: + tuned = dtype == torch.bfloat16 + plans = [(mode, tuned, 144 if tuned else 36) for mode in [0, 1, 0, 1, 1, 0]] + _check_native_shape(helper, 3, 64, dtype, device, plans) + for sequence in [32000, 128000, 250000]: + modes = [0, 1] if sequence == 250000 else [0, 1, 0, 1, 1, 0] + _check_native_shape(helper, 4, sequence // 4, torch.bfloat16, device, + [(mode, True, 144) for mode in modes]) + assert helper._nbytes <= 1024**3 + with envs.FASTVIDEO_ULYSSES_A2A_LONG_TRAINING.override('chunked'): + _check_native_shape(helper, 4, 250000 // 4, torch.bfloat16, device, + [(mode, True, 144) for mode in [0, 1, 0, 1, 1, 0]]) + x = torch.randn(3, 8000, 56, 128, dtype=torch.bfloat16, device=device, requires_grad=True) + first = comm.all_to_all_4D(x, 2, 1) + saved = first.detach().clone() + # A later call and inverse cannot overwrite an earlier autograd output. + second = comm.all_to_all_4D(x.detach() + 1, 2, 1) + comm.all_to_all_4D(second, 1, 2) + assert torch.equal(first, saved) + grad = torch.randn_like(first) + first.backward(grad) + expected = DeviceCommunicatorBase.all_to_all_4D(comm, grad, 1, 2) + assert torch.equal(x.grad, expected) + + # A rank with an older launch capability must make the whole group decline. + helper._h3_tuning_available = rank != 0 + assert helper.try_all_to_all_4D(x.detach(), 2, 1) is None + helper._h3_tuning_available = True + assert helper.try_all_to_all_4D(x.detach(), 2, 1) is not None + contract, _ = helper._call_signature(x.detach(), 2, 1) + assert contract[10:] == (1, 144) + dist.barrier() + if rank == 0: + print('NATIVE_H3_OK dtypes=3 lengths=32000,128000,250000 exact_gradients=True ownership=True recovery=True ' + 'execution_plans=True', + flush=True) + cleanup_dist_env_and_memory() + + +if __name__ == '__main__': + main() diff --git a/fastvideo/tests/distributed/test_ulysses_fault_injection.py b/fastvideo/tests/distributed/test_ulysses_fault_injection.py index dbc7efc33..d81cdc83a 100644 --- a/fastvideo/tests/distributed/test_ulysses_fault_injection.py +++ b/fastvideo/tests/distributed/test_ulysses_fault_injection.py @@ -9,6 +9,7 @@ import socket import subprocess import sys from pathlib import Path +from unittest.mock import patch import pytest import torch @@ -38,6 +39,87 @@ FAULT_CASES = [ (2, 0, "lifecycle"), (4, 3, "lifecycle"), ] +FAULT_CASES += [(world, fault_rank, stage) + for world, fault_rank in [(2, 1), (4, 0)] + for stage in ("constructor", "backend", "pynccl", "hostname", "late_layout", "late_capture", + "late_configuration", "late_capacity", "late_lifecycle", "cached_shape", "cached_mode", + "decline_then_recover", "late_real_capture", "late_mixed_capture")] + + +def _check_after_warmup(helper, rank: int, fault_rank: int, stage: str, device: torch.device) -> None: + """Exercise divergence after an armed signature has already succeeded twice.""" + from fastvideo.distributed.device_communicators import ulysses_a2a + + assert helper is not None, "this regression must exercise an available fused helper" + x = torch.randn(3, 64, 8, 64, device=device, dtype=torch.bfloat16) + if stage == "decline_then_recover": + # Populate a declined contract before any window or successful call. + operand = torch.stack((x, x), dim=-1)[..., 0] if rank == fault_rank else x + assert helper.try_all_to_all_4D(operand, 2, 1) is None + assert helper._handle is None + for _ in range(3): + y = helper.try_all_to_all_4D(x, 2, 1) + assert y is not None + assert torch.equal(helper.try_all_to_all_4D(y, 1, 2), x) + alt = x.reshape(3, 32, 16, 64) + for _ in range(2): + assert helper.try_all_to_all_4D(alt, 2, 1) is not None + + if stage in ("late_real_capture", "late_mixed_capture"): + capture_stream = torch.cuda.Stream(device=device) + graph = torch.cuda.CUDAGraph() + # Synchronize warmup before starting the real, default-global capture. + torch.cuda.synchronize(device) + if stage == "late_real_capture" or rank == fault_rank: + with torch.cuda.graph(graph, stream=capture_stream): + graph_output = x + 1 + assert helper.try_all_to_all_4D(x, 2, 1) is None + graph.replay() + torch.cuda.synchronize(device) + assert torch.equal(graph_output, x + 1) + else: + assert helper.try_all_to_all_4D(x, 2, 1) is None + recovered = helper.try_all_to_all_4D(x, 2, 1) + assert recovered is not None + assert torch.equal(helper.try_all_to_all_4D(recovered, 1, 2), x) + print(f"RANK_DONE rank={rank} recovered=True", flush=True) + if rank == 0: + print(f"ALL_RANKS_COMPLETED world={helper.world_size} stage={stage}", flush=True) + return + + dims = (2, 1) + operand = x + if rank == fault_rank: + if stage in ("late_layout", "decline_then_recover"): + operand = torch.stack((x, x), dim=-1)[..., 0] + elif stage == "cached_shape": + operand = alt + elif stage == "cached_mode": + operand, dims = y, (1, 2) + elif stage == "late_lifecycle": + helper._nbytes += 1 + + # Both the rank with the changed input and the ranks with familiar inputs + # must decline together. Calling NCCL with mismatched shapes/modes would + # itself be invalid, so check the helper result before the fallback path. + with patch.object(torch.cuda, "is_current_stream_capturing", + return_value=rank == fault_rank and stage == "late_capture"), \ + patch.object(ulysses_a2a, "is_enabled", + return_value=not (rank == fault_rank and stage == "late_configuration")), \ + patch.object(ulysses_a2a, "MAX_WINDOW_BYTES", + 1 if rank == fault_rank and stage == "late_capacity" else 1024**3): + assert helper.try_all_to_all_4D(operand, *dims) is None, "a peer entered the fused path alone" + + if stage == "late_lifecycle": + assert helper._handle is None and helper._disabled_reason is not None + else: + # A transient decline must not poison a signature when peers recover. + recovered = helper.try_all_to_all_4D(x, 2, 1) + assert recovered is not None + assert torch.equal(helper.try_all_to_all_4D(recovered, 1, 2), x) + print(f"RANK_DONE rank={rank} recovered=True", flush=True) + if rank == 0: + print(f"ALL_RANKS_COMPLETED world={helper.world_size} stage={stage}", flush=True) def _worker() -> None: @@ -67,6 +149,31 @@ def _worker() -> None: lambda self: (False, "injected capability failure")) elif rank == fault_rank and fault_stage == "configuration": envs.FASTVIDEO_ULYSSES_A2A.set("off") + elif rank == fault_rank and fault_stage == "constructor": + from fastvideo.distributed.device_communicators.ulysses_a2a import UlyssesA2AHelper + + def _fail_constructor(*args, **kwargs): + raise RuntimeError("injected helper constructor failure") + + UlyssesA2AHelper.__init__ = _fail_constructor + elif rank == fault_rank and fault_stage == "backend": + from fastvideo_kernel import comm_ops + + comm_ops.is_available = lambda: False + elif rank == fault_rank and fault_stage == "pynccl": + from fastvideo.distributed.device_communicators.pynccl import PyNcclCommunicator + + real_init = PyNcclCommunicator.__init__ + + def _disable_pynccl(self, *args, **kwargs): + real_init(self, *args, **kwargs) + self.disabled = True + + PyNcclCommunicator.__init__ = _disable_pynccl + elif rank == fault_rank and fault_stage == "hostname": + from fastvideo.distributed.device_communicators import ulysses_a2a + + ulysses_a2a._host_identity = lambda: "injected-other-host" maybe_init_distributed_environment_and_model_parallel(1, world) communicator = get_sp_group().device_communicator @@ -74,11 +181,15 @@ def _worker() -> None: try: if helper is None: - if fault_stage not in ("capability", "configuration"): + if fault_stage not in ("capability", "configuration", "constructor", "backend", "pynccl", "hostname"): if rank == 0: print("UNAVAILABLE helper not created", flush=True) return + if fault_stage.startswith("late_") or fault_stage in ("cached_shape", "cached_mode", "decline_then_recover"): + _check_after_warmup(helper, rank, fault_rank, fault_stage, device) + return + if helper is not None and rank == fault_rank and fault_stage == "allocation": helper._allocate = lambda nbytes: (_ for _ in ()).throw( RuntimeError(f"injected allocation failure for {nbytes} bytes")) @@ -195,8 +306,11 @@ def test_rank_local_setup_failure_falls_back_group_wide(world: int, fault_rank: pytest.skip(process.stdout.strip()) assert process.returncode == 0 and "ALL_RANKS_COMPLETED" in process.stdout, ( f"stdout:\n{process.stdout}\nstderr:\n{process.stderr[-6000:]}") - armed = dict(re.findall(r"RANK_DONE rank=(\d+) armed=(True|False)", process.stdout)) - assert len(armed) == world and set(armed.values()) == {"False"} + if fault_stage.startswith("late_") or fault_stage in ("cached_shape", "cached_mode", "decline_then_recover"): + assert len(re.findall(r"RANK_DONE rank=\d+ recovered=True", process.stdout)) == world + else: + armed = dict(re.findall(r"RANK_DONE rank=(\d+) armed=(True|False)", process.stdout)) + assert len(armed) == world and set(armed.values()) == {"False"} if __name__ == "__main__" and "--worker" in sys.argv: diff --git a/fastvideo/tests/distributed/test_ulysses_h3_policy.py b/fastvideo/tests/distributed/test_ulysses_h3_policy.py new file mode 100644 index 000000000..f83eb7331 --- /dev/null +++ b/fastvideo/tests/distributed/test_ulysses_h3_policy.py @@ -0,0 +1,142 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Policy, rank agreement and saved backward-plan regressions without a GPU.""" +import sys +from types import ModuleType, SimpleNamespace +from unittest.mock import Mock, mock_open, patch + +import pytest +import torch + +import fastvideo.envs as envs +from fastvideo.distributed.device_communicators import ulysses_a2a as ulysses + + +@pytest.fixture(autouse=True) +def _default_long_training_policy(env_overrides): + # A caller may run these tests from a shell configured for long training. + env_overrides.enter_context(envs.FASTVIDEO_ULYSSES_A2A_LONG_TRAINING.override('auto')) + + +def _helper(): + return ulysses.UlyssesA2AHelper(object(), object(), 4, torch.device('cuda:0'), object()) + + +def _operand(planes, global_sequence, mode, requires_grad): + shape = (planes, global_sequence // 4, 56, 128) if mode == 0 else (planes, global_sequence, 14, 128) + x = Mock(shape=shape, dtype=torch.bfloat16, device=torch.device('cuda:0'), is_cuda=True, + requires_grad=requires_grad) + x.dim.return_value = 4 + x.numel.return_value = planes * global_sequence * 14 * 128 + x.element_size.return_value = 2 + x.is_contiguous.return_value = True + return x + + +@pytest.mark.parametrize('sequence,large_pack', [(32000, False), (64000, False), (128000, False), + (149796, False), (149800, True), (250000, True)]) +@pytest.mark.parametrize('planes,mode,training', [(3, 0, False), (4, 0, True), (1, 1, True), (4, 1, True)]) +def test_h3_policy_caps_planes_and_preserves_long_training_default(sequence, large_pack, planes, mode, training): + helper = _helper() + x = _operand(planes, sequence, mode, training) + with patch.object(helper, '_execution_plan', return_value=(True, 144)), \ + patch.object(ulysses, 'is_enabled', return_value=True), \ + patch('torch.cuda.is_current_stream_capturing', return_value=False), torch.set_grad_enabled(training): + signature, _ = helper._call_signature(x, *((2, 1) if mode == 0 else (1, 2))) + declined = training and planes > 1 and large_pack + assert signature[0] == (0 if declined else 1) + original_plan = training and large_pack + assert signature[8] == sequence * 14 * 128 * 2 * (planes if original_plan else 1) + assert signature[10:] == ((0, 36) if original_plan else (1, 144)) + + +def test_long_training_chunk_opt_in_keeps_a_bounded_window(): + helper = _helper() + x = _operand(4, 250000, 0, True) + with patch.object(helper, '_execution_plan', return_value=(True, 144)), \ + patch.object(ulysses, 'is_enabled', return_value=True), \ + envs.FASTVIDEO_ULYSSES_A2A_LONG_TRAINING.override('chunked'), \ + patch('torch.cuda.is_current_stream_capturing', return_value=False), torch.enable_grad(): + signature, _ = helper._call_signature(x, 2, 1) + assert signature[0] == 1 + assert signature[8] == 896000000 + assert signature[8] <= ulysses.MAX_WINDOW_BYTES + assert signature[10:] == (1, 144) + + +def _install_older_kernel_wheel(monkeypatch): + """Stand in for a fastvideo-kernel wheel built before the tuned H3 launch. + + Its comm_ops has no supports_tuned_launch(). The installed wheel may predate + comm_ops entirely or already ship the tuned entrypoint, so it is not imported. + """ + package = ModuleType('fastvideo_kernel') + comm_ops = ModuleType('fastvideo_kernel.comm_ops') + package.comm_ops = comm_ops # type: ignore[attr-defined] + monkeypatch.setitem(sys.modules, 'fastvideo_kernel', package) + monkeypatch.setitem(sys.modules, 'fastvideo_kernel.comm_ops', comm_ops) + + +def test_older_wheel_keeps_original_launch_and_capacity_policy(monkeypatch): + helper = _helper() + x = _operand(3, 250000, 0, False) + _install_older_kernel_wheel(monkeypatch) + props = SimpleNamespace(name='NVIDIA GB200', major=10, minor=0, multi_processor_count=152) + with patch('torch.cuda.get_device_properties', return_value=props), \ + patch.object(ulysses, 'is_enabled', return_value=True), \ + patch('torch.cuda.is_current_stream_capturing', return_value=False): + signature, _ = helper._call_signature(x, 2, 1) + assert signature[0] == 0 + assert signature[8] == x.numel() * 2 + assert signature[10:] == (0, 36) + + +def test_rank_disagreement_on_launch_or_chunking_declines_and_recovers(): + helper = _helper() + signature = (1, 0, 0, 2, 4, 32000, 56, 128, 458752000, 0, 1, 144) + for changed_index, value in [(10, 0), (11, 36), (0, 0)]: + peer = list(signature) + peer[changed_index] = value + + def gather(output, local, **kwargs): + output.copy_(torch.tensor([*signature, *peer, *signature, *signature])) + + with patch.object(ulysses.dist, 'all_gather_into_tensor', side_effect=gather): + assert helper._agree_call(signature) == (False, False, True) + + def unanimous(output, local, **kwargs): + output.copy_(local.repeat(4)) + + with patch.object(ulysses.dist, 'all_gather_into_tensor', side_effect=unanimous): + assert helper._agree_call(signature) == (True, False, True) + + +def test_host_identity_is_namespace_independent_with_a_hostname_fallback(): + # Ranks in separate containers share the kernel boot id but not the UTS + # hostname, so the single-host agreement must not be keyed on the hostname. + with patch('builtins.open', mock_open(read_data='boot-123\n')): + assert ulysses._host_identity() == 'boot-123' + with patch('builtins.open', mock_open(read_data='')), \ + patch.object(ulysses.socket, 'gethostname', return_value='pod-a'): + assert ulysses._host_identity() == 'pod-a' + with patch('builtins.open', side_effect=OSError), \ + patch.object(ulysses.socket, 'gethostname', return_value='pod-b'): + assert ulysses._host_identity() == 'pod-b' + + +def test_backward_uses_its_own_forward_plan_after_another_call(): + calls = [] + + class Helper: + def run_armed(self, x, mode, chunked, blocks): + calls.append((mode, chunked, blocks)) + return x.clone() + + helper = Helper() + x = torch.randn(4, requires_grad=True) + first = ulysses._FusedUlyssesA2A.apply(helper, x, 0, True, 144) + second = ulysses._FusedUlyssesA2A.apply(helper, x, 1, False, 36) + first.sum().backward() + assert calls == [(0, True, 144), (1, False, 36), (1, True, 144)] + assert torch.equal(first, second) + assert first.data_ptr() != second.data_ptr() + assert torch.equal(x.grad, torch.ones_like(x)) diff --git a/fastvideo/tests/distributed/ulysses_native_test_support.py b/fastvideo/tests/distributed/ulysses_native_test_support.py new file mode 100644 index 000000000..3850e54ae --- /dev/null +++ b/fastvideo/tests/distributed/ulysses_native_test_support.py @@ -0,0 +1,36 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Load a separately built communication extension for development validation.""" +import importlib.util +import runpy +import sys +from pathlib import Path + +import torch # noqa: F401 - load the extension's libtorch dependencies first + + +def load_module(name: str, path: Path): + spec = importlib.util.spec_from_file_location(name, path) + if spec is None or spec.loader is None: + raise ImportError(str(path)) + module = importlib.util.module_from_spec(spec) + sys.modules[name] = module + spec.loader.exec_module(module) + return module + + +def install_native_ops(source: Path, extension: Path): + import fastvideo_kernel + + native = load_module('ulysses_native', extension) + wrapper = load_module('fastvideo_kernel.comm_ops', + source / 'fastvideo-kernel/python/fastvideo_kernel/comm_ops.py') + wrapper._ops = native + fastvideo_kernel.comm_ops = wrapper + return wrapper + + +if __name__ == '__main__': + source, extension, target = map(Path, sys.argv[1:4]) + install_native_ops(source, extension) + sys.argv = [str(target), *sys.argv[4:]] + runpy.run_path(str(target), run_name='__main__')