Suy luận phân tán

Suy luận phân tán chia nhỏ khối lượng công việc trên nhiều GPU. Đây là một kỹ thuật hữu ích để nạp các mô hình lớn hơn vào bộ nhớ và có thể xử lý nhiều prompt để đạt thông lượng cao hơn.

Hướng dẫn này sẽ cho bạn thấy cách sử dụng Accelerate và PyTorch Distributed để suy luận phân tán.

Accelerate

Accelerate là một thư viện được thiết kế để đơn giản hóa suy luận và huấn luyện trên nhiều bộ tăng tốc bằng cách xử lý việc thiết lập, cho phép người dùng tập trung vào mã PyTorch của họ.

Cài đặt Accelerate bằng lệnh sau.

uv pip install accelerate

Khởi tạo lớp accelerate.PartialState trong một tệp Python để tạo môi trường phân tán. Lớp accelerate.PartialState quản lý việc quản lý tiến trình, kiểm soát và phân phối thiết bị, cũng như phối hợp giữa các tiến trình.

Di chuyển DiffusionPipeline sang accelerate.PartialState.device để gán một GPU cho mỗi tiến trình.

import torch
from accelerate import PartialState
from diffusers import DiffusionPipeline

pipeline = DiffusionPipeline.from_pretrained(
    "Qwen/Qwen-Image", dtype=torch.float16
)
distributed_state = PartialState()
pipeline.to(distributed_state.device)

Sử dụng tiện ích split_between_processes như một trình quản lý ngữ cảnh để tự động phân phối các prompt giữa các quy trình.

with distributed_state.split_between_processes(["a dog", "a cat"]) as prompt:
    result = pipeline(prompt).images[0]
    result.save(f"result_{distributed_state.process_index}.png")

Gọi accelerate launch để chạy script và sử dụng đối số --num_processes để đặt số lượng GPU cần dùng.

accelerate launch run_distributed.py --num_processes=2

Tham khảo script tối giản này để chạy suy luận trên nhiều GPU. Để tìm hiểu thêm, hãy xem hướng dẫn Distributed Inference with 🤗 Accelerate.

Phân tán PyTorch

PyTorch DistributedDataParallel cho phép data parallelism, tức là nhân bản cùng một mô hình trên mỗi thiết bị, để xử lý các lô dữ liệu khác nhau song song.

Nhập torch.distributed và torch.multiprocessing vào một tệp Python để thiết lập nhóm tiến trình phân tán và khởi tạo các tiến trình cho suy luận trên mỗi GPU.

import torch
import torch.distributed as dist
import torch.multiprocessing as mp

from diffusers import DiffusionPipeline

pipeline = DiffusionPipeline.from_pretrained(
    "Qwen/Qwen-Image", dtype=torch.float16,
)

Tạo một hàm cho suy luận với init_process_group. Phương thức này tạo ra một môi trường phân tán với loại backend, rank của tiến trình hiện tại và world_size hoặc số lượng tiến trình tham gia (ví dụ, 2 GPU sẽ là world_size=2).

Chuyển pipeline sang rank và sử dụng get_rank để gán một GPU cho mỗi tiến trình. Mỗi tiến trình xử lý một prompt khác nhau.

def run_inference(rank, world_size):
    dist.init_process_group("nccl", rank=rank, world_size=world_size)

    pipeline.to(rank)

    if torch.distributed.get_rank() == 0:
        prompt = "a dog"
    elif torch.distributed.get_rank() == 1:
        prompt = "a cat"

    image = sd(prompt).images[0]
    image.save(f"./{'_'.join(prompt)}.png")

Sử dụng mp.spawn để tạo ra số lượng tiến trình được định nghĩa trong world_size.

def main():
    world_size = 2
    mp.spawn(run_inference, args=(world_size,), nprocs=world_size, join=True)


if __name__ == "__main__":
    main()

Gọi torchrun để chạy script suy luận và sử dụng đối số --nproc_per_node để đặt số lượng GPU cần dùng.

torchrun --nproc_per_node=2 run_distributed.py

device_map

Tham số device_map cho phép suy luận phân tán bằng cách tự động đặt các thành phần của mô hình lên các GPU riêng biệt. Điều này đặc biệt hữu ích khi một mô hình không vừa trên một GPU duy nhất. Bạn có thể sử dụng device_map để tải và gỡ các thành phần mô hình cần thiết ở một giai đoạn nhất định, như minh họa trong ví dụ bên dưới (giả sử có hai GPU khả dụng).

