Function benchmark_inference

examples/inference/llama/benchmark_llama3.py:111–191

(args)

def benchmark_inference(args):
    coordinator = DistCoordinator()

    torch_dtype = TORCH_DTYPE_MAP.get(args.dtype, None)
    config = CONFIG_MAP[args.model]
    config.torch_dtype = torch_dtype
    config.pad_token_id = config.eos_token_id

    if args.model_path is not None:
        model = transformers.LlamaForCausalLM.from_pretrained(args.model_path, torch_dtype=torch_dtype)
        tokenizer = AutoTokenizer.from_pretrained(args.model_path)
    else:
        # Random weights
        model = transformers.LlamaForCausalLM(config)
        tokenizer = AutoTokenizer.from_pretrained("hf-internal-testing/llama-tokenizer")
    if args.dtype == "fp16":
        model = model.half()
    elif args.dtype == "bf16":
        model = model.to(torch.bfloat16)

    inference_config = InferenceConfig(
        dtype=args.dtype,
        max_batch_size=args.batch_size,
        max_input_len=args.max_seq_len,
        max_output_len=args.max_output_len,
        prefill_ratio=1.2,
        block_size=32,
        tp_size=args.tp_size,
        use_cuda_kernel=True,
    )
    engine = InferenceEngine(model, tokenizer, inference_config, verbose=True)

    data = data_gen(args.batch_size, args.max_seq_len)
    generation_config = GenerationConfig(
        pad_token_id=tokenizer.pad_token_id,
        max_length=args.max_seq_len + args.max_output_len,
        # max_new_tokens=args.max_output_len,
    )
    coordinator.print_on_master(f"Generation Config: \n{generation_config.to_dict()}")

    ctx = (
        torch.profiler.profile(
            record_shapes=True,
            with_stack=True,
            with_modules=True,
            activities=[
                torch.profiler.ProfilerActivity.CPU,
                torch.profiler.ProfilerActivity.CUDA,
            ],
            schedule=torch.profiler.schedule(wait=0, warmup=N_WARMUP_STEPS, active=1),
            on_trace_ready=torch.profiler.tensorboard_trace_handler(
                f"./tb_log_{args.batch_size}_{args.max_seq_len}_{args.max_output_len}"
            ),
        )
        if args.profile
        else nullcontext()
    )
    with ctx: