Update utils.py

This commit is contained in:
Daniel Han 2025-10-17 20:51:54 -07:00
commit 6046fbcd54

View file

@ -112,7 +112,7 @@ if DEVICE_COUNT > 1:
else:
from contextlib import nullcontext
def torch_gpu_device(device): return nullcontext()
pass
pass
# INTEL GPU Specific Logic
if DEVICE_TYPE == "xpu":
@ -212,15 +212,21 @@ def get_lora_parameters(proj):
W = base_layer.weight
# Optionally apply fake quantization to base layer weights for QAT
weight_fake_quantizer = getattr(base_layer, "weight_fake_quantizer", None)
if weight_fake_quantizer is not None:
W = weight_fake_quantizer(W)
if hasattr(base_layer, "weight_fake_quantizer"):
weight_fake_quantizer = getattr(base_layer, "weight_fake_quantizer", None)
if weight_fake_quantizer is not None:
W = weight_fake_quantizer(W)
W_quant = next((x for x in [getattr(W, "quant_state", None), getattr(base_layer, "weight_scale_inv", None), getattr(base_layer, "weight_scale", None)] if x is not None), None)
# Get quant state for 4bit or FP8
W_quant = getattr(W, "quant_state", None)
if W_quant is None:
W_quant = getattr(base_layer, "weight_scale_inv", None)
if W_quant is None:
W_quant = getattr(base_layer, "weight_scale", None)
if getattr(base_layer, 'quant_method', None) == 'fp8':
if getattr(base_layer, "quant_method", None) == "fp8":
# we need to somehow store and pass this information :)
W.block_size = getattr(base_layer, 'block_size', [128, 128])
W.block_size = getattr(base_layer, "block_size", [128, 128])
W_quant.block_size = W.block_size
# if not hasattr(proj, "disable_adapters") or proj.disable_adapters or proj.merged:
@ -235,14 +241,16 @@ def get_lora_parameters(proj):
# Optionally apply fake quantization to lora weights for QAT
lora_A_linear = proj.lora_A[adapter]
lora_B_linear = proj.lora_B[adapter]
lora_A_fake_quantizer = getattr(lora_A_linear, "weight_fake_quantizer", None)
lora_B_fake_quantizer = getattr(lora_B_linear, "weight_fake_quantizer", None)
A = lora_A_linear.weight
B = lora_B_linear.weight
if lora_A_fake_quantizer is not None:
A = lora_A_fake_quantizer(A)
if lora_B_fake_quantizer is not None:
B = lora_B_fake_quantizer(B)
if hasattr(lora_A_linear, "weight_fake_quantizer"):
lora_A_fake_quantizer = getattr(lora_A_linear, "weight_fake_quantizer", None)
if lora_A_fake_quantizer is not None:
A = lora_A_fake_quantizer(A)
if hasattr(lora_B_linear, "weight_fake_quantizer"):
lora_B_fake_quantizer = getattr(lora_B_linear, "weight_fake_quantizer", None)
if lora_B_fake_quantizer is not None:
B = lora_B_fake_quantizer(B)
return (
W,
@ -259,16 +267,21 @@ def get_lora_parameters_bias(proj):
base_layer = getattr(proj, "base_layer", proj) # (proj.base_layer if hasattr(proj, "base_layer") else proj)
W = base_layer.weight
W_quant = next((x for x in [getattr(W, "quant_state", None), getattr(base_layer, "weight_scale_inv", None), getattr(base_layer, "weight_scale", None)] if x is not None), None)
# Get quant state for 4bit or FP8
W_quant = getattr(W, "quant_state", None)
if W_quant is None:
W_quant = getattr(base_layer, "weight_scale_inv", None)
if W_quant is None:
W_quant = getattr(base_layer, "weight_scale", None)
# if not hasattr(proj, "disable_adapters") or proj.disable_adapters or proj.merged:
if getattr(proj, "disable_adapters", True) or proj.merged:
return W, W_quant, None, None, None, base_layer.bias
pass
if getattr(base_layer, 'quant_method', None) == 'fp8':
if getattr(base_layer, "quant_method", None) == "fp8":
# we need to somehow store and pass this information :)
W.block_size = getattr(base_layer, 'block_size', [128, 128])
W.block_size = getattr(base_layer, "block_size", [128, 128])
W_quant.block_size = W.block_size
adapter = getattr(proj, "active_adapters", None)