Buckets:
| # AWS Trainium | |
| [TorchNeuron](https://github.com/aws-neuron/torch-neuronx)은 Trainium2와 Trainium3의 네이티브 PyTorch 백엔드입니다. Diffusers는 `torch_neuronx` 패키지를 통해 이러한 인스턴스에서 실행되며, 이 패키지는 `"neuron"` 장치를 등록하여 파이프라인이 다른 어떤 장치와 동일한 코드로 실행되도록 합니다. | |
| 두 가지 실행 모드가 있습니다: | |
| | 모드 | 활성화 방법 | The input is a single word "Notes". | |
| 참고 | | |
| |---|---|---| | |
| | Eager | `pipe.to("neuron")` | Operations dispatched one at a time, asynchronous | | |
| | 컴파일 | `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")`만 변경하면 됩니다. `TORCH_DEVICE_BACKEND_AUTOLOAD=0`로 백엔드 자동 로딩을 비활성화한 경우에만 `import torch_neuronx`를 추가하세요. | |
| 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` 대상과 동일한 방식으로 컴파일됩니다. 첫 번째 호출(워밍업)은 컴파일을 수행하기 때문에 느리지만, 동일한 shape을 가진 이후의 호출은 컴파일된 그래프를 재사용합니다. | |
| > [!IMPORTANT] | |
| > TorchNeuron은 *정적 shape*를 요구하므로 `dynamic=False`를 전달하세요. 새로운 `height` 또는 `width`는 해당 shape에 대해 다시 컴파일됩니다. 이미 확인된 shape는 재사용됩니다. `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") | |
| ``` | |
| 반복 블록을 가진 transformers의 경우, `compile_repeated_blocks(...)`를 사용하여 해당 블록을 개별적으로 컴파일합니다. 이 블록은 재사용되며 컴파일 시간을 상당히 줄입니다. | |
| ## Tensor 병렬성 | |
| 모델 로딩, 샤딩 플랜, 제한 사항에 대해서는 [Tensor parallelism](../training/distributed_inference#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`를 transformer의 [from_pretrained()](/docs/diffusers/main/ko/api/models/overview#diffusers.ModelMixin.from_pretrained)에 전달하고 나머지 파이프라인 구성 요소에는 `device`를 사용하세요. | |
| NeuronCore마다 프로세스 하나를 시작합니다. 샤딩할 NeuronCore의 수를 `--nproc_per_node`에 설정합니다. | |
| ```bash | |
| torchrun --nproc_per_node=8 tensor_parallel_neuron.py | |
| ``` | |
Xet Storage Details
- Size:
- 5.09 kB
- Xet hash:
- a192fe5869427e6f8263465e80bd56c5e8203285d3cdf4d29fd42e862bb8e424
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.