Function benchmark_inference

benchmark.py:190–236

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

def benchmark_inference(model, dataloader, num_warmup_batches=10, gpu_idx=0):
    model.eval()
    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)
        with torch.no_grad():
            for i, batch in enumerate(dataloader):
                if i >= num_warmup_batches:
                    break
                input_ids, attention_mask, _ = [t.to(device) for t in batch]
                with torch.cuda.amp.autocast(dtype=torch.bfloat16):
                    _ = model(input_ids, attention_mask=attention_mask)
                progress.update(warmup_task, advance=1)

        inference_task = progress.add_task("[cyan]Inference", total=len(dataloader))
        total_time = 0
        with torch.no_grad():
            run_start_time = time.time()
            for i, batch in enumerate(dataloader):
                input_ids, attention_mask, _ = [t.to(device) for t in batch]
                with torch.cuda.amp.autocast(dtype=torch.bfloat16):
                    _ = model(input_ids, attention_mask=attention_mask)
                progress.update(inference_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())
            run_end_time = time.time()
            total_time += run_end_time - run_start_time

    avg_run_time = total_time
    avg_power = np.mean(power_readings)
    max_power = np.max(power_readings)
    return avg_run_time, avg_power, max_power, max_allocated_memory, max_reserved_memory


def create_dummy_data(num_samples, seq_length, vocab_size, model_type):