Use unique_ptr and shared_ptr properly for runner components (#10338)

Summary:
Pull Request resolved: #10338

The ownership of these components need some clarification.

* `Module` should be shared by `TextDecoderRunner` and potentially `TextPrefiller` (or `ImagePrefiller` in multimodal runner).
* `TextDecoderRunner` should be shared by the `TextPrefiller` and `TextTokenGenerator`.
* `Tokenizer` should be owned by the `Runner` as well as `TextTokenGenerator`.

Reviewed By: kirklandsign

Differential Revision: D73399600

Author: larryliu0820

examples/mediatek/executor_runner/mtk_llama_executor_runner.cpp

@@ -287,7 +287,7 @@ std::unique_ptr<Tokenizer> load_tokenizer() {
   if (FLAGS_tokenizer_type == "bpe") {
     tokenizer = std::make_unique<BPETokenizer>();
   } else if (FLAGS_tokenizer_type == "tiktoken") {
-    tokenizer = example::get_tiktoken_for_llama();
+    tokenizer = example::get_tiktoken_for_llama<decltype(tokenizer)>();
   }
   ET_CHECK_MSG(
       tokenizer, "Invalid tokenizer type: %s", FLAGS_tokenizer_type.c_str());

examples/mediatek/executor_runner/mtk_llama_runner.cpp

@@ -292,7 +292,7 @@ Error MTKLlamaRunner::inference(
 std::unique_ptr<Tokenizer> MTKLlamaRunner::load_tokenizer() {
   std::unique_ptr<Tokenizer> tokenizer;
   // Assumes that tokenizer type is Tiktoken
-  tokenizer = example::get_tiktoken_for_llama();
+  tokenizer = example::get_tiktoken_for_llama<decltype(tokenizer)>();
   tokenizer->load(modelpaths_.tokenizer_path);
   return tokenizer;
 }

examples/models/llama/runner/runner.cpp

@@ -54,10 +54,10 @@ Runner::Runner(
           {kUseSDPAWithKVCache, false},
       }) {
   if (data_path.has_value()) {
-    module_ = std::make_unique<Module>(
+    module_ = std::make_shared<Module>(
         model_path, data_path.value(), Module::LoadMode::File);
   } else {
-    module_ = std::make_unique<Module>(model_path, Module::LoadMode::File);
+    module_ = std::make_shared<Module>(model_path, Module::LoadMode::File);
   }
   ET_LOG(
       Info,
@@ -89,7 +89,7 @@ Error Runner::load() {
   ET_CHECK_OK_OR_RETURN_ERROR(module_->load_method("forward"));
   // load tokenizer. Assuming tiktoken is the default tokenizer
   tokenizer_ = nullptr;
-  tokenizer_ = get_tiktoken_for_llama();
+  tokenizer_ = get_tiktoken_for_llama<decltype(tokenizer_)>();
   ::tokenizers::Error err = tokenizer_->load(tokenizer_path_);
   // Rely on tiktoken to throw error if the artifact is incompatible. Then we
   // fallback to BPE tokenizer.
@@ -99,7 +99,7 @@ Error Runner::load() {
         "Failed to load %s as a Tiktoken artifact, trying BPE tokenizer",
         tokenizer_path_.c_str());
     tokenizer_.reset();
-    tokenizer_ = std::make_unique<::tokenizers::Llama2cTokenizer>();
+    tokenizer_ = std::make_shared<::tokenizers::Llama2cTokenizer>();
     err = tokenizer_->load(tokenizer_path_);
     ET_CHECK_TK_OK_OR_RETURN_ERROR(
         err,
@@ -143,20 +143,21 @@ Error Runner::load() {
     }
   }
   // @lint-ignore CLANGTIDY facebook-hte-Deprecated
-  text_decoder_runner_ = std::make_unique<llm::TextDecoderRunner>(
-      module_.get(), metadata_.at(kUseKVCache));
+  text_decoder_runner_ = std::make_shared<llm::TextDecoderRunner>(
+      module_, metadata_.at(kUseKVCache));
   text_prefiller_ = std::make_unique<llm::TextPrefiller>(
-      text_decoder_runner_.get(),
+      text_decoder_runner_,
       metadata_.at(kUseKVCache),
       metadata_.at(kEnableDynamicShape),
       metadata_.at(kMaxSeqLen));
 
+  stats_ = std::make_shared<llm::Stats>();
   text_token_generator_ = std::make_unique<llm::TextTokenGenerator>(
-      tokenizer_.get(),
-      text_decoder_runner_.get(),
+      tokenizer_,
+      text_decoder_runner_,
       metadata_.at(kUseKVCache),
       std::move(eos_ids),
-      &stats_);
+      stats_);
 
   return Error::Ok;
 }
@@ -178,9 +179,9 @@ Error Runner::generate(
   // Use ones-initialized inputs.
   ET_CHECK_MSG(!prompt.empty(), "Prompt cannot be null");
   if (!is_loaded()) {
-    stats_.model_load_start_ms = llm::time_in_ms();
+    stats_->model_load_start_ms = llm::time_in_ms();
     ET_CHECK_OK_OR_RETURN_ERROR(load());
-    stats_.model_load_end_ms = llm::time_in_ms();
+    stats_->model_load_end_ms = llm::time_in_ms();
   }
 
   if (config.warming) {
@@ -206,7 +207,7 @@ Error Runner::generate(
   // First token time only measures the time it takes to encode the prompt and
   // return a response token.
 
-  stats_.inference_start_ms = llm::time_in_ms();
+  stats_->inference_start_ms = llm::time_in_ms();
   shouldStop_ = false;
 
   ::tokenizers::Result<std::vector<uint64_t>> encode_res = tokenizer_->encode(
@@ -247,8 +248,8 @@ Error Runner::generate(
   auto prefill_res = text_prefiller_->prefill(prompt_tokens, pos);
   ET_CHECK_OK_OR_RETURN_ERROR(prefill_res.error());
   uint64_t cur_token = prefill_res.get();
-  stats_.first_token_ms = llm::time_in_ms();
-  stats_.prompt_eval_end_ms = llm::time_in_ms();
+  stats_->first_token_ms = llm::time_in_ms();
+  stats_->prompt_eval_end_ms = llm::time_in_ms();
 
   // print the first token from prefill. No prev_token so use cur_token for it.
   wrapped_callback(
@@ -269,7 +270,7 @@ Error Runner::generate(
       temperature_ == -1.0f ? config.temperature : temperature_,
       wrapped_callback));
 
-  stats_.inference_end_ms = llm::time_in_ms();
+  stats_->inference_end_ms = llm::time_in_ms();
   if (!config.warming) {
     printf("\n");
   }
@@ -282,17 +283,17 @@ Error Runner::generate(
     RUNNER_ET_LOG(config.warming, "Max new tokens %i reached!", max_new_tokens);
   }
 
-  stats_.num_prompt_tokens = num_prompt_tokens;
-  stats_.num_generated_tokens = num_generated_tokens;
+  stats_->num_prompt_tokens = num_prompt_tokens;
+  stats_->num_generated_tokens = num_generated_tokens;
 
   if (config.warming) {
     ET_LOG(Info, "Warmup run finished!");
   } else {
     // Do not print report during warmup
-    ::executorch::llm::print_report(stats_);
+    ::executorch::llm::print_report(*stats_);
   }
   if (stats_callback) {
-    stats_callback(stats_);
+    stats_callback(*stats_);
   }
 
   return Error::Ok;
@@ -307,7 +308,7 @@ Error Runner::warmup(const std::string& prompt, int32_t max_new_tokens) {
   Error err = generate(prompt, config);
 
   // Reset stats after warmup
-  stats_.reset();
+  stats_->reset();
   return err;
 }
 

examples/models/llama/runner/runner.h

@@ -60,18 +60,18 @@ class ET_EXPERIMENTAL Runner : public executorch::extension::llm::IRunner {
   bool shouldStop_{false};
 
   // model
-  std::unique_ptr<::executorch::extension::Module> module_;
+  std::shared_ptr<::executorch::extension::Module> module_;
   std::string tokenizer_path_;
-  std::unique_ptr<::tokenizers::Tokenizer> tokenizer_;
+  std::shared_ptr<::tokenizers::Tokenizer> tokenizer_;
   std::unordered_map<std::string, int64_t> metadata_;
-  std::unique_ptr<::executorch::extension::llm::TextDecoderRunner>
+  std::shared_ptr<::executorch::extension::llm::TextDecoderRunner>
       text_decoder_runner_;
   std::unique_ptr<::executorch::extension::llm::TextPrefiller> text_prefiller_;
   std::unique_ptr<::executorch::extension::llm::TextTokenGenerator>
       text_token_generator_;
 
   // stats
-  ::executorch::extension::llm::Stats stats_;
+  std::shared_ptr<::executorch::extension::llm::Stats> stats_;
 
   // temperature.
   // Deprecated, we should rely on the temperature in GenerationConfig instead.

examples/models/llama/tokenizer/llama_tiktoken.cpp

@@ -11,51 +11,85 @@
 namespace example {
 
 using ::tokenizers::Tiktoken;
+using ::tokenizers::Tokenizer;
 
 namespace {
 static constexpr int32_t kSpecialTokensSize = 256;
 static constexpr size_t kBOSTokenIndex = 0;
 static constexpr size_t kEOSTokenIndex = 1;
 
-static inline std::unique_ptr<std::vector<std::string>>
-_get_default_special_tokens() {
-  auto special_tokens =
-      std::make_unique<std::vector<std::string>>(std::vector<std::string>{
-          "<|begin_of_text|>",
-          "<|end_of_text|>",
-          "<|reserved_special_token_0|>",
-          "<|reserved_special_token_1|>",
-          "<|finetune_right_pad_id|>",
-          "<|step_id|>",
-          "<|start_header_id|>",
-          "<|end_header_id|>",
-          "<|eom_id|>",
-          "<|eot_id|>",
-          "<|python_tag|>"});
-  // pad the rest of the special tokens with reserved tokens
-  ssize_t reserved_special_token_num = 2;
-  while (special_tokens->size() < kSpecialTokensSize) {
-    special_tokens->emplace_back(
-        "<|reserved_special_token_" +
-        std::to_string(reserved_special_token_num++) + "|>");
+// Compile-time special tokens selection using templates
+template <Version V>
+struct SpecialTokensSelector {
+  static std::unique_ptr<std::vector<std::string>> get();
+};
+
+// Compile-time special tokens selection using templates
+template <>
+struct SpecialTokensSelector<Version::Default> {
+  static std::unique_ptr<std::vector<std::string>> get() {
+    auto special_tokens =
+        std::make_unique<std::vector<std::string>>(std::vector<std::string>{
+            "<|begin_of_text|>",
+            "<|end_of_text|>",
+            "<|reserved_special_token_0|>",
+            "<|reserved_special_token_1|>",
+            "<|finetune_right_pad_id|>",
+            "<|step_id|>",
+            "<|start_header_id|>",
+            "<|end_header_id|>",
+            "<|eom_id|>",
+            "<|eot_id|>",
+            "<|python_tag|>"});
+    // pad the rest of the special tokens with reserved tokens
+    ssize_t reserved_special_token_num = 2;
+    while (special_tokens->size() < kSpecialTokensSize) {
+      special_tokens->emplace_back(
+          "<|reserved_special_token_" +
+          std::to_string(reserved_special_token_num++) + "|>");
+    }
+    return special_tokens;
   }
-  return special_tokens;
-}
+};
 
-std::unique_ptr<std::vector<std::string>> _get_special_tokens(Version version) {
-  switch (version) {
-    case Version::Multimodal:
-      return get_multimodal_special_tokens();
-    default:
-      return _get_default_special_tokens();
+// Specialization for Multimodal version
+template <>
+struct SpecialTokensSelector<Version::Multimodal> {
+  static std::unique_ptr<std::vector<std::string>> get() {
+    return get_multimodal_special_tokens();
   }
-}
+};
 
 } // namespace
 
-std::unique_ptr<Tiktoken> get_tiktoken_for_llama(Version version) {
-  return std::make_unique<Tiktoken>(
-      _get_special_tokens(version), kBOSTokenIndex, kEOSTokenIndex);
+namespace detail {
+// Helper function to create a Tiktoken with the given version
+template <typename PtrType, Version V>
+PtrType create_tiktoken() {
+  std::unique_ptr<std::vector<std::string>> special_tokens =
+      example::SpecialTokensSelector<V>::get();
+  if constexpr (is_shared_ptr_of_tokenizer<PtrType>()) {
+    return std::make_shared<Tiktoken>(
+        std::move(special_tokens), kBOSTokenIndex, kEOSTokenIndex);
+  } else if constexpr (is_unique_ptr_of_tokenizer<PtrType>()) {
+    return std::make_unique<Tiktoken>(
+        std::move(special_tokens), kBOSTokenIndex, kEOSTokenIndex);
+  } else {
+    static_assert(
+        is_shared_ptr_of_tokenizer<PtrType>() ||
+            is_unique_ptr_of_tokenizer<PtrType>(),
+        "PtrType must be either std::shared_ptr<Tiktoken> or std::unique_ptr<Tiktoken>");
+    // This line is never reached due to the static_assert, but needed for
+    // compilation
+    return PtrType{};
+  }
+}
+} // namespace detail
+
+// Function that returns a shared_ptr
+template <typename PtrType, Version V>
+PtrType get_tiktoken_for_llama() {
+  return detail::create_tiktoken<PtrType, V>();
 }
 
 std::unique_ptr<std::vector<std::string>> get_multimodal_special_tokens() {
@@ -87,4 +121,29 @@ std::unique_ptr<std::vector<std::string>> get_multimodal_special_tokens() {
   return special_tokens;
 }
 
+// specialization
+
+template std::shared_ptr<Tiktoken>
+get_tiktoken_for_llama<std::shared_ptr<Tiktoken>, Version::Multimodal>();
+
+template std::unique_ptr<Tiktoken>
+get_tiktoken_for_llama<std::unique_ptr<Tiktoken>, Version::Multimodal>();
+
+template std::shared_ptr<Tiktoken>
+get_tiktoken_for_llama<std::shared_ptr<Tiktoken>, Version::Default>();
+
+template std::unique_ptr<Tiktoken>
+get_tiktoken_for_llama<std::unique_ptr<Tiktoken>, Version::Default>();
+
+template std::shared_ptr<Tokenizer>
+get_tiktoken_for_llama<std::shared_ptr<Tokenizer>, Version::Multimodal>();
+
+template std::unique_ptr<Tokenizer>
+get_tiktoken_for_llama<std::unique_ptr<Tokenizer>, Version::Multimodal>();
+
+template std::shared_ptr<Tokenizer>
+get_tiktoken_for_llama<std::shared_ptr<Tokenizer>, Version::Default>();
+
+template std::unique_ptr<Tokenizer>
+get_tiktoken_for_llama<std::unique_ptr<Tokenizer>, Version::Default>();
 } // namespace example

examples/models/llama/tokenizer/llama_tiktoken.h

@@ -17,9 +17,27 @@ enum class Version {
   Multimodal,
 };
 
-std::unique_ptr<::tokenizers::Tiktoken> get_tiktoken_for_llama(
-    Version version = Version::Default);
+// Type traits to check if a type is a shared_ptr or unique_ptr of Tokenizer or
+// a derived class
+template <typename T>
+struct is_shared_ptr_of_tokenizer : std::false_type {};
 
+template <typename T>
+struct is_shared_ptr_of_tokenizer<std::shared_ptr<T>>
+    : std::is_base_of<::tokenizers::Tokenizer, T> {};
+
+template <typename T>
+struct is_unique_ptr_of_tokenizer : std::false_type {};
+
+template <typename T>
+struct is_unique_ptr_of_tokenizer<std::unique_ptr<T>>
+    : std::is_base_of<::tokenizers::Tokenizer, T> {};
+
+// Template version that can return either shared_ptr or unique_ptr
+template <typename PtrType, Version V = Version::Default>
+PtrType get_tiktoken_for_llama();
+
+// For backward compatibility
 std::unique_ptr<std::vector<std::string>> get_multimodal_special_tokens();
 
 } // namespace example

examples/models/llava/runner/llava_image_prefiller.h

@@ -18,7 +18,7 @@ namespace example {
 class ET_EXPERIMENTAL LlavaImagePrefiller
     : public ::executorch::extension::llm::ImagePrefiller {
  public:
-  LlavaImagePrefiller(::executorch::extension::Module* module)
+  LlavaImagePrefiller(std::shared_ptr<::executorch::extension::Module> module)
       : ImagePrefiller(module){};
   /**
    * Prefill an LLM Module with the given image input.