From 357f8738cdc050e3435e920ff3968a0bc101d4bb Mon Sep 17 00:00:00 2001 From: yzhautouskay Date: Fri, 2 Oct 2026 21:49:42 +0200 Subject: [PATCH] [Bugfix] Fix Cosmos3 FP8 tensor parallelism (#14921) Fix Cosmos3 FP8 tensor parallelism --- examples/cosmos3/cosmos_parallel.py | 10 +++++++++- 1 file changed, 9 insertions(+), 1 deletion(-) diff --git a/examples/cosmos3/cosmos_parallel.py b/examples/cosmos3/cosmos_parallel.py index 6175005ab0..663dd91838 100644 --- a/examples/cosmos3/cosmos_parallel.py +++ b/examples/cosmos3/cosmos_parallel.py @@ -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())