Compare commits
3 Commits
08fc7de87e
...
feat/torch
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
970b2a6f2f | ||
|
|
1f7f5e7c26 | ||
|
|
60c0a828cc |
@@ -15,7 +15,7 @@ from torch import nn
|
|||||||
from torch.distributed.tensor import DTensor
|
from torch.distributed.tensor import DTensor
|
||||||
|
|
||||||
from .geglu import geglu_backward, geglu_forward
|
from .geglu import geglu_backward, geglu_forward
|
||||||
from .quantize import dequantize
|
from .quantize import dequantize_weight
|
||||||
from .swiglu import swiglu_backward, swiglu_forward
|
from .swiglu import swiglu_backward, swiglu_forward
|
||||||
from .utils import torch_amp_custom_bwd, torch_amp_custom_fwd
|
from .utils import torch_amp_custom_bwd, torch_amp_custom_fwd
|
||||||
|
|
||||||
@@ -46,6 +46,12 @@ def get_lora_parameters(
|
|||||||
W = base_layer.weight
|
W = base_layer.weight
|
||||||
b = base_layer.bias
|
b = base_layer.bias
|
||||||
|
|
||||||
|
# Unwrap DTensor if FSDP2 left the weight wrapped -- DTensor does not proxy
|
||||||
|
# attribute access to the underlying tensor subclass, so torchao methods like
|
||||||
|
# .dequantize() or .get_original_weight() would not be visible.
|
||||||
|
if isinstance(W, DTensor):
|
||||||
|
W = W.full_tensor()
|
||||||
|
|
||||||
if not hasattr(proj, "disable_adapters") or proj.disable_adapters or proj.merged:
|
if not hasattr(proj, "disable_adapters") or proj.disable_adapters or proj.merged:
|
||||||
quant_state = getattr(W, "quant_state", None)
|
quant_state = getattr(W, "quant_state", None)
|
||||||
return W, b, quant_state, None, None, None
|
return W, b, quant_state, None, None, None
|
||||||
@@ -86,6 +92,7 @@ def matmul_lora(
|
|||||||
B: torch.Tensor | None,
|
B: torch.Tensor | None,
|
||||||
s: float | None,
|
s: float | None,
|
||||||
out: torch.Tensor | None = None,
|
out: torch.Tensor | None = None,
|
||||||
|
transpose: bool = True,
|
||||||
) -> torch.Tensor:
|
) -> torch.Tensor:
|
||||||
"""
|
"""
|
||||||
Efficient fused matmul + LoRA computation.
|
Efficient fused matmul + LoRA computation.
|
||||||
@@ -98,12 +105,15 @@ def matmul_lora(
|
|||||||
B: LoRA B matrix [out_features, rank]
|
B: LoRA B matrix [out_features, rank]
|
||||||
s: LoRA scaling factor
|
s: LoRA scaling factor
|
||||||
out: Optional output tensor for inplace operations
|
out: Optional output tensor for inplace operations
|
||||||
|
transpose: If True (default), transpose W before matmul (forward path).
|
||||||
|
Set to False for backward paths where W is already in the correct layout.
|
||||||
|
|
||||||
Returns:
|
Returns:
|
||||||
Result of X @ W + X @ A @ B
|
Result of X @ W + X @ A @ B
|
||||||
"""
|
"""
|
||||||
dtype = X.dtype
|
dtype = X.dtype
|
||||||
W = dequantize(W.t(), W_quant)
|
is_quantized = W_quant is not None or type(W) is not torch.Tensor
|
||||||
|
W = dequantize_weight(W, W_quant, transpose=transpose)
|
||||||
|
|
||||||
reshape = False
|
reshape = False
|
||||||
if X.dim() == 3:
|
if X.dim() == 3:
|
||||||
@@ -112,7 +122,7 @@ def matmul_lora(
|
|||||||
reshape = True
|
reshape = True
|
||||||
|
|
||||||
out = torch.matmul(X, W, out=out)
|
out = torch.matmul(X, W, out=out)
|
||||||
if W_quant is not None:
|
if is_quantized:
|
||||||
del W
|
del W
|
||||||
|
|
||||||
if A is not None:
|
if A is not None:
|
||||||
@@ -292,15 +302,16 @@ class LoRA_MLP(torch.autograd.Function):
|
|||||||
up = up.view(-1, up.shape[-1])
|
up = up.view(-1, up.shape[-1])
|
||||||
dtype = X.dtype
|
dtype = X.dtype
|
||||||
|
|
||||||
# Down projection
|
# Down projection (backward: no transpose needed, W is already [out, in])
|
||||||
grad_down = matmul_lora(
|
grad_down = matmul_lora(
|
||||||
grad_output,
|
grad_output,
|
||||||
down_weight.t(),
|
down_weight,
|
||||||
None,
|
None,
|
||||||
down_quant,
|
down_quant,
|
||||||
down_B,
|
down_B,
|
||||||
down_A,
|
down_A,
|
||||||
down_scale,
|
down_scale,
|
||||||
|
transpose=False,
|
||||||
)
|
)
|
||||||
|
|
||||||
# Activation backward
|
# Activation backward
|
||||||
@@ -332,7 +343,7 @@ class LoRA_MLP(torch.autograd.Function):
|
|||||||
|
|
||||||
if dX is not None:
|
if dX is not None:
|
||||||
# Up projection gradients
|
# Up projection gradients
|
||||||
up_weight = dequantize(up_weight.t(), up_quant)
|
up_weight = dequantize_weight(up_weight, up_quant, transpose=True)
|
||||||
if ctx.inplace:
|
if ctx.inplace:
|
||||||
dX = torch.matmul(grad_up, up_weight.t(), out=X)
|
dX = torch.matmul(grad_up, up_weight.t(), out=X)
|
||||||
else:
|
else:
|
||||||
@@ -344,7 +355,7 @@ class LoRA_MLP(torch.autograd.Function):
|
|||||||
dX += grad_up @ up_B.to(dtype).t() @ (up_scale * up_A.to(dtype).t())
|
dX += grad_up @ up_B.to(dtype).t() @ (up_scale * up_A.to(dtype).t())
|
||||||
|
|
||||||
# Gate projection gradients
|
# Gate projection gradients
|
||||||
gate_weight = dequantize(gate_weight, gate_quant)
|
gate_weight = dequantize_weight(gate_weight, gate_quant)
|
||||||
dX += grad_gate @ gate_weight
|
dX += grad_gate @ gate_weight
|
||||||
del gate_weight
|
del gate_weight
|
||||||
|
|
||||||
@@ -631,7 +642,7 @@ class LoRA_QKV(torch.autograd.Function):
|
|||||||
out_buffer = X if ctx.inplace else None
|
out_buffer = X if ctx.inplace else None
|
||||||
|
|
||||||
# Q path
|
# Q path
|
||||||
q_weight_t = dequantize(q_weight, q_quant)
|
q_weight_t = dequantize_weight(q_weight, q_quant)
|
||||||
grad_X = torch.mm(q_grad, q_weight_t, out=out_buffer)
|
grad_X = torch.mm(q_grad, q_weight_t, out=out_buffer)
|
||||||
del q_weight
|
del q_weight
|
||||||
del q_weight_t
|
del q_weight_t
|
||||||
@@ -639,7 +650,7 @@ class LoRA_QKV(torch.autograd.Function):
|
|||||||
grad_X.addmm_(q_grad, torch.mm(B_q_scaled, A_q_scaled))
|
grad_X.addmm_(q_grad, torch.mm(B_q_scaled, A_q_scaled))
|
||||||
|
|
||||||
# K path
|
# K path
|
||||||
k_weight_t = dequantize(k_weight, k_quant)
|
k_weight_t = dequantize_weight(k_weight, k_quant)
|
||||||
grad_X.addmm_(k_grad, k_weight_t)
|
grad_X.addmm_(k_grad, k_weight_t)
|
||||||
del k_weight
|
del k_weight
|
||||||
del k_weight_t
|
del k_weight_t
|
||||||
@@ -647,7 +658,7 @@ class LoRA_QKV(torch.autograd.Function):
|
|||||||
grad_X.addmm_(k_grad, torch.mm(B_k_scaled, A_k_scaled))
|
grad_X.addmm_(k_grad, torch.mm(B_k_scaled, A_k_scaled))
|
||||||
|
|
||||||
# V path
|
# V path
|
||||||
v_weight_t = dequantize(v_weight, v_quant)
|
v_weight_t = dequantize_weight(v_weight, v_quant)
|
||||||
grad_X.addmm_(v_grad, v_weight_t)
|
grad_X.addmm_(v_grad, v_weight_t)
|
||||||
del v_weight
|
del v_weight
|
||||||
del v_weight_t
|
del v_weight_t
|
||||||
@@ -810,7 +821,7 @@ class LoRA_O(torch.autograd.Function):
|
|||||||
d_B = s * A @ dY_X
|
d_B = s * A @ dY_X
|
||||||
|
|
||||||
# Get derivative for dX
|
# Get derivative for dX
|
||||||
W = dequantize(W.t(), W_quant)
|
W = dequantize_weight(W, W_quant, transpose=True)
|
||||||
dX = dY @ W.t()
|
dX = dY @ W.t()
|
||||||
del W
|
del W
|
||||||
|
|
||||||
|
|||||||
@@ -146,3 +146,43 @@ def dequantize(
|
|||||||
# Handle transposed data
|
# Handle transposed data
|
||||||
is_transposed: bool = W.shape[0] == 1
|
is_transposed: bool = W.shape[0] == 1
|
||||||
return out.t() if is_transposed else out
|
return out.t() if is_transposed else out
|
||||||
|
|
||||||
|
|
||||||
|
def dequantize_weight(
|
||||||
|
W: torch.Tensor,
|
||||||
|
quant_state: QuantState | list | None = None,
|
||||||
|
transpose: bool = False,
|
||||||
|
) -> torch.Tensor:
|
||||||
|
"""Unified dequantization for both torchao and bnb quantized weights.
|
||||||
|
|
||||||
|
For torchao tensor subclasses (AffineQuantizedTensor, NF4Tensor), dequantizes
|
||||||
|
using the appropriate instance method. For bnb Params4bit, delegates to the
|
||||||
|
optimized CUDA kernel in ``dequantize``.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
W: Quantized weight tensor ``[out_features, in_features]``.
|
||||||
|
quant_state: bnb ``QuantState`` (None for torchao / unquantized).
|
||||||
|
transpose: If True, return ``[in_features, out_features]``.
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
Dequantized float tensor, optionally transposed.
|
||||||
|
"""
|
||||||
|
# torchao path: tensor subclass with embedded quantization state
|
||||||
|
if quant_state is None and type(W) is not torch.Tensor:
|
||||||
|
result = None
|
||||||
|
# NF4Tensor (check first — NF4Tensor.dequantize is a static method)
|
||||||
|
if hasattr(W, "get_original_weight"):
|
||||||
|
result = W.get_original_weight()
|
||||||
|
else:
|
||||||
|
# AffineQuantizedTensor (INT4, etc.)
|
||||||
|
try:
|
||||||
|
result = W.dequantize()
|
||||||
|
except (TypeError, RuntimeError):
|
||||||
|
pass
|
||||||
|
if result is not None:
|
||||||
|
return result.t() if transpose else result
|
||||||
|
|
||||||
|
# bnb path: transpose input before the CUDA kernel (existing convention)
|
||||||
|
if transpose:
|
||||||
|
return dequantize(W.t(), quant_state)
|
||||||
|
return dequantize(W, quant_state)
|
||||||
|
|||||||
@@ -23,6 +23,7 @@ from axolotl.loaders.utils import get_linear_embedding_layers
|
|||||||
from axolotl.telemetry.errors import send_errors
|
from axolotl.telemetry.errors import send_errors
|
||||||
from axolotl.utils.dict import DictDefault
|
from axolotl.utils.dict import DictDefault
|
||||||
from axolotl.utils.logging import get_logger
|
from axolotl.utils.logging import get_logger
|
||||||
|
from axolotl.utils.schemas.enums import TorchAOQuantDType
|
||||||
|
|
||||||
LOG = get_logger(__name__)
|
LOG = get_logger(__name__)
|
||||||
|
|
||||||
@@ -134,11 +135,13 @@ def load_lora(
|
|||||||
|
|
||||||
rank = int(os.environ.get("LOCAL_RANK", 0))
|
rank = int(os.environ.get("LOCAL_RANK", 0))
|
||||||
|
|
||||||
|
is_torchao = cfg.peft and cfg.peft.backend == "torchao"
|
||||||
if (
|
if (
|
||||||
cfg.fsdp_config
|
cfg.fsdp_config
|
||||||
and cfg.adapter
|
and cfg.adapter
|
||||||
and cfg.fsdp_config.cpu_ram_efficient_loading
|
and cfg.fsdp_config.cpu_ram_efficient_loading
|
||||||
and rank != 0
|
and rank != 0
|
||||||
|
and not is_torchao
|
||||||
):
|
):
|
||||||
setup_quantized_meta_for_peft(model)
|
setup_quantized_meta_for_peft(model)
|
||||||
|
|
||||||
@@ -146,6 +149,15 @@ def load_lora(
|
|||||||
if cfg.peft_autocast_adapter_dtype is not None:
|
if cfg.peft_autocast_adapter_dtype is not None:
|
||||||
model_kwargs["autocast_adapter_dtype"] = cfg.peft_autocast_adapter_dtype
|
model_kwargs["autocast_adapter_dtype"] = cfg.peft_autocast_adapter_dtype
|
||||||
|
|
||||||
|
# Patch PEFT's torchao dispatch before any model creation/loading.
|
||||||
|
# Must happen before both get_peft_model and PeftModel.from_pretrained,
|
||||||
|
# as both trigger LoRA layer dispatch that would fail for INT4/NF4 weights.
|
||||||
|
# INT8 is natively supported by PEFT's TorchaoLoraLinear, so skip the patch.
|
||||||
|
if is_torchao and cfg.peft.weight_dtype != TorchAOQuantDType.int8:
|
||||||
|
from axolotl.monkeypatch.peft.utils import patch_peft_torchao_dispatch
|
||||||
|
|
||||||
|
patch_peft_torchao_dispatch()
|
||||||
|
|
||||||
if cfg.lora_model_dir:
|
if cfg.lora_model_dir:
|
||||||
LOG.debug("Loading pretrained PEFT - LoRA")
|
LOG.debug("Loading pretrained PEFT - LoRA")
|
||||||
if cfg.lora_on_cpu:
|
if cfg.lora_on_cpu:
|
||||||
@@ -172,6 +184,7 @@ def load_lora(
|
|||||||
and cfg.adapter
|
and cfg.adapter
|
||||||
and cfg.fsdp_config.cpu_ram_efficient_loading
|
and cfg.fsdp_config.cpu_ram_efficient_loading
|
||||||
and rank != 0
|
and rank != 0
|
||||||
|
and not is_torchao
|
||||||
):
|
):
|
||||||
setup_quantized_peft_meta_for_training(model)
|
setup_quantized_peft_meta_for_training(model)
|
||||||
|
|
||||||
|
|||||||
@@ -158,6 +158,15 @@ class ModelLoader:
|
|||||||
"""Property that determines if FSDP with QLoRA is enabled."""
|
"""Property that determines if FSDP with QLoRA is enabled."""
|
||||||
return self.is_fsdp_enabled and self.cfg.adapter == "qlora"
|
return self.is_fsdp_enabled and self.cfg.adapter == "qlora"
|
||||||
|
|
||||||
|
@property
|
||||||
|
def is_torchao_qlora(self):
|
||||||
|
"""Property that determines if torchao backend is used for QLoRA."""
|
||||||
|
return (
|
||||||
|
self.cfg.adapter == "qlora"
|
||||||
|
and self.cfg.peft
|
||||||
|
and self.cfg.peft.backend == "torchao"
|
||||||
|
)
|
||||||
|
|
||||||
@send_errors
|
@send_errors
|
||||||
def load(self) -> tuple[PreTrainedModel | PeftModelForCausalLM, PeftConfig | None]:
|
def load(self) -> tuple[PreTrainedModel | PeftModelForCausalLM, PeftConfig | None]:
|
||||||
"""Load and prepare the model with all configurations and patches.
|
"""Load and prepare the model with all configurations and patches.
|
||||||
@@ -491,8 +500,9 @@ class ModelLoader:
|
|||||||
|
|
||||||
# FSDP requires control over device placement, so don't set device_map when FSDP is enabled
|
# FSDP requires control over device placement, so don't set device_map when FSDP is enabled
|
||||||
if self.is_fsdp_enabled:
|
if self.is_fsdp_enabled:
|
||||||
# For QLoRA + FSDP, we still need to set device_map to "auto" for proper initialization
|
# For QLoRA + FSDP with bnb, we still need to set device_map for proper initialization
|
||||||
if self.is_qlora_and_fsdp_enabled:
|
# torchao tensors work natively with FSDP2, no device_map override needed
|
||||||
|
if self.is_qlora_and_fsdp_enabled and not self.is_torchao_qlora:
|
||||||
self.model_kwargs["device_map"] = {
|
self.model_kwargs["device_map"] = {
|
||||||
"": int(os.environ.get("LOCAL_RANK", 0))
|
"": int(os.environ.get("LOCAL_RANK", 0))
|
||||||
}
|
}
|
||||||
@@ -561,6 +571,44 @@ class ModelLoader:
|
|||||||
self.model_kwargs["quantization_config"] = BitsAndBytesConfig(
|
self.model_kwargs["quantization_config"] = BitsAndBytesConfig(
|
||||||
**self.model_config.quantization_config
|
**self.model_config.quantization_config
|
||||||
)
|
)
|
||||||
|
elif (
|
||||||
|
self.cfg.adapter == "qlora"
|
||||||
|
and self.cfg.peft
|
||||||
|
and self.cfg.peft.backend == "torchao"
|
||||||
|
and not self.cfg.merge_lora
|
||||||
|
):
|
||||||
|
from transformers import TorchAoConfig
|
||||||
|
|
||||||
|
from axolotl.utils.schemas.enums import TorchAOQuantDType
|
||||||
|
|
||||||
|
weight_dtype = self.cfg.peft.weight_dtype
|
||||||
|
if weight_dtype == TorchAOQuantDType.int4:
|
||||||
|
group_size = self.cfg.peft.group_size or 128
|
||||||
|
self.model_kwargs["quantization_config"] = TorchAoConfig(
|
||||||
|
quant_type="int4_weight_only",
|
||||||
|
group_size=group_size,
|
||||||
|
)
|
||||||
|
elif weight_dtype == TorchAOQuantDType.int8:
|
||||||
|
group_size = self.cfg.peft.group_size or 128
|
||||||
|
self.model_kwargs["quantization_config"] = TorchAoConfig(
|
||||||
|
quant_type="int8_weight_only",
|
||||||
|
group_size=group_size,
|
||||||
|
)
|
||||||
|
elif weight_dtype == TorchAOQuantDType.nf4:
|
||||||
|
from torchao.dtypes._nf4tensor_api import NF4WeightOnlyConfig
|
||||||
|
|
||||||
|
block_size = self.cfg.peft.group_size or 64
|
||||||
|
self.model_kwargs["quantization_config"] = TorchAoConfig(
|
||||||
|
quant_type=NF4WeightOnlyConfig(
|
||||||
|
block_size=block_size,
|
||||||
|
scaler_block_size=256,
|
||||||
|
),
|
||||||
|
)
|
||||||
|
else:
|
||||||
|
raise ValueError(
|
||||||
|
f"Unsupported torchao weight_dtype for QLoRA: {weight_dtype}. "
|
||||||
|
"Supported: int4, int8, nf4"
|
||||||
|
)
|
||||||
elif self.cfg.adapter == "qlora" and self.cfg.load_in_4bit:
|
elif self.cfg.adapter == "qlora" and self.cfg.load_in_4bit:
|
||||||
bnb_config = {
|
bnb_config = {
|
||||||
"load_in_4bit": True,
|
"load_in_4bit": True,
|
||||||
@@ -860,6 +908,10 @@ class ModelLoader:
|
|||||||
# Make sure everything is in the same dtype
|
# Make sure everything is in the same dtype
|
||||||
skip_prepare_model_for_kbit_training = True
|
skip_prepare_model_for_kbit_training = True
|
||||||
|
|
||||||
|
# torchao quantized models don't use Params4bit and don't need kbit preparation
|
||||||
|
if self.is_torchao_qlora:
|
||||||
|
skip_prepare_model_for_kbit_training = True
|
||||||
|
|
||||||
if (
|
if (
|
||||||
not skip_prepare_model_for_kbit_training
|
not skip_prepare_model_for_kbit_training
|
||||||
and self.cfg.adapter in ["lora", "qlora"]
|
and self.cfg.adapter in ["lora", "qlora"]
|
||||||
|
|||||||
@@ -348,10 +348,12 @@ class PatchManager:
|
|||||||
|
|
||||||
def _apply_fsdp2_bnb_patches(self):
|
def _apply_fsdp2_bnb_patches(self):
|
||||||
"""Apply FSDP2 BNB patches."""
|
"""Apply FSDP2 BNB patches."""
|
||||||
|
is_torchao = self.cfg.peft and self.cfg.peft.backend == "torchao"
|
||||||
if (
|
if (
|
||||||
self.cfg.fsdp_config
|
self.cfg.fsdp_config
|
||||||
and str(self.cfg.fsdp_version) == "2"
|
and str(self.cfg.fsdp_version) == "2"
|
||||||
and self.cfg.adapter == "qlora"
|
and self.cfg.adapter == "qlora"
|
||||||
|
and not is_torchao
|
||||||
):
|
):
|
||||||
from axolotl.monkeypatch.fsdp2_qlora import (
|
from axolotl.monkeypatch.fsdp2_qlora import (
|
||||||
apply_init_sharded_param_patch,
|
apply_init_sharded_param_patch,
|
||||||
|
|||||||
@@ -78,3 +78,30 @@ def patch_peft_prep_code():
|
|||||||
axolotl.loaders.model.prepare_model_for_kbit_training = (
|
axolotl.loaders.model.prepare_model_for_kbit_training = (
|
||||||
fixed_prepare_model_for_kbit_training
|
fixed_prepare_model_for_kbit_training
|
||||||
)
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def patch_peft_torchao_dispatch():
|
||||||
|
"""Skip PEFT's TorchaoLoraLinear for non-INT8 torchao weights.
|
||||||
|
|
||||||
|
PEFT's dispatch_torchao() matches AffineQuantizedTensor but then errors in
|
||||||
|
_check_dtype_supported() because it only allows INT8. Our LoRA kernels handle
|
||||||
|
dequantization explicitly, so we bypass PEFT's torchao dispatch entirely and
|
||||||
|
let it fall back to standard Linear LoRA layers.
|
||||||
|
"""
|
||||||
|
try:
|
||||||
|
from peft.tuners.lora import torchao as peft_torchao
|
||||||
|
except ImportError:
|
||||||
|
LOG.warning("Could not import peft.tuners.lora.torchao for patching")
|
||||||
|
return
|
||||||
|
|
||||||
|
if getattr(peft_torchao, "_axolotl_patched", False):
|
||||||
|
return
|
||||||
|
|
||||||
|
def patched_dispatch(target, adapter_name, lora_config, **kwargs):
|
||||||
|
# Return None so PEFT falls back to standard Linear LoRA layers.
|
||||||
|
# Our LoRA kernels handle torchao dequantization explicitly.
|
||||||
|
return None
|
||||||
|
|
||||||
|
peft_torchao.dispatch_torchao = patched_dispatch
|
||||||
|
peft_torchao._axolotl_patched = True
|
||||||
|
LOG.info("Patched PEFT dispatch_torchao to skip TorchaoLoraLinear")
|
||||||
|
|||||||
@@ -8,6 +8,7 @@ import torch
|
|||||||
class TorchAOQuantDType(Enum):
|
class TorchAOQuantDType(Enum):
|
||||||
int4 = torch.int4
|
int4 = torch.int4
|
||||||
int8 = torch.int8
|
int8 = torch.int8
|
||||||
|
nf4 = "nf4"
|
||||||
float8_e4m3fn = torch.float8_e4m3fn
|
float8_e4m3fn = torch.float8_e4m3fn
|
||||||
nvfp4 = "nvfp4"
|
nvfp4 = "nvfp4"
|
||||||
|
|
||||||
@@ -16,6 +17,8 @@ class TorchAOQuantDType(Enum):
|
|||||||
return TorchAOQuantDType.int4
|
return TorchAOQuantDType.int4
|
||||||
if str == "int8":
|
if str == "int8":
|
||||||
return TorchAOQuantDType.int8
|
return TorchAOQuantDType.int8
|
||||||
|
if str == "nf4":
|
||||||
|
return TorchAOQuantDType.nf4
|
||||||
if str in ["float8_e4m3fn", "fp8", "float8"]:
|
if str in ["float8_e4m3fn", "fp8", "float8"]:
|
||||||
return TorchAOQuantDType.float8_e4m3fn
|
return TorchAOQuantDType.float8_e4m3fn
|
||||||
if str == "nvfp4":
|
if str == "nvfp4":
|
||||||
|
|||||||
@@ -1,9 +1,12 @@
|
|||||||
"""Pydantic models for PEFT-related configuration"""
|
"""Pydantic models for PEFT-related configuration"""
|
||||||
|
|
||||||
from typing import Any
|
from typing import Any, Literal
|
||||||
|
|
||||||
from pydantic import BaseModel, Field, field_validator, model_validator
|
from pydantic import BaseModel, Field, field_validator, model_validator
|
||||||
|
|
||||||
|
from axolotl.utils.schemas.enums import TorchAOQuantDType
|
||||||
|
from axolotl.utils.schemas.quantization import validate_ao_dtype
|
||||||
|
|
||||||
|
|
||||||
class LoftQConfig(BaseModel):
|
class LoftQConfig(BaseModel):
|
||||||
"""LoftQ configuration subset"""
|
"""LoftQ configuration subset"""
|
||||||
@@ -15,7 +18,7 @@ class LoftQConfig(BaseModel):
|
|||||||
|
|
||||||
|
|
||||||
class PeftConfig(BaseModel):
|
class PeftConfig(BaseModel):
|
||||||
"""peftq configuration subset"""
|
"""PEFT configuration subset"""
|
||||||
|
|
||||||
loftq_config: LoftQConfig | None = Field(
|
loftq_config: LoftQConfig | None = Field(
|
||||||
default=None,
|
default=None,
|
||||||
@@ -23,6 +26,29 @@ class PeftConfig(BaseModel):
|
|||||||
"description": "Configuration options for loftq initialization for LoRA"
|
"description": "Configuration options for loftq initialization for LoRA"
|
||||||
},
|
},
|
||||||
)
|
)
|
||||||
|
backend: Literal["bnb", "torchao"] | None = Field(
|
||||||
|
default=None,
|
||||||
|
json_schema_extra={
|
||||||
|
"description": "Quantization backend for QLoRA. 'bnb' for bitsandbytes (default), 'torchao' for torchao."
|
||||||
|
},
|
||||||
|
)
|
||||||
|
weight_dtype: TorchAOQuantDType | None = Field(
|
||||||
|
default=None,
|
||||||
|
json_schema_extra={
|
||||||
|
"description": "Weight quantization dtype (int4, int8, or nf4). Also used with bnb backend to auto-configure quantization."
|
||||||
|
},
|
||||||
|
)
|
||||||
|
group_size: int | None = Field(
|
||||||
|
default=None,
|
||||||
|
json_schema_extra={
|
||||||
|
"description": "Group size for quantization. Defaults to 128 for int4, 64 for nf4."
|
||||||
|
},
|
||||||
|
)
|
||||||
|
|
||||||
|
@field_validator("weight_dtype", mode="before")
|
||||||
|
@classmethod
|
||||||
|
def validate_weight_dtype(cls, v):
|
||||||
|
return validate_ao_dtype(v)
|
||||||
|
|
||||||
|
|
||||||
class LoraConfig(BaseModel):
|
class LoraConfig(BaseModel):
|
||||||
@@ -156,6 +182,56 @@ class LoraConfig(BaseModel):
|
|||||||
|
|
||||||
merge_lora: bool | None = None
|
merge_lora: bool | None = None
|
||||||
|
|
||||||
|
@model_validator(mode="before")
|
||||||
|
@classmethod
|
||||||
|
def auto_detect_qlora(cls, data):
|
||||||
|
"""Auto-set adapter type and quantization flags from peft config.
|
||||||
|
|
||||||
|
When peft.backend and peft.weight_dtype are set, this infers the correct
|
||||||
|
adapter type and internal flags (load_in_4bit, load_in_8bit) so users
|
||||||
|
don't need to set them manually.
|
||||||
|
"""
|
||||||
|
peft = data.get("peft")
|
||||||
|
if not isinstance(peft, dict):
|
||||||
|
return data
|
||||||
|
|
||||||
|
backend = peft.get("backend")
|
||||||
|
weight_dtype = peft.get("weight_dtype")
|
||||||
|
|
||||||
|
# Validate: weight_dtype requires backend
|
||||||
|
if weight_dtype and not backend:
|
||||||
|
raise ValueError(
|
||||||
|
"peft.backend is required when peft.weight_dtype is set. "
|
||||||
|
"Use 'torchao' or 'bnb'."
|
||||||
|
)
|
||||||
|
|
||||||
|
if not weight_dtype:
|
||||||
|
return data
|
||||||
|
|
||||||
|
adapter = data.get("adapter")
|
||||||
|
|
||||||
|
if backend == "torchao":
|
||||||
|
# torchao: any quantized weight_dtype means qlora
|
||||||
|
if adapter == "lora":
|
||||||
|
data["adapter"] = "qlora"
|
||||||
|
|
||||||
|
elif backend == "bnb":
|
||||||
|
if weight_dtype == "nf4":
|
||||||
|
# bnb nf4 = qlora with load_in_4bit
|
||||||
|
if adapter == "lora":
|
||||||
|
data["adapter"] = "qlora"
|
||||||
|
data.setdefault("load_in_4bit", True)
|
||||||
|
elif weight_dtype == "int8":
|
||||||
|
# bnb int8 = lora with load_in_8bit
|
||||||
|
data.setdefault("load_in_8bit", True)
|
||||||
|
else:
|
||||||
|
raise ValueError(
|
||||||
|
f"peft.weight_dtype '{weight_dtype}' is not supported with bnb backend. "
|
||||||
|
"Supported: nf4, int8."
|
||||||
|
)
|
||||||
|
|
||||||
|
return data
|
||||||
|
|
||||||
@model_validator(mode="before")
|
@model_validator(mode="before")
|
||||||
@classmethod
|
@classmethod
|
||||||
def validate_adapter(cls, data):
|
def validate_adapter(cls, data):
|
||||||
@@ -173,6 +249,8 @@ class LoraConfig(BaseModel):
|
|||||||
@model_validator(mode="after")
|
@model_validator(mode="after")
|
||||||
def validate_qlora(self):
|
def validate_qlora(self):
|
||||||
if self.adapter == "qlora":
|
if self.adapter == "qlora":
|
||||||
|
is_torchao = self.peft and self.peft.backend == "torchao"
|
||||||
|
|
||||||
if self.merge_lora:
|
if self.merge_lora:
|
||||||
# can't merge qlora if loaded in 8bit or 4bit
|
# can't merge qlora if loaded in 8bit or 4bit
|
||||||
if self.load_in_8bit:
|
if self.load_in_8bit:
|
||||||
@@ -184,7 +262,20 @@ class LoraConfig(BaseModel):
|
|||||||
if self.load_in_4bit:
|
if self.load_in_4bit:
|
||||||
raise ValueError("Can't merge qlora if loaded in 4bit")
|
raise ValueError("Can't merge qlora if loaded in 4bit")
|
||||||
|
|
||||||
|
elif is_torchao:
|
||||||
|
# torchao backend: validate torchao-specific requirements
|
||||||
|
if self.load_in_4bit or self.load_in_8bit:
|
||||||
|
raise ValueError(
|
||||||
|
"load_in_4bit/load_in_8bit are for bitsandbytes. "
|
||||||
|
"With peft.backend: torchao, quantization is handled by torchao."
|
||||||
|
)
|
||||||
|
if not self.peft.weight_dtype:
|
||||||
|
raise ValueError(
|
||||||
|
"peft.weight_dtype is required when peft.backend is 'torchao'"
|
||||||
|
)
|
||||||
|
|
||||||
else:
|
else:
|
||||||
|
# Default bnb path
|
||||||
if self.load_in_8bit:
|
if self.load_in_8bit:
|
||||||
raise ValueError("Can't load qlora in 8bit")
|
raise ValueError("Can't load qlora in 8bit")
|
||||||
|
|
||||||
|
|||||||
@@ -16,6 +16,8 @@ def validate_ao_dtype(v: Any) -> TorchAOQuantDType | None:
|
|||||||
return TorchAOQuantDType.int4
|
return TorchAOQuantDType.int4
|
||||||
if v == "int8":
|
if v == "int8":
|
||||||
return TorchAOQuantDType.int8
|
return TorchAOQuantDType.int8
|
||||||
|
if v == "nf4":
|
||||||
|
return TorchAOQuantDType.nf4
|
||||||
if v in ["float8_e4m3fn", "fp8", "float8"]:
|
if v in ["float8_e4m3fn", "fp8", "float8"]:
|
||||||
return TorchAOQuantDType.float8_e4m3fn
|
return TorchAOQuantDType.float8_e4m3fn
|
||||||
if v == "nvfp4":
|
if v == "nvfp4":
|
||||||
|
|||||||
@@ -3,7 +3,7 @@
|
|||||||
import torch
|
import torch
|
||||||
from bitsandbytes.functional import QuantState
|
from bitsandbytes.functional import QuantState
|
||||||
|
|
||||||
from axolotl.kernels.quantize import dequantize
|
from axolotl.kernels.quantize import dequantize, dequantize_weight
|
||||||
|
|
||||||
|
|
||||||
def test_dequantize_null_state():
|
def test_dequantize_null_state():
|
||||||
@@ -100,3 +100,18 @@ def test_dequantize_output_tensor():
|
|||||||
|
|
||||||
result = dequantize(W, quant_state, out=out)
|
result = dequantize(W, quant_state, out=out)
|
||||||
assert result is out
|
assert result is out
|
||||||
|
|
||||||
|
|
||||||
|
def test_dequantize_weight_plain_tensor():
|
||||||
|
"""Test that dequantize_weight passes through unquantized tensors unchanged"""
|
||||||
|
W = torch.randn(32, 64)
|
||||||
|
result = dequantize_weight(W, quant_state=None, transpose=False)
|
||||||
|
assert torch.equal(result, W)
|
||||||
|
|
||||||
|
|
||||||
|
def test_dequantize_weight_plain_tensor_transpose():
|
||||||
|
"""Test that dequantize_weight transposes unquantized tensors"""
|
||||||
|
W = torch.randn(32, 64)
|
||||||
|
result = dequantize_weight(W, quant_state=None, transpose=True)
|
||||||
|
assert result.shape == (64, 32)
|
||||||
|
assert torch.equal(result, W.t())
|
||||||
|
|||||||
@@ -3,6 +3,14 @@ import pytest
|
|||||||
from axolotl.utils.config import validate_config
|
from axolotl.utils.config import validate_config
|
||||||
from axolotl.utils.dict import DictDefault
|
from axolotl.utils.dict import DictDefault
|
||||||
|
|
||||||
|
BASE_CFG = {
|
||||||
|
"datasets": [{"path": "dummy_dataset", "type": "alpaca"}],
|
||||||
|
"micro_batch_size": 1,
|
||||||
|
"gradient_accumulation_steps": 1,
|
||||||
|
"learning_rate": 1e-5,
|
||||||
|
"base_model": "dummy_model",
|
||||||
|
}
|
||||||
|
|
||||||
|
|
||||||
class TestLoRAConfigValidation:
|
class TestLoRAConfigValidation:
|
||||||
"""Test suite for LoRA/QLoRA configuration validation"""
|
"""Test suite for LoRA/QLoRA configuration validation"""
|
||||||
@@ -149,3 +157,195 @@ class TestLoRAConfigValidation:
|
|||||||
result = validate_config(valid_config)
|
result = validate_config(valid_config)
|
||||||
assert result["lora_qkv_kernel"] is True
|
assert result["lora_qkv_kernel"] is True
|
||||||
assert result["trust_remote_code"] is None
|
assert result["trust_remote_code"] is None
|
||||||
|
|
||||||
|
|
||||||
|
class TestTorchaoQLoRAConfigValidation:
|
||||||
|
"""Test suite for torchao QLoRA auto-detection and validation"""
|
||||||
|
|
||||||
|
# --- Auto-detection: torchao ---
|
||||||
|
|
||||||
|
@pytest.mark.parametrize("weight_dtype", ["int4", "int8", "nf4"])
|
||||||
|
def test_torchao_auto_detect_from_lora(self, weight_dtype):
|
||||||
|
"""adapter: lora + peft.backend: torchao auto-upgrades to qlora"""
|
||||||
|
cfg = DictDefault(
|
||||||
|
{
|
||||||
|
"adapter": "lora",
|
||||||
|
"peft": {"backend": "torchao", "weight_dtype": weight_dtype},
|
||||||
|
**BASE_CFG,
|
||||||
|
}
|
||||||
|
)
|
||||||
|
result = validate_config(cfg)
|
||||||
|
assert result["adapter"] == "qlora"
|
||||||
|
assert result["peft"]["backend"] == "torchao"
|
||||||
|
|
||||||
|
def test_torchao_explicit_qlora(self):
|
||||||
|
"""adapter: qlora + peft.backend: torchao works directly"""
|
||||||
|
cfg = DictDefault(
|
||||||
|
{
|
||||||
|
"adapter": "qlora",
|
||||||
|
"peft": {"backend": "torchao", "weight_dtype": "int4"},
|
||||||
|
**BASE_CFG,
|
||||||
|
}
|
||||||
|
)
|
||||||
|
result = validate_config(cfg)
|
||||||
|
assert result["adapter"] == "qlora"
|
||||||
|
|
||||||
|
# --- Auto-detection: bnb ---
|
||||||
|
|
||||||
|
def test_bnb_nf4_auto_detect_from_lora(self):
|
||||||
|
"""adapter: lora + peft.backend: bnb + weight_dtype: nf4 → qlora + load_in_4bit"""
|
||||||
|
cfg = DictDefault(
|
||||||
|
{
|
||||||
|
"adapter": "lora",
|
||||||
|
"peft": {"backend": "bnb", "weight_dtype": "nf4"},
|
||||||
|
**BASE_CFG,
|
||||||
|
}
|
||||||
|
)
|
||||||
|
result = validate_config(cfg)
|
||||||
|
assert result["adapter"] == "qlora"
|
||||||
|
assert result["load_in_4bit"] is True
|
||||||
|
|
||||||
|
def test_bnb_int8_auto_detect_from_lora(self):
|
||||||
|
"""adapter: lora + peft.backend: bnb + weight_dtype: int8 → lora + load_in_8bit"""
|
||||||
|
cfg = DictDefault(
|
||||||
|
{
|
||||||
|
"adapter": "lora",
|
||||||
|
"peft": {"backend": "bnb", "weight_dtype": "int8"},
|
||||||
|
**BASE_CFG,
|
||||||
|
}
|
||||||
|
)
|
||||||
|
result = validate_config(cfg)
|
||||||
|
assert result["adapter"] == "lora"
|
||||||
|
assert result["load_in_8bit"] is True
|
||||||
|
|
||||||
|
def test_bnb_nf4_explicit_qlora_auto_sets_load_in_4bit(self):
|
||||||
|
"""adapter: qlora + peft.backend: bnb + weight_dtype: nf4 auto-sets load_in_4bit"""
|
||||||
|
cfg = DictDefault(
|
||||||
|
{
|
||||||
|
"adapter": "qlora",
|
||||||
|
"peft": {"backend": "bnb", "weight_dtype": "nf4"},
|
||||||
|
**BASE_CFG,
|
||||||
|
}
|
||||||
|
)
|
||||||
|
result = validate_config(cfg)
|
||||||
|
assert result["adapter"] == "qlora"
|
||||||
|
assert result["load_in_4bit"] is True
|
||||||
|
|
||||||
|
# --- Backward compat ---
|
||||||
|
|
||||||
|
def test_old_style_qlora_unchanged(self):
|
||||||
|
"""Old-style adapter: qlora + load_in_4bit: true still works"""
|
||||||
|
cfg = DictDefault(
|
||||||
|
{
|
||||||
|
"adapter": "qlora",
|
||||||
|
"load_in_4bit": True,
|
||||||
|
**BASE_CFG,
|
||||||
|
}
|
||||||
|
)
|
||||||
|
result = validate_config(cfg)
|
||||||
|
assert result["adapter"] == "qlora"
|
||||||
|
assert result["load_in_4bit"] is True
|
||||||
|
|
||||||
|
def test_old_style_lora_8bit_unchanged(self):
|
||||||
|
"""Old-style adapter: lora + load_in_8bit: true still works"""
|
||||||
|
cfg = DictDefault(
|
||||||
|
{
|
||||||
|
"adapter": "lora",
|
||||||
|
"load_in_8bit": True,
|
||||||
|
**BASE_CFG,
|
||||||
|
}
|
||||||
|
)
|
||||||
|
result = validate_config(cfg)
|
||||||
|
assert result["adapter"] == "lora"
|
||||||
|
assert result["load_in_8bit"] is True
|
||||||
|
|
||||||
|
def test_plain_lora_unchanged(self):
|
||||||
|
"""adapter: lora without peft block stays as lora"""
|
||||||
|
cfg = DictDefault(
|
||||||
|
{
|
||||||
|
"adapter": "lora",
|
||||||
|
**BASE_CFG,
|
||||||
|
}
|
||||||
|
)
|
||||||
|
result = validate_config(cfg)
|
||||||
|
assert result["adapter"] == "lora"
|
||||||
|
|
||||||
|
# --- Validation errors ---
|
||||||
|
|
||||||
|
def test_torchao_with_load_in_4bit_errors(self):
|
||||||
|
"""peft.backend: torchao + load_in_4bit is a conflict"""
|
||||||
|
cfg = DictDefault(
|
||||||
|
{
|
||||||
|
"adapter": "qlora",
|
||||||
|
"load_in_4bit": True,
|
||||||
|
"peft": {"backend": "torchao", "weight_dtype": "int4"},
|
||||||
|
**BASE_CFG,
|
||||||
|
}
|
||||||
|
)
|
||||||
|
with pytest.raises(ValueError, match="load_in_4bit.*bitsandbytes"):
|
||||||
|
validate_config(cfg)
|
||||||
|
|
||||||
|
def test_torchao_with_load_in_8bit_errors(self):
|
||||||
|
"""peft.backend: torchao + load_in_8bit is a conflict"""
|
||||||
|
cfg = DictDefault(
|
||||||
|
{
|
||||||
|
"adapter": "qlora",
|
||||||
|
"load_in_8bit": True,
|
||||||
|
"peft": {"backend": "torchao", "weight_dtype": "int4"},
|
||||||
|
**BASE_CFG,
|
||||||
|
}
|
||||||
|
)
|
||||||
|
with pytest.raises(ValueError, match="load_in_4bit.*bitsandbytes"):
|
||||||
|
validate_config(cfg)
|
||||||
|
|
||||||
|
def test_torchao_without_weight_dtype_errors(self):
|
||||||
|
"""peft.backend: torchao without weight_dtype errors"""
|
||||||
|
cfg = DictDefault(
|
||||||
|
{
|
||||||
|
"adapter": "qlora",
|
||||||
|
"peft": {"backend": "torchao"},
|
||||||
|
**BASE_CFG,
|
||||||
|
}
|
||||||
|
)
|
||||||
|
with pytest.raises(ValueError, match="peft.weight_dtype is required"):
|
||||||
|
validate_config(cfg)
|
||||||
|
|
||||||
|
def test_weight_dtype_without_backend_errors(self):
|
||||||
|
"""peft.weight_dtype without peft.backend errors"""
|
||||||
|
cfg = DictDefault(
|
||||||
|
{
|
||||||
|
"adapter": "lora",
|
||||||
|
"peft": {"weight_dtype": "int4"},
|
||||||
|
**BASE_CFG,
|
||||||
|
}
|
||||||
|
)
|
||||||
|
with pytest.raises(ValueError, match="peft.backend is required"):
|
||||||
|
validate_config(cfg)
|
||||||
|
|
||||||
|
def test_bnb_unsupported_weight_dtype_errors(self):
|
||||||
|
"""peft.backend: bnb + unsupported weight_dtype errors"""
|
||||||
|
cfg = DictDefault(
|
||||||
|
{
|
||||||
|
"adapter": "lora",
|
||||||
|
"peft": {"backend": "bnb", "weight_dtype": "int4"},
|
||||||
|
**BASE_CFG,
|
||||||
|
}
|
||||||
|
)
|
||||||
|
with pytest.raises(ValueError, match="not supported with bnb"):
|
||||||
|
validate_config(cfg)
|
||||||
|
|
||||||
|
# --- Redundant flags don't conflict ---
|
||||||
|
|
||||||
|
def test_bnb_nf4_with_explicit_load_in_4bit(self):
|
||||||
|
"""peft.backend: bnb + weight_dtype: nf4 + load_in_4bit: true is fine (redundant)"""
|
||||||
|
cfg = DictDefault(
|
||||||
|
{
|
||||||
|
"adapter": "lora",
|
||||||
|
"load_in_4bit": True,
|
||||||
|
"peft": {"backend": "bnb", "weight_dtype": "nf4"},
|
||||||
|
**BASE_CFG,
|
||||||
|
}
|
||||||
|
)
|
||||||
|
result = validate_config(cfg)
|
||||||
|
assert result["adapter"] == "qlora"
|
||||||
|
assert result["load_in_4bit"] is True
|
||||||
|
|||||||
Reference in New Issue
Block a user