Buckets:
| # AWS Trainium | |
| [TorchNeuron](https://github.com/aws-neuron/torch-neuronx) adalah backend PyTorch native untuk Trainium2 dan Trainium3. Diffusers berjalan pada instance ini melalui paket `torch_neuronx`, yang mendaftarkan perangkat `"neuron"`, sehingga pipeline berjalan dengan kode yang sama seperti pada perangkat lainnya. | |
| Dua mode eksekusi tersedia: | |
| | Mode | Cara mengaktifkan | Catatan | | |
| |---|---|---| | |
| | Eager | `pipe.to("neuron")` | Operasi yang dikirim satu per satu, asinkron | | |
| | Compile | `torch.compile(module, backend="neuron")` | Kompilasi grafik utuh, bentuk statis | | |
| Ikuti [dokumentasi AWS Neuron](https://awsdocs-neuron.readthedocs-hosted.com/en/latest/) untuk mengatur SDK Neuron dan menginstal `torch_neuronx`. Setelah terinstal, `import torch` akan memuatnya secara otomatis dan mendaftarkan perangkat `"neuron"`, sehingga `pipe.to("neuron")` adalah satu-satunya perubahan yang diperlukan. Tambahkan `import torch_neuronx` hanya jika Anda menonaktifkan pemuatan otomatis backend dengan `TORCH_DEVICE_BACKEND_AUTOLOAD=0`. | |
| Dalam mode eager, setiap operasi dikompilasi pada eksekusi pertamanya. Dalam mode compiled, panggilan pertama mengompilasi graf. Kedua mode tersebut meng-cache format eksekutabel Neuron yang dikompilasi (NEFF) di bawah `tmp/neff_cache` secara default, dan direktori itu dibersihkan saat reboot. Atur `TORCH_NEURONX_NEFF_CACHE_DIR` ke direktori persisten untuk mempertahankan cache. | |
| ## Mode Eager | |
| FLUX.2-klein-4B muat di satu NeuronCore, jadi pindahkan seluruh pipeline ke perangkat `"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") | |
| ``` | |
| Jika model terlalu besar untuk satu NeuronCore, atau Anda memiliki beberapa core dan ingin latensi yang lebih rendah, shard model di antara core. Lihat bagian [Tensor parallelism](#tensor-parallelism). | |
| ## Mode terkompilasi | |
| TorchNeuron mendaftarkan `"neuron"` sebagai nama backend untuk `torch.compile`, sehingga komponen dikompilasi seperti target `torch.compile` lainnya. Panggilan pertama (warmup) lambat karena proses kompilasi; panggilan berikutnya dengan bentuk (shapes) yang sama akan menggunakan ulang graf yang telah dikompilasi. | |
| > [!IMPORTANT] | |
| > TorchNeuron memerlukan *bentuk statis*, jadi berikan `dynamic=False`. Nilai `height` atau `width` baru akan dikompilasi ulang untuk bentuk tersebut. Bentuk yang sudah pernah dilihat akan digunakan kembali. Mengubah `num_inference_steps` tidak memicu kompilasi ulang. | |
| ```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") | |
| ``` | |
| Untuk transformer dengan blok berulang, gunakan `compile_repeated_blocks(...)` untuk mengompilasi blok-blok tersebut secara individual. Blok ini digunakan kembali dan mengurangi waktu kompilasi secara signifikan. | |
| ## Paralelisme tensor | |
| Ikuti panduan [Tensor parallelism](./tensor_parallelism) untuk pemuatan model, rencana sharding, dan keterbatasannya. Untuk menyesuaikan contohnya untuk Neuron, ganti setup process-group dan perangkat CUDA dengan berikut. | |
| Pada Neuron, inisialisasi grup proses dengan `backend="neuron"` dan bangun mesh dengan `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) | |
| ``` | |
| Lewatkan `parallel_config` ini ke [from_pretrained()](/docs/diffusers/v0.43.0/id/api/models/overview#diffusers.ModelMixin.from_pretrained) transformer dan gunakan `device` untuk komponen pipeline lainnya. | |
| Luncurkan satu proses per NeuronCore. Atur `--nproc_per_node` ke jumlah NeuronCore yang ingin Anda shard. | |
| ```bash | |
| torchrun --nproc_per_node=8 tensor_parallel_neuron.py | |
| ``` | |
Xet Storage Details
- Size:
- 4.66 kB
- Xet hash:
- eab87d944f8faf57bd81648f4f1198ab371fa17f30a39b9bfa9e86b37c3fa0d0
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.