Buckets:

Quantumbraid
/

trainingdata

Files

xet

Quantumbraid/trainingdata / toolkit /util /ostris_quant.py

Quantumbraid

about 1 month ago

download

raw

17.1 kB

"""
Quantization-agnostic custom quantized linear.
OstrisLinear is a drop-in nn.Linear replacement whose weight is held by a pluggable
quantizer backend (OstrisQuantizer). Backends own the quantized representation
(buffers + per-module attributes) and how the forward pass computes W x from it; the
module and the rest of the toolkit stay backend agnostic. The first backend is
OrbitQuant (toolkit/util/orbit_quant.py) via the orbit2/orbit3/orbit4 qtypes; add new
backends by implementing OstrisQuantizer and resolving them in get_ostris_quantizer.
Modules are converted in place by convert_linear_to_ostris via class swap, so the
original module object (and any references to it, e.g. LoRA org_module or parent
containers) stays valid.
"""

from typing import Dict, Optional

import torch
import torch.nn.functional as F


class OstrisQuantizer:
    """Base class for weight quantization backends used by OstrisLinear.
    Backends are stateless with respect to tensors: everything tensor-shaped must be
    registered as a buffer on the module inside quantize_ (so device moves and dtype
    casts through nn.Module._apply keep working), and read back off the module in the
    other methods. One backend instance may be shared by many modules.
    """

    # the qtype string this instance was resolved from (stamped by
    # get_ostris_quantizer); quantized saves need it to restore the backend
    qtype: Optional[str] = None

    # backends that quantize in the weight's own dtype can set this False to
    # receive the raw weight tensor in quantize_ instead of a float32 copy,
    # avoiding a 2x-weight-size allocation during model quantization
    wants_fp32_weight: bool = True

    def can_quantize(self, module: torch.nn.Linear) -> bool:
        """Whether this backend can quantize the given linear (e.g. shape constraints)."""
        return True

    def quantize_(self, module: torch.nn.Linear, weight_fp32: torch.Tensor) -> None:
        """Build the quantized representation of weight_fp32 and attach it to the
        module (register_buffer for tensors, plain attributes for scalars). Called
        while the module is still an nn.Linear, before the weight param is removed."""
        raise NotImplementedError

    def dequantize(self, module: "OstrisLinear") -> torch.Tensor:
        """Reconstruct the full weight in the original basis, in float32."""
        raise NotImplementedError

    def dequantize_folded(self, module: "OstrisLinear") -> torch.Tensor:
        """The full weight with any ACTIVATION-side transform folded in, i.e. a
        weight that computes the same output on raw activations. Used when
        re-quantizing into a different backend (which won't know about this
        backend's activation transforms). Default: same as dequantize."""
        return self.dequantize(module)

    def requantize_(self, module: "OstrisLinear", fp_weight: torch.Tensor) -> None:
        """Re-quantize in place from a full precision weight in the original basis
        (used by the continuous merge/reset method)."""
        raise NotImplementedError

    def forward(self, module: "OstrisLinear", x: torch.Tensor) -> torch.Tensor:
        # default: dequantize per forward and run a plain linear. backends can
        # override with a cheaper formulation. the weight is frozen, so build it
        # outside autograd; gradients still flow to x through the matmul
        with torch.no_grad():
            w = self.dequantize(module).to(x.dtype)
        return F.linear(x, w, module.bias)


_wrong_device_warned = False


