Song song tensor cho suy luận

Tensor parallelism chia một lớp của mô hình thành các phần để nhiều bộ tăng tốc phần cứng xử lý chúng đồng thời. Điều này cho phép bạn chạy các mô hình vượt quá dung lượng bộ nhớ của một GPU duy nhất và đạt được thông lượng cao hơn. Bạn sẽ cần liên lạc nhanh trong node vì các GPU trao đổi kết quả trung gian ở mỗi lớp.

Một model hỗ trợ tensor parallelism nếu config của nó định nghĩa base_model_tp_plan. Kiểm tra một model đã được nạp bằng thuộc tính supports_tp_plan.

from transformers import AutoModelForCausalLM

model = AutoModelForCausalLM.from_pretrained("meta-llama/Meta-Llama-3-8B-Instruct")
print(model.supports_tp_plan)

Hướng dẫn này bao gồm cách bật song song tensor trong Transformers và các chiến lược phân chia có sẵn.

Chia nhỏ một mô hình

Cấu hình số lượng thiết bị song song tensor bằng tp_size trong DistributedConfig.

• Đặt DistributedConfig(tp_size=N) để sử dụng kế hoạch được định nghĩa trước của mô hình.
• Định nghĩa tp_plan thủ công và truyền nó vào DistributedConfig cùng với tp_size.

Bạn cũng có thể đặt tp_plan="auto" để yêu cầu kế hoạch được định nghĩa trước một cách rõ ràng. Khi tp_size bị bỏ qua và một tp_plan được đặt, tp_size sẽ được suy ra từ WORLD_SIZE chia cho các kích thước song song khác. Việc truyền tp_plan trực tiếp vào from_pretrained() đã bị ngừng hỗ trợ và sẽ bị loại bỏ trong v5.18.

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, DistributedConfig

# model_id = "meta-llama/Llama-4-Scout-17B-16E-Instruct" # better to visualize all the possible strategies
distributed_config = DistributedConfig(tp_size=4)
model = AutoModelForCausalLM.from_pretrained(
    "meta-llama/Meta-Llama-3-8B-Instruct",
    dtype=torch.bfloat16,
    distributed_config=distributed_config,
)
print(model.tp_plan)

tokenizer = AutoTokenizer.from_pretrained("meta-llama/Meta-Llama-3-8B-Instruct")
prompt = "Can I help"
inputs = tokenizer(prompt, return_tensors="pt").input_ids.to(model.device)

# distributed run
outputs = model(inputs)

torchrun --nproc-per-node 4 demo.py

from transformers import AutoModelForCausalLM, DistributedConfig

tp_plan = {
    "model.layers.*.self_attn.q_proj": "colwise",
    "model.layers.*.self_attn.k_proj": "colwise",
    "model.layers.*.self_attn.v_proj": "colwise",
    "model.layers.*.self_attn.o_proj": "rowwise",
    ...
}

distributed_config = DistributedConfig(tp_size=4, tp_plan=tp_plan)
model = AutoModelForCausalLM.from_pretrained(
    "meta-llama/Meta-Llama-3-8B-Instruct",
    dtype="auto",
    distributed_config=distributed_config,
)
print(model.tp_plan)

Chiến lược phân vùng

ALL_PARALLEL_STYLES, một ParallelInterface được định nghĩa trong distributed/parallel_utils.py, ánh xạ mỗi tên chiến lược mà bạn có thể sử dụng trong một kế hoạch đến một instance chiến lược đã được cấu hình. Bạn không tương tác trực tiếp với nó để shard một model, nhưng đây là danh sách chính thức các tên khả dụng. distributed/tensor_parallel.py đăng ký các chiến lược tensor parallel vào đó:

