address Shreyas' comments, add lpips test

Signed-off-by: Brenden Elgarten <belgarten@nvidia.com>

Author: belgarten-nv

tensorrt_llm/_torch/modules/gated_mlp.py

@@ -71,6 +71,14 @@ def __init__(
                                                     mapping.tp_rank + 1)
         local_intermediate_size = local_intermediate_end - local_intermediate_start
 
+        self._uneven_tp_blocks_lora = (mapping.tp_size > 1
+                                       and self.intermediate_size %
+                                       mapping.tp_size != 0)
+
+        # gateup_shard_indices_mapping is the local offset and size for each sub-weight
+        # in this rank's concatenated (gate || up) buffer.
+        # override_tp_sharding is the absolute range of the global weight from which
+        # this rank pulls each sub-weight.
         gateup_shard_indices_mapping = {
             'gate': (0, local_intermediate_size),
             'up': (local_intermediate_size, local_intermediate_size),
@@ -297,6 +305,10 @@ def forward_lora(
     ) -> torch.Tensor:
         assert lora_params is not None
         assert self.layer_idx is not None, "layer_idx is required for lora"
+        if self._uneven_tp_blocks_lora:
+            raise NotImplementedError(
+                "LoRA is not supported with uneven TP for GatedMLP "
+                "(intermediate_size not divisible by tp_size).")
 
         h1 = self.gate_up_proj(x)
 

tensorrt_llm/_torch/modules/linear.py

@@ -106,6 +106,12 @@ def load_weight_shard(
     device: torch.device = torch.device('cpu'),
     return_slice_indices: bool = False,
 ) -> torch.Tensor:
+    """Legacy weight shard helper using ceil-divide sharding.
+
+    `Linear.load_shard` is preferred — it respects uneven-TP overrides, fused
+    QKV/gate-up sharding dicts, and quant-specific scale/packing semantics that
+    this function does not.
+    """
     # Skip device transfers on integrated GPUs to conserve shared memory
     if weight.device.type != device.type and is_device_integrated():
         # For integrated GPU systems (e.g., DGX Spark), CPU and GPU share limited physical memory.
@@ -2923,8 +2929,11 @@ def __init__(
             'cutlass', 'cublaslt', 'cuda_core'
         ]
 
-        assert self.tp_mode in (TensorParallelMode.ROW,
-                                TensorParallelMode.COLUMN, None)
+        if self.tp_mode not in (TensorParallelMode.ROW,
+                                TensorParallelMode.COLUMN, None):
+            raise ValueError(
+                f"Invalid tp_mode {self.tp_mode!r}; expected ROW, COLUMN, or None."
+            )
 
         # Init TP sharding either from override or auto generated
         _uneven_tp_unsupported = {QuantAlgo.NVFP4_ARC}
@@ -3006,8 +3015,10 @@ def _calc_shard(total, tp_size, rank):
     def _auto_tp_sharding(self, features, quant_config):
         """Auto-generate tp_sharding tuple based on quant alignment requirements.
 
-        For VANILLA mode only. Fused modes with non-divisible dims require
-        explicit override_tp_sharding from the model layer.
+        VANILLA mode only. Fused modes (FUSED_QKV, FUSED_GATE_UP) require explicit
+        override_tp_sharding because individual sub-weight sizes (Q vs K vs V; gate
+        vs up) are not knowable here — they aren't always equal (e.g. GQA), and
+        cross-rank consistency must be decided by the caller.
         """
         assert self.weights_loading_config.weight_mode == WeightMode.VANILLA, (
             f"_auto_tp_sharding only supports VANILLA mode, got "
@@ -3049,20 +3060,22 @@ def _calculate_local_features_helper(self, features):
             return end - start
 
     def calculate_local_in_features(self, in_features):
+        """Local input feature count after TP sharding (full size if not row-parallel)."""
         if self.tp_mode != TensorParallelMode.ROW:
             return in_features
 
         return self._calculate_local_features_helper(in_features)
 
     def calculate_local_out_features(self, out_features):
+        """Local output feature count after TP sharding (full size if not column-parallel)."""
         if self.tp_mode != TensorParallelMode.COLUMN:
             return out_features
 
         return self._calculate_local_features_helper(out_features)
 
     def load_shard(
         self,
-        weights: Dict,
+        weights: Union[Dict, torch.Tensor],
         label: Optional[str] = None,
         device: torch.device = torch.device('cpu'),
         name: Optional[str] = None,
@@ -3072,6 +3085,14 @@ def load_shard(
         # 2 are packed in each 8 bit element of the tensor
         elm_packing: int = 1,
     ) -> torch.Tensor:
+        """Slice a weight tensor for this rank's TP shard.
+
+        Unified entry point for module-aware weight loading: respects
+        `self.tp_sharding` (uneven TP overrides, fused QKV/gate-up dicts) and
+        quant-specific knobs (`scale_span`, `elm_packing`). Pass `weights` as a
+        dict with `label` to pick the entry, or as a bare tensor when there's
+        only one. Supersedes the free function `load_weight_shard`.
+        """
         if label:
             if label not in weights:
                 return None

tensorrt_llm/_torch/visual_gen/models/flux/joint_proj.py

@@ -63,7 +63,7 @@ def __init__(
         self.has_bias = bias
         self.attn_shard = attn_shard
 
-        assert attn_dim % self.tp_size == 0 or self.attn_shard, (
+        assert attn_dim % self.tp_size == 0 or self.attn_shard is not None, (
             "Explicit attention sharding required for uneven TP"
         )
 
@@ -197,6 +197,9 @@ def __init__(
             self.local_qkv_dim = q_dim + 2 * kv_dim
             self.local_mlp_dim = mlp_dim
         else:
+            assert override_qkv_sharding is not None, (
+                "override_qkv_sharding required when tp_size > 1"
+            )
 
             def range_size(r):
                 return r[1] - r[0]

tensorrt_llm/_torch/visual_gen/models/flux/transformer_flux.py

@@ -495,6 +495,21 @@ def __init__(
             attn_shard=(self.attn.local_q_dim_start, self.attn.local_q_dim_end),
         )
 
+        # MLP + Attn Output projection, requires special handling for TP
+        self.proj_out = FluxJointAttnMLPProj(
+            attn_dim=self.attn.q_dim,
+            mlp_dim=self.mlp_hidden_dim,
+            out_dim=dim,
+            bias=True,
+            dtype=dtype,
+            quant_config=quant_config,
+            skip_create_weights_in_init=skip_create_weights,
+            force_dynamic_quantization=force_dynamic_quant,
+            config=config,
+            # need explicit shard because we are aligned on head boundaries
+            attn_shard=(self.attn.local_q_dim_start, self.attn.local_q_dim_end),
+        )
+
     def forward(
         self,
         hidden_states: torch.Tensor,

tensorrt_llm/_torch/visual_gen/modules/attention.py

@@ -109,7 +109,7 @@ def __init__(
         self.q_dim = self.num_attention_heads * self.head_dim
         self.kv_dim = self.num_key_value_heads * self.head_dim
 
-        self._calculate_tp_parameters(ulysses_size if enable_ulysses else None)
+        self._calculate_tp_parameters(ulysses_size if enable_sequence_parallel else None)
         self._init_qkv_proj()
 
         # Structural eligibility for SEPARATE_QKV self-attn quantize dedup.
@@ -239,18 +239,6 @@ def __init__(
             async_ulysses=use_ulysses and async_ulysses,
         )
 
-    @staticmethod
-    def _qualified_module_name(
-        component_name: Optional[str],
-        module_name: Optional[str],
-    ) -> Optional[str]:
-        if module_name is None:
-            return None
-        if component_name is None:
-            return module_name
-        prefix = f"{component_name}."
-        return module_name if module_name.startswith(prefix) else f"{prefix}{module_name}"
-
     def _calculate_tp_parameters(self, ulysses_size: Optional[int]):
         assert self.num_attention_heads % self.num_key_value_heads == 0
         gqa_ratio = self.num_attention_heads // self.num_key_value_heads
@@ -259,19 +247,14 @@ def _calculate_tp_parameters(self, ulysses_size: Optional[int]):
             ulysses_size = 1
 
         assert self.num_key_value_heads % ulysses_size == 0
-        # Note: this is intentionally stronger than `num_kv_head >= ulysses_size * tp_size`
         assert self.num_key_value_heads // ulysses_size >= self.tp_size
 
-        def _calc_shard(full, size, rank):
-            full //= ulysses_size
-            shard = (full // size) * rank + min(full % size, rank)
-            return shard * ulysses_size
-
-        self.local_key_value_head_start = _calc_shard(
-            self.num_key_value_heads, self.tp_size, self.tp_rank
+        kv_heads_per_ulysses = self.num_key_value_heads // ulysses_size
+        self.local_key_value_head_start = (
+            Linear._calc_shard(kv_heads_per_ulysses, self.tp_size, self.tp_rank) * ulysses_size
         )
-        self.local_key_value_head_end = _calc_shard(
-            self.num_key_value_heads, self.tp_size, self.tp_rank + 1
+        self.local_key_value_head_end = (
+            Linear._calc_shard(kv_heads_per_ulysses, self.tp_size, self.tp_rank + 1) * ulysses_size
         )
         self.local_num_key_value_heads = (
             self.local_key_value_head_end - self.local_key_value_head_start

tests/integration/defs/examples/visual_gen/test_visual_gen_multi_gpu.py

@@ -59,6 +59,7 @@
 
 WAN22_LPIPS_TP_VARIANTS = [
     ("tp2", {"tp_size": 2}),
+    ("tp3", {"tp_size": 3}),
     ("cfg2_tp2", {"cfg_size": 2, "tp_size": 2}),
     ("tp2_ulysses2", {"tp_size": 2, "ulysses_size": 2}),
 ]