Đặt device_map="balanced" để phân bổ đều các text encoder trên tất cả các GPU khả dụng. Bạn có thể sử dụng đối số max_memory để phân bổ một lượng bộ nhớ tối đa cho mỗi text encoder. Không nạp bất kỳ thành phần pipeline nào khác để tránh sử dụng bộ nhớ.

from diffusers import FluxPipeline
import torch

prompt = """
cinematic film still of a cat sipping a margarita in a pool in Palm Springs, California
highly detailed, high budget hollywood movie, cinemascope, moody, epic, gorgeous, film grain
"""

pipeline = FluxPipeline.from_pretrained(
    "black-forest-labs/FLUX.1-dev",
    transformer=None,
    vae=None,
    device_map="balanced",
    max_memory={0: "16GB", 1: "16GB"},
    dtype=torch.bfloat16
)
with torch.no_grad():
    print("Encoding prompts.")
    prompt_embeds, pooled_prompt_embeds, text_ids = pipeline.encode_prompt(
        prompt=prompt, prompt_2=None, max_sequence_length=512
    )

Sau khi tính toán xong các vector nhúng văn bản, hãy xóa chúng khỏi GPU để tạo không gian cho bộ biến đổi khuếch tán.

import gc 

def flush():
    gc.collect()
    torch.cuda.empty_cache()
    torch.cuda.reset_max_memory_allocated()
    torch.cuda.reset_peak_memory_stats()

del pipeline.text_encoder
del pipeline.text_encoder_2
del pipeline.tokenizer
del pipeline.tokenizer_2
del pipeline

flush()

Đặt device_map="auto" để tự động phân phối mô hình trên hai GPU. Chiến lược này đặt mô hình lên thiết bị nhanh nhất trước, sau đó mới đặt lên thiết bị chậm hơn như CPU hoặc ổ cứng nếu cần. Sự đánh đổi khi lưu trữ các tham số mô hình trên các thiết bị chậm hơn là độ trễ suy luận sẽ chậm hơn.

from diffusers import AutoModel
import torch 

transformer = AutoModel.from_pretrained(
    "black-forest-labs/FLUX.1-dev", 
    subfolder="transformer",
    device_map="auto",
    dtype=torch.bfloat16
)

Chạy pipeline.hf_device_map để xem các model khác nhau được phân bổ trên các thiết bị như thế nào. Điều này hữu ích để theo dõi vị trí đặt model trên thiết bị. Bạn cũng có thể gọi hf_device_map trên model transformer để xem nó được phân bổ như thế nào.

Thêm mô hình transformer vào pipeline và đặt output_type="latent" để tạo ra các latent.

pipeline = FluxPipeline.from_pretrained(
    "black-forest-labs/FLUX.1-dev",
    text_encoder=None,
    text_encoder_2=None,
    tokenizer=None,
    tokenizer_2=None,
    vae=None,
    transformer=transformer,
    dtype=torch.bfloat16
)

print("Running denoising.")
height, width = 768, 1360
latents = pipeline(
    prompt_embeds=prompt_embeds,
    pooled_prompt_embeds=pooled_prompt_embeds,
    num_inference_steps=50,
    guidance_scale=3.5,
    height=height,
    width=width,
    output_type="latent",
).images

Gỡ pipeline và transformer khỏi bộ nhớ và nạp VAE để giải mã các latents. VAE thường đủ nhỏ để có thể được nạp trên một thiết bị duy nhất.

import torch
from diffusers import AutoencoderKL
from diffusers.image_processor import VaeImageProcessor

vae = AutoencoderKL.from_pretrained(ckpt_id, subfolder="vae", dtype=torch.bfloat16).to("cuda")  # or "mps", "xpu", "cpu"
vae_scale_factor = 2 ** (len(vae.config.block_out_channels) - 1)
image_processor = VaeImageProcessor(vae_scale_factor=vae_scale_factor)

with torch.no_grad():
    print("Running decoding.")
    latents = FluxPipeline._unpack_latents(latents, height, width, vae_scale_factor)
    latents = (latents / vae.config.scaling_factor) + vae.config.shift_factor

    image = vae.decode(latents, return_dict=False)[0]
    image = image_processor.postprocess(image, output_type="pil")
    image[0].save("split_transformer.png")

