diff --git a/.github/workflows/npu_unit_tests.yml b/.github/workflows/npu_unit_tests.yml index 1c76084e0..3b48d93e0 100644 --- a/.github/workflows/npu_unit_tests.yml +++ b/.github/workflows/npu_unit_tests.yml @@ -232,6 +232,9 @@ jobs: uv run --frozen pytest -v -s -x tests/ops/test_comp.py uv run --frozen pytest -v -s -x tests/ops/test_moe_hw_gate.py uv run --frozen pytest -v -s -x tests/ops/test_npu_kernels.py + # A missing native runtime/device must fail Ascend CI, not skip the parity gate. + uv run --frozen python -c "import torch; import torch_npu; assert torch.npu.is_available(), 'Ascend parity gate requires a live NPU'" + VEOMNI_REQUIRE_NPU=1 uv run --frozen pytest -v -s -x tests/ops/test_npu_rms_norm_gated.py # This numerical gate must execute on NPU; fail instead of silently skipping # when the vendor extension is missing from the CI image. uv run --frozen python -c "import cann_ops_transformer" diff --git a/tests/ops/test_npu_kernels.py b/tests/ops/test_npu_kernels.py index 875a3e86f..8046f14a2 100644 --- a/tests/ops/test_npu_kernels.py +++ b/tests/ops/test_npu_kernels.py @@ -328,8 +328,8 @@ def test_partial_matches_eager_bf16(self, B, H, S, D, rotary_dim): class TestNPURmsNormGated: """Tests for the ``rms_norm_gated`` NPU kernel (NPUFusedRMSNormGated).""" - @pytest.mark.parametrize("batch,seq,hidden,ffn_dim", [(2, 16, 128, 256), (1, 8, 64, 128)]) - def test_matches_eager_bf16(self, batch, seq, hidden, ffn_dim): + @pytest.mark.parametrize("batch,seq,hidden", [(2, 16, 128), (1, 8, 64)]) + def test_matches_eager_bf16(self, batch, seq, hidden): slot = OpSlot("rms_norm_gated", "standard") slot.bind("npu") # The bound kernel is the NPUFusedRMSNormGated class; instantiate it. @@ -338,7 +338,7 @@ def test_matches_eager_bf16(self, batch, seq, hidden, ffn_dim): fused_module = fused_cls(hidden_size=hidden, eps=1e-6).to(device=DEVICE, dtype=torch.bfloat16) hidden_states = torch.randn(batch, seq, hidden, device=DEVICE, dtype=torch.bfloat16) - gate = torch.randn(batch, seq, ffn_dim, device=DEVICE, dtype=torch.bfloat16) + gate = torch.randn_like(hidden_states) out_fused = fused_module(hidden_states, gate=gate) out_eager = _eager_rms_norm_gated(hidden_states, fused_module.weight, fused_module.variance_epsilon, gate) diff --git a/tests/ops/test_npu_kernels_extended.py b/tests/ops/test_npu_kernels_extended.py index 8a4076a9b..2cdfdb48e 100644 --- a/tests/ops/test_npu_kernels_extended.py +++ b/tests/ops/test_npu_kernels_extended.py @@ -267,34 +267,27 @@ def test_partial_pass_through_preserved(self): class TestNPURmsNormGatedExtended: """Coverage for the NPU FusedRMSNormGated module beyond PR #818.""" - @pytest.mark.parametrize("ffn_dim_factor", [2, 4]) - def test_gating_output_shape_and_dtype(self, ffn_dim_factor): - """Output shape is (B, S, (ffn_dim + hidden) // 2). - - NPUFusedRMSNormGated concatenates [gate, rms_norm(hidden)] - along the last dim (so the concat length is ffn_dim + - hidden), then npu_swiglu splits in half via SiLU(x[:h]) - * x[h:]. The result is shape (B, S, (ffn_dim + hidden) / 2). - """ + @pytest.mark.parametrize("hidden", [32, 128]) + def test_gating_output_shape_and_dtype(self, hidden): + """Elementwise gating preserves the hidden-state shape and dtype.""" from veomni.ops.kernels.gated_delta_rule.npu_rms_norm_gated import NPUFusedRMSNormGated - batch, seq, hidden, ffn_dim = 1, 4, 32, 32 * ffn_dim_factor - expected_dim = (ffn_dim + hidden) // 2 + batch, seq = 1, 4 fused_module = NPUFusedRMSNormGated(hidden_size=hidden, eps=1e-6).to(device=DEVICE, dtype=torch.bfloat16) hidden_states = torch.randn(batch, seq, hidden, device=DEVICE, dtype=torch.bfloat16) - gate = torch.randn(batch, seq, ffn_dim, device=DEVICE, dtype=torch.bfloat16) + gate = torch.randn(batch, seq, hidden, device=DEVICE, dtype=torch.bfloat16) out = fused_module(hidden_states, gate=gate) - assert out.shape == (batch, seq, expected_dim) + assert out.shape == (batch, seq, hidden) assert out.dtype == torch.bfloat16 def test_gating_zero_gate_zeros_output(self): """When gate is zero, the output is 0 * normed = 0.""" from veomni.ops.kernels.gated_delta_rule.npu_rms_norm_gated import NPUFusedRMSNormGated - batch, seq, hidden, ffn_dim = 1, 4, 32, 64 + batch, seq, hidden = 1, 4, 32 fused_module = NPUFusedRMSNormGated(hidden_size=hidden, eps=1e-6).to(device=DEVICE, dtype=torch.bfloat16) hidden_states = torch.randn(batch, seq, hidden, device=DEVICE, dtype=torch.bfloat16) - gate = torch.zeros(batch, seq, ffn_dim, device=DEVICE, dtype=torch.bfloat16) + gate = torch.zeros(batch, seq, hidden, device=DEVICE, dtype=torch.bfloat16) out = fused_module(hidden_states, gate=gate) # SiLU(0) * normed -> 0, so the output is bit-exact zero. assert out.abs().max().item() == 0.0, f"Output should be exactly zero, got max abs {out.abs().max().item()}" @@ -303,10 +296,10 @@ def test_gating_zero_gate_zeros_output(self): def test_gating_eps_parameter_used(self, eps): from veomni.ops.kernels.gated_delta_rule.npu_rms_norm_gated import NPUFusedRMSNormGated - batch, seq, hidden, ffn_dim = 1, 4, 32, 64 + batch, seq, hidden = 1, 4, 32 fused_module = NPUFusedRMSNormGated(hidden_size=hidden, eps=eps).to(device=DEVICE, dtype=torch.bfloat16) hidden_states = torch.randn(batch, seq, hidden, device=DEVICE, dtype=torch.bfloat16) - gate = torch.randn(batch, seq, ffn_dim, device=DEVICE, dtype=torch.bfloat16) + gate = torch.randn(batch, seq, hidden, device=DEVICE, dtype=torch.bfloat16) out_fused = fused_module(hidden_states, gate=gate) out_eager = _eager_rms_norm_gated(hidden_states, fused_module.weight, eps, gate) # Compound op (RMSNorm + concat + SiLU gate) with bf16 rounding. diff --git a/tests/ops/test_npu_rms_norm_gated.py b/tests/ops/test_npu_rms_norm_gated.py new file mode 100644 index 000000000..92ed99024 --- /dev/null +++ b/tests/ops/test_npu_rms_norm_gated.py @@ -0,0 +1,133 @@ +"""CPU contract checks and native Ascend forward/backward parity.""" + +import importlib.util +import os +import sys +import types +from pathlib import Path +from unittest.mock import patch + +import pytest +import torch +import torch.nn.functional as F + + +def load_wrapper(): + def rms(x, weight, eps): + normalized = x.float() * torch.rsqrt(x.float().square().mean(-1, keepdim=True) + eps) + return (normalized.to(x.dtype) * weight,) + + def swiglu(x, dim): + gate, value = x.chunk(2, dim=dim) + return F.silu(gate) * value + + fake = types.SimpleNamespace(npu_rms_norm=rms, npu_swiglu=swiglu) + path = Path(__file__).parents[2] / "veomni/ops/kernels/gated_delta_rule/npu_rms_norm_gated.py" + spec = importlib.util.spec_from_file_location("_gated_norm_test", path) + module = importlib.util.module_from_spec(spec) + with patch.dict(sys.modules, {"torch_npu": fake}): + spec.loader.exec_module(module) + return module.NPUFusedRMSNormGated + + +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float32]) +def test_sigmoid_reference_and_gradients(dtype): + torch.manual_seed(73) + module = load_wrapper()(128, dtype=dtype, activation="sigmoid") + with torch.no_grad(): + module.weight.copy_(torch.linspace(0.25, 1.75, 128).to(dtype)) + x = torch.randn(7, 128, dtype=dtype).requires_grad_() + z = torch.linspace(-20, 20, 896).reshape(7, 128).to(dtype).requires_grad_() + r = x.float() * torch.rsqrt(x.float().square().mean(-1, keepdim=True) + module.variance_epsilon) + expected = (module.weight * r.to(dtype) * z.float().sigmoid()).to(dtype) + actual = module(x, z) + torch.testing.assert_close(actual, expected, atol=0, rtol=0) + dy = torch.randn_like(actual) + got = torch.autograd.grad(actual, (x, z, module.weight), dy) + ref = torch.autograd.grad(expected, (x, z, module.weight), dy) + for a, b in zip(got, ref, strict=True): + torch.testing.assert_close(a, b, atol=0, rtol=0) + + +def test_zero_gate_is_half_normalized_and_silu_unchanged(): + cls = load_wrapper() + x = torch.ones(2, 128) + z = torch.zeros_like(x) + torch.testing.assert_close(cls(128, activation="sigmoid")(x, z), x * 0.5 / (1 + 1e-6) ** 0.5) + torch.testing.assert_close(cls(128)(x, z), torch.zeros_like(x)) + with pytest.raises(ValueError, match="activation"): + cls(128, activation="unsupported_gate") + with pytest.raises(ValueError, match="gate"): + cls(128, activation="sigmoid")(x) + + +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float32]) +@pytest.mark.parametrize("activation", ["sigmoid", "silu"]) +@pytest.mark.parametrize("autocast_enabled", [False, True]) +def test_native_npu_reference_and_gradients(dtype, activation, autocast_enabled): + # Load the real wrapper without the CPU tests' torch_npu substitution. + if os.environ.get("VEOMNI_REQUIRE_NPU") == "1": + import torch_npu # noqa: F401 + + assert torch.npu.is_available(), "Ascend parity gate requires a live NPU" + else: + pytest.importorskip("torch_npu") + if not torch.npu.is_available(): + pytest.skip("Native Ascend device is required") + path = Path(__file__).parents[2] / "veomni/ops/kernels/gated_delta_rule/npu_rms_norm_gated.py" + spec = importlib.util.spec_from_file_location("_native_gated_norm_test", path) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + torch.manual_seed(73) + norm = module.NPUFusedRMSNormGated(128, device="npu", dtype=dtype, activation=activation) + with torch.no_grad(): + norm.weight.copy_(torch.linspace(0.25, 1.75, 128).to(device="npu", dtype=dtype)) + x = torch.randn(2, 7, 128, device="npu", dtype=dtype).requires_grad_() + z = torch.linspace(-20, 20, x.numel()).reshape(x.shape).to(device="npu", dtype=dtype).requires_grad_() + rx, rz, rw = (tensor.detach().clone().requires_grad_() for tensor in (x, z, norm.weight)) + # An independent torch reference preserves the architecture's cast boundary. + with torch.autocast("npu", enabled=False): + normalized = rx.float() * torch.rsqrt(rx.float().square().mean(-1, keepdim=True) + norm.variance_epsilon) + if activation == "sigmoid": + expected = (rw * normalized.to(dtype) * rz.float().sigmoid()).to(dtype) + else: + expected = F.silu(rz) * (normalized * rw.float()).to(dtype) + with torch.autocast("npu", dtype=torch.bfloat16, enabled=autocast_enabled): + actual = norm(x, z) + assert actual.dtype == dtype + dy = torch.randn_like(actual) + got = torch.autograd.grad(actual, (x, z, norm.weight), dy) + ref = torch.autograd.grad(expected, (rx, rz, rw), dy) + # Fused CANN reductions need tolerance; BF16 roundoff is not bitwise FP32. + rtol, atol = (3e-2, 2e-2) if dtype == torch.bfloat16 else (3e-5, 3e-5) + for a, b in zip((actual, *got), (expected, *ref), strict=True): + assert torch.isfinite(a).all() + torch.testing.assert_close(a, b, atol=atol, rtol=rtol) + + +@pytest.mark.parametrize("activation", ["sigmoid", "silu"]) +def test_gate_must_match_input_shape(activation): + norm = load_wrapper()(32, activation=activation) + x = torch.ones(1, 4, 32) + for gate in [None, torch.zeros(1, 4, 64), torch.zeros(1, 1, 32)]: + with pytest.raises(ValueError, match="gate with the input shape"): + norm(x, gate) + + +@pytest.mark.parametrize("activation", ["relu", "gelu", "tanh", "swish"]) +def test_upstream_activation_fallback_is_preserved(activation): + from transformers.activations import ACT2FN + + torch.manual_seed(94) + norm = load_wrapper()(32, activation=activation) + x = torch.randn(2, 32, requires_grad=True) + z = torch.randn_like(x, requires_grad=True) + normalized = x * torch.rsqrt(x.square().mean(-1, keepdim=True) + norm.variance_epsilon) + expected = normalized * norm.weight * ACT2FN[activation](z) + actual = norm(x, z) + torch.testing.assert_close(actual, expected) + for got, ref in zip( + torch.autograd.grad(actual.sum(), (x, z, norm.weight)), + torch.autograd.grad(expected.sum(), (x, z, norm.weight)), + ): + torch.testing.assert_close(got, ref) diff --git a/veomni/ops/kernels/gated_delta_rule/npu_rms_norm_gated.py b/veomni/ops/kernels/gated_delta_rule/npu_rms_norm_gated.py index 5c38b1141..be7514453 100644 --- a/veomni/ops/kernels/gated_delta_rule/npu_rms_norm_gated.py +++ b/veomni/ops/kernels/gated_delta_rule/npu_rms_norm_gated.py @@ -24,29 +24,30 @@ from transformers.activations import ACT2FN -# ``npu_swiglu`` hard-codes a SiLU gate, so it can only stand in for the gate -# activation when the model asks for SiLU. Qwen3.5 does (``hidden_act``), but -# Qwen4-Exp sets ``output_gate_type="sigmoid"`` — fusing there would silently -# compute the wrong thing. -_FUSED_GATE_ACTIVATIONS = frozenset({"silu", "swish"}) - - class NPUFusedRMSNormGated(nn.Module): - """Ascend NPU RMSNorm with an activation-gated output projection.""" - - def __init__(self, hidden_size, eps=1e-6, activation: str = "silu", **kwargs): - """Initialize the norm weight and gate activation.""" + def __init__(self, hidden_size, eps=1e-6, activation="silu", device=None, dtype=None, **kwargs): super().__init__() - self.weight = nn.Parameter(torch.ones(hidden_size)) + if activation not in ACT2FN: + raise ValueError(f"Unsupported NPU gated RMSNorm activation: {activation!r}") + self.weight = nn.Parameter(torch.ones(hidden_size, device=device, dtype=dtype)) self.variance_epsilon = eps self.activation = activation def forward(self, hidden_states, gate=None): - """Apply RMSNorm to ``hidden_states`` and multiply by the activated ``gate``.""" - normed = torch_npu.npu_rms_norm(hidden_states, self.weight, self.variance_epsilon)[0] - if self.activation in _FUSED_GATE_ACTIVATIONS: - return torch_npu.npu_swiglu(torch.cat([gate, normed], dim=-1), dim=-1) - - # Still fuses the norm; only the gate falls back to the eager formula - # (float32 activation, then cast back) that the reference modules use. - return (normed * ACT2FN[self.activation](gate.to(torch.float32))).to(hidden_states.dtype) + if gate is None or gate.shape != hidden_states.shape: + raise ValueError("NPU gated RMSNorm requires a gate with the input shape") + if self.activation == "sigmoid": + # Match Qwen4-Exp's rounding boundary: normalize, cast to the input + # dtype, then apply the learned weight. Fusing the learned weight + # into RMSNorm would move that boundary for BF16 compute weights. + unit_weight = torch.ones(hidden_states.shape[-1], device=hidden_states.device, dtype=hidden_states.dtype) + normalized = torch_npu.npu_rms_norm(hidden_states, unit_weight, self.variance_epsilon)[0] + weighted = self.weight * normalized.to(hidden_states.dtype) + return (weighted * gate.float().sigmoid()).to(hidden_states.dtype) + # Preserve the existing SiLU route for models such as Qwen3.6. + normalized = torch_npu.npu_rms_norm(hidden_states, self.weight, self.variance_epsilon)[0] + if self.activation in {"silu", "swish"}: + return torch_npu.npu_swiglu(torch.cat([gate, normalized], dim=-1), dim=-1) + # Preserve upstream's generic activation fallback without changing the + # Qwen4 sigmoid cast boundary or the established SiLU fused path. + return (normalized * ACT2FN[self.activation](gate.float())).to(hidden_states.dtype)