aitraining llm --train \
  --model meta-llama/Llama-3.2-8B \
  --data-path ./data.jsonl \
  --project-name quantized-model \
  --peft \
  --quantization int4

from autotrain.trainers.clm.params import LLMTrainingParams

params = LLMTrainingParams(
    model="meta-llama/Llama-3.2-8B",
    data_path="./data.jsonl",
    project_name="quantized-model",

    peft=True,
    quantization="int4",  # or "int8"
    lora_r=16,
)

params = LLMTrainingParams(
    ...
    quantization="int4",
)

params = LLMTrainingParams(
    ...
    quantization="int8",
)

params = LLMTrainingParams(
    ...
    peft=True,  # Required for quantized training
    quantization="int4",
)

params = LLMTrainingParams(
    ...
    peft=True,
    quantization="int4",
    lr=2e-4,  # Higher LR works well with QLoRA
)

params = LLMTrainingParams(
    ...
    quantization="int4",
    use_flash_attention_2=True,  # Requires Linux + CUDA + flash-attn package
)

import torch
from transformers import AutoModelForCausalLM, BitsAndBytesConfig

# 4-bit config (matches AITraining defaults)
bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_compute_dtype=torch.float16,
    bnb_4bit_use_double_quant=False,
)

model = AutoModelForCausalLM.from_pretrained(
    "meta-llama/Llama-3.2-8B",
    quantization_config=bnb_config,
)

# 在 Apple Silicon 上 - 使用 LoRA 而不使用量化以获得 MPS 加速
aitraining llm --train \
  --model google/gemma-3-270m \
  --data-path ./data.jsonl \
  --project-name mac-training \
  --peft \
  --lora-r 16