ALL_PARALLEL_STYLES.register("embedding_rowwise", EmbeddingParallel(shard_dim=0))
ALL_PARALLEL_STYLES.register("embedding_colwise", EmbeddingParallel(shard_dim=1))
ALL_PARALLEL_STYLES.register("colwise_gather_output", ColwiseParallel(input_layouts=Replicate(), output_layouts=Replicate()))
ALL_PARALLEL_STYLES.register("colwise_rep", ColwiseParallel(input_layouts=Replicate(), output_layouts=Replicate()))
ALL_PARALLEL_STYLES.register("colwise", ColwiseParallel(input_layouts=Replicate(), output_layouts=Shard(-1)))
ALL_PARALLEL_STYLES.register("rowwise", RowwiseParallel(input_layouts=Shard(-1), output_layouts=Replicate()))
ALL_PARALLEL_STYLES.register("rowwise_split_input", RowwiseParallel(input_layouts=Replicate(), output_layouts=Replicate()))
ALL_PARALLEL_STYLES.register("rowwise_rep", RowwiseParallel(input_layouts=Replicate(), output_layouts=Replicate()))
ALL_PARALLEL_STYLES.register("packed_colwise", PackedColwiseParallel())
ALL_PARALLEL_STYLES.register("packed_rowwise", PackedRowwiseParallel())
ALL_PARALLEL_STYLES.register("sequence_parallel", SequenceParallel(use_local_output=True))
ALL_PARALLEL_STYLES.register("replicated_with_grad_allreduce", ReplicatedWithGradAllReduce())
ALL_PARALLEL_STYLES.register("mla_kv_a_proj", MlaKvAProjParallel())
ALL_PARALLEL_STYLES.register("all_reduce", AllReduceParallel())
ALL_PARALLEL_STYLES.register("moe_tp_experts", MoeExpertsParallel())
ALL_PARALLEL_STYLES.register("moe_identity_expert", MoeIdentityParallel())
ALL_PARALLEL_STYLES.register("local_params", LocalParamsParallel())

Mọi chiến lược song song tensor đều là lớp con của TensorParallelLayer trong tệp đó. Các chiến lược song song chuyên gia, chẳng hạn như ep_dispatch_experts, nằm trong distributed/expert_parallel.py, kế thừa từ ExpertParallelLayer, có cùng các hook, và được đăng ký vào cùng một ALL_PARALLEL_STYLES. Hai chiến lược bạn sẽ dùng thường xuyên nhất, ColwiseParallel và RowwiseParallel, nhận các bố trí input_layouts và output_layouts, đây là cách một lớp bao phủ nhiều tên kế hoạch. colwise giữ đầu ra của nó được chia mảnh trên chiều cuối để một lớp rowwise tiếp theo tiêu thụ, trong khi colwise_gather_output thu thập tất cả (all-gather) nó trở lại thành một tensor đầy đủ.

Bảng dưới đây mô tả từng chiến lược.

Chiến lược đóng gói

Weight packing kết hợp nhiều lớp tuyến tính thành một lớp lớn hơn duy nhất. Các chiến lược PackedColwiseParallel và PackedRowwiseParallel chia nhỏ (shard) các trọng lượng đã đóng gói một cách chính xác. Các chiến lược ColwiseParallel hoặc RowwiseParallel cơ bản sẽ chia nhỏ các trọng lượng đã đóng gói một cách không chính xác.

Ví dụ dưới đây đóng gói up_proj và gate_proj thành một module gate_up_proj duy nhất và yêu cầu chiến lược packed_rowwise để chia nhỏ gate_up_proj.

class Llama4TextExperts(nn.Module):
    ...
    self.gate_up_proj = nn.Parameter(torch.zeros(self.num_experts, self.hidden_size, 2 * self.expert_dim))

Sử dụng phép nhân ma trận theo lô trong bước forward để tính toán đầu ra của module gate_up_proj.

def forward(self, hidden_states):
    ...
    gate_up = torch.bmm(hidden_states, self.gate_up_proj) # Compute the output of the gate_up_proj module
    gate, up = gate_up.chunk(2, dim=-1) # Split the output into gate and up

Một Shard thông thường chia chiều đó thành các khối liên tục, do đó một rank sẽ nhận được phần cuối của gate và phần đầu của up thay vì một lát cắt của từng phần. Các chiến lược đóng gói sử dụng _StridedShard với split_factor (mặc định là 2) để xen kẽ việc chia, đảm bảo mỗi rank nhận được một lát cắt tương ứng của cả hai nửa để chunk vẫn khớp nhau sau khi sharding.

