add peak memory log for tvm as well

octoml · Feb 20, 2024 · f716851 · f716851
1 parent cf0813d
commit f716851
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/serve/mlc_serve/model/tvm_model.py b/serve/mlc_serve/model/tvm_model.py
@@ -253,6 +253,8 @@ def profile_memory_usage(self, seq_lens):
 
         vm_alloc_after = self.get_used_memory()
 
+        LOG.info(f"peak memory during profling: {(vm_alloc_after - vm_alloc_before) / 1e9} GB")
+
         return self.get_param_nbytes() + (vm_alloc_after - vm_alloc_before)
 
     def generate_multi_query(