Song song tensor

Tensor parallelism chia nhỏ các ma trận trọng số của mô hình trên các thiết bị. Mỗi thiết bị giữ một lát cắt theo cột ("colwise") hoặc theo hàng ("rowwise") của mỗi lớp và tính toán một kết quả cục bộ. Một AllReduce hoặc AllGather tại ranh giới lớp sẽ kết hợp các kết quả cục bộ thành đầu ra đầy đủ. Khác với context parallelism, nó giảm bộ nhớ weight trên mỗi thiết bị, hữu ích khi mô hình không vừa trên một thiết bị. Xem Distributed inference để so sánh với các chiến lược khác.

Tải một mô hình phân mảnh

Truyền một TensorParallelConfig vào tham số parallel_config của from_pretrained() của mô hình. tp_degree là số lượng thiết bị để chia nhỏ và phải chia hết cho số lượng đầu chú ý của mô hình. Lớp của mô hình phải định nghĩa _tp_plan, ánh xạ tên module sang kiểu chia nhỏ. Để kiểm tra xem một mô hình có hỗ trợ song song tensor hay không, hãy tìm thuộc tính này trên lớp của nó, chẳng hạn như Flux2Transformer2DModel.

Để thêm một kế hoạch vào một mô hình khác, hãy xem Thêm song song tensor vào một mô hình.

Việc nạp theo cách này sẽ chia checkpoint thành các mảnh khi đọc. Mỗi rank chỉ đọc phần riêng của mình trong mỗi trọng số và đặt trực tiếp lên GPU của nó, do đó không rank nào giữ toàn bộ trọng số ở kích thước đầy đủ. Bộ nhớ trên mỗi rank giảm khi tp_degree tăng.

import os
import torch
from torch import distributed as dist
from diffusers import DiffusionPipeline, Flux2Transformer2DModel, TensorParallelConfig

def main():
    dist.init_process_group(backend="nccl")
    rank, world_size = dist.get_rank(), dist.get_world_size()
    # torchrun sets LOCAL_RANK to this process's GPU index on its node
    local_rank = int(os.environ["LOCAL_RANK"])
    device = torch.device(f"cuda:{local_rank}")
    torch.cuda.set_device(device)

    # Each rank reads only its own shard of every planned weight, straight onto its own GPU.
    transformer = Flux2Transformer2DModel.from_pretrained(
        "black-forest-labs/FLUX.2-dev",
        subfolder="transformer",
        dtype=torch.bfloat16,
        parallel_config=TensorParallelConfig(tp_degree=world_size),
    )

    pipeline = DiffusionPipeline.from_pretrained(
        "black-forest-labs/FLUX.2-dev", transformer=transformer, dtype=torch.bfloat16
    )
    # The transformer is already on its device; move the remaining components individually. Do not call
    # `pipeline.to(device)` — that would move every rank's shards onto the same device.
    pipeline.text_encoder.to(device)
    pipeline.vae.to(device)

    generator = torch.Generator().manual_seed(42)
    image = pipeline(prompt="a cat holding a sign that says hello", generator=generator).images[0]
    if rank == 0:
        image.save("output.png")
    dist.destroy_process_group()

if __name__ == "__main__":
    main()

Lưu script dưới tên tensor_parallel_flux2.py và chạy nó bằng torchrun.

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

Script đặt tp_degree bằng world_size, do đó --nproc-per-node 4 sẽ chia transformer thành 4 phần trên 4 GPU.

Chỉ có transformer được chia nhỏ. Mỗi rank giữ một bản sao đầy đủ của text encoder, do đó mỗi GPU cần đủ bộ nhớ cho text encoder cùng với phần transformer của nó.

Hiệu năng tải

Việc nạp với parallel_config nhanh hơn và sử dụng ít bộ nhớ CPU hơn trên mỗi rank so với việc nạp toàn bộ mô hình rồi gọi enable_parallelism(). Bảng dưới đây cho thấy kết quả với black-forest-labs/FLUX.2-dev (chỉ transformer, 32B tham số, bf16) trên 4x A10G (23GB).

Hạn chế

Tensor parallelism không hoạt động với một số tùy chọn và tính năng tải, và bạn không thể lưu một mô hình đã được chia mảnh.

• Việc tải với parallel_config tensor-parallel không hoạt động cùng với device_map, quantization_config, low_cpu_mem_usage=False, use_flashpack=True, hoặc trọng lượng không phải safetensors. Mỗi trường hợp sẽ gây ra lỗi thay vì âm thầm tải toàn bộ checkpoint.
• Việc kết hợp tensor parallelism với quantization, offloading hoặc các adapter LoRA cũng sẽ gây ra lỗi, bất kể mô hình được sharded theo cách nào.
• Để chia nhỏ (shard) một model đã nằm trong bộ nhớ, hãy gọi enable_parallelism() với cùng một config. Hàm này sẽ tải toàn bộ model trước rồi mới chia nhỏ, do đó mỗi rank cần bộ nhớ cho toàn bộ checkpoint. Việc tải bằng parallel_config giúp tránh chi phí đó.
• Việc lưu một mô hình song song tensor sẽ gây ra lỗi trong save_pretrained(). Hãy lưu mô hình trước khi chia nhỏ (sharding) nó.

Thêm song song tensor vào một mô hình