Bằng cách tải và gỡ tải có chọn lọc các mô hình cần thiết ở từng giai đoạn cụ thể và chia nhỏ các mô hình lớn nhất trên nhiều GPU, bạn có thể chạy suy luận với các mô hình lớn trên GPU dành cho người tiêu dùng.

Song song ngữ cảnh

Song song ngữ cảnh chia các chuỗi đầu vào trên nhiều GPU để giảm mức sử dụng bộ nhớ. Mỗi GPU xử lý phần riêng của chuỗi.

Sử dụng set_attention_backend() để chuyển sang backend attention được tối ưu hơn. Tham khảo bảng này để xem danh sách đầy đủ các backend khả dụng.

Hầu hết các backend attention đều tương thích với context parallelism. Hãy mở một issue nếu một backend không tương thích.

Ring Attention

Các biểu diễn Key (K) và value (V) giao tiếp giữa các thiết bị bằng cách sử dụng Ring Attention. Điều này đảm bảo mỗi phần chia nhìn thấy K/V của mọi token khác. Mỗi GPU tính toán attention cho K/V cục bộ của nó và truyền nó sang GPU tiếp theo trong vòng. Không có GPU nào giữ toàn bộ chuỗi, điều này làm giảm độ trễ truyền thông.

Truyền một ContextParallelConfig vào tham số parallel_config của mô hình transformer. Cấu hình này hỗ trợ tham số ring_degree để xác định số lượng thiết bị cần sử dụng cho Ring Attention.

import torch
from torch import distributed as dist
from diffusers import DiffusionPipeline, ContextParallelConfig

def setup_distributed():
    if not dist.is_initialized():
        dist.init_process_group(backend="nccl")
    rank = dist.get_rank()
    device = torch.device(f"cuda:{rank}")
    torch.cuda.set_device(device)
    return device

def main():
    device = setup_distributed()
    world_size = dist.get_world_size()

    pipeline = DiffusionPipeline.from_pretrained(
        "black-forest-labs/FLUX.1-dev", dtype=torch.bfloat16
    ).to(device)
    pipeline.transformer.set_attention_backend("_native_cudnn")

    cp_config = ContextParallelConfig(ring_degree=world_size)
    pipeline.transformer.enable_parallelism(config=cp_config)

    prompt = """
    cinematic film still of a cat sipping a margarita in a pool in Palm Springs, California
    highly detailed, high budget hollywood movie, cinemascope, moody, epic, gorgeous, film grain
    """

    # Must specify generator so all ranks start with same latents (or pass your own)
    generator = torch.Generator().manual_seed(42)
    image = pipeline(
        prompt,
        guidance_scale=3.5,
        num_inference_steps=50,
        generator=generator,
    ).images[0]

    if dist.get_rank() == 0:
        image.save(f"output.png")

    if dist.is_initialized():
        dist.destroy_process_group()


if __name__ == "__main__":
    main()

Mảnh mã ở trên cần được chạy bằng trình khởi chạy phân tán, chẳng hạn như torchrun, tương thích với PyTorch. --nproc-per-node được đặt bằng số lượng GPU khả dụng.

torchrun --nproc-per-node 2 above_script.py

Ulysses Attention

Ulysses Attention chia một chuỗi (sequence) trên các GPU và thực hiện giao tiếp all-to-all (mọi thiết bị gửi/nhận dữ liệu tới mọi thiết bị khác). Mỗi GPU cuối cùng sẽ có tất cả các token cho một tập con các đầu chú ý (attention heads). Mỗi GPU tính toán chú ý cục bộ trên tất cả các token cho đầu của nó, sau đó thực hiện một lần all-to-all khác để nhóm lại kết quả theo token cho lớp tiếp theo.

ContextParallelConfig hỗ trợ Ulysses Attention thông qua đối số ulysses_degree. Đối số này xác định số lượng thiết bị được sử dụng cho Ulysses Attention.

Truyền ContextParallelConfig vào enable_parallelism().

# Depending on the number of GPUs available.
pipeline.transformer.enable_parallelism(config=ContextParallelConfig(ulysses_degree=2))

Unified Attention

Unified Sequence Parallelism kết hợp Ring Attention và Ulysses Attention thành một phương pháp duy nhất để xử lý chuỗi dài hiệu quả. Phương pháp này trước tiên áp dụng giao tiếp all-to-all của Ulysses để phân phối lại các head và token chuỗi, sau đó sử dụng Ring Attention để xử lý dữ liệu đã được phân phối lại, và cuối cùng đảo ngược all-to-all để khôi phục bố cục ban đầu.

