Fine-Tuning LLM Tốc Độ Cao với Unsloth trên Linux: Tiết Kiệm 80% VRAM trên GPU Phổ Thông

AI tutorial - IT technology blog
AI tutorial - IT technology blog

Bối Cảnh & Lý Do: Thực Tế Khi Fine-Tuning Trên GPU Phổ Thông

Nếu đã từng thử chạy PEFT/LoRA tiêu chuẩn của Hugging Face trên các dòng card đồ họa phổ thông, gần như chắc chắn bạn sẽ gặp phải bức tường quen thuộc: lỗi tràn bộ nhớ CUDA (Out-of-Memory). Việc fine-tuning một mô hình 8B như Llama 3 hay Qwen 2.5 trên RTX 3060 (12GB) hoặc RTX 4070 (12GB) thường đòi hỏi phải offload dữ liệu liên tục, hạ batch size xuống mức cực nhỏ, hoặc cắn răng thuê máy chủ đám mây đắt đỏ. Quá trình huấn luyện diễn ra chậm chạp và những rắc rối khi thiết lập môi trường cứ thế chồng chất.

Unsloth giải quyết triệt để vấn đề này ngay từ tầng gốc rễ. Thay vì chỉ vá víu PyTorch ở tầng API bậc cao, nhóm phát triển đã viết lại hoàn toàn các compute kernel cốt lõi bằng ngôn ngữ Triton của OpenAI. Họ tự tay tối ưu thủ công các lượt lan truyền ngược (backward pass) cho RoPE embedding, cross-entropy loss và các phép chiếu MLP. Nhờ vậy, engine này loại bỏ được việc lưu trữ hàng loạt tensor trung gian cồng kềnh trong lượt lan truyền xuôi (forward pass).

Trên thực tế, điều này mang lại những lợi ích vượt trội:

  • Tiết kiệm tới 80% bộ nhớ: Một mô hình 8B tham số có thể huấn luyện mượt mà chỉ với 6GB đến 8GB VRAM mà không bị tràn sang bộ nhớ RAM hệ thống chậm chạp.
  • Tốc độ huấn luyện nhanh gấp 2x đến 5x: Giảm lượng dữ liệu luân chuyển qua bus PCIe đồng nghĩa với việc băng thông bộ nhớ không còn là nút thắt cổ chai kìm hãm nhân GPU.
  • Độ chính xác toán học tuyệt đối: Hoàn toàn không cắt tỉa (pruning) mất mát dữ liệu hay dùng phép tính xấp xỉ thô. Gradient hoàn toàn khớp từng bit so với huấn luyện PyTorch FP16 chuẩn.

Nếu bạn đang phát triển các mô hình chuyên biệt cho từng lĩnh vực, bước chuyển dịch này thay đổi hoàn toàn cuộc chơi. Giờ đây bạn có thể thử nghiệm, lặp lại quy trình trực tiếp trên máy tính để bàn thay vì phải trả chi phí thuê từng giờ cho các cụm máy chủ nhiều card A100.

Cài Đặt: Thiết Lập Môi Trường Linux Chuẩn

Một môi trường dependencies sạch sẽ giúp loại bỏ đến 90% các lỗi đau đầu liên quan đến CUDA. Trước khi cài đặt bất kỳ thứ gì, hãy khởi tạo một môi trường Conda hoặc venv riêng biệt trên máy Linux của bạn.

1. Kiểm Tra Driver NVIDIA & CUDA

Unsloth phụ thuộc rất nhiều vào các custom kernel viết bằng Triton. Các kernel này chạy tối ưu nhất trên các kiến trúc NVIDIA hiện đại, bao gồm Turing (RTX 2000), Ampere (RTX 3000) và Ada Lovelace (RTX 4000). Hãy bắt đầu bằng cách kiểm tra trạng thái driver:

nvidia-smi

Hãy đảm bảo driver đã cài đặt hỗ trợ CUDA 12.1 hoặc 12.4 (phiên bản driver 535.xx trở lên).