Phần này dành cho những người đóng góp muốn thêm hỗ trợ song song tensor cho một mô hình Diffusers.

Tensor parallelism chỉ hoạt động trên các model có định nghĩa _tp_plan, một thuộc tính lớp ánh xạ tên module sang kiểu sharding. Việc viết một _tp_plan chủ yếu là ghép mỗi projection mở rộng chiều hidden với projection thu gọn nó lại.

Mỗi key có thể chứa tối đa một *, và phần trước * phải trỏ đến một nn.ModuleList, do đó một entry sẽ bao phủ mọi block. Key không có * sẽ áp dụng cho chính model. Các đường dẫn là tương đối so với model.

Colwise và rowwise

Luôn ghép "colwise" và "rowwise" theo đúng thứ tự đó. Một phép chiếu "colwise" để đầu ra của nó ở dạng sharded, phép chiếu "rowwise" tiếp theo tiêu thụ shard đó trực tiếp, và một AllReduce duy nhất tại ranh giới khối tái tạo kết quả. Việc chia shard cặp này theo bất kỳ cách nào khác sẽ buộc phải thực hiện gather ở giữa và truyền thông nhiều hơn rất nhiều.

Trong cơ chế chú ý, các phép chiếu Q, K và V theo cột cấp cho mỗi rank một tập con các đầu, do đó tp_degree phải chia hết cho số lượng đầu. Các bản sao của luồng encoder (add_q_proj, to_add_out, ff_context) tuân theo cùng một mẫu như các phần tử tương ứng của luồng ảnh.

Các phép chiếu đã hợp nhất

Khi một Linear đóng gói nhiều tensor logic dọc theo chiều đang được chia nhỏ, việc cắt đơn giản theo "colwise"/"rowwise" xuyên qua phần ghép nối sẽ làm lệch các mảnh. Hãy dùng PackedColwiseParallel/PackedRowwiseParallel thay vào đó, vì chúng chia nhỏ từng khối đã đóng gói một cách độc lập.

# in src/diffusers/models/transformers/your_model.py
from ...hooks.tensor_parallel import PackedColwiseParallel, PackedRowwiseParallel

Truyền blocks, một danh sách các số nguyên tỷ lệ có tổng chia hết cho chiều đã đóng gói. Ví dụ, sử dụng [1, 1] cho phép và phép chiếu lên của SwiGLU có hai nửa bằng nhau, hoặc [1, 1, 1, 3, 3] cho phép chiếu Q, K, V, gate và up đã hợp nhất với mlp_ratio=3.

"transformer_blocks.*.ff.linear_in": PackedColwiseParallel([1, 1]),

Khi kích thước khối chỉ được biết từ cấu hình, hãy bỏ qua blocks và thay vào đó lưu các kích thước tuyệt đối trên Linear trong __init__ dưới dạng _tp_packed_col_blocks hoặc _tp_packed_row_blocks.

# in the attention module's __init__
self.to_out._tp_packed_row_blocks = [self.inner_dim, self.mlp_hidden_dim]

# in the model's _tp_plan
"single_transformer_blocks.*.attn.to_out": PackedRowwiseParallel(),

Kế hoạch trong transformer_flux2.py sử dụng cả hai dạng. Cặp SwiGLU double-stream truyền blocks trực tiếp. Dự toán QKV và MLP single-stream đã hợp nhất lưu kích thước của chúng trên Linear trong Flux2ParallelSelfAttention.__init__, vì chúng phụ thuộc vào mlp_ratio và mlp_mult_factor.

Những gì cần loại bỏ

Các module không được đưa vào kế hoạch sẽ được nhân bản trên mọi rank. Hãy bỏ qua các lớp chuẩn hóa, embedding patch và text, cùng với norm_out và proj_out cuối cùng. Chúng nhỏ, nên việc chia nhỏ chúng chỉ tiết kiệm được ít bộ nhớ mà lại làm tăng giao tiếp. Điều chế AdaLN (img_mod và txt_mod) thường đủ nhỏ để cũng có thể bỏ qua. Khi nó lớn, như adaln_proj.linear của MiniMax-H3, hãy chia nhỏ nó bằng "rowwise_split_input".

Ràng buộc và xác minh

• tp_degree phải chia hết cho config.num_attention_heads.
• Mỗi khối được đóng gói phải riêng lẻ chia hết cho tp_degree, không chỉ tổng của chúng.

from_pretrained() và enable_parallelism() kiểm tra cả hai trước khi tải hoặc chia nhỏ bất kỳ trọng số nào.

Kiểm tra một kế hoạch mới về mặt số học thay vì bằng mắt: tạo với một hạt (seed) cố định trên một thiết bị duy nhất, sau đó tạo lại dưới tensor parallelism, và so sánh các đầu ra. Một "colwise"/"rowwise" đặt sai vị trí thường vẫn chạy và tạo ra một hình ảnh trông hợp lý nhưng sai.

Bắt đầu từ một kế hoạch có sẵn cho một kiến trúc tương tự.

• transformer_qwenimage.py không có các projection đã được hợp nhất, do đó mọi mục đều là "colwise" hoặc "rowwise" thông thường.
• transformer_flux.py thêm một phép chiếu hàng dạng đóng gói.
• transformer_flux2.py sử dụng cả hai phong cách đóng gói.