Buckets:
| # AWS Trainium | |
| [TorchNeuron](https://github.com/aws-neuron/torch-neuronx) é o backend nativo do PyTorch para Trainium2 e Trainium3. O Diffusers roda nessas instâncias por meio do pacote `torch_neuronx`, que registra um dispositivo `"neuron"`, permitindo que os pipelines sejam executados com o mesmo código de qualquer outro dispositivo. | |
| Dois modos de execução estão disponíveis: | |
| | Modo | Como ativar | Observações | | |
| |---|---|---| | |
| | Eager | `pipe.to("neuron")` | Operações despachadas uma de cada vez, assíncronas | | |
| | Compilar | `torch.compile(module, backend="neuron")` | Compilação de grafo inteiro, formas estáticas | | |
| Siga a [documentação do AWS Neuron](https://awsdocs-neuron.readthedocs-hosted.com/en/latest/) para configurar o Neuron SDK e instalar o `torch_neuronx`. Uma vez instalado, `import torch` o carrega automaticamente e registra o dispositivo `"neuron"`, de modo que `pipe.to("neuron")` é a única alteração necessária. Adicione `import torch_neuronx` apenas se você desativou o carregamento automático do backend com `TORCH_DEVICE_BACKEND_AUTOLOAD=0`. | |
| No modo eager, cada operação é compilada na primeira vez que é executada. No modo compilado, a primeira chamada compila o grafo. Ambos os modos fazem o cache do formato executável Neuron compilado (NEFF) em `tmp/neff_cache` por padrão, e esse diretório é limpo na reinicialização. Defina `TORCH_NEURONX_NEFF_CACHE_DIR` para um diretório persistente para manter o cache. | |
| ## Modo Eager | |
| FLUX.2-klein-4B cabe em um único NeuronCore, então mova o pipeline inteiro para o dispositivo `"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") | |
| ``` | |
| Se um modelo for grande demais para um único NeuronCore, ou se você tiver vários núcleos e quiser menor latência, distribua os modelos entre os núcleos. Consulte a seção [Paralelismo de tensor](#tensor-parallelism). | |
| ## Modo compilado | |
| TorchNeuron registra `"neuron"` como um nome de backend para `torch.compile`, de modo que os componentes são compilados como qualquer outro alvo de `torch.compile`. A primeira chamada (aquecimento) é lenta porque realiza a compilação; chamadas subsequentes com as mesmas formas reutilizam o gráfico compilado. | |
| > [!IMPORTANT] | |
| > TorchNeuron requer *formas estáticas*, então passe `dynamic=False`. Um novo `height` ou `width` compila novamente para essa forma. Formas já vistas são reutilizadas. Alterar `num_inference_steps` não recompila. | |
| ```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") | |
| ``` | |
| Para transformers com blocos repetidos, use `compile_repeated_blocks(...)` para compilar esses blocos individualmente. Esse bloco é reutilizado e reduz o tempo de compilação consideravelmente. | |
| ## Paralelismo de tensor | |
| Siga o guia de [Paralelismo de tensor](./tensor_parallelism) para carregamento de modelos, planos de sharding e limitações. Para adaptar seu exemplo para Neuron, substitua a configuração de grupo de processos e dispositivo CUDA pela seguinte. | |
| No Neuron, inicialize o grupo de processos com `backend="neuron"` e construa a malha com `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) | |
| ``` | |
| Passe este `parallel_config` para o [from_pretrained()](/docs/diffusers/main/pt/api/models/overview#diffusers.ModelMixin.from_pretrained) do transformer e use `device` para os componentes restantes do pipeline. | |
| Inicie um processo por NeuronCore. Defina `--nproc_per_node` para o número de NeuronCores que você deseja dividir. | |
| ```bash | |
| torchrun --nproc_per_node=8 tensor_parallel_neuron.py | |
| ``` | |
Xet Storage Details
- Size:
- 4.73 kB
- Xet hash:
- ea4c90cd57c1d96a40bf27f6e0fe3a4357e096b4529ac746d70194d4fdb8e46f
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.