Function train

glue.py:434–588

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
    """
    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
    round_1_task_names = {
        # glue:
        *{"cola", "sst2", "qqp", "qnli", "mnli"},
        # superglue:
        *{"boolq", "cb", "multirc", "wic"},
        # misc:
        *{"swag", "eurlex"},
    }
    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)

    # Builds up the information needed to run the second round, starting from the MNLI checkpoints
    checkpoint_paths = {}
    for job_name, output_dict in round_1_results.items():
        job_results = output_dict["result"]
        job_values = get_values_from_path(job_name, separator="_")