Function build_from_hf

tests/microbenchmarks/build_time_benchmark.py:179–258

Build model and init executor using huggingface model config and fake weights, useful for benchmarking

(args,
                  model_tag,
                  hf_model_dir,
                  dtype,
                  load_weights,
                  tp,
                  pp,
                  rank=0)

def build_from_hf(args,
                  model_tag,
                  hf_model_dir,
                  dtype,
                  load_weights,
                  tp,
                  pp,
                  rank=0):
    '''Build model and init executor using huggingface model config and fake weights, useful for benchmarking
    '''
    status, = cudart.cudaSetDevice(rank)
    assert status == cudart.cudaError_t.cudaSuccess, f"cuda set device to {rank} errored: {status}"
    log_level(args)
    mpi_barrier()
    world_size = tp * pp
    # TODO: Only build 1 rank for now, all the ranks shall have similar build time
    # shall we build all ranks in parallel?
    mapping = Mapping(world_size=world_size, rank=rank, tp_size=tp, pp_size=pp)

    quant_config = None
    if args.quant == 'fp8':
        quant_config = QuantConfig(QuantAlgo.FP8)

    phase_and_time = []
    if load_weights:
        quant_output_dir = tempfile.TemporaryDirectory(model_tag)
        start = time.time()
        if args.quant is None:
            trtllm_model = AutoModelForCausalLM.from_hugging_face(
                hf_model_dir,
                dtype,
                mapping,
                load_model_on_cpu=args.load_to_cpu)
        else:
            model_cls = AutoModelForCausalLM.get_trtllm_model_class(
                hf_model_dir)
            if rank == 0:
                model_cls.quantize(hf_model_dir,
                                   output_dir=quant_output_dir.name,
                                   dtype=args.dtype,
                                   mapping=mapping,
                                   quant_config=quant_config)
            mpi_barrier(
            )  # every rank must wait rank 0 to get the correct quantized checkpoint
            trtllm_model = model_cls.from_checkpoint(quant_output_dir.name)
        phase_and_time.append(('load_and_convert', time.time() - start))
        quant_output_dir.cleanup()

    else:  # fake weights
        trtllm_config = AutoConfig.from_hugging_face(hf_model_dir,
                                                     dtype,
                                                     mapping,
                                                     quant_config,
                                                     trust_remote_code=True)
        trtllm_model = AutoModelForCausalLM.get_trtllm_model_class(
            hf_model_dir)(trtllm_config)

    start = time.time()