Buckets:
| # AWS Trainium | |
| [TorchNeuron](https://github.com/aws-neuron/torch-neuronx) es el backend nativo de PyTorch para Trainium2 y Trainium3. Diffusers se ejecuta en estas instancias a través del paquete `torch_neuronx`, que registra un dispositivo `"neuron"`, de modo que los pipelines se ejecutan con el mismo código que en cualquier otro dispositivo. | |
| Hay dos modos de ejecución disponibles: | |
| | Modo | Cómo activar | Notas | | |
| |---|---|---| | |
| | Eager | `pipe.to("neuron")` | Operaciones despachadas una a la vez, asíncronas | | |
| | Compilar | `torch.compile(module, backend="neuron")` | Compilación de todo el grafo, formas estáticas | | |
| Siga la [documentación de AWS Neuron](https://awsdocs-neuron.readthedocs-hosted.com/en/latest/) para configurar el SDK de Neuron e instalar `torch_neuronx`. Una vez instalado, `import torch` lo carga automáticamente y registra el dispositivo `"neuron"`, por lo que `pipe.to("neuron")` es el único cambio necesario. Añada `import torch_neuronx` solo si deshabilitó la carga automática del backend con `TORCH_DEVICE_BACKEND_AUTOLOAD=0`. | |
| En modo eager, cada operación se compila la primera vez que se ejecuta. En modo compilado, la primera llamada compila el grafo. Ambos modos almacenan en caché el formato ejecutable de Neuron (NEFF) compilado en `tmp/neff_cache` de forma predeterminada, y ese directorio se borra al reiniciar. Establezca `TORCH_NEURONX_NEFF_CACHE_DIR` en un directorio persistente para conservar la caché. | |
| ## Modo eager | |
| FLUX.2-klein-4B cabe en un solo NeuronCore, por lo que mueva todo el pipeline al 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") | |
| ``` | |
| Si un modelo es demasiado grande para un solo NeuronCore, o si tienes varios núcleos y deseas una latencia más baja, divide los modelos entre los núcleos. Consulta la sección [Tensor parallelism](#tensor-parallelism). | |
| ## Modo compilado | |
| TorchNeuron registra `"neuron"` como un nombre de backend para `torch.compile`, por lo que los componentes se compilan como cualquier otro destino de `torch.compile`. La primera llamada (calentamiento) es lenta porque compila; las llamadas posteriores con las mismas formas reutilizan el gráfico compilado. | |
| > [!IMPORTANT] | |
| > TorchNeuron requiere *formas estáticas*, por lo que debe pasar `dynamic=False`. Un nuevo `height` o `width` vuelve a compilar para esa forma. Las formas ya vistas se reutilizan. Cambiar `num_inference_steps` no vuelve a compilar. | |
| ```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 los transformers con bloques repetidos, use `compile_repeated_blocks(...)` para compilar esos bloques de forma individual. Este bloque se reutiliza y reduce considerablemente el tiempo de compilación. | |
| ## Paralelismo de tensores | |
| Siga la guía de [Paralelismo de tensores](./tensor_parallelism) para la carga de modelos, los planes de particionamiento y las limitaciones. Para adaptar su ejemplo a Neuron, reemplace la configuración del grupo de procesos y el dispositivo de CUDA con lo siguiente. | |
| En Neuron, inicializa el grupo de procesos con `backend="neuron"` y construye la malla con `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) | |
| ``` | |
| Pasa este `parallel_config` al [from_pretrained()](/docs/diffusers/main/es/api/models/overview#diffusers.ModelMixin.from_pretrained) del transformer y usa `device` para los componentes restantes del pipeline. | |
| Inicie un proceso por NeuronCore. Establezca `--nproc_per_node` en el número de NeuronCores que desea dividir. | |
| ```bash | |
| torchrun --nproc_per_node=8 tensor_parallel_neuron.py | |
| ``` | |
Xet Storage Details
- Size:
- 4.75 kB
- Xet hash:
- d0fd25342951ca0ff0a24c2031593528245351e0031ea764bccbbc042094be28
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.