Skip to content
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

Cherry-pick static LLM pipeline changes #654

Merged
merged 2 commits into from
Jul 22, 2024
Merged

Cherry-pick static LLM pipeline changes #654

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
9efde00
Static shape LLM pipeline out-of-the-box (#576)
TolyaTalamanov Jul 11, 2024
12b30af
Support chat conversation for StaticLLMPipeline (#580)
TolyaTalamanov Jul 17, 2024
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
2 changes: 1 addition & 1 deletion samples/cpp/chat_sample/chat_sample.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -10,7 +10,7 @@ int main(int argc, char* argv[]) try {
std::string prompt;
std::string model_path = argv[1];

std::string device = "CPU"; // GPU can be used as well
std::string device = "CPU"; // GPU, NPU can be used as well
ov::genai::LLMPipeline pipe(model_path, "CPU");

ov::genai::GenerationConfig config;
Expand Down
88 changes: 68 additions & 20 deletions src/cpp/src/llm_pipeline_static.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,8 @@
#include "text_callback_streamer.hpp"
#include "utils.hpp"

#include <openvino/pass/stateful_to_stateless.hpp>

namespace {

std::shared_ptr<ov::Model> add_slices_to_kvcache_inputs(const std::shared_ptr<ov::Model>& model) {
Expand Down Expand Up @@ -75,25 +77,42 @@ void reshape_to_static(std::shared_ptr<ov::Model> model,
model->reshape(new_shapes);
}

void fill_tensor(ov::Tensor tensor, int64_t fill_val) {
void fill_tensor(ov::Tensor tensor, int64_t fill_val, size_t offset = 0u) {
int64_t* tensor_data = tensor.data<int64_t>();
std::fill(tensor_data, tensor_data + tensor.get_size(), fill_val);
std::fill(tensor_data + offset, tensor_data + tensor.get_size(), fill_val);
}

void copy_with_left_offset(const ov::Tensor& orig, ov::Tensor& padded) {
const auto orig_size = orig.get_size();
const auto padded_size = padded.get_size();
const auto kLeftOffset = padded_size - orig_size;
void copy_with_offset(const ov::Tensor& orig, const int32_t offset, ov::Tensor& padded) {
int64_t* orig_data = orig.data<int64_t>();
int64_t* padded_data = padded.data<int64_t>();
std::copy(orig_data, orig_data + orig_size, padded_data + kLeftOffset);
std::copy(orig_data, orig_data + orig.get_size(), padded_data + offset);
}

ov::AnyMap extract_config_or_empty(const ov::AnyMap& config, const std::string& config_name) {
ov::AnyMap extract_config_or_default(const ov::AnyMap& config, const std::string& config_name) {
ov::AnyMap stage_cfg;
if (auto it = config.find(config_name); it != config.end()) {
const auto& map = it->second.as<std::map<std::string, std::string>>();
stage_cfg = { map.begin(), map.end() };
} else if (config_name == "PREFILL_CONFIG") {
std::map<std::string, std::string> prefill_config = {
{ "NPU_USE_NPUW", "YES" },
{ "NPUW_FOLD", "YES" },
{ "NPUW_DCOFF_TYPE", "f16" },
{ "NPUW_DCOFF_SCALE", "YES" },
{ "NPUW_ONLINE_AVOID", "P:RMSNorm/NPU" }
};
stage_cfg.insert(prefill_config.begin(), prefill_config.end());
} else if (config_name == "GENERATE_CONFIG") {
std::map<std::string, std::string> generate_config = {
{ "NPU_USE_NPUW", "YES" },
{ "NPUW_FOLD", "YES" },
{ "NPUW_DCOFF_TYPE", "f16" },
{ "NPUW_DCOFF_SCALE", "YES" },
{ "NPU_COMPILATION_MODE_PARAMS", "compute-layers-with-higher-precision=Sqrt,Power,ReduceMean,Add" },
{ "NPUW_PARALLEL_COMPILE", "YES" },
{ "NPUW_FUNCALL_ASYNC", "YES" }
};
stage_cfg.insert(generate_config.begin(), generate_config.end());
}
return stage_cfg;
}
Expand Down Expand Up @@ -126,7 +145,8 @@ StaticLLMPipeline::StaticLLMPipeline(
ov::Core core;
// (1) Read the template model - this will be kvcache model
auto kvcache_model = core.read_model(path / "openvino_model.xml");
// (2) TODO: Expose KV-cache input and output layers from kvcache model
// (2) Expose KV-cache input and output layers from kvcache model
ov::pass::StatefulToStateless().run_on_model(kvcache_model);
// (3) Clone the model - this will be prefill
auto prefill_model = kvcache_model->clone();
prefill_model->set_friendly_name(kvcache_model->get_friendly_name() + "_prefill");
Expand All @@ -140,10 +160,10 @@ StaticLLMPipeline::StaticLLMPipeline(
kvcache_model = add_slices_to_kvcache_inputs(kvcache_model);
// (6) Compile both model
m_prefill_request = core.compile_model(
prefill_model, device, extract_config_or_empty(config, "PREFILL_CONFIG")
prefill_model, device, extract_config_or_default(config, "PREFILL_CONFIG")
).create_infer_request();
m_kvcache_request = core.compile_model(
kvcache_model, device, extract_config_or_empty(config, "GENERATE_CONFIG")
kvcache_model, device, extract_config_or_default(config, "GENERATE_CONFIG")
).create_infer_request();
// (7) Initialize tensors
prepare_for_new_conversation();
Expand All @@ -156,6 +176,18 @@ StaticLLMPipeline::StaticLLMPipeline(
) : StaticLLMPipeline(path, path.string(), device, config) {
}

void StaticLLMPipeline::start_chat(const std::string& system_message) {
if (!system_message.empty()) {
m_history.push_back({{"role", "system"}, {"content", system_message}});
}
m_is_chat_conversation = true;
};

void StaticLLMPipeline::finish_chat() {
m_is_chat_conversation = false;
m_history.clear();
};

void StaticLLMPipeline::prepare_for_new_conversation() {
fill_tensor(m_prefill_request.get_tensor("input_ids"), m_tokenizer.get_pad_token_id());
fill_tensor(m_prefill_request.get_tensor("position_ids"), 0u);
Expand All @@ -175,9 +207,23 @@ DecodedResults StaticLLMPipeline::generate(
}

OPENVINO_ASSERT(std::holds_alternative<std::string>(inputs));
auto tokenized_input = m_tokenizer.encode(std::get<std::string>(inputs));
auto& prompt = std::get<std::string>(inputs);

if (m_is_chat_conversation) {
m_history.push_back({{"role", "user"}, {"content", prompt}});
constexpr bool add_generation_prompt = true;
prompt = m_tokenizer.apply_chat_template(m_history, add_generation_prompt);
}

auto tokenized_input = m_tokenizer.encode(prompt);
auto encoded_results = generate(tokenized_input, config, streamer);
return {m_tokenizer.decode(encoded_results.tokens), encoded_results.scores};
DecodedResults decoded_results = {m_tokenizer.decode(encoded_results.tokens), encoded_results.scores};

if (m_is_chat_conversation) {
auto answer = decoded_results.texts[0];
m_history.push_back({{"role", "assistant"}, {"content", answer}});
}
return decoded_results;
}

EncodedResults StaticLLMPipeline::generate(
Expand Down Expand Up @@ -222,22 +268,25 @@ EncodedResults StaticLLMPipeline::generate(
ov::genai::EncodedResults results;
// NB: Only batch=1 is supported now
results.scores.resize(1u);
results.scores[0] = 0u;
results.tokens.resize(1u);

// NB: Check if input prompt less than maximum size
// NB: Check if there is enough space in KV-cache to process input prompt
auto prompt_len = input_ids.get_size();
if (prompt_len > m_kvcache_desc.total_size) {
OPENVINO_THROW("Currently static pipeline only process up to " + std::to_string(m_kvcache_desc.total_size) + " tokens");
}

// NB: Reset tensors on every generate call - chat conversation isn't supported yet!
// NB: From the "generate" perspective, every call is treated as start of new conversation,
// but if continuation is needed, prompt contains information about the entire conversation.
prepare_for_new_conversation();

auto padded_input_ids = m_prefill_request.get_tensor("input_ids");
copy_with_left_offset(input_ids, padded_input_ids);
const size_t offset = padded_input_ids.get_size() - input_ids.get_size();
copy_with_offset(input_ids, offset, padded_input_ids);

auto padded_attention_mask = m_prefill_request.get_tensor("attention_mask");
copy_with_left_offset(attention_mask, padded_attention_mask);
fill_tensor(padded_attention_mask, 1u, offset);

auto padded_position_ids = m_prefill_request.get_tensor("position_ids");
auto* padded_pos_data = padded_position_ids.data<int64_t>();
Expand All @@ -248,13 +297,13 @@ EncodedResults StaticLLMPipeline::generate(
// NB: Now there are prompt_len tokens in KV-cache
m_kvcache_desc.num_stored_tokens += prompt_len;
int64_t last_token = utils::argmax(m_prefill_request.get_tensor("logits"), 0);
results.tokens[0].push_back(last_token);
if (streamer_ptr && streamer_ptr->put(last_token)) {
return results;
}

padded_attention_mask.copy_to(m_kvcache_request.get_tensor("attention_mask"));


// Inputs: input_ids, attention_mask, position_ids, ...
// Outputs: logits, ...
const auto kStartInputKVCacheLayers = 3u;
Expand Down Expand Up @@ -286,13 +335,12 @@ EncodedResults StaticLLMPipeline::generate(

last_token = utils::argmax(m_kvcache_request.get_tensor("logits"), 0);
results.tokens[0].push_back(last_token);
results.scores[0] = 0u;

if (streamer_ptr && streamer_ptr->put(last_token)) {
break;
}

if (last_token == m_generation_config.eos_token_id) {
if (last_token == config.eos_token_id && !config.ignore_eos) {
break;
}

Expand Down
12 changes: 5 additions & 7 deletions src/cpp/src/llm_pipeline_static.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -35,13 +35,8 @@ class StaticLLMPipeline final : public LLMPipelineImplBase {
StreamerVariant streamer
) override;

void start_chat(const std::string& system_message) override {
OPENVINO_THROW("Currently chat conversation mode isn't supported");
};
void finish_chat() override {
OPENVINO_THROW("Currently chat conversation mode isn't supported");
};

void start_chat(const std::string& system_message) override;
void finish_chat() override;
private:
void prepare_for_new_conversation();

Expand All @@ -54,6 +49,9 @@ class StaticLLMPipeline final : public LLMPipelineImplBase {
KVCacheDesc m_kvcache_desc;
ov::InferRequest m_kvcache_request;
ov::InferRequest m_prefill_request;

bool m_is_chat_conversation = false;
ChatHistory m_history;
};

} // namespace genai
Expand Down
Loading