diff --git a/nemo/collections/llm/tools/auto_configurator/autoconfig/search_config.py b/nemo/collections/llm/tools/auto_configurator/autoconfig/search_config.py index 59b6cffa20a6..53d7b68c492c 100644 --- a/nemo/collections/llm/tools/auto_configurator/autoconfig/search_config.py +++ b/nemo/collections/llm/tools/auto_configurator/autoconfig/search_config.py @@ -42,7 +42,7 @@ def search_config(cfg: dict): :param Optional[str] hydra_args: hydra override arguments in string format. :return: None """ - + # Read config nodes = cfg.get("num_nodes") gpus_per_node = cfg.get("gpus_per_node") @@ -57,7 +57,7 @@ def search_config(cfg: dict): seq_length = cfg.get("seq_length") log_dir = cfg.get("log_dir") custom_cfg = None - + print(cfg) print(model_name) assert model_name in SUPPORTED_MODELS, f"model must be set to one of {SUPPORTED_MODELS}/" diff --git a/nemo/collections/llm/tools/auto_configurator/main_copy.py b/nemo/collections/llm/tools/auto_configurator/main_copy.py index f1ac5ddff667..091d121ea0f2 100644 --- a/nemo/collections/llm/tools/auto_configurator/main_copy.py +++ b/nemo/collections/llm/tools/auto_configurator/main_copy.py @@ -25,17 +25,22 @@ def get_args(): parser.add_argument('--gpus_per_node', required=False, default=8, type=int, help="Number of GPUs per node") parser.add_argument('--gpu_memory_gb', required=False, default=80, type=int, help="GPU memory size") parser.add_argument('--max_training_days', required=False, default=2, type=int, help="Path to data file") - parser.add_argument('--max_minutes_per_run', required=False, default=30, type=int, help="Max minutes per job on cluster") + parser.add_argument( + '--max_minutes_per_run', required=False, default=30, type=int, help="Max minutes per job on cluster" + ) parser.add_argument('--model_type', required=True, type=str, help="Model size in billions") parser.add_argument('--model_size_in_b', required=True, type=int, help="Model size in billions") parser.add_argument('--vocab_size', required=False, default=32000, type=int, help="Size of tokenizer vocab") parser.add_argument('--tflops_per_gpu', required=False, default=140, type=int, help="Estimated tflops per GPU") - parser.add_argument('--num_tokens_in_b', required=False, default=300, type=int, help="Number of tokens in dataset in billions") + parser.add_argument( + '--num_tokens_in_b', required=False, default=300, type=int, help="Number of tokens in dataset in billions" + ) parser.add_argument('--seq_length', required=False, default=2048, type=int, help="Model sequence length") parser.add_argument('--log_dir', required=True, type=str, help="Path to results directory") return parser.parse_args() + def main(): """ Main function in the entire pipeline, it reads the config using hydra and calls search_config.