Buckets:
| # AWS Trainium | |
| [TorchNeuron](https://github.com/aws-neuron/torch-neuronx) là backend PyTorch native cho Trainium2 và Trainium3. Diffusers chạy trên các instance này thông qua gói `torch_neuronx`, gói này đăng ký một thiết bị `"neuron"`, do đó các pipeline chạy với cùng mã nguồn như trên bất kỳ thiết bị nào khác. | |
| Có hai chế độ thực thi: | |
| | Chế độ | Cách kích hoạt | Ghi chú | | |
| |---|---|---| | |
| | Eager | `pipe.to("neuron")` | Các thao tác được thực hiện lần lượt một tại một thời điểm, bất đồng bộ | | |
| | Biên dịch | `torch.compile(module, backend="neuron")` | Biên dịch toàn đồ thị, hình dạng tĩnh | | |
| Làm theo [tài liệu AWS Neuron](https://awsdocs-neuron.readthedocs-hosted.com/en/latest/) để thiết lập Neuron SDK và cài đặt `torch_neuronx`. Sau khi cài đặt, `import torch` sẽ tự động nạp và đăng ký thiết bị `"neuron"`, do đó `pipe.to("neuron")` là thay đổi duy nhất cần thực hiện. Chỉ thêm `import torch_neuronx` nếu bạn đã tắt tự động nạp backend bằng `TORCH_DEVICE_BACKEND_AUTOLOAD=0`. | |
| Trong chế độ eager, mỗi thao tác được biên dịch lần đầu tiên khi chạy. Trong chế độ compiled, lần gọi đầu tiên sẽ biên dịch đồ thị. Cả hai chế độ đều lưu bộ nhớ đệm định dạng thực thi Neuron đã biên dịch (NEFF) dưới `tmp/neff_cache` theo mặc định, và thư mục này sẽ bị xóa khi khởi động lại. Đặt `TORCH_NEURONX_NEFF_CACHE_DIR` thành một thư mục bền vững để giữ lại bộ nhớ đệm. | |
| ## Chế độ Eager | |
| FLUX.2-klein-4B vừa đủ trên một NeuronCore, vì vậy hãy di chuyển toàn bộ pipeline sang thiết bị `"neuron"`. | |
| ```python | |
| import torch | |
| from diffusers import Flux2KleinPipeline | |
| pipe = Flux2KleinPipeline.from_pretrained("black-forest-labs/FLUX.2-klein-4B", dtype=torch.bfloat16) | |
| pipe.to("neuron") | |
| image = pipe( | |
| prompt="a golden retriever surfing a wave, photorealistic", | |
| height=1024, | |
| width=1024, | |
| num_inference_steps=4, | |
| ).images[0] | |
| image.save("output.png") | |
| ``` | |
| Nếu một model quá lớn cho một NeuronCore, hoặc bạn có nhiều core và muốn giảm độ trễ, hãy chia model thành các shard trên các core thay vì chạy trên một core. Xem phần [Tensor parallelism](#tensor-parallelism). | |
| ## Chế độ biên dịch | |
| TorchNeuron đăng ký `"neuron"` làm tên backend của `torch.compile`, do đó các thành phần được biên dịch giống như mọi mục tiêu `torch.compile` khác. Lệnh gọi đầu tiên (giai đoạn làm nóng) sẽ chậm vì quá trình biên dịch; các lệnh gọi sau với cùng các hình dạng sẽ tái sử dụng đồ thị đã biên dịch. | |
| > [!IMPORTANT] | |
| > TorchNeuron yêu cầu *hình dạng tĩnh*, vì vậy hãy truyền `dynamic=False`. Một `height` hoặc `width` mới sẽ được biên dịch lại cho hình dạng đó. Các hình dạng đã được nhìn thấy sẽ được tái sử dụng. Việc thay đổi `num_inference_steps` sẽ không biên dịch lại. | |
| ```python | |
| import torch | |
| from diffusers import Flux2KleinPipeline | |
| pipe = Flux2KleinPipeline.from_pretrained("black-forest-labs/FLUX.2-klein-4B", dtype=torch.bfloat16) | |
| pipe.to("neuron") | |
| pipe.transformer.compile(backend="neuron", fullgraph=True, dynamic=False) | |
| # Warmup — triggers static graph compilation. | |
| pipe(prompt="warmup", height=1024, width=1024, num_inference_steps=4) | |
| # Later calls with the same shapes reuse the compiled graph. | |
| image = pipe( | |
| prompt="a golden retriever surfing a wave, photorealistic", | |
| height=1024, | |
| width=1024, | |
| num_inference_steps=4, | |
| ).images[0] | |
| image.save("output.png") | |
| ``` | |
| Đối với các transformer có các khối lặp lại, hãy sử dụng `compile_repeated_blocks(...)` để biên dịch các khối đó riêng lẻ. Khối này được tái sử dụng và giúp giảm đáng kể thời gian biên dịch. | |
| ## Song song tensor | |
| Làm theo hướng dẫn [Tensor parallelism](./tensor_parallelism) để tải mô hình, lập kế hoạch chia nhỏ và các hạn chế. Để điều chỉnh ví dụ của nó cho Neuron, hãy thay thế việc thiết lập nhóm tiến trình và thiết bị CUDA bằng các bước sau. | |
| Trên Neuron, khởi tạo nhóm tiến trình với `backend="neuron"` và xây dựng lưới bằng `DeviceMesh("neuron", ...)`. | |
| ```python | |
| import torch | |
| import torch.distributed as dist | |
| from torch.distributed.device_mesh import DeviceMesh | |
| from diffusers import TensorParallelConfig | |
| dist.init_process_group(backend="neuron") | |
| mesh = DeviceMesh("neuron", list(range(dist.get_world_size()))) | |
| device = torch.neuron.current_device() | |
| parallel_config = TensorParallelConfig(mesh=mesh) | |
| ``` | |
| Truyền `parallel_config` này vào [from_pretrained()](/docs/diffusers/main/vi/api/models/overview#diffusers.ModelMixin.from_pretrained) của transformer và sử dụng `device` cho các thành phần pipeline còn lại. | |
| Khởi chạy một tiến trình cho mỗi NeuronCore. Đặt `--nproc_per_node` bằng số NeuronCore mà bạn muốn chia nhỏ (shard) trên đó. | |
| ```bash | |
| torchrun --nproc_per_node=8 tensor_parallel_neuron.py | |
| ``` | |
Xet Storage Details
- Size:
- 5.33 kB
- Xet hash:
- e47e29ac20c33519b754ca08bd1c1d63247f2737e44123a166058adefef29083
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.