W' = W + BA

import numpy as np
from scipy.stats import norm

def compute_nf4_levels() -> np.ndarray:
    """
    计算 NF4 的 16 个量化级别
    NF4 假设权重归一化到 [-1, 1] 且服从标准正态分布
    """
    # 16 个等概率分位点
    quantiles = np.linspace(0, 1, 17)[1:-1]  # 去掉 0 和 1
    levels = norm.ppf(quantiles)              # 正态分位函数
    # 归一化到 [-1, 1]
    levels = levels / np.max(np.abs(levels))
    # 添加 -1 和 1 端点
    levels = np.concatenate([[-1.0], levels, [1.0]])
    return levels

nf4_levels = compute_nf4_levels()
print(f"NF4 levels ({len(nf4_levels)}): {nf4_levels}")
# 16 个非均匀分布的量化级别，在 0 附近密集，两端稀疏

原始权重 (FP16, 2 bytes) → 4bit 量化 (0.5 bytes) + 量化常数 (FP32, 4 bytes per block)
→ 双重量化: 量化常数也做 8bit (0.5 bytes per block)

pip install transformers==4.41.0
pip install peft==0.11.0
pip install bitsandbytes==0.43.0
pip install accelerate==0.30.0
pip install datasets==2.19.0
pip install trl==0.9.0

import json
from datasets import Dataset, DatasetDict

def load_alpaca_format_data(path: str) -> Dataset:
    """
    加载 Alpaca 格式数据集
    预期格式: [{"instruction": "...", "input": "...", "output": "..."}, ...]
    """
    with open(path, 'r', encoding='utf-8') as f:
        data = json.load(f)

    def format_prompt(example):
        if example.get('input'):
            text = (
                "Below is an instruction that describes a task, paired with an input. "
                "Write a response that appropriately completes the request.\n\n"
                f"### Instruction:\n{example['instruction']}\n\n"
                f"### Input:\n{example['input']}\n\n"
                f"### Response:\n{example['output']}"
            )
        else:
            text = (
                "Below is an instruction that describes a task. "
                "Write a response that appropriately completes the request.\n\n"
                f"### Instruction:\n{example['instruction']}\n\n"
                f"### Response:\n{example['output']}"
            )
        return {"text": text}

    dataset = Dataset.from_list(data)
    dataset = dataset.map(format_prompt, remove_columns=dataset.column_names)
    return dataset

# 示例：加载训练集和验证集
train_dataset = load_alpaca_format_data("train.json")
eval_dataset = load_alpaca_format_data("eval.json")

import torch
from transformers import (
    AutoModelForCausalLM,
    AutoTokenizer,
    BitsAndBytesConfig,
    TrainingArguments,
)
from peft import (
    LoraConfig,
    get_peft_model,
    prepare_model_for_kbit_training,
    TaskType,
)

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

# ===== 加载模型 =====
model_name = "meta-llama/Meta-Llama-3-8B-Instruct"
tokenizer = AutoTokenizer.from_pretrained(model_name)
tokenizer.pad_token = tokenizer.eos_token

model = AutoModelForCausalLM.from_pretrained(
    model_name,
    quantization_config=bnb_config,
    device_map="auto",
    torch_dtype=torch.bfloat16,
    trust_remote_code=True,
)

# 准备 kbit 训练
model = prepare_model_for_kbit_training(model)

# ===== LoRA 配置 =====
lora_config = LoraConfig(
    task_type=TaskType.CAUSAL_LM,
    r=64,                                  # LoRA 秩
    lora_alpha=16,                         # 缩放系数
    lora_dropout=0.05,                     # Dropout
    bias="none",
    target_modules=[
        "q_proj", "k_proj", "v_proj", "o_proj",  # Attention
        "gate_proj", "up_proj", "down_proj",      # MLP
    ],
)

model = get_peft_model(model, lora_config)
model.print_trainable_parameters()
# 输出: trainable params: 41,943,040 || all params: 8,072,204,288 || trainable%: 0.52%

from trl import SFTTrainer, SFTConfig

# ===== 训练参数 =====
training_args = SFTConfig(
    output_dir="./output/llama3-8b-qlora",
    num_train_epochs=3,
    per_device_train_batch_size=4,
    per_device_eval_batch_size=4,
    gradient_accumulation_steps=4,         # 有效 batch = 4×4 = 16
    gradient_checkpointing=True,           # 梯度检查点，省显存
    optim="paged_adamw_8bit",             # 分页 8bit AdamW
    learning_rate=2e-4,
    weight_decay=0.01,
    max_grad_norm=0.3,
    warmup_ratio=0.03,
    lr_scheduler_type="cosine",
    logging_steps=10,
    save_strategy="epoch",
    eval_strategy="epoch",
    bf16=True,                            # 使用 BF16 混合精度
    max_seq_length=512,                   # 最大序列长度
    dataset_text_field="text",
    packing=True,                         # 打包短序列提高利用率
)

