# pip install peft transformers bitsandbytes datasets

from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training
from datasets import Dataset
import torch

# 1. 加载模型（4-bit 量化）
model_id = "Qwen/Qwen3.8-Flash-Next"  # 以 HF 模型卡实际发布为准
bnb_config = BitsAndBytesConfig(load_in_4bit=True, bnb_4bit_compute_dtype=torch.float16)
model = AutoModelForCausalLM.from_pretrained(model_id, quantization_config=bnb_config)
tok = AutoTokenizer.from_pretrained(model_id)

# 2. 配置 LoRA
lora_config = LoraConfig(
    r=16,               # LoRA 秩，越大表达能力越强、显存越多
    lora_alpha=32,      # 缩放系数，通常为 r 的 2 倍
    target_modules=["q_proj", "v_proj"],  # 目标模块
    lora_dropout=0.05,
    task_type="CAUSAL_LM",
)
model = prepare_model_for_kbit_training(model)
model = get_peft_model(model, lora_config)

# 3. 准备训练数据（instruction-tuning 格式）
train_data = Dataset.from_dict({
    "text": [
        "### 指令：\n解释什么是 KV cache\n\n### 回答：\nKV cache 是...",
        # 你的真实训练样本
    ]
})

# 4. 训练（用 transformers Trainer 或 TRL 库）
from transformers import TrainingArguments, Trainer
training_args = TrainingArguments(
    output_dir="./output",
    num_train_epochs=3,
    per_device_train_batch_size=2,
    gradient_accumulation_steps=4,
    learning_rate=2e-4,
    fp16=True,
    logging_steps=10,
)
trainer = Trainer(model=model, args=training_args, train_dataset=train_data)
trainer.train()

# 5. 保存 LoRA 适配层
model.save_pretrained("./my-lora-adapter")

{
  "instruction": "解释什么是上下文窗口",
  "input": "",
  "output": "上下文窗口是模型一次能处理的最大 token 数量..."
}