class OstrisLinear(torch.nn.Linear):
    """A linear layer whose weight is quantized by an OstrisQuantizer backend.
    Never instantiate directly: created in place by convert_linear_to_ostris. The
    weight parameter is removed; the quantized representation lives in backend-owned
    buffers, plus:
      ostris_quantizer   the backend instance
      ostris_orig_dtype  dtype of the original weight (used for dequantized views)
    """

    is_ostris_quantized = True

    @torch.no_grad()
    def dequantize_weight(self) -> torch.Tensor:
        """Reconstruct the weight in the original basis and dtype."""
        return self.ostris_quantizer.dequantize(self).to(self.ostris_orig_dtype)

    @property
    def weight(self):
        # materializes the full dequantized weight. kept for code that inspects the
        # weight (shape/dtype/device) and for the network merge paths, which detect
        # the marker via toolkit.util.quantize.is_quantized_tensor
        w = self.dequantize_weight()
        w._is_ostris_weight = True
        return w

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        if x.is_cuda and not hasattr(self, "_layer_memory_manager"):
            # a module left behind on the wrong device (usually cpu after a
            # low_vram load) would run its dequant/matmul on cpu threads and
            # silently hammer the cpu. move it to the input's device once.
            # memory-managed modules are excluded: they keep cpu buffers by
            # design and stage them to the gpu per forward.
            buf = next((b for b in self._buffers.values() if b is not None), None)
            if buf is not None and buf.device != x.device:
                global _wrong_device_warned
                if not _wrong_device_warned:
                    _wrong_device_warned = True
                    from toolkit.print import print_acc
                    print_acc(
                        f"OstrisLinear: quantized weights found on {buf.device} while the "
                        f"input is on {x.device}; moving them to {x.device}. This usually "
                        f"means something left the model behind after a low_vram load."
                    )
                self.to(x.device)
        return self.ostris_quantizer.forward(self, x)

    @torch.no_grad()
    def requantize_(self, fp_weight: torch.Tensor) -> None:
        self.ostris_quantizer.requantize_(self, fp_weight)

    def _save_to_state_dict(self, destination, prefix, keep_vars):
        # emit a lazy stand-in instead of the materialized weight: a full-model
        # state_dict() would otherwise hold every layer's dequantized weight on the
        # gpu at once (OOM on large models). save loops dequantize per key via
        # dequantize_if_quantized, mirroring how torchao state dicts are consumed
        destination[prefix + "weight"] = OstrisLazyWeight(self)
        if self.bias is not None:
            destination[prefix + "bias"] = (
                self.bias if keep_vars else self.bias.detach()
            )


class OstrisLazyWeight(torch.Tensor):
    """Lazy weight stand-in emitted by OstrisLinear._save_to_state_dict.
    Reports the real shape/dtype/device but holds no data; .dequantize()
    materializes the full weight from the module's quantized buffers. This is a
    live view of the module (not a snapshot), which is fine for the save paths
    state dicts feed: they consume each key once via dequantize_if_quantized.
    Any other tensor op falls through __torch_dispatch__ and materializes first.
    """

    @staticmethod
    def __new__(cls, module: "OstrisLinear"):
        buf = next(b for b in module._buffers.values() if b is not None)
        r = torch.Tensor._make_wrapper_subclass(
            cls,
            (module.out_features, module.in_features),
            dtype=module.ostris_orig_dtype,
            device=buf.device,
            requires_grad=False,
        )
        r._ostris_module = module
        # routes is_quantized_tensor/dequantize_if_quantized (toolkit/util/quantize.py)
        r._is_ostris_weight = True
        return r

    def dequantize(self) -> torch.Tensor:
        return self._ostris_module.dequantize_weight()

    def __repr__(self):
        return (
            f"OstrisLazyWeight(shape={tuple(self.shape)}, dtype={self.dtype}, "
            f"device={self.device})"
        )

    @classmethod
    def __torch_dispatch__(cls, func, types, args=(), kwargs=None):
        from torch.utils._pytree import tree_map

        def unwrap(t):
            return t._ostris_module.dequantize_weight() if isinstance(t, cls) else t

        return func(*tree_map(unwrap, args), **tree_map(unwrap, kwargs or {}))


