Repository navigation
fix(moe): refuse e8m0 scale code 0xFF when MXFP4 expert pieces are read - #656
Avicennasis wants to merge 1 commit into
Conversation
e8m0 block scales have one code with no finite value: 0xFF is NaN by the OCP MX spec, every other code is 2**(code-127). The fused ds_fp4 and mxfp4 expert kernels fold the scale straight into the GEMM, so one 0xFF byte makes its 32-wide block NaN and the NaN reaches the logits with nothing in the log. Count 0xFF bytes in every e8m0 scale tensor as the DeepSeek-V4 and gpt-oss readers materialise it (serial and parallel), one pass per tensor, and raise ValueError naming the tensor, the count and the shard; the engine's weight load context reports it as WeightLoadError. Clean tensors pass through unchanged.
176d38c to
0a2e1e8
Compare
|
Rebased onto current Re-ran on a CPU-only Ubuntu 24.04 x86_64 box (AMD Ryzen 7 5800X, 16 threads, 125 GiB, no GPU; Python 3.12.3, torch 2.14.1+cpu, triton 3.8.0, transformers 5.18.0, flashlib 0.4.0, numpy 2.5.2, safetensors 0.8.0, pytest 9.1.1;
The scan still reads each e8m0 scale tensor once, as the DeepSeek-V4 and gpt-oss readers materialise it (serial and parallel paths), and raises |
Fixes #655. One commit on
main@2d8b9cc; no kernel, no model family, no GPU involved.What this is
e8m0 has one code with no finite value:
0xFFis NaN by the OCP MX spec, every other code is2**(code-127). Theds_fp4and gpt-oss expert kernels fold that straight into the GEMM, so one0xFFbyte in a scale tensor makes its 32-wide block NaN, and the NaN reaches the logits with nothing in the log naming the file. Neither MXFP4 reader looks for it today.Change
moe/expert_pieces.py:check_e8m0_scales(name, tensor, file)counts the bytes equal to0xFFin one e8m0 scale tensor (torch.count_nonzero(tensor.view(torch.uint8) == 0xFF), one pass,U8orF8_E8M0storage alike, strided slices included) and raisesValueErrornaming the tensor, the count and the shard:e8m0 scale tensor 'layers.0.ffn.experts.1.w2.scale' in 'model-00001-of-00002.safetensors' carries 1 NaN scale code(s) (0xFF) out of 64: the checkpoint is corrupt. A clean tensor is returned unchanged.ValueErroris what the sibling checks in this module raise (Incomplete expert tensors in checkpoint), and the engine's_weight_load_contextreports it asWeightLoadErroron the log and/healthlike any other unreadable checkpoint.models/deepseek_v4/weight.py::iter_expert_pieces: every*.scaletensor passes through the check as it is read, beforeper_expert_piecesgroups it, on the serial and the parallel reader (the shard name comes from the index's weight map, which the serial reader already opens).models/gpt_oss/weight.py::iter_expert_pieces: every*_scalestensor passes through the check as it is read, serial (the open shard) and parallel (this rank's intermediate slice; shard name fromsafetensors_weight_map).Refusal rather than zeroing: a NaN scale is never a valid value, and the published checkpoints do not carry it (below), so a
0xFFis a corrupt file and the operator should see which one. Not covered: an FTW checkpoint converted before this change keeps whatever the raw checkpoint had, since the banks are read back without the raw reader; NVFP4 experts carry fp8 (e4m3) block scales and are not e8m0, so they are not scanned.How tested
CPU-only box (Ubuntu 24.04, AMD Ryzen 7 5800X, no GPU), Python 3.12.3, torch 2.11.0+cu130, safetensors 0.8.0, flashlib 0.3.0, pytest 9.1.1,
PYTHONPATH=python, C++ extension not built, no flashinfer. The readers stop at the piece, so nothing here needs a kernel.tests/moe/test_e8m0_scale_check.py, 12 cases:0xFFbyte naming the tensor, the count (1 ... out of 32) and the file, asU8and asF8_E8M0, and counts two; every finite code0x00..0xFEpasses and the same object comes back, including a strided view;ds_fp4synthetic checkpoint (two shards, one MoE layer of two experts plus an MTP layer the reader skips,F8_E8M0scales asdeepseek-ai/DeepSeek-V4-Flashstores them): one0xFFinlayers.0.ffn.experts.1.w2.scaleis refused by that name and its shard, serial and parallel; the clean checkpoint yields both expert pieces with the six roles, andtorch.count_nonzerois called exactly six times, once per scale tensor and each over exactly that tensor's bytes (the weights and the MTP layer are never scanned);U8stacked scales asopenai/gpt-oss-20bstores them): one0xFFingate_up_proj_scalesand two indown_proj_scalesare refused by name, shard and count, serial and parallel; the clean checkpoint yields the layer piece.The same file against unmodified
main@2d8b9cc: 10 failed, 2 passed (the two clean gpt-oss cases pass on both; the refusals and the per-tensor count fail becausemaindoes not scan).Full suite on this box:
pytest -q testson this branch 8 failed / 1528 passed / 635 skipped (2153 collected, 9:25); unmodifiedmain@2d8b9ccin the same venv 8 failed / 1516 passed / 635 skipped (2153 collected, 9:20). Failure sets identical, all environment-bound (twotest_cache_budgetand onetest_offloadcase that need flashinfer, the swiglu-clamp anddeepseek_v41/test_engramcases that need the C++ extensions, threetest_quant_configcases that probe this box's HF cache); +12 is exactly the new file.Real checkpoints, read over HTTP range requests of the scale tensors alone (no full download): all 48 expert-scale tensors of
openai/gpt-oss-20b(597,196,800 bytes) and 9 of the 46 shards (6,939 of the 33,924 expert-scale tensors, 1,811,953,152 bytes) ofdeepseek-ai/DeepSeek-V4-Flashcarry zero0xFFbytes, so the refusal does not touch either published checkpoint.git diff --checkclean, ruff 0.17.0 on the touched files: no new findings againstmain.Not run on a GPU: the change does not reach a kernel; it ends at the piece the banks are packed from.