coffee/DeepInfer
1
0
Fork 0
mirror of https://github.com/deepseek-ai/DeepInfer.git synced 2026-10-31 08:46:21 +08:00
Code Issues Packages Projects Releases Wiki Activity

Neuron: concatenate FP8 checkpoint tensors as bytes (fixes MoE on device)

Browse Source 
torch.cat of float8 tensors (e4m3 and e8m0) returns wrong data on the
neuron device (torch-neuronx 2.14). The HF loader concatenates w1/w3 (and
their UE8M0 scales) into the fused gate_up weights, so on the NeuronCore
every routed and shared expert gate_up scale was zero and the MoE output
vanished (issue #1). Concatenate 1-byte float dtypes through uint8 views,
which is bit-identical on every platform.

tests/neuron/device: create the neuron process group before the runtime
initializes (torch-neuronx requirement) and add the float8 cat checks.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
This commit is contained in:
whn09 2026-10-21 10:25:13 +08:00
parent a7117d3380
commit e10220bed0
2 changed files with 21 additions and 5 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

9
deepinfer/hugging_face/weights.py
View File

@ -242,7 +242,14 @@ class _TensorReader(ExitStack):
def read_concat(self, names: Iterable[str], device: torch.device) -> torch.Tensor:
"""Read one tensor or concatenate multiple tensors along dim 0."""
parts = [self.read(name, device) for name in names]
return parts[0] if len(parts) == 1 else torch.cat(parts, dim=0)
if len(parts) == 1:
return parts[0]
dtype = parts[0].dtype
if dtype.is_floating_point and dtype.itemsize == 1:
# Concatenate FP8 (e4m3 / e8m0) as raw bytes: bit-identical everywhere, and torch.cat of float8
# tensors is wrong on the Neuron device (torch-neuronx 2.14).
return torch.cat([part.view(torch.uint8) for part in parts], dim=0).view(dtype)
return torch.cat(parts, dim=0)
def _local_expert_ids(model: Any, params: dict[str, torch.nn.Parameter]) -> dict[str, tuple[int, ...]]:

17
tests/neuron/device/test_device_primitives.py
View File

@ -38,6 +38,12 @@ def same(dev: torch.Tensor, host: torch.Tensor) -> bool:
def main() -> int:
# torch-neuronx requires the neuron process group to exist before the runtime initializes (first device op).
os.environ.setdefault("MASTER_ADDR", "127.0.0.1")
os.environ.setdefault("MASTER_PORT", "29761")
dist.init_process_group("gloo", rank=0, world_size=1)
g = dist.new_group([0], backend="neuron")
torch.manual_seed(0)
x = torch.randn(11, 512, dtype=torch.bfloat16)
xd = x.to(DEV)
@ -48,6 +54,12 @@ def main() -> int:
check("fp8 -> bf16 upcast on device", "ok", lambda: same(f8.to(DEV).to(torch.bfloat16), f8.to(torch.bfloat16)))
e8 = torch.randint(100, 140, (8, 8), dtype=torch.uint8).view(torch.float8_e8m0fnu)
check("e8m0 H2D bit-exact", "ok", lambda: same(e8.to(DEV), e8))
f8b = torch.randn(4, 8).clamp(-400, 400).to(torch.float8_e4m3fn)
e8b = torch.randint(110, 130, (4, 8), dtype=torch.uint8).view(torch.float8_e8m0fnu)
check("torch.cat of fp8 e4m3 on device", "known", lambda: same(torch.cat([f8.to(DEV)[:4, :8], f8b.to(DEV)]), torch.cat([f8[:4, :8], f8b])))
check("torch.cat of e8m0 on device", "known", lambda: same(torch.cat([e8.to(DEV)[:4], e8b.to(DEV)]), torch.cat([e8[:4], e8b])))
check("torch.cat of fp8 via uint8 views", "ok", lambda: same(
torch.cat([e8.to(DEV)[:4].view(torch.uint8), e8b.to(DEV).view(torch.uint8)]).view(torch.float8_e8m0fnu), torch.cat([e8[:4], e8b])))
b = torch.randint(-128, 127, (4, 16), dtype=torch.int8)
check("int8 nibble unpack", "ok", lambda: same(b.to(DEV) & 0xF, b & 0xF) and same((b.to(DEV) >> 4) & 0xF, (b >> 4) & 0xF))
s = torch.randint(118, 130, (4, 16), dtype=torch.uint8)
@ -100,11 +112,8 @@ def main() -> int:
check("bf16 topk indices == host", "ok", lambda: torch.equal(
torch.topk(xd[0], 64).indices.cpu().sort().values, torch.topk(x[0], 64).indices.sort().values))
# single-rank process group (EP=1 bring-up)
os.environ.setdefault("MASTER_ADDR", "127.0.0.1")
os.environ.setdefault("MASTER_PORT", "29761")
dist.init_process_group("gloo", rank=0, world_size=1)
g = dist.new_group([0], backend="neuron")
check("neuron PG size 1: all_reduce", "ok", lambda: (lambda t: (dist.all_reduce(t, group=g), same(t, x))[1])(xd.clone()))
check("neuron PG size 1: all_gather_into_tensor", "ok",
lambda: (lambda o: (dist.all_gather_into_tensor(o, xd, group=g), same(o, x))[1])(torch.empty_like(xd)))