DeepSpeed

10e358b8 - Make wait() idempotent on AllGatherHandle and NoGatherHandle (#8487)

Make wait() idempotent on AllGatherHandle and NoGatherHandle (#8487) # Make `wait()` idempotent on `AllGatherHandle` and `NoGatherHandle` (ZeRO-3) Fix https://github.com/deepspeedai/DeepSpeed/issues/8486 ## The problem `PartitionedParameterCoordinator.__all_gather_params_` registers **one** handle object under **every** parameter of a gather group: ```python handle = param_group[0].all_gather_coalesced(param_group) for param in param_group: self.__inflight_param_registry[param] = handle # same object, many keys ``` `_fetch_sub_module_impl` pops and waits those keys independently, one per module, so a group's handle is waited once per parameter in the group. Every multi-parameter handle class already expects that and returns early on a completion flag. The two single-parameter classes do not: | class | idempotent before this PR | |---|---| | `AllGatherCoalescedHandle` | yes | | `NoGatherCoalescedHandle` | yes | | `AllReduceCoalescedHandle` | yes | | `AllGatherHandle` | **no** | | `NoGatherHandle` | **no** | If the parameter is released between two of those waits, the second wait runs the body again and its last statement sets `ds_status = AVAILABLE` over storage that `free_param` already replaced with `torch.empty(0)`. Fetching is driven purely by that flag, so every later fetch skips the all-gather and the next matmul receives an empty weight: ``` RuntimeError: size mismatch, got input (2048), mat (2048x4096), vec (0) deepspeed/runtime/zero/linear.py:92 in forward ``` `AllGatherHandle` cannot repair itself on the common path even in principle: `_all_gather_sequential` points `param.data` at the all-gather output buffer at *launch* time and does not pass that buffer to the handle, so `wait()` has nothing to rebuild from. `AllGatherCoalescedHandle` keeps its slices in `self.partitions` and rebuilds `param.data` on every wait, which is why it was never affected. ## The change Ten lines. `AllGatherHandle` and `NoGatherHandle` each record completion and return early on a repeat wait, matching what the three multi-parameter classes already do. ```python def wait(self, handle_dependency=True) -> None: + if self.__complete: + return + instrument_w_nvtx(self.__handle.wait)() ... self.__param.ds_status = ZeroParamStatus.AVAILABLE + self.__complete = True ``` This is a no-op for any program that waits each handle once. A handle covers exactly one launched gather; once that gather is waited there is nothing left for a second wait to do. ## Why a completion flag rather than a storage check An alternative is to return early when the parameter is not `INFLIGHT` and its storage is empty. That was tried first and it does fix the originally reported crash, but it is narrower in three ways: 1. It has to be repeated in every affected class. Applied to `AllGatherHandle` alone it leaves `NoGatherHandle` broken, which is the class an expert-parallel placement reaches. 2. It allows the second wait to run its body, so it depends on the corrupted state being recognisable from the outside. The completion flag stops the second wait from happening at all. 3. It does not restore the class contract. Making `wait()` idempotent everywhere is the invariant the rest of the code already assumes. ## Which configurations were affected Two independent routes lead to a single-parameter handle inside a multi-parameter registration. Both are reproduced below. **`stage3_allgather_sequential: true`.** The flag makes `all_gather_coalesced` take the sequential path for every group, so each parameter gets its own `AllGatherHandle`. No MoE and no expert parallelism are involved. This affects plain dense models. **A gather group spanning two process groups.** `all_gather_coalesced` splits by `id(get_partition_dp_group(param))` and recurses, so a parameter alone in its bucket takes the sequential path. DeepSpeed's AutoEP produces this at stock settings, because expert parameters are partitioned over `_get_expert_data_parallel_group(...)` and everything else over the data-parallel group. Those groups differ for any `ep_size > 1`. Parameters at or below `stage3_param_persistence_threshold` are never released, so the window never opens for them. The exposure is parameters above the threshold. ## Tests Added `tests/unit/runtime/zero/test_zero_stale_handle_wait.py`, two tests on 2 ranks, one per affected class. Each drives the coordinator's registration pattern by hand: build the shared handle, wait it, release the victim, wait again, and assert the victim is still `NOT_AVAILABLE` with zero elements. ``` cd tests && python -m pytest -q unit/runtime/zero/test_zero_stale_handle_wait.py master: 2 failed AssertionError: a second wait() on a released parameter marked it AVAILABLE this PR: 2 passed ``` Two standalone handle-level reproductions were also run on 2 GPUs against unmodified `master` and against this change. Victims are ordinary 4096x4096 `nn.Linear` weights unless stated. | case | handle built | master | this PR | |---|---|---|---| | `stage3_allgather_sequential` on, one process group, 2 weights | `MultipleAllGatherHandles[AllGatherHandle, AllGatherHandle]` | corrupt | ok | | sequential on, victim shares its bucket with 2 others | three `AllGatherHandle` | corrupt | ok | | two process groups, ordinary weight alone in its bucket | `MultipleAllGatherHandles[MultipleAllGatherHandles[AllGatherHandle], NoGatherHandle]` | corrupt | ok | | two process groups, victim is the size-1-group parameter | same, victim on the `NoGatherHandle` side | corrupt | ok | | control: 2 weights, one process group | `AllGatherCoalescedHandle` | ok | ok | | control: 3 weights, one process group | `AllGatherCoalescedHandle` | ok | ok | | control: mixed bf16 and fp32, one process group | one `AllGatherCoalescedHandle` per dtype | ok | ok | "corrupt" means the victim ended `AVAILABLE` with `data.numel() == 0` after the second wait. The mixed-dtype control matters: the per-dtype split inside `_all_gather_coalesced` still produces coalesced handles, so it is not a third route into this. ## Origin Found while running ZeRO-3 with AutoEP on an 8xH200 node, where a router gate weight above the persistence threshold was prefetched together with an expert tensor from a different process group. The crash appeared in the backward pass inside activation-checkpoint recompute, but activation checkpointing is not required: it only affects whether the parameter is released between the two waits. Signed-off-by: pengdurice <pengduhit@gmail.com>