# ===== 启动训练 =====
trainer = SFTTrainer(
    model=model,
    args=training_args,
    train_dataset=train_dataset,
    eval_dataset=eval_dataset,
    tokenizer=tokenizer,
)

trainer.train()

# 保存 LoRA adapter
model.save_pretrained("./output/llama3-8b-qlora/adapter")
tokenizer.save_pretrained("./output/llama3-8b-qlora/adapter")

from peft import PeftModel

# 加载基座模型 + LoRA adapter 推理
base_model = AutoModelForCausalLM.from_pretrained(
    model_name,
    quantization_config=bnb_config,
    device_map="auto",
)
model = PeftModel.from_pretrained(base_model, "./output/llama3-8b-qlora/adapter")

# 生成
inputs = tokenizer("### Instruction:\n解释量子纠缠\n\n### Response:\n", return_tensors="pt").to(model.device)
outputs = model.generate(**inputs, max_new_tokens=256, temperature=0.7)
print(tokenizer.decode(outputs[0], skip_special_tokens=True))

# ===== 合并为完整模型（可选，需要足够显存）=====
# merged_model = model.merge_and_unload()
# merged_model.save_pretrained("./output/llama3-8b-merged")

# 仅 Attention（参数少，效果一般）
target_modules = ["q_proj", "v_proj"]

# Attention 全量（推荐起步配置）
target_modules = ["q_proj", "k_proj", "v_proj", "o_proj"]

# Attention + MLP（效果最好，参数稍多）
target_modules = ["q_proj", "k_proj", "v_proj", "o_proj",
                  "gate_proj", "up_proj", "down_proj"]

# 方式一：在 TrainingArguments 中启用
training_args = TrainingArguments(
    gradient_checkpointing=True,
    # ...
)

# 方式二：手动启用
model.gradient_checkpointing_enable()
model.enable_input_require_grads()  # 必须配合使用

model = AutoModelForCausalLM.from_pretrained(
    model_name,
    quantization_config=bnb_config,
    attn_implementation="flash_attention_2",  # 启用 Flash Attention 2
    device_map="auto",
)

总显存 ≈ 模型权重(4bit) + LoRA参数(FP32) + 梯度 + 优化器状态 + 激活值

≈ N × 0.5 bytes           (4bit 权重)
+ N × r/d × 4 bytes × 3   (LoRA: 参数 + 梯度 + Adam状态)
+ batch_size × seq_len × hidden_dim × 2 bytes  (激活值，开启梯度检查点后大幅减少)

# 排查步骤
# 1. 减小 batch_size
per_device_train_batch_size=1
# 2. 增大 gradient_accumulation_steps 补偿
gradient_accumulation_steps=16
# 3. 确认梯度检查点已开启
gradient_checkpointing=True
# 4. 减小 max_seq_length
max_seq_length=256
# 5. 使用 paged_adamw_8bit（而非 paged_adamw_32bit）
optim="paged_adamw_8bit"

# 检查项
# 1. learning_rate：QLoRA 推荐 1e-4 ~ 3e-4
# 2. lora_alpha：太小导致更新幅度不足，尝试 alpha = 2*r
# 3. target_modules：仅 q_proj/v_proj 可能不够，加 MLP
# 4. 数据格式：确保 prompt 和 completion 格式正确
# 5. 检查是否忘记 prepare_model_for_kbit_training

# 1. 检查数据质量：人工抽检训练样本
# 2. 调整生成参数
generation_config = {
    "temperature": 0.7,
    "top_p": 0.9,
    "repetition_penalty": 1.1,
    "max_new_tokens": 256,
}
# 3. 增加 epoch（但注意过拟合）
# 4. 增大 LoRA rank

# deepspeed_zero2_config.json
{
  "bf16": {
    "enabled": true
  },
  "zero_optimization": {
    "stage": 2,
    "offload_optimizer": {
      "device": "cpu",
      "pin_memory": true
    }
  },
  "gradient_accumulation_steps": "auto",
  "gradient_clipping": "auto",
  "train_batch_size": "auto",
  "train_micro_batch_size_per_gpu": "auto"
}

# 启动命令
# accelerate launch --num_processes 2 --use_deepspeed \
#   --deepspeed_config_file deepspeed_zero2_config.json train.py

数据准备 → 格式化（Alpaca/ChatML） → 分词
    ↓
模型加载（4bit NF4） → prepare_for_kbit_training
    ↓
LoRA 配置（r, alpha, target_modules） → get_peft_model
    ↓
SFTTrainer 配置（batch, lr, scheduler, paged_adamw_8bit）
    ↓
训练（gradient_checkpointing + bf16）
    ↓
评估 → 保存 adapter
    ↓
推理测试 / 合并导出