Function main

src/infer/infer.py:106–191

(
    dp_size, 
    dp_rank, 
    dp_master_ip, 
    dp_master_port, 
    GPUs_per_dp_rank, 
    model_name_or_path, 
    samples,
    output_path,
    model_name,
    temperature,
    lock,
)

def main(
    dp_size, 
    dp_rank, 
    dp_master_ip, 
    dp_master_port, 
    GPUs_per_dp_rank, 
    model_name_or_path, 
    samples,
    output_path,
    model_name,
    temperature,
    lock,
):
    os.environ["VLLM_DP_RANK"] = str(dp_rank)
    os.environ["VLLM_DP_SIZE"] = str(dp_size)
    os.environ["VLLM_DP_MASTER_IP"] = dp_master_ip
    os.environ["VLLM_DP_MASTER_PORT"] = str(dp_master_port)
    # set devices for each dp_rank
    os.environ["CUDA_VISIBLE_DEVICES"] = ",".join(
        str(i) 
        for i in range(dp_rank * GPUs_per_dp_rank, (dp_rank + 1) * GPUs_per_dp_rank)
    )
    print(f"DP rank {dp_rank} needs to process {len(samples)} prompts")

    # Create a sampling params object.
    # since we are doing data parallel, every rank can have different
    # sampling params. here we set different max_tokens for different
    # ranks for demonstration.
    sampling_params = SamplingParams(
        temperature=temperature,
        top_p=1.0,
        # top_k=50,
        max_tokens=4096
    )

    # Create an LLM.
    llm = LLM(
        model=model_name_or_path,
        tensor_parallel_size=GPUs_per_dp_rank,
        enforce_eager=False,
        # max_model_len=32768,
        gpu_memory_utilization=0.8,
        dtype='bfloat16',
    )
    convs = []
    for sample in samples:
        ctx = []
        for i in range(0, len(sample['context']), 2):
            ctx.append(f'Q{i // 2 + 1}: {sample["context"][i]}')
            ctx.append(f'A{i // 2 + 1}: {sample["context"][i + 1]}')
        convs.append([
            {
                'role': 'user',
                'content': INSTRUCTION.format(
                    context='\n'.join(ctx),
                    query=sample['question']
                )
            }