Tensor Parallelism — Cắt weight matrix cho nhiều GPU

Khi OpenAI huấn luyện GPT-3 175B năm 2020, không có chiếc GPU nào trên thị trường chứa nổi toàn bộ model (cần ~350GB chỉ cho weights). Trong khi Pipeline Parallelism tạo ra "bubble" chết thời gian và Data Parallelism bị kẹt cứng bởi giới hạn VRAM, Tensor Parallelism mở ra lối thoát thứ ba: cắt nhỏ từng ma trận weight bên trong một layer, phân tán qua nhiều GPU như thể đang xếp lego. Đây là kỹ thuật cốt lõi giúp train model trăm tỷ tham số trở nên khả thi, biến bottleneck từ "bộ nhớ" thành "băng thông liên lạc".

Vấn đề

Trước năm 2019, có hai cách để chia nhỏ model khổng lồ trên nhiều GPU:

Pipeline Parallelism (PP) chia theo chiều dọc: GPU 0 giữ layers 1-10, GPU 1 giữ layers 11-20, v.v. Vấn đề là tính tuần tự—GPU 1 phải chờ GPU 0 hoàn thành mới bắt đầu được, tạo ra "pipeline bubble" (thời gian GPU ngồi chơi) lên đến 30-50% tổng thời gian.

Data Parallelism (DP) nhân bản toàn bộ model lên mỗi GPU. Nhưng với GPT-3 175B, ngay cả FP16 cũng cần 350GB cho weights, chưa kể gradients và optimizer states (Adam cần thêm 2x nữa). Không có GPU nào đủ VRAM. Bạn không thể nhân bản cái gì đó không vừa một nơi để bắt đầu.

Cần một cách chia nhỏ bên trong từng layer—nơi chiếm đa số parameters (attention và FFN)—để mỗi GPU chỉ giữ một mảnh nhỏ của weight matrix thay vì toàn bộ layer.

Ý tưởng cốt lõi

Tensor Parallelism (TP) dựa trên insight cơ bản: phép nhân ma trận chỉ là phép cộng và nhân phân tán được. Nếu bạn chia ma trận trọng số $W$ đúng cách, mỗi GPU có thể tính một phần kết quả, sau đó gộp lại bằng các phép toán tập thể (collective operations).

Column-wise vs Row-wise: Hai cách cắt bánh

Hãy tưởng tượng một layer Linear $Y = XW$ với $X$ là input (batch × input_dim) và $W$ là weight (input_dim × output_dim).

Column-wise sharding (chia theo cột): Cắt $W$ thành $[W_1 | W_2]$ theo chiều output. GPU 1 tính $XW_1$, GPU 2 tính $XW_2$. Kết quả là hai nửa output, chỉ cần ghép lại (concatenate) — không cần tính toán thêm, chỉ cần giao tiếp nhỏ ở cuối (All-Gather).

Row-wise sharding (chia theo hàng): Cắt $W$ thành $\begin{bmatrix} W_1 \\ W_2 \end{bmatrix}$ và $X$ tương ứng thành $[X_1 | X_2]$. GPU 1 tính $X_1W_1$, GPU 2 tính $X_2W_2$. Bây giờ hai kết quả là partial sums—phải cộng lại bằng All-Reduce để có output cuối.

Sự khác biệt then chốt: Column-wise cần ghép (concat), Row-wise cần cộng (sum). Row-wise tốn kém hơn chút vì All-Reduce phức tạp hơn All-Gather, nhưng cả hai đều khả thi.

Chiêu Megatron: Đổi chiều cắt giữa các layer

Insight làm nên tên tuổi của Megatron-LM (NVIDIA) là nhận ra transformer block có hai phép nhân ma trận liên tiếp:

1. Q/K/V projection (input → attention heads)
2. Output projection (attention heads → output)

Nếu cả hai đều dùng cùng một kiểu sharding, bạn sẽ cần All-Reduce ở giữa chúng để "sửa lại" dữ liệu. Nhưng Megatron làm một việc tinh tế:

• Q/K/V dùng Column-wise: Output bị chia thành các chunks theo chiều head (mỗi GPU giữ một số heads).
• Output projection dùng Row-wise: Input cần là chunks theo chiều head để nhân với $W_\text{out}$.

Kết quả: Output của bước 1 đã được sharded đúng theo format mà bước 2 cần! Không cần All-Reduce ở giữa. Toàn bộ transformer block chỉ cần một lần All-Reduce duy nhất (ở cuối block) thay vì hai lần. Đây là "free lunch" của communication optimization.

Hình ảnh trực quan

