[ExecuTorch][WebGPU] Tests for the q4gsw_train WebGPU ops

backends/vulkan/custom_ops_lib.py

@@ -1153,6 +1153,108 @@ def rms_norm_impl(
 rms_norm_op = getattr(getattr(torch.ops, namespace), name)
 
 
+########################
+## fused_ce (training) ##
+########################
+
+
+def fused_ce_impl(
+    logits: torch.Tensor,
+    labels: torch.Tensor,
+    n_valid: float,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    mask = labels >= 0
+    safe = labels.clamp(min=0).long()
+    lse = torch.logsumexp(logits, dim=-1)
+    picked = logits.gather(-1, safe[:, None]).squeeze(-1)
+    loss = torch.where(mask, (lse - picked) / n_valid, torch.zeros_like(lse)).sum()
+    softmax = torch.softmax(logits, dim=-1)
+    onehot = torch.nn.functional.one_hot(safe, logits.shape[-1]).to(logits.dtype)
+    dlogits = torch.where(
+        mask[:, None], (softmax - onehot) / n_valid, torch.zeros_like(softmax)
+    )
+    return loss, dlogits
+
+
+def fused_ce_meta(
+    logits: torch.Tensor,
+    labels: torch.Tensor,
+    n_valid: float,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    return logits.new_empty([]), torch.empty_like(logits)
+
+
+def fused_ce_setup_context(ctx, inputs, output) -> None:
+    ctx.save_for_backward(output[1])
+
+
+def fused_ce_backward(ctx, grad_loss, grad_dlogits):
+    (dlogits,) = ctx.saved_tensors
+    return grad_loss * dlogits, None, None
+
+
+name = "fused_ce"
+lib.define(f"{name}(Tensor logits, Tensor labels, float n_valid) -> (Tensor, Tensor)")
+lib.impl(name, fused_ce_impl, "CompositeExplicitAutograd")
+lib.impl(name, fused_ce_meta, "Meta")
+torch.library.register_autograd(
+    f"{namespace}::{name}", fused_ce_backward, setup_context=fused_ce_setup_context
+)
+fused_ce_op = getattr(getattr(torch.ops, namespace), name)
+
+
+###########################
+## adamw_step (training) ##
+###########################
+
+
+def adamw_step_impl(
+    param: torch.Tensor,
+    m: torch.Tensor,
+    v: torch.Tensor,
+    grad: torch.Tensor,
+    lr: float,
+    beta1: float,
+    beta2: float,
+    eps: float,
+    weight_decay: float,
+    bias_correction1: float,
+    bias_correction2: float,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    param.mul_(1.0 - lr * weight_decay)
+    m.mul_(beta1).add_(grad, alpha=1.0 - beta1)
+    v.mul_(beta2).addcmul_(grad, grad, value=1.0 - beta2)
+    mhat = m / bias_correction1
+    denom = (v / bias_correction2).sqrt() + eps
+    param.addcdiv_(mhat, denom, value=-lr)
+    return param, m, v
+
+
+def adamw_step_meta(
+    param: torch.Tensor,
+    m: torch.Tensor,
+    v: torch.Tensor,
+    grad: torch.Tensor,
+    lr: float,
+    beta1: float,
+    beta2: float,
+    eps: float,
+    weight_decay: float,
+    bias_correction1: float,
+    bias_correction2: float,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    return param, m, v
+
+
+name = "adamw_step"
+lib.define(
+    f"{name}(Tensor(a!) param, Tensor(b!) m, Tensor(c!) v, Tensor grad, float lr, float beta1, float beta2, float eps, float weight_decay, float bias_correction1, float bias_correction2) -> (Tensor(a!), Tensor(b!), Tensor(c!))"
+)
+lib.impl(name, adamw_step_impl, "CompositeExplicitAutograd")
+lib.impl(name, adamw_step_meta, "Meta")
+adamw_step_op = getattr(getattr(torch.ops, namespace), name)
+
+
 # STE weight gradient d_out^T @ x through the frozen 4-bit linear_q4gsw base.
 def linear_dW_impl(
     d_out: torch.Tensor,

backends/vulkan/op_registry.py

@@ -1755,6 +1755,20 @@ def register_rms_norm():
     )
 
 
+# =============================================================================
+# FusedCe.cpp (training)
+# =============================================================================
+
+
+@update_features(exir_ops.edge.et_vk.fused_ce.default)
+def register_fused_ce():
+    return OpFeatures(
+        inputs_storage=utils.CONTIGUOUS_ANY,
+        inputs_dtypes=[utils.FP_T, utils.INT_T, utils.NONE_T],
+        outputs_dtypes=[utils.FP_T, utils.FP_T],
+    )
+
+
 @update_features(
     [
         exir_ops.edge.aten.ne.Scalar,
@@ -1784,6 +1798,14 @@ def register_logical_not():
     )
 
 
+@update_features("et_vk::adamw_step")
+def register_adamw_step():
+    return OpFeatures(
+        inputs_storage=utils.CONTIGUOUS_ANY,
+        inputs_dtypes=utils.FP_T,
+    )
+
+
 @update_features(exir_ops.edge.et_vk.linear_dW.default)
 def register_linear_dW():
     return OpFeatures(

backends/vulkan/runtime/graph/ops/glsl/adamw_step.glsl

@@ -0,0 +1,54 @@
+/*
+ * 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.
+ */
+
+#version 450 core
+
+#define PRECISION ${PRECISION}
+
+#define T ${buffer_scalar_type(DTYPE)}
+
+${define_active_storage_type(STORAGE)}
+
+${define_required_extensions(STORAGE, DTYPE)}
+
+layout(std430) buffer;
+
+${layout_declare_tensor(B, "rw", "t_param", DTYPE, STORAGE)}
+${layout_declare_tensor(B, "rw", "t_m", DTYPE, STORAGE)}
+${layout_declare_tensor(B, "rw", "t_v", DTYPE, STORAGE)}
+${layout_declare_tensor(B, "r", "t_grad", DTYPE, STORAGE)}
+
+layout(push_constant) uniform restrict Block {
+  int numel;
+  float lr;
+  float beta1;
+  float beta2;
+  float eps;
+  float weight_decay;
+  float bias_correction1;
+  float bias_correction2;
+};
+
+layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in;
+
+void main() {
+  const int i = int(gl_GlobalInvocationID.x);
+  if (i >= numel) {
+    return;
+  }
+  T g = t_grad[i];
+  T p = t_param[i];
+  p = p - lr * weight_decay * p;
+  T m = beta1 * t_m[i] + (1.0 - beta1) * g;
+  T v = beta2 * t_v[i] + (1.0 - beta2) * g * g;
+  t_m[i] = m;
+  t_v[i] = v;
+  T mhat = m / bias_correction1;
+  T vhat = v / bias_correction2;
+  t_param[i] = p - lr * mhat / (sqrt(vhat) + eps);
+}

backends/vulkan/runtime/graph/ops/glsl/adamw_step.yaml

@@ -0,0 +1,17 @@
+# 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.
+
+adamw_step:
+  parameter_names_with_default_values:
+    DTYPE: float
+    STORAGE: buffer
+  generate_variant_forall:
+    STORAGE:
+      - VALUE: buffer
+    DTYPE:
+      - VALUE: float
+  shader_variants:
+    - NAME: adamw_step

backends/vulkan/runtime/graph/ops/glsl/fused_ce.glsl

@@ -0,0 +1,110 @@
+/*
+ * 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.
+ */
+
+#version 450 core
+
+${define_required_extensions(STORAGE, DTYPE)}
+
+#define PRECISION ${PRECISION}
+#define T ${buffer_scalar_type(DTYPE)}
+
+${define_active_storage_type(STORAGE)}
+
+layout(std430) buffer;
+
+${layout_declare_tensor(B, "w", "t_dlogits", DTYPE, "buffer")}
+${layout_declare_tensor(B, "w", "t_loss_partial", DTYPE, "buffer")}
+${layout_declare_tensor(B, "r", "t_logits", DTYPE, "buffer")}
+${layout_declare_tensor(B, "r", "t_labels", "int", "buffer")}
+
+${layout_declare_ubo(B, "int", "vocab")}
+${layout_declare_ubo(B, "int", "n_rows")}
+${layout_declare_ubo(B, "float", "n_valid")}
+
+layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in;
+
+#define NWORKERS 64
+
+shared float red_m[NWORKERS];
+shared float red_l[NWORKERS];
+
+// Fused cross-entropy: one workgroup per row cooperatively reduces the vocab
+// dimension with a single-pass online softmax (running max + rescaled running
+// sum), then writes per-row loss and dlogits. Rows with label < 0 are masked.
+void main() {
+  const uint row = gl_GlobalInvocationID.y;
+  if (int(row) >= n_rows) {
+    return;
+  }
+
+  const uint tid = gl_LocalInvocationID.x;
+  const uint V = uint(vocab);
+  const uint base = row * V;
+  const int lbl = t_labels[row];
+
+  if (lbl < 0) {
+    for (uint j = tid; j < V; j += NWORKERS) {
+      t_dlogits[base + j] = T(0);
+    }
+    if (tid == 0u) {
+      t_loss_partial[row] = T(0);
+    }
+    return;
+  }
+
+  // Single read pass: maintain a running max m and the sum l of exp(x - m),
+  // rescaling l whenever a larger value updates m. Finite -3.4e38 init.
+  float m = -3.4e38;
+  float l = 0.0;
+  for (uint j = tid; j < V; j += NWORKERS) {
+    float x = float(t_logits[base + j]);
+    if (x > m) {
+      l = l * exp(m - x) + 1.0;
+      m = x;
+    } else {
+      l = l + exp(x - m);
+    }
+  }
+  red_m[tid] = m;
+  red_l[tid] = l;
+  memoryBarrierShared();
+  barrier();
+
+  // Tree-combine the (m, l) pairs: m becomes the max, l is rescaled to it.
+  for (uint s = NWORKERS / 2u; s > 0u; s >>= 1u) {
+    if (tid < s) {
+      float ma = red_m[tid];
+      float la = red_l[tid];
+      float mb = red_m[tid + s];
+      float lb = red_l[tid + s];
+      float mm = max(ma, mb);
+      red_m[tid] = mm;
+      red_l[tid] = la * exp(ma - mm) + lb * exp(mb - mm);
+    }
+    memoryBarrierShared();
+    barrier();
+  }
+
+  const float row_max = red_m[0];
+  const float denom = red_l[0];
+  const float inv = 1.0 / denom;
+  const float scale = 1.0 / n_valid;
+
+  if (tid == 0u) {
+    const float lse = row_max + log(denom);
+    t_loss_partial[row] = T((lse - float(t_logits[base + uint(lbl)])) * scale);
+  }
+
+  for (uint j = tid; j < V; j += NWORKERS) {
+    float g = exp(float(t_logits[base + j]) - row_max) * inv * scale;
+    if (j == uint(lbl)) {
+      g = g - scale;
+    }
+    t_dlogits[base + j] = T(g);
+  }
+}

backends/vulkan/runtime/graph/ops/glsl/fused_ce.yaml

@@ -0,0 +1,15 @@
+# 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_ce:
+  parameter_names_with_default_values:
+    DTYPE: float
+    STORAGE: buffer
+  generate_variant_forall:
+    DTYPE:
+      - VALUE: float
+  shader_variants:
+    - NAME: fused_ce_buffer