Repository: lm-sys/FastChat
File: fastchat/model/apply_delta.py (Python)

Code before:


    DEFAULT_PAD_TOKEN = "[PAD]"
    base_tokenizer = AutoTokenizer.from_pretrained(base_model_path, use_fast=False)
    num_new_tokens = base_tokenizer.add_special_tokens(dict(pad_token=DEFAULT_PAD_TOKEN))
    print("num_new_tokens: ",num_new_tokens)

    delta_num = 0
    delta_state_dict = torch.load(f"{delta_path}/pytorch_model-0000{str(delta_num%3+1)}-of-00003.bin", map_location=torch.device('cpu'))
    config = AutoConfig.from_pretrained(base_model_path)

    

    print("Applying the delta")
    for i in range(42):
        if i < 10:
            file_name = f"pytorch_model-0000{str(i)}-of-00041.bin"
        else:
            file_name = f"pytorch_model-000{str(i)}-of-00041.bin"
        state_dict_path = f"{base_model_path}/{file_name}"

        state_dict = torch.load(state_dict_path, map_location=torch.device('cpu'))

        for name, param in tqdm(state_dict.items(), desc="Applying delta"):

            while(name not in delta_state_dict):
                delta_num = (delta_num + 1) % 3
                delta_state_dict = torch.load(f"{delta_path}/pytorch_model-0000{str(delta_num%3+1)}-of-00003.bin", map_location=torch.device('cpu'))
                gc.collect()

            if param.shape != delta_state_dict[name].shape:
                new_embeddings = torch.zeros(len(base_tokenizer),state_dict[name].shape[1])
                new_embeddings[:len(base_tokenizer)-num_new_tokens] = state_dict[name]
                state_dict[name] = new_embeddings
                
            state_dict[name] += delta_state_dict[name]

        # 保存更新后的 state_dict
        torch.save(state_dict, f"{target_model_path}/{file_name}")
        state_dict = None 
        gc.collect()
        del state_dict


    print(f"Saving the target model to {target_model_path}")
    delta_tokenizer.save_pretrained(target_model_path)
    config.save_pretrained(target_model_path)
    shutil.copyfile(f"{base_model_path}/pytorch_model.bin.index.json", f"{target_model_path}/pytorch_model.bin.index.json")


if __name__ == "__main__":

Review comment (suggestion):
```suggestion
            while name not in delta_state_dict:
```
style

Code after:


GB = 1 << 30


def split_files(model_path, tmp_path, split_size):
    if not os.path.exists(model_path):
        model_path = snapshot_download(repo_id=model_path)
    if not os.path.exists(tmp_path):
        os.makedirs(tmp_path)

    file_pattern = os.path.join(model_path, "pytorch_model-*.bin")
    files = glob.glob(file_pattern)

    part = 0
    for file_path in tqdm(files):
        state_dict = torch.load(file_path)
        new_state_dict = {}

        current_size = 0

        for name, param in state_dict.items():
            param_size = param.numel() * param.element_size()

            if current_size + param_size > split_size:
                new_file_name = f"pytorch_model-{part}.bin"
                new_file_path = os.path.join(tmp_path, new_file_name)
                torch.save(new_state_dict, new_file_path)
                current_size = 0
                new_state_dict = None
                gc.collect()
                new_state_dict = {}
                part += 1

            new_state_dict[name] = param
            current_size += param_size

        new_file_name = f"pytorch_model-{part}.bin"
        new_file_path = os.path.join(tmp_path, new_file_name)
        torch.save(new_state_dict, new_file_path)
        new_state_dict = None
        gc.collect()
        new_state_dict = {}
        part += 1


def apply_delta_low_cpu_mem(base_model_path, target_model_path, delta_path):
    base_tokenizer = AutoTokenizer.from_pretrained(base_model_path, use_fast=False)
    base_config = AutoConfig.from_pretrained(base_model_path)

    if os.path.exists(target_model_path):