diff --git a/nemo/collections/llm/tools/auto_configurator/autoconfig/base_config.py b/nemo/collections/llm/tools/auto_configurator/autoconfig/base_config.py index 79c90b633b06..c35ed5b42f62 100644 --- a/nemo/collections/llm/tools/auto_configurator/autoconfig/base_config.py +++ b/nemo/collections/llm/tools/auto_configurator/autoconfig/base_config.py @@ -526,7 +526,8 @@ def generate_base_config( :return: base config object for the given model. :rtype: dict """ - base_cfg = utils.generic_base_config(model_name=model_name, model_version=model_version, model_size_in_b=model_size_in_b, cfg=cfg) - + base_cfg = utils.generic_base_config( + model_name=model_name, model_version=model_version, model_size_in_b=model_size_in_b, cfg=cfg + ) return base_cfg diff --git a/nemo/collections/llm/tools/auto_configurator/autoconfig/training_config.py b/nemo/collections/llm/tools/auto_configurator/autoconfig/training_config.py index b8cd4798280e..eb61bc3acca9 100644 --- a/nemo/collections/llm/tools/auto_configurator/autoconfig/training_config.py +++ b/nemo/collections/llm/tools/auto_configurator/autoconfig/training_config.py @@ -43,9 +43,9 @@ def search_training_config( # Generate candidate configs. configs = generate_grid_search_configs(base_cfg, train_cfg, model_size_in_b, model_name) # Launch candidate configs. - #job_ids = launch_grid_search_configs(base_dir, results_cfgs, model_name, cfg) + # job_ids = launch_grid_search_configs(base_dir, results_cfgs, model_name, cfg) # Measure and compare throughputs for each config. - #launch_throughput_measure(job_ids, model_name, model_size_in_b, num_nodes, hydra_args, cfg) + # launch_throughput_measure(job_ids, model_name, model_size_in_b, num_nodes, hydra_args, cfg) return configs diff --git a/nemo/collections/llm/tools/auto_configurator/autoconfig/utils.py b/nemo/collections/llm/tools/auto_configurator/autoconfig/utils.py index 77ce9ba94edc..3e5bf1f40c4c 100644 --- a/nemo/collections/llm/tools/auto_configurator/autoconfig/utils.py +++ b/nemo/collections/llm/tools/auto_configurator/autoconfig/utils.py @@ -18,9 +18,8 @@ from nemo.collections.llm.tools.auto_configurator import base_configs -MODULES = { - "llama": "Llama" -} +MODULES = {"llama": "Llama"} + def _calculate_model_size( vocab_size: int = None, @@ -308,7 +307,9 @@ def calculate_model_size_params( raise Exception("Number of layers not found, config is not possible.") -def generic_base_config(model_name: str = "llama", model_version: int = 2, model_size_in_b: int = 7, cfg: dict = {}) -> dict: +def generic_base_config( + model_name: str = "llama", model_version: int = 2, model_size_in_b: int = 7, cfg: dict = {} +) -> dict: """ Generates a base config dictionary from a base config yaml file. :param omegaconf.dictconfig.DictConfig cfg: hydra-like config object for the HP tool. diff --git a/nemo/collections/llm/tools/auto_configurator/main.py b/nemo/collections/llm/tools/auto_configurator/main.py index 7332dda1b424..0529811d5f1e 100644 --- a/nemo/collections/llm/tools/auto_configurator/main.py +++ b/nemo/collections/llm/tools/auto_configurator/main.py @@ -51,5 +51,6 @@ def main(): args = get_args() configs = search_config(cfg=vars(args)) + if __name__ == "__main__": main()