From 9e473c10d54bfd9b99e024dec04f8042f319759a Mon Sep 17 00:00:00 2001 From: "s.malakhov" Date: Fri, 31 Jul 2026 14:31:20 +0300 Subject: [PATCH] [quantization] Add BOS token This commit adds BOS token to calibration dataset for llama model. TICO-DCO-1.0-Signed-off-by: s.malakhov --- tico/quantization/recipes/data/llm.py | 8 +++++++- .../wrapq/examples/quantize_full_qmodel_with_gptq.py | 10 ++++++++-- 2 files changed, 15 insertions(+), 3 deletions(-) diff --git a/tico/quantization/recipes/data/llm.py b/tico/quantization/recipes/data/llm.py index 3e4ae6b5..54b49ecf 100644 --- a/tico/quantization/recipes/data/llm.py +++ b/tico/quantization/recipes/data/llm.py @@ -45,7 +45,13 @@ def build_wikitext_calibration_inputs( random.seed(seed) samples: list[torch.Tensor] = [] + bos = getattr(tokenizer, "bos_token_id", None) for _ in range(n_samples): i = random.randint(0, token_ids.shape[1] - seq_len - 1) - samples.append(token_ids[:, i : i + seq_len].cpu()) + w = token_ids[:, i : i + seq_len] + if bos is not None: + w = torch.cat( + [torch.tensor([[bos]], device=w.device), w[:, : seq_len - 1]], dim=1 + ) + samples.append(w.cpu()) return samples diff --git a/tico/quantization/wrapq/examples/quantize_full_qmodel_with_gptq.py b/tico/quantization/wrapq/examples/quantize_full_qmodel_with_gptq.py index 4b5f958f..59ffb485 100644 --- a/tico/quantization/wrapq/examples/quantize_full_qmodel_with_gptq.py +++ b/tico/quantization/wrapq/examples/quantize_full_qmodel_with_gptq.py @@ -1295,10 +1295,16 @@ def build_calibration_inputs( random.seed(args.seed) calib_inputs = [] + bos = getattr(tokenizer, "bos_token_id", None) for _ in range(nsamples): i = random.randint(0, train_ids.shape[1] - seqlen - 1) - j = i + seqlen - calib_inputs.append(train_ids[:, i:j].cpu()) + w = train_ids[:, i : i + seqlen].cpu() + if bos is not None: + w = torch.cat( + [torch.tensor([[bos]], device=w.device), w[:, : seqlen - 1]], dim=1 + ) + + calib_inputs.append(w) return calib_inputs