Repository: lm-sys/FastChat
File: fastchat/serve/inference.py (Python)

Code before:

    return model, tokenizer


def load_compress_model(model_path, device):

    # partially load model
    tokenizer = AutoTokenizer.from_pretrained(model_path, use_fast=False)
    base_pattern = os.path.join(model_path, "pytorch_model-*.bin")
    files = glob.glob(base_pattern)
    
    with init_empty_weights():
        config = AutoConfig.from_pretrained(model_path, low_cpu_mem_usage=True, device_map='auto',torch_dtype=torch.float16)
        model = AutoModelForCausalLM.from_config(config)
        linear_weights = get_compressed_list(model)


    compressed_state_dict = {}

    for file in files:
        tmp_state_dict = torch.load(file, map_location="cpu")
        for name in tqdm(tmp_state_dict):
            if name in linear_weights:
                compressed_state_dict[name] = tmp_state_dict[name].to(device)
                tensor = compressed_state_dict[name].data
                compressed_state_dict[name] = compress(tensor, default_compression_config)
            else:
                compressed_state_dict[name] = tmp_state_dict[name].to(device)
            tmp_state_dict[name] = None
            tensor = None
            gc.collect()
            torch.cuda.empty_cache()

    for name in model.state_dict():
        if name not in linear_weights:
            set_module_tensor_to_device(model, name, device, value=compressed_state_dict[name])
    apply_compressed_weight(model, compressed_state_dict, device)

    model.to(device)

    return model, tokenizer

@torch.inference_mode()
def generate_stream(model, tokenizer, params, device,
                    context_len=2048, stream_interval=2):
    prompt = params["prompt"]
    l_prompt = len(prompt)
    temperature = float(params.get("temperature", 1.0))
    max_new_tokens = int(params.get("max_new_tokens", 256))
    stop_str = params.get("stop", None)
    if stop_str == tokenizer.eos_token:

Review comment (suggestion):
```suggestion
                tensor = tmp_state_dict[name].to(device).data
                compressed_state_dict[name] = compress(tensor, default_compression_config)
```

Code after:
        kwargs = {"torch_dtype": torch.float16}
        # Avoid bugs in mps backend by not using in-place operations.
        replace_llama_attn_with_non_inplace_operations()
    else:
        raise ValueError(f"Invalid device: {device}")
    
    if load_8bit:
        if num_gpus != 1 and num_gpus != "1":
            warnings.warn("8-bit quantization is not supported for multi-gpu inference.")
        else:
            return load_compress_model(model_path=model_path, device=device, torch_dtype=kwargs["torch_dtype"])

    if "chatglm" in model_path:
        tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
        model = AutoModel.from_pretrained(
            model_path, trust_remote_code=True, **kwargs
        ).cuda()
    elif "google/flan-t5" in model_path:
        tokenizer = AutoTokenizer.from_pretrained(model_path, use_fast=False)
        model = AutoModelForSeq2SeqLM.from_pretrained(
            model_path, low_cpu_mem_usage=True, **kwargs
        )
    elif "dolly" in model_path:
        tokenizer = AutoTokenizer.from_pretrained(model_path, use_fast=True)
        model = AutoModelForCausalLM.from_pretrained(
            model_path, low_cpu_mem_usage=True, **kwargs
        )
        # 50277 means "### End"
        tokenizer.eos_token_id = 50277
    elif "pythia" in model_path or "stablelm" in model_path:
        tokenizer = AutoTokenizer.from_pretrained(model_path, use_fast=True)
        model = AutoModelForCausalLM.from_pretrained(
            model_path, low_cpu_mem_usage=True, **kwargs
        )
    else:
        tokenizer = AutoTokenizer.from_pretrained(model_path, use_fast=False)
        model = AutoModelForCausalLM.from_pretrained(
            model_path, low_cpu_mem_usage=True, **kwargs
        )
        raise_warning_for_old_weights(model_path, model)


    if (device == "cuda" and num_gpus == 1) or device == "mps":
        model.to(device)

    if debug:
        print(model)

    return model, tokenizer