Skip to content

Commit 62809e7

Browse files
JCNTHJCNTH
committed
[ExecuTorch][Vulkan] Tests for et_vk.adamw_step
Pull Request resolved: #20938 Vulkan op-test golden for `et_vk.adamw_step` (ATen library-computed reference of the CPU-eager update; in-place param/m/v vs the kernel). 3 cases incl. non-zero weight-decay + a large 1-D tail. Wires `adamw_step_test` in `targets.bzl` + `CMakeLists.txt`. ghstack-source-id: 405030805 @exported-using-ghexport Differential Revision: [D111761775](https://our.internmc.facebook.com/intern/diff/D111761775/)
1 parent 21d5828 commit 62809e7

3 files changed

Lines changed: 200 additions & 0 deletions

File tree

‎backends/vulkan/test/op_tests/CMakeLists.txt‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -116,6 +116,9 @@ if(TARGET vulkan_backend AND LIB_TORCH)
116116
quantized_linear_test ${CMAKE_CURRENT_SOURCE_DIR}/quantized_linear_test.cpp
117117
test_utils
118118
)
119+
vulkan_op_test(
120+
adamw_step_test ${CMAKE_CURRENT_SOURCE_DIR}/adamw_step_test.cpp test_utils
121+
)
119122

120123
# Only build generated op tests if a path to tags.yaml and
121124
# native_functions.yaml is provided. These files are required for codegen.

‎backends/vulkan/test/op_tests/adamw_step_test.cpp‎

