Add model warmup flag into cli (#197)

vivianrwu · web-flow · commit fe22a9f70545 · 2024-11-06T11:54:52.000-08:00
add model warmup flag into cli
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()

Original file line number	Diff line number	Diff line change
`@@ -36,6 +36,7 @@`
`36`	`36`	`flags.DEFINE_bool(`
`37`	`37`	`"internal_use_local_tokenizer", 0, "Use local tokenizer if set to True"`
`38`	`38`	`)`
	`39`	`+flags.DEFINE_bool("enable_model_warmup", False, "enable model warmup")`
`39`	`40`
`40`	`41`
`41`	`42`	`def shard_weights(env, weights, weight_shardings):`
`@@ -111,6 +112,7 @@ def serve():`
`111`	`112`	`config=server_config,`
`112`	`113`	`devices=devices,`
`113`	`114`	`metrics_server_config=metrics_server_config,`
	`115`	`+ enable_model_warmup=FLAGS.enable_model_warmup,`
`114`	`116`	`)`
`115`	`117`	`print("Started jetstream_server....")`
`116`	`118`	`jetstream_server.wait_for_termination()`