Skip to content

Commit 961e293

Browse files
committed
convert-hf : simplify BitNet pre-quantization
This still results in the exact same tensor weights and scales, but it reveals some weirdness in the current algorithm.
1 parent 89dc3b2 commit 961e293

File tree

1 file changed

+12
-8
lines changed

1 file changed

+12
-8
lines changed

convert-hf-to-gguf.py

Lines changed: 12 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -265,7 +265,10 @@ def write_tensors(self):
265265
break
266266

267267
for new_name, data in ((n, d.squeeze().numpy()) for n, d in self.modify_tensors(data_torch, name, bid)):
268-
data: np.ndarray = data # type hint
268+
data: np.ndarray # type hint
269+
if len(data.shape) == 0:
270+
# otherwise single-value tensors get squeezed
271+
data = data.reshape((1,))
269272
n_dims = len(data.shape)
270273
data_dtype = data.dtype
271274
data_qtype: gguf.GGMLQuantizationType | None = None
@@ -336,7 +339,7 @@ def write_tensors(self):
336339
shape = gguf.quant_shape_from_byte_shape(data.shape, data_qtype) if data.dtype == np.uint8 else data.shape
337340

338341
# reverse shape to make it similar to the internal ggml dimension order
339-
shape_str = f"{{{', '.join(str(n) for n in reversed(shape)) or '1'}}}"
342+
shape_str = f"{{{', '.join(str(n) for n in reversed(shape))}}}"
340343

341344
# n_dims is implicit in the shape
342345
logger.info(f"{f'%-{max_name_len}s' % f'{new_name},'} {old_dtype} --> {data_qtype.name}, shape = {shape_str}")
@@ -1446,12 +1449,13 @@ def set_gguf_parameters(self):
14461449
def weight_quant(self, weight):
14471450
dtype = weight.dtype
14481451
weight = weight.float()
1449-
s = 1 / weight.abs().mean().clamp(min=1e-5)
1450-
weight = (weight * s).round().clamp(-1, 1) / s
1451-
scale = weight.abs().max().unsqueeze(0)
1452-
weight = torch.where(weight.abs().less(1e-6), 0, weight).type(dtype)
1453-
weight = torch.sign(weight).type(dtype)
1454-
return weight.type(dtype), scale.type(torch.float32)
1452+
scale = weight.abs().mean().clamp(min=1e-5)
1453+
iscale = 1 / scale
1454+
weight = (weight * iscale).round().clamp(-1, 1)
1455+
# TODO: use the scale directly instead of inverting it twice
1456+
# (this is also unnecessarily doubly inverted upstream)
1457+
# ref: https://huggingface.co/1bitLLM/bitnet_b1_58-3B/blob/af89e318d78a70802061246bf037199d2fb97020/utils_quant.py#L10
1458+
return weight.type(dtype), (1 / iscale).type(torch.float32)
14551459

14561460
def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:
14571461
new_name = self.map_tensor_name(name)

0 commit comments

Comments
 (0)