2. Tạo Môi Trường Biệt Lập

# Tạo và kích hoạt môi trường conda mới
conda create -n unsloth_env python=3.10 -y
conda activate unsloth_env

3. Cài Đặt PyTorch và Unsloth

Tải bản dựng PyTorch dành riêng cho CUDA 12.1, sau đó cài đặt Unsloth cùng các thư viện phụ thuộc cốt lõi:

# Cài đặt PyTorch hỗ trợ CUDA 12.1
pip install --upgrade --no-cache-dir torch torchvision torchaudio \
    --index-url https://download.pytorch.org/whl/cu121

# Cài đặt Unsloth và các thư viện tăng tốc cần thiết
pip install --no-deps "unsloth[colab-new] @ git+https://github.com/unslothai/unsloth.git"
pip install --no-deps trl peft accelerate bitsandbytes datasets xformers

Cấu Hình: Tinh Chỉnh Mô Hình Llama 3 và Qwen

Unsloth tích hợp mượt mà vào hệ sinh thái TRL (Transformer Reinforcement Learning) của Hugging Face. Đoạn mã bên dưới hướng đến mô hình Llama 3 (8B) hoặc Qwen 2.5 (7B) sử dụng kỹ thuật lượng tử hóa 4-bit NormalFloat (NF4) kết hợp với adapter LoRA.

1. Tải Mô Hình Nền Đã Lượng Tử Hóa 4-bit

Lưu đoạn mã cấu hình sau vào tệp train_lora.py:

import torch
from unsloth import FastLanguageModel

max_seq_length = 2048 # Điều chỉnh tùy theo nhu cầu tập dữ liệu của bạn
dtype = None          # None = Tự động nhận diện (Float16 cho card đời cũ, Bfloat16 cho Ampere trở lên)
load_in_4bit = True   # Giữ dung lượng mô hình nền dưới mức 6GB

# Chuyển sang unsloth/Qwen2.5-7B-Instruct-bnb-4bit nếu làm việc với Qwen
model_name = "unsloth/Meta-Llama-3-8B-Instruct-bnb-4bit"

model, tokenizer = FastLanguageModel.from_pretrained(
    model_name=model_name,
    max_seq_length=max_seq_length,
    dtype=dtype,
    load_in_4bit=load_in_4bit,
)

2. Gắn Adapter LoRA Tốc Độ Cao

Tiếp theo, tiến hành patch các lớp linear projection. Unsloth can thiệp trực tiếp vào các lớp attention và MLP projection mà không làm phát sinh thêm overhead của Python:

model = FastLanguageModel.get_peft_model(
    model,
    r=16, # Chiều rank (16 hoặc 32 mang lại năng lực suy luận rất tốt)
    target_modules=["q_proj", "k_proj", "v_proj", "o_proj",
                    "gate_proj", "up_proj", "down_proj"],
    lora_alpha=16,
    lora_dropout=0,       # Tối ưu về mức 0% để tăng tốc với Triton kernel
    bias="none",
    use_gradient_checkpointing="unsloth", # Cắt giảm bộ nhớ activation khoảng 30%
    random_state=3407,
)

3. Định Dạng Tập Dữ Liệu và Thực Thi Trainer

Bây giờ hãy chuyển đổi cấu trúc hội thoại của bạn theo đúng template gốc của mô hình. Chúng ta sẽ sử dụng SFTTrainer của Hugging Face để điều phối vòng lặp huấn luyện:

from datasets import load_dataset
from trl import SFTTrainer
from transformers import TrainingArguments
from unsloth import is_bfloat16_supported

# Ánh xạ các bản ghi vào chat template gốc của mô hình
def formatting_prompts_func(examples):
    convos = examples["conversations"]
    texts = [tokenizer.apply_chat_template(convo, tokenize=False, add_generation_prompt=False) for convo in convos]
    return {"text": texts}

dataset = load_dataset("philschmid/dolly-15k-oai-style", split="train")
dataset = dataset.map(formatting_prompts_func, batched=True)

