Pipeline Parallelism — Mỗi GPU 1 nhóm layer, pipeline mini-batch

Khi kiến trúc AI vượt ngưỡng 70B parameters, bạn đối mặt với một sự thật phũ phàng: ngay cả GPU A100 80GB cũng không đủ để chứa toàn bộ model, optimizer states và gradients cùng lúc. Data Parallelism chỉ đơn thuần nhân bản model sang nhiều GPU — nhưng nếu model đã không vừa một card thì nhân bản cũng vô nghĩa. Pipeline Parallelism là giải pháp chia nhỏ chính kiến trúc model theo chiều dọc — mỗi GPU nắm giữ một nhóm layer nhất định, và dữ liệu chảy qua như một dây chuyền lắp ráp ô tô, cho phép huấn luyện những mô hình nghìn tỷ tham số trên cụm GPU thông thường.

Vấn đề

Với một model 175B parameters ở độ chính xác FP16, bạn cần 350GB chỉ để lưu weights. Thêm optimizer states (Adam cần 2 bản sao của weights) và gradients, con số bùng nổ lên gần 1.5TB — gấp 18 lần VRAM của một chiếc A100. Giải pháp đầu tiên mọi người nghĩ đến là Data Parallelism kết hợp với ZeRO để shard optimizer states, nhưng giới hạn vẫn tồn tại: bạn vẫn phải giữ toàn bộ weights và activations trên mỗi GPU trong quá trình forward/backward.

Tensor Parallelism — cắt nhỏ từng weight matrix ra nhiều GPU — giải quyết được vấn đề này, nhưng lại tạo ra bottleneck mới: all-reduce communication sau mỗi layer, đòi hỏi bandwidth cực cao giữa các GPU (thường giới hạn trong 1 node duy nhất). Khi cần train model trên hàng trăm GPU cross-node, tensor parallelism trở nên kém hiệu quả do latency mạng.

Cần một cách tiếp cận "chia theo chiều dọc" thay vì "chia theo chiều ngang" — tách model thành các đoạn pipeline, mỗi GPU xử lý một giai đoạn tuần tự, giống như các trạm trong dây chuyền sản xuất.

Ý tưởng cốt lõi

Chia layer, không chia matrix. Thay vì cắt nhỏ từng lớp Linear thành nhiều mảnh trên các GPU khác nhau, Pipeline Parallelism chia toàn bộ model theo chiều sâu: GPU 0 giữ layers 1-6, GPU 1 giữ layers 7-12, GPU 2 giữ layers 13-18, v.v. Một sample đi qua GPU 0 → GPU 1 → GPU 2 như một chiếc xe đi qua các trạm rửa xe.

The "Aha" Moment — Thời gian là không gian: Điểm mấu chốt là nhận ra rằng trong khi layer phụ thuộc tuần tự (không thể tính layer 2 trước khi có output layer 1), thì các sample lại độc lập hoàn toàn. Khi GPU 1 đang xử lý layer 2 của Sample A, GPU 0 hoàn toàn có thể bắt đầu layer 1 của Sample B. Bằng cách chia mini-batch thành các micro-batch nhỏ hơn (ví dụ: batch 64 → 8 micro-batch của 8 sample), chúng ta biến chiều thời gian thành chiều không gian — tận dụng song song ở cấp độ sample thay vì cấp độ layer.

1F1B Schedule — Bí quyết giảm memory: Vấn đề còn lại là activation memory. Nếu chờ tất cả forward xong mới làm backward (GPipe style), các GPU đầu pipeline phải giữ activations cho toàn bộ micro-batch, gây nổ bộ nhớ. Giải pháp là One-Forward-One-Backward (1F1B): ngay sau khi forward một micro-batch xong, ta lập tức chạy backward cho nó, giải phóng activation trước khi nhận micro-batch tiếp theo. Điều này tạo ra trạng thái ổn định (steady state) nơi tất cả GPU luân phiên forward và backward liên tục, không có GPU nào nghỉ sau khi pipeline được lấp đầy.

Bubble overhead — Cái giá phải trả: Khi khởi động, GPU cuối pipeline phải chờ micro-batch đầu tiên đi qua tất cả các stage trước — gọi là "bubble". Nhưng với công thức đơn giản (N-1)/M (N = số stage, M = số micro-batch), nếu bạn chia đủ nhỏ (M >> N), bubble chỉ chiếm vài phần trăm thời gian. Ví dụ: 4 GPU (N=4) với 16 micro-batch (M=16) chỉ lãng phí 3/16 = 18.75% thời gian ở startup, sau đó đạt 100% utilization.

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

