Buckets:
| # Google TPU | |
| Tensor Processing Units (TPU) — это ускорители машинного обучения от Google, доступные в виде Cloud TPU (v6e, v7x). Diffusers работает на них через [TorchTPU](https://github.com/google-pytorch/torch_tpu/), бэкенд PyTorch для TPU, с минимальными изменениями кода. | |
| Доступны три режима выполнения: | |
| | Режим | Постоянная | Как активировать | Примечания | | |
| |---|---|---|---| | |
| | Строгий режим eager (по умолчанию) | `EagerMode.DEFER_NEVER` | `pipe.to("tpu")` | Операции, отправляемые по одной, асинхронные | | |
| | Оптимизированный режим eager | `EagerMode.DEFER_AND_FUSE` | `TORCH_TPU_DEFER_AND_FUSE=1` | Операции откладываются и объединяются в более крупные графы; быстрее строгого режима eager, но время компиляции дольше | | |
| | Компиляция | — | `torch.compile(module, backend="tpu")` | AOT-компиляция с `TpuBackend` | | |
| Следуйте [инструкции по установке TorchTPU](https://github.com/google-pytorch/torch_tpu/). После установки `import torch` автоматически загружает его и регистрирует устройство `"tpu"`, поэтому единственным необходимым изменением является `pipe.to("tpu")`. Добавляйте `import torch_tpu` только в том случае, если вы отключили автоматическую загрузку бэкенда с помощью `TORCH_DEVICE_BACKEND_AUTOLOAD=0`. | |
| ## Режим Eager | |
| FLUX.1-schnell не помещается целиком на одном чипе v6e, поэтому используйте [enable_model_cpu_offload()](/docs/diffusers/main/ru/api/pipelines/overview#diffusers.DiffusionPipeline.enable_model_cpu_offload), чтобы перемещать каждую модель на TPU только во время её выполнения. Устройство `"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`. Первый вызов (прогрев) выполняется медленно, так как происходит компиляция; последующие вызовы с теми же формами повторно используют скомпилированный граф. | |
| > [!IMPORTANT] | |
| > TorchTPU требует **статических форм**, поэтому передайте `dynamic=False`. Новая `height` или `width` приводит к повторной компиляции для этой формы, но только один раз; уже встреченные формы переиспользуются. Изменение `num_inference_steps` не вызывает повторной компиляции. | |
| Если весь пайплайн помещается на один чип, перенесите его на TPU и скомпилируйте полный трансформер. Stable Diffusion 3.5 Medium (~15 ГБ в bf16) помещается на один чип 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(...)` работает так же и с шардированным трансформером. | |
| - При использовании [enable_model_cpu_offload()](/docs/diffusers/main/ru/api/pipelines/overview#diffusers.DiffusionPipeline.enable_model_cpu_offload) хуки offload не могут быть отслежены `torch.compile`. Компилируйте только повторяющиеся блоки трансформера с помощью `pipe.transformer.compile_repeated_blocks(backend="tpu", fullgraph=True, dynamic=False)`. | |
| ## Параллелизм тензоров | |
| Размещайте модели, слишком большие для одного чипа, на нескольких. Текстовый энкодер (~48 ГБ) и трансформер (~64 ГБ) FLUX.2-dev по отдельности превышают объём одного чипа, поэтому в примере ниже оба компонента распределяются по нескольким чипам: | |
| - трансформер с [TensorParallelConfig](/docs/diffusers/main/ru/api/parallel#diffusers.TensorParallelConfig), переданным в аргумент `parallel_config` метода [from_pretrained()](/docs/diffusers/main/ru/api/models/overview#diffusers.ModelMixin.from_pretrained). Каждый ранг читает только свой срез каждого шардированного веса, поэтому полная модель никогда не материализуется. Общие сведения о TP (`_tp_plan`, colwise/rowwise) см. в руководстве [Tensor parallelism](./tensor_parallelism). | |
| - текстовый энкодер с собственным [тензорным параллелизмом](https://huggingface.co/docs/transformers/perf_infer_gpu_multi) Transformers, передавая [DistributedConfig](https://huggingface.co/docs/transformers/main/en/expert_parallelism#transformers.DistributedConfig) и ту же сетку. | |
| На 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") | |
| ``` | |
| Запустите один процесс на чип. Установите `--nproc_per_node` в значение, равное общему числу чипов TPU на вашем хосте. | |
| ```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:
- 9.5 kB
- Xet hash:
- df5466e7cd5b1d63ef35a4f753b9213fa1b9e0dd921fb76febcfd1f91f123fab
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.