Skip to content

Custom KV input names & types #3

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

Merged
Ryan Hill (RyanUnderhill) merged 1 commit into from
Jan 18, 2024
Merged

Custom KV input names & types #3

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
85 changes: 69 additions & 16 deletions src/config.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,68 @@

namespace Generators {

ONNXTensorElementDataType TranslateTensorType(std::string_view value) {
if (value == "float32")
return ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT;
if (value == "float16")
return ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT16;

throw std::runtime_error("Invalid tensor type: " + std::string(value));
}

struct Model_Element : JSON::Element {
Model_Element(Config::Model& model) : model_{model} {}

void OnString(std::string_view name, std::string_view value) override {
if (name == "type")
model_.type = value;
else if (name == "decoder")
model_.decoder = value;
else if (name == "encoder_decoder_init")
model_.encoder_decoder_init = value;
else if (name == "past_names_key")
model_.past_names_key = value;
else if (name == "past_names_value")
model_.past_names_value = value;
else if (name == "present_names_key")
model_.present_names_key = value;
else if (name == "present_names_value")
model_.present_names_value = value;
else if (name == "past_names")
model_.past_names = value;
else if (name == "present_names")
model_.present_names = value;
else if (name == "cross_past_names_key")
model_.cross_past_names_key = value;
else if (name == "cross_past_names_value")
model_.cross_past_names_value = value;
else if (name == "cross_present_names_key")
model_.cross_present_names_key = value;
else if (name == "cross_present_names_value")
model_.cross_present_names_value = value;
else if (name == "logits_type")
model_.logits_type = TranslateTensorType(value);
else if (name == "kv_type")
model_.kv_type = TranslateTensorType(value);
else
throw std::runtime_error("Unknown name: " + std::string(name));
}

void OnNumber(std::string_view name, double value) override {
if (name == "vocab_size")
model_.vocab_size = static_cast<int>(value);
else if (name == "hidden_size" || name == "n_embed")
model_.hidden_size = static_cast<int>(value);
else if (name == "num_attention_heads" || name == "num_heads" || name == "n_head")
model_.num_attention_heads = static_cast<int>(value);
else if (name == "num_hidden_layers" || name == "num_layers" || name == "n_layer")
model_.num_hidden_layers = static_cast<int>(value);
}

private:
Config::Model& model_;
};

struct Root_Element : JSON::Element {
Root_Element(Config& config) : config_{config} {}

Expand All @@ -12,12 +74,6 @@ struct Root_Element : JSON::Element {
config_.tokenizer_class = value;
else if (name == "prefix")
config_.prefix = value;
else if (name == "model_type")
config_.model_type = value;
else if (name == "ogai_model_decoder")
config_.model_decoder = value;
else if (name == "ogai_model_encoder_decoder_init")
config_.model_encoder_decoder_init = value;
}

void OnNumber(std::string_view name, double value) override {
Expand Down Expand Up @@ -50,24 +106,21 @@ struct Root_Element : JSON::Element {
config_.decoder_start_token_id = static_cast<int>(value);
else if (name == "sep_token_id")
config_.sep_token_id = static_cast<int>(value);

// Model Class Attributes
else if (name == "vocab_size")
config_.vocab_size = static_cast<int>(value);
else if (name == "hidden_size" || name == "n_embed")
config_.hidden_size = static_cast<int>(value);
else if (name == "num_attention_heads" || name == "num_heads" || name == "n_head")
config_.num_attention_heads = static_cast<int>(value);
else if (name == "num_hidden_layers" || name == "num_layers" || name == "n_layer")
config_.num_hidden_layers = static_cast<int>(value);
}

void OnBool(std::string_view name, bool value) override {
if (name == "early_stopping")
config_.early_stopping = value;
}

Element& OnObject(std::string_view name) {
if (name == "model")
return model_element_;
return Element::OnObject(name);
}

Config& config_;
Model_Element model_element_{config_.model};
};

struct RootObject_Element : JSON::Element {
Expand Down
32 changes: 24 additions & 8 deletions src/config.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -27,14 +27,30 @@ struct Config {
int decoder_start_token_id{}; // If an encoder-decoder model starts decoding with a different token than bos, the id of that token.
int sep_token_id{}; // The id of the separation token.

// Model Class Attributes
std::string model_decoder;
std::string model_encoder_decoder_init;
std::string model_type;
int vocab_size{};
int hidden_size{};
int num_attention_heads{};
int num_hidden_layers{};
struct Model {
std::string decoder;
std::string encoder_decoder_init;
std::string type;

ONNXTensorElementDataType logits_type{ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT}; // float16/float32 are the valid types
ONNXTensorElementDataType kv_type{ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT}; // float16/float32 are the valid types

int vocab_size{};
int hidden_size{};
int num_attention_heads{};
int num_hidden_layers{};

// KV_Cache names (will be a string with a %d in it, like "past_key_self_%d", "past_value_self_%d")
std::string past_names_key, past_names_value;
std::string present_names_key, present_names_value;

// KV_Cache_Combined names where the kv key/value are merged into one tensor
std::string past_names, present_names;

// Cross_Cache for models like whisper
std::string cross_past_names_key, cross_past_names_value;
std::string cross_present_names_key, cross_present_names_value;
} model;
};

} // namespace Generators
2 changes: 1 addition & 1 deletion src/generators.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -44,7 +44,7 @@ float Float16ToFloat32(uint16_t v) {
SearchParams::SearchParams(const Model& model)
: pad_token_id{model.config_->pad_token_id},
eos_token_id{model.config_->eos_token_id},
vocab_size{model.config_->vocab_size},
vocab_size{model.config_->model.vocab_size},
max_length{model.config_->max_length},
length_penalty{model.config_->length_penalty},
early_stopping{model.config_->early_stopping},
Expand Down
3 changes: 1 addition & 2 deletions src/models/gpt.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5,10 +5,9 @@ namespace Generators {

Gpt_Model::Gpt_Model(std::unique_ptr<Config> config, OrtEnv& ort_env, const ProviderOptions* provider_options)
: Model{std::move(config), ort_env, provider_options} {
session_decoder_ = OrtSession::Create(ort_env, (config_->config_path / config_->model_decoder).c_str(), session_options_.get());
session_decoder_ = OrtSession::Create(ort_env, (config_->config_path / config_->model.decoder).c_str(), session_options_.get());

InitDeviceAllocator(*session_decoder_);
InitLogits(*session_decoder_->GetOutputTypeInfo(0));
}

std::unique_ptr<State> Gpt_Model::CreateState(RoamingArray<int32_t> sequence_lengths, const SearchParams& params) {
Expand Down
88 changes: 42 additions & 46 deletions src/models/kv_cache.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,21 +7,21 @@ namespace Generators {
KV_Cache_Combined::KV_Cache_Combined(Model& model, State& state)
: model_{model},
state_{state},
layer_count_{model.config_->num_hidden_layers},
shape_{2, state_.search_params_.batch_size * state_.search_params_.num_beams, model.config_->num_attention_heads, 0, model.config_->hidden_size},
empty_past_{OrtValue::CreateTensor(*model_.allocator_device_, shape_, model_.score_type_)} {
layer_count_{model.config_->model.num_hidden_layers},
shape_{2, state_.search_params_.batch_size * state_.search_params_.num_beams, model.config_->model.num_attention_heads, 0, model.config_->model.hidden_size},
empty_past_{OrtValue::CreateTensor(*model_.allocator_device_, shape_, model_.config_->model.kv_type)} {
pasts_.resize(layer_count_);
presents_.reserve(layer_count_);

shape_[3] = state_.search_params_.sequence_length;
for (int i = 0; i < layer_count_; ++i) {
presents_.push_back(OrtValue::CreateTensor(*model.allocator_device_, shape_, model_.score_type_));
presents_.push_back(OrtValue::CreateTensor(*model.allocator_device_, shape_, model_.config_->model.kv_type));

char string[32];
snprintf(string, std::size(string), past_name_, i);
char string[64];
snprintf(string, std::size(string), model.config_->model.past_names.c_str(), i);
input_name_strings_.push_back(string);

snprintf(string, std::size(string), present_name_, i);
snprintf(string, std::size(string), model.config_->model.present_names.c_str(), i);
output_name_strings_.push_back(string);
}
}
Expand Down Expand Up @@ -50,7 +50,7 @@ void KV_Cache_Combined::Update(std::span<const int32_t> beam_indices, int curren

shape_[3] = current_length;
for (int i = 0; i < layer_count_; i++) {
presents_[i] = OrtValue::CreateTensor(*model_.allocator_device_, shape_, model_.score_type_);
presents_[i] = OrtValue::CreateTensor(*model_.allocator_device_, shape_, model_.config_->model.kv_type);
state_.inputs_[input_index_ + i] = pasts_[i].get();
state_.outputs_[output_index_ + i] = presents_[i].get();
}
Expand Down Expand Up @@ -99,39 +99,37 @@ void KV_Cache_Combined::PickPastState(std::span<const int32_t> beam_indices, int
}

void KV_Cache_Combined::PickPastState(std::span<const int32_t> beam_indices, int index) {
if (model_.score_type_ == Ort::TypeToTensorType<float>::type)
if (model_.config_->model.kv_type == Ort::TypeToTensorType<float>::type)
PickPastState<float>(beam_indices, index);
else
PickPastState<Ort::Float16_t>(beam_indices, index);
}

KV_Cache::KV_Cache(Model& model, State& state,
std::span<const char*> past_names, std::span<const char*> present_names)
KV_Cache::KV_Cache(Model& model, State& state)
: model_{model},
state_{state},
layer_count_{model_.config_->num_hidden_layers},
past_names_{past_names},
present_names_{present_names},
shape_{state_.search_params_.batch_size * state_.search_params_.num_beams, model.config_->num_attention_heads, 0, model.config_->hidden_size},
empty_past_{OrtValue::CreateTensor(*model_.allocator_device_, shape_, model_.score_type_)} {
layer_count_{model_.config_->model.num_hidden_layers},
shape_{state_.search_params_.batch_size * state_.search_params_.num_beams, model.config_->model.num_attention_heads, 0, model.config_->model.hidden_size},
empty_past_{OrtValue::CreateTensor(*model_.allocator_device_, shape_, model_.config_->model.kv_type)} {
pasts_.resize(layer_count_ * 2);
presents_.reserve(layer_count_ * 2);

shape_[2] = state_.search_params_.sequence_length; // Set this after empty_past_ has been created with 0 for this field

for (int i = 0; i < layer_count_; ++i) {
presents_.push_back(OrtValue::CreateTensor(*model_.allocator_device_, shape_, model_.score_type_));
presents_.push_back(OrtValue::CreateTensor(*model_.allocator_device_, shape_, model_.score_type_));
presents_.push_back(OrtValue::CreateTensor(*model_.allocator_device_, shape_, model_.config_->model.kv_type));
presents_.push_back(OrtValue::CreateTensor(*model_.allocator_device_, shape_, model_.config_->model.kv_type));

char string[32];
for (auto* name : past_names_) {
snprintf(string, std::size(string), name, i);
input_name_strings_.push_back(string);
}
for (auto* name : present_names_) {
snprintf(string, std::size(string), name, i);
output_name_strings_.push_back(string);
}
char string[64];
snprintf(string, std::size(string), model.config_->model.past_names_key.c_str(), i);
input_name_strings_.push_back(string);
snprintf(string, std::size(string), model.config_->model.past_names_value.c_str(), i);
input_name_strings_.push_back(string);

snprintf(string, std::size(string), model.config_->model.present_names_key.c_str(), i);
output_name_strings_.push_back(string);
snprintf(string, std::size(string), model.config_->model.present_names_value.c_str(), i);
output_name_strings_.push_back(string);
}
}

Expand Down Expand Up @@ -167,7 +165,7 @@ void KV_Cache::Update(std::span<const int32_t> beam_indices, int current_length)

shape_[2] = current_length;
for (int i = 0; i < layer_count_ * 2; i++) {
presents_[i] = OrtValue::CreateTensor(*model_.allocator_device_, shape_, model_.score_type_);
presents_[i] = OrtValue::CreateTensor(*model_.allocator_device_, shape_, model_.config_->model.kv_type);
state_.outputs_[output_index_ + i] = presents_[i].get();
}
}
Expand Down Expand Up @@ -206,35 +204,33 @@ void KV_Cache::PickPastState(std::span<const int32_t> beam_indices, int index) {
}

void KV_Cache::PickPastState(std::span<const int32_t> beam_indices, int index) {
if (model_.score_type_ == Ort::TypeToTensorType<float>::type)
if (model_.config_->model.kv_type == Ort::TypeToTensorType<float>::type)
PickPastState<float>(beam_indices, index);
else
PickPastState<Ort::Float16_t>(beam_indices, index);
}

Cross_Cache::Cross_Cache(Model& model, State& state,
std::span<const char*> past_names, std::span<const char*> present_names)
Cross_Cache::Cross_Cache(Model& model, State& state)
: model_{model},
state_{state},
layer_count_{model_.config_->num_hidden_layers},
past_names_{past_names},
present_names_{present_names},
shape_{state_.search_params_.batch_size * state_.search_params_.num_beams, model.config_->num_attention_heads, 1500, model.config_->hidden_size} {
layer_count_{model_.config_->model.num_hidden_layers},
shape_{state_.search_params_.batch_size * state_.search_params_.num_beams, model.config_->model.num_attention_heads, 1500, model.config_->model.hidden_size} {
values_.reserve(layer_count_ * 2);

for (int i = 0; i < layer_count_; ++i) {
char string[32];
values_.push_back(OrtValue::CreateTensor(*model_.allocator_device_, shape_, model_.score_type_));
values_.push_back(OrtValue::CreateTensor(*model_.allocator_device_, shape_, model_.score_type_));
values_.push_back(OrtValue::CreateTensor(*model_.allocator_device_, shape_, model_.config_->model.kv_type));
values_.push_back(OrtValue::CreateTensor(*model_.allocator_device_, shape_, model_.config_->model.kv_type));

for (auto* name : past_names_) {
snprintf(string, std::size(string), name, i);
input_name_strings_.push_back(string);
}
for (auto* name : present_names_) {
snprintf(string, std::size(string), name, i);
output_name_strings_.push_back(string);
}
char string[64];
snprintf(string, std::size(string), model.config_->model.cross_past_names_key.c_str(), i);
input_name_strings_.push_back(string);
snprintf(string, std::size(string), model.config_->model.cross_past_names_value.c_str(), i);
input_name_strings_.push_back(string);

snprintf(string, std::size(string), model.config_->model.cross_present_names_key.c_str(), i);
output_name_strings_.push_back(string);
snprintf(string, std::size(string), model.config_->model.cross_present_names_value.c_str(), i);
output_name_strings_.push_back(string);
}
}

Expand Down
13 changes: 2 additions & 11 deletions src/models/kv_cache.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -13,10 +13,6 @@ struct KV_Cache_Combined {
void PickPastState(std::span<const int32_t> beam_indices, int index);
void PickPastState(std::span<const int32_t> beam_indices, int index);

// KV combined
const char *past_name_{"past_%d"};
const char *present_name_{"present_%d"};

private:

Model& model_;
Expand All @@ -32,7 +28,7 @@ struct KV_Cache_Combined {
};

struct KV_Cache {
KV_Cache(Model& model, State& state, std::span<const char*> past_names, std::span<const char*> present_names);
KV_Cache(Model& model, State& state);

void AddEncoder(); // If model has an initial encoder step, this is used
void Add();
Expand All @@ -47,9 +43,6 @@ struct KV_Cache {
int layer_count_;
size_t input_index_{~0U}, output_index_{~0U};

std::span<const char*> past_names_; // past key name/past value name
std::span<const char*> present_names_; // present key name/present value name

std::array<int64_t, 4> shape_;

std::unique_ptr<OrtValue> empty_past_;
Expand All @@ -59,7 +52,7 @@ struct KV_Cache {

// Very similar to the KV_Cache, but is only created once at the encoder step, then used without modification for every decoder step
struct Cross_Cache {
Cross_Cache(Model& model, State& state, std::span<const char*> past_names, std::span<const char*> present_names);
Cross_Cache(Model& model, State& state);

void AddOutputs();
void AddInputs();
Expand All @@ -69,8 +62,6 @@ struct Cross_Cache {
State& state_;
int layer_count_;

std::span<const char*> past_names_, present_names_;

std::array<int64_t, 4> shape_;

std::vector<std::unique_ptr<OrtValue>> values_;
Expand Down
3 changes: 1 addition & 2 deletions src/models/llama.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5,10 +5,9 @@ namespace Generators {

Llama_Model::Llama_Model(std::unique_ptr<Config> config, OrtEnv& ort_env, const ProviderOptions* provider_options)
: Model{std::move(config), ort_env, provider_options} {
session_decoder_ = OrtSession::Create(ort_env, (config_->config_path / config_->model_decoder).c_str(), session_options_.get());
session_decoder_ = OrtSession::Create(ort_env, (config_->config_path / config_->model.decoder).c_str(), session_options_.get());

InitDeviceAllocator(*session_decoder_);
InitLogits(*session_decoder_->GetOutputTypeInfo(0));
}

std::unique_ptr<State> Llama_Model::CreateState(RoamingArray<int32_t> sequence_lengths, const SearchParams& params) {
Expand Down
Loading