(model_path, device, torch_dtype)
| 73 | |
| 74 | |
| 75 | def load_compress_model(model_path, device, torch_dtype): |
| 76 | # partially load model |
| 77 | tokenizer = AutoTokenizer.from_pretrained(model_path, use_fast=False) |
| 78 | base_pattern = os.path.join(model_path, "pytorch_model-*.bin") |
| 79 | files = glob.glob(base_pattern) |
| 80 | |
| 81 | config = AutoConfig.from_pretrained( |
| 82 | model_path, low_cpu_mem_usage=True, torch_dtype=torch_dtype |
| 83 | ) |
| 84 | model = AutoModelForCausalLM.from_config(config) |
| 85 | linear_weights = get_compressed_list(model) |
| 86 | |
| 87 | compressed_state_dict = {} |
| 88 | |
| 89 | for filename in files: |
| 90 | tmp_state_dict = torch.load(filename) |
| 91 | for name in tmp_state_dict: |
| 92 | if name in linear_weights: |
| 93 | tensor = tmp_state_dict[name].to(device).data.to(torch_dtype) |
| 94 | compressed_state_dict[name] = compress( |
| 95 | tensor, default_compression_config |
| 96 | ) |
| 97 | else: |
| 98 | compressed_state_dict[name] = tmp_state_dict[name].to(device) |
| 99 | tmp_state_dict[name] = None |
| 100 | tensor = None |
| 101 | torch.cuda.empty_cache() |
| 102 | |
| 103 | for name, param in model.named_parameters(): |
| 104 | if name not in linear_weights: |
| 105 | param.data = compressed_state_dict[name] |
| 106 | apply_compressed_weight(model, compressed_state_dict, device) |
| 107 | |
| 108 | model.to(device) |
| 109 | |
| 110 | return model, tokenizer |
| 111 | |
| 112 | |
| 113 | def compress(tensor, config): |
no test coverage detected