Function train

pretrain/main.py:439–638

Trains the model

(
    fabric: lightning.Fabric,
    state: TrainingState,
    train_dataloader: DataLoader,
    val_dataloader: Optional[DataLoader],
    monitor: Monitor,
    resume: Union[Path, bool],
    training_params: TrainingParams,
    ignore_index=-100,
)

def train(
    fabric: lightning.Fabric,
    state: TrainingState,
    train_dataloader: DataLoader,
    val_dataloader: Optional[DataLoader],
    monitor: Monitor,
    resume: Union[Path, bool],
    training_params: TrainingParams,
    ignore_index=-100,
):
    """Trains the model"""

    model = state['model']
    optimizer = state['optimizer']

    if val_dataloader is not None:
        validate(
            fabric, model, val_dataloader, training_params['eval_iters']
        )  # sanity check

    with torch.device('meta'):
        meta_model = GPT(model.config)
        # "estimated" is not as precise as "measured". Estimated is optimistic but widely used in the wild.
        # When comparing MFU or FLOP numbers with other projects that use estimated FLOPs,
        # consider passing `SpeedMonitor(flops_per_batch=estimated_flops)` instead
        estimated_flops = (
            estimate_flops(meta_model) * training_params['micro_batch_size']
        )

        fabric.print(
            f'Estimated TFLOPs: { estimated_flops * fabric.world_size / 1e12:.2f}'
        )
        x = torch.randint(
            0, 1, (training_params['micro_batch_size'], model.config.block_size)
        )
        # measured_flos run in meta. Will trigger fusedRMSNorm error
        # measured_flops = measure_flops(meta_model, x)
        # fabric.print(f"Measured TFLOPs: {measured_flops * fabric.world_size / 1e12:.2f}")
        del meta_model, x

    total_lengths = 0
    total_t0 = time.perf_counter()

    if fabric.device.type == 'xla':
        import torch_xla.core.xla_model as xm

        xm.mark_step()

    initial_iter = state['iter_num']
    curr_iter = 0

    loss_func = FusedCrossEntropyLoss(ignore_index=ignore_index)
    train_dataloader: Union[CycleIterator, DataLoader] = CycleIterator(train_dataloader)

    for train_data in train_dataloader:
        if resume:
            if curr_iter < initial_iter:
                curr_iter += 1