Buckets:
| # 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 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](./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](../training/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](/docs/diffusers/v0.43.0/vi/api/parallel#diffusers.TensorParallelConfig) vào tham số `parallel_config` của [from_pretrained()](/docs/diffusers/v0.43.0/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 đầu chú ý (attention heads) của model. Lớp của model phải định nghĩa `_tp_plan`, ánh xạ tên module sang một kiểu chia nhỏ. Để kiểm tra xem một model có hỗ trợ tensor parallelism 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](/docs/diffusers/v0.43.0/vi/api/models/flux2_transformer#diffusers.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](#add-tensor-parallelism-to-a-model). | |
| 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. | |
| ```py | |
| 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`. | |
| ```shell | |
| 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](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 | | |
| ## 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()](/docs/diffusers/v0.43.0/vi/api/models/overview#diffusers.ModelMixin.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`](https://pytorch.org/docs/stable/generated/torch.nn.ModuleList.html), 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 | |
| | 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 tổng một phần từ lát cắt của đầu vào được nhân bản | một phép chiếu độc lập có đầu vào được nhân bản, chẳng hạn như phép chiếu điều biến AdaLN đọc nhúng bước thời gian | | |
| 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. | |
| ```py | |
| # 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`. | |
| ```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 `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`. | |
| ```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(), | |
| ``` | |
| Kế hoạch trong [`transformer_flux2.py`](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/transformers/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 nằm trong kế hoạch sẽ được nhân bản trên mọi rank. Hãy loại bỏ 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ể loại bỏ. 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()](/docs/diffusers/v0.43.0/vi/api/models/overview#diffusers.ModelMixin.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. | |
| > [!TIP] | |
| > 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`](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/transformers/transformer_qwenimage.py) không có các phép chiếu đã hợp nhất, do đó mọi mục đều là `"colwise"` hoặc `"rowwise"` thuần túy. - [`transformer_flux.py`](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/transformers/transformer_flux.py) thêm một phép chiếu row-wise đã đóng gói. - [`transformer_flux2.py`](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/transformers/transformer_flux2.py) sử dụng cả hai kiểu đóng gói. | |
Xet Storage Details
- Size:
- 12.3 kB
- Xet hash:
- 5a688380dea0de8f5bc6705b8dce0076b112f003c049ead44f7a6ba356dd0f20
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.