Fabio
c36510c21e
Some checks failed
Linux Build / Linux cpu (push) Has been cancelled
Linux Build / Linux cpu-clang (push) Has been cancelled
Linux Build / Linux vulkan (push) Has been cancelled
mac-build / macOS CPU (push) Has been cancelled
Nix Build / Nix (cpu) (push) Has been cancelled
Nix Build / Nix (vulkan) (push) Has been cancelled
windows-build / Windows CPU (push) Has been cancelled
359 lines
16 KiB
C++
359 lines
16 KiB
C++
#include "test_assert.h"
|
|
|
|
#include "engine/framework/assets/tensor_source.h"
|
|
#include "engine/framework/modules/speech_encoders/wav2vec2_encoder.h"
|
|
|
|
#include <ggml-backend.h>
|
|
#include <ggml.h>
|
|
|
|
#include <cmath>
|
|
#include <cstdint>
|
|
#include <cstdio>
|
|
#include <cstring>
|
|
#include <filesystem>
|
|
#include <iostream>
|
|
#include <memory>
|
|
#include <optional>
|
|
#include <stdexcept>
|
|
#include <string>
|
|
#include <unordered_map>
|
|
#include <vector>
|
|
|
|
namespace {
|
|
|
|
using engine::test::require;
|
|
using engine::test::require_close;
|
|
|
|
constexpr int64_t kHidden = 8;
|
|
constexpr int64_t kIntermediate = 16;
|
|
constexpr int64_t kHeads = 2;
|
|
constexpr int64_t kPosConvGroups = 2;
|
|
constexpr int64_t kPosConvKernel = 5;
|
|
|
|
// Deterministic per-index values, distinct from zeros.
|
|
float fill_value(size_t index) {
|
|
return static_cast<float>((index * 7 + 3) % 11) / 5.0F - 0.4F;
|
|
}
|
|
|
|
class FakeTensorSource final : public engine::assets::TensorSource {
|
|
public:
|
|
struct Tensor {
|
|
std::vector<int64_t> shape;
|
|
std::vector<float> values;
|
|
};
|
|
|
|
void add(std::string name, std::vector<int64_t> shape, std::vector<float> values) {
|
|
tensors_.emplace(std::move(name), Tensor{std::move(shape), std::move(values)});
|
|
}
|
|
|
|
const std::filesystem::path & source_path() const noexcept override {
|
|
static const std::filesystem::path kPath = std::filesystem::path("fake");
|
|
return kPath;
|
|
}
|
|
bool has_tensor(std::string_view name) const noexcept override {
|
|
return tensors_.find(std::string(name)) != tensors_.end();
|
|
}
|
|
engine::assets::TensorMetadata require_metadata(std::string_view name) const override {
|
|
return {std::string(name), "f32", require(name)->shape};
|
|
}
|
|
std::vector<engine::assets::TensorMetadata> tensors() const override {
|
|
std::vector<engine::assets::TensorMetadata> out;
|
|
for (const auto & [name, tensor] : tensors_) {
|
|
out.push_back({name, "f32", tensor.shape});
|
|
}
|
|
return out;
|
|
}
|
|
void release_storage() const override {}
|
|
engine::assets::RawTensorData require_tensor_data(std::string_view name) const override {
|
|
const auto & tensor = *require(name);
|
|
engine::assets::RawTensorData raw;
|
|
raw.metadata = {std::string(name), "f32", tensor.shape};
|
|
raw.bytes.resize(tensor.values.size() * sizeof(float));
|
|
std::memcpy(raw.bytes.data(), tensor.values.data(), raw.bytes.size());
|
|
return raw;
|
|
}
|
|
std::vector<float> require_f32(
|
|
std::string_view name,
|
|
const std::optional<std::vector<int64_t>> & expected_shape) const override {
|
|
const auto & tensor = *require(name);
|
|
if (expected_shape.has_value() && *expected_shape != tensor.shape) {
|
|
throw std::runtime_error("shape mismatch for " + std::string(name));
|
|
}
|
|
return tensor.values;
|
|
}
|
|
std::optional<std::vector<float>> optional_f32(
|
|
std::string_view name,
|
|
const std::optional<std::vector<int64_t>> & expected_shape) const override {
|
|
if (!has_tensor(name)) {
|
|
return std::nullopt;
|
|
}
|
|
return require_f32(name, expected_shape);
|
|
}
|
|
int64_t require_i64_scalar(std::string_view name) const override {
|
|
throw std::runtime_error("unexpected i64 scalar request: " + std::string(name));
|
|
}
|
|
|
|
private:
|
|
const Tensor * require(std::string_view name) const {
|
|
const auto it = tensors_.find(std::string(name));
|
|
if (it == tensors_.end()) {
|
|
throw std::runtime_error("missing tensor: " + std::string(name));
|
|
}
|
|
return &it->second;
|
|
}
|
|
|
|
std::unordered_map<std::string, Tensor> tensors_;
|
|
};
|
|
|
|
engine::modules::Wav2Vec2EncoderConfig tiny_mms_config() {
|
|
engine::modules::Wav2Vec2EncoderConfig mms_config;
|
|
mms_config.hidden_size = kHidden;
|
|
mms_config.intermediate_size = kIntermediate;
|
|
mms_config.num_hidden_layers = 1;
|
|
mms_config.output_hidden_layer = 1;
|
|
mms_config.num_attention_heads = kHeads;
|
|
mms_config.conv_dim = {kHidden};
|
|
mms_config.conv_kernel = {3};
|
|
mms_config.conv_stride = {2};
|
|
mms_config.layer_norm_eps = 1.0e-5F;
|
|
mms_config.num_conv_pos_embeddings = kPosConvKernel;
|
|
mms_config.num_conv_pos_embedding_groups = kPosConvGroups;
|
|
mms_config.feature_extractor_norm = engine::modules::Wav2Vec2FeatureExtractorNorm::LayerNormEveryLayer;
|
|
mms_config.encoder_layer_norm_order = engine::modules::Wav2Vec2EncoderLayerNormOrder::PreNorm;
|
|
return mms_config;
|
|
}
|
|
|
|
// Populates every non-positional tensor the tiny config needs. Returns the
|
|
// effective (out, in, kernel) positional-conv kernel computed in the same way
|
|
// the MMS encoder runtime folds weight-norm layouts.
|
|
std::vector<float> populate_common_tensors(FakeTensorSource & source) {
|
|
const auto fill = [](const std::vector<int64_t> & shape) {
|
|
size_t count = 1;
|
|
for (const int64_t dim : shape) {
|
|
count *= static_cast<size_t>(dim);
|
|
}
|
|
std::vector<float> values(count, 0.0F);
|
|
for (size_t i = 0; i < values.size(); ++i) {
|
|
values[i] = fill_value(i + 1);
|
|
}
|
|
return values;
|
|
};
|
|
// feature_extractor.conv_layers.0
|
|
source.add("feature_extractor.conv_layers.0.conv.weight", {kHidden, 1, 3}, fill({kHidden, 1, 3}));
|
|
source.add("feature_extractor.conv_layers.0.conv.bias", {kHidden}, std::vector<float>(static_cast<size_t>(kHidden), 0.1F));
|
|
source.add("feature_extractor.conv_layers.0.layer_norm.weight", {kHidden}, std::vector<float>(static_cast<size_t>(kHidden), 1.0F));
|
|
source.add("feature_extractor.conv_layers.0.layer_norm.bias", {kHidden}, std::vector<float>(static_cast<size_t>(kHidden), 0.0F));
|
|
// feature_projection
|
|
source.add("feature_projection.layer_norm.weight", {kHidden}, std::vector<float>(static_cast<size_t>(kHidden), 1.0F));
|
|
source.add("feature_projection.layer_norm.bias", {kHidden}, std::vector<float>(static_cast<size_t>(kHidden), 0.0F));
|
|
source.add("feature_projection.projection.weight", {kHidden, kHidden}, fill({kHidden, kHidden}));
|
|
source.add("feature_projection.projection.bias", {kHidden}, std::vector<float>(static_cast<size_t>(kHidden), 0.0F));
|
|
// encoder.layer_norm (apply_final_layer_norm)
|
|
source.add("encoder.layer_norm.weight", {kHidden}, std::vector<float>(static_cast<size_t>(kHidden), 1.0F));
|
|
source.add("encoder.layer_norm.bias", {kHidden}, std::vector<float>(static_cast<size_t>(kHidden), 0.0F));
|
|
// encoder.layers.0 attention
|
|
for (const char * proj : {"q_proj", "k_proj", "v_proj", "out_proj"}) {
|
|
const std::string prefix = std::string("encoder.layers.0.attention.") + proj;
|
|
source.add(prefix + ".weight", {kHidden, kHidden}, fill({kHidden, kHidden}));
|
|
source.add(prefix + ".bias", {kHidden}, std::vector<float>(static_cast<size_t>(kHidden), 0.0F));
|
|
}
|
|
// encoder.layers.0 PreNorm layer norm
|
|
source.add("encoder.layers.0.layer_norm.weight", {kHidden}, std::vector<float>(static_cast<size_t>(kHidden), 1.0F));
|
|
source.add("encoder.layers.0.layer_norm.bias", {kHidden}, std::vector<float>(static_cast<size_t>(kHidden), 0.0F));
|
|
// encoder.layers.0 feed forward
|
|
source.add("encoder.layers.0.feed_forward.intermediate_dense.weight", {kIntermediate, kHidden}, fill({kIntermediate, kHidden}));
|
|
source.add("encoder.layers.0.feed_forward.intermediate_dense.bias", {kIntermediate}, std::vector<float>(static_cast<size_t>(kIntermediate), 0.0F));
|
|
source.add("encoder.layers.0.feed_forward.output_dense.weight", {kHidden, kIntermediate}, fill({kHidden, kIntermediate}));
|
|
source.add("encoder.layers.0.feed_forward.output_dense.bias", {kHidden}, std::vector<float>(static_cast<size_t>(kHidden), 0.0F));
|
|
// encoder.layers.0 final layer norm
|
|
source.add("encoder.layers.0.final_layer_norm.weight", {kHidden}, std::vector<float>(static_cast<size_t>(kHidden), 1.0F));
|
|
source.add("encoder.layers.0.final_layer_norm.bias", {kHidden}, std::vector<float>(static_cast<size_t>(kHidden), 0.0F));
|
|
// positional-conv bias
|
|
source.add("encoder.pos_conv_embed.conv.bias", {kHidden}, std::vector<float>(static_cast<size_t>(kHidden), 0.0F));
|
|
|
|
// Effective kernel from fairseq-style g/v folding, mirroring the component.
|
|
std::vector<float> v(static_cast<size_t>(kHidden * (kHidden / kPosConvGroups) * kPosConvKernel));
|
|
for (size_t i = 0; i < v.size(); ++i) {
|
|
v[i] = fill_value(i + 1000);
|
|
}
|
|
std::vector<float> g(static_cast<size_t>(kPosConvKernel));
|
|
for (int64_t k = 0; k < kPosConvKernel; ++k) {
|
|
g[static_cast<size_t>(k)] = 0.5F + 0.1F * static_cast<float>(k);
|
|
}
|
|
std::vector<float> weight(v.size());
|
|
const int64_t out_channels = kHidden;
|
|
const int64_t in_channels = kHidden / kPosConvGroups;
|
|
for (int64_t k = 0; k < kPosConvKernel; ++k) {
|
|
double sum = 0.0;
|
|
for (int64_t out = 0; out < out_channels; ++out) {
|
|
for (int64_t in = 0; in < in_channels; ++in) {
|
|
const size_t index = static_cast<size_t>((out * in_channels + in) * kPosConvKernel + k);
|
|
sum += static_cast<double>(v[index]) * static_cast<double>(v[index]);
|
|
}
|
|
}
|
|
const double norm = std::sqrt(sum);
|
|
const float scale = static_cast<float>(static_cast<double>(g[static_cast<size_t>(k)]) / norm);
|
|
for (int64_t out = 0; out < out_channels; ++out) {
|
|
for (int64_t in = 0; in < in_channels; ++in) {
|
|
const size_t index = static_cast<size_t>((out * in_channels + in) * kPosConvKernel + k);
|
|
weight[index] = v[index] * scale;
|
|
}
|
|
}
|
|
}
|
|
return weight;
|
|
}
|
|
|
|
std::vector<float> read_backend_tensor(const engine::core::TensorValue & tensor) {
|
|
require(tensor.type == GGML_TYPE_F32, "expected f32 backend tensor");
|
|
const size_t count = static_cast<size_t>(ggml_nelements(tensor.tensor));
|
|
std::vector<float> out(count);
|
|
ggml_backend_tensor_get(tensor.tensor, out.data(), 0, count * sizeof(float));
|
|
return out;
|
|
}
|
|
|
|
std::vector<float> load_and_read_pos_conv(std::shared_ptr<const engine::assets::TensorSource> source) {
|
|
const engine::core::BackendConfig backend{engine::core::BackendType::Cpu, 0, 1};
|
|
(void) engine::core::init_backend(backend);
|
|
const auto mms_encoder = engine::modules::Wav2Vec2EncoderRuntime::load_from_tensor_source(
|
|
std::move(source),
|
|
backend,
|
|
tiny_mms_config());
|
|
const auto & tensors = mms_encoder.weights()->tensors;
|
|
const auto it = tensors.find("encoder.pos_conv_embed.conv.weight");
|
|
require(it != tensors.end(), "positional-conv logical tensor missing");
|
|
return read_backend_tensor(it->second);
|
|
}
|
|
|
|
void test_prefolded_layout() {
|
|
auto source = std::make_shared<FakeTensorSource>();
|
|
const auto expected = populate_common_tensors(*source);
|
|
source->add("encoder.pos_conv_embed.conv.weight", {kHidden, kHidden / kPosConvGroups, kPosConvKernel}, expected);
|
|
const auto actual = load_and_read_pos_conv(std::move(source));
|
|
require(actual.size() == expected.size(), "prefolded kernel size");
|
|
for (size_t i = 0; i < expected.size(); ++i) {
|
|
require_close(actual[i], expected[i], 1.0e-6F, "prefolded kernel element");
|
|
}
|
|
}
|
|
|
|
void test_legacy_weight_norm_layout() {
|
|
auto source = std::make_shared<FakeTensorSource>();
|
|
const auto expected = populate_common_tensors(*source);
|
|
const auto v = [&] {
|
|
std::vector<float> out(static_cast<size_t>(kHidden * (kHidden / kPosConvGroups) * kPosConvKernel));
|
|
for (size_t i = 0; i < out.size(); ++i) {
|
|
out[i] = fill_value(i + 1000);
|
|
}
|
|
return out;
|
|
}();
|
|
std::vector<float> g(static_cast<size_t>(kPosConvKernel));
|
|
for (int64_t k = 0; k < kPosConvKernel; ++k) {
|
|
g[static_cast<size_t>(k)] = 0.5F + 0.1F * static_cast<float>(k);
|
|
}
|
|
source->add("encoder.pos_conv_embed.conv.weight_g", {1, 1, kPosConvKernel}, g);
|
|
source->add("encoder.pos_conv_embed.conv.weight_v", {kHidden, kHidden / kPosConvGroups, kPosConvKernel}, v);
|
|
const auto actual = load_and_read_pos_conv(std::move(source));
|
|
require(actual.size() == expected.size(), "legacy kernel size");
|
|
for (size_t i = 0; i < expected.size(); ++i) {
|
|
require_close(actual[i], expected[i], 1.0e-6F, "legacy kernel element");
|
|
}
|
|
}
|
|
|
|
void test_parametrized_layout() {
|
|
auto source = std::make_shared<FakeTensorSource>();
|
|
const auto expected = populate_common_tensors(*source);
|
|
const auto v = [&] {
|
|
std::vector<float> out(static_cast<size_t>(kHidden * (kHidden / kPosConvGroups) * kPosConvKernel));
|
|
for (size_t i = 0; i < out.size(); ++i) {
|
|
out[i] = fill_value(i + 1000);
|
|
}
|
|
return out;
|
|
}();
|
|
std::vector<float> g(static_cast<size_t>(kPosConvKernel));
|
|
for (int64_t k = 0; k < kPosConvKernel; ++k) {
|
|
g[static_cast<size_t>(k)] = 0.5F + 0.1F * static_cast<float>(k);
|
|
}
|
|
source->add(
|
|
"encoder.pos_conv_embed.conv.parametrizations.weight.original0",
|
|
{1, 1, kPosConvKernel},
|
|
g);
|
|
source->add(
|
|
"encoder.pos_conv_embed.conv.parametrizations.weight.original1",
|
|
{kHidden, kHidden / kPosConvGroups, kPosConvKernel},
|
|
v);
|
|
const auto actual = load_and_read_pos_conv(std::move(source));
|
|
require(actual.size() == expected.size(), "parametrized kernel size");
|
|
for (size_t i = 0; i < expected.size(); ++i) {
|
|
require_close(actual[i], expected[i], 1.0e-6F, "parametrized kernel element");
|
|
}
|
|
}
|
|
|
|
bool load_rejected(std::shared_ptr<const engine::assets::TensorSource> source) {
|
|
const engine::core::BackendConfig backend{engine::core::BackendType::Cpu, 0, 1};
|
|
(void) engine::core::init_backend(backend);
|
|
try {
|
|
(void) engine::modules::Wav2Vec2EncoderRuntime::load_from_tensor_source(
|
|
std::move(source),
|
|
backend,
|
|
tiny_mms_config());
|
|
} catch (const std::exception &) {
|
|
return true;
|
|
}
|
|
return false;
|
|
}
|
|
|
|
void test_missing_scale_pair_rejected() {
|
|
auto source = std::make_shared<FakeTensorSource>();
|
|
(void) populate_common_tensors(*source);
|
|
std::vector<float> g(static_cast<size_t>(kPosConvKernel), 1.0F);
|
|
source->add("encoder.pos_conv_embed.conv.weight_g", {1, 1, kPosConvKernel}, g);
|
|
require(load_rejected(std::move(source)), "half-pair must be rejected");
|
|
}
|
|
|
|
void test_wrong_scale_shape_rejected() {
|
|
auto source = std::make_shared<FakeTensorSource>();
|
|
(void) populate_common_tensors(*source);
|
|
std::vector<float> g(4, 1.0F);
|
|
source->add("encoder.pos_conv_embed.conv.weight_g", {1, 1, 4}, g);
|
|
std::vector<float> v(static_cast<size_t>(kHidden * (kHidden / kPosConvGroups) * kPosConvKernel), 1.0F);
|
|
source->add("encoder.pos_conv_embed.conv.weight_v", {kHidden, kHidden / kPosConvGroups, kPosConvKernel}, v);
|
|
require(load_rejected(std::move(source)), "wrong scale shape must be rejected");
|
|
}
|
|
|
|
void test_zero_norm_rejected() {
|
|
auto source = std::make_shared<FakeTensorSource>();
|
|
(void) populate_common_tensors(*source);
|
|
std::vector<float> g(static_cast<size_t>(kPosConvKernel), 1.0F);
|
|
source->add("encoder.pos_conv_embed.conv.weight_g", {1, 1, kPosConvKernel}, g);
|
|
std::vector<float> v(static_cast<size_t>(kHidden * (kHidden / kPosConvGroups) * kPosConvKernel), 0.0F);
|
|
source->add("encoder.pos_conv_embed.conv.weight_v", {kHidden, kHidden / kPosConvGroups, kPosConvKernel}, v);
|
|
require(load_rejected(std::move(source)), "zero norm must be rejected");
|
|
}
|
|
|
|
void test_prefolded_shape_mismatch_rejected() {
|
|
auto source = std::make_shared<FakeTensorSource>();
|
|
const auto expected = populate_common_tensors(*source);
|
|
source->add(
|
|
"encoder.pos_conv_embed.conv.weight",
|
|
{kHidden, kHidden / kPosConvGroups, kPosConvKernel - 1},
|
|
std::vector<float>(expected.size() - static_cast<size_t>(kHidden * (kHidden / kPosConvGroups)), 0.5F));
|
|
require(load_rejected(std::move(source)), "prefolded shape mismatch must be rejected");
|
|
}
|
|
|
|
} // namespace
|
|
|
|
int main() {
|
|
try {
|
|
test_prefolded_layout();
|
|
test_legacy_weight_norm_layout();
|
|
test_parametrized_layout();
|
|
test_missing_scale_pair_rejected();
|
|
test_wrong_scale_shape_rejected();
|
|
test_zero_norm_rejected();
|
|
test_prefolded_shape_mismatch_rejected();
|
|
std::cout << "mms_wav2vec2_positional_conv_test passed\n";
|
|
return 0;
|
|
} catch (const std::exception & error) {
|
|
std::fprintf(stderr, "mms_wav2vec2_positional_conv_test: %s\n", error.what());
|
|
return 1;
|
|
}
|
|
}
|