diff --git a/backends/vulkan/test/op_tests/CMakeLists.txt b/backends/vulkan/test/op_tests/CMakeLists.txt index 5e8991f8e50..47170373185 100644 --- a/backends/vulkan/test/op_tests/CMakeLists.txt +++ b/backends/vulkan/test/op_tests/CMakeLists.txt @@ -116,6 +116,9 @@ if(TARGET vulkan_backend AND LIB_TORCH) quantized_linear_test ${CMAKE_CURRENT_SOURCE_DIR}/quantized_linear_test.cpp test_utils ) + vulkan_op_test( + adamw_step_test ${CMAKE_CURRENT_SOURCE_DIR}/adamw_step_test.cpp test_utils + ) # Only build generated op tests if a path to tags.yaml and # native_functions.yaml is provided. These files are required for codegen. diff --git a/backends/vulkan/test/op_tests/adamw_step_test.cpp b/backends/vulkan/test/op_tests/adamw_step_test.cpp new file mode 100644 index 00000000000..6affd44b159 --- /dev/null +++ b/backends/vulkan/test/op_tests/adamw_step_test.cpp @@ -0,0 +1,191 @@ +/* + * 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 "test_utils.h" + +#include +#include + +// +// Reference implementation: an ATen transcription of the CPU-eager +// adamw_step_impl golden in backends/vulkan/custom_ops_lib.py. +// + +std::tuple adamw_reference_impl( + const at::Tensor& param_in, + const at::Tensor& m_in, + const at::Tensor& v_in, + const at::Tensor& grad, + const double lr, + const double beta1, + const double beta2, + const double eps, + const double weight_decay, + const double bias_correction1, + const double bias_correction2) { + at::Tensor param = param_in.clone(); + at::Tensor m = m_in.clone(); + at::Tensor v = v_in.clone(); + + param.mul_(1.0 - lr * weight_decay); + m.mul_(beta1).add_(grad, 1.0 - beta1); + v.mul_(beta2).addcmul_(grad, grad, 1.0 - beta2); + at::Tensor mhat = m / bias_correction1; + at::Tensor denom = (v / bias_correction2).sqrt() + eps; + param.addcdiv_(mhat, denom, -lr); + + return std::make_tuple(param, m, v); +} + +// +// Test harness +// + +void test_vulkan_adamw_step( + const std::vector& sizes, + const double lr, + const double beta1, + const double beta2, + const double eps, + const double weight_decay, + const int step) { + torch::manual_seed(0); + + const double bias_correction1 = 1.0 - std::pow(beta1, step); + const double bias_correction2 = 1.0 - std::pow(beta2, step); + + const auto options = at::device(at::kCPU).dtype(at::kFloat); + at::Tensor param = at::randn(sizes, options); + at::Tensor m = at::randn(sizes, options); + at::Tensor v = at::rand(sizes, options); // second moment is non-negative + at::Tensor grad = at::randn(sizes, options); + + at::Tensor ref_param; + at::Tensor ref_m; + at::Tensor ref_v; + std::tie(ref_param, ref_m, ref_v) = adamw_reference_impl( + param, + m, + v, + grad, + lr, + beta1, + beta2, + eps, + weight_decay, + bias_correction1, + bias_correction2); + + using namespace vkcompute; + + GraphConfig config; + ComputeGraph graph(config); + + ValueRef r_param = graph.add_tensor(sizes, vkapi::kFloat, utils::kBuffer); + ValueRef r_m = graph.add_tensor(sizes, vkapi::kFloat, utils::kBuffer); + ValueRef r_v = graph.add_tensor(sizes, vkapi::kFloat, utils::kBuffer); + ValueRef r_grad = graph.add_tensor(sizes, vkapi::kFloat, utils::kBuffer); + + ValueRef staging_param = graph.set_input_tensor(r_param); + ValueRef staging_m = graph.set_input_tensor(r_m); + ValueRef staging_v = graph.set_input_tensor(r_v); + ValueRef staging_grad = graph.set_input_tensor(r_grad); + + VK_GET_OP_FN("et_vk.adamw_step.default") + (graph, + { + r_param, + r_m, + r_v, + r_grad, + graph.add_scalar(lr), + graph.add_scalar(beta1), + graph.add_scalar(beta2), + graph.add_scalar(eps), + graph.add_scalar(weight_decay), + graph.add_scalar(bias_correction1), + graph.add_scalar(bias_correction2), + }); + + ValueRef staging_param_out = graph.set_output_tensor(r_param); + ValueRef staging_m_out = graph.set_output_tensor(r_m); + ValueRef staging_v_out = graph.set_output_tensor(r_v); + + graph.prepare(); + graph.prepack(); + + graph.maybe_cast_and_copy_into_staging( + staging_param, param.const_data_ptr(), param.numel(), vkapi::kFloat); + graph.maybe_cast_and_copy_into_staging( + staging_m, m.const_data_ptr(), m.numel(), vkapi::kFloat); + graph.maybe_cast_and_copy_into_staging( + staging_v, v.const_data_ptr(), v.numel(), vkapi::kFloat); + graph.maybe_cast_and_copy_into_staging( + staging_grad, grad.const_data_ptr(), grad.numel(), vkapi::kFloat); + + graph.execute(); + + at::Tensor vk_param = at::empty_like(param); + at::Tensor vk_m = at::empty_like(m); + at::Tensor vk_v = at::empty_like(v); + + graph.maybe_cast_and_copy_from_staging( + staging_param_out, + vk_param.mutable_data_ptr(), + vk_param.numel(), + vkapi::kFloat); + graph.maybe_cast_and_copy_from_staging( + staging_m_out, vk_m.mutable_data_ptr(), vk_m.numel(), vkapi::kFloat); + graph.maybe_cast_and_copy_from_staging( + staging_v_out, vk_v.mutable_data_ptr(), vk_v.numel(), vkapi::kFloat); + + ASSERT_TRUE(at::allclose(ref_param, vk_param, /*rtol=*/1e-4, /*atol=*/1e-4)); + ASSERT_TRUE(at::allclose(ref_m, vk_m, /*rtol=*/1e-4, /*atol=*/1e-4)); + ASSERT_TRUE(at::allclose(ref_v, vk_v, /*rtol=*/1e-4, /*atol=*/1e-4)); +} + +TEST(VulkanAdamwStepTest, test_adamw_step_no_weight_decay) { + test_vulkan_adamw_step( + /*sizes=*/{4, 8}, + /*lr=*/1e-3, + /*beta1=*/0.9, + /*beta2=*/0.999, + /*eps=*/1e-8, + /*weight_decay=*/0.0, + /*step=*/1); +} + +TEST(VulkanAdamwStepTest, test_adamw_step_with_weight_decay) { + test_vulkan_adamw_step( + /*sizes=*/{3, 5, 7}, + /*lr=*/1e-2, + /*beta1=*/0.9, + /*beta2=*/0.999, + /*eps=*/1e-8, + /*weight_decay=*/0.01, + /*step=*/10); +} + +TEST(VulkanAdamwStepTest, test_adamw_step_large_1d) { + test_vulkan_adamw_step( + /*sizes=*/{1000}, + /*lr=*/1e-3, + /*beta1=*/0.9, + /*beta2=*/0.999, + /*eps=*/1e-8, + /*weight_decay=*/0.05, + /*step=*/5); +} diff --git a/backends/vulkan/test/op_tests/targets.bzl b/backends/vulkan/test/op_tests/targets.bzl index 0a9b0743415..8aab725417a 100644 --- a/backends/vulkan/test/op_tests/targets.bzl +++ b/backends/vulkan/test/op_tests/targets.bzl @@ -180,6 +180,12 @@ def define_common_targets(is_fbcode = False): ":test_utils", ] ) + define_test_targets( + "adamw_step_test", + extra_deps = [ + ":test_utils", + ] + ) define_test_targets( "rms_norm_test", extra_deps = [