CLI Arg Groups: Add more structure to the arg parser for Model specification and Model configuration (pytorch#937)

* Suppress the occurance of unused args in  README subcommands

* Move GGUF into model specification exclusive require group

* Further group args into parsing groups

* Function Description Typos

Author: Jack-Khuu

cli.py

@@ -28,8 +28,11 @@
 # Subcommands related to downloading and managing model artifacts
 INVENTORY_VERBS = ["download", "list", "remove", "where"]
 
+# Subcommands related to generating inference output based on user prompts
+GENERATION_VERBS = ["browser", "chat", "generate", "server"] 
+
 # List of all supported subcommands in torchchat
-KNOWN_VERBS = ["chat", "browser", "generate", "server", "eval", "export"] + INVENTORY_VERBS
+KNOWN_VERBS = GENERATION_VERBS + ["eval", "export"] + INVENTORY_VERBS
 
 
 # Handle CLI arguments that are common to a majority of subcommands.
@@ -46,61 +49,97 @@ def check_args(args, verb: str) -> None:
 
 # Given a arg parser and a subcommand (verb), add the appropriate arguments
 # for that subcommand.
+#
+# Note the use of argparse.SUPPRESS to hide arguments from --help due to 
+# legacy CLI arg parsing. See https://github.com/pytorch/torchchat/issues/932
 def add_arguments_for_verb(parser, verb: str) -> None:
     # Argument closure for inventory related subcommands
     if verb in INVENTORY_VERBS:
         _configure_artifact_inventory_args(parser, verb)
         _add_cli_metadata_args(parser)
         return
 
-    # Model specification
-    # A model can be specified using a positional model name or checkpoint path
-    parser.add_argument(
+    # Add argument groups for model specification (what base model to use)
+    _add_model_specification_args(parser)
+
+    # Add argument groups for exported model path IO
+    _add_exported_input_path_args(parser, verb)
+    _add_export_output_path_args(parser, verb)
+    
+    # Add argument groups for model configuration (compilation, quant, etc)
+    _add_model_config_args(parser, verb)
+
+    # Add thematic argument groups based on the subcommand
+    if verb in ["browser", "chat", "generate", "server"]:
+        _add_generation_args(parser, verb)
+    if verb == "eval":
+        _add_evaluation_args(parser)
+
+    # Add CLI Args related to downloading of model artifacts (if not already downloaded)
+    _add_jit_downloading_args(parser)
+
+    # Add CLI Args that are general to subcommand cli execution
+    _add_cli_metadata_args(parser)
+
+    # WIP Features (suppressed from --help)
+    _add_distributed_args(parser)
+    _add_custom_model_args(parser)
+    _add_speculative_execution_args(parser)
+
+
+# Add CLI Args related to model specification (what base model to use)
+def _add_model_specification_args(parser) -> None:
+    model_specification_parser = parser.add_argument_group("Model Specification", "(REQUIRED) Specify the base model. Args are mutually exclusive.")
+    exclusive_parser = model_specification_parser.add_mutually_exclusive_group(required=True)
+    exclusive_parser.add_argument(
         "model",
         type=str,
         nargs="?",
         default=None,
         help="Model name for well-known models",
     )
-    parser.add_argument(
+    exclusive_parser.add_argument(
         "--checkpoint-path",
         type=Path,
         default="not_specified",
         help="Use the specified model checkpoint path",
     )
+    # See _add_custom_model_args() for more details
+    exclusive_parser.add_argument(
+        "--gguf-path",
+        type=Path,
+        default=None,
+        help=argparse.SUPPRESS,
+        # "Use the specified GGUF model file",
+    )
 
-    # Add thematic argument groups based on the subcommand
-    if verb in ["browser", "chat", "generate", "server"]:
-        _add_generation_args(parser)
-    if verb == "eval":
-        _add_evaluation_args(parser)
-
-    # Add argument groups for exported model path IO
-    _add_exported_input_path_args(parser)
-    _add_export_output_path_args(parser)
-
-    parser.add_argument(
+    model_specification_parser.add_argument(
         "--is-chat-model",
         action="store_true",
-        help="Indicate that the model was trained to support chat functionality",
+        # help="Indicate that the model was trained to support chat functionality",
+        help=argparse.SUPPRESS,
     )
-    parser.add_argument(
+
+# Add CLI Args related to model configuration (compilation, quant, etc)
+def _add_model_config_args(parser, verb: str) -> None:
+    model_config_parser = parser.add_argument_group("Model Configuration", "Specify model configurations")
+    model_config_parser.add_argument(
         "--compile",
         action="store_true",
         help="Whether to compile the model with torch.compile",
     )
-    parser.add_argument(
+    model_config_parser.add_argument(
         "--compile-prefill",
         action="store_true",
         help="Whether to compile the prefill. Improves prefill perf, but has higher compile times.",
     )
-    parser.add_argument(
+    model_config_parser.add_argument(
         "--dtype",
         default="fast",
         choices=allowable_dtype_names(),
         help="Override the dtype of the model (default is the checkpoint dtype). Options: bf16, fp16, fp32, fast16, fast",
     )
-    parser.add_argument(
+    model_config_parser.add_argument(
         "--quantize",
         type=str,
         default="{ }",
@@ -109,81 +148,84 @@ def add_arguments_for_verb(parser, verb: str) -> None:
             + "modes are: embedding, linear:int8, linear:int4, linear:a8w4dq, precision."
         ),
     )
-    parser.add_argument(
+    model_config_parser.add_argument(
         "--device",
         type=str,
         default=default_device,
         choices=["fast", "cpu", "cuda", "mps"],
         help="Hardware device to use. Options: cpu, cuda, mps",
     )
-    parser.add_argument(
-        "--hf-token",
-        type=str,
-        default=None,
-        help="A HuggingFace API token to use when downloading model artifacts",
-    )
-    parser.add_argument(
-        "--model-directory",
-        type=Path,
-        default=default_model_dir,
-        help=f"The directory to store downloaded model artifacts. Default: {default_model_dir}",
-    )
-    parser.add_argument(
-        "--profile",
-        type=Path,
-        default=None,
-        help="Profile path.",
-    )
-    _add_cli_metadata_args(parser)
-
-    # WIP Features (suppressed from --help)
-    _add_distributed_args(parser)
-    _add_custom_model_args(parser)
-    _add_speculative_execution_args(parser)
 
+# Add CLI Args representing output paths of exported model files
+def _add_export_output_path_args(parser, verb: str) -> None:
+    is_export = verb == "export"
 
-# Add CLI Args representing user provided exported model files
-def _add_export_output_path_args(parser) -> None:
     output_path_parser = parser.add_argument_group(
-        "Export Output Path Args",
-        "Specify the output path for the exported model files",
+        "Export Output Path" if is_export else None,
+        "Specify the output path for the exported model files" if is_export else None,
     )
-    output_path_parser.add_argument(
+    exclusive_parser = output_path_parser.add_mutually_exclusive_group()
+    exclusive_parser.add_argument(
         "--output-pte-path",
         type=str,
         default=None,
-        help="Output to the specified ExecuTorch .pte model file",
+        help="Output to the specified ExecuTorch .pte model file" if is_export else argparse.SUPPRESS,
     )
-    output_path_parser.add_argument(
+    exclusive_parser.add_argument(
         "--output-dso-path",
         type=str,
         default=None,
-        help="Output to the specified AOT Inductor .dso model file",
+        help="Output to the specified AOT Inductor .dso model file" if is_export else argparse.SUPPRESS,
     )
 
 
 # Add CLI Args representing user provided exported model files
-def _add_exported_input_path_args(parser) -> None:
+def _add_exported_input_path_args(parser, verb: str) -> None:
+    is_generation_verb = verb in GENERATION_VERBS
+
     exported_model_path_parser = parser.add_argument_group(
-        "Exported Model Path Args",
-        "Specify the path of the exported model files to ingest",
+        "Exported Model Path" if is_generation_verb else None,
+        "Specify the path of the exported model files to ingest" if is_generation_verb else None,
     )
-    exported_model_path_parser.add_argument(
+    exclusive_parser = exported_model_path_parser.add_mutually_exclusive_group()
+    exclusive_parser.add_argument(
         "--dso-path",
         type=Path,
         default=None,
-        help="Use the specified AOT Inductor .dso model file",
+        help="Use the specified AOT Inductor .dso model file" if is_generation_verb else argparse.SUPPRESS,
     )
-    exported_model_path_parser.add_argument(
+    exclusive_parser.add_argument(
         "--pte-path",
         type=Path,
         default=None,
-        help="Use the specified ExecuTorch .pte model file",
+        help="Use the specified ExecuTorch .pte model file" if is_generation_verb else argparse.SUPPRESS,
     )
 
+# Add CLI Args related to JIT downloading of model artifacts
+def _add_jit_downloading_args(parser) -> None:
+    jit_downloading_parser = parser.add_argument_group("Model Downloading", "Specify args for model downloading (if model is not downloaded)",)
+    jit_downloading_parser.add_argument(
+        "--hf-token",
+        type=str,
+        default=None,
+        help="A HuggingFace API token to use when downloading model artifacts",
+    )
+    jit_downloading_parser.add_argument(
+        "--model-directory",
+        type=Path,
+        default=default_model_dir,
+        help=f"The directory to store downloaded model artifacts. Default: {default_model_dir}",
+    )
 
 # Add CLI Args that are general to subcommand cli execution
 def _add_cli_metadata_args(parser) -> None:
+    parser.add_argument(
+        "--profile",
+        type=Path,
+        default=None,
+        # help="Profile path.",
+        help=argparse.SUPPRESS,
+    )
     parser.add_argument(
         "-v",
         "--verbose",
@@ -227,25 +269,27 @@ def _configure_artifact_inventory_args(parser, verb: str) -> None:
 
 
 # Add CLI Args specific to user prompted generation
-def _add_generation_args(parser) -> None:
+def _add_generation_args(parser, verb: str) -> None:
     generator_parser = parser.add_argument_group(
-        "Generation Args", "Configs for generating output based on provided prompt"
+        "Generation", "Configs for generating output based on provided prompt"
     )
     generator_parser.add_argument(
         "--prompt",
         type=str,
         default="Hello, my name is",
-        help="Input prompt for manual output generation",
+        help="Input prompt for manual output generation" if verb == "generate" else argparse.SUPPRESS,
     )
     generator_parser.add_argument(
         "--chat",
         action="store_true",
-        help="Whether to start an interactive chat session",
+        # help="Whether to start an interactive chat session",
+        help=argparse.SUPPRESS,
     )
     generator_parser.add_argument(
         "--gui",
         action="store_true",
-        help="Whether to use a web UI for an interactive chat session",
+        # help="Whether to use a web UI for an interactive chat session",
+        help=argparse.SUPPRESS,
     )
     generator_parser.add_argument(
         "--num-samples",
@@ -271,14 +315,15 @@ def _add_generation_args(parser) -> None:
     generator_parser.add_argument(
         "--sequential-prefill",
         action="store_true",
-        help="Whether to perform prefill sequentially. Only used for model debug.",
+        # help="Whether to perform prefill sequentially. Only used for model debug.",
+        help=argparse.SUPPRESS,
     )
 
 
 # Add CLI Args specific to Model Evaluation
 def _add_evaluation_args(parser) -> None:
     eval_parser = parser.add_argument_group(
-        "Evaluation Args", "Configs for evaluating model performance"
+        "Evaluation", "Configs for evaluating model performance"
     )
     eval_parser.add_argument(
         "--tasks",
@@ -337,13 +382,6 @@ def _add_custom_model_args(parser) -> None:
         help=argparse.SUPPRESS,
         # "Use the specified parameter file",
     )
-    parser.add_argument(
-        "--gguf-path",
-        type=Path,
-        default=None,
-        help=argparse.SUPPRESS,
-        # "Use the specified GGUF model file",
-    )
     parser.add_argument(
         "--tokenizer-path",
         type=Path,