diff --git a/deepinfer/model_executor/models/v4_modules/mhc_ops.py b/deepinfer/model_executor/models/v4_modules/mhc_ops.py index b4ab925..e68c79c 100644 --- a/deepinfer/model_executor/models/v4_modules/mhc_ops.py +++ b/deepinfer/model_executor/models/v4_modules/mhc_ops.py @@ -39,13 +39,15 @@ def mhc_post( post_layer_mix: torch.Tensor, comb_res_mix: torch.Tensor, ) -> torch.Tensor: - if nki_hc.active(x): + if nki_hc.active(x, "post"): return nki_hc.post(x, residual, post_layer_mix, comb_res_mix) + if nki_hc.active(x): # post routed to torch (DEEPINFER_NEURON_HC_TORCH): the kernel's operation order + return nki_hc.post_torch(x, residual, post_layer_mix, comb_res_mix) return mhc_post_fwd(x, residual, post_layer_mix, comb_res_mix) def mhc_pre_apply_mix(x: torch.Tensor, mix: torch.Tensor) -> torch.Tensor: - if nki_hc.active(x): + if nki_hc.active(x, "collapse"): return nki_hc.pre_apply_mix(x, mix) h = x.shape[-1] mhc = mix.shape[-2] diff --git a/deepinfer/neuron_ops/nki_hc.py b/deepinfer/neuron_ops/nki_hc.py index 527900e..3344a4c 100644 --- a/deepinfer/neuron_ops/nki_hc.py +++ b/deepinfer/neuron_ops/nki_hc.py @@ -30,8 +30,15 @@ import torch ENABLED = os.environ.get("DEEPINFER_NEURON_NKI_HC", "1") == "1" -def active(t: torch.Tensor) -> bool: - return ENABLED and t.device.type == "neuron" +def active(t: torch.Tensor, op: str | None = None) -> bool: + return ENABLED and t.device.type == "neuron" and op not in _TORCH_OPS + + +#: DEEPINFER_NEURON_HC_TORCH: mHC ops that take torch even with the NKI kernels enabled. collapse (the default) is +#: bit-exact and compiles to a faster graph than its kernel. post is opt-in: post_torch is bit-equal to the kernel +#: standalone, but fused into the layer graph it rounds differently (tiny golden 5/9 on the long prompt). The mixes +#: keep the kernel path (the shim's fp32 fn GEMM runs on bf16-narrowed operands in a graph, see above). +_TORCH_OPS = set(filter(None, os.environ.get("DEEPINFER_NEURON_HC_TORCH", "collapse").split(","))) _KERNELS: dict = {} @@ -80,6 +87,20 @@ def post(x: torch.Tensor, residual: torch.Tensor, post_layer_mix: torch.Tensor, ) +def post_torch(x: torch.Tensor, residual: torch.Tensor, post_layer_mix: torch.Tensor, comb_res_mix: torch.Tensor) -> torch.Tensor: + """``post`` in torch, in ``hc_attn_out``'s fp32 operation order (``comb`` terms first, ``post * x`` last), so it + compiles to plain elementwise work instead of the kernel (DEEPINFER_NEURON_HC_TORCH=post).""" + b, n, hc, h = residual.shape + res = residual.float() + xf = x.reshape(b, n, 1, h).float() + comb = comb_res_mix.reshape(b, n, hc, hc) + acc = res[:, :, 0:1, :] * comb[:, :, 0, :].unsqueeze(-1) # [b, n, j, h]: comb[0, j] * res[0] + for i in range(1, hc): + acc = res[:, :, i : i + 1, :] * comb[:, :, i, :].unsqueeze(-1) + acc + out = xf * post_layer_mix.reshape(b, n, hc, 1) + acc + return out.to(residual.dtype) + + def split_fn(fn: torch.Tensor) -> None: """Attach ``fn``'s bf16 ``hi`` / ``lo`` halves (``hi + lo == fn`` to ~2**-17); eager only, see module doc.""" hi = fn.detach().to(torch.bfloat16) @@ -150,4 +171,8 @@ def pre_big_fuse( ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: """``mhc_pre_big_fuse``: ``(post [..., hc, 1], comb [..., hc, hc], layer_input [..., h])``.""" pre, post_mix, comb = mixes(residual, fn, scale, base) + if "collapse" in _TORCH_OPS: + from deepinfer.model_executor.models.v4_modules.mhc_ops import mhc_pre_apply_mix + + return post_mix, comb, mhc_pre_apply_mix(residual, pre) return post_mix, comb, pre_apply_mix(residual, pre) diff --git a/deepinfer/utils.py b/deepinfer/utils.py index 3db7e9e..62facdf 100644 --- a/deepinfer/utils.py +++ b/deepinfer/utils.py @@ -141,7 +141,7 @@ def allow_cuda_sync() -> AbstractContextManager[None]: def _neuron_bind_cpus(local_rank: int) -> list[int] | None: - """CPUs for the rank on logical NeuronCore ``local_rank`` (of CUDA_VISIBLE_DEVICES), on its device's NUMA node. + """CPUs for the rank on logical NeuronCore ``local_rank`` (the physical core id), on its device's NUMA node. The generic split (physical cores // ranks) gave each of 64 ranks one logical CPU, numbered by rank: the step's dispatch thread and the runtime's threads shared it, and devices 0-3 (NUMA 1 on trn2.48xlarge) got CPUs of NUMA 0. @@ -158,8 +158,7 @@ def _neuron_bind_cpus(local_rank: int) -> list[int] | None: except (OSError, subprocess.SubprocessError, ValueError): logger.warning("neuron-ls unavailable: generic CPU binding") return None - visible = os.environ.get("CUDA_VISIBLE_DEVICES") - core = int(visible.split(",")[local_rank]) if visible else local_rank + core = local_rank # bind_cpu gets the worker's gpu_id: already the physical core (CUDA_VISIBLE_DEVICES entry) device = next((d for d in devices if core in d["neuroncore_ids"]), None) if device is None: return None