Skip to content

Commit 427d25d

Browse files
JCNTHJCNTH
committed
[ExecuTorch][WebGPU] Add et_vk.adamw_step (on-GPU AdamW optimizer step) to the WebGPU backend
Pull Request resolved: #20935 Add `et_vk.adamw_step` — an on-GPU AdamW optimizer step (delegated in-place parameter update) for on-device training. Key changes: - `runtime/ops/adamw/` — elementwise AdamW-update WGSL kernel + handler - `custom_ops_lib.py` — `adamw_step` custom-op def (mutating, `Tensor(a!)`) - `op_registry.py` — `OpFeatures` (string key `et_vk::adamw_step`) - `CMakeLists.txt` `WEBGPU_SRCS` — wire the source Delegating the optimizer step to the GPU is new to ExecuTorch — its optimizers are host-side C++ (`extension/training/optimizer/adamw.cpp`). Training-only WebGPU custom op under the shared Vulkan partitioner; no Vulkan kernel yet. Co-authored-with: Claude Code. ghstack-source-id: 405026081 @exported-using-ghexport Differential Revision: [D111755116](https://our.internmc.facebook.com/intern/diff/D111755116/)
1 parent ef983fe commit 427d25d

6 files changed

Lines changed: 348 additions & 0 deletions

File tree

‎backends/vulkan/custom_ops_lib.py‎

Lines changed: 52 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1203,6 +1203,58 @@ def fused_ce_backward(ctx, grad_loss, grad_dlogits):
12031203
fused_ce_op = getattr(getattr(torch.ops, namespace), name)
12041204

12051205

1206+
###########################
1207+
## adamw_step (training) ##
1208+
###########################
1209+
1210+
1211+
def adamw_step_impl(
1212+
param: torch.Tensor,
1213+
m: torch.Tensor,
1214+
v: torch.Tensor,
1215+
grad: torch.Tensor,
1216+
lr: float,
1217+
beta1: float,
1218+
beta2: float,
1219+
eps: float,
1220+
weight_decay: float,
1221+
bias_correction1: float,
1222+
bias_correction2: float,
1223+
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
1224+
param.mul_(1.0 - lr * weight_decay)
1225+
m.mul_(beta1).add_(grad, alpha=1.0 - beta1)
1226+
v.mul_(beta2).addcmul_(grad, grad, value=1.0 - beta2)
1227+
mhat = m / bias_correction1
1228+
denom = (v / bias_correction2).sqrt() + eps
1229+
param.addcdiv_(mhat, denom, value=-lr)
1230+
return param, m, v
1231+
1232+
1233+
def adamw_step_meta(
1234+
param: torch.Tensor,
1235+
m: torch.Tensor,
1236+
v: torch.Tensor,
1237+
grad: torch.Tensor,
1238+
lr: float,
1239+
beta1: float,
1240+
beta2: float,
1241+
eps: float,
1242+
weight_decay: float,
1243+
bias_correction1: float,
1244+
bias_correction2: float,
1245+
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
1246+
return param, m, v
1247+
1248+
1249+
name = "adamw_step"
1250+
lib.define(
1251+
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!))"
1252+
)
1253+
lib.impl(name, adamw_step_impl, "CompositeExplicitAutograd")
1254+
lib.impl(name, adamw_step_meta, "Meta")
1255+
adamw_step_op = getattr(getattr(torch.ops, namespace), name)
1256+
1257+
12061258
# STE weight gradient d_out^T @ x through the frozen 4-bit linear_q4gsw base.
12071259
def linear_dW_impl(
12081260
d_out: torch.Tensor,

‎backends/vulkan/op_registry.py‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1798,6 +1798,14 @@ def register_logical_not():
17981798
)
17991799

18001800

