Move tokenizer info into pte for et

Author: Yuhan Guo

runner/run.cpp

@@ -53,6 +53,12 @@ using executorch::extension::TensorPtr;
 using torch::executor::EValue;
 using torch::executor::Module;
 using torch::executor::Result;
+using executorch::runtime::MemoryManager;
+using executorch::runtime::HierarchicalAllocator;
+using executorch::runtime::MemoryAllocator;
+using executorch::runtime::Span;
+using executorch::runtime::MethodMeta;
+using executorch::runtime::Error;
 #endif
 
 using tokenizers::SPTokenizer;
@@ -867,6 +873,37 @@ int main(int argc, char *argv[]) {
                     : torch::Device(torch::kCUDA);
   ModelType model_type = get_model_type(std::stoi(aoti_metadata["tokenizer_type"]));
 #else // __ET_MODEL__
+  transformer.runner->load();
+  auto num_methods = transformer.runner->program()->num_methods();
+  fprintf(stderr, "num_methods: %zu\n", num_methods);
+
+  Result<MethodMeta> method_meta = transformer.runner->program()->method_meta("tokenizer_type");
+
+  static std::array<uint8_t, 4 * 1024U * 1024U> method_allocator_pool; // 4MB
+  MemoryAllocator method_allocator{MemoryAllocator(
+      sizeof(method_allocator_pool), method_allocator_pool.data())};
+      std::vector<std::unique_ptr<uint8_t[]>> planned_buffers; // Owns the memory
+  std::vector<Span<uint8_t>> planned_spans; // Passed to the allocator
+  size_t num_memory_planned_buffers = method_meta->num_memory_planned_buffers();
+  for (size_t id = 0; id < num_memory_planned_buffers; ++id) {
+    // .get() will always succeed because id < num_memory_planned_buffers.
+    size_t buffer_size =
+        static_cast<size_t>(method_meta->memory_planned_buffer_size(id).get());
+    planned_buffers.push_back(std::make_unique<uint8_t[]>(buffer_size));
+    planned_spans.push_back({planned_buffers.back().get(), buffer_size});
+  }
+  HierarchicalAllocator planned_memory(
+      {planned_spans.data(), planned_spans.size()});
+  MemoryManager memory_manager(&method_allocator, &planned_memory);
+  auto tokenizer_method = transformer.runner->program()->load_method("tokenizer_type", &memory_manager);
+
+  Error err = tokenizer_method->execute();
+  fprintf(stderr, "err: %u\n", err);
+
+  std::vector<EValue> outputs(tokenizer_method->outputs_size());
+  Error status = tokenizer_method->get_outputs(outputs.data(), outputs.size());
+  fprintf(stderr, "status: %u\n", status);
+  
   ModelType model_type = get_model_type(llama_ver);
 #endif
 

torchchat/export.py

@@ -313,7 +313,7 @@ def export_to_edge(
             core_aten_ep, edge_constant_methods, edge_compile_config, verbose=verbose
         )
 
-    def export_for_et(model, device, output_path) -> str:
+    def export_for_et(model, device, output_path, edge_constant_methods) -> str:
 
         input = (
             torch.tensor([[1]], dtype=torch.long, device=device),
@@ -344,12 +344,15 @@ def export_for_et(model, device, output_path) -> str:
         with torch.nn.attention.sdpa_kernel(
             [torch.nn.attention.SDPBackend.MATH]
         ), torch.no_grad():
-            m = export_for_training(model, input, dynamic_shapes=dynamic_shapes).module()
+            m = export_for_training(
+                model, input, dynamic_shapes=dynamic_shapes
+            ).module()
 
             edge_manager = export_to_edge(
                 m,
                 input,
                 dynamic_shapes=dynamic_shapes,
+                edge_constant_methods=edge_constant_methods,
                 edge_compile_config=edge_config,
             )
         edge_manager = edge_manager.to_backend(XnnpackDynamicallyQuantizedPartitioner())
@@ -365,6 +368,7 @@ def export_for_et(model, device, output_path) -> str:
         )
 
         print("The methods are: ", export_program.methods)
+        print("The config methods are: ", export_program.config_methods)
         with open(output_path, "wb") as f:
             export_program.write_to_file(f)
 
@@ -407,7 +411,9 @@ def main(args):
             f"Warning! ExecuTorch export target is controlled by export recipe, not device setting. Ignoring device={builder_args.device} setting."
         )
         builder_args.device = "cpu"
-    elif (output_pte_path or output_dso_path or output_aoti_package_path) and "mps" in builder_args.device:
+    elif (
+        output_pte_path or output_dso_path or output_aoti_package_path
+    ) and "mps" in builder_args.device:
         print("Warning! Device MPS not supported for export. Exporting for device CPU.")
         builder_args.device = "cpu"
 
@@ -473,13 +479,26 @@ def main(args):
                 support_tensor_subclass=False,
             )
             _unset_gguf_kwargs(builder_args)
- 
+
+    if tokenizer_args is None:
+        tokenizer_type = "0"
+    elif tokenizer_args.is_sentencepiece:
+        tokenizer_type = "2"  # Corresponding to llama2
+    else:
+        tokenizer_type = "3"  # Corresponding to llama3
+
     with torch.no_grad():
         if output_pte_path:
             output_pte_path = str(os.path.abspath(output_pte_path))
             if executorch_export_available:
                 print(f"Exporting model using ExecuTorch to {output_pte_path}")
-                export_for_et(model_to_pte, builder_args.device, args.output_pte_path)
+                print(f"Tokenizer type is {tokenizer_type}")
+                export_for_et(
+                    model_to_pte,
+                    builder_args.device,
+                    args.output_pte_path,
+                    {"tokenizer_type": tokenizer_type},
+                )
             else:
                 print(
                     "Export with executorch requested but ExecuTorch could not be loaded"
@@ -503,13 +522,6 @@ def main(args):
         if output_aoti_package_path:
             output_aoti_package_path = str(os.path.abspath(output_aoti_package_path))
 
-            if tokenizer_args is None:
-                tokenizer_type = "0"
-            elif tokenizer_args.is_sentencepiece:
-                tokenizer_type = "2"  # Corresponding to llama2
-            else:
-                tokenizer_type = "3"  # Corresponding to llama3
-
             metadata = {"tokenizer_type": tokenizer_type}
             print(
                 "Exporting model using AOT Inductor to " f"{output_aoti_package_path}."