Tăng tốc ND-Parallel- Hướng dẫn đào tạo đa GPU hiệu quả

August 11, 2025
13 min read

Accelerate ND-Parallel: Hướng dẫn huấn luyện Multi-GPU hiệu quả

Việc huấn luyện các mô hình lớn trên nhiều GPU có thể gặp nhiều thách thức do sự phức tạp của các chiến lược song song khác nhau. Trong Accelerate, cùng với Axolotl, chúng tôi đã tích hợp một cách nhanh chóng và dễ dàng để sử dụng bất kỳ sự kết hợp nào của các chiến lược song song trong script huấn luyện của bạn!

Đây là cách thêm nó vào script huấn luyện của bạn:

from transformers import AutoModelForCausalLM
from accelerate import Accelerator
from accelerate.parallelism_config import ParallelismConfig
from accelerate.utils import FullyShardedDataParallelPlugin

# cấu hình các parallelisms mong muốn của bạn tại đây - cấu hình cụ thể này yêu cầu ít nhất 2 nodes với 8 GPUs mỗi node.
# đặt bất kỳ bậc parallelism nào thành 1 sẽ tắt nó, ví dụ: dp_replicate_size=1 tắt DP.
pc = ParallelismConfig(
    dp_shard_size=2,  # Bậc Fully Sharded Data Parallel
    dp_replicate_size=2,  # Bậc Data Parallel
    cp_size=2,  # Bậc Context Parallel
    tp_size=2,  # Bậc Tensor Parallel
)

fsdp_plugin = FullyShardedDataParallelPlugin(
    fsdp_version=2,
    auto_wrap_policy="transformer_based_wrap",
    transformer_cls_names_to_wrap=["LlamaDecoderLayer"],
    state_dict_type="SHARDED_STATE_DICT",
)

accelerator = Accelerator(
    parallelism_config=pc,
    fsdp_plugin=fsdp_plugin
)

model = AutoModelForCausalLM.from_pretrained(
    "NousResearch/Hermes-3-Llama-3.1-8B",
    device_mesh=accelerator.torch_device_mesh
)

model = accelerator.prepare(model)

Chúng tôi cũng đã bao gồm một script huấn luyện toàn diện hơn từ đầu đến cuối trong Accelerate repo, script này trình bày cách thiết lập dataloader, optimizer và training loop của bạn, đồng thời cách lưu mô hình của bạn sau khi huấn luyện.

Để hợp lý hóa hơn nữa các mô hình tinh chỉnh ở quy mô lớn và kết hợp các chiến lược song song với nhiều kỹ thuật tinh chỉnh, chúng tôi cũng đã tích hợp kỹ thuật này vào Axolotl. Để giúp bạn bắt đầu ngay lập tức, chúng tôi đã thử nghiệm một số cấu hình ví dụ mà bạn có thể sửa đổi cho phù hợp với nhu cầu của mình - hãy thử một cấu hình với:

# lưu ý: điều này yêu cầu kích thước world tối thiểu là 16
axolotl train examples/distributed-parallel/llama-3_1-8b-hsdp-tp.yaml

Bạn cũng có thể xem tài liệu Axolotl ND-Parallelism để biết thêm chi tiết - việc thêm các kỹ thuật song song ND vào các cấu hình hiện có của bạn rất đơn giản, chỉ cần thêm một hoặc nhiều trường sau vào tệp cấu hình Axolotl của bạn:

# Bậc Fully Sharded Data Parallel (lưu ý: cũng yêu cầu trường fsdp_config)
# xem https://docs.axolotl.ai/docs/multi-gpu.html#sec-fsdp để biết thêm chi tiết
dp_shard_size: 2
# Bậc Data Parallel
dp_replicate_size: 2
# Bậc Context Parallel
context_parallel_size: 2
# Bậc Tensor Parallel
tensor_parallel_size: 2

Chúng tôi đã giúp bạn dễ dàng định cấu hình các bậc của các chiến lược song song khác nhau và cách chúng được kết hợp thông qua lớp ParallelismConfig trong Accelerate hoặc thông qua các trường cấu hình trong Axolotl, nhưng làm cách nào để chúng ta biết cấu hình nào sẽ hoạt động tốt nhất cho trường hợp sử dụng của chúng ta? Khi chúng ta mở rộng quy mô để huấn luyện các mô hình với hàng chục hoặc thậm chí hàng trăm tỷ tham số, thách thức chính đến từ việc hiểu các chiến lược song song khác nhau và cách chúng tương tác để giảm thiểu chi phí truyền thông trên các thiết bị. Trong bài đăng này, chúng ta sẽ xem xét cách các chiến lược song song khác nhau hoạt động, đồng thời khi nào và cách bạn có thể muốn kết hợp chúng.