Các chiến lược phân vùng tùy chỉnh

Kế thừa từ TensorParallelLayer trong distributed/tensor_parallel.py để tạo chiến lược phân vùng tùy chỉnh. Chỉ ghi đè các hook mà chiến lược của bạn cần, vì mỗi hook đều có mặc định không thực hiện gì (no-op).

Ví dụ dưới đây đi qua một phiên bản rút gọn của ColwiseParallel.

1. Kế thừa từ TensorParallelLayer và lưu trữ các vị trí (placements) mà chiến lược làm việc. Lớp cơ sở không định nghĩa __init__, vì vậy không có gì để gọi super().

   Đã sao chép

   class ColwiseParallel(TensorParallelLayer):
       def __init__(self, *, input_layouts=None, output_layouts=None, use_local_output: bool = True):
           self.input_layouts = input_layouts or Replicate()
           self.output_layouts = output_layouts if output_layouts is not None else Shard(-1)
           self.use_local_output = use_local_output

2. Thực hiện shard_param để bọc một tham số thành một placeholder DTensor. Nó chạy trên các meta tensor, do đó distribute_tensor chỉ xây dựng metadata và không di chuyển dữ liệu. Truyền src_data_rank=None vì hiện tại chưa có tensor đầy đủ để phân tán.

   Đã sao chép

   def shard_param(self, module, param, mesh):
       meta = module._parameters.get(param)
       if meta is None:
           return
       # Output features live on dim 0 for a 2D weight and dim -1 for a 1D bias
       placement = Shard(1) if isinstance(module, torch.nn.Embedding) else Shard(meta.ndim - 2)
       module._parameters[param] = torch.nn.Parameter(
           distribute_tensor(meta, mesh, [placement], src_data_rank=None),
           requires_grad=meta.requires_grad,
       )

3. Thực hiện các phép biến đổi đầu vào và đầu ra. install_forward gọi chúng bao quanh hàm forward gốc của module, vì vậy chúng chỉ cần di chuyển tensor giữa các bố cục. Việc phân vùng theo cột kỳ vọng đầu vào được nhân bản và tạo ra đầu ra được chia mảnh trên chiều cuối.

   Đã sao chép

   def transform_inputs_pre_forward(self, module, args, kwargs, mesh):
       x = args[0]
       if not isinstance(x, DTensor):
           x = DTensor.from_local(x, mesh, [self.input_layouts], run_check=False)
       if x.placements != (Replicate(),):
           x = x.redistribute(placements=[Replicate()])
       return (x,) + args[1:], kwargs
   
   def transform_output_post_forward(self, module, output, mesh):
       if not isinstance(output, DTensor):
           output = DTensor.from_local(output, mesh, [Shard(-1)], run_check=False)
       if output.placements != (self.output_layouts,):
           output = output.redistribute(placements=[self.output_layouts])
       return output.to_local() if self.use_local_output else output

   ColwiseParallel được cung cấp sẵn bổ sung các đường dẫn nhanh trên cơ sở này, bỏ qua vòng lặp DTensor cho suy luận nn.Linear thông thường và cho các module lượng tử hóa cần tensor cục bộ. Hãy đọc mã nguồn trước khi sao chép nếu chiến lược của bạn cần xử lý tương tự.

4. Đăng ký chiến lược để tp_plan có thể gọi tên nó. Việc đăng ký nhận một instance của chiến lược, do đó các placements bạn truyền vào __init__ sẽ được lưu trữ cùng với tên. Lưu ý () bên dưới.

   Đã sao chép

   import torch
   
   from transformers import AutoModelForCausalLM, DistributedConfig
   from transformers.distributed.parallel_utils import ALL_PARALLEL_STYLES
   
   ALL_PARALLEL_STYLES.register("colwise_custom", ColwiseParallel())
   tp_plan = {
       "model.layers.*.self_attn.q_proj": "colwise_custom",
       ...
   }
   distributed_config = DistributedConfig(tp_size=4, tp_plan=tp_plan)
   model = AutoModelForCausalLM.from_pretrained(
       "meta-llama/Meta-Llama-3-8B-Instruct",
       dtype=torch.bfloat16,
       distributed_config=distributed_config,
   )