1801+
@update_features("et_vk::adamw_step")
1802+
def register_adamw_step():
1803+
return OpFeatures(
1804+
inputs_storage=utils.CONTIGUOUS_ANY,
1805+
inputs_dtypes=utils.FP_T,
1806+
)
1807+
1808+
18011809
@update_features(exir_ops.edge.et_vk.linear_dW.default)
18021810
def register_linear_dW():
18031811
return OpFeatures(

‎backends/webgpu/CMakeLists.txt‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -69,6 +69,7 @@ set(WEBGPU_SRCS
6969
runtime/ops/dim_order/DimOrder.cpp
7070
runtime/ops/linear/Linear.cpp
7171
runtime/ops/embedding/Embedding.cpp
72+
runtime/ops/adamw/AdamwStep.cpp
7273
runtime/ops/quantized_linear/LinearDw.cpp
7374
runtime/ops/quantized_linear/QuantizedLinearRequant.cpp
7475
)

‎backends/webgpu/runtime/ops/adamw/AdamwStep.cpp‎

Lines changed: 181 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,181 @@
1+
/*
2+
* Copyright (c) Meta Platforms, Inc. and affiliates.
3+
* All rights reserved.
4+
*
5+
* This source code is licensed under the BSD-style license found in the
6+
* LICENSE file in the root directory of this source tree.
7+
*/
8+
9+
#include <executorch/backends/webgpu/runtime/WebGPUGraph.h>
10+
#include <executorch/backends/webgpu/runtime/WebGPUUtils.h>
11+
#include <executorch/backends/webgpu/runtime/ops/OperatorRegistry.h>
12+
#include <executorch/backends/webgpu/runtime/ops/adamw/adamw_step_wgsl.h>
13+
14+
#include <webgpu/webgpu.h>
15+
16+
#include <cstdint>
17+
#include <cstring>
18+
#include <stdexcept>
19+
20+
namespace executorch::backends::webgpu {
21+
22+
namespace {
23+
24+
struct AdamwStepParams {
25+
uint32_t numel;
26+
uint32_t _pad0;
27+
uint32_t _pad1;
28+
uint32_t _pad2;
29+
float lr;
30+
float beta1;
31+
float beta2;
32+
float eps;
33+
float weight_decay;
34+
float bias_correction1;
35+
float bias_correction2;
36+
float _pad3;
37+
};
38+
static_assert(sizeof(AdamwStepParams) == 48, "params must be 48 bytes");
39+
40+
// AdamW step over an fp32 latent (elementwise, in place); mirrors torch AdamW.
41+
void adamw_step_impl(WebGPUGraph& graph, const std::vector<int>& args) {
42+
const int param_id = args.at(0);
43+
const int m_id = args.at(1);
44+
const int v_id = args.at(2);
45+
const int grad_id = args.at(3);
46+
47+
WGPUDevice device = graph.device();
48+
const auto& param = graph.get_tensor(param_id);
49+
const auto& m = graph.get_tensor(m_id);
50+
const auto& v = graph.get_tensor(v_id);
51+
const auto& grad = graph.get_tensor(grad_id);
52+
53+
uint64_t numel = 1;
54+
for (int64_t d : param.dims) {
55+
numel *= static_cast<uint64_t>(d);
56+
}
57+
if (param.dims.empty() || numel == 0) {
58+
throw std::runtime_error("adamw_step: empty param");
59+
}
60+
const uint64_t bytes = numel * sizeof(float);
61+
if (param.nbytes != bytes || m.nbytes != bytes || v.nbytes != bytes ||
62+
grad.nbytes != bytes) {
63+
throw std::runtime_error(
64+
"adamw_step: param/m/v/grad must be fp32 and same numel");
65+
}
66+
if (numel > UINT32_MAX) {
67+
throw std::runtime_error("adamw_step: numel exceeds u32");
68+
}
69+
70+
auto scalar = [&](int id, const char* name) -> float {
71+
if (graph.get_value_type(id) != WebGPUGraph::ValueType::Double) {
72+
throw std::runtime_error(
73+
std::string("adamw_step: ") + name + " must be a float scalar");
74+
}
75+
return static_cast<float>(graph.get_double(id));
76+
};
77+
78+
AdamwStepParams params = {};
79+
params.numel = static_cast<uint32_t>(numel);
80+
params.lr = scalar(args.at(4), "lr");
81+
params.beta1 = scalar(args.at(5), "beta1");
82+
params.beta2 = scalar(args.at(6), "beta2");
83+
params.eps = scalar(args.at(7), "eps");
84+
params.weight_decay = scalar(args.at(8), "weight_decay");
85+
params.bias_correction1 = scalar(args.at(9), "bias_correction1");
86+
params.bias_correction2 = scalar(args.at(10), "bias_correction2");
87+
if (params.bias_correction1 == 0.0f || params.bias_correction2 == 0.0f) {
88+
throw std::runtime_error("adamw_step: bias corrections must be non-zero");
89+
}
90+
91+
const uint32_t wg_size =
92+
utils::clamp_workgroup_size(device, kAdamwStepWorkgroupSizeX);
93+
const uint32_t workgroup_count = utils::compute_1d_workgroup_count(
94+
device, params.numel, wg_size, "adamw_step");
95+
96+
WGPUBuffer uniform_buffer =
97+
utils::make_uniform(device, &params, sizeof(params));
98+
graph.add_uniform_buffer_bytes(sizeof(params));
99+
100+
WGPUShaderSourceWGSL wgsl_desc = {};
101+
wgsl_desc.chain.sType = WGPUSType_ShaderSourceWGSL;
102+
wgsl_desc.code = {kAdamwStepWGSL, WGPU_STRLEN};
103+
WGPUShaderModuleDescriptor shader_desc = {};
104+
shader_desc.nextInChain = &wgsl_desc.chain;
105+
WGPUShaderModule shader = wgpuDeviceCreateShaderModule(device, &shader_desc);
106+
107+
WGPUBindGroupLayoutEntry entries[5] = {};
108+
for (uint32_t i = 0; i <= 2; i++) {
109+
entries[i].binding = i;
110+
entries[i].visibility = WGPUShaderStage_Compute;
111+
entries[i].buffer.type = WGPUBufferBindingType_Storage;
112+
}
113+
entries[3].binding = 3;
114+
entries[3].visibility = WGPUShaderStage_Compute;
115+
entries[3].buffer.type = WGPUBufferBindingType_ReadOnlyStorage;
116+
entries[4].binding = 4;
117+
entries[4].visibility = WGPUShaderStage_Compute;
118+
entries[4].buffer.type = WGPUBufferBindingType_Uniform;
119+
120+
WGPUBindGroupLayoutDescriptor bgl_desc = {};
121+
bgl_desc.entryCount = 5;
122+
bgl_desc.entries = entries;
123+
WGPUBindGroupLayout bgl = wgpuDeviceCreateBindGroupLayout(device, &bgl_desc);
124+
125+
WGPUPipelineLayoutDescriptor pl_desc = {};
126+
pl_desc.bindGroupLayoutCount = 1;
127+
pl_desc.bindGroupLayouts = &bgl;
128+
WGPUPipelineLayout pipeline_layout =
129+
wgpuDeviceCreatePipelineLayout(device, &pl_desc);
130+
131+
WGPUConstantEntry wg_size_constant = {};
132+
wg_size_constant.key = {"wg_size", WGPU_STRLEN};
133+
wg_size_constant.value = static_cast<double>(wg_size);
134+
135+
WGPUComputePipelineDescriptor pipeline_desc = {};
136+
pipeline_desc.layout = pipeline_layout;
137+
pipeline_desc.compute.module = shader;
138+
pipeline_desc.compute.entryPoint = {"main", WGPU_STRLEN};
139+
pipeline_desc.compute.constantCount = 1;
140+
pipeline_desc.compute.constants = &wg_size_constant;
141+
WGPUComputePipeline pipeline =
142+
wgpuDeviceCreateComputePipeline(device, &pipeline_desc);
143+
144+
WGPUBindGroupEntry bg_entries[5] = {};
145+
bg_entries[0].binding = 0;
146+
bg_entries[0].buffer = param.buffer;
147+
bg_entries[0].size = param.nbytes;
148+
bg_entries[1].binding = 1;
149+
bg_entries[1].buffer = m.buffer;
150+
bg_entries[1].size = m.nbytes;
151+
bg_entries[2].binding = 2;
152+
bg_entries[2].buffer = v.buffer;
153+
bg_entries[2].size = v.nbytes;
154+
bg_entries[3].binding = 3;
155+
bg_entries[3].buffer = grad.buffer;
156+
bg_entries[3].size = grad.nbytes;
157+
bg_entries[4].binding = 4;
158+
bg_entries[4].buffer = uniform_buffer;
159+
bg_entries[4].size = sizeof(params);
160+
161+
WGPUBindGroupDescriptor bg_desc = {};
162+
bg_desc.layout = bgl;
163+
bg_desc.entryCount = 5;
164+
bg_desc.entries = bg_entries;
165+
WGPUBindGroup bind_group = wgpuDeviceCreateBindGroup(device, &bg_desc);
166+
167+
graph.add_dispatch({pipeline, bind_group, workgroup_count, "adamw_step"});
168+
169+
wgpuShaderModuleRelease(shader);
170+
wgpuBindGroupLayoutRelease(bgl);
171+
wgpuPipelineLayoutRelease(pipeline_layout);
172+
graph.own_uniform_buffer(uniform_buffer);
173+
}
174+
175+
} // namespace
176+
177+
WEBGPU_REGISTER_OPERATORS {
178+
WEBGPU_REGISTER_OP(et_vk.adamw_step.default, adamw_step_impl);
179+
}
180+
181+
} // namespace executorch::backends::webgpu

‎backends/webgpu/runtime/ops/adamw/adamw_step.wgsl‎

Lines changed: 41 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,41 @@
1+
2+
@group(0) @binding(0) var<storage, read_write> t_param: array<f32>;
3+
@group(0) @binding(1) var<storage, read_write> t_m: array<f32>;
4+
@group(0) @binding(2) var<storage, read_write> t_v: array<f32>;
5+
@group(0) @binding(3) var<storage, read> t_grad: array<f32>;
6+
7+
struct Params {
8+
numel: u32,
9+
_pad0: u32,
10+
_pad1: u32,
11+
_pad2: u32,
12+
lr: f32,
13+
beta1: f32,
14+
beta2: f32,
15+
eps: f32,
16+
weight_decay: f32,
17+
bias_correction1: f32,
18+
bias_correction2: f32,
19+
_pad3: f32,
20+
}
21+
@group(0) @binding(4) var<uniform> params: Params;
22+
23+
override wg_size: u32 = 64u;
24+
25+
@compute @workgroup_size(wg_size, 1, 1)
26+
fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
27+
let i = gid.x;
28+
if (i >= params.numel) {
29+
return;
30+
}
31+
let g = t_grad[i];
32+
var p = t_param[i];
33+
p = p - params.lr * params.weight_decay * p;
34+
let m = params.beta1 * t_m[i] + (1.0 - params.beta1) * g;
35+
let v = params.beta2 * t_v[i] + (1.0 - params.beta2) * g * g;
36+
t_m[i] = m;
37+
t_v[i] = v;
38+
let mhat = m / params.bias_correction1;
39+
let vhat = v / params.bias_correction2;
40+
t_param[i] = p - params.lr * mhat / (sqrt(vhat) + params.eps);
41+
}