Cách tiếp cận lai này tận dụng những điểm mạnh của cả hai phương pháp:

• Ulysses Attention song song hóa hiệu quả trên các đầu chú ý
• Ring Attention xử lý các chuỗi rất dài với chi phí bộ nhớ tối thiểu
• Cùng nhau, chúng cho phép song song hóa 2D trên cả chiều đầu (heads) và chiều chuỗi (sequence).

ContextParallelConfig hỗ trợ Unified Attention bằng cách chỉ định cả ulysses_degree và ring_degree. Tổng số thiết bị được sử dụng là ulysses_degree * ring_degree, được sắp xếp thành một lưới 2D trong đó các nhóm Ulysses và Ring là trực giao (không chồng lấn). Truyền ContextParallelConfig với cả ulysses_degree và ring_degree được đặt lớn hơn 1 vào enable_parallelism().

pipeline.transformer.enable_parallelism(config=ContextParallelConfig(ulysses_degree=2, ring_degree=2))

Unified Attention được sử dụng khi có đủ thiết bị để sắp xếp thành lưới 2D (ít nhất 4 thiết bị).

Chúng tôi đã chạy một bài benchmark với Ulysess, Ring và Unified Attention bằng script này trên một node gồm 4 GPU H100. Kết quả được tóm tắt như sau:

Từ bảng trên, có thể thấy rõ rằng Ulysses cung cấp thông lượng tốt hơn, nhưng số lượng thiết bị mà nó có thể sử dụng vẫn bị giới hạn ở số lượng đầu chú ý, một hạn chế được giải quyết bởi unified attention.

Ulysses Anything Attention

Cơ chế Ulysses Attention mặc định yêu cầu độ dài chuỗi của hidden states phải chia hết cho số lượng thiết bị. Điều này đặt ra những hạn chế đáng kể đối với việc ứng dụng thực tế của Ulysses Attention. Ulysses Anything Attention là một biến thể của Ulysses Attention, hỗ trợ độ dài chuỗi tùy ý và số lượng attention heads tùy ý, từ đó nâng cao tính linh hoạt của Ulysses Attention trong sử dụng thực tế.

ContextParallelConfig hỗ trợ Ulysses Anything Attention bằng cách chỉ định cả ulysses_degree và ulysses_anything. Vui lòng lưu ý rằng Ulysses Anything Attention hiện chưa được hỗ trợ bởi Unified Attention. Truyền ContextParallelConfig với cả ulysses_degree được đặt lớn hơn 1 và ulysses_anything=True vào enable_parallelism().

pipeline.transformer.enable_parallelism(config=ContextParallelConfig(ulysses_degree=2, ulysses_anything=True))

Để tránh nhiều lần đồng bộ CUDA bắt buộc do các thao tác chuyển H2D và D2H, vui lòng thêm backend gloo trong init_process_group. Điều này sẽ giảm đáng kể độ trễ truyền thông.

Chúng tôi đã chạy benchmark cho FLUX.1-dev với Ulysses, Ring, Unified Attention và Ulysses Anything Attention bằng script này trên một node gồm 4 GPU L20. Kết quả được tóm tắt như sau:

Từ bảng trên, có thể thấy rằng Ulysses Anything Attention cung cấp khả năng tương thích tốt hơn với các độ dài chuỗi tùy ý trong khi vẫn duy trì hiệu suất tương đương với Ulysses Attention tiêu chuẩn.

Ring Anything Attention

Ring Attention mặc định yêu cầu độ dài chuỗi của hidden states phải chia hết đều cho ring degree. Ring Anything Attention là một biến thể của Ring Attention hỗ trợ độ dài chuỗi tùy ý (không chia hết đều). Nó đệm KV cục bộ của mỗi rank lên độ dài chuỗi tối đa toàn cục, all-gather bộ đệm KV đã đệm, và cắt lại về độ dài thực của từng rank trước khi chạy attention.

Ring Anything Attention không được hỗ trợ bởi Unified Attention. Đặt ring_degree > 1 và ring_anything=True để bật Ring Anything Attention.

pipeline.transformer.enable_parallelism(config=ContextParallelConfig(ring_degree=2, ring_anything=True))

