TRiP: Động cơ AI Transformer hoàn chỉnh viết bằng ngôn ngữ C từ con số 0

TRiP: Động cơ AI Transformer hoàn chỉnh viết bằng ngôn ngữ C từ con số 0

TRiP (viết tắt của TRansformer in Progress) là một dự án phần mềm thú vị, cung cấp một engine hoàn chỉnh dành cho các mô hình AI Transformer, được viết hoàn toàn bằng ngôn ngữ C thuần túy. Dự án này bao gồm mọi khía cạnh từ suy luận (inference), đào tạo (training), tạo tokenizer, chat cho đến xử lý thị giác, tất cả được gói gọn trong một vài tệp tin nhỏ gọn.

Được phát triển trong vòng 18 tháng (từ tháng 3 năm 2024 đến tháng 8 năm 2025) vào những giờ nghỉ trưa và cuối tuần của tác giả, TRiP ra đời không phải để cạnh tranh với các công cụ tối ưu hóa hiệu năng như llama.cpp, mà với mục đích giáo dục sâu sắc. Tác giả muốn xây dựng nó để thực sự hiểu rõ nội dung của kiến trúc Transformer, từ những phép nhân ma trận cơ bản nhất.

Mục đích và Triết lý phát triển

Mục tiêu chính của TRiP là phục vụ việc học tập, dành cho chính tác giả và bất kỳ ai muốn tìm hiểu sâu về cách thức hoạt động của Transformer. Dự án hỗ trợ các mô hình phổ biến như Gemma 1, Llama 2, PaliGemma và GPT-2 với khả năng suy luận và đào tạo đầy đủ.

Một điểm đáng chú ý là mức độ "tự tay làm" (hand-coded) của dự án. Mặc dù một số phần nhỏ như trình phân tích cú pháp JSON, chức năng lưu checkpoint safetensors, hoặc phần quản lý JPEG-X11 được tạo bởi AI, nhưng phần lớn logic cốt lõi—đặc biệt là các phép toán và kiến trúc mô hình—được tác giả viết thủ công. Điều này đảm bảo mục tiêu ban đầu là đạt được sự hiểu biết toàn diện ("full-stack understanding") về nội bộ của Transformer.

Các tính năng chính và Mô hình hỗ trợ

TRiP không chỉ là một công cụ demo đơn giản mà là một hệ thống khá hoàn chỉnh với các tính năng:

• Kiến trúc hỗ trợ: Llama2, Gemma 1.0/1.1, PaliGemma 1 (kết hợp thị giác và ngôn ngữ), GPT-2.
• Định dạng checkpoint: SafeTensors (HuggingFace), định dạng llama2.c và gpt2 của Karpathy.
• Kiểu dữ liệu trọng số: bf16, float16, float32.
• Đào tạo: Hỗ trợ đầy đủ lan truyền ngược (backpropagation) với trình tối ưu hóa AdamW, điều chỉnh tốc độ học cosine annealing và cắt gradient (gradient clipping).
• Tokenizer: Hỗ trợ BPE (tương thích SentencePiece) với khả năng tạo từ vựng từ đầu.
• Suy luận: Các phương pháp lấy mẫu như greedy, top-k và nucleus (top-p).
• Chat: Tương tác chat trực tiếp với các mẫu chat của Llama, Gemma và TinyLlama.
• Thị giác: Suy luận đa phương thức (multimodal) với PaliGemma (nhập JPEG, hiển thị qua X11).
• Bộ nhớ: Chế độ tối ưu hóa RAM thông qua mmap cho các mô hình lớn trên phần cứng hạn chế.

Cấu trúc mã nguồn và Triển khai

TRiP được tổ chức gọn gàng trong 7 tệp tin C, không phụ thuộc vào các framework nặng nề như PyTorch hay TensorFlow, mà chỉ sử dụng đại số tuyến tính trên các mảng số thực.

• trip.h: Bản đồ tổng thể, chứa mọi kiểu dữ liệu, cấu trúc và khai báo.
• math.c: Các phép toán tensor, với mỗi hàm forward đều đi kèm hàm backward ngay bên dưới (matmul, softmax, layernorm, attention, v.v.).
• forward.c: Điều phối quá trình truyền tiến và lấy mẫu token.
• backward.c: Quá trình truyền ngược, trình tối ưu hóa AdamW và quản lý gradient.
• model.c: Nhập/xuất checkpoint, khởi tạo mô hình, quản lý bộ nhớ, tokenizer và xử lý hình ảnh.
• utils.c: Ghi log, phân tích JSON, vào/ra terminal và xử lý ảnh JPEG/X11.
• main.c: Phân tích đối số dòng lệnh, vòng lặp chat, đào tạo và suy luận.

Cách cài đặt và Sử dụng

Để biên dịch TRiP, bạn chỉ cần một trình biên dịch gcc (khuyến nghị bản 13 trở lên để hỗ trợ bfloat16) và một vài thư viện hệ thống. Không cần CMake hay Python.

Trên Debian/Ubuntu, cài đặt các phụ thuộc:

sudo apt install build-essential libomp-dev libjpeg-dev libx11-dev

Biên dịch dự án chỉ với một lệnh:

make

Ví dụ sử dụng

Để chat với mô hình Gemma:

./trip --chat \
--checkpoint gemma-2b-it/model.safetensors \
--tokenizer gemma-2b-it/tokenizer.json \
--chat_scheme GEMMA

Để chạy suy luận trên một đoạn văn bản:

./trip --decode \
--input_text "Thủ đô của nước Ý là" \
--checkpoint gemma-2b-it/model.safetensors \
--tokenizer gemma-2b-it/tokenizer.json

Để đào tạo một mô hình mới:

./trip --train \
--checkpoint my_model/model.safetensors \
--tokenizer my_model/tokenizer.json \
--train_data my_dataset.txt \
--train_config training_args.json

Hiểu sâu về cơ chế hoạt động

TRiP triển khai Transformer từ các nguyên lý đầu tiên (first principles). Khái niệm trung tâm là "dòng chảy dư" (residual stream)—một vectơ chảy qua mô hình giống như dữ liệu trên một bus. Mỗi lớp đọc từ nó, xử lý qua attention và mạng feed-forward, sau đó ghi lại vào đó.

Tác giả đã chú thích mã nguồn rất kỹ lưỡng, biến TRiP thành một cuốn sách giáo khoa có chú thích về Transformer. Mỗi phép toán trong math.c đều được viết dưới dạng cặp forward/backward, cho phép người đọc thấy chính xác gradient chảy ngược qua phép tính đó như thế nào.

Dự án được cấp phép theo CC BY-NC 4.0, cho phép sử dụng, nghiên cứu và sửa đổi miễn phí cho mục đích phi thương mại. Đây là một tài nguyên quý giá cho những ai muốn hiểu rõ "bụng" của các mô hình AI hiện đại mà không bị che khuất bởi các abstraction layer của các thư viện cao cấp.