Các Gradient cục bộ Nhân lên
Forward Pass
Forward pass của ANDREA-120M đưa input qua một chuỗi các phép toán:
x = embed(token_ids) # nhúng token
for layer in 12_layers:
x = x + attn(LN(x)) # tầng con attention
x = x + mlp(LN(x)) # tầng con MLP
logits = LN(x) @ embed.T # chiếu đầu ra chung
loss = cross_entropy(logits, targets)Mỗi hoạt động đọc các tensor đầu vào & tạo ra các tensor đầu ra. Quá trình lan truyền thuận kết thúc bằng một số vô hướng duy nhất: hàm mất mát cross-entropy cho batch này.
Quá trình Lan truyền Ngược
Việc huấn luyện cập nhật các trọng số theo hướng làm giảm mất mát. Để lấy hướng cập nhật, engine cần:
dL/dW cho mọi W có thể học được trong mô hìnhQuy tắc chuỗi cho điều này. Với một chuỗi loss = f(g(h(x))):
dL/dx = (dL/df) * (df/dg) * (dg/dh) * (dh/dx)Mỗi yếu tố là một gradient cục bộ: cách đầu ra của một phép toán thay đổi khi đầu vào của nó thay đổi một lượng nhỏ. Nhân các gradient cục bộ ngược qua đồ thị sẽ lan truyền tín hiệu mất mát đến mọi trọng số.
Phân biệt ngược (Reverse-Mode Differentiation)
Backprop tính gradient theo thứ tự ngược: bắt đầu từ dL/dlogits = 1, sau đó đi ngược qua cross-entropy, rồi projection đầu ra, rồi layer norm, rồi mười hai khối transformer, rồi embeddings. Tại mỗi bước, nhân gradient đến với Jacobian cục bộ.
Chế độ ngược hiệu quả khi đầu ra là một số vô hướng duy nhất (mất mát) & có nhiều đầu vào (các trọng số). Một lần truyền ngược tạo gradient cho mọi trọng số trong mô hình. Chế độ tiến (forward-mode) cần một lần truyền cho mỗi trọng số; với ANDREA-120M có ~120M trọng số, forward-mode là không khả thi.
Tại sao Reverse-Mode
Mỗi Phép Toán Thuận Được Ghép Với Phép Toán Ngược
Quy Tắc Ghép Đôi
microgpt_cuda.cu cung cấp hai kernel CUDA cho mọi phép toán: một kernel tính toán đầu ra thuận, một kernel tính toán gradient đầu vào dựa trên gradient đầu ra. Sự ghép đôi là một-một:
| Kernel tiến | Kernel lùi | Thao tác |
|---|---|---|
k_embed_fwd | k_embed_bwd | Tra cứu embedding token |
k_layernorm_fwd | k_layernorm_bwd | Chuẩn hóa lớp |
k_attn_qkv_fwd | k_attn_qkv_bwd | Chiếu Q, K, V |
k_attn_fwd | k_attn_bwd | Attention scaled dot-product |
k_attn_out_fwd | k_attn_out_bwd | Chiếu đầu ra W_O |
k_mlp_fwd | k_mlp_bwd | MLP (với GELU) |
k_residual_add | k_residual_add_bwd | Kết nối dư |
k_loss_fwd | k_loss_bwd | Hàm mất mát Cross-entropy |
Tám cặp phép toán bao phủ toàn bộ transformer. Cộng thêm một vài kernel tiện ích: k_grad_norm_partial, k_grad_norm_final, k_grad_scale cho việc cắt gradient (xem hoạt động 75).
Nhiệm vụ của một Backward Kernel
Cho gradient chảy vào từ các lớp sau (grad_output), một backward kernel tính toán:
1. grad_input: gradient liên quan đến tensor đầu vào của phép toán. Nó được truyền tiếp ngược trở lại.
2. grad_weight: gradient liên quan đến các tham số có thể học được trong phép toán. Nó được đưa vào trạng thái của optimizer.
Cả hai đều được tính toán trong một lần khởi chạy kernel. Các thread CUDA hợp tác song song trên các tile của tensor gradient.
Các Tensor Đã Lưu
Phép tính ngược thường cần các giá trị từ forward pass. Ví dụ, k_layernorm_bwd cần mean & variance được tính trong forward; k_mlp_bwd cần pre-activation GELU. Engine huấn luyện lưu trữ chúng trong các buffer dành riêng trong forward, sau đó đọc chúng trong backward.
Chi phí bộ nhớ: gần giống hình dạng của đầu ra forward cho mỗi tensor được lưu. Với ANDREA-120M, batch=8, seq=1024, d_model=768, một tensor được lưu là 8 × 1024 × 768 × 4 bytes = 25 MB. Qua 12 lớp & nhiều tensor được lưu mỗi lớp, activations chiếm ưu thế VRAM trong quá trình huấn luyện (~5-10 GB trên card 24 GB).
Theo dõi Một Bước Backward
Nơi Gradient Được Lưu Trữ Trong Bộ Nhớ
Một Tensor Gradient Cho Mỗi Tensor Trọng Số
Mỗi tensor trọng số có thể học được trong ANDREA-120M đều có một tensor gradient tương ứng với hình dạng giống hệt. Đối với mỗi block:
W_Q [768, 768] ↔ grad_W_Q [768, 768]
W_K [768, 768] ↔ grad_W_K [768, 768]
W_V [768, 768] ↔ grad_W_V [768, 768]
W_O [768, 768] ↔ grad_W_O [768, 768]
W_1 [768, 3072] ↔ grad_W_1 [768, 3072]
W_2 [3072, 768] ↔ grad_W_2 [3072, 768]
LN1.gamma [768] ↔ grad_LN1.gamma [768]
LN1.beta [768] ↔ grad_LN1.beta [768]
LN2.gamma [768] ↔ grad_LN2.gamma [768]
LN2.beta [768] ↔ grad_LN2.beta [768]Cộng với token embeddings, position embeddings, & một layer norm cuối cùng. Tổng bộ nhớ buffer gradient khớp với bộ nhớ weight: ~120M floats, ~480 MB tại FP32, ~240 MB tại FP16.
Tích lũy Qua Các Microbatch
batch_size = 8 của ANDREA vừa với VRAM tại FP16. Các batch hiệu quả lớn hơn yêu cầu gradient accumulation: chạy nhiều forward+backward passes trên các batch nhỏ, cộng dồn gradients vào cùng buffer, sau đó thực hiện một bước optimizer.
for microbatch in range(n_microbatches):
forward(microbatch)
backward() # THÊM vào bộ đệm grad, không ghi đè
scale_grads(1.0 / n_microbatches) # tính trung bình qua các microbatches
optimizer_step()
zero_grads() # đặt lại cho bước huấn luyện tiếp theoCác kernel backward sử dụng ngữ nghĩa +=, không phải =. Mỗi lần gọi thêm đóng góp gradient vào bộ đệm hiện có; bộ đệm giữ tổng chạy cho đến khi zero_grads() xóa nó.
Trạng thái của Optimizer
AdamW (hoạt động 73) giữ hai buffer nữa cho mỗi trọng số: moment thứ nhất m & moment thứ hai v. Tổng bộ nhớ thời gian huấn luyện:
trọng số: 1× số lượng trọng số
gradient: 1× số lượng trọng số
Adam m: 1× số lượng trọng số
Adam v: 1× số lượng trọng số
các acts đã lưu: ~2-4× tùy thuộc vào số lớp & batch
──────────────────────────────────────────
tổng cộng: ~6-8× số lượng trọng sốANDREA-120M ở FP16: ~240 MB × 4 buffers (trọng số, grad, m, v) + ~5-10 GB activations = ~10-12 GB tổng cộng. Thoải mái dưới giới hạn 24 GB của RTX 4090. ANDREA-12M huấn luyện trong 1.4 GB; việc mở rộng quy mô tham số 10× mang lại ~10× bộ nhớ.
Kích thước Buffer Gradient
Kiểm soát hoàn toàn Bộ nhớ & Độ chính xác
Chi phí của các Framework chung
PyTorch & JAX làm autograd tiện lợi: viết code Python, nhận gradient tự động. Chi phí: một lớp dispatch chung giữa code của bạn & CUDA. Mọi thao tác đều qua overhead trình thông dịch Python, sổ sách framework, & chọn kernel động. Khi huấn luyện một mô hình ngôn ngữ nhỏ trên một GPU, overhead đó quan trọng.
Chi phí cụ thể mà ANDREA tránh được:
1. Độ trễ trình thông dịch Python. Mỗi phép toán PyTorch đều vượt qua ranh giới Python/C++. Với ~100 lần khởi chạy kernel mỗi bước huấn luyện, ~9 bước/phút, đó là ~900 lần vượt ranh giới mỗi phút. Phân phối ở mức C loại bỏ điều này.
2. Tính không dự đoán được của bộ cấp phát framework. Bộ cấp phát bộ nhớ đệm của PyTorch mang lại thông lượng tốt trung bình nhưng bộ nhớ đỉnh không dự đoán được. Công cụ huấn luyện của ANDREA cấp phát trước mọi bộ đệm lúc khởi động; không cấp phát lại trong quá trình huấn luyện, không phân mảnh, không OOM bất ngờ ở bước 100K.
3. Lựa chọn kernel chung chung. PyTorch chọn kernel lúc chạy qua các heuristic. ANDREA chọn kernel lúc biên dịch, được điều chỉnh cho kích thước ô tensor core RTX 4090.
4. Đường ống xử lý độ chính xác hỗn hợp. Đường dẫn cuBLAS FP16 của ANDREA-120M & các thí nghiệm tensor core FP8 E4M3 của ANDREA yêu cầu kiểm soát chính xác tensor nào sống ở độ chính xác nào. Các framework chung tiếp xúc kiểm soát này qua các API phân tầng; CUDA tùy chỉnh viết trực tiếp.
Sự đánh đổi
Chi phí CUDA tùy chỉnh: nhiều code hơn để viết, nhiều bug hơn để tìm, không có hệ sinh thái cộng đồng. microgpt_cuda.cu của ANDREA là ~6000 dòng CUDA viết tay mất hàng tháng để debug. Mỗi hoạt động mới yêu cầu viết kernel forward, kernel backward, & tests.
Những gì ANDREA đạt được:
- Tái lập hoàn toàn. Đường ống huấn luyện là một tệp nhị phân C cộng với một proxy Python. Không có sự lệch phiên bản qua các bản phát hành PyTorch, không có sự không khớp phiên bản CUDA với các wheel của framework.
- Tiếp tục chính xác từng bit. SIGTERM kích hoạt việc ghi checkpoint nắm bắt mọi tensor chính xác như GPU nhìn thấy. Tiếp tục nhặt lại cùng quỹ đạo loss mà lần chạy đang theo.
- Bộ nhớ có thể dự đoán. ANDREA-120M được huấn luyện 200K bước mà không có OOM. Bộ nhớ được tính toán ngay từ lúc khởi động engine.
- Truy cập phần cứng trực tiếp. Kích thước tile tensor core, cài đặt FP8 E4M3, sao chép bộ nhớ bất đồng bộ: tất cả đều có thể truy cập trực tiếp trong CUDA, mờ mịt trong các framework chung chung.
Tái lập Như Sứ Mệnh
Phần 9 của whitepaper ANDREA liệt kê toàn bộ stack tái lập:
**Động cơ huấn luyện:** microgpt/microgpt_cuda.cu
**Proxy huấn luyện:** microgpt/training_proxy.py
**Cấu hình thí nghiệm:** experiments/ANDREA-*-TRAIN.json
**Đường ống dữ liệu:** scripts/pull-hermes3.py, scripts/prep-megachat.py
**Bảng điều khiển:** scripts/live-loss-dashboard.html
**Thông số Bandit:** docs/FIREHOSE-BANDIT.md
**Tài liệu mô hình:** docs/ANDREA.mdYêu cầu phần cứng: một GPU NVIDIA với ≥8 GB VRAM (RTX 3060 hoặc tốt hơn). Bất kỳ ai cũng có thể tái tạo ANDREA-12M từ các artifacts này. Đường dẫn CUDA tùy chỉnh là một phần lý do: không có phiên bản framework bị đóng băng, không có bất ngờ về dependency sau năm năm nữa.
Tín hiệu & Checkpoints
Vòng lặp huấn luyện CUDA phản hồi với hai tín hiệu POSIX:
- SIGTERM: ghi checkpoint ngay lập tức, sau đó thoát. Sử dụng khi dừng huấn luyện một cách sạch sẽ.
- SIGUSR1: ghi checkpoint ngay lập tức, tiếp tục huấn luyện. Được sử dụng trong pivot polish ở v3 để capture trạng thái mà không gián đoạn run.
Định dạng checkpoint: [int32 step][int32 n_params][n_params × float32 weights][n_params × float32 m][n_params × float32 v]. Bộ đếm step, số lượng weight, sau đó là weights theo sau bởi Adam moments. Tiếp tục bit-exactly. Proxy lưu trữ .samples.json & .state.json riêng biệt trên polish; .loss.json không bao giờ được lưu trữ (nó tích lũy toàn bộ lịch sử huấn luyện).