# 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](https://huggingface.co/docs/accelerate/index) và [PyTorch Distributed](https://pytorch.org/tutorials/beginner/dist_overview.html) để 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. ```bash uv pip install accelerate ``` Khởi tạo lớp [accelerate.PartialState](https://huggingface.co/docs/accelerate/main/en/package_reference/state#accelerate.PartialState) trong một tệp Python để tạo môi trường phân tán. Lớp [accelerate.PartialState](https://huggingface.co/docs/accelerate/main/en/package_reference/state#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](/docs/diffusers/main/vi/api/pipelines/overview#diffusers.DiffusionPipeline) sang `accelerate.PartialState.device` để gán một GPU cho mỗi tiến trình. ```py 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](https://huggingface.co/docs/accelerate/main/en/package_reference/state#accelerate.PartialState.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. ```py 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. ```bash accelerate launch run_distributed.py --num_processes=2 ``` > [!TIP] > Tham khảo [script](https://gist.github.com/sayakpaul/cfaebd221820d7b43fae638b4dfa01ba) 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](https://huggingface.co/docs/accelerate/en/usage_guides/distributed_inference#distributed-inference-with-accelerate). ## Phân tán PyTorch PyTorch [DistributedDataParallel](https://pytorch.org/docs/stable/generated/torch.nn.parallel.DistributedDataParallel.html) cho phép [data parallelism](https://huggingface.co/spaces/nanotron/ultrascale-playbook?section=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. ```py 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](https://pytorch.org/docs/stable/distributed.html?highlight=init_process_group#torch.distributed.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. ```py 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](https://pytorch.org/docs/stable/multiprocessing.html#torch.multiprocessing.spawn) để tạo ra số lượng tiến trình được định nghĩa trong `world_size`. ```py 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. ```bash 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 phối đề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 tải bất kỳ thành phần pipeline nào khác để tránh sử dụng bộ nhớ. ```py 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. ```py 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. ```py from diffusers import AutoModel import torch transformer = AutoModel.from_pretrained( "black-forest-labs/FLUX.1-dev", subfolder="transformer", device_map="auto", dtype=torch.bfloat16 ) ``` > [!TIP] > 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. ```py 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. ```py 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](https://huggingface.co/spaces/nanotron/ultrascale-playbook?section=context_parallelism) 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()](/docs/diffusers/main/vi/api/models/overview#diffusers.ModelMixin.set_attention_backend) để chuyển sang backend attention được tối ưu hơn. Tham khảo [bảng](../optimization/attention_backends#available-backends) 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](https://github.com/huggingface/diffusers/issues/new) nếu một backend không tương thích. ### Ring Attention Key Đại diện (K) và (V) giao tiếp giữa các thiết bị bằng [Ring Attention](https://huggingface.co/papers/2310.01889). Đ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 chú ý 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](/docs/diffusers/main/vi/api/parallel#diffusers.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. ```py 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](https://docs.pytorch.org/docs/stable/elastic/run.html), tương thích với PyTorch. `--nproc-per-node` được đặt bằng số lượng GPU khả dụng. ```shell torchrun --nproc-per-node 2 above_script.py ``` ### Ulysses Attention [Ulysses Attention](https://huggingface.co/papers/2309.14509) 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](/docs/diffusers/main/vi/api/parallel#diffusers.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](/docs/diffusers/main/vi/api/parallel#diffusers.ContextParallelConfig) vào `enable_parallelism()`. ```py # Depending on the number of GPUs available. pipeline.transformer.enable_parallelism(config=ContextParallelConfig(ulysses_degree=2)) ``` ### Unified Attention [Unified Sequence Parallelism](https://huggingface.co/papers/2405.07719) 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](/docs/diffusers/main/vi/api/parallel#diffusers.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](/docs/diffusers/main/vi/api/parallel#diffusers.ContextParallelConfig) với cả `ulysses_degree` và `ring_degree` được đặt lớn hơn 1 vào `enable_parallelism()`. ```py pipeline.transformer.enable_parallelism(config=ContextParallelConfig(ulysses_degree=2, ring_degree=2)) ``` > [!TIP] > 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](https://github.com/huggingface/diffusers/pull/12693#issuecomment-3694727532) trên một node gồm 4 GPU H100. Kết quả được tóm tắt như sau: | CP Backend | Thời gian / Lặp (ms) | Bước / Giây | Bộ nhớ đỉnh (GB) | |--------------------|------------------|-------------|------------------| | ulysses | 6670.789 | 7.50 | 33.85 | | ring | 13076.492 | 3.82 | 56.02 | | unified_balanced | 11068.705 | 4.52 | 33.85 | 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](https://github.com/huggingface/diffusers/pull/12996) 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](/docs/diffusers/main/vi/api/parallel#diffusers.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](/docs/diffusers/main/vi/api/parallel#diffusers.ContextParallelConfig) với cả `ulysses_degree` được đặt lớn hơn 1 và `ulysses_anything=True` vào `enable_parallelism()`. ```py pipeline.transformer.enable_parallelism(config=ContextParallelConfig(ulysses_degree=2, ulysses_anything=True)) ``` > [!TIP] Để tránh nhiều lần đồng bộ hóa CUDA bị ép do các thao tác chuyển dữ liệu 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](https://github.com/huggingface/diffusers/pull/12996#issuecomment-3797695999) trên một node gồm 4 GPU L20. Kết quả được tóm tắt như sau: | CP Backend | Thời gian / Lặp (ms) | Bước / Giây | Bộ nhớ đỉnh (GB) | Hình dạng (HxW)| |--------------------|------------------|-------------|------------------|------------| | ulysses | 281.07 | 3.56 | 37.11 | 1024x1024 | | ring | 351.34 | 2.85 | 37.01 | 1024x1024 | | unified_balanced | 324.37 | 3.08 | 37.16 | 1024x1024 | | ulysses_anything | 280.94 | 3.56 | 37.11 | 1024x1024 | | ulysses | thất bại | thất bại | thất bại | 1008x1008 | | ring | thất bại | thất bại | thất bại | 1008x1008 | | unified_balanced | thất bại | thất bại | thất bại | 1008x1008 | | ulysses_anything | 278.40 | 3.59 | 36.99 | 1008x1008 | 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](https://huggingface.co/papers/2310.01889) 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](https://github.com/huggingface/diffusers/pull/13545#issuecomment-4302195582) 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 mỗi 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. ```py pipeline.transformer.enable_parallelism(config=ContextParallelConfig(ring_degree=2, ring_anything=True)) ``` > [!TIP] > Thêm backend `gloo` vào [init_process_group](https://docs.pytorch.org/docs/stable/distributed.html#torch.distributed.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. ```py import torch.distributed as dist dist.init_process_group(backend="cpu:gloo,cuda:nccl") ``` > [!NOTE] > 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). | CP Backend | Thời gian / Lặp (ms) | Bước / Giây | Bộ nhớ đỉnh (GB) | Hình dạng (HxW)| |--------------------|------------------|-------------|------------------|------------| | ulysses | 259.07 | 3.86 | 33.83 | 1024x1024 | | ring | 338.98 | 2.95 | 33.83 | 1024x1024 | | unified_balanced | 321.54 | 3.11 | 33.83 | 1024x1024 | | ulysses_anything | 259.07 | 3.86 | 33.83 | 1024x1024 | | ring_anything | 340.14 | 2.94 | 33.83 | 1024x1024 | | ulysses | thất bại | thất bại | thất bại | 1008x1008 | | ring | thất bại | thất bại | thất bại | 1008x1008 | | unified_balanced | thất bại | thất bại | thất bại | 1008x1008 | | ulysses_anything | 253.16 | 3.95 | 33.75 | 1008x1008 | | ring_anything | 335.57 | 2.98 | 33.75 | 1008x1008 | 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](https://github.com/huggingface/diffusers/pull/13545#issuecomment-4304104462). ### parallel_config Truyền `parallel_config` trong quá trình khởi tạo mô hình để bật context parallelism. ```py 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](https://huggingface.co/spaces/nanotron/ultrascale-playbook?section=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](/docs/diffusers/main/vi/api/parallel#diffusers.TensorParallelConfig) vào tham số `parallel_config` của [from_pretrained()](/docs/diffusers/main/vi/api/models/overview#diffusers.ModelMixin.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](https://huggingface.co/black-forest-labs/FLUX.2-dev) (chỉ transformer, 32B tham số, bf16) trên 4x A10G (23GB). | tp_degree | phương thức | thời gian tải | CPU tối đa/mỗi rank | |---|---|---|---| | 4 | `from_pretrained(parallel_config=...)` | 12.5s | 6.8GB | | 4 | `from_pretrained` + `enable_parallelism` | 30.4s | 64.1GB | ```py 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() ``` ```shell 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ị. > [!CAUTION] > 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 đó chúng 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. Chưa hỗ trợ lưu mô hình song song tensor, và [save_pretrained()](/docs/diffusers/main/vi/api/models/overview#diffusers.ModelMixin.save_pretrained) sẽ gây lỗi nếu 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`](https://pytorch.org/docs/stable/generated/torch.nn.ModuleList.html) để 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 | Phong cách | Các phân mảnh | Mỗi rank | Dùng cho | |---|---|---|---| | `"colwise"` | đặc trưng đầu ra (`weight` chiều 0) | tính toán một lát cắt của đầu ra | `to_q`, `to_k`, `to_v`, FFN in-projection | | `"rowwise"` | đặc trưng đầu vào (chiều 1 của `weight`) | tính tổng một phần | `to_out.0`, chiếu ra đầu ra của FFN | | `"rowwise_split_input"` | đặc trưng đầu vào (chiều 1 của `weight`) | tính toán một kết quả một phần bằng cách sử dụng lát cắt của đầu vào và trọng số của nó; `AllReduce` cộng các kết quả này lại để mọi rank đều nhận được tất cả các đặc trưng đầu ra | một phép chiếu độc lập đọc một kích hoạt được nhân bản, ví dụ như một phép chiếu điều biến AdaLN từ nhúng bước thời gian | 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. ```py # 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`. ```py "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`. ```py # in the attention module's __init__ self.to_out._tp_packed_row_blocks = [self.inner_dim, self.mlp_hidden_dim] ``` ```py # 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`](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/transformers/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()](/docs/diffusers/main/vi/api/models/overview#diffusers.ModelMixin.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. > [!TIP] > Bắt đầu từ một kế hoạch có sẵn cho kiến trúc tương tự. [QwenImageTransformer2DModel](/docs/diffusers/main/vi/api/models/qwenimage_transformer2d#diffusers.QwenImageTransformer2DModel) hoàn toàn không hợp nhất và mọi mục đều là `"colwise"`/`"rowwise"`, [FluxTransformer2DModel](/docs/diffusers/main/vi/api/models/flux_transformer#diffusers.FluxTransformer2DModel) thêm một phép chiếu hàng theo kiểu đóng gói đơn, và [Flux2Transformer2DModel](/docs/diffusers/main/vi/api/models/flux2_transformer#diffusers.Flux2Transformer2DModel) bao phủ cả hai kiểu đóng gói — xem [`transformer_qwenimage.py`](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/transformers/transformer_qwenimage.py), [`transformer_flux.py`](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/transformers/transformer_flux.py), và [`transformer_flux2.py`](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/transformers/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. | Chiến lược | Các phần chia | Giảm | Độ trễ cho một prompt | Phù hợp nhất khi | |---|---|---|---|---| | [Accelerate](#accelerate) / [DDP](#pytorch-distributed) | prompts giữa các bản sao | không có gì — mỗi thiết bị giữ một bản sao đầy đủ | không thay đổi | mô hình đã vừa vặn và bạn có nhiều prompt | | [`device_map`](#devicemap) | các thành phần trên các thiết bị | bộ nhớ trọng số | kém hơn một chút | mô hình không vừa và kết nối giữa các thành phần chậm | | [Song song ngữ cảnh](#context-parallelism) | chuỗi đầu vào | bộ nhớ kích hoạt | thấp hơn | các chuỗi dài — độ phân giải cao hoặc video | | [Song song tensor](#tensor-parallelism) | các ma trận trọng số | bộ nhớ trọng số | thấp hơn | trọng số của một thành phần không vừa và kết nối giữa các thành phần rất nhanh | 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](#ulysses-attention) ở 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](/docs/diffusers/main/vi/api/parallel#diffusers.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.