Add model warmup flag into cli (#197)

add model warmup flag into cli
AI-Hypercomputer · Nov 6, 2024 · fe22a9f · fe22a9f
1 parent 02927c9
commit fe22a9f
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/jetstream_pt/cli.py b/jetstream_pt/cli.py
@@ -36,6 +36,7 @@
 flags.DEFINE_bool(
     "internal_use_local_tokenizer", 0, "Use local tokenizer if set to True"
 )
+flags.DEFINE_bool("enable_model_warmup", False, "enable model warmup")
 
 
 def shard_weights(env, weights, weight_shardings):
@@ -111,6 +112,7 @@ def serve():
       config=server_config,
       devices=devices,
       metrics_server_config=metrics_server_config,
+      enable_model_warmup=FLAGS.enable_model_warmup,
   )
   print("Started jetstream_server....")
   jetstream_server.wait_for_termination()