Buckets:

hf-doc-build
/

doc

Files

xet

hf-doc-build/doc / diffusers /main /ko /optimization /tpu.md

mishig

about 10 hours ago

preview code

download

raw

7.73 kB

# Google TPU

Tensor Processing Units(TPU)는 Google의 머신러닝 가속기로, Cloud TPU(v6e, v7x)로 제공됩니다. Diffusers는 TPU용 PyTorch 백엔드인 [TorchTPU](https://github.com/google-pytorch/torch_tpu/)를 통해 TPU에서 최소한의 코드 변경으로 실행됩니다.

세 가지 실행 모드가 사용할 수 있습니다:

| 모드 | 상수 | 활성화 방법 | The input is a single word "Notes".

참고 |
|---|---|---|---|
| 엄격한 eager (기본값) | `EagerMode.DEFER_NEVER` | `pipe.to("tpu")` | Operations dispatched one at a time, asynchronous |
| 최적화된 eager | `EagerMode.DEFER_AND_FUSE` | `TORCH_TPU_DEFER_AND_FUSE=1` | 작업이 지연되고 더 큰 그래프로 융합됩니다. 엄격한 eager보다 빠르지만 컴파일이 더 오래 걸립니다 |
| 컴파일 | — | `torch.compile(module, backend="tpu")` | `TpuBackend`를 사용한 AOT 컴파일 |

[TorchTPU 설치 가이드](https://github.com/google-pytorch/torch_tpu/)를 따르세요. 설치 후 `import torch`를 실행하면 자동으로 로드되고 `"tpu"` 디바이스가 등록되므로, `pipe.to("tpu")`만 변경하면 됩니다. `TORCH_DEVICE_BACKEND_AUTOLOAD=0`로 백엔드 자동 로딩을 비활성화한 경우에만 `import torch_tpu`를 추가하세요.

## Eager 모드

FLUX.1-schnell은 단일 v6e 칩에 한 번에 모두 들어가지 않으므로, 각 모델이 실행되는 동안에만 TPU로 이동하도록 [enable_model_cpu_offload()](/docs/diffusers/main/ko/api/pipelines/overview#diffusers.DiffusionPipeline.enable_model_cpu_offload)를 사용하세요. `"tpu"` 장치를 자동으로 감지합니다.

```python
import torch
from diffusers import FluxPipeline
pipe = FluxPipeline.from_pretrained("black-forest-labs/FLUX.1-schnell", dtype=torch.bfloat16)
pipe.enable_model_cpu_offload()
image = pipe(
    prompt="a golden retriever surfing a wave, photorealistic",
    height=1024,
    width=1024,
    num_inference_steps=4,
    guidance_scale=0.0,
).images[0]
image.save("output.png")
```

모델이 단일 칩에 너무 크거나, 여러 칩을 보유하고 있어 낮은 지연 시간을 원한다면, 대신 모델을 칩에 걸쳐 샤딩하세요. [Tensor parallelism](#tensor-parallelism) 섹션을 참조하세요.

## 컴파일 모드

TorchTPU는 `"tpu"`를 `torch.compile` 백엔드 이름으로 등록합니다(내부적으로는 `TpuBackend`를 사용하므로), 따라서 각 구성 요소는 다른 `torch.compile` 대상과 동일한 방식으로 컴파일됩니다. 첫 번째 호출(워밍업)은 컴파일을 수행하기 때문에 느리지만, 동일한 shape을 가진 후속 호출은 컴파일된 그래프를 재사용합니다.

> [!IMPORTANT]
> TorchTPU는 **정적 shape**를 요구하므로 `dynamic=False`를 전달해야 합니다. 새로운 `height` 또는 `width`는 해당 shape에 대해 한 번만 다시 컴파일되며, 이미 확인된 shape는 재사용됩니다. `num_inference_steps`를 변경해도 다시 컴파일되지 않습니다.

파이프라인 전체가 하나의 칩에 들어갈 수 있는 경우, 이를 TPU로 이동하고 전체 transformer를 컴파일합니다. Stable Diffusion 3.5 Medium(bf16 기준 약 15GB)은 단일 v6e 칩에 들어갑니다.

```python
import torch
from diffusers import StableDiffusion3Pipeline
pipe = StableDiffusion3Pipeline.from_pretrained("stabilityai/stable-diffusion-3.5-medium", dtype=torch.bfloat16)
pipe.to("tpu")
pipe.transformer.compile(backend="tpu", fullgraph=True, dynamic=False)
# Warmup — triggers static graph compilation.
pipe(prompt="warmup", height=1024, width=1024, num_inference_steps=40, guidance_scale=4.5)
# 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=40,
    guidance_scale=4.5,
).images[0]
image.save("output.png")
```

파이프라인이 하나의 칩에 맞지 않는 경우:

- 여러 개의 칩을 사용하는 경우, 대신 [tensor parallelism](#tensor-parallelism)으로 샤딩하세요. 모든 것이 TPU에 유지되며, `pipe.transformer.compile(...)`은 샤딩된 transformer에서도 동일한 방식으로 작동합니다.
- [enable_model_cpu_offload()](/docs/diffusers/main/ko/api/pipelines/overview#diffusers.DiffusionPipeline.enable_model_cpu_offload)를 사용하면 오프로드 훅을 `torch.compile`로 추적할 수 없습니다. 대신 `pipe.transformer.compile_repeated_blocks(backend="tpu", fullgraph=True, dynamic=False)`를 사용하여 transformer의 반복 블록만 컴파일하세요.

## Tensor 병렬성

하나의 칩에 너무 큰 모델을 여러 칩에 걸쳐 샤딩합니다. FLUX.2-dev의 텍스트 인코더(~48GB)와 트랜스포머(~64GB)는 각각 단일 칩을 초과하므로, 아래 예시에서는 둘 다 샤딩합니다:

- `parallel_config` 인수로 전달되는 [from_pretrained()](/docs/diffusers/main/ko/api/models/overview#diffusers.ModelMixin.from_pretrained)에 [TensorParallelConfig](/docs/diffusers/main/ko/api/parallel#diffusers.TensorParallelConfig)를 갖춘 transformer입니다. 각 rank는 샤딩된 가중치의 자기 자신의 슬라이스만 읽으므로 전체 모델은 절대 메모리에 로드되지 않습니다. 일반적인 TP 세부 사항(`_tp_plan`, colwise/rowwise)은 [Tensor parallelism](../training/distributed_inference#tensor-parallelism) 가이드를 참조하세요.
- Transformers 자체의 [tensor parallelism](https://huggingface.co/docs/transformers/perf_infer_gpu_multi)를 사용하여 텍스트 인코더를 구성하고, [DistributedConfig](https://huggingface.co/docs/transformers/main/en/expert_parallelism#transformers.DistributedConfig)와 동일한 mesh를 전달합니다.

TPU에서는 `backend="tpu_dist"`로 프로세스 그룹을 초기화하고 `DeviceMesh("tpu", ...)`로 메시를 구성합니다.

```python
import torch
import torch.distributed as dist
from torch.distributed.device_mesh import DeviceMesh
from transformers import DistributedConfig, Mistral3ForConditionalGeneration
from diffusers import Flux2Pipeline, Flux2Transformer2DModel, TensorParallelConfig
dist.init_process_group(backend="tpu_dist")
mesh = DeviceMesh("tpu", list(range(dist.get_world_size())))
repo_id = "black-forest-labs/FLUX.2-dev"
text_encoder = Mistral3ForConditionalGeneration.from_pretrained(
    repo_id,
    subfolder="text_encoder",
    dtype=torch.bfloat16,
    distributed_config=DistributedConfig(tp_plan="auto"),
    device_mesh=mesh,
)
transformer = Flux2Transformer2DModel.from_pretrained(
    repo_id, subfolder="transformer", dtype=torch.bfloat16, parallel_config=TensorParallelConfig(mesh=mesh)
)
pipe = Flux2Pipeline.from_pretrained(
    repo_id, text_encoder=text_encoder, transformer=transformer, dtype=torch.bfloat16
)
pipe.vae.to("tpu")
image = pipe(
    prompt="a golden retriever surfing a wave, photorealistic",
    num_inference_steps=28,
    generator=torch.Generator("cpu").manual_seed(0),
).images[0]
if dist.get_rank() == 0:
    image.save("output.png")
```

칩당 하나의 프로세스를 시작합니다. 호스트의 모든 TPU 칩 수를 사용하려면 `--nproc_per_node`를 설정합니다.

```bash
eval $(python -m torch_tpu._internal.distributed.launchers.singlehost_wrapper | sed 's/^/export /')
torchrun --nproc_per_node=8 flux2_tp.py
```

> [!WARNING]
> 위 시작 명령은 TorchTPU 내부 모듈(`torch_tpu._internal`)을 사용하며, 이는 안정화된 API가 아니므로 향후 TorchTPU 릴리스에서 변경될 수 있습니다. 멀티 칩 작업을 시작하는 최신 방법은 [TorchTPU 저장소](https://github.com/google-pytorch/torch_tpu/)에서 확인하세요.

Xet Storage Details

Size:: 7.73 kB
Xet hash:: 54de90521e8932b94fbcdbeab5a3026cc0d9e09ff8a462b8466d8455f8266946

Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.

	# Google TPU

	Tensor Processing Units(TPU)는 Google의 머신러닝 가속기로, Cloud TPU(v6e, v7x)로 제공됩니다. Diffusers는 TPU용 PyTorch 백엔드인 [TorchTPU](https://github.com/google-pytorch/torch_tpu/)를 통해 TPU에서 최소한의 코드 변경으로 실행됩니다.

	세 가지 실행 모드가 사용할 수 있습니다:

	\| 모드 \| 상수 \| 활성화 방법 \| The input is a single word "Notes".

	참고 \|
	\|---\|---\|---\|---\|
	\| 엄격한 eager (기본값) \| `EagerMode.DEFER_NEVER` \| `pipe.to("tpu")` \| Operations dispatched one at a time, asynchronous \|
	\| 최적화된 eager \| `EagerMode.DEFER_AND_FUSE` \| `TORCH_TPU_DEFER_AND_FUSE=1` \| 작업이 지연되고 더 큰 그래프로 융합됩니다. 엄격한 eager보다 빠르지만 컴파일이 더 오래 걸립니다 \|
	\| 컴파일 \| — \| `torch.compile(module, backend="tpu")` \| `TpuBackend`를 사용한 AOT 컴파일 \|

	[TorchTPU 설치 가이드](https://github.com/google-pytorch/torch_tpu/)를 따르세요. 설치 후 `import torch`를 실행하면 자동으로 로드되고 `"tpu"` 디바이스가 등록되므로, `pipe.to("tpu")`만 변경하면 됩니다. `TORCH_DEVICE_BACKEND_AUTOLOAD=0`로 백엔드 자동 로딩을 비활성화한 경우에만 `import torch_tpu`를 추가하세요.

	## Eager 모드

	FLUX.1-schnell은 단일 v6e 칩에 한 번에 모두 들어가지 않으므로, 각 모델이 실행되는 동안에만 TPU로 이동하도록 [enable_model_cpu_offload()](/docs/diffusers/main/ko/api/pipelines/overview#diffusers.DiffusionPipeline.enable_model_cpu_offload)를 사용하세요. `"tpu"` 장치를 자동으로 감지합니다.

	```python
	import torch

	from diffusers import FluxPipeline

	pipe = FluxPipeline.from_pretrained("black-forest-labs/FLUX.1-schnell", dtype=torch.bfloat16)
	pipe.enable_model_cpu_offload()

	image = pipe(
	prompt="a golden retriever surfing a wave, photorealistic",
	height=1024,
	width=1024,
	num_inference_steps=4,
	guidance_scale=0.0,
	).images[0]

	image.save("output.png")
	```

	모델이 단일 칩에 너무 크거나, 여러 칩을 보유하고 있어 낮은 지연 시간을 원한다면, 대신 모델을 칩에 걸쳐 샤딩하세요. [Tensor parallelism](#tensor-parallelism) 섹션을 참조하세요.

	## 컴파일 모드

	TorchTPU는 `"tpu"`를 `torch.compile` 백엔드 이름으로 등록합니다(내부적으로는 `TpuBackend`를 사용하므로), 따라서 각 구성 요소는 다른 `torch.compile` 대상과 동일한 방식으로 컴파일됩니다. 첫 번째 호출(워밍업)은 컴파일을 수행하기 때문에 느리지만, 동일한 shape을 가진 후속 호출은 컴파일된 그래프를 재사용합니다.

	> [!IMPORTANT]
	> TorchTPU는 정적 shape를 요구하므로 `dynamic=False`를 전달해야 합니다. 새로운 `height` 또는 `width`는 해당 shape에 대해 한 번만 다시 컴파일되며, 이미 확인된 shape는 재사용됩니다. `num_inference_steps`를 변경해도 다시 컴파일되지 않습니다.

	파이프라인 전체가 하나의 칩에 들어갈 수 있는 경우, 이를 TPU로 이동하고 전체 transformer를 컴파일합니다. Stable Diffusion 3.5 Medium(bf16 기준 약 15GB)은 단일 v6e 칩에 들어갑니다.

	```python
	import torch

	from diffusers import StableDiffusion3Pipeline

	pipe = StableDiffusion3Pipeline.from_pretrained("stabilityai/stable-diffusion-3.5-medium", dtype=torch.bfloat16)
	pipe.to("tpu")
	pipe.transformer.compile(backend="tpu", fullgraph=True, dynamic=False)

	# Warmup — triggers static graph compilation.
	pipe(prompt="warmup", height=1024, width=1024, num_inference_steps=40, guidance_scale=4.5)

	# 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=40,
	guidance_scale=4.5,
	).images[0]

	image.save("output.png")
	```

	파이프라인이 하나의 칩에 맞지 않는 경우:

	- 여러 개의 칩을 사용하는 경우, 대신 [tensor parallelism](#tensor-parallelism)으로 샤딩하세요. 모든 것이 TPU에 유지되며, `pipe.transformer.compile(...)`은 샤딩된 transformer에서도 동일한 방식으로 작동합니다.
	- [enable_model_cpu_offload()](/docs/diffusers/main/ko/api/pipelines/overview#diffusers.DiffusionPipeline.enable_model_cpu_offload)를 사용하면 오프로드 훅을 `torch.compile`로 추적할 수 없습니다. 대신 `pipe.transformer.compile_repeated_blocks(backend="tpu", fullgraph=True, dynamic=False)`를 사용하여 transformer의 반복 블록만 컴파일하세요.

	## Tensor 병렬성

	하나의 칩에 너무 큰 모델을 여러 칩에 걸쳐 샤딩합니다. FLUX.2-dev의 텍스트 인코더(~48GB)와 트랜스포머(~64GB)는 각각 단일 칩을 초과하므로, 아래 예시에서는 둘 다 샤딩합니다:

	- `parallel_config` 인수로 전달되는 [from_pretrained()](/docs/diffusers/main/ko/api/models/overview#diffusers.ModelMixin.from_pretrained)에 [TensorParallelConfig](/docs/diffusers/main/ko/api/parallel#diffusers.TensorParallelConfig)를 갖춘 transformer입니다. 각 rank는 샤딩된 가중치의 자기 자신의 슬라이스만 읽으므로 전체 모델은 절대 메모리에 로드되지 않습니다. 일반적인 TP 세부 사항(`_tp_plan`, colwise/rowwise)은 [Tensor parallelism](../training/distributed_inference#tensor-parallelism) 가이드를 참조하세요.
	- Transformers 자체의 [tensor parallelism](https://huggingface.co/docs/transformers/perf_infer_gpu_multi)를 사용하여 텍스트 인코더를 구성하고, [DistributedConfig](https://huggingface.co/docs/transformers/main/en/expert_parallelism#transformers.DistributedConfig)와 동일한 mesh를 전달합니다.

	TPU에서는 `backend="tpu_dist"`로 프로세스 그룹을 초기화하고 `DeviceMesh("tpu", ...)`로 메시를 구성합니다.

	```python
	import torch
	import torch.distributed as dist
	from torch.distributed.device_mesh import DeviceMesh
	from transformers import DistributedConfig, Mistral3ForConditionalGeneration

	from diffusers import Flux2Pipeline, Flux2Transformer2DModel, TensorParallelConfig

	dist.init_process_group(backend="tpu_dist")
	mesh = DeviceMesh("tpu", list(range(dist.get_world_size())))

	repo_id = "black-forest-labs/FLUX.2-dev"
	text_encoder = Mistral3ForConditionalGeneration.from_pretrained(
	repo_id,
	subfolder="text_encoder",
	dtype=torch.bfloat16,
	distributed_config=DistributedConfig(tp_plan="auto"),
	device_mesh=mesh,
	)
	transformer = Flux2Transformer2DModel.from_pretrained(
	repo_id, subfolder="transformer", dtype=torch.bfloat16, parallel_config=TensorParallelConfig(mesh=mesh)
	)
	pipe = Flux2Pipeline.from_pretrained(
	repo_id, text_encoder=text_encoder, transformer=transformer, dtype=torch.bfloat16
	)
	pipe.vae.to("tpu")

	image = pipe(
	prompt="a golden retriever surfing a wave, photorealistic",
	num_inference_steps=28,
	generator=torch.Generator("cpu").manual_seed(0),
	).images[0]
	if dist.get_rank() == 0:
	image.save("output.png")
	```

	칩당 하나의 프로세스를 시작합니다. 호스트의 모든 TPU 칩 수를 사용하려면 `--nproc_per_node`를 설정합니다.

	```bash
	eval $(python -m torch_tpu._internal.distributed.launchers.singlehost_wrapper \| sed 's/^/export /')
	torchrun --nproc_per_node=8 flux2_tp.py
	```

	> [!WARNING]
	> 위 시작 명령은 TorchTPU 내부 모듈(`torch_tpu._internal`)을 사용하며, 이는 안정화된 API가 아니므로 향후 TorchTPU 릴리스에서 변경될 수 있습니다. 멀티 칩 작업을 시작하는 최신 방법은 [TorchTPU 저장소](https://github.com/google-pytorch/torch_tpu/)에서 확인하세요.