trainer = SFTTrainer(
    model=model,
    tokenizer=tokenizer,
    train_dataset=dataset,
    dataset_text_field="text",
    max_seq_length=max_seq_length,
    dataset_num_proc=2,
    packing=False, # Đặt thành True để ghép các mẫu ngắn lại với nhau giúp tăng tốc độ
    args=TrainingArguments(
        per_device_train_batch_size=2,
        gradient_accumulation_steps=4,
        warmup_steps=10,
        max_steps=60,
        learning_rate=2e-4,
        fp16=not is_bfloat16_supported(),
        bf16=is_bfloat16_supported(),
        logging_steps=5,
        optim="adamw_8bit",
        weight_decay=0.01,
        lr_scheduler_type="linear",
        seed=3407,
        output_dir="outputs",
    ),
)

trainer_stats = trainer.train()

Kiểm Thử & Giám Sát: Theo Dõi VRAM & Xuất Trọng Số

1. Giám Sát Phần Cứng Theo Thời Gian Thực

Khởi chạy script huấn luyện và mở một cửa sổ terminal thứ hai để theo dõi trực tiếp mức tiêu thụ phần cứng:

# Lấy thông số GPU định kỳ mỗi giây
watch -n 1 nvidia-smi

Hãy chú ý đến biểu đồ bộ nhớ. Trên RTX 3090, quá trình huấn luyện mô hình 8B với context length 2.048 token chỉ tiêu tốn đỉnh điểm khoảng 6.8GB đến 7.4GB VRAM. Trong khi đó, phương pháp LoRA tiêu chuẩn chưa tối ưu dễ dàng vượt quá 22GB cho cùng một thiết lập tham số. Khoảng trống bộ nhớ dư dả này cho phép bạn tự tin tăng gấp đôi batch size hoặc mở rộng context length mà không phải lo sợ lỗi OOM.

2. Kiểm Thử Chất Lượng Suy Luận Ngay Tại Local

Khi quá trình huấn luyện hoàn tất, hãy kiểm tra khả năng suy luận trực tiếp trong script của bạn. Đảm bảo mô hình phản hồi chính xác theo prompt khi đã tích hợp adapter vào bộ nhớ:

# Kích hoạt engine suy luận tích hợp sẵn, nhanh gấp 2 lần
FastLanguageModel.for_inference(model)

messages = [
    {"role": "user", "content": "Hãy giải thích bằng từ ngữ đơn giản cách Unsloth giảm VRAM trong quá trình huấn luyện."}
]
inputs = tokenizer.apply_chat_template(
    messages,
    tokenize=True,
    add_generation_prompt=True,
    return_tensors="pt"
).to("cuda")

outputs = model.generate(input_ids=inputs, max_new_tokens=128, use_cache=True)
print(tokenizer.batch_decode(outputs))

3. Xuất Ra Định Dạng GGUF Hoặc Hugging Face 16-bit

Bạn đã sẵn sàng triển khai mô hình? Hãy xuất adapter trực tiếp sang môi trường runtime mong muốn mà không cần bất kỳ script chuyển đổi Python rời nào:

# Xuất sang định dạng Hugging Face 16-bit đã gộp trọng số dành cho vLLM hoặc TGI
model.save_pretrained_merged("llama3_custom_model", tokenizer, save_method="merged_16bit")

# Hoặc xuất trực tiếp sang định dạng GGUF lượng tử hóa cho Ollama và llama.cpp
model.save_pretrained_gguf("llama3_gguf", tokenizer, quantization_method="q4_k_m")

Nhờ các Triton kernel tối ưu sâu cùng kỹ thuật lượng tử hóa 4-bit, việc fine-tuning các mô hình ngôn ngữ lớn hiện đại giờ đây không còn đòi hỏi hàng ngàn USD chi phí hạ tầng đám mây. Chiếc card đồ họa trên máy tính để bàn của bạn hoàn toàn có thể đảm đương xuất sắc công việc này.

Share: