Skip to content

[executorch][native-vulkan] Lower rms_norm without a weight on the Vulkan engine #25097

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
digantdesai wants to merge 1 commit into
base: gh/digantdesai/213/base
Choose a base branch
from
Open

[executorch][native-vulkan] Lower rms_norm without a weight on the Vulkan engine #25097

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
38 changes: 26 additions & 12 deletions backends/native/runtime/vulkan/passes/LowerRMSNorm.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -28,9 +28,27 @@
value.role == ValueRole::ConstantTensor;
}

// A None weight scales by one; a present weight must be a constant that matches
// the normalized shape and the input dtype.
bool is_valid_weight(
const Graph& graph,
const Argument& weight,
const std::vector<int64_t>& shape,
ScalarType dtype) {
if (weight.kind() == ArgKind::None) {
return true;
}
if (weight.kind() != ArgKind::Tensor || !valid(weight.as_tensor().id)) {
return false;
}
const Value& value = graph.value(weight.as_tensor().id);
return value.is_tensor() && is_constant(value) &&
value.tensor_meta().sizes == shape && value.tensor_meta().dtype == dtype;
}

} // namespace

size_t lower_rms_norms(Graph& graph) {

Check warning on line 51 in backends/native/runtime/vulkan/passes/LowerRMSNorm.cpp

View workflow job for this annotation

GitHub Actions / lintrunner

CPPCHECK unusedFunction 

The function 'lower_rms_norms' is never used.
CWE: CWE-561
size_t count = 0;
for (const NodeId node_id : graph.schedule) {
Node& node = graph.node(node_id);
Expand All @@ -46,37 +64,33 @@
if (input == nullptr || input->kind() != ArgKind::Tensor ||
normalized_shape == nullptr ||
normalized_shape->kind() != ArgKind::IntList || weight == nullptr ||
weight->kind() != ArgKind::Tensor || eps == nullptr ||
eps->kind() != ArgKind::Float || valid(eps->as_float().id) ||
!std::isfinite(eps->as_float().value) || eps->as_float().value < 0) {
eps == nullptr || eps->kind() != ArgKind::Float ||
valid(eps->as_float().id) || !std::isfinite(eps->as_float().value) ||
eps->as_float().value < 0) {
continue;
}

const ValueId input_id = input->as_tensor().id;
const ValueId weight_id = weight->as_tensor().id;
const ValueId output_id = node.outputs[0].value_id;
if (!valid(input_id) || !valid(weight_id)) {
if (!valid(input_id)) {
continue;
}
const Value& input_value = graph.value(input_id);
const Value& weight_value = graph.value(weight_id);
const Value& output_value = graph.value(output_id);
if (!input_value.is_tensor() || !weight_value.is_tensor() ||
!output_value.is_tensor() || !is_constant(weight_value)) {
if (!input_value.is_tensor() || !output_value.is_tensor()) {
continue;
}

const TensorMeta& input_meta = input_value.tensor_meta();
const TensorMeta& weight_meta = weight_value.tensor_meta();
const TensorMeta& output_meta = output_value.tensor_meta();
const IntListArg& shape = normalized_shape->as_int_list();
if (input_meta.sizes.empty() || !shape.ids.empty() ||
if (input_meta.sizes.empty() || input_meta.sizes.back() < 1 ||
!shape.ids.empty() ||
shape.values != std::vector<int64_t>{input_meta.sizes.back()} ||
weight_meta.sizes != shape.values ||
(input_meta.dtype != ScalarType::Float &&
input_meta.dtype != ScalarType::Half &&
input_meta.dtype != ScalarType::BFloat16) ||
weight_meta.dtype != input_meta.dtype ||
!is_valid_weight(graph, *weight, shape.values, input_meta.dtype) ||
output_meta.dtype != input_meta.dtype ||
output_meta.sizes != input_meta.sizes) {
continue;
Expand Down
58 changes: 53 additions & 5 deletions backends/native/test/runtime/vulkan/test_lower_rms_norm.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -22,14 +22,14 @@
namespace ptn::vulkan {
namespace {

Method make_rms_norm(ScalarType dtype = ScalarType::Float) {
Method make_rms_norm(ScalarType dtype = ScalarType::Float, int64_t width = 4) {
Method method;
method.name = "forward";
Graph& graph = method.graph;
graph.values = {
Value("input", dtype, {2, 4}),
Value("weight", dtype, {4}),
Value("output", dtype, {2, 4}),
Value("input", dtype, {2, width}),
Value("weight", dtype, {width}),
Value("output", dtype, {2, width}),
};
graph.value(1).role = ValueRole::Parameter;
graph.nodes = {
Expand All @@ -49,7 +49,7 @@ Method make_rms_norm(ScalarType dtype = ScalarType::Float) {
.inputs =
{
{.name = "input", .arg = TensorArg{0}},
{.name = "normalized_shape", .arg = IntListArg{{4}}},
{.name = "normalized_shape", .arg = IntListArg{{width}}},
{.name = "weight", .arg = TensorArg{1}},
{.name = "eps", .arg = FloatArg{1e-5}},
},
Expand Down Expand Up @@ -104,6 +104,20 @@ TEST(LowerRMSNormTest, RewritesBFloat16Norm) {
EXPECT_NO_THROW(validate_graph(method.graph));
}

TEST(LowerRMSNormTest, RewritesNormWithoutWeight) {
Method method = make_rms_norm();
method.graph.node(2).inputs[2].arg = NoneArg{};
method.graph.rebuild_def_use();

EXPECT_EQ(lower_rms_norms(method.graph), 1);

const Node& norm = method.graph.node(2);
EXPECT_EQ(norm.target, "torch.ops.et_vk.rms_norm.default");
ASSERT_EQ(norm.inputs.size(), 3);
EXPECT_EQ(norm.inputs[1].arg.kind(), ArgKind::None);
EXPECT_NO_THROW(validate_graph(method.graph));
}

TEST(LowerRMSNormTest, RejectsMismatchedNormalizedShape) {
Method method = make_rms_norm();
method.graph.node(2).inputs[1].arg = IntListArg{{8}};
Expand All @@ -112,5 +126,39 @@ TEST(LowerRMSNormTest, RejectsMismatchedNormalizedShape) {
EXPECT_EQ(method.graph.node(2).target, "torch.ops.aten.rms_norm.default");
}

TEST(LowerRMSNormTest, RejectsInvalidWeight) {
struct Case {
const char* name;
ValueRole role;
ScalarType dtype;
std::vector<int64_t> sizes;
};
const std::vector<Case> cases = {
{"intermediate", ValueRole::Intermediate, ScalarType::Float, {4}},
{"user input", ValueRole::UserInput, ScalarType::Float, {4}},
{"wrong shape", ValueRole::Parameter, ScalarType::Float, {8}},
{"wrong dtype", ValueRole::Parameter, ScalarType::Half, {4}},
};
for (const Case& c : cases) {
SCOPED_TRACE(c.name);
Method method = make_rms_norm();
method.graph.value(1) = Value("weight", c.dtype, c.sizes);
method.graph.value(1).role = c.role;
method.graph.rebuild_def_use();

EXPECT_EQ(lower_rms_norms(method.graph), 0);
EXPECT_EQ(method.graph.node(2).target, "torch.ops.aten.rms_norm.default");
}
}

TEST(LowerRMSNormTest, RejectsZeroWidthNorm) {
Method method = make_rms_norm(ScalarType::Float, /*width=*/0);
method.graph.node(2).inputs[2].arg = NoneArg{};
method.graph.rebuild_def_use();

EXPECT_EQ(lower_rms_norms(method.graph), 0);
EXPECT_EQ(method.graph.node(2).target, "torch.ops.aten.rms_norm.default");
}

} // namespace
} // namespace ptn::vulkan
50 changes: 41 additions & 9 deletions backends/native/test/runtime/vulkan/test_vulkan_engine.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2833,12 +2833,15 @@ std::vector<uint8_t> stored_floats(const std::vector<StoredFloats>& constants) {
return tensors;
}

// aten.rms_norm over the last dim of a [2, 4] input with a stored weight.
std::vector<uint8_t> make_rms_norm_program(fbs::ScalarType dtype) {
// aten.rms_norm over the last dim of a [2, 4] input, with a stored weight if
// `affine` and a None weight otherwise.
std::vector<uint8_t> make_rms_norm_program(
fbs::ScalarType dtype,
bool affine = true) {
flatbuffers::FlatBufferBuilder builder;
const auto meta = create_tensor_meta(builder, dtype, {2, 4});
const auto weight_meta = create_tensor_meta(builder, dtype, {4});
const std::vector<flatbuffers::Offset<fbs::TensorValue>> tensor_values = {
std::vector<flatbuffers::Offset<fbs::TensorValue>> tensor_values = {
fbs::CreateTensorValueDirect(builder, "input", meta),
fbs::CreateTensorValueDirect(builder, "weight", weight_meta),
fbs::CreateTensorValueDirect(builder, "normed", meta),
Expand All @@ -2856,7 +2859,10 @@ std::vector<uint8_t> make_rms_norm_program(fbs::ScalarType dtype) {
fbs::CreateNamedArgumentDirect(
builder, "normalized_shape", create_int_list_arg(builder, {4})),
fbs::CreateNamedArgumentDirect(
builder, "weight", create_tensor_arg(builder, "weight")),
builder,
"weight",
affine ? create_tensor_arg(builder, "weight")
: create_none_arg(builder)),
fbs::CreateNamedArgumentDirect(
builder,
"eps",
Expand All @@ -2866,7 +2872,7 @@ std::vector<uint8_t> make_rms_norm_program(fbs::ScalarType dtype) {
const std::vector<flatbuffers::Offset<fbs::NamedArgument>> output_inputs = {
fbs::CreateNamedArgumentDirect(
builder, "", create_tensor_arg(builder, "normed"))};
const std::vector<flatbuffers::Offset<fbs::Node>> nodes = {
std::vector<flatbuffers::Offset<fbs::Node>> nodes = {
fbs::CreateNodeDirect(
builder,
"input",
Expand All @@ -2891,16 +2897,22 @@ std::vector<uint8_t> make_rms_norm_program(fbs::ScalarType dtype) {
fbs::CreateNodeDirect(
builder, "output", fbs::OpKind::OUTPUT, "", &output_inputs),
};
std::vector<flatbuffers::Offset<fbs::NamedTensorRef>> constants = {
fbs::CreateNamedTensorRefDirect(
builder, "weight", "weight", weight_meta, fbs::InputKind::PARAMETER),
};
if (!affine) {
// Drop the weight value, its placeholder and its constant.
tensor_values.erase(tensor_values.begin() + 1);
nodes.erase(nodes.begin() + 1);
constants.clear();
}
const auto graph = fbs::CreateGraph(
builder,
builder.CreateVector(nodes),
create_strings(builder, {"input"}),
create_strings(builder, {"normed"}),
builder.CreateVector(tensor_values));
const std::vector<flatbuffers::Offset<fbs::NamedTensorRef>> constants = {
fbs::CreateNamedTensorRefDirect(
builder, "weight", "weight", weight_meta, fbs::InputKind::PARAMETER),
};
return finish_program(
builder,
graph,
Expand Down Expand Up @@ -3587,6 +3599,26 @@ TEST(VulkanEngineTest, ExecutesBFloat16RMSNormWithStoredWeight) {
0x3F00, 0xBFC0, 0xC000, 0xC040, 0x3F00, 0x3FC0, 0x4000, 0x4040}));
}

TEST(VulkanEngineTest, ExecutesBFloat16RMSNormWithoutWeight) {
CompiledProgram compiled = compile_program(
make_rms_norm_program(fbs::ScalarType::BFLOAT16, /*affine=*/false));
// bf16 bits: rows {2, -2, 2, -2} and {4, 4, -4, 4}, each with an RMS that
// normalizes it to +-1.
const std::vector<uint16_t> input = {
0x4000, 0xC000, 0x4000, 0xC000, 0x4080, 0x4080, 0xC080, 0x4080};
compiled.executable->set_input(
0, input.data(), input.size(), ScalarType::BFloat16);
compiled.executable->execute();
std::vector<uint16_t> output(8);
compiled.executable->get_output(
0, output.data(), output.size(), ScalarType::BFloat16);
// {1, -1, 1, -1} and {1, 1, -1, 1}.
EXPECT_EQ(
output,
(std::vector<uint16_t>{
0x3F80, 0xBF80, 0x3F80, 0xBF80, 0x3F80, 0x3F80, 0xBF80, 0x3F80}));
}

TEST(VulkanEngineTest, ExecutesMetadataOnlyAliasAfterAliasLowering) {
CompiledProgram compiled = compile_program(make_alias_program());
const std::vector<float> input = {1.0f, 2.0f, 3.0f, 4.0f};
Expand Down
Loading