mirror of
https://github.com/huggingface/diffusers.git
synced 2026-10-24 20:59:27 +03:00
1d5d056ecf
* feat: add torchtpu * feat:draft TorchTPU support * fix: wan overflow issue + compile mode error on sdxl * doc: enhance with TorchTPU doc * doc: enhance with TorchTPU doc * style: remove unused imports flagged by ruff Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com> * docs: remove Debug Eager and Fused Eager sections from tpu.md Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com> * feat: add TP for TPU * style: fix import sorting in TPU test scripts Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com> * style: ruff format TPU test scripts Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com> * fix: style * fix: test for native 4 devices * fix: propagate TPU device fixes to Flux/Flux2/Wan-family copies; register tpu.md; drop redundant execution_device check - Propagate the text_encoder.device-based fix (introduced for TPU CPU-offload support) from FluxPipeline/Flux2KleinPipeline/WanPipeline into their `# Copied from` copies (flux/*, flux2_klein_inpaint, visualcloze, anyflow, chronoedit, lucy_edit, skyreels_v2/*). SDXL-family copies of StableDiffusionXLPipeline.encode_prompt are intentionally left untouched; they'll be handled in a follow-up PR that fixes device placement for every pipeline component (not just text encoders). - Register docs/source/en/optimization/tpu.md in _toctree.yml (was breaking the docs build: "not present in the table of contents"). - Remove the redundant "prefer non-CPU, non-meta component" loop from DiffusionPipeline._execution_device: PR #14383 already fixed this in DiffusionPipeline.device, which _execution_device falls back to. Verified on TPU hardware that _execution_device still resolves correctly for a split-placement pipeline after the removal. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com> * tests: cleanup * fix: fix compile mode * Update docs/source/en/optimization/tpu.md Co-authored-by: Steven Liu <59462357+stevhliu@users.noreply.github.com> * Update docs/source/en/optimization/tpu.md Co-authored-by: Steven Liu <59462357+stevhliu@users.noreply.github.com> * Update docs/source/en/optimization/tpu.md Co-authored-by: Steven Liu <59462357+stevhliu@users.noreply.github.com> * Update docs/source/en/optimization/tpu.md Co-authored-by: Steven Liu <59462357+stevhliu@users.noreply.github.com> * Update docs/source/en/optimization/tpu.md Co-authored-by: Steven Liu <59462357+stevhliu@users.noreply.github.com> * Update docs/source/en/optimization/tpu.md Co-authored-by: Steven Liu <59462357+stevhliu@users.noreply.github.com> * Update docs/source/en/optimization/tpu.md Co-authored-by: Steven Liu <59462357+stevhliu@users.noreply.github.com> * Update docs/source/en/optimization/tpu.md Co-authored-by: Steven Liu <59462357+stevhliu@users.noreply.github.com> * Update docs/source/en/optimization/tpu.md Co-authored-by: Steven Liu <59462357+stevhliu@users.noreply.github.com> * Update docs/source/en/optimization/tpu.md Co-authored-by: Steven Liu <59462357+stevhliu@users.noreply.github.com> * Update docs/source/en/optimization/tpu.md Co-authored-by: Steven Liu <59462357+stevhliu@users.noreply.github.com> * test: remove flux2 e2e test * doc: apply suggestions * doc: apply suggestions * review: remove monkey patch * review: revert neuron-specific changes in the tests * review: apply suggestions * Update docs/source/en/optimization/tpu.md Co-authored-by: Steven Liu <59462357+stevhliu@users.noreply.github.com> * doc: add tpu to tp doc * review: remove unnecessary for tpu * test: assert every _tp_plan parameter is sharded after a TP load Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * removal: delete redundant tp shard helpers * removal: drop TPU workarounds no longer needed after recent main changes * doc: keep the FLUX.2-dev text encoder off a single TPU chip in the TP example * doc: encode the prompt on CPU in the TPU TP example so it runs end to end * test: tighten TPU TP tolerance and simplify the TPU TP worker * doc: use enable_model_cpu_offload() in the TPU eager example * doc: shard both FLUX.2-dev text encoder and transformer in the TPU TP example * doc: run the text encoders on TPU in the compiled example * test: drop the redundant TPU sync in the TP worker * test: trim comments in the TPU TP tests * test: run the TPU TP test with mp.spawn like the CUDA one * test: pin the TPU TP test to 4 chips * Update docs/source/en/optimization/tpu.md Co-authored-by: Steven Liu <59462357+stevhliu@users.noreply.github.com> * doc: use DistributedConfig and dtype, and drop the torch_tpu import in the TPU examples * test: skip the TPU TP test for models without a _tp_plan * doc: correct the warning of recompilation * fix: clone the SD3 timestep so torch.compile doesn't recompile every step, thus we can use it for the tpu.md compile mode snippet * doc: update advice for compile mode when one chip doesn't fit * test: pass the TPU TP tolerances as test arguments * review: apply suggestions to the doc * review: view compilation patch for tpu-only --------- Co-authored-by: Claude Sonnet 4.6 <noreply@anthropic.com> Co-authored-by: JingyaHuang <huang_jingya@outlok.com> Co-authored-by: Steven Liu <59462357+stevhliu@users.noreply.github.com> Co-authored-by: Sayak Paul <spsayakpaul@gmail.com>
116 lines
3.4 KiB
Python
116 lines
3.4 KiB
Python
from .attention import AttentionBackendTesterMixin, AttentionTesterMixin
|
|
from .cache import (
|
|
CacheTesterMixin,
|
|
FasterCacheConfigMixin,
|
|
FasterCacheTesterMixin,
|
|
FirstBlockCacheConfigMixin,
|
|
FirstBlockCacheTesterMixin,
|
|
MagCacheConfigMixin,
|
|
MagCacheTesterMixin,
|
|
PyramidAttentionBroadcastConfigMixin,
|
|
PyramidAttentionBroadcastTesterMixin,
|
|
SeaCacheConfigMixin,
|
|
SeaCacheTesterMixin,
|
|
TaylorSeerCacheConfigMixin,
|
|
TaylorSeerCacheTesterMixin,
|
|
)
|
|
from .common import BaseModelTesterConfig, ModelTesterMixin
|
|
from .compile import TorchCompileTesterMixin
|
|
from .ip_adapter import IPAdapterTesterMixin
|
|
from .lokr import LoKrTesterMixin
|
|
from .lora import LoraHotSwappingForModelTesterMixin, LoraTesterMixin
|
|
from .memory import CPUOffloadTesterMixin, GroupOffloadTesterMixin, LayerwiseCastingTesterMixin, MemoryTesterMixin
|
|
from .parallelism import (
|
|
ContextParallelAttentionBackendsTesterMixin,
|
|
ContextParallelTesterMixin,
|
|
TensorParallelTesterMixin,
|
|
TensorParallelTPUTesterMixin,
|
|
)
|
|
from .quantization import (
|
|
AutoRoundCompileTesterMixin,
|
|
AutoRoundConfigMixin,
|
|
AutoRoundTesterMixin,
|
|
BitsAndBytesCompileTesterMixin,
|
|
BitsAndBytesConfigMixin,
|
|
BitsAndBytesTesterMixin,
|
|
GGUFCompileTesterMixin,
|
|
GGUFConfigMixin,
|
|
GGUFTesterMixin,
|
|
ModelOptCompileTesterMixin,
|
|
ModelOptConfigMixin,
|
|
ModelOptTesterMixin,
|
|
NunchakuLiteCompileTesterMixin,
|
|
NunchakuLiteConfigMixin,
|
|
NunchakuLiteTesterMixin,
|
|
QuantizationCompileTesterMixin,
|
|
QuantizationTesterMixin,
|
|
SDNQCompileTesterMixin,
|
|
SDNQConfigMixin,
|
|
SDNQTesterMixin,
|
|
TorchAoCompileTesterMixin,
|
|
TorchAoConfigMixin,
|
|
TorchAoTesterMixin,
|
|
)
|
|
from .single_file import SingleFileTesterMixin
|
|
from .training import TrainingTesterMixin
|
|
from .utils import run_nondeterministic
|
|
|
|
|
|
__all__ = [
|
|
"AttentionBackendTesterMixin",
|
|
"AttentionTesterMixin",
|
|
"AutoRoundConfigMixin",
|
|
"AutoRoundTesterMixin",
|
|
"BaseModelTesterConfig",
|
|
"BitsAndBytesCompileTesterMixin",
|
|
"BitsAndBytesConfigMixin",
|
|
"BitsAndBytesTesterMixin",
|
|
"CacheTesterMixin",
|
|
"ContextParallelTesterMixin",
|
|
"ContextParallelAttentionBackendsTesterMixin",
|
|
"TensorParallelTesterMixin",
|
|
"TensorParallelTPUTesterMixin",
|
|
"CPUOffloadTesterMixin",
|
|
"FasterCacheConfigMixin",
|
|
"FasterCacheTesterMixin",
|
|
"FirstBlockCacheConfigMixin",
|
|
"FirstBlockCacheTesterMixin",
|
|
"MagCacheConfigMixin",
|
|
"MagCacheTesterMixin",
|
|
"GGUFCompileTesterMixin",
|
|
"GGUFConfigMixin",
|
|
"GGUFTesterMixin",
|
|
"GroupOffloadTesterMixin",
|
|
"IPAdapterTesterMixin",
|
|
"LayerwiseCastingTesterMixin",
|
|
"LoKrTesterMixin",
|
|
"LoraHotSwappingForModelTesterMixin",
|
|
"LoraTesterMixin",
|
|
"MemoryTesterMixin",
|
|
"ModelOptCompileTesterMixin",
|
|
"ModelOptConfigMixin",
|
|
"ModelOptTesterMixin",
|
|
"ModelTesterMixin",
|
|
"NunchakuLiteCompileTesterMixin",
|
|
"NunchakuLiteConfigMixin",
|
|
"NunchakuLiteTesterMixin",
|
|
"PyramidAttentionBroadcastConfigMixin",
|
|
"PyramidAttentionBroadcastTesterMixin",
|
|
"SeaCacheConfigMixin",
|
|
"SeaCacheTesterMixin",
|
|
"TaylorSeerCacheConfigMixin",
|
|
"TaylorSeerCacheTesterMixin",
|
|
"QuantizationCompileTesterMixin",
|
|
"QuantizationTesterMixin",
|
|
"run_nondeterministic",
|
|
"SDNQCompileTesterMixin",
|
|
"SDNQConfigMixin",
|
|
"SDNQTesterMixin",
|
|
"SingleFileTesterMixin",
|
|
"TorchAoCompileTesterMixin",
|
|
"TorchAoConfigMixin",
|
|
"TorchAoTesterMixin",
|
|
"TorchCompileTesterMixin",
|
|
"TrainingTesterMixin",
|
|
]
|