From dcd5d8bbdd64c8d8ec09e0d63d1bf65887ab4e5c Mon Sep 17 00:00:00 2001
From: Pujitha Paladugu <10557236+pujitha24@users.noreply.github.com>
Date: Sun, 6 Sep 2026 13:26:09 -0700
Subject: [PATCH] fix: emit real I2_S tensors for offline-quantized LlamaModel
 checkpoints

Motivation:
`convert-hf-to-gguf-bitnet.py --outtype i2_s` mishandled offline-quantized
LlamaForCausalLM BitNet checkpoints (ternary weights packed as uint8 plus a
separate `<name>.weight_scale` tensor): the scale division was applied
unconditionally, and the I2_S quantization dispatch only had TL1/TL2
branches, so these tensors fell through to the F16 default instead of
producing real I2_S data.

Approach:
In `LlamaModel`, skip the unconditional `scale_map` division when the
requested ftype is I2_S, so the ternary values stay undivided. Add an I2_S
branch to the quantization dispatch, gated on the tensor name being present
in `scale_map`, which calls `quantize_to_i2_s` with the corresponding
scale as `override_scale`. Tensors not present in `scale_map` (ordinary,
non-offline-quantized checkpoints) still fall through to the existing
F16 path, so the division and dispatch changes only affect genuinely
offline-quantized tensors.

Report: https://github.com/microsoft/BitNet/issues/621
Signed-off-by: Pujitha Paladugu <10557236+pujitha24@users.noreply.github.com>
Assisted-by: claude-sonnet-5 (via Claude Code)
---
 utils/convert-hf-to-gguf-bitnet.py | 13 +++++++++++--
 1 file changed, 11 insertions(+), 2 deletions(-)

diff --git a/utils/convert-hf-to-gguf-bitnet.py b/utils/convert-hf-to-gguf-bitnet.py
index b11e831b9..107f29ae3 100644
--- a/utils/convert-hf-to-gguf-bitnet.py
+++ b/utils/convert-hf-to-gguf-bitnet.py
@@ -804,7 +804,9 @@ def write_tensors(self):
                 data_torch = data_torch.unsqueeze(0).expand((4, *origin_shape)) >> shift
                 data_torch = data_torch & 3
                 data_torch = (data_torch.float() - 1).reshape((origin_shape[0] * 4, *origin_shape[1:]))
-                data_torch = data_torch / scale_map[name.replace(".weight", "")].float()
+                # For I2_S output: keep as ternary {-1,0,1}, scale is passed separately to quantize_to_i2_s
+                if self.ftype != gguf.GGMLQuantizationType.I2_S:
+                    data_torch = data_torch / scale_map[name.replace(".weight", "")].float()
 
             # use the first number-like part of the tensor name as the block id
             bid = None
@@ -866,7 +868,14 @@ def write_tensors(self):
 
                 i2_scale = None
                 if self.ftype != gguf.GGMLQuantizationType.F32 and extra_f16 and not extra_f32:
-                    if self.ftype == gguf.GGMLQuantizationType.TL1 and suit_i2:
+                    if self.ftype == gguf.GGMLQuantizationType.I2_S and suit_i2 and name.replace(".weight", "") in scale_map:
+                        # Only tensors from offline-quantized checkpoints (uint8 + weight_scale) are
+                        # already ternary; a plain LlamaForCausalLM checkpoint has no such scale and
+                        # must not be blindly ternarized here, so it falls through to the F16 branch below.
+                        data_qtype = gguf.GGMLQuantizationType.I2_S
+                        override_scale = scale_map[name.replace(".weight", "")].item()
+                        data = quantize_to_i2_s(data, override_scale=override_scale)
+                    elif self.ftype == gguf.GGMLQuantizationType.TL1 and suit_i2:
                         data, i2_scale = transform_to_tl1(data)
                         assert data.dtype == np.uint8
                         assert i2_scale.dtype == np.float32