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/models/llama/runner/runner.cpp

@@ -52,7 +52,9 @@ Runner::Runner(
           {kMaxContextLen, 128},
           {kUseKVCache, true},
           {kUseSDPAWithKVCache, false},
-      }) {
+      }),
+      // @lint-ignore CLANGTIDY facebook-hte-Deprecated
+      stats_(std::make_unique<llm::Stats>()) {
   if (data_path.has_value()) {
     module_ = std::make_unique<Module>(
         model_path, data_path.value(), Module::LoadMode::File);
@@ -99,6 +101,7 @@ Error Runner::load() {
         "Failed to load %s as a Tiktoken artifact, trying BPE tokenizer",
         tokenizer_path_.c_str());
     tokenizer_.reset();
+    // @lint-ignore CLANGTIDY facebook-hte-Deprecated
     tokenizer_ = std::make_unique<::tokenizers::Llama2cTokenizer>();
     err = tokenizer_->load(tokenizer_path_);
     ET_CHECK_TK_OK_OR_RETURN_ERROR(
@@ -156,7 +159,7 @@ Error Runner::load() {
       text_decoder_runner_.get(),
       metadata_.at(kUseKVCache),
       std::move(eos_ids),
-      &stats_);
+      stats_.get());
 
   return Error::Ok;
 }
@@ -178,9 +181,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 +209,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 +250,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 +272,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 +285,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 +310,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

@@ -71,7 +71,7 @@ class ET_EXPERIMENTAL Runner : public executorch::extension::llm::IRunner {
       text_token_generator_;
 
   // stats
-  ::executorch::extension::llm::Stats stats_;
+  std::unique_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,71 @@
 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>> create();
+};
+
+// Compile-time special tokens selection using templates
+template <>
+struct SpecialTokensSelector<Version::Default> {
+  static std::unique_ptr<std::vector<std::string>> create() {
+    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>> create() {
+    return get_multimodal_special_tokens();
   }
-}
+};
 
 } // namespace
 
-std::unique_ptr<Tiktoken> get_tiktoken_for_llama(Version version) {
+namespace detail {
+// Helper function to create a Tiktoken with the given version
+template <Version V>
+std::unique_ptr<Tiktoken> create_tiktoken() {
+  std::unique_ptr<std::vector<std::string>> special_tokens =
+      example::SpecialTokensSelector<V>::create();
   return std::make_unique<Tiktoken>(
-      _get_special_tokens(version), kBOSTokenIndex, kEOSTokenIndex);
+      std::move(special_tokens), kBOSTokenIndex, kEOSTokenIndex);
+}
+} // namespace detail
+
+template <Version V = Version::Default>
+std::unique_ptr<::tokenizers::Tiktoken> get_tiktoken_for_llama() {
+  return detail::create_tiktoken<V>();
 }
 
 std::unique_ptr<std::vector<std::string>> get_multimodal_special_tokens() {
@@ -87,4 +107,10 @@ std::unique_ptr<std::vector<std::string>> get_multimodal_special_tokens() {
   return special_tokens;
 }
 
+// specialization
+
+template std::unique_ptr<Tiktoken>
+get_tiktoken_for_llama<Version::Multimodal>();
+
+template std::unique_ptr<Tiktoken> get_tiktoken_for_llama<Version::Default>();
 } // namespace example

examples/models/llama/tokenizer/llama_tiktoken.h

@@ -17,9 +17,10 @@ enum class Version {
   Multimodal,
 };
 
-std::unique_ptr<::tokenizers::Tiktoken> get_tiktoken_for_llama(
-    Version version = Version::Default);
+template <Version V = Version::Default>
+std::unique_ptr<::tokenizers::Tiktoken> 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)
+  explicit LlavaImagePrefiller(::executorch::extension::Module* module)
       : ImagePrefiller(module){};
   /**
    * Prefill an LLM Module with the given image input.

examples/models/llava/runner/llava_runner.cpp

@@ -40,7 +40,7 @@ Error LlavaRunner::load() {
   if (is_loaded()) {
     return Error::Ok;
   }
-  stats_.model_load_start_ms = llm::time_in_ms();
+  stats_->model_load_start_ms = llm::time_in_ms();
 
   // Load the tokenizer
   tokenizer_ = std::make_unique<tokenizers::Llama2cTokenizer>();
@@ -71,9 +71,9 @@ Error LlavaRunner::load() {
       /*use_kv_cache=*/true,
       std::make_unique<std::unordered_set<uint64_t>>(
           std::unordered_set<uint64_t>{tokenizer_->eos_tok()}),
-      &stats_);
+      stats_.get());
 
-  stats_.model_load_end_ms = llm::time_in_ms();
+  stats_->model_load_end_ms = llm::time_in_ms();
   return Error::Ok;
 }
 
