Function train

eval.py:500–671

Main training logic. Args: config (DictConfig): Configuration composed by OmegaConf

(config: om.DictConfig)

def train(config: om.DictConfig) -> None:
    """Main training logic.

    Args:
        config (DictConfig): Configuration composed by OmegaConf
    """
    # these subtasks require the parent task to have been run
    round_2_task_names = config.get(
        "round_2_task_names",
        {
            "mnli": {"rte", "mrpc", "stsb"},
            "swag": {"copa"},
        },
    )

    start_time = time.time()

    # Initial default seed
    reproducibility.seed_all(config.default_seed)

    # Quiet down WandB
    os.environ["WANDB_SILENT"] = "true"

    # Set tokenizer parallelism
    os.environ["TOKENIZERS_PARALLELISM"] = "false"

    # Confirm GPUs if parallel=True
    if config.parallel:
        assert (
            torch.cuda.device_count() > 0
        ), "Can only use parallel mode if GPUs are available. Please set parallel=False."

    # Downloads the starting checkpoint ahead of time so that
    # the different tasks don't all try to download it at the same time
    if config.get("starting_checkpoint_load_path", None):
        local_pretrain_checkpoint_path = download_starting_checkpoint(
            config.starting_checkpoint_load_path,
            config.local_pretrain_checkpoint_folder,
        )
    else:
        local_pretrain_checkpoint_path = None

    # Builds round 1 configs and runs them by first filtering out all round 2 tasks
    if round_2_task_names:
        round_2_tasks = [task for tasks in round_2_task_names.values() for task in tasks]
    else:
        round_2_tasks = []
    round_1_task_names = [task for task in TASK_NAME_TO_CLASS.keys() if task not in round_2_tasks]

    round_1_job_configs = create_job_configs(config, round_1_task_names, local_pretrain_checkpoint_path)

    round_1_results = {}
    if len(round_1_job_configs) > 0:
        if config.parallel:
            round_1_results = run_jobs_parallel(round_1_job_configs)
        else:
            round_1_results = run_jobs_serial(round_1_job_configs)