From 291987113b816dcd45c0900c71a67d4a1fe30a78 Mon Sep 17 00:00:00 2001
From: Scott Roy <161522778+metascroy@users.noreply.github.com>
Date: Wed, 1 Oct 2025 19:14:32 -0700
Subject: [PATCH] up (#3391)

---
 unsloth/models/_utils.py | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/unsloth/models/_utils.py b/unsloth/models/_utils.py
index 902d3ee6a8..3079196e68 100644
--- a/unsloth/models/_utils.py
+++ b/unsloth/models/_utils.py
@@ -1660,11 +1660,12 @@ def _prepare_model_for_qat(model: torch.nn.Module, qat_scheme: str) -> torch.nn.
     from torchao.quantization import (
         Float8DynamicActivationInt4WeightConfig,
         Float8DynamicActivationFloat8WeightConfig,
-        Int8DynamicActivationInt4WeightConfig,
+        Int8DynamicActivationIntxWeightConfig,
         Int4WeightOnlyConfig,
         PerRow,
         quantize_,
     )
+    from torchao.quantization.granularity import PerGroup
     from torchao.quantization.qat import QATConfig
     filter_fn = None
     if qat_scheme == "fp8-int4":
@@ -1675,7 +1676,7 @@ def _prepare_model_for_qat(model: torch.nn.Module, qat_scheme: str) -> torch.nn.
         base_config = Float8DynamicActivationFloat8WeightConfig(granularity=PerRow())
     elif qat_scheme == "int8-int4":
         group_size = 32
-        base_config = Int8DynamicActivationInt4WeightConfig(group_size=group_size)
+        base_config = Int8DynamicActivationIntxWeightConfig(weight_dtype=torch.int4, weight_granularity=PerGroup(group_size))
         filter_fn = lambda m, _: isinstance(m, torch.nn.Linear) and m.in_features >= group_size
     elif qat_scheme == "int4":
         group_size = 128