Thêm backend gloo vào init_process_group để tránh nhiều lần đồng bộ hóa CUDA bắt buộc từ các thao tác chuyển dữ liệu H2D và D2H.

import torch.distributed as dist

dist.init_process_group(backend="cpu:gloo,cuda:nccl")

Ring Anything Attention hiện chỉ hỗ trợ suy luận và không hỗ trợ các mặt nạ chú ý khác None. attn_mask phải là None.

Xem các kết quả benchmark của FLUX.1-dev bên dưới trên một node gồm 4 GPU RTX 4090 (48GB).

Từ bảng trên, Ring Anything Attention cung cấp khả năng tương thích với các độ dài chuỗi tùy ý trong khi duy trì hiệu suất tương đương với Ring Attention tiêu chuẩn. Để biết thêm chi tiết về động lực và các đánh đổi của Ring Anything Attention, hãy xem bình luận này.

parallel_config

Truyền parallel_config trong quá trình khởi tạo mô hình để bật context parallelism.

CKPT_ID = "black-forest-labs/FLUX.1-dev"

cp_config = ContextParallelConfig(ring_degree=2)
transformer = AutoModel.from_pretrained(
    CKPT_ID, 
    subfolder="transformer", 
    dtype=torch.bfloat16,
    parallel_config=cp_config
)

pipeline = DiffusionPipeline.from_pretrained(
    CKPT_ID, transformer=transformer, dtype=torch.bfloat16,
).to(device)

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, tính toán một kết quả một phần, và một AllReduce/AllGather tại ranh giới lớp tái tạo toàn bộ đầu ra. Khác với context parallelism, nó giảm bộ nhớ weight trên mỗi thiết bị, hữu ích cho các mô hình không vừa trên một thiết bị duy nhất.

Truyền một TensorParallelConfig vào tham số parallel_config của from_pretrained() của model. tp_degree là số lượng thiết bị để chia nhỏ (shard) và phải chia hết cho số lượng attention heads của model. Model phải định nghĩa một _tp_plan (một ánh xạ phẳng từ các glob tên module sang kiểu "colwise"/"rowwise").

Việc nạp theo cách này sẽ chia checkpoint thành các mảnh trong khi đọc: mỗi rank chỉ đọc phần riêng của mình trong mỗi trọng lượng đã chia mảnh và đặt trực tiếp lên thiết bị của nó. Không có gì kích thước đầy đủ bao giờ được tạo ra, nên bộ nhớ trên mỗi rank giảm khi tp_degree tăng.