Lines changed: 191 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,191 @@
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 <gtest/gtest.h>
10+
11+
#include <ATen/ATen.h>
12+
13+
#include <executorch/backends/vulkan/runtime/api/api.h>
14+
#include <executorch/backends/vulkan/runtime/graph/ComputeGraph.h>
15+
#include <executorch/backends/vulkan/runtime/graph/ops/OperatorRegistry.h>
16+
17+
#include "test_utils.h"
18+
19+
#include <cmath>
20+
#include <tuple>
21+
22+
//
23+
// Reference implementation: an ATen transcription of the CPU-eager
24+
// adamw_step_impl golden in backends/vulkan/custom_ops_lib.py.
25+
//
26+
27+
std::tuple<at::Tensor, at::Tensor, at::Tensor> adamw_reference_impl(
28+
const at::Tensor& param_in,
29+
const at::Tensor& m_in,
30+
const at::Tensor& v_in,
31+
const at::Tensor& grad,
32+
const double lr,
33+
const double beta1,
34+
const double beta2,
35+
const double eps,
36+
const double weight_decay,
37+
const double bias_correction1,
38+
const double bias_correction2) {
39+
at::Tensor param = param_in.clone();
40+
at::Tensor m = m_in.clone();
41+
at::Tensor v = v_in.clone();
42+
43+
param.mul_(1.0 - lr * weight_decay);
44+
m.mul_(beta1).add_(grad, 1.0 - beta1);
45+
v.mul_(beta2).addcmul_(grad, grad, 1.0 - beta2);
46+
at::Tensor mhat = m / bias_correction1;
47+
at::Tensor denom = (v / bias_correction2).sqrt() + eps;
48+
param.addcdiv_(mhat, denom, -lr);
49+
50+
return std::make_tuple(param, m, v);
51+
}
52+
53+
//
54+
// Test harness
55+
//
56+
57+
void test_vulkan_adamw_step(
58+
const std::vector<int64_t>& sizes,
59+
const double lr,
60+
const double beta1,
61+
const double beta2,
62+
const double eps,
63+
const double weight_decay,
64+
const int step) {
65+
at::manual_seed(0);
66+
67+
const double bias_correction1 = 1.0 - std::pow(beta1, step);
68+
const double bias_correction2 = 1.0 - std::pow(beta2, step);
69+
70+
const auto options = at::device(at::kCPU).dtype(at::kFloat);
71+
at::Tensor param = at::randn(sizes, options);
72+
at::Tensor m = at::randn(sizes, options);
73+
at::Tensor v = at::rand(sizes, options); // second moment is non-negative
74+
at::Tensor grad = at::randn(sizes, options);
75+
76+
at::Tensor ref_param;
77+
at::Tensor ref_m;
78+
at::Tensor ref_v;
79+
std::tie(ref_param, ref_m, ref_v) = adamw_reference_impl(
80+
param,
81+
m,
82+
v,
83+
grad,
84+
lr,
85+
beta1,
86+
beta2,
87+
eps,
88+
weight_decay,
89+
bias_correction1,
90+
bias_correction2);
91+
92+
using namespace vkcompute;
93+
94+
GraphConfig config;
95+
ComputeGraph graph(config);
96+
97+
ValueRef r_param = graph.add_tensor(sizes, vkapi::kFloat, utils::kBuffer);
98+
ValueRef r_m = graph.add_tensor(sizes, vkapi::kFloat, utils::kBuffer);
99+
ValueRef r_v = graph.add_tensor(sizes, vkapi::kFloat, utils::kBuffer);
100+
ValueRef r_grad = graph.add_tensor(sizes, vkapi::kFloat, utils::kBuffer);
101+
102+
ValueRef staging_param = graph.set_input_tensor(r_param);
103+
ValueRef staging_m = graph.set_input_tensor(r_m);
104+
ValueRef staging_v = graph.set_input_tensor(r_v);
105+
ValueRef staging_grad = graph.set_input_tensor(r_grad);
106+
107+
VK_GET_OP_FN("et_vk.adamw_step.default")
108+
(graph,
109+
{
110+
r_param,
111+
r_m,
112+
r_v,
113+
r_grad,
114+
graph.add_scalar<double>(lr),
115+
graph.add_scalar<double>(beta1),
116+
graph.add_scalar<double>(beta2),
117+
graph.add_scalar<double>(eps),
118+
graph.add_scalar<double>(weight_decay),
119+
graph.add_scalar<double>(bias_correction1),
120+
graph.add_scalar<double>(bias_correction2),
121+
});
122+
123+
ValueRef staging_param_out = graph.set_output_tensor(r_param);
124+
ValueRef staging_m_out = graph.set_output_tensor(r_m);
125+
ValueRef staging_v_out = graph.set_output_tensor(r_v);
126+
127+
graph.prepare();
128+
graph.prepack();
129+
130+
graph.maybe_cast_and_copy_into_staging(
131+
staging_param, param.const_data_ptr(), param.numel(), vkapi::kFloat);
132+
graph.maybe_cast_and_copy_into_staging(
133+
staging_m, m.const_data_ptr(), m.numel(), vkapi::kFloat);
134+
graph.maybe_cast_and_copy_into_staging(
135+
staging_v, v.const_data_ptr(), v.numel(), vkapi::kFloat);
136+
graph.maybe_cast_and_copy_into_staging(
137+
staging_grad, grad.const_data_ptr(), grad.numel(), vkapi::kFloat);
138+
139+
graph.execute();
140+
141+
at::Tensor vk_param = at::empty_like(param);
142+
at::Tensor vk_m = at::empty_like(m);
143+
at::Tensor vk_v = at::empty_like(v);
144+
145+
graph.maybe_cast_and_copy_from_staging(
146+
staging_param_out,
147+
vk_param.mutable_data_ptr(),
148+
vk_param.numel(),
149+
vkapi::kFloat);
150+
graph.maybe_cast_and_copy_from_staging(
151+
staging_m_out, vk_m.mutable_data_ptr(), vk_m.numel(), vkapi::kFloat);
152+
graph.maybe_cast_and_copy_from_staging(
153+
staging_v_out, vk_v.mutable_data_ptr(), vk_v.numel(), vkapi::kFloat);
154+
155+
ASSERT_TRUE(at::allclose(ref_param, vk_param, /*rtol=*/1e-4, /*atol=*/1e-4));
156+
ASSERT_TRUE(at::allclose(ref_m, vk_m, /*rtol=*/1e-4, /*atol=*/1e-4));
157+
ASSERT_TRUE(at::allclose(ref_v, vk_v, /*rtol=*/1e-4, /*atol=*/1e-4));
158+
}
159+
160+
TEST(VulkanAdamwStepTest, test_adamw_step_no_weight_decay) {
161+
test_vulkan_adamw_step(
162+
/*sizes=*/{4, 8},
163+
/*lr=*/1e-3,
164+
/*beta1=*/0.9,
165+
/*beta2=*/0.999,
166+
/*eps=*/1e-8,
167+
/*weight_decay=*/0.0,
168+
/*step=*/1);
169+
}
170+
171+
TEST(VulkanAdamwStepTest, test_adamw_step_with_weight_decay) {
172+
test_vulkan_adamw_step(
173+
/*sizes=*/{3, 5, 7},
174+
/*lr=*/1e-2,
175+
/*beta1=*/0.9,
176+
/*beta2=*/0.999,
177+
/*eps=*/1e-8,
178+
/*weight_decay=*/0.01,
179+
/*step=*/10);
180+
}
181+
182+
TEST(VulkanAdamwStepTest, test_adamw_step_large_1d) {
183+
test_vulkan_adamw_step(
184+
/*sizes=*/{1000},
185+
/*lr=*/1e-3,
186+
/*beta1=*/0.9,
187+
/*beta2=*/0.999,
188+
/*eps=*/1e-8,
189+
/*weight_decay=*/0.05,
190+
/*step=*/5);
191+
}

‎backends/vulkan/test/op_tests/targets.bzl‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -180,6 +180,12 @@ def define_common_targets(is_fbcode = False):
180180
":test_utils",
181181
]
182182
)
183+
define_test_targets(
184+
"adamw_step_test",
185+
extra_deps = [
186+
":test_utils",
187+
]
188+
)
183189
define_test_targets(
184190
"rms_norm_test",
185191
extra_deps = [

0 commit comments

Comments
 (0)