Function main

pretrain/main.py:335–436

(
    fabric: lightning.Fabric,
    data_dir: Path,
    resume: Union[Path, Literal[False]],
    training_params: TrainingParams,
    ignore_index=-100,
)

def main(
    fabric: lightning.Fabric,
    data_dir: Path,
    resume: Union[Path, Literal[False]],
    training_params: TrainingParams,
    ignore_index=-100,
):
    monitor = Monitor(
        fabric,
        window_size=2,
        time_unit='seconds',
        log_iter_interval=training_params['log_iter_interval'],
    )

    # Create model out folder only for the first node.
    if fabric.global_rank == 0:
        training_params['out_dir'].mkdir(parents=True, exist_ok=True)

    config = Config.from_name(training_params['model_name'])

    train_dataloader, val_dataloader = create_dataloaders(
        break_into_chunks=training_params['devices'],
        batch_size=training_params['micro_batch_size'],
        block_size=config.block_size,
        fabric=fabric,
        data_dir=data_dir,
        seed=3407,
    )

    if train_dataloader is not None:
        fabric.print('Train dataloader prepared successfully')

    if val_dataloader is None:
        fabric.print('Setting up fabric dataloader for training only...')
        train_dataloader = fabric.setup_dataloaders(train_dataloader)
    else:
        fabric.print('Validation dataloader prepared successfully.')
        fabric.print('Setting up fabric dataloaders for training and validation...')
        train_dataloader, val_dataloader = fabric.setup_dataloaders(
            train_dataloader, val_dataloader
        )

    # same seed for every process to init model (FSDP)
    fabric.seed_everything(3407)

    fabric.print(f'Loading model with {config.__dict__}')

    t0 = time.perf_counter()

    with fabric.init_module(empty_init=False):
        model = GPT(config)
        if (
            resume is False
        ):  # this means we are training from scratch, hence need to initalize the weights for all model layers
            model.apply(partial(model._init_weights, n_layer=config.n_layer))

    instantiation_time = time.perf_counter() - t0