Benchmarks

Song song tensor tăng tốc độ suy luận đáng kể, đặc biệt là với kích thước lô lớn hoặc các chuỗi dài.

Biểu đồ này cho thấy tốc độ tăng tốc dự kiến cho một lần truyền tiến đơn lẻ trên Llama với độ dài chuỗi là 512.

Thiết kế triển khai

Transformers triển khai tensor parallelism theo cách độc lập với framework. Nó dựa vào DeviceMesh và DTensor từ torch.distributed để cung cấp một giao diện đơn giản, có thể mở rộng.

DeviceMesh

DeviceMesh tạo ra một lưới thiết bị nhiều chiều giao tiếp với nhau. Các chiến lược song song hóa khác nhau đòi hỏi các mẫu giao tiếp khác nhau. Hãy tạo một DeviceMesh với nhiều lưới con để xử lý các mẫu này.

import torch
from torch.distributed.device_mesh import init_device_mesh

# Create a 1D mesh of 4 accelerators
device_type = torch.accelerator.current_accelerator().type
device_mesh = init_device_mesh(device_type, (4,), mesh_dim_names=["tp"])

Hầu hết các chiến lược song song hóa torch.distributed được áp dụng cho chính lưới (mesh) hoặc các lưới con của nó. Lưới tự động xử lý các mẫu giao tiếp.

DTensor

DTensor (Distributed Tensor) xử lý logic phân tán trên các thao tác tensor thông thường. Trọng số mô hình dưới tensor parallelism được lưu trữ dưới dạng DTensor, điều này cho phép một chiến lược mô tả giao tiếp như một sự thay đổi layout thay vì một tập thể hợp (collective) tường minh.

Thuộc tính placements cho PyTorch biết cách một tensor được bố trí trên các thiết bị trong một DeviceMesh. Nó chấp nhận các giá trị sau:

• Shard(dimension) chia một DTensor trên một chiều cho trước theo DeviceMesh mà nó được xây dựng. Việc phân vùng theo cột chia chiều output-feature của trọng số và chiều duy nhất của bias.

  Đã sao chép

  weight = DTensor.from_local(weight, device_mesh["tp"], placements=[Shard(0)]) # Shard the output features
  bias = DTensor.from_local(bias, device_mesh["tp"], placements=[Shard(-1)]) # Shard the ONLY dimension

  Phân vùng theo hàng (row-wise partitioning) chia nhỏ chiều input-feature thay vào đó và nhân bản bias, vì bias được cộng một lần sau khi thực hiện phép giảm (reduction) chứ không phải trên mỗi rank.

  Đã sao chép

  weight = DTensor.from_local(weight, device_mesh["tp"], placements=[Shard(-1)]) # Shard the input features
  bias = DTensor.from_local(bias, device_mesh["tp"], placements=[Replicate()]) # Replicate bias across all GPUs

• Replicate() nhân bản một DTensor trên DeviceMesh, tạo ra một bản sao đầy đủ của tensor trên mỗi thiết bị.

• Partial() đánh dấu một tensor đang chờ giảm (reduction). Đầu ra forward của một lớp row-wise là Partial(), và việc phân phối lại nó sang Replicate() chính là thao tác phát lệnh all-reduce.

Tài nguyên

• Phần Ultra-Scale Playbook về song song tensor cung cấp thêm chi tiết.

• Kiểm tra hướng dẫn expert parallelism nếu bạn đang sử dụng mô hình mixture-of-experts (MoE). Các mô hình này hỗ trợ tensor parallelism và expert parallelism.

• Đọc bài đăng trên blog Tensor Parallelism (TP) in Transformers: 5 Minutes to Understand để có cái nhìn tổng quan nhanh về tensor parallelism và tìm hiểu cách các cấu hình song song cột và hàng khác nhau.

• Xem hướng dẫn huấn luyện Tensor parallelism để tìm hiểu cách sử dụng nó trong môi trường huấn luyện.