Fix multi-device mxfp4 dequantization race in _convert_moe_packed_tensors#47423
Conversation
`_convert_moe_packed_tensors` Signed-off-by: kaixuanliu <[email protected]>
Signed-off-by: kaixuanliu <[email protected]>
|
I tried on CUDA device as well. This bug also exists on cuda, and this PR can solve it. |
kaixuanliu
changed the title
_convert_moe_packed_tensors_convert_moe_packed_tensors
CI recapDashboard: View test results in Grafana |
| # With device_map="auto", tensors sitting on a non-current accelerator device are not | ||
| # ordered after their async H2D copy, so the compute below may read garbage and emit | ||
| # out-of-bounds `lut` indices (illegal memory access on CUDA, indexing abort on XPU). | ||
| # Aligning the active device with the tensor's device orders it correctly (no-op on CPU). | ||
| with on_device(blk.device): | ||
| # This vector is only used to index into `lut`, but is hugeee in GPU memory so we delete it immediately | ||
| idx_lo = (blk & 0x0F).to(torch.int) | ||
| sub[:, 0::2] = lut[idx_lo] | ||
| del idx_lo |
There was a problem hiding this comment.
are you sure this is what's happening ? wdym by "tensors .. are not ordered" ? blk and lut are already on the same device at this point, the device context changes where new tensors with no device are created, like a torch.zeros or arange, why would it fix this ?
There was a problem hiding this comment.
Hi, Good Questions. Let me update more background:
Here are facts from my side:
- With device_map="auto" across multiple accelerators, loading gpt-oss-20b aborts inside this loop: on XPU as an index out of bounds in the indexing kernel, on CUDA as an illegal memory access.
- During debugging, it shows that for a blk living on a non-current device (e.g. blk on xpu:1/cuda:1 while the process' current device is device 0), idx_lo = blk & 0x0F transiently produces values far outside [0, 15] (I saw 0–255 then garbage like [-2147483648, 2139160448]), which then indexes lut out of bounds. For tensors on the current device it is always correct.
- Wrapping only this compute in on_device(blk.device) makes the corruption disappear and load+generation succeed, on both XPU and CUDA.
By "not ordered" I mean a cross-device stream ordering gap, not tensor allocation. The loader copies blk H2D on a worker thread using its own device's stream, but the elementwise kernels here are launched while the process' current device is still device 0. Since kernel launch / default-stream selection follows the current device, these kernels aren't guaranteed to be ordered after the H2D copy that fills blk on device 1 — so they can read not-yet-valid memory. Setting the current device to blk.device makes them launch on that device's stream, which is properly ordered after the copy.
|
Sample code to reproduce: |
SunMarc
left a comment
SunMarc
left a comment
There was a problem hiding this comment.
Thanks ! If use_kernels=False, this doesn't happen ?
Yes, if we set |
|
The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update. |
* upstream/main: (39 commits) Remove deprecated training args and `is_fast` property (huggingface#46917) Consistent output shape from `get_image_features` (huggingface#46405) Fix multi-device mxfp4 dequantization race in `_convert_moe_packed_tensors` (huggingface#47423) fix failed test cases for qwen3_omni_moe model (huggingface#47449) Fix Hunyuan-VL PIL image resize parity with reference preprocessing (huggingface#47233) Move `value` padding into the attention interfaces that need it (huggingface#47451) Simplify function dispatch for linear attention (huggingface#47450) [cache] Allow sliding window layers to be roll-backed for speculative decoding (huggingface#47447) Fix double-shifted training loss in GitForCausalLM (huggingface#47395) Fix CohereASR training-loss double-shift (same as Moonshine fix huggingface#46784) (huggingface#46895) Warn when `group_by_length` is silently ignored for iterable datasets (huggingface#47379) Update bug report list (huggingface#46607) Fix shape mismatch in KyutaiSpeechToText `generate()` last window (huggingface#46952) Optimize flash attention max seqlen computation in vision attention (huggingface#47170) fix: remove unreachable return in special token builder (huggingface#47420) Add Harry to slow CI (huggingface#47454) BLT: vectorize patch length processing (huggingface#47385) Fix `TrackioCallback` fails to log evaluation metrics after training ends (huggingface#46935) [Kimi] add integration tests (huggingface#47383) Fix typo in `MusicgenForCausalLM.generate()` (huggingface#46974) ...
4 participants
When loading an mxfp4 model (e.g. openai/gpt-oss-20b) with device_map="auto" across multiple Intel XPUs and use_kernels=True, loading aborts during dequantization:
Root cause:
the multi-threaded weight loader copies each shard to its target device; for a tensor on a non-current XPU device (e.g. xpu:1 while the active device is xpu:0), the subsequent blk & 0x0F compute is not ordered after that copy, reads not-yet-valid memory, and produces out-of-bounds lut indices. Instrumentation confirmed idx_lo on xpu:1 goes 0–15 → 0–255 → garbage [-2147483648, 2139160448], while xpu:0 is always correct.
Fix:
wrap the per-chunk dequant compute in the tensor's own XPU device context with
on_devicecontext manager. Verified: load + generation of gpt-oss-20b across XPUs now succeeds; no assertions.Pls help review, thx! @SunMarc @IlyasMoutawwakil