Data Parallelism

Diagram for Data Parallel Hình: Distributed Data Parallel sao chép toàn bộ mô hình trên mỗi thiết bị và chia đều dữ liệu thành các lô con cho mỗi thiết bị. (Nguồn: Martynas Šubonis).

Data parallelism (DP) là kỹ thuật phổ biến nhất để huấn luyện các mô hình trên nhiều GPU và bao gồm việc sao chép mô hình, gradient và trạng thái optimizer trên mỗi thiết bị, đồng thời phân phối đều các lô dữ liệu giữa các GPU và đồng bộ hóa gradient trên các thiết bị trước khi cập nhật các tham số. Điều này có thể làm tăng đáng kể thông lượng so với huấn luyện trên một thiết bị, nhưng yêu cầu mô hình của bạn có thể phù hợp trên một thiết bị duy nhất.

Chúng ta có thể kiểm soát số lượng bản sao của mô hình bằng tham số dp_replicate_size trong ParallelismConfig của Accelerate hoặc trường cấu hình trong Axolotl. Điều đáng chú ý là DP là một chiến lược song song cấp cao nhất, có nghĩa là nếu chúng ta sử dụng dp_replicate_size=2 và chúng ta kết hợp nó với các chiến lược song song khác, sẽ có 2 bản sao của mô hình, mỗi bản cũng bị ảnh hưởng bởi các chiến lược song song khác. Ví dụ: nếu chúng ta sử dụng dp_replicate_size=2 và tp_size=2, chúng ta sẽ có 2 bản sao của mô hình, mỗi bản có 2 shards tensor parallel.

Chúng tôi sử dụng thuật ngữ shard để mô tả dữ liệu trên một thiết bị duy nhất là một phân vùng của một phần dữ liệu lớn hơn.

Fully Sharded Data Parallelism

Diagram for Fully Sharded Data Parallel Hình: Fully Sharded Data Parallel chia đều từng tham số của mô hình trên mỗi thiết bị và, giống như DDP, chia đều dữ liệu thành các lô con cho mỗi thiết bị. Để hoàn thành một lượt forward và backward, FSDP phải thu thập trọng số của từng tham số trước lượt forward/backward để mỗi thiết bị có được một bản sao đầy đủ của tham số. (Nguồn: Martynas Šubonis).

Điều gì sẽ xảy ra nếu mô hình của chúng ta quá lớn để phù hợp trên một thiết bị duy nhất? Fully sharded data parallel (FSDP) giải quyết vấn đề này bằng cách sharding (phân phối đều) trọng số, gradient và trạng thái optimizer của mô hình trên các GPU (điều này được lấy cảm hứng từ ZeRO-3 của DeepSpeed), trong khi mỗi thiết bị vẫn nhận được phần lô dữ liệu đầy đủ của nó. Như bạn có thể nhận thấy từ sơ đồ trên, thay vì yêu cầu một bản sao đầy đủ của toàn bộ mô hình trên mỗi thiết bị, chúng ta chỉ thu thập trọng số cho một lớp duy nhất tại một thời điểm trước lượt forward, sau đó trọng số có thể được sharded lại.

Theo cách này, chúng ta đánh đổi việc sử dụng bộ nhớ để lấy chi phí truyền thông của việc thu thập các tham số sharded trước mỗi lượt forward và backward, đồng thời giảm phân tán gradient cục bộ. Chúng ta có thể kiểm soát sự đánh đổi này trong FSDP bằng cách điều chỉnh độ chi tiết mà tại đó các tham số được thu thập. Ở một thái cực, chúng ta có thể thu thập và re-shard mọi lớp của mô hình, điều này sẽ dẫn đến việc sử dụng bộ nhớ đỉnh thấp nhất, nhưng phát sinh chi phí truyền thông cao nhất. Trong thực tế, một cách tiếp cận phổ biến là thu thập trọng số cho toàn bộ khối giải mã transformer tại một thời điểm.

Mặc dù chúng ta có thể thực hiện thêm các giao dịch đánh đổi giữa bộ nhớ và tính toán và chuyển các tham số và gradient mô hình sang CPU để huấn luyện các mô hình lớn hơn, nhưng điều này có thể chậm một cách cấm đoán. Thay vào đó, hãy xem xét cách chúng ta có thể sử dụng hiệu quả nhiều thiết bị hơn nữa để huấn luyện các mô hình lớn hơn trong khi vẫn duy trì thông lượng dữ liệu cao.

