DeepSpeed

32e301ff - stop allocating per-element temporaries in the overflow check (#8325)

stop allocating per-element temporaries in the overflow check (#8325) Currently there could be quite a surge in memory usage to just check if a tensor has inf or nan. This PR introduces a new general purpose util: `_has_inf_or_nan`, which check if a tensor contains a non-finite value -- and it runs once per gradient on the ZeRO-2 CPU-offload path (`deepspeed/runtime/zero/stage_1_and_2.py:1524`) and per partition on the ZeRO-3 path. Every original implementation of it answered that question elementwise, allocating device memory proportional to the input: - ZeRO-2 built an fp32 copy, three boolean masks, and a second fp32 tensor, all live at once: **11 bytes of device memory for every 2 bytes of a bf16 gradient**. - ZeRO-3 and the dynamic loss scaler built one fp32 copy and inferred overflow from the sum: **4 bytes per element**, plus a false positive whenever a finite tensor's sum exceeds the accumulator range. Maximum and minimum already answer it. `max` propagates NaN and exposes `+inf`, `min` exposes `-inf`, so testing finiteness on those two scalars is equivalent and allocates nothing proportional to the input. **This matters most where a single tensor is large. On an 8B-parameter MoE model the largest gradient is a fused expert weight of 1.27e9 elements, and scanning it allocates 13.02 GiB on top of the gradient itself.** ## One implementation instead of four The predicate existed in four places: one form written three times and a fourth rewritten independently. | location | before | after | | ---------------------------------------------- | ------------------------ | ------------------------------- | | `deepspeed/runtime/utils.py` | sum-based, `(x, i)` | holds the single implementation | | `deepspeed/runtime/fp16/loss_scaler.py:244` | sum-based, `(x)` | delegates | | `deepspeed/runtime/zero/stage3.py:2765` | sum-based, `(x, j=None)` | delegates | | `deepspeed/runtime/zero/stage_1_and_2.py:2487` | elementwise fp32 + masks | delegates | `has_inf_or_nan` in `deepspeed/runtime/utils.py:314` is now the only implementation, and the four call sites delegate to it. The static loss scaler's `_has_inf_or_nan` at `deepspeed/runtime/fp16/loss_scaler.py:184` still returns `False` unconditionally; that is its defined behaviour, not a copy of the predicate. Net effect on the diff is 41 fewer lines. ## Measurements [bench_has_inf_or_nan.py](https://github.com/user-attachments/files/31487233/bench_has_inf_or_nan.py) measures peak device memory allocated beyond the input tensor and checks that every implementation agrees on finite, NaN, `+inf`, and `-inf` inputs. On an H200 with torch 2.11.0+cu130, bf16: | numel | gradient | ZeRO-2 now | ZeRO-3 now | this PR | B/elem ZeRO-2 | B/elem ZeRO-3 | B/elem this PR | | ------: | --------: | ----------: | ---------: | ------: | ------------: | ------------: | -------------: | | 16,384 | 32,768 B | 180,736 B | 66,048 B | 3,072 B | 11.03 | 4.03 | 0.19 | | 65,536 | 131,072 B | 721,408 B | 262,656 B | 3,072 B | 11.01 | 4.01 | 0.05 | | 131,072 | 262,144 B | 1,443,328 B | 525,824 B | 3,072 B | 11.01 | 4.01 | 0.02 | The cost after this change is a constant 3,072 B, independent of tensor size, because only scalars are materialized. End to end, an 8-GPU ZeRO-2 job with optimizer and activation offload, sequence parallelism 8, 16,384 tokens per rank, on a 4-layer Qwen3.5-MoE build: peak allocated per rank falls from **26.643 GiB to 20.185 GiB**, a 6.46 GiB (24%) reduction, with identical losses across three steps. The end-to-end saving is smaller than the 13.02 GiB the scan allocates because removing that peak exposes the next-highest allocation. ## Equivalence `max`/`min` and the elementwise forms return the same answer on every input class the check exists to catch: | input | ZeRO-2 now | ZeRO-3 now | this PR | | ------------- | ---------: | ---------: | ------: | | all finite | False | False | False | | contains NaN | True | True | True | | contains +inf | True | True | True | | contains -inf | True | True | True | Empty tensors are handled explicitly, since `amax` on an empty tensor raises. For ZeRO-3 and the dynamic loss scaler this is also a correctness fix: summation reports overflow for a finite tensor whose total exceeds the accumulator range. The `try`/`except RuntimeError` those forms carried is gone with them, since no scalar conversion of a possibly-infinite value takes place. ## Testing - `bench_has_inf_or_nan.py`: equivalence and memory on H200 / torch 2.11.0+cu130 / bf16, as tabulated above. - End-to-end ZeRO-2 with CPU offload in the configuration described above, comparing peak allocated with and without the change and confirming identical losses across three steps. - All four call sites checked on CUDA and CPU, confirming they agree with each other and that the shared helper introduces no import cycle between `runtime/utils.py`, `runtime/fp16/loss_scaler.py`, and the two ZeRO stages. --------- Signed-off-by: Stas Bekman <stas@stason.org>