AI/diffusers
1
0
Fork 0
mirror of https://github.com/huggingface/diffusers.git synced 2026-10-23 09:09:01 +03:00
Code Activity

[Bugfix] Fix Cosmos3 FP8 tensor parallelism (#14921)

Browse Source 
Fix Cosmos3 FP8 tensor parallelism
This commit is contained in:
yzhautouskay
2026-10-02 21:49:42 +02:00
committed by GitHub
parent 41566305ce
commit 357f8738cd
1 changed files with 9 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

10
examples/cosmos3/cosmos_parallel.py
View File

@@ -431,7 +431,15 @@ def enable_cosmos3_tensor_parallel(transformer, tp_mesh):
own processor) on a 2-D ``(tp, cp)`` mesh, or with ``enable_cosmos3_flash_attention``
for TP without CP.
"""
from torch.distributed.tensor.parallel import ColwiseParallel, RowwiseParallel, parallelize_module
from torch.distributed.tensor.parallel import ColwiseParallel, parallelize_module
from torch.distributed.tensor.parallel import RowwiseParallel as TorchRowwiseParallel
# Skip ModelOpt quantizer children to avoid:
# AttributeError: 'TensorQuantizer' object has no attribute 'weight'
class RowwiseParallel(TorchRowwiseParallel):
def _partition_linear_fn(self, name, module, device_mesh):
if isinstance(module, torch.nn.Linear):
return super()._partition_linear_fn(name, module, device_mesh)
tp = tp_mesh.size()
dev = torch.device("cuda", torch.cuda.current_device())