diff --git a/backends/vulkan/custom_ops_lib.py b/backends/vulkan/custom_ops_lib.py index cf67e4a69d2..89a0d077c78 100644 --- a/backends/vulkan/custom_ops_lib.py +++ b/backends/vulkan/custom_ops_lib.py @@ -1203,6 +1203,58 @@ def fused_ce_backward(ctx, grad_loss, grad_dlogits): 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_q4gsw_dw_impl( d_out: torch.Tensor, diff --git a/backends/vulkan/op_registry.py b/backends/vulkan/op_registry.py index 99f01c30652..8b7484731b5 100644 --- a/backends/vulkan/op_registry.py +++ b/backends/vulkan/op_registry.py @@ -1797,6 +1797,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_q4gsw_dw.default) def register_linear_q4gsw_dw(): return OpFeatures( diff --git a/backends/webgpu/CMakeLists.txt b/backends/webgpu/CMakeLists.txt index d1d002970b2..dc92fcc1947 100644 --- a/backends/webgpu/CMakeLists.txt +++ b/backends/webgpu/CMakeLists.txt @@ -70,6 +70,7 @@ set(WEBGPU_SRCS runtime/ops/linear/Linear.cpp runtime/ops/embedding/Embedding.cpp runtime/ops/logical_not/LogicalNot.cpp + runtime/ops/adamw/AdamwStep.cpp runtime/ops/quantized_linear/QuantizedLinearDw.cpp runtime/ops/quantized_linear/QuantizedLinearRequant.cpp ) diff --git a/backends/webgpu/runtime/ops/adamw/AdamwStep.cpp b/backends/webgpu/runtime/ops/adamw/AdamwStep.cpp new file mode 100644 index 00000000000..4a78694d40b --- /dev/null +++ b/backends/webgpu/runtime/ops/adamw/AdamwStep.cpp @@ -0,0 +1,181 @@ +/* + * 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. + */ + +#include +#include +#include +#include + +#include + +#include +#include +#include + +namespace executorch::backends::webgpu { + +namespace { + +struct AdamwStepParams { + uint32_t numel; + uint32_t _pad0; + uint32_t _pad1; + uint32_t _pad2; + float lr; + float beta1; + float beta2; + float eps; + float weight_decay; + float bias_correction1; + float bias_correction2; + float _pad3; +}; +static_assert(sizeof(AdamwStepParams) == 48, "params must be 48 bytes"); + +// AdamW step over an fp32 latent (elementwise, in place); mirrors torch AdamW. +void adamw_step_impl(WebGPUGraph& graph, const std::vector& args) { + const int param_id = args.at(0); + const int m_id = args.at(1); + const int v_id = args.at(2); + const int grad_id = args.at(3); + + WGPUDevice device = graph.device(); + const auto& param = graph.get_tensor(param_id); + const auto& m = graph.get_tensor(m_id); + const auto& v = graph.get_tensor(v_id); + const auto& grad = graph.get_tensor(grad_id); + + uint64_t numel = 1; + for (int64_t d : param.dims) { + numel *= static_cast(d); + } + if (param.dims.empty() || numel == 0) { + throw std::runtime_error("adamw_step: empty param"); + } + const uint64_t bytes = numel * sizeof(float); + if (param.nbytes != bytes || m.nbytes != bytes || v.nbytes != bytes || + grad.nbytes != bytes) { + throw std::runtime_error( + "adamw_step: param/m/v/grad must be fp32 and same numel"); + } + if (numel > UINT32_MAX) { + throw std::runtime_error("adamw_step: numel exceeds u32"); + } + + auto scalar = [&](int id, const char* name) -> float { + if (graph.get_value_type(id) != WebGPUGraph::ValueType::Double) { + throw std::runtime_error( + std::string("adamw_step: ") + name + " must be a float scalar"); + } + return static_cast(graph.get_double(id)); + }; + + AdamwStepParams params = {}; + params.numel = static_cast(numel); + params.lr = scalar(args.at(4), "lr"); + params.beta1 = scalar(args.at(5), "beta1"); + params.beta2 = scalar(args.at(6), "beta2"); + params.eps = scalar(args.at(7), "eps"); + params.weight_decay = scalar(args.at(8), "weight_decay"); + params.bias_correction1 = scalar(args.at(9), "bias_correction1"); + params.bias_correction2 = scalar(args.at(10), "bias_correction2"); + if (params.bias_correction1 == 0.0f || params.bias_correction2 == 0.0f) { + throw std::runtime_error("adamw_step: bias corrections must be non-zero"); + } + + const uint32_t wg_size = + utils::clamp_workgroup_size(device, kAdamwStepWorkgroupSizeX); + const uint32_t workgroup_count = utils::compute_1d_workgroup_count( + device, params.numel, wg_size, "adamw_step"); + + WGPUBuffer uniform_buffer = + utils::make_uniform(device, ¶ms, sizeof(params)); + graph.add_uniform_buffer_bytes(sizeof(params)); + + WGPUShaderSourceWGSL wgsl_desc = {}; + wgsl_desc.chain.sType = WGPUSType_ShaderSourceWGSL; + wgsl_desc.code = {kAdamwStepWGSL, WGPU_STRLEN}; + WGPUShaderModuleDescriptor shader_desc = {}; + shader_desc.nextInChain = &wgsl_desc.chain; + WGPUShaderModule shader = wgpuDeviceCreateShaderModule(device, &shader_desc); + + WGPUBindGroupLayoutEntry entries[5] = {}; + for (uint32_t i = 0; i <= 2; i++) { + entries[i].binding = i; + entries[i].visibility = WGPUShaderStage_Compute; + entries[i].buffer.type = WGPUBufferBindingType_Storage; + } + entries[3].binding = 3; + entries[3].visibility = WGPUShaderStage_Compute; + entries[3].buffer.type = WGPUBufferBindingType_ReadOnlyStorage; + entries[4].binding = 4; + entries[4].visibility = WGPUShaderStage_Compute; + entries[4].buffer.type = WGPUBufferBindingType_Uniform; + + WGPUBindGroupLayoutDescriptor bgl_desc = {}; + bgl_desc.entryCount = 5; + bgl_desc.entries = entries; + WGPUBindGroupLayout bgl = wgpuDeviceCreateBindGroupLayout(device, &bgl_desc); + + WGPUPipelineLayoutDescriptor pl_desc = {}; + pl_desc.bindGroupLayoutCount = 1; + pl_desc.bindGroupLayouts = &bgl; + WGPUPipelineLayout pipeline_layout = + wgpuDeviceCreatePipelineLayout(device, &pl_desc); + + WGPUConstantEntry wg_size_constant = {}; + wg_size_constant.key = {"wg_size", WGPU_STRLEN}; + wg_size_constant.value = static_cast(wg_size); + + WGPUComputePipelineDescriptor pipeline_desc = {}; + pipeline_desc.layout = pipeline_layout; + pipeline_desc.compute.module = shader; + pipeline_desc.compute.entryPoint = {"main", WGPU_STRLEN}; + pipeline_desc.compute.constantCount = 1; + pipeline_desc.compute.constants = &wg_size_constant; + WGPUComputePipeline pipeline = + wgpuDeviceCreateComputePipeline(device, &pipeline_desc); + + WGPUBindGroupEntry bg_entries[5] = {}; + bg_entries[0].binding = 0; + bg_entries[0].buffer = param.buffer; + bg_entries[0].size = param.nbytes; + bg_entries[1].binding = 1; + bg_entries[1].buffer = m.buffer; + bg_entries[1].size = m.nbytes; + bg_entries[2].binding = 2; + bg_entries[2].buffer = v.buffer; + bg_entries[2].size = v.nbytes; + bg_entries[3].binding = 3; + bg_entries[3].buffer = grad.buffer; + bg_entries[3].size = grad.nbytes; + bg_entries[4].binding = 4; + bg_entries[4].buffer = uniform_buffer; + bg_entries[4].size = sizeof(params); + + WGPUBindGroupDescriptor bg_desc = {}; + bg_desc.layout = bgl; + bg_desc.entryCount = 5; + bg_desc.entries = bg_entries; + WGPUBindGroup bind_group = wgpuDeviceCreateBindGroup(device, &bg_desc); + + graph.add_dispatch({pipeline, bind_group, workgroup_count, "adamw_step"}); + + wgpuShaderModuleRelease(shader); + wgpuBindGroupLayoutRelease(bgl); + wgpuPipelineLayoutRelease(pipeline_layout); + graph.own_uniform_buffer(uniform_buffer); +} + +} // namespace + +WEBGPU_REGISTER_OPERATORS { + WEBGPU_REGISTER_OP(et_vk.adamw_step.default, adamw_step_impl); +} + +} // namespace executorch::backends::webgpu diff --git a/backends/webgpu/runtime/ops/adamw/adamw_step.wgsl b/backends/webgpu/runtime/ops/adamw/adamw_step.wgsl new file mode 100644 index 00000000000..f6a5ad58943 --- /dev/null +++ b/backends/webgpu/runtime/ops/adamw/adamw_step.wgsl @@ -0,0 +1,41 @@ + +@group(0) @binding(0) var t_param: array; +@group(0) @binding(1) var t_m: array; +@group(0) @binding(2) var t_v: array; +@group(0) @binding(3) var t_grad: array; + +struct Params { + numel: u32, + _pad0: u32, + _pad1: u32, + _pad2: u32, + lr: f32, + beta1: f32, + beta2: f32, + eps: f32, + weight_decay: f32, + bias_correction1: f32, + bias_correction2: f32, + _pad3: f32, +} +@group(0) @binding(4) var params: Params; + +override wg_size: u32 = 64u; + +@compute @workgroup_size(wg_size, 1, 1) +fn main(@builtin(global_invocation_id) gid: vec3) { + let i = gid.x; + if (i >= params.numel) { + return; + } + let g = t_grad[i]; + var p = t_param[i]; + p = p - params.lr * params.weight_decay * p; + let m = params.beta1 * t_m[i] + (1.0 - params.beta1) * g; + let v = params.beta2 * t_v[i] + (1.0 - params.beta2) * g * g; + t_m[i] = m; + t_v[i] = v; + let mhat = m / params.bias_correction1; + let vhat = v / params.bias_correction2; + t_param[i] = p - params.lr * mhat / (sqrt(vhat) + params.eps); +} diff --git a/backends/webgpu/runtime/ops/adamw/adamw_step_wgsl.h b/backends/webgpu/runtime/ops/adamw/adamw_step_wgsl.h new file mode 100644 index 00000000000..38cbda98c6b --- /dev/null +++ b/backends/webgpu/runtime/ops/adamw/adamw_step_wgsl.h @@ -0,0 +1,65 @@ +/* + * 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. + */ + +#pragma once + +#include + +namespace executorch::backends::webgpu { + +// @generated from adamw_step.wgsl - DO NOT EDIT. +// wgsl-sha256: 0957c04168872db5e2b39cf5f26beefba27f3b5514ec69cece4de5145e97156f +inline constexpr const char* kAdamwStepWGSL = R"( + +@group(0) @binding(0) var t_param: array; +@group(0) @binding(1) var t_m: array; +@group(0) @binding(2) var t_v: array; +@group(0) @binding(3) var t_grad: array; + +struct Params { + numel: u32, + _pad0: u32, + _pad1: u32, + _pad2: u32, + lr: f32, + beta1: f32, + beta2: f32, + eps: f32, + weight_decay: f32, + bias_correction1: f32, + bias_correction2: f32, + _pad3: f32, +} +@group(0) @binding(4) var params: Params; + +override wg_size: u32 = 64u; + +@compute @workgroup_size(wg_size, 1, 1) +fn main(@builtin(global_invocation_id) gid: vec3) { + let i = gid.x; + if (i >= params.numel) { + return; + } + let g = t_grad[i]; + var p = t_param[i]; + p = p - params.lr * params.weight_decay * p; + let m = params.beta1 * t_m[i] + (1.0 - params.beta1) * g; + let v = params.beta2 * t_v[i] + (1.0 - params.beta2) * g * g; + t_m[i] = m; + t_v[i] = v; + let mhat = m / params.bias_correction1; + let vhat = v / params.bias_correction2; + t_param[i] = p - params.lr * mhat / (sqrt(vhat) + params.eps); +} +)"; + +inline constexpr uint32_t kAdamwStepWorkgroupSizeX = 64; +inline constexpr uint32_t kAdamwStepWorkgroupSizeY = 1; +inline constexpr uint32_t kAdamwStepWorkgroupSizeZ = 1; + +} // namespace executorch::backends::webgpu