Pipeline Parallelism tồn tại được nhờ hai đặc tính toán học quan trọng:

1. Communication giảm thiểu: Khác với Tensor Parallelism phải all-reduce sau mỗi layer, Pipeline Parallelism chỉ cần truyền activation tensors ở ranh giới giữa các stage. Với 24 layer chia 4 stage, chỉ có 3 điểm cần gửi/nhận dữ liệu giữa GPU, thay vì 24 lần all-reduce. Điều này làm cho pipeline parallelism lý tưởng cho multi-node training qua InfiniBand/Ethernet, nơi bandwidth thấp hơn NVLink trong node.

2. Tính chất tuyến tính của gradient: Gradient accumulation trong pipeline hoạt động tương tự Gradient Accumulation — mỗi stage tích lũy gradient locally cho weights của mình, chỉ cần đồng bộ hóa khi kết thúc mini-batch. Không cần all-reduce khổng lồ qua toàn bộ model mỗi bước.

Pseudocode minh họa 1F1B:

# Mỗi GPU giữ một stage (group of layers)
for micro_batch in split(batch, M):
    # Forward: gửi output sang GPU kế tiếp
    output = forward(micro_batch)
    send_to_next_gpu(output)
    
    # Nhận gradient từ GPU sau và backward ngay lập tức
    grad = recv_from_next_gpu()
    backward(grad)
    # Activation được giải phóng ở đây, không chờ tất cả forward xong

Ý nghĩa thực tế

Hiệu suất thực tế: GPipe (Google, 2019) đạt 3.14× speedup trên AmoebaNet với 4 GPU. Với số lượng micro-batch đủ lớn, bubble overhead tiệm cận 0 — bạn gần như linear scaling về throughput khi thêm GPU vào pipeline.

Ai đang dùng:

• DeepSpeed và Megatron-LM (NVIDIA) — chuẩn công nghiệp cho việc train GPT-3, PaLM, và các model nghìn tỷ tham số
• vLLM — áp dụng pipeline parallelism cho inference, chia model lớn qua nhiều GPU khi serving
• SageMaker Model Parallelism — AWS managed service tự động partition model cho khách hàng

Giới hạn cần nhớ:

• Pipeline Bubble: Vẫn cần M >= 4N để ẩn latency, tốn thêm memory lưu nhiều activation (nếu không dùng checkpointing)
• Load Imbalance: Các layer đầu/cuối (embedding và softmax) thường nặng hơn middle layers (FFN đơn giản), dễ gây "nghẽn cổ chai" nếu chia đều số layer
• Chỉ phù hợp batch size lớn: Với batch size nhỏ, không đủ micro-batch để lấp đầy pipeline, GPU sẽ nghỉ nhiều
• Không giảm memory cho activations: Mỗi stage vẫn phải lưu activations cho tất cả micro-batch đang "trong ống", thường kết hợp với activation checkpointing để trade compute lấy memory

Trong thực tế, pipeline parallelism hiếm khi đứng riêng lẻ. Nó thường được kết hợp với data parallelism (3D parallelism: Data + Pipeline + Tensor) để tận dụng cả cụm hàng trăm GPU — pipeline chia model theo chiều sâu, tensor parallelism chia layer theo chiều rộng, và data parallelism nhân bản toàn bộ cụm này cho nhiều batch song song.

Đào sâu hơn

Paper gốc:

• GPipe: Easy Scaling with Micro-Batch Pipeline Parallelism (2019) — Google, giới thiệu micro-batch pipelining cơ bản
• PipeDream: Fast and Efficient Pipeline Parallel DNN Training (2019) — Microsoft, 1F1B schedule và weight stashing
• Pipeline Parallelism with Controllable Memory (2024) — Phân tích các scheduling strategies như repeating blocks

Cùng cụm (Training Efficiency):

• Tensor Parallelism — Cắt weight matrix theo chiều ngang, thường dùng kết hợp với Pipeline Parallelism trong 3D parallelism
• ZeRO — Tối ưu memory cho Data Parallelism, bổ sung cho pipeline khi cần scale ra nhiều node
• QLoRA — Khi không có cụm GPU lớn, fine-tune 65B model trên 1 GPU consumer bằng quantization + LoRA

Đọc tiếp:

• Training at Scale — Foundations về distributed training, data parallelism và gradient synchronization (Level 0)
• Alignment Basics — Sau khi train xong model lớn, làm sao align với human preferences (Level 0)
• Alignment Frontier — Các phương pháp alignment mới nhất cho model đã được train với pipeline parallelism (Level 2)