Force text generation to exit if checkpoint fails to load

zarzen · Jul 11, 2023 · cd961e1 · cd961e1
1 parent 7861224
commit cd961e1
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/tools/run_text_generation_server.py b/tools/run_text_generation_server.py
@@ -53,6 +53,9 @@ def add_text_generate_args(parser):
     if args.num_layers_per_virtual_pipeline_stage is not None:
         print("Interleaved pipeline schedule is not yet supported for text generation.")
         exit()
+    print_rank_0("WARNING: Forcing exit_on_missing_checkpoint to True for text "
+                 "generation.")
+    args.exit_on_missing_checkpoint = True
     # Set up model and load checkpoint
     model = get_model(model_provider, wrap_with_ddp=False)