bitsandbytes 量化库

import bitsandbytes as bnb

# 创建 8-bit 线性层
linear_8bit = bnb.nn.Linear8bitLt(
    input_features=4096,      # 输入维度
    output_features=4096,     # 输出维度
    bias=True,                # 是否包含偏置
    has_fp16_weights=False,   # 是否保留fp16权重（False表示纯8-bit存储）
    threshold=6.0,            # 离群值检测阈值
    index=None,               # 内部索引（通常不需设置）
)

import torch
import bitsandbytes as bnb

# 先创建 fp16 线性层，再替换为 8-bit
fp16_linear = torch.nn.Linear(4096, 4096)
# 替换关键参数
int8_linear = bnb.nn.Linear8bitLt(
    fp16_linear.in_features,
    fp16_linear.out_features,
    fp16_linear.bias is not None,
    has_fp16_weights=False,
    threshold=6.0,
)
# 将 fp16 权重量化为 int8
int8_linear.weight = bnb.nn.Int8Params(
    fp16_linear.weight.data.cpu(),
    has_fp16_weights=False,
    requires_grad=False,
).to(fp16_linear.weight.dtype)

# 前向传播（自动处理混合精度分解）
x = torch.randn(2, 4096, dtype=torch.float16, device='cuda')
out = int8_linear(x.cuda())  # 输出为 fp16

import bitsandbytes as bnb

# 创建 4-bit 线性层
linear_4bit = bnb.nn.Linear4bit(
    input_features=4096,          # 输入维度
    output_features=4096,         # 输出维度
    bias=True,                    # 是否包含偏置
    compute_dtype=torch.bfloat16, # 计算时反量化的数据类型
    compress_statistics=True,     # 是否使用双重量化（Double Quantization）
    quant_type='nf4',             # 量化类型：'nf4' 或 'fp4'
    quant_storage=torch.uint8,    # 量化存储类型
)

import torch
import bitsandbytes as bnb

# 替换模型中的 Linear 层为 4-bit
def replace_with_4bit(model):
    for name, module in model.named_modules():
        if isinstance(module, torch.nn.Linear):
            # 创建对应的 4-bit 层
            new_module = bnb.nn.Linear4bit(
                module.in_features,
                module.out_features,
                module.bias is not None,
                compute_dtype=torch.bfloat16,
                compress_statistics=True,
                quant_type='nf4',
            )
            # 设置名称以便后续访问
            setattr(model, name.split('.')[-1], new_module)
    return model

from bitsandbytes import prepare_model_for_kbit_training

model = prepare_model_for_kbit_training(
    model,                          # 已量化的模型
    use_gradient_checkpointing=True, # 是否启用梯度检查点
    gradient_checkpointing_kwargs=None,  # 梯度检查点额外参数
)

from transformers import AutoModelForCausalLM
from bitsandbytes import prepare_model_for_kbit_training

# 加载 4-bit 量化模型
model = AutoModelForCausalLM.from_pretrained(
    "meta-llama/Llama-2-7b-hf",
    quantization_config=BitsAndBytesConfig(load_in_4bit=True),
    device_map="auto",
)

# 准备训练
model = prepare_model_for_kbit_training(model)

# 添加 LoRA 适配器
from peft import LoraConfig, get_peft_model
lora_config = LoraConfig(
    r=16,
    lora_alpha=32,
    target_modules=["q_proj", "v_proj"],
    lora_dropout=0.05,
    bias="none",
    task_type="CAUSAL_LM",
)
model = get_peft_model(model, lora_config)
model.print_trainable_parameters()
# 输出: trainable params: 13,107,200 || all params: 6,738,415,616 || trainable%: 0.1944

from transformers import BitsAndBytesConfig

# 4-bit NF4 量化配置（QLoRA 推荐配置）
bnb_config_4bit = BitsAndBytesConfig(
    load_in_4bit=True,                          # 启用 4-bit 量化
    bnb_4bit_quant_type="nf4",                  # 量化类型："nf4" 或 "fp4"
    bnb_4bit_compute_dtype=torch.bfloat16,      # 计算精度
    bnb_4bit_use_double_quantization=True,       # 启用双重量化
)

# 8-bit 量化配置
bnb_config_8bit = BitsAndBytesConfig(
    load_in_8bit=True,              # 启用 8-bit 量化
    llm_int8_threshold=6.0,         # 离群值阈值
    llm_int8_skip_modules=None,     # 跳过量化的模块列表
    llm_int8_enable_fp32_cpu_offload=False,  # CPU 卸载
)

import bitsandbytes as bnb

# 8-bit AdamW 优化器
optimizer = bnb.optim.AdamW8bit(
    model.parameters(),
    lr=2e-5,
    betas=(0.9, 0.999),
    eps=1e-8,
    weight_decay=0.01,
    optim_bits=8,      # 优化器状态的位数：8 或 32
    is_paged=True,     # 是否启用分页优化器
)

# 8-bit Adam 优化器
optimizer = bnb.optim.Adam8bit(
    model.parameters(),
    lr=2e-5,
)

# PagedOptimizer — 自动将优化器状态分页到 CPU 内存
optimizer = bnb.optim.PagedAdamW8bit(
    model.parameters(),
    lr=2e-5,
)

import torch
import bitsandbytes as bnb

