diff --git a/backends/vulkan/runtime/graph/ops/glsl/adamw_step.glsl b/backends/vulkan/runtime/graph/ops/glsl/adamw_step.glsl new file mode 100644 index 00000000000..ccb7b34bcb7 --- /dev/null +++ b/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); +} diff --git a/backends/vulkan/runtime/graph/ops/glsl/adamw_step.yaml b/backends/vulkan/runtime/graph/ops/glsl/adamw_step.yaml new file mode 100644 index 00000000000..114a257cfad --- /dev/null +++ b/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 diff --git a/backends/vulkan/runtime/graph/ops/impl/AdamwStep.cpp b/backends/vulkan/runtime/graph/ops/impl/AdamwStep.cpp new file mode 100644 index 00000000000..489e95a773d --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/impl/AdamwStep.cpp @@ -0,0 +1,126 @@ +/* + * 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 + +namespace vkcompute { + +void check_adamw_step_args( + ComputeGraph& graph, + const ValueRef param, + const ValueRef m, + const ValueRef v, + const ValueRef grad) { + VK_CHECK_COND(graph.dtype_of(param) == vkapi::kFloat); + VK_CHECK_COND(graph.dtype_of(m) == vkapi::kFloat); + VK_CHECK_COND(graph.dtype_of(v) == vkapi::kFloat); + VK_CHECK_COND(graph.dtype_of(grad) == vkapi::kFloat); + + const int32_t numel = graph.numel_of(param); + VK_CHECK_COND(graph.numel_of(m) == numel); + VK_CHECK_COND(graph.numel_of(v) == numel); + VK_CHECK_COND(graph.numel_of(grad) == numel); +} + +void add_adamw_step_node( + ComputeGraph& graph, + const ValueRef param, + const ValueRef m, + const ValueRef v, + const ValueRef grad, + const ValueRef lr, + const ValueRef beta1, + const ValueRef beta2, + const ValueRef eps, + const ValueRef weight_decay, + const ValueRef bias_correction1, + const ValueRef bias_correction2) { + check_adamw_step_args(graph, param, m, v, grad); + + const float bc1 = graph.extract_scalar(bias_correction1); + const float bc2 = graph.extract_scalar(bias_correction2); + VK_CHECK_COND(bc1 != 0.0f); + VK_CHECK_COND(bc2 != 0.0f); + + // Split into <=16-byte push-constant entries (PushConstantData.h limit). + const std::array pc0 = { + graph.extract_scalar(lr), + graph.extract_scalar(beta1), + graph.extract_scalar(beta2), + graph.extract_scalar(eps)}; + const std::array pc1 = { + graph.extract_scalar(weight_decay), bc1, bc2}; + + std::string kernel_name("adamw_step"); + add_storage_type_suffix(kernel_name, graph.storage_type_of(param)); + add_dtype_suffix(kernel_name, graph.dtype_of(param)); + + graph.execute_nodes().emplace_back(new DynamicDispatchNode( + graph, + VK_KERNEL_FROM_STR(kernel_name), + default_pick_global_wg_size, + default_pick_local_wg_size, + // Inputs and Outputs + {{param, vkapi::kReadWrite}, + {m, vkapi::kReadWrite}, + {v, vkapi::kReadWrite}, + {grad, vkapi::kRead}}, + // Shader params buffers + {}, + // Push Constants + {graph.numel_pc_of(param), + PushConstantDataInfo(pc0.data(), sizeof(pc0)), + PushConstantDataInfo(pc1.data(), sizeof(pc1))}, + // Specialization Constants + {}, + // Resize Args + {}, + // Resizing Logic + nullptr)); +} + +void adamw_step(ComputeGraph& graph, const std::vector& args) { + int arg_idx = 0; + const ValueRef param = args[arg_idx++]; + const ValueRef m = args[arg_idx++]; + const ValueRef v = args[arg_idx++]; + const ValueRef grad = args[arg_idx++]; + const ValueRef lr = args[arg_idx++]; + const ValueRef beta1 = args[arg_idx++]; + const ValueRef beta2 = args[arg_idx++]; + const ValueRef eps = args[arg_idx++]; + const ValueRef weight_decay = args[arg_idx++]; + const ValueRef bias_correction1 = args[arg_idx++]; + const ValueRef bias_correction2 = args[arg_idx++]; + add_adamw_step_node( + graph, + param, + m, + v, + grad, + lr, + beta1, + beta2, + eps, + weight_decay, + bias_correction1, + bias_correction2); +} + +REGISTER_OPERATORS { + VK_REGISTER_OP(et_vk.adamw_step.default, adamw_step); +} + +} // namespace vkcompute