@@ -113,9 +113,9 @@ Error LlavaRunner::generate_from_pos(
 
   uint64_t prefill_next_token =
       ET_UNWRAP(prefill_prompt(prompt, start_pos, /*bos=*/0, /*eos*/ 0));
-  stats_.first_token_ms = llm::time_in_ms();
-  stats_.prompt_eval_end_ms = llm::time_in_ms();
-  stats_.num_prompt_tokens = start_pos;
+  stats_->first_token_ms = llm::time_in_ms();
+  stats_->prompt_eval_end_ms = llm::time_in_ms();
+  stats_->num_prompt_tokens = start_pos;
 
   // Generate tokens
   int64_t num_generated_tokens = ET_UNWRAP(text_token_generator_->generate(
@@ -126,9 +126,9 @@ Error LlavaRunner::generate_from_pos(
       /*token_callback=*/token_callback));
 
   // Bookkeeping
-  stats_.num_generated_tokens = num_generated_tokens;
+  stats_->num_generated_tokens = num_generated_tokens;
   if (stats_callback) {
-    stats_callback(stats_);
+    stats_callback(*stats_);
   }
   return Error::Ok;
 }
@@ -161,7 +161,7 @@ Error LlavaRunner::generate(
       };
 
   int64_t pos = 0;
-  stats_.inference_start_ms = llm::time_in_ms();
+  stats_->inference_start_ms = llm::time_in_ms();
 
   // prefill preset prompt
   prefill_prompt(kPresetPrompt, pos, /*bos=*/1, /*eos*/ 0);
@@ -178,8 +178,8 @@ Error LlavaRunner::generate(
   Error err = generate_from_pos(
       prompt, seq_len, pos, wrapped_callback, stats_callback, echo);
 
-  stats_.inference_end_ms = llm::time_in_ms();
-  ::executorch::llm::print_report(stats_);
+  stats_->inference_end_ms = llm::time_in_ms();
+  ::executorch::llm::print_report(*stats_);
 
   ET_LOG(
       Info,

extension/llm/runner/multimodal_runner.h

@@ -41,7 +41,8 @@ class ET_EXPERIMENTAL MultimodalRunner {
       const float temperature = 0.8f)
       : temperature_(temperature),
         module_(std::make_unique<Module>(model_path, Module::LoadMode::File)),
-        tokenizer_path_(tokenizer_path) {
+        tokenizer_path_(tokenizer_path),
+        stats_(std::make_unique<llm::Stats>()) {
     ET_LOG(
         Info,
         "Creating Multimodal LLM runner: model_path=%s, tokenizer_path=%s",
@@ -132,7 +133,7 @@ class ET_EXPERIMENTAL MultimodalRunner {
   std::unique_ptr<::tokenizers::Tokenizer> tokenizer_;
 
   // stats
-  Stats stats_;
+  std::unique_ptr<Stats> stats_;
 };
 
 } // namespace llm

extension/llm/runner/text_decoder_runner.h

@@ -14,7 +14,6 @@
 #include <executorch/extension/module/module.h>
 #include <executorch/extension/tensor/tensor.h>
 #include <executorch/runtime/platform/compiler.h>
-#include <functional>
 
 namespace executorch {
 namespace extension {
@@ -94,7 +93,13 @@ class ET_EXPERIMENTAL TextDecoderRunner {
   }
 
  protected:
-  // TODO: use shared_ptr for module
+  /**
+   * Note: TextDecoderRunner does not own the Module instance. It is expected
+   * that the outer class (likely Runner) manages the lifecycle of the Module.
+   * This means that the responsibility for creating, maintaining, and
+   * destroying the Module lies outside of TextDecoderRunner. Ensure that the
+   * Module remains valid for the duration of TextDecoderRunner's usage.
+   */
   Module* module_;
   bool use_kv_cache_;
   bool should_stop_{false};

extension/llm/runner/text_prefiller.h

@@ -24,6 +24,8 @@ class ET_EXPERIMENTAL TextPrefiller {
       bool use_kv_cache_,
       bool enable_parallel_prefill,
       int64_t max_seq_len = 128);
+
+  virtual ~TextPrefiller() = default;
   /**
    * Prefill an LLM Module with the given text input.
    * @param prompt_tokens The text prompt tokens to the LLM Module. Encoded by
@@ -32,7 +34,7 @@ class ET_EXPERIMENTAL TextPrefiller {
    * Module.
    * @return The next token of the LLM Module after prefill.
    */
-  ::executorch::runtime::Result<uint64_t> prefill(
+  virtual ::executorch::runtime::Result<uint64_t> prefill(
       std::vector<uint64_t>& prompt_tokens,
       int64_t& start_pos);
 
@@ -48,6 +50,12 @@ class ET_EXPERIMENTAL TextPrefiller {
       int64_t& start_pos);
 
  private:
+  /**
+   * Note: TextPrefiller does not own the TextDecoderRunner instance.
+   * The responsibility of managing the lifecycle of TextDecoderRunner
+   * lies with the outer class or entity (likely Runner) that creates
+   * and passes the TextDecoderRunner instance to TextPrefiller.
+   */
   TextDecoderRunner* text_decoder_runner_;
   bool use_kv_cache_;
   bool enable_parallel_prefill_;