Chúng tôi sử dụng thuật ngữ node để chỉ một máy duy nhất lưu trữ nhiều GPU (tối đa là 8), với các kênh truyền thông trong node nhanh chóng sử dụng, ví dụ: NVLink giữa các GPU. Khi sử dụng nhiều node để huấn luyện, chúng ta dựa vào các kênh truyền thông giữa các node tương đối chậm hơn giữa các máy sử dụng, ví dụ: Infiniband. Chúng tôi cũng đề cập đến tổng số thiết bị trong nhóm quy trình là kích thước world - ví dụ: một node duy nhất có 8 GPU thể hiện kích thước world là 8 và 4 node sẽ thể hiện kích thước world là 32.

Khi sử dụng FSDP trên nhiều node, chúng ta coi toàn bộ tập hợp các thiết bị trên các node như thể chúng ta đang huấn luyện trên một node duy nhất. Ví dụ: với 4 node chứa 8 GPU mỗi node, chúng ta thực hiện sharding trên 32 thiết bị và thực hiện các hoạt động all-reduce và reduce-scatter tập thể của chúng ta bằng cách sử dụng cả các backend truyền thông giữa và trong node. Bằng cách này, chỉ riêng FSDP đã có thể mở rộng quy mô lên một số lượng GPU đáng kể với kích thước lô toàn cầu lớn để tăng thông lượng dữ liệu. Tuy nhiên, đến một điểm mà một số thách thức phát sinh có thể yêu cầu kết hợp FSDP với các kỹ thuật song song khác. Chúng ta thường cố gắng tránh thực hiện FSDP trên nhiều hơn một node đầy đủ, vì chi phí truyền thông có thể trở nên quá cao, chúng ta sẽ nói về cách giải quyết vấn đề này trong phần về Hybrid Sharded Data Parallelism.

Bạn có thể sử dụng tham số dp_shard_size trong ParallelismConfig của Accelerate cùng với FullyShardedDataParallelPlugin đã chuẩn bị hoặc đặt trường cấu hình dp_shard_size trong Axolotl để đặt mức độ FSDP được áp dụng cho mô hình của bạn.

Tensor Parallelism

Diagram for Tensor Parallel Hình: Tensor Parallelism chia các lớp tuyến tính lớn trên các thiết bị, thường sử dụng sharding theo cột cho lớp đầu tiên và sharding theo hàng cho lớp tiếp theo. Phương pháp này chỉ yêu cầu một hoạt động truyền thông AllReduce duy nhất để kết hợp các đầu ra sharded, giảm thiểu chi phí truyền thông trong khi phân phối cả bộ nhớ và tính toán trên các thiết bị trong một node.

Tensor Parallel (TP) là một loại kỹ thuật model parallelism, trong đó các shards của mô hình thường trú trên các thiết bị riêng biệt và trái ngược với các kỹ thuật data parallel, mỗi thiết bị nhận được một lô dữ liệu giống hệt nhau. TP hoạt động bằng cách phân phối tính toán của các lớp tuyến tính trên các thiết bị, do đó mỗi thiết bị chỉ tính toán một phần của phép nhân ma trận. Kỹ thuật này hoạt động tốt nhất khi có các lớp tuyến tính lớn, chẳng hạn như các lớp feed-forward trong các mô hình transformer, có thể được chia trên các thiết bị. Chúng ta cũng có thể sử dụng TP trên từng phép chiếu query, key, value và output trong các lớp attention mà hầu như không tốn thêm chi phí truyền thông.

Để đạt được hiệu suất tốt nhất, các tham số của các lớp liên tiếp có thể được phân phối theo một cách cụ thể, giảm thiểu truyền thông cần thiết. Khi làm việc với các cặp lớp tuyến tính, chúng ta có thể chia lớp đầu tiên theo cột và lớp tiếp theo theo hàng, cho phép chúng ta tính toán đầu ra chỉ với một hoạt động all-reduce duy nhất để kết hợp các đầu ra sharded.