def get_ostris_quantizer(qtype: str) -> Optional[OstrisQuantizer]:
    """Resolve a qtype string to a quantizer backend instance, or None if the qtype
    does not belong to a custom backend. Add new backends here."""
    from toolkit.util.convrot_quant import CONVROT_QTYPES, get_convrot_quantizer
    from toolkit.util.float8_quant import FLOAT8_QTYPES, Float8Quantizer
    from toolkit.util.nvfp4_quant import NVFP4_QTYPES, Nvfp4Quantizer
    from toolkit.util.orbit_quant import ORBIT_QTYPES, OrbitQuantizer
    from toolkit.util.orbit_vq_quant import ORBIT_VQ_QTYPES, OrbitVQQuantizer
    from toolkit.util.uintx_quant import UINTX_QTYPES, UIntXQuantizer

    quantizer = None
    if qtype in ORBIT_QTYPES:
        quantizer = OrbitQuantizer(ORBIT_QTYPES[qtype])
    elif qtype in ORBIT_VQ_QTYPES:
        quantizer = OrbitVQQuantizer(**ORBIT_VQ_QTYPES[qtype])
    elif qtype in CONVROT_QTYPES:
        quantizer = get_convrot_quantizer(qtype)
    elif qtype in NVFP4_QTYPES:
        quantizer = Nvfp4Quantizer()
    elif qtype in FLOAT8_QTYPES:
        quantizer = Float8Quantizer()
    elif qtype in UINTX_QTYPES:
        quantizer = UIntXQuantizer(UINTX_QTYPES[qtype])
    if quantizer is not None:
        # quantized saves read this back to restore the backend on load
        quantizer.qtype = qtype
    return quantizer


# metadata key shared with toolkit/models/classes/_mixin.py save_quantized
QUANT_LAYERS_METADATA_KEY = "aitk_quantization"


@torch.no_grad()
def save_quantized_layers(
    modules: Dict[str, "OstrisLinear"],
    file_path: str,
    metadata: Optional[Dict[str, str]] = None,
    extra_tensors: Optional[Dict[str, torch.Tensor]] = None,
) -> None:
    """Save a set of quantized linears (keyed by their submodule path in the
    target model) as a single safetensors file: backend buffers + bias per
    module, with the qtype/attrs needed to restore them recorded in the file
    metadata. Same layout as OstrisModelMixin.save_quantized, but partial —
    apply with load_quantized_layers."""
    import json
    from safetensors.torch import save_file

    quant_map = {}
    state_dict = {}
    for name, module in modules.items():
        if module.ostris_quantizer.qtype is None:
            raise ValueError(
                f"Cannot save quantized module '{name}': its quantizer has no "
                "qtype recorded (was it created through get_ostris_quantizer?)."
            )
        entry = {
            "qtype": module.ostris_quantizer.qtype,
            "dtype": str(module.ostris_orig_dtype).replace("torch.", ""),
            "buffers": [],
            "attrs": {},
        }
        if module.bias is not None:
            state_dict[f"{name}.bias"] = module.bias
        for buf_name, buf in module._buffers.items():
            if buf is None:
                continue
            state_dict[f"{name}.{buf_name}"] = buf
            entry["buffers"].append(buf_name)
        for attr, value in vars(module).items():
            if attr.startswith("_") or attr in (
                "training", "in_features", "out_features",
                "ostris_quantizer", "ostris_orig_dtype",
            ):
                continue
            if isinstance(value, (bool, int, float, str)):
                entry["attrs"][attr] = value
        quant_map[name] = entry

    if extra_tensors:
        # training-state extras (e.g. QAT master weights); load_quantized_layers
        # ignores them, so deployment loads are unaffected
        state_dict.update(extra_tensors)
    state_dict = {
        k: v.detach().to("cpu", copy=True).contiguous() for k, v in state_dict.items()
    }
    meta = dict(metadata or {})
    meta[QUANT_LAYERS_METADATA_KEY] = json.dumps(
        {"modules": quant_map, "layers_only": True}
    )
    save_file(state_dict, file_path, metadata=meta)


