From 6e2ceae93bcd5434983147841c029fffaf660c1c Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Tue, 14 Jul 2026 14:49:24 -0700 Subject: [PATCH] Update [ghstack-poisoned] --- backends/webgpu/test/ops/test_fused_ce.py | 130 ++++++++++++++++++++++ 1 file changed, 130 insertions(+) create mode 100644 backends/webgpu/test/ops/test_fused_ce.py diff --git a/backends/webgpu/test/ops/test_fused_ce.py b/backends/webgpu/test/ops/test_fused_ce.py new file mode 100644 index 00000000000..631a15c379c --- /dev/null +++ b/backends/webgpu/test/ops/test_fused_ce.py @@ -0,0 +1,130 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +"""Fused cross-entropy training op (`et_vk.fused_ce`) export + fp64 golden. + +`fused_ce(logits[M,V], labels[M], n_valid) -> (loss, dlogits[M,V])` computes the +mean-over-valid CE loss and its gradient in one op (labels < 0 are ignored/pad). +Golden is the fp64 reference (`logsumexp - picked`, `softmax - onehot`), the same math +torch's cross_entropy uses; the native test reconstructs the deterministic inputs. +""" + +import os +import unittest +from dataclasses import dataclass + +import numpy as np +import torch + +from executorch.backends.vulkan import VulkanPartitioner +from executorch.exir import to_edge_transform_and_lower + + +@dataclass(frozen=True) +class CeConfig: + name: str + m: int # rows (valid + pad positions) + v: int # vocab + n_pad: int = 0 # trailing rows with label = -1 (ignored) + + +# Mirrored by the C++ kFusedCeConfigs table. Llama-3.2-1B vocab = 128256. +CONFIGS = [ + CeConfig("tiny", 4, 32), + CeConfig("masked", 8, 128, n_pad=3), # some ignored labels + CeConfig("llama_vocab", 16, 128256), # real vocab width +] + + +def _inputs(cfg: CeConfig): + """Deterministic logits [M,V] + labels [M] (last n_pad = -1); reconstructable in C++.""" + flat = np.arange(cfg.m * cfg.v, dtype=np.int64) + logits = torch.from_numpy( + (((flat % 23) - 11).astype(np.float32) / np.float32(8.0)).reshape(cfg.m, cfg.v) + ) + labels = torch.from_numpy((np.arange(cfg.m, dtype=np.int64) * 7 + 3) % cfg.v) + if cfg.n_pad: + labels[cfg.m - cfg.n_pad :] = -1 + n_valid = float(max(1, cfg.m - cfg.n_pad)) + return logits, labels, n_valid + + +def _fp64_golden(logits: torch.Tensor, labels: torch.Tensor, n_valid: float): + mask = labels >= 0 + safe = labels.clamp(min=0).long() + lg = logits.double() + lse = torch.logsumexp(lg, dim=-1) + picked = lg.gather(-1, safe[:, None]).squeeze(-1) + loss = torch.where(mask, (lse - picked) / n_valid, torch.zeros_like(lse)).sum() + softmax = torch.softmax(lg, dim=-1) + onehot = torch.nn.functional.one_hot(safe, logits.shape[-1]).double() + dlogits = torch.where( + mask[:, None], (softmax - onehot) / n_valid, torch.zeros_like(softmax) + ) + return loss.to(torch.float32), dlogits.to(torch.float32) + + +class _CeModule(torch.nn.Module): + def forward(self, logits, labels, n_valid): + return torch.ops.et_vk.fused_ce(logits, labels, n_valid) + + +def _export(logits, labels, n_valid): + ep = torch.export.export(_CeModule(), (logits, labels, n_valid)) + return to_edge_transform_and_lower( + ep, partitioner=[VulkanPartitioner()] + ).to_executorch() + + +class TestFusedCe(unittest.TestCase): + def test_export_delegates(self) -> None: + for cfg in CONFIGS: + if cfg.v > 1024: # width-independent; skip the 128k fixture + continue + with self.subTest(config=cfg.name): + logits, labels, n_valid = _inputs(cfg) + et = _export(logits, labels, n_valid) + found = any( + d.id == "VulkanBackend" + for plan in et.executorch_program.execution_plan + for d in plan.delegates + ) + self.assertTrue(found, f"no VulkanBackend delegate in {cfg.name}") + + def test_op_matches_fp64_golden(self) -> None: + for cfg in CONFIGS: + if cfg.v > 1024: + continue + with self.subTest(config=cfg.name): + logits, labels, n_valid = _inputs(cfg) + loss, dlogits = torch.ops.et_vk.fused_ce(logits, labels, n_valid) + g_loss, g_dlogits = _fp64_golden(logits, labels, n_valid) + torch.testing.assert_close(loss, g_loss, atol=5e-4, rtol=1e-3) + torch.testing.assert_close(dlogits, g_dlogits, atol=5e-4, rtol=1e-3) + + +def export_fused_ce_model(cfg: CeConfig, pte_path: str, golden_path: str) -> None: + logits, labels, n_valid = _inputs(cfg) + et = _export(logits, labels, n_valid) + with open(pte_path, "wb") as f: + f.write(et.buffer) + g_loss, g_dlogits = _fp64_golden(logits, labels, n_valid) + # loss scalar then dlogits, both raw LE fp32 + np.concatenate( + [g_loss.reshape(1).numpy(), g_dlogits.reshape(-1).numpy()] + ).astype(" None: + for cfg in CONFIGS: + pte = os.path.join(out_dir, f"fused_ce_{cfg.name}.pte") + golden = os.path.join(out_dir, f"fused_ce_{cfg.name}.golden.bin") + export_fused_ce_model(cfg, pte, golden) + + +if __name__ == "__main__": + unittest.main()