Function benchmark_training

benchmark.py:134–187

(model, dataloader, num_warmup_batches=10, gpu_idx=0)

def benchmark_training(model, dataloader, num_warmup_batches=10, gpu_idx=0):
    model.train()
    optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4)
    device = next(model.parameters()).device

    torch.cuda.reset_peak_memory_stats()

    power_readings = []
    max_allocated_memory = 0
    max_reserved_memory = 0

    with Progress(
        TextColumn("[progress.description]{task.description}"),
        BarColumn(),
        TextColumn("[progress.percentage]{task.percentage:>3.0f}%"),
        TimeRemainingColumn(),
        TimeElapsedColumn(),
    ) as progress:
        warmup_task = progress.add_task("[yellow]Warmup", total=num_warmup_batches)
        for i, batch in enumerate(dataloader):
            if i >= num_warmup_batches:
                break
            input_ids, attention_mask, labels = [t.to(device) for t in batch]
            with torch.cuda.amp.autocast(dtype=torch.bfloat16):
                outputs = model(input_ids, attention_mask=attention_mask, labels=labels)
                loss = outputs.loss
            loss.backward()
            optimizer.step()
            optimizer.zero_grad()
            progress.update(warmup_task, advance=1)

        train_task = progress.add_task("[green]Training", total=len(dataloader))
        total_time = 0
        epoch_start_time = time.time()
        for i, batch in enumerate(dataloader):
            input_ids, attention_mask, labels = [t.to(device) for t in batch]
            with torch.cuda.amp.autocast(dtype=torch.bfloat16):
                outputs = model(input_ids, attention_mask=attention_mask, labels=labels)
                loss = outputs.loss
            loss.backward()
            optimizer.step()
            optimizer.zero_grad()
            progress.update(train_task, advance=1)
            if i % 10 == 0:
                power_readings.append(get_gpu_power(gpu_idx))
                max_allocated_memory = max(max_allocated_memory, torch.cuda.max_memory_allocated())
                max_reserved_memory = max(max_reserved_memory, torch.cuda.max_memory_reserved())
        epoch_end_time = time.time()
        total_time += epoch_end_time - epoch_start_time

    avg_epoch_time = total_time
    avg_power = np.mean(power_readings)
    max_power = np.max(power_readings)
    return avg_epoch_time, avg_power, max_power, max_allocated_memory, max_reserved_memory, loss.item()


def benchmark_inference(model, dataloader, num_warmup_batches=10, gpu_idx=0):