Buckets:
| # AWS Trainium | |
| [TorchNeuron](https://github.com/aws-neuron/torch-neuronx) — это нативный бэкенд PyTorch для Trainium2 и Trainium3. Diffusers работает на этих инстансах через пакет `torch_neuronx`, который регистрирует устройство `"neuron"`, благодаря чему пайплайны выполняются с тем же кодом, что и на любом другом устройстве. | |
| Доступны два режима выполнения: | |
| | Режим | Как активировать | Примечания | | |
| |---|---|---| | |
| | Eager | `pipe.to("neuron")` | Операции, отправляемые по одной, асинхронные | | |
| | Компиляция | `torch.compile(module, backend="neuron")` | Компиляция всего графа, статические формы | | |
| Следуйте [документации AWS Neuron](https://awsdocs-neuron.readthedocs-hosted.com/en/latest/), чтобы настроить Neuron SDK и установить `torch_neuronx`. После установки `import torch` автоматически загружает его и регистрирует устройство `"neuron"`, поэтому единственным необходимым изменением является `pipe.to("neuron")`. Добавляйте `import torch_neuronx` только в том случае, если вы отключили автоматическую загрузку бэкенда с помощью `TORCH_DEVICE_BACKEND_AUTOLOAD=0`. | |
| В режиме eager каждая операция компилируется при первом запуске. В режиме compiled первый вызов компилирует граф. Оба режима по умолчанию кэшируют скомпилированный исполняемый формат Neuron (NEFF) в `tmp/neff_cache`, и этот каталог очищается при перезагрузке. Установите `TORCH_NEURONX_NEFF_CACHE_DIR` в постоянный каталог, чтобы сохранить кэш. | |
| ## Режим Eager | |
| FLUX.2-klein-4B помещается на один NeuronCore, поэтому переместите весь пайплайн на устройство `"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") | |
| ``` | |
| Если модель слишком велика для одного NeuronCore или у вас несколько ядер и требуется меньшая задержка, распределите модели по ядрам. См. раздел [Tensor parallelism](#tensor-parallelism). | |
| ## Скомпилированный режим | |
| TorchNeuron регистрирует `"neuron"` как имя бэкенда для `torch.compile`, поэтому компоненты компилируются так же, как и любые другие цели `torch.compile`. Первый вызов (прогрев) выполняется медленно, так как происходит компиляция; последующие вызовы с теми же формами повторно используют скомпилированный граф. | |
| > [!IMPORTANT] | |
| > TorchNeuron требует *статических форм*, поэтому передайте `dynamic=False`. Новые значения `height` или `width` вызывают повторную компиляцию для соответствующей формы. Уже встреченные формы переиспользуются. Изменение `num_inference_steps` не приводит к повторной компиляции. | |
| ```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") | |
| ``` | |
| Для трансформеров с повторяющимися блоками используйте `compile_repeated_blocks(...)`, чтобы компилировать эти блоки по отдельности. Такой блок переиспользуется, что значительно сокращает время компиляции. | |
| ## Параллелизм тензоров | |
| Следуйте руководству [Tensor parallelism](./tensor_parallelism) для загрузки модели, планов шардирования и ограничений. Чтобы адаптировать его пример для Neuron, замените настройку процесс-группы CUDA и устройств на следующую. | |
| На Neuron инициализируйте группу процессов с `backend="neuron"` и создайте сетку с помощью `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) | |
| ``` | |
| Передайте этот `parallel_config` в метод [from_pretrained()](/docs/diffusers/main/ru/api/models/overview#diffusers.ModelMixin.from_pretrained) трансформера и используйте `device` для остальных компонентов конвейера. | |
| Запустите один процесс на каждый NeuronCore. Установите значение `--nproc_per_node` равным количеству NeuronCores, на которые вы хотите распределить модель. | |
| ```bash | |
| torchrun --nproc_per_node=8 tensor_parallel_neuron.py | |
| ``` | |
Xet Storage Details
- Size:
- 6.62 kB
- Xet hash:
- db8e25c036b2f58ddba786dc790060be2a73c5bd8daa1013a1bde6fe18662e99
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.