Function run_benchmark

benchmarking/inference_benchmark.py:120–163

(args, config, batch_size)

def run_benchmark(args, config, batch_size):
    launcher_config = ProcessConfig(device_isolation=True, device_isolation_action="warn", start_method="spawn")
    scenario_config = InferenceConfig(
        latency=True,
        memory=True,
        input_shapes={"batch_size": batch_size, "sequence_length": args.input_length},
        iterations=args.iterations,
        warmup_runs=args.warmup_runs,
        # set duration to 0 to disable the duration-based stopping criterion
        # this is IMPORTANT to ensure that all benchmarks run the same number of operations, regardless of hardware speed/bottlenecks
        duration=0,
        # for consistent results, set a fixed min and max for output tokens
        generate_kwargs={"min_new_tokens": args.output_length, "max_new_tokens": args.output_length},
        forward_kwargs={"min_new_tokens": args.output_length, "max_new_tokens": args.output_length},
    )

    backend_config = PyTorchConfig(
        device="cuda",
        device_ids="0",
        device_map="auto",
        no_weights=False,
        model=args.model_id,
        **WEIGHTS_CONFIGS[config],
    )

    test_name = (
        f"benchmark-{config}"
        f"-bsz-{batch_size}"
        f"-isz-{args.input_length}"
        f"-osz-{args.output_length}"
        f"-iter-{args.iterations}"
        f"-wrmup-{args.warmup_runs}"
    )
    benchmark_config = BenchmarkConfig(
        name=test_name,
        scenario=scenario_config,
        launcher=launcher_config,
        backend=backend_config,
    )

    out_path = out_dir / (test_name + ".json")
    print(f"[{test_name}] Starting:")
    benchmark_report = Benchmark.launch(benchmark_config)
    benchmark_report.save_json(out_path)


if __name__ == "__main__":