已合并
fix qkv-extra in load ckpt in 2.0.0 #2705
温一盏创建于 2025年5月24日
fix qkv-extra in load ckpt in 2.0.0 #2705
已合并
从refs/pull/2705/head合入到2.0.0
共 1 个文件变更+5-0
| @@ -470,6 +470,11 @@ class LinearNoTP(torch.nn.Linear): | |||
| 470 | torch.nn.init.xavier_uniform_(self.weight) | 470 | torch.nn.init.xavier_uniform_(self.weight) |
| 471 | torch.random.manual_seed(current_seed) | 471 | torch.random.manual_seed(current_seed) |
| 472 | 472 | ||
| 473 | + self._register_load_state_dict_pre_hook( | ||
| 474 | + lambda state_dict, prefix, *args, **kwargs: state_dict.pop( | ||
| 475 | + f'{prefix}_extra_state', None) | ||
| 476 | + ) | ||
| 477 | + | ||
| 473 | def forward(self, input_): | 478 | def forward(self, input_): |
| 474 | if hasattr(self.weight, "quant_state"): | 479 | if hasattr(self.weight, "quant_state"): |
| 475 | output = bnb.matmul_4bit(input_, self.weight.t(), self.weight.quant_state, bias=self.bias) | 480 | output = bnb.matmul_4bit(input_, self.weight.t(), self.weight.quant_state, bias=self.bias) |