def replace_linear_with_8bit(model, threshold=6.0):
    """将模型中所有 nn.Linear 替换为 Linear8bitLt"""
    for name, module in model.named_children():
        if isinstance(module, torch.nn.Linear):
            # 创建 8-bit 替代层
            new_module = bnb.nn.Linear8bitLt(
                module.in_features,
                module.out_features,
                module.bias is not None,
                has_fp16_weights=False,
                threshold=threshold,
            )
            # 量化和替换
            new_module.weight = bnb.nn.Int8Params(
                module.weight.data.cpu(),
                has_fp16_weights=False,
                requires_grad=False,
            ).to(module.weight.dtype)
            if module.bias is not None:
                new_module.bias = module.bias
            setattr(model, name, new_module)
        else:
            # 递归处理子模块
            replace_linear_with_8bit(module, threshold)
    return model

def replace_linear_with_4bit(model, compute_dtype=torch.bfloat16,
                              quant_type='nf4', double_quant=True):
    """将模型中所有 nn.Linear 替换为 Linear4bit"""
    for name, module in model.named_children():
        if isinstance(module, torch.nn.Linear):
            new_module = bnb.nn.Linear4bit(
                module.in_features,
                module.out_features,
                module.bias is not None,
                compute_dtype=compute_dtype,
                compress_statistics=double_quant,
                quant_type=quant_type,
            )
            new_module.weight = bnb.nn.Params4bit(
                module.weight.data.cpu(),
                requires_grad=False,
                compress_statistics=double_quant,
                quant_type=quant_type,
            ).to(module.weight.dtype)
            if module.bias is not None:
                new_module.bias = module.bias
            setattr(model, name, new_module)
        else:
            replace_linear_with_4bit(module, compute_dtype, quant_type, double_quant)
    return model

import torch
from transformers import (
    AutoModelForCausalLM,
    AutoTokenizer,
    BitsAndBytesConfig,
    TrainingArguments,
    Trainer,
)
from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training
from datasets import load_dataset

# ========== 1. 配置 4-bit 量化 ==========
bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,                              # 启用 4-bit 加载
    bnb_4bit_quant_type="nf4",                      # NF4 量化格式
    bnb_4bit_compute_dtype=torch.bfloat16,          # 计算精度
    bnb_4bit_use_double_quantization=True,           # 双重量化
)

# ========== 2. 加载量化模型 ==========
model_name = "meta-llama/Llama-2-7b-hf"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    quantization_config=bnb_config,                 # 量化配置
    device_map="auto",                              # 自动分配设备
    torch_dtype=torch.bfloat16,
)

# ========== 3. 准备训练 ==========
model = prepare_model_for_kbit_training(
    model,
    use_gradient_checkpointing=True,                # 减少激活值显存
)

# ========== 4. 配置 LoRA ==========
lora_config = LoraConfig(
    r=16,                   # LoRA 秩
    lora_alpha=32,          # LoRA 缩放因子
    target_modules=["q_proj", "k_proj", "v_proj", "o_proj"],  # 目标层
    lora_dropout=0.05,      # Dropout 概率
    bias="none",
    task_type="CAUSAL_LM",
)
model = get_peft_model(model, lora_config)

# ========== 5. 准备数据 ==========
dataset = load_dataset("tatsu-lab/alpaca", split="train[:1000]")

def tokenize_function(examples):
    return tokenizer(
        examples["instruction"],
        truncation=True,
        max_length=512,
        padding="max_length",
    )

tokenized_dataset = dataset.map(tokenize_function, batched=True)

# ========== 6. 训练 ==========
training_args = TrainingArguments(
    output_dir="./qlora-output",
    per_device_train_batch_size=4,
    gradient_accumulation_steps=4,
    learning_rate=2e-4,
    num_train_epochs=3,
    logging_steps=10,
    save_strategy="epoch",
    bf16=True,                     # 使用 bf16 混合精度
    optim="paged_adamw_8bit",      # 使用 8-bit 分页优化器
)

trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=tokenized_dataset,
)
trainer.train()

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig

# 配置 8-bit 量化
bnb_config = BitsAndBytesConfig(
    load_in_8bit=True,
    llm_int8_threshold=6.0,        # 离群值阈值
)

# 加载模型
model_name = "meta-llama/Llama-2-7b-hf"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    quantization_config=bnb_config,
    device_map="auto",
)

# 推理
prompt = "Explain quantum computing in simple terms:"
inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
outputs = model.generate(
    **inputs,
    max_new_tokens=200,
    temperature=0.7,
    do_sample=True,
)
print(tokenizer.decode(outputs[0], skip_special_tokens=True))

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig

bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_compute_dtype=torch.bfloat16,
)

model = AutoModelForCausalLM.from_pretrained(
    "meta-llama/Llama-2-7b-hf",
    quantization_config=bnb_config,
    device_map="auto",
)
tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-2-7b-hf")

# 推理
inputs = tokenizer("Hello, how are you?", return_tensors="pt").to(model.device)
with torch.no_grad():
    outputs = model.generate(**inputs, max_new_tokens=50)
print(tokenizer.decode(outputs[0], skip_special_tokens=True))

import torch
import bitsandbytes as bnb
from transformers import AutoModelForCausalLM, AutoTokenizer

model = AutoModelForCausalLM.from_pretrained("gpt2-medium").cuda()
tokenizer = AutoTokenizer.from_pretrained("gpt2-medium")

# 使用 8-bit AdamW 优化器，节省优化器状态显存
optimizer = bnb.optim.PagedAdamW8bit(
    model.parameters(),
    lr=5e-5,
    weight_decay=0.01,
)

# 标准训练循环
model.train()
for batch in dataloader:
    inputs = tokenizer(batch["text"], return_tensors="pt", padding=True, truncation=True)
    inputs = {k: v.cuda() for k, v in inputs.items()}
    outputs = model(**inputs, labels=inputs["input_ids"])
    loss = outputs.loss
    loss.backward()
    optimizer.step()
    optimizer.zero_grad()