@torch.no_grad()
def load_quantized_layers(root: torch.nn.Module, file_path: str) -> int:
    """Apply a file written by save_quantized_layers onto a model: restores the
    backend buffers (and bias) of each recorded module. Target modules may
    already be OstrisLinear (model quantized on load — buffers are replaced) or
    still plain nn.Linear (converted in place, no quantization math needed).
    Returns the number of modules restored."""
    import json
    from safetensors import safe_open
    from safetensors.torch import load_file

    with safe_open(file_path, framework="pt", device="cpu") as f:
        meta = f.metadata() or {}
    if QUANT_LAYERS_METADATA_KEY not in meta:
        raise ValueError(f"{file_path} has no quantized-layer metadata")
    quant_map = json.loads(meta[QUANT_LAYERS_METADATA_KEY])["modules"]
    state_dict = load_file(file_path)

    for name, entry in quant_map.items():
        module = root.get_submodule(name)
        quantizer = get_ostris_quantizer(entry["qtype"])
        if quantizer is None:
            raise ValueError(f"Unknown qtype '{entry['qtype']}' in {file_path}")
        # figure out the device the module currently lives on
        ref = next(
            (t for t in module._parameters.values() if t is not None),
            next((t for t in module._buffers.values() if t is not None), None),
        )
        device = ref.device if ref is not None else torch.device("cpu")
        if isinstance(module, OstrisLinear):
            # replacing an existing quantized state (possibly another backend):
            # its buffers are exclusively backend state, drop them all
            module._buffers.clear()
        else:
            if "weight" in module._parameters:
                del module._parameters["weight"]
            module.ostris_orig_dtype = getattr(torch, entry["dtype"])
            module.__class__ = OstrisLinear
        for buf_name in entry["buffers"]:
            module.register_buffer(
                buf_name,
                state_dict.pop(f"{name}.{buf_name}").to(device),
                persistent=False,
            )
        for attr, value in entry.get("attrs", {}).items():
            setattr(module, attr, value)
        module.ostris_quantizer = quantizer
        bias_key = f"{name}.bias"
        if bias_key in state_dict and module.bias is not None:
            module.bias.data.copy_(
                state_dict.pop(bias_key).to(device, module.bias.dtype)
            )
        if module.bias is not None:
            module.bias.requires_grad_(False)
    return len(quant_map)


@torch.no_grad()
def convert_linear_to_ostris(
    module: torch.nn.Linear, quantizer: OstrisQuantizer
) -> bool:
    """Quantize an nn.Linear in place (class swap). Returns True if the module was
    converted (or already was), False if it is not a candidate.
    A module that is ALREADY quantized (e.g. loaded from a pre-quantized
    checkpoint) is re-quantized into the requested backend when the qtypes
    differ: the weight is dequantized and re-quantized layer by layer, so the
    full-precision transient never exceeds one layer's weight. Same qtype is
    a no-op (the shipped quantization is kept)."""
    if isinstance(module, OstrisLinear):
        current_qtype = getattr(module.ostris_quantizer, "qtype", None)
        if quantizer.qtype is None or current_qtype == quantizer.qtype:
            return True
        if not quantizer.can_quantize(module):
            return True  # keep the existing quantization rather than dropping it
        # fold any activation-side transform (e.g. an AWQ pre_quant_scale) into
        # the weight so the new backend computes the same function on raw inputs
        weight = module.ostris_quantizer.dequantize_folded(module).to(
            module.ostris_orig_dtype
        )
        # backend state lives exclusively in buffers; leftover scalar attrs from
        # the old backend are inert
        module._buffers.clear()
        if quantizer.wants_fp32_weight:
            quantizer.quantize_(module, weight.to(torch.float32))
        else:
            quantizer.quantize_(module, weight)
        del weight
        module.ostris_quantizer = quantizer
        return True
    weight = getattr(module, "weight", None)
    if not isinstance(weight, torch.nn.Parameter) or not weight.dtype.is_floating_point:
        return False
    if type(weight.data) is not torch.Tensor:
        # already holds a quantized tensor subclass (e.g. torchao)
        return False
    if not quantizer.can_quantize(module):
        return False
    if quantizer.wants_fp32_weight:
        quantizer.quantize_(module, weight.data.to(torch.float32))
    else:
        quantizer.quantize_(module, weight.data)
    module.ostris_quantizer = quantizer
    module.ostris_orig_dtype = weight.dtype
    del module._parameters["weight"]
    if module.bias is not None:
        module.bias.requires_grad_(False)
    module.__class__ = OstrisLinear
    return True

Xet Storage Details

