From e545f04546b1fb944fa51666ab791b9c4d6652d0 Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Wed, 15 Oct 2025 14:28:19 -0700
Subject: [PATCH] Update _utils.py

---
 unsloth/models/_utils.py | 23 +++++++++++++++++++++++
 1 file changed, 23 insertions(+)

diff --git a/unsloth/models/_utils.py b/unsloth/models/_utils.py
index efb6e28345..94f07485aa 100644
--- a/unsloth/models/_utils.py
+++ b/unsloth/models/_utils.py
@@ -1651,6 +1651,14 @@ def error_out_no_vllm(*args, **kwargs):
     raise NotImplementedError("Unsloth: vLLM is not yet supported for fast inference for this model! Please use `.generate` instead")
 
 
+class TorchAOMetadata:
+    def __init__(qat_scheme, base_config, filter_fn, group_size):
+        self.qat_scheme = qat_scheme
+        self.base_config = base_config
+        self.filter_fn = filter_fn
+        self.group_size = group_size
+pass
+
 def _prepare_model_for_qat(model: torch.nn.Module, qat_scheme: str) -> torch.nn.Module:
     """
     Transform a model for Quantization-Aware Training (QAT) during fine-tuning.
@@ -1673,6 +1681,8 @@ def _prepare_model_for_qat(model: torch.nn.Module, qat_scheme: str) -> torch.nn.
     from torchao.quantization.granularity import PerGroup
     from torchao.quantization.qat import QATConfig
     filter_fn = None
+    group_size = None
+    base_config = None
     if qat_scheme == "fp8-int4":
         group_size = 128
         base_config = Float8DynamicActivationInt4WeightConfig()
@@ -1690,6 +1700,19 @@ def _prepare_model_for_qat(model: torch.nn.Module, qat_scheme: str) -> torch.nn.
     else:
         raise ValueError(f"Unexpected QAT scheme {qat_scheme}")
     pass
+    # Save TorchAO schemes
+    torchao_metadata = TorchAOMetadata(
+        qat_scheme = qat_scheme,
+        base_config = base_config,
+        filter_fn = filter_fn,
+        group_size = group_size,
+    )
+    inner_model = model
+    while hasttr(model, "model"):
+        model._torchao_metadata = torchao_metadata
+        model = model.model
+    if hasattr(model, "model"):
+        model._torchao_metadata = torchao_metadata
     quantize_(model, QATConfig(base_config, step="prepare"), filter_fn=filter_fn)
     return model
 pass