Không giống như hành vi sharding động của FSDP, TP tạo ra các phân vùng bộ nhớ tĩnh, dẫn đến giảm sử dụng bộ nhớ không đổi tỷ lệ với kích thước nhóm TP. Điều này trở nên quan trọng đối với các mô hình lớn, trong đó ngay cả một lớp giải mã duy nhất cũng quá lớn để phù hợp với bộ nhớ trong quá trình all-gather FSDP (hãy nhớ lại rằng thực hành phổ biến trong FSDP là thu thập trọng số của toàn bộ lớp giải mã tại một thời điểm). Tuy nhiên, không giống như FSDP mở rộng quy mô tương đối tuyến tính trên các node (cho đến một điểm - ~512 GPU trên một cụm đồng nhất, ít hơn đáng kể trên các kết nối băng thông thấp hơn), TP chỉ hiệu quả trong phạm vi một node duy nhất. TP yêu cầu đồng bộ hóa kích hoạt thường xuyên giữa các thiết bị trong quá trình tính toán, vì mỗi thiết bị chỉ tính toán một phần của đầu ra, yêu cầu các đầu ra từ các thiết bị khác phải được truyền đạt trước khi tiếp tục lượt forward. Do đó, nếu chúng ta muốn sử dụng TP trong thiết lập nhiều node, chúng ta phải xem xét việc kết hợp TP với các kỹ thuật song song khác, trong khi chỉ giữ TP trong một node duy nhất. Do chi phí truyền thông lớn, TP không được khuyến nghị cho các GPU được liên kết PCIe.

Trong Accelerate, kích thước TP được định cấu hình thông qua tp_size trong ParallelismConfig, trong khi trong Axolotl, bạn có thể sử dụng trường cấu hình tensor_parallel_size.

Context Parallelism

Gần đây, khả năng suy luận trong LLM đã dẫn đến độ dài chuỗi tăng vọt khi các mô hình sử dụng ngày càng nhiều token để giải quyết các tác vụ phức tạp. Để đạt được hành vi này thông qua tinh chỉnh, chúng ta cần một cách để huấn luyện các mô hình trên độ dài chuỗi rất lớn - đôi khi có thể đạt tới một triệu token!

Vì hoạt động attention trong transformers mở rộng quy mô theo bậc hai với độ dài ngữ cảnh, điều này trở nên không thể trên một GPU duy nhất. Ví dụ: khi tinh chỉnh một mô hình tương đối nhỏ như Mistral-7B (sử dụng 32 đầu attention), nếu chúng ta sử dụng độ dài chuỗi là 128k, một ma trận attention duy nhất sẽ sử dụng 128k * 128k * 2 bytes * num_heads=32 = ~32GB * 32 = ~1TB bộ nhớ kích hoạt! Mặc dù ví dụ này không thực tế khi sử dụng các triển khai attention được tối ưu hóa như FlashAttention, nhưng nó giúp minh họa sự tăng trưởng về yêu cầu bộ nhớ từ việc tăng độ dài ngữ cảnh.

Với context parallelism (CP), chúng ta có thể shard các đầu vào theo chiều chuỗi, dẫn đến mỗi thiết bị chỉ xử lý một phần của ngữ cảnh đầy đủ và tính toán một phần nhỏ hơn của ma trận attention đầy đủ, lớn đến mức cấm đoán. Để xem cách thức hoạt động của điều này, hãy nhớ lại rằng tính toán attention được mô tả bằng phương trình: Attention(Q,K,V)=softmax(QKT)V \text{Attention}(Q, K, V) = \text{softmax}(QK^T)V Attention(Q,K,V)=softmax(QKT)V Trong đó Q Q Q, K K K, và V V V là các ma trận query, key và value tương ứng. Mỗi vector query (hàng hoặc nhúng đầu vào) của Q Q Q phải tính toán điểm attention với mọi vector key của K K K trong toàn bộ chuỗi để áp dụng chính xác phép chuẩn hóa softmax. Các điểm attention này sau đó được trọng số với tất cả các vector value trong V V V.

Chi tiết quan trọng ở đây nằm ở chỗ mỗi hàng trong Q Q Q có thể tính toán điểm attention của nó một cách độc lập với nhau, nhưng mỗi vector query vẫn yêu cầu các ma trận K K K và V V V đầy đủ. Nói cách khác, với một đầu vào có độ dài chuỗi $n$, chúng ta có thể mở rộng phương trình attention ở trên của mình như sau:

Attention(Q,K,V)1=softmax(Q1KT)VAttention(Q,K,V)2=softmax(Q2KT)V⋮Attention(Q,K,V)n=softmax(QnKT)V \begin{align} \text{Attention}(Q, K, V)_1 &= \text{softmax}(Q_1 K^T) V \ \text{Attention}(Q, K, V)_2 &= \text{softmax}(Q_2 K^T) V \ &\vdots \ \text{Attention}(Q, K, V)_n &= \text{softmax}(Q_n K^T) V \end{align} Attention(Q,K,V)1Attention(Q,K,V)2Attention(Q,K,V)n<span class=“pstrut” style=“hei

AI Today - SkyAI