Repository: hpcaitech/ColossalAI
File: examples/inference/benchmark_llama.py (Python)

Code before:
        if args.dtype in ["fp16", "bf16"]:
            num_bytes = 2
        else:
            num_bytes = 4

        msg += f"Whole batch end2end time: {whole_end2end * 1000:.2f} ms\n"
        msg += f"Whole batch per token latency: {whole_avg_latency * 1000:.2f} ms\n"
        msg += f"Throughput: {total_tokens_num / whole_end2end:.2f} tokens/s\n"
        msg += f"Flops: {num_parameters * num_bytes / whole_avg_latency / 1e12:.2f} TFLOPS\n"

    if torch.cuda.is_available():
        msg += f"-------Memory Summary Device:{get_accelerator().current_device()}-------\n"
        msg += f"Max memory allocated: {get_accelerator().max_memory_allocated() / GIGABYTE:.2f} GB\n"
        msg += f"Max memory reserved: {get_accelerator().max_memory_reserved() / GIGABYTE:.2f} GB\n"

    print(msg)


def benchmark_inference(args):
    with torch.no_grad():
        config = CONFIG_MAP[args.model]
        config.pad_token_id = config.eos_token_id
        # model = transformers.LlamaForCausalLM(config).cuda()
        model = transformers.LlamaForCausalLM.from_pretrained("/home/caidi/llama_model/").cuda()
        model = model.eval()
        tokenizer = AutoTokenizer.from_pretrained("/home/caidi/llama_model/")

        if args.dtype == "fp16":
            model = model.half()
        elif args.dtype == "bf16":
            model = model.to(torch.bfloat16)

        if args.continous_batching:
            mbsz = args.mbsz
        else:
            mbsz = args.batch_size
        if args.mode == "caiinference":
            inference_config = InferenceConfig(
                dtype=args.dtype,
                micro_batch_size=args.mb_size,
                max_batch_size=mbsz,
                max_input_len=args.seq_len,
                max_output_len=args.output_len,
                prefill_ratio=1.2,
                block_size=32,
            )
            engine = InferenceEngine(model, tokenizer, inference_config, verbose=True)
        elif args.mode == "vllm":
            engine = LLM(
                model="/home/caidi/llama_model/",
                max_num_seqs=mbsz,

Review comment (suggestion):
Do not use your actual path. Make it an argument passed via command line.

e.g. 

```python
Model.from_pretrained(args.model_path)
```

Code after:
        if args.dtype in ["fp16", "bf16"]:
            num_bytes = 2
        else:
            num_bytes = 4

        msg += f"Whole batch end2end time: {whole_end2end * 1000:.2f} ms\n"
        msg += f"Whole batch per token latency: {whole_avg_latency * 1000:.2f} ms\n"
        msg += f"Throughput: {total_token_num / whole_end2end:.2f} tokens/s\n"
        msg += f"Flops: {num_parameters * num_bytes / whole_avg_latency / 1e12:.2f} TFLOPS\n"

    if torch.cuda.is_available():
        msg += f"-------Memory Summary Device:{get_accelerator().current_device()}-------\n"
        msg += f"Max memory allocated: {get_accelerator().max_memory_allocated() / GIGABYTE:.2f} GB\n"
        msg += f"Max memory reserved: {get_accelerator().max_memory_reserved() / GIGABYTE:.2f} GB\n"

    print(msg)


def benchmark_inference(args):
    with torch.no_grad():
        config = CONFIG_MAP[args.model]
        config.pad_token_id = config.eos_token_id
        if args.test_random_weight:
            model = transformers.LlamaForCausalLM(config).cuda()
            tokenizer = AutoTokenizer.from_pretrained("hf-internal-testing/llama-tokenizer")
        else:
            assert args.model_path, "When testing pretrained weights, the model path must be provided.'"
            model = transformers.LlamaForCausalLM.from_pretrained(args.model_path).cuda()
            tokenizer = AutoTokenizer.from_pretrained(args.model_path)

        model = model.eval()

        if args.dtype == "fp16":
            model = model.half()
        elif args.dtype == "bf16":
            model = model.to(torch.bfloat16)

        if args.continous_batching:
            mbsz = args.mbsz
        else:
            mbsz = args.batch_size
        if args.mode == "colossalai":
            inference_config = InferenceConfig(
                dtype=args.dtype,
                micro_batch_size=args.mb_size,
                max_batch_size=mbsz,
                max_input_len=args.seq_len,
                max_output_len=args.output_len,
                prefill_ratio=1.2,
                block_size=32,
            )