Size:: 17.1 kB
Xet hash:: 197b4939f24686727b405804b031e55f3216555bb334e7090c4e914001581d50

Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.

	"""
	Quantization-agnostic custom quantized linear.

	OstrisLinear is a drop-in nn.Linear replacement whose weight is held by a pluggable
	quantizer backend (OstrisQuantizer). Backends own the quantized representation
	(buffers + per-module attributes) and how the forward pass computes W x from it; the
	module and the rest of the toolkit stay backend agnostic. The first backend is
	OrbitQuant (toolkit/util/orbit_quant.py) via the orbit2/orbit3/orbit4 qtypes; add new
	backends by implementing OstrisQuantizer and resolving them in get_ostris_quantizer.

	Modules are converted in place by convert_linear_to_ostris via class swap, so the
	original module object (and any references to it, e.g. LoRA org_module or parent
	containers) stays valid.
	"""

	from typing import Dict, Optional

	import torch
	import torch.nn.functional as F


	class OstrisQuantizer:
	"""Base class for weight quantization backends used by OstrisLinear.

	Backends are stateless with respect to tensors: everything tensor-shaped must be
	registered as a buffer on the module inside quantize_ (so device moves and dtype
	casts through nn.Module._apply keep working), and read back off the module in the
	other methods. One backend instance may be shared by many modules.
	"""

	# the qtype string this instance was resolved from (stamped by
	# get_ostris_quantizer); quantized saves need it to restore the backend
	qtype: Optional[str] = None

	# backends that quantize in the weight's own dtype can set this False to
	# receive the raw weight tensor in quantize_ instead of a float32 copy,
	# avoiding a 2x-weight-size allocation during model quantization
	wants_fp32_weight: bool = True

	def can_quantize(self, module: torch.nn.Linear) -> bool:
	"""Whether this backend can quantize the given linear (e.g. shape constraints)."""
	return True

	def quantize_(self, module: torch.nn.Linear, weight_fp32: torch.Tensor) -> None:
	"""Build the quantized representation of weight_fp32 and attach it to the
	module (register_buffer for tensors, plain attributes for scalars). Called
	while the module is still an nn.Linear, before the weight param is removed."""
	raise NotImplementedError

	def dequantize(self, module: "OstrisLinear") -> torch.Tensor:
	"""Reconstruct the full weight in the original basis, in float32."""
	raise NotImplementedError

	def dequantize_folded(self, module: "OstrisLinear") -> torch.Tensor:
	"""The full weight with any ACTIVATION-side transform folded in, i.e. a
	weight that computes the same output on raw activations. Used when
	re-quantizing into a different backend (which won't know about this
	backend's activation transforms). Default: same as dequantize."""
	return self.dequantize(module)

	def requantize_(self, module: "OstrisLinear", fp_weight: torch.Tensor) -> None:
	"""Re-quantize in place from a full precision weight in the original basis
	(used by the continuous merge/reset method)."""
	raise NotImplementedError

	def forward(self, module: "OstrisLinear", x: torch.Tensor) -> torch.Tensor:
	# default: dequantize per forward and run a plain linear. backends can
	# override with a cheaper formulation. the weight is frozen, so build it
	# outside autograd; gradients still flow to x through the matmul
	with torch.no_grad():
	w = self.dequantize(module).to(x.dtype)
	return F.linear(x, w, module.bias)


	_wrong_device_warned = False


	class OstrisLinear(torch.nn.Linear):
	"""A linear layer whose weight is quantized by an OstrisQuantizer backend.

	Never instantiate directly: created in place by convert_linear_to_ostris. The
	weight parameter is removed; the quantized representation lives in backend-owned
	buffers, plus:
	ostris_quantizer the backend instance
	ostris_orig_dtype dtype of the original weight (used for dequantized views)
	"""

	is_ostris_quantized = True

	@torch.no_grad()
	def dequantize_weight(self) -> torch.Tensor:
	"""Reconstruct the weight in the original basis and dtype."""
	return self.ostris_quantizer.dequantize(self).to(self.ostris_orig_dtype)

	@property
	def weight(self):
	# materializes the full dequantized weight. kept for code that inspects the
	# weight (shape/dtype/device) and for the network merge paths, which detect
	# the marker via toolkit.util.quantize.is_quantized_tensor
	w = self.dequantize_weight()
	w._is_ostris_weight = True
	return w

	def forward(self, x: torch.Tensor) -> torch.Tensor:
	if x.is_cuda and not hasattr(self, "_layer_memory_manager"):
	# a module left behind on the wrong device (usually cpu after a
	# low_vram load) would run its dequant/matmul on cpu threads and
	# silently hammer the cpu. move it to the input's device once.
	# memory-managed modules are excluded: they keep cpu buffers by
	# design and stage them to the gpu per forward.
	buf = next((b for b in self._buffers.values() if b is not None), None)
	if buf is not None and buf.device != x.device:
	global _wrong_device_warned
	if not _wrong_device_warned:
	_wrong_device_warned = True
	from toolkit.print import print_acc
	print_acc(
	f"OstrisLinear: quantized weights found on {buf.device} while the "
	f"input is on {x.device}; moving them to {x.device}. This usually "
	f"means something left the model behind after a low_vram load."
	)
	self.to(x.device)
	return self.ostris_quantizer.forward(self, x)

	@torch.no_grad()
	def requantize_(self, fp_weight: torch.Tensor) -> None:
	self.ostris_quantizer.requantize_(self, fp_weight)

	def _save_to_state_dict(self, destination, prefix, keep_vars):
	# emit a lazy stand-in instead of the materialized weight: a full-model
	# state_dict() would otherwise hold every layer's dequantized weight on the
	# gpu at once (OOM on large models). save loops dequantize per key via
	# dequantize_if_quantized, mirroring how torchao state dicts are consumed
	destination[prefix + "weight"] = OstrisLazyWeight(self)
	if self.bias is not None:
	destination[prefix + "bias"] = (
	self.bias if keep_vars else self.bias.detach()
	)


	class OstrisLazyWeight(torch.Tensor):
	"""Lazy weight stand-in emitted by OstrisLinear._save_to_state_dict.

	Reports the real shape/dtype/device but holds no data; .dequantize()
	materializes the full weight from the module's quantized buffers. This is a
	live view of the module (not a snapshot), which is fine for the save paths
	state dicts feed: they consume each key once via dequantize_if_quantized.
	Any other tensor op falls through __torch_dispatch__ and materializes first.
	"""

	@staticmethod
	def __new__(cls, module: "OstrisLinear"):
	buf = next(b for b in module._buffers.values() if b is not None)
	r = torch.Tensor._make_wrapper_subclass(
	cls,
	(module.out_features, module.in_features),
	dtype=module.ostris_orig_dtype,
	device=buf.device,
	requires_grad=False,
	)
	r._ostris_module = module
	# routes is_quantized_tensor/dequantize_if_quantized (toolkit/util/quantize.py)
	r._is_ostris_weight = True
	return r

	def dequantize(self) -> torch.Tensor:
	return self._ostris_module.dequantize_weight()

	def __repr__(self):
	return (
	f"OstrisLazyWeight(shape={tuple(self.shape)}, dtype={self.dtype}, "
	f"device={self.device})"
	)

	@classmethod
	def __torch_dispatch__(cls, func, types, args=(), kwargs=None):
	from torch.utils._pytree import tree_map

	def unwrap(t):
	return t._ostris_module.dequantize_weight() if isinstance(t, cls) else t

	return func(tree_map(unwrap, args), *tree_map(unwrap, kwargs or {}))


	def get_ostris_quantizer(qtype: str) -> Optional[OstrisQuantizer]:
	"""Resolve a qtype string to a quantizer backend instance, or None if the qtype
	does not belong to a custom backend. Add new backends here."""
	from toolkit.util.convrot_quant import CONVROT_QTYPES, get_convrot_quantizer
	from toolkit.util.float8_quant import FLOAT8_QTYPES, Float8Quantizer
	from toolkit.util.nvfp4_quant import NVFP4_QTYPES, Nvfp4Quantizer
	from toolkit.util.orbit_quant import ORBIT_QTYPES, OrbitQuantizer
	from toolkit.util.orbit_vq_quant import ORBIT_VQ_QTYPES, OrbitVQQuantizer
	from toolkit.util.uintx_quant import UINTX_QTYPES, UIntXQuantizer

	quantizer = None
	if qtype in ORBIT_QTYPES:
	quantizer = OrbitQuantizer(ORBIT_QTYPES[qtype])
	elif qtype in ORBIT_VQ_QTYPES:
	quantizer = OrbitVQQuantizer(**ORBIT_VQ_QTYPES[qtype])
	elif qtype in CONVROT_QTYPES:
	quantizer = get_convrot_quantizer(qtype)
	elif qtype in NVFP4_QTYPES:
	quantizer = Nvfp4Quantizer()
	elif qtype in FLOAT8_QTYPES:
	quantizer = Float8Quantizer()
	elif qtype in UINTX_QTYPES:
	quantizer = UIntXQuantizer(UINTX_QTYPES[qtype])
	if quantizer is not None:
	# quantized saves read this back to restore the backend on load
	quantizer.qtype = qtype
	return quantizer


	# metadata key shared with toolkit/models/classes/_mixin.py save_quantized
	QUANT_LAYERS_METADATA_KEY = "aitk_quantization"


	@torch.no_grad()
	def save_quantized_layers(
	modules: Dict[str, "OstrisLinear"],
	file_path: str,
	metadata: Optional[Dict[str, str]] = None,
	extra_tensors: Optional[Dict[str, torch.Tensor]] = None,
	) -> None:
	"""Save a set of quantized linears (keyed by their submodule path in the
	target model) as a single safetensors file: backend buffers + bias per
	module, with the qtype/attrs needed to restore them recorded in the file
	metadata. Same layout as OstrisModelMixin.save_quantized, but partial —
	apply with load_quantized_layers."""
	import json
	from safetensors.torch import save_file

	quant_map = {}
	state_dict = {}
	for name, module in modules.items():
	if module.ostris_quantizer.qtype is None:
	raise ValueError(
	f"Cannot save quantized module '{name}': its quantizer has no "
	"qtype recorded (was it created through get_ostris_quantizer?)."
	)
	entry = {
	"qtype": module.ostris_quantizer.qtype,
	"dtype": str(module.ostris_orig_dtype).replace("torch.", ""),
	"buffers": [],
	"attrs": {},
	}
	if module.bias is not None:
	state_dict[f"{name}.bias"] = module.bias
	for buf_name, buf in module._buffers.items():
	if buf is None:
	continue
	state_dict[f"{name}.{buf_name}"] = buf
	entry["buffers"].append(buf_name)
	for attr, value in vars(module).items():
	if attr.startswith("_") or attr in (
	"training", "in_features", "out_features",
	"ostris_quantizer", "ostris_orig_dtype",
	):
	continue
	if isinstance(value, (bool, int, float, str)):
	entry["attrs"][attr] = value
	quant_map[name] = entry

	if extra_tensors:
	# training-state extras (e.g. QAT master weights); load_quantized_layers
	# ignores them, so deployment loads are unaffected
	state_dict.update(extra_tensors)
	state_dict = {
	k: v.detach().to("cpu", copy=True).contiguous() for k, v in state_dict.items()
	}
	meta = dict(metadata or {})
	meta[QUANT_LAYERS_METADATA_KEY] = json.dumps(
	{"modules": quant_map, "layers_only": True}
	)
	save_file(state_dict, file_path, metadata=meta)


	@torch.no_grad()
	def load_quantized_layers(root: torch.nn.Module, file_path: str) -> int:
	"""Apply a file written by save_quantized_layers onto a model: restores the
	backend buffers (and bias) of each recorded module. Target modules may
	already be OstrisLinear (model quantized on load — buffers are replaced) or
	still plain nn.Linear (converted in place, no quantization math needed).
	Returns the number of modules restored."""
	import json
	from safetensors import safe_open
	from safetensors.torch import load_file

	with safe_open(file_path, framework="pt", device="cpu") as f:
	meta = f.metadata() or {}
	if QUANT_LAYERS_METADATA_KEY not in meta:
	raise ValueError(f"{file_path} has no quantized-layer metadata")
	quant_map = json.loads(meta[QUANT_LAYERS_METADATA_KEY])["modules"]
	state_dict = load_file(file_path)

	for name, entry in quant_map.items():
	module = root.get_submodule(name)
	quantizer = get_ostris_quantizer(entry["qtype"])
	if quantizer is None:
	raise ValueError(f"Unknown qtype '{entry['qtype']}' in {file_path}")
	# figure out the device the module currently lives on
	ref = next(
	(t for t in module._parameters.values() if t is not None),
	next((t for t in module._buffers.values() if t is not None), None),
	)
	device = ref.device if ref is not None else torch.device("cpu")
	if isinstance(module, OstrisLinear):
	# replacing an existing quantized state (possibly another backend):
	# its buffers are exclusively backend state, drop them all
	module._buffers.clear()
	else:
	if "weight" in module._parameters:
	del module._parameters["weight"]
	module.ostris_orig_dtype = getattr(torch, entry["dtype"])
	module.__class__ = OstrisLinear
	for buf_name in entry["buffers"]:
	module.register_buffer(
	buf_name,
	state_dict.pop(f"{name}.{buf_name}").to(device),
	persistent=False,
	)
	for attr, value in entry.get("attrs", {}).items():
	setattr(module, attr, value)
	module.ostris_quantizer = quantizer
	bias_key = f"{name}.bias"
	if bias_key in state_dict and module.bias is not None:
	module.bias.data.copy_(
	state_dict.pop(bias_key).to(device, module.bias.dtype)
	)
	if module.bias is not None:
	module.bias.requires_grad_(False)
	return len(quant_map)


	@torch.no_grad()
	def convert_linear_to_ostris(
	module: torch.nn.Linear, quantizer: OstrisQuantizer
	) -> bool:
	"""Quantize an nn.Linear in place (class swap). Returns True if the module was
	converted (or already was), False if it is not a candidate.

	A module that is ALREADY quantized (e.g. loaded from a pre-quantized
	checkpoint) is re-quantized into the requested backend when the qtypes
	differ: the weight is dequantized and re-quantized layer by layer, so the
	full-precision transient never exceeds one layer's weight. Same qtype is
	a no-op (the shipped quantization is kept)."""
	if isinstance(module, OstrisLinear):
	current_qtype = getattr(module.ostris_quantizer, "qtype", None)
	if quantizer.qtype is None or current_qtype == quantizer.qtype:
	return True
	if not quantizer.can_quantize(module):
	return True # keep the existing quantization rather than dropping it
	# fold any activation-side transform (e.g. an AWQ pre_quant_scale) into
	# the weight so the new backend computes the same function on raw inputs
	weight = module.ostris_quantizer.dequantize_folded(module).to(
	module.ostris_orig_dtype
	)
	# backend state lives exclusively in buffers; leftover scalar attrs from
	# the old backend are inert
	module._buffers.clear()
	if quantizer.wants_fp32_weight:
	quantizer.quantize_(module, weight.to(torch.float32))
	else:
	quantizer.quantize_(module, weight)
	del weight
	module.ostris_quantizer = quantizer
	return True
	weight = getattr(module, "weight", None)
	if not isinstance(weight, torch.nn.Parameter) or not weight.dtype.is_floating_point:
	return False
	if type(weight.data) is not torch.Tensor:
	# already holds a quantized tensor subclass (e.g. torchao)
	return False
	if not quantizer.can_quantize(module):
	return False
	if quantizer.wants_fp32_weight:
	quantizer.quantize_(module, weight.data.to(torch.float32))
	else:
	quantizer.quantize_(module, weight.data)
	module.ostris_quantizer = quantizer
	module.ostris_orig_dtype = weight.dtype
	del module._parameters["weight"]
	if module.bias is not None:
	module.bias.requires_grad_(False)
	module.__class__ = OstrisLinear
	return True