So với việc tải toàn bộ mô hình rồi gọi enable_parallelism(), cách này tải nhanh hơn và dùng ít bộ nhớ CPU hơn trên mỗi rank, với khoảng cách ngày càng lớn khi tp_degree tăng. Các con số bên dưới là cho black-forest-labs/FLUX.2-dev (chỉ transformer, 32B tham số, bf16) trên 4x A10G (23GB).

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()
    device = torch.device(f"cuda:{rank}")
    torch.cuda.set_device(device)

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

    pipeline = DiffusionPipeline.from_pretrained(
        "black-forest-labs/FLUX.2-dev", transformer=transformer, torch_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()

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

tp_degree được lấy từ world_size ở trên, do đó --nproc-per-node 4 sẽ chia transformer thành 4 phần trên 4 thiết bị.

Việc tải với parallel_config tensor-parallel hiện chưa được hỗ trợ 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 quay lại 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 chưa được hỗ trợ, do đó những tính năng này sẽ gây ra lỗi bất kể model được sharded theo cách nào.

Chúng tôi khuyến nghị nạp các mô hình tensor-parallel theo cách này. Bạn cũng có thể chia nhỏ (shard) một mô hình đã nằm trong bộ nhớ bằng cách gọi enable_parallelism() với cùng cấu hình, nhưng cách này sẽ nạp toàn bộ trước rồi mới chia nhỏ lại, do đó tốn bộ nhớ checkpoint đầy đủ trên mỗi rank.

Việc lưu mô hình song song tensor (tensor-parallel) hiện chưa được hỗ trợ, và save_pretrained() sẽ gây lỗi nếu được gọi trên một mô hình như vậy. Hãy lưu mô hình trước khi chia nhỏ (sharding) nó.

Viết kế hoạch song song hóa tensor

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 phẳng ánh xạ các glob tên module sang một 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 tiền tố trước nó phải trỏ đến một nn.ModuleList để một mục duy nhất 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 chúng theo đúng thứ tự đó. Một phép chiếu "colwise" giữ đầu ra ở dạng phân mảnh, phép chiếu "rowwise" tiếp theo tiêu thụ trực tiếp mảnh đó, và một lần AllReduce duy nhất tại ranh giới khối tái tạo kết quả. Phân mảnh cặp này theo bất kỳ cách nào khác sẽ buộc phải thu thập dữ liệu ở giữa và truyền thông nhiều hơn đáng kể.

Đối với attention, điều này có nghĩa là mỗi rank sở hữu một tập con các heads, đó là lý do tại sao tp_degree phải chia hết cho số lượng heads. Các bản sao của encoder-stream (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 image-stream.

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

blocks là một danh sách các số nguyên tỷ lệ có tổng chia hết cho chiều đã đóng gói — [1, 1] cho phép chiếu gate+up của SwiGLU với hai nửa bằng nhau, hoặc [1, 1, 1, 3, 3] cho phép chiếu Q+K+V+gate+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 đối số và lưu các kích thước tuyệt đối trên Linear trong __init__ thay vào đó, 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(),

Cả hai dạng đều xuất hiện trong một kế hoạch trong transformer_flux2.py: cặp SwiGLU double-stream truyền blocks một cách trực tiếp, trong khi phép chiếu QKV+MLP single-stream đã hợp nhất lưu kích thước của nó 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ỏ

Bất kỳ thành phần nào không có trong kế hoạch sẽ được nhân bản trên mọi rank, đây là lựa chọn đúng đắn cho các lớp chuẩn hóa, điều biến AdaLN (img_mod/txt_mod), nhúng patch và văn bản, cùng với norm_out/proj_out cuối cùng. Những thành phần này có kích thước nhỏ, nên việc chia nhỏ chúng chỉ tiết kiệm được ít bộ nhớ trong khi lại làm tăng chi phí truyền thông.

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.

Cả hai đều được xác thực bởi from_pretrained() và enable_parallelism() trước khi bất kỳ trọng số nào được tải hoặc chia nhỏ.

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 kiến trúc tương tự. QwenImageTransformer2DModel hoàn toàn không hợp nhất và mọi mục đều là "colwise"/"rowwise", FluxTransformer2DModel thêm một phép chiếu hàng theo kiểu đóng gói đơn, và Flux2Transformer2DModel bao phủ cả hai kiểu đóng gói — xem transformer_qwenimage.py, transformer_flux.py, và transformer_flux2.py tương ứng.

Chọn một chiến lược

Các chiến lược ở trên giải quyết những vấn đề khác nhau, và câu hỏi hữu ích không phải là chiến lược nào nhanh nhất về mặt lý thuyết, mà là bạn đang thiếu hụt tài nguyên nào.

Một số hướng dẫn thực tế:

• Thông lượng trên nhiều prompt, model đã vừa vặn. Hãy dùng data parallelism. Đây là chiến lược duy nhất ở đây có thể mở rộng thông lượng tuyến tính mà không đụng đến model, và giữ nguyên độ trễ của single-prompt.
• Trọng số của một thành phần đơn lẻ không vừa. Hãy ưu tiên dùng tensor parallelism trước, vì nó giảm cả bộ nhớ lẫn độ trễ. Phương pháp này trao đổi dữ liệu tại mỗi ranh giới khối, nên cần một liên kết nhanh như NVLink; qua PCIe, lưu lượng truyền tải theo từng lớp có thể lớn hơn lượng tính toán mà nó tiết kiệm được, và device_map sẽ là lựa chọn tốt hơn. device_map cũng xử lý được trường hợp các thành phần riêng lẻ đều ổn nhưng tổng thể lại quá lớn.
• Vấn đề nằm ở activations, không phải weights. Đây là chế độ chuỗi dài — ảnh lớn, nhiều khung hình — và context parallelism là câu trả lời trực tiếp. Hãy chọn một backend từ Ulysses/Ring benchmarks ở trên.
• Cả trọng số và chuỗi đều quá lớn. Kết hợp song song tensor và song song ngữ cảnh. TensorParallelConfig nhận một đối số mesh để cả hai có thể chia sẻ một lưới thiết bị duy nhất. Sự kết hợp này chưa được xác thực, vì vậy hãy coi nó là thử nghiệm.