diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/indexed_order.py b/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/indexed_order.py new file mode 100644 index 00000000000..d7c9c63ed0c --- /dev/null +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/indexed_order.py @@ -0,0 +1,53 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Ordered sequence with indexed item lookup.""" + +from collections.abc import Iterator +from typing import Generic, TypeVar + +T = TypeVar("T") + + +class IndexedOrder(Generic[T]): + """Insertion order with constant-time successor lookup by item.""" + + def __init__(self) -> None: + """Create an empty indexed order.""" + self._items: list[T] = [] + self._index_by_item: dict[T, int] = {} + + def append(self, item: T) -> None: + """Append ``item`` to the order. + + Args: + item: Item to append. + + Raises: + ValueError: If ``item`` is already present in the order. + """ + if item in self._index_by_item: + raise ValueError("IndexedOrder does not support duplicate items.") + self._index_by_item[item] = len(self._items) + self._items.append(item) + + def __iter__(self) -> Iterator[T]: + """Iterate over items in order.""" + return iter(self._items) + + def next_item(self, item: T) -> T | None: + """Return the item that follows ``item``, if any.""" + index = self._index_by_item[item] + next_index = index + 1 + return self._items[next_index] if next_index < len(self._items) else None diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/module.py b/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/module.py index 71144b311d9..a1c4999d141 100644 --- a/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/module.py +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/module.py @@ -14,8 +14,6 @@ """Module mixin for the minimal Megatron-FSDP path.""" -import dataclasses -from collections import deque from collections.abc import Callable from typing import Literal, cast @@ -24,28 +22,26 @@ from torch.distributed import DeviceMesh from ..mixed_precision import MixedPrecisionPolicy +from .indexed_order import IndexedOrder from .parameter_group import FsdpParameterGroup, contained_in_parameter_group from .placement import MeshAxis, Placements -@dataclasses.dataclass(frozen=True) -class DelayedRelease: - """A module whose unsharded storage can be released after its consumer event.""" - - consumer_event: torch.cuda.Event | None - module: "FsdpModule" - - class FsdpContext: - """Runtime state, stream, and release scheduler shared by one FSDP subtree.""" + """Runtime stream and prefetch state shared by one FSDP subtree.""" allgather_stream: torch.cuda.Stream - delayed_releases: deque[DelayedRelease] + reduce_scatter_stream: torch.cuda.Stream # HFSDP/HSDP need explicit last-microbatch state. First-microbatch state is # unnecessary because it can be detected when ``model_weight``, after syncing # from ``main_weight``, has placements different from ``Placements.optimizer``. is_last_microbatch: bool root_module: "FsdpModule" + # Static orders used to drive all-gather prefetch. We may want to switch to + # capturing runtime order if static module order proves too fragile. Each + # FsdpModule tracks its own materialized state via ``FsdpModule._unshard_event``. + forward_order: IndexedOrder["FsdpModule"] + backward_order: IndexedOrder["FsdpModule"] def __init__(self, device: torch.device, root_module: "FsdpModule") -> None: """Create rank-local runtime state for a root FSDP subtree. @@ -56,26 +52,29 @@ def __init__(self, device: torch.device, root_module: "FsdpModule") -> None: """ self.root_module = root_module self.is_last_microbatch = True - self.delayed_releases = deque() + self.forward_order = IndexedOrder() + self.backward_order = IndexedOrder() with torch.cuda.device(device): self.allgather_stream = torch.cuda.Stream() + self.reduce_scatter_stream = torch.cuda.Stream() + + def current_stream(self) -> torch.cuda.Stream: + """Current stream on this context's device.""" + return torch.cuda.current_stream(self.allgather_stream.device) + + def register_post_backward_final_callback(self) -> None: + """Register this root context's final callback for the current backward. - def enqueue_release(self, module: "FsdpModule") -> None: - """Queue a module's unsharded storage for delayed release.""" - consumer_event = torch.cuda.current_stream(self.allgather_stream.device).record_event() - self.delayed_releases.append(DelayedRelease(consumer_event=consumer_event, module=module)) + Root ``post_backward()`` means only that root-owned parameters have + accumulated gradients; it may run before descendant reductions, or not + run at all when the root owns no trainable parameters. Waiting at + autograd completion orders consumers after every descendant reduction. + """ - def drain_delayed_releases(self, target_length: int) -> None: - """Release queued module storages FIFO until the queue reaches ``target_length``.""" - if target_length < 0: - raise ValueError(f"target_length must be non-negative, got {target_length}.") + def post_backward_final_callback() -> None: + self.current_stream().wait_stream(self.reduce_scatter_stream) - while len(self.delayed_releases) > target_length: - delayed_release = self.delayed_releases.popleft() - with torch.cuda.stream(self.allgather_stream): - if delayed_release.consumer_event is not None: - self.allgather_stream.wait_event(delayed_release.consumer_event) - delayed_release.module.release_unsharded_storage() + torch.autograd.Variable._execution_engine.queue_callback(post_backward_final_callback) class FsdpModule: @@ -88,6 +87,10 @@ class FsdpModule: _context: FsdpContext | None _ready_grad_parameters: set[nn.Parameter] _num_trainable_parameters: int + # Event recorded after this FsdpModule's full parameters are materialized. + # ``None`` lets pre_forward enqueue an all-gather unless an earlier FsdpModule + # already prefetched this module. + _unshard_event: torch.cuda.Event | None def __init__( self, @@ -99,6 +102,7 @@ def __init__( """Initialize FSDP runtime state on an already-constructed module.""" self._context = None self._name = None + self._unshard_event = None owned_parameters = _collect_owned_parameters(self) axis_indices = tuple(_axis_index(mesh, axis) for axis in placements.dp_axes) assert axis_indices == tuple( @@ -147,8 +151,11 @@ def _lazy_init_context(self) -> None: if self._context is not None: return + root_module = cast(nn.Module, self) context = FsdpContext(device=self._parameter_groups[0].main_weight.device, root_module=self) - for submodule_name, submodule in cast(nn.Module, self).named_modules(): + # named_modules() yields FsdpModules in registration order, which is the static + # forward execution order used to prefetch the next FsdpModule's all-gather. + for submodule_name, submodule in root_module.named_modules(): if not isinstance(submodule, FsdpModule): continue if submodule._context is not None: @@ -158,6 +165,11 @@ def _lazy_init_context(self) -> None: ) submodule._context = context submodule._name = submodule_name + context.forward_order.append(submodule) + + # Backward starts from the root pre-backward hook before visiting child + # subtrees in reverse module order. + _collect_backward_order(root_module, context.backward_order) @property def context(self) -> FsdpContext: @@ -207,22 +219,43 @@ def grad_hook(_parameter: nn.Parameter) -> None: return grad_hook def pre_forward(self) -> None: - """Prepare full parameters for forward compute.""" + """Prepare full parameters for forward compute and prefetch the next FsdpModule. + + While this FsdpModule computes, we issue the next FsdpModule's all-gather + on the comm stream, so ``AG_{i+1}`` is launched before ``F_i`` finishes. + """ self._lazy_init_context() torch.cuda.nvtx.range_push(self._nvtx_label("forward")) self._ready_grad_parameters.clear() + context = self.context + allgather_stream = context.allgather_stream + current_stream = context.current_stream() + if self.is_root(): - allgather_stream = self.context.allgather_stream - allgather_stream.wait_stream(torch.cuda.current_stream(allgather_stream.device)) + allgather_stream.wait_stream(current_stream) + self._unshard_parameter_groups(sync_model_weight=True) + assert self._unshard_event is not None + # Compute waits only for this FsdpModule's all-gather (the prefetch below is + # issued afterwards, so it is free to run concurrently with this FsdpModule). + current_stream.wait_event(self._unshard_event) + + next_module = context.forward_order.next_item(self) + if next_module is not None: + next_module._unshard_parameter_groups(sync_model_weight=True) def _unshard_parameter_groups(self, *, sync_model_weight: bool) -> None: - """Materialize full parameters for this FsdpModule.""" - self.context.drain_delayed_releases(target_length=1) + """Unshard this FsdpModule's parameter groups on the all-gather stream. - allgather_stream = self.context.allgather_stream - current_stream = torch.cuda.current_stream(allgather_stream.device) + If ``_unshard_event`` is already set, this FsdpModule was already + unsharded or prefetched and this method is a no-op. Otherwise, this + method records ``_unshard_event`` after materialization so compute + can wait without depending on later release work. + """ + if self._unshard_event is not None: + return + allgather_stream = self.context.allgather_stream with torch.cuda.stream(allgather_stream): for group in self._parameter_groups: if sync_model_weight: @@ -230,41 +263,81 @@ def _unshard_parameter_groups(self, *, sync_model_weight: bool) -> None: # optimizer post-step hook instead of running it every microbatch. group.sync_model_weight_from_main_weight() group.unshard_parameters() - current_stream.wait_stream(allgather_stream) + self._unshard_event = allgather_stream.record_event() def post_forward(self) -> None: """Return parameters to their sharded resting state after forward compute.""" self._reshard_parameter_groups() - self.context.enqueue_release(self) - if self.is_root(): - self.context.drain_delayed_releases(target_length=0) torch.cuda.nvtx.range_pop() def _reshard_parameter_groups(self) -> None: + """Reshard parameter groups and release unsharded storage after compute. + + This method clears ``_unshard_event`` after queuing the release, so + future users enqueue a fresh all-gather. + """ for group in self._parameter_groups: group.reshard_parameters() + allgather_stream = self.context.allgather_stream + allgather_stream.wait_stream(self.context.current_stream()) + # Release on the all-gather stream where unsharded storage was allocated, + # so no record_stream() call is required for the storage. + with torch.cuda.stream(allgather_stream): + for group in self._parameter_groups: + group.release_unsharded_storage() + self._unshard_event = None + def pre_backward(self) -> None: - """Prepare full parameters for backward compute.""" + """Prepare full parameters and prefetch the next FsdpModule in backward order.""" torch.cuda.nvtx.range_push(self._nvtx_label("backward")) + context = self.context + current_stream = context.current_stream() + if self.is_root(): + context.register_post_backward_final_callback() + # Fork the reduce-scatter stream from the current stream once, at the + # start of backward, so every module's post-backward reduce-scatter is + # part of any active CUDA-graph capture. A stream only joins the + # capture via this wait_stream edge; without it the first allocation on + # the reduce-scatter stream falls back to a raw cudaMalloc, which is + # illegal during capture. Later modules are covered by the post-copy + # fork each preceding module issues before its collective. + context.reduce_scatter_stream.wait_stream(current_stream) + self._unshard_parameter_groups(sync_model_weight=False) + assert self._unshard_event is not None + current_stream.wait_event(self._unshard_event) + + next_module = context.backward_order.next_item(self) + if next_module is not None: + next_module._unshard_parameter_groups(sync_model_weight=False) def post_backward(self) -> None: """Reduce gradients and return parameters to their sharded resting state.""" - for group in self._parameter_groups: - if group.requires_grad: - group.reduce_gradients() + self._reduce_gradient_groups() self._reshard_parameter_groups() - self.context.enqueue_release(self) - if self.is_root(): - self.context.drain_delayed_releases(target_length=0) self._ready_grad_parameters.clear() torch.cuda.nvtx.range_pop() - def release_unsharded_storage(self) -> None: - """Release unsharded storage owned by this FsdpModule.""" + def _reduce_gradient_groups(self) -> None: + """Pack gradients and immediately launch their reduce-scatters.""" + context = self.context + reduce_scatter_stream = context.reduce_scatter_stream + current_stream = context.current_stream() + for group in self._parameter_groups: - group.release_unsharded_storage() + if not group.requires_grad: + continue + + with torch.cuda.stream(reduce_scatter_stream): + partial_grad = group.allocate_partial_grad_buffer() + + current_stream.wait_stream(reduce_scatter_stream) + group.copy_gradients_to_partial_buffer(partial_grad) + + reduce_scatter_stream.wait_stream(current_stream) + with torch.cuda.stream(reduce_scatter_stream): + group.reduce_partial_gradients(partial_grad) @property def parameter_groups(self) -> tuple[FsdpParameterGroup, ...]: @@ -276,6 +349,15 @@ def _nvtx_label(self, phase: Literal["forward", "backward"]) -> str: return f"MFSDP {name} {phase}" +def _collect_backward_order(module: nn.Module, order: IndexedOrder["FsdpModule"]) -> None: + """Collect FsdpModules in static backward prefetch order.""" + if isinstance(module, FsdpModule): + order.append(module) + + for child in reversed(list(module.children())): + _collect_backward_order(child, order) + + def _axis_index(mesh: DeviceMesh, axis: MeshAxis) -> int: if isinstance(axis, int): axis_index = axis diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/parameter_group.py b/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/parameter_group.py index eeec848416b..3e5b7a49392 100644 --- a/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/parameter_group.py +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/parameter_group.py @@ -14,7 +14,6 @@ """Parameter-group runtime state for the minimal Megatron-FSDP path.""" -from collections.abc import Iterable from contextlib import nullcontext import torch @@ -245,11 +244,39 @@ def release_unsharded_storage(self) -> None: # so keep the shared storage-release path. self._unsharded_model_weight.release_storage() - def reduce_gradients(self) -> None: - """Reduce full local gradients into sharded parameter gradients.""" + def allocate_partial_grad_buffer(self) -> DBuffer: + """Allocate the unreduced reduce-scatter input buffer.""" assert self.main_grad is not None - def has_grad(parameters: Iterable[nn.Parameter]) -> bool: + # NCCL symmetric-memory reduce-scatter only selects the symmetric kernel for SUM today. + # Preserve AVG semantics by reducing SUM and scaling the output below. + partial_op = dist.ReduceOp.AVG if self._symm_mem_pool is None else dist.ReduceOp.SUM + grads: list[torch.Tensor] = [] + for name, parameter in zip(self.parameter_names, self.unsharded_parameters, strict=True): + if parameter.grad is None: + raise RuntimeError(f"Missing gradient for FSDP parameter {name!r}.") + grads.append(parameter.grad) + with self._symmetric_memory_context(): + return DBuffer( + mesh=self.mesh, + placements=[Partial(partial_op)] * self.mesh.ndim, + tensor_shapes=tuple(grad.shape for grad in grads), + dtype=grads[0].dtype, + device=grads[0].device, + ) + + def copy_gradients_to_partial_buffer(self, partial_grad: DBuffer) -> None: + """Pack full local gradients into an existing reduce-scatter input buffer.""" + # A future fused-wgrad path can write directly into these buffer views. + for index, parameter in enumerate(self.unsharded_parameters): + partial_grad.get_local_tensor(index).copy_(parameter.grad) + parameter.grad = None + + def reduce_partial_gradients(self, partial_grad: DBuffer) -> None: + """Reduce a packed partial gradient buffer into sharded parameter gradients.""" + assert self.main_grad is not None + + def has_grad(parameters: tuple[nn.Parameter, ...]) -> bool: has_any_grad = False has_any_missing_grad = False for parameter in parameters: @@ -261,20 +288,6 @@ def has_grad(parameters: Iterable[nn.Parameter]) -> bool: raise RuntimeError("FSDP sharded gradients must be either all set or all None.") return has_any_grad - grads: list[torch.Tensor] = [] - for name, parameter in zip(self.parameter_names, self.unsharded_parameters, strict=True): - if parameter.grad is None: - raise RuntimeError(f"Missing gradient for FSDP parameter {name!r}.") - grads.append(parameter.grad) - - # NCCL symmetric-memory reduce-scatter only selects the symmetric kernel for SUM today. - # Preserve AVG semantics by reducing SUM and scaling the output below. - partial_op = dist.ReduceOp.AVG if self._symm_mem_pool is None else dist.ReduceOp.SUM - with self._symmetric_memory_context(): - partial_grad = DBuffer.distribute_tensors( - grads, mesh=self.mesh, placements=[Partial(partial_op)] * self.mesh.ndim - ) - # zero_grad(set_to_none=True) clears sharded parameter grads, so the next # backward can reduce directly into main_grad. zero_grad(set_to_none=False) # leaves sharded grads installed, so this backward accumulates into main_grad. @@ -285,7 +298,8 @@ def has_grad(parameters: Iterable[nn.Parameter]) -> bool: reduce_axis = changed_mesh_axis(partial_grad.placements, self.main_grad.placements) if reduce_axis is None: raise RuntimeError("FSDP gradient reduction requires a changed placement axis.") - grad_divisor = self.mesh.size(reduce_axis) if partial_op == dist.ReduceOp.SUM else 1 + partial_reduce_op = partial_grad.placements[reduce_axis].reduce_op + grad_divisor = self.mesh.size(reduce_axis) if partial_reduce_op == dist.ReduceOp.SUM else 1 if self._symm_mem_pool is not None: partial_grad.rendezvous(reduce_axis) if can_reduce_into_main_grad: @@ -305,9 +319,6 @@ def has_grad(parameters: Iterable[nn.Parameter]) -> bool: for index, parameter in enumerate(self.sharded_parameters): parameter.grad = self.main_grad.get_dtensor(index) - for parameter in self.unsharded_parameters: - parameter.grad = None - def _get_parameter_owner(module: nn.Module, name: str) -> tuple[nn.Module, str]: """Resolve a root-module-relative parameter FQN to its direct owner.""" diff --git a/tests/unit_tests/distributed/mfsdp_v2/test_context.py b/tests/unit_tests/distributed/mfsdp_v2/test_context.py index 6a99a2d19a8..102b2fde332 100644 --- a/tests/unit_tests/distributed/mfsdp_v2/test_context.py +++ b/tests/unit_tests/distributed/mfsdp_v2/test_context.py @@ -42,6 +42,33 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: return x +class BranchModel(nn.Module): + """Nested branch with its own child FsdpModule.""" + + def __init__(self, dim: int) -> None: + super().__init__() + self.bias = nn.Parameter(torch.ones(dim)) + self.inner = nn.Linear(dim, dim, bias=False) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + """Run the nested branch.""" + return torch.relu(self.inner(x) + self.bias) + + +class NestedSiblingModel(nn.Module): + """Model with a nested left subtree and a right sibling.""" + + def __init__(self, dim: int) -> None: + super().__init__() + self.bias = nn.Parameter(torch.ones(dim)) + self.left = BranchModel(dim) + self.right = nn.Linear(dim, dim, bias=False) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + """Run the nested subtree before the right sibling.""" + return self.right(self.left(x) + self.bias) + + def _flat_placements() -> Placements: return Placements(dp_axes=[0], parameter=[Flat()], gradient=[Flat()], optimizer=[Flat()]) @@ -98,3 +125,23 @@ def test_sibling_roots_without_parent_keep_separate_contexts(distributed_setup): assert model.layers[0].context is not model.layers[1].context assert model.layers[0].is_root() assert model.layers[1].is_root() + + +def test_nested_prefetch_orders_use_dfs(distributed_setup): + """Nested FsdpModules should use DFS orders for one-step prefetch.""" + device = distributed_setup.device + + mesh = init_device_mesh(device.type, (distributed_setup.world_size,)) + model = NestedSiblingModel(dim=4).to(device) + + fully_shard(model.left.inner, mesh=mesh, placements=_flat_placements()) + fully_shard(model.left, mesh=mesh, placements=_flat_placements()) + fully_shard(model.right, mesh=mesh, placements=_flat_placements()) + fully_shard(model, mesh=mesh, placements=_flat_placements()) + + with torch.no_grad(): + model(torch.ones(2, 4, device=device)) + + context = model.context + assert list(context.forward_order) == [model, model.left, model.left.inner, model.right] + assert list(context.backward_order) == [model, model.right, model.left, model.left.inner] diff --git a/tests/unit_tests/distributed/mfsdp_v2/test_fully_shard.py b/tests/unit_tests/distributed/mfsdp_v2/test_fully_shard.py index 20686b0e57d..c26a2ab9130 100644 --- a/tests/unit_tests/distributed/mfsdp_v2/test_fully_shard.py +++ b/tests/unit_tests/distributed/mfsdp_v2/test_fully_shard.py @@ -311,8 +311,8 @@ def test_root_backward_returns_to_resting_memory(distributed_setup): ) -def test_overlaps_all_gather_and_compute(distributed_setup): - """A shared root context should let child all-gathers overlap GEMM compute.""" +def test_overlaps_communication_and_compute(distributed_setup): + """Forward and backward communication should overlap GEMM compute.""" world_size = distributed_setup.world_size device = distributed_setup.device if world_size < 2: @@ -352,6 +352,12 @@ def train_one_iteration() -> None: for event in cuda_events if "nccl" in event.name.lower() and "allgather" in event.name.lower() ] + reduce_scatter_events = [ + event + for event in cuda_events + if "nccl" in event.name.lower() + and ("reducescatter" in event.name.lower() or "reduce_scatter" in event.name.lower()) + ] # GEMM device-kernel names vary across CUDA/cuBLAS versions and GPU archs # (e.g. "*gemm*", "cutlass*", "cublas*", and cuBLASLt's Hopper "nvjet_sm90_*"). gemm_events = [ @@ -360,29 +366,38 @@ def train_one_iteration() -> None: if any(token in event.name.lower() for token in ("gemm", "cutlass", "cublas", "nvjet")) ] assert all_gather_events, [event.name for event in cuda_events] + assert reduce_scatter_events, [event.name for event in cuda_events] assert gemm_events, [event.name for event in cuda_events] all_gather_streams = {event.device_resource_id for event in all_gather_events} + reduce_scatter_streams = {event.device_resource_id for event in reduce_scatter_events} gemm_streams = {event.device_resource_id for event in gemm_events} assert len(all_gather_streams) == 1 + assert len(reduce_scatter_streams) == 1 + assert all_gather_streams.isdisjoint(reduce_scatter_streams) assert all_gather_streams.isdisjoint(gemm_streams) + assert reduce_scatter_streams.isdisjoint(gemm_streams) - overlap_count = sum( + all_gather_overlap_count = sum( any(_events_overlap(all_gather_event, gemm_event) for gemm_event in gemm_events) for all_gather_event in all_gather_events ) - # This profiles a full forward/backward iteration, so backward all-gathers are - # included in all_gather_events. The expected overlap count is from the forward - # child pipeline: each child after the first can all-gather while the previous - # child computes, giving num_children - 1 overlaps. Backward does not overlap - # in this all-gather-only path because gradient reduction is not delayed: - # each module synchronously reduces gradients in post_backward before autograd - # reaches the next module's pre_backward all-gather. The next PR addresses - # this by delaying gradient reduction. - expected_overlap_count = num_children - 1 - assert overlap_count >= expected_overlap_count, ( - f"Expected at least {expected_overlap_count} all-gather events to overlap compute, " - f"got {overlap_count}/{len(all_gather_events)}." + reduce_scatter_overlap_count = sum( + any(_events_overlap(reduce_scatter_event, gemm_event) for gemm_event in gemm_events) + for reduce_scatter_event in reduce_scatter_events + ) + # Communication overlaps compute only partially due to SM contention (the + # SM-based NCCL collectives share SMs with the GEMMs) and the count varies + # run to run, so assert only that overlap meaningfully happens rather than the + # theoretical maximum (2*(num_children - 1) / num_children - 1). Symmetric-memory + # collectives (use_symm_mem, ~SM-free) would let these thresholds be tightened. + assert all_gather_overlap_count >= 2, ( + f"Expected all-gather to overlap compute, " + f"got {all_gather_overlap_count}/{len(all_gather_events)}." + ) + assert reduce_scatter_overlap_count >= 1, ( + f"Expected reduce-scatter to overlap compute, " + f"got {reduce_scatter_overlap_count}/{len(reduce_scatter_events)}." )