‎backends/webgpu/runtime/ops/adamw/adamw_step_wgsl.h‎

Lines changed: 65 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,65 @@
1+
/*
2+
* Copyright (c) Meta Platforms, Inc. and affiliates.
3+
* All rights reserved.
4+
*
5+
* This source code is licensed under the BSD-style license found in the
6+
* LICENSE file in the root directory of this source tree.
7+
*/
8+
9+
#pragma once
10+
11+
#include <cstdint>
12+
13+
namespace executorch::backends::webgpu {
14+
15+
// @generated from adamw_step.wgsl - DO NOT EDIT.
16+
// wgsl-sha256: 0957c04168872db5e2b39cf5f26beefba27f3b5514ec69cece4de5145e97156f
17+
inline constexpr const char* kAdamwStepWGSL = R"(
18+
19+
@group(0) @binding(0) var<storage, read_write> t_param: array<f32>;
20+
@group(0) @binding(1) var<storage, read_write> t_m: array<f32>;
21+
@group(0) @binding(2) var<storage, read_write> t_v: array<f32>;
22+
@group(0) @binding(3) var<storage, read> t_grad: array<f32>;
23+
24+
struct Params {
25+
numel: u32,
26+
_pad0: u32,
27+
_pad1: u32,
28+
_pad2: u32,
29+
lr: f32,
30+
beta1: f32,
31+
beta2: f32,
32+
eps: f32,
33+
weight_decay: f32,
34+
bias_correction1: f32,
35+
bias_correction2: f32,
36+
_pad3: f32,
37+
}
38+
@group(0) @binding(4) var<uniform> params: Params;
39+
40+
override wg_size: u32 = 64u;
41+
42+
@compute @workgroup_size(wg_size, 1, 1)
43+
fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
44+
let i = gid.x;
45+
if (i >= params.numel) {
46+
return;
47+
}
48+
let g = t_grad[i];
49+
var p = t_param[i];
50+
p = p - params.lr * params.weight_decay * p;
51+
let m = params.beta1 * t_m[i] + (1.0 - params.beta1) * g;
52+
let v = params.beta2 * t_v[i] + (1.0 - params.beta2) * g * g;
53+
t_m[i] = m;
54+
t_v[i] = v;
55+
let mhat = m / params.bias_correction1;
56+
let vhat = v / params.bias_correction2;
57+
t_param[i] = p - params.lr * mhat / (sqrt(vhat) + params.eps);
58+
}
59+
)";
60+
61+
inline constexpr uint32_t kAdamwStepWorkgroupSizeX = 64;
62+
inline constexpr uint32_t kAdamwStepWorkgroupSizeY = 1;
63+
inline constexpr uint32_t kAdamwStepWorkgroupSizeZ = 1;
64+
65+
} // namespace executorch::backends::webgpu

0 commit comments

Comments
 (0)