Hãy tưởng tượng hai người dịch thuật:

• Column-wise: Hai người cùng đọc một câu tiếng Việt, người A dịch nửa đầu câu, người B dịch nửa sau. Họ không cần nói chuyện trong lúc dịch, chỉ cần ghép bản dịch lại ở cuối.
• Row-wise: Người A đọc nửa đầu câu gốc, người B đọc nửa sau. Mỗi người dịch phần của mình thành một bản thảo tạm, rồi họ cộng gộp ý nghĩa lại để ra câu hoàn chỉnh.

Megatron để hai người làm việc liên tục: Người A dịch xong nửa câu đầu, đưa cho người B (đã có nửa câu sau) để họ hoàn thiện toàn bộ câu mà không cần họp bàn ở giữa.

Tại sao nó hoạt động

Mathematically, điều này hoạt động vì tính chất phân phối của đại số tuyến tính:

$X \cdot [W_1 | W_2] = [XW_1 | XW_2] \quad \text{(Column-wise)}$ $[X_1 | X_2] \cdot \begin{bmatrix} W_1 \\ W_2 \end{bmatrix} = X_1W_1 + X_2W_2 \quad \text{(Row-wise)}$

Quan trọng hơn là communication pattern. All-Reduce trên NVLink/InfiniBand rất nhanh khi trong cùng một node (thường 8 GPUs), với bandwidth tới vài trăm GB/s. Trong khi đó, memory bandwidth của một GPU chỉ ~2 TB/s nhưng bị giới hạn bởi kích thước vật lý của VRAM.

Bằng cách giảm memory footprint xuống $1/TP\_degree$ (với TP=8, mỗi GPU chỉ giữ 12.5% weight), chúng ta trade memory (giải quyết được) lấy communication (chấp nhận được trong single-node). Đây là lý do Tensor Parallelism thường chỉ dùng trong node (intra-node), còn cross-node thường dùng Pipeline Parallelism.

Ý nghĩa thực tế

Impact thực tế:

• Scale: Megatron-LM báo cáo 76% scaling efficiency trên 512 GPUs khi train GPT-3—nghĩa là thêm GPU gần như tuyến tính tăng tốc độ.
• Kiến trúc: Mọi model lớn hiện đại (GPT-4, Llama 3 405B, PaLM) đều dùng TP=8 hoặc TP=16 trong node để fit weights vào H100/A100.

So sánh với các kỹ thuật khác:

Giới hạn:

• Băng thông: TP đòi hỏi interconnect cực nhanh (NVLink). Trên Ethernet hoặc cross-node, All-Reduce trở thành bottleneck.
• Ràng buộc kích thước: TP degree phải chia hết hidden dimension (ví dụ: 12288 chiều hidden chia được cho TP=8, nhưng không chia cho TP=7).
• Overhead nhỏ batch: Với batch size nhỏ, thời gian communication có thể chiếm 50-70% thời gian, làm TP chậm hơn cả Data Parallelism đơn thuần.

Misconception cần phá: TP không làm giảm tổng lượng tính toán (FLOPs), chỉ giải phóng memory. Bạn vẫn phải chạy qua toàn bộ weights, chỉ là giờ chạy song song trên nhiều GPU.

Đào sâu hơn

Paper gốc:

• "Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism" (Shoeybi et al., 2019) — giới thiệu column/row-wise sharding và chiến lược xen kẽ cho transformer.
• "DeepSpeed Inference: Enabling Efficient Inference of Transformer Models at Unprecedented Scale" (2022) — áp dụng TP cho inference autoregressive với KV-cache sharding.

Cùng cụm (Training Efficiency):

• PEFT — Family of methods: LoRA, prefix tuning, adapters — khi model quá lớn để full fine-tune, dùng PEFT thay vì TP.
• QLoRA — nếu chỉ cần inference hoặc fine-tune nhỏ trên consumer GPU, quantize 4-bit hiệu quả hơn TP.
• ZeRO — cách khác để giải quyết memory optimizer states mà không cần sharding weights như TP.
• Pipeline Parallelism — kết hợp với TP tạo thành 3D parallelism (Data + Pipeline + Tensor).
• Gradient Accumulation — kỹ thuật đơn giản hơn để giả lập batch lớn khi không có nhiều GPU.

Đọc tiếp:

• alignment-basics — sau khi train xong model khổng lồ bằng TP, làm sao để align nó với human preference.
• training-at-scale — các khái niệm cơ bản về distributed training trước khi đi sâu vào TP.
• alignment-frontier — các phương pháp alignment mới nhất cho model đã được train bằng TP ở quy mô lớn.