TorchAO: Unsloth 및 Axolotl 통합과 PARQ를 통한 양자화 인식 학습(QAT) 가속화

from torchao.quantization import quantize_, Int4WeightOnlyConfig
from torchao.quantization.qat import QATConfig

# PTQ와 동일한 설정 사용
base_config = Int4WeightOnlyConfig(group_size=32)

# 준비 단계: 모델을 "가짜 양자화" 상태로 만들고 학습 준비
quantize_(model, QATConfig(base_config, step="prepare"))
train(model)

# 변환 단계: 모델을 실제 양자화하여 추론 준비
quantize_(model, QATConfig(base_config, step="convert"))

from unsloth import FastLanguageModel

model, tokenizer = FastLanguageModel.from_pretrained(
    model_name = "unsloth/gemma3-12b-it",
    max_seq_length = 2048,
    load_in_16bit = True,
)

model = FastLanguageModel.get_peft_model(
    model,
    r = 16,
    target_modules = ["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"],
    lora_alpha = 32,
    # fp8-int4, fp8-fp8, int4, int8-int4 등 지원
    qat_scheme = "int4",
)

from torchao.prototype.parq.api import QuantConfig, create_optimizer

def linear_filter_fn(module, fqn):
    return isinstance(module, torch.nn.Linear) and fqn.endswith("weight")

linear_config = QuantConfig(bitwidth=2, group_size=None)
quant_configs_and_filter_fns = [(linear_config, linear_filter_fn)]

# optimizer.step()에서 양자화를 수행하는 QuantOptimizer 적용
optimizer = create_optimizer(
    model,
    quant_configs_and_filter_fns,
    base_optimizer_cls=torch.optim.AdamW,
    base_optimizer_kwargs={"weight_decay": 1e-2},
    quant_per_channel=True,
)