From e10220bed01f521e997245574a64eac1975b9a0f Mon Sep 17 00:00:00 2001 From: whn09 Date: Wed, 21 Oct 2026 10:25:13 +0800 Subject: [PATCH] Neuron: concatenate FP8 checkpoint tensors as bytes (fixes MoE on device) 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 --- deepinfer/hugging_face/weights.py | 9 ++++++++- tests/neuron/device/test_device_primitives.py | 17 +++++++++++++---- 2 files changed, 21 insertions(+), 5 deletions(-) diff --git a/deepinfer/hugging_face/weights.py b/deepinfer/hugging_face/weights.py index 8b61ade..e440d22 100644 --- a/deepinfer/hugging_face/weights.py +++ b/deepinfer/hugging_face/weights.py @@ -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, ...]]: diff --git a/tests/neuron/device/test_device_primitives.py b/tests/neuron/device/test_device_primitives.py index f24190c..7295b87 100644 --- a/tests/neuron/device/test_device_primitives.py +++ b/tests/neuron/device/test_device_primitives.py @@ -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)))