DeepSpeed
18f92828 - Guard the shm inference all-reduce on the host ISA (#8701)

Commit
3 days ago
Guard the shm inference all-reduce on the host ISA (#8701) ## Problem The shm comm op's bf16/fp16 reduce kernels are compiled with `target("avx512bw")` and use `_mm512_*` conversions; the fp32 kernels use 256-bit AVX. Executing them on a host without the ISA faults, and because the reduce synchronizes through shm flags, the fault surfaces downstream as a **permanent hang of the peers instead of an error** — on the multi-rank CPU CI this froze an xdist worker until the phase timeout. ## Change - Export `vector_kernel_level()` from the op: `2` = avx512bw, `1` = avx2, `0` = neither; non-x86 hosts report `2` (their kernels are JIT-built for the host ISA). - Gate the dispatch in `inference_all_reduce`: fall back to the process-group `all_reduce` when the tensor dtype needs a level the host lacks (fp32 requires level 1, bf16/fp16 require level 2). ## Validation - **avx2-only host**: fp32 keeps the shm path, bf16/fp16 now fall back; all dtype parametrizations of `TestDistInferenceAllReduce` pass (previously segfaulted/hung), and a standalone 4-rank repro that crashed before returns correct sums. - **AVX512 host**: `vector_kernel_level()` reports 2 and every dtype still dispatches to the original shm kernels — the guard is a pure init-time probe with no change to the kernel code, so behavior on AVX512 machines is unchanged. `inference_all_reduce` results remain numerically equivalent to the process-group `all_reduce` across fp32/bf16/fp16 and small through large message sizes (crossing the 1MB distributed-path threshold), and the shm latency advantage over gloo is preserved (microsecond-scale small messages). --------- Signed-off-by: Guokai Ma <guokai.ma@intel.com>
Author
Parents
Loading