Optimize LTX2 feed-forward tensor parallelism (#23221)

2026-04-21 16:29:23 +08:00
parent e22dfe8fc2
commit 0d69012ef8
1 changed files with 7 additions and 3 deletions
@@ -764,11 +764,15 @@ class LTX2FeedForward(nn.Module):
        inner_dim = int(dim * mult)

        self.proj_in = ColumnParallelLinear(
-            dim, inner_dim, bias=True, gather_output=True, quant_config=quant_config
+            dim, inner_dim, bias=True, gather_output=False, quant_config=quant_config
        )
        self.act = nn.GELU(approximate="tanh")
-        self.proj_out = ColumnParallelLinear(
-            inner_dim, dim_out, bias=True, gather_output=True, quant_config=quant_config
+        self.proj_out = RowParallelLinear(
+            inner_dim,
+            dim_out,
+            bias=True,
+            input_is_parallel=True,
+            quant_config=quant_config,
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor: