Repository: huggingface/transformers
Pull request: fix `Dtensor` and `tensor` mismatch

# Bug

`local_rowise` or `local_colwise` is calling `RowiseParallel(use_dtensor=False)` (resp. `ColwiseParallel(use_dtensor=False)`. Issue was first noticed in https://github.com/huggingface/transformers/pull/42356 , quoting 

> we would like to not have Dtensor logic in the modeling. For example, sinks are supposed to use local_rowwise (cf [main/src/transformers/models/gpt_oss/configuration_gpt_oss.py#L41](https://github.com/huggingface/transformers/blob/main/src/transformers/models/gpt_oss/configuration_gpt_oss.py?rgh-link-date=2025-11-24T12%3A36%3A48Z#L41)) which is supposed to not return a Dtensor (cf [main/src/transformers/integrations/tensor_parallel.py#L1171](https://github.com/huggingface/transformers/blob/main/src/transformers/integrations/tensor_parallel.py?rgh-link-date=2025-11-24T12%3A36%3A48Z#L1171)) but somehow doesnt work

# Fix

```python
def convert_and_load_state_dict_in_model(
	#...
	tp_layer = ALL_PARALLEL_STYLES[model.tp_plan[matched_tp_pattern]].__class__
```

was creating the bug because `.__class__` will reuse the class default `use_dtensor` value. Thus, overwritting the value we specifed in `local_rowise/colwise`.

The fix makes sure to properly use the `use_dtensor` value and thus no more `Dtensor` and `tensor` mismatch 

 



Changed files:
filename: src/transformers/integrations/tensor_parallel.py
patch: @@ -758,6 +758,15 @@ def partition_tensor(self, param, empty_param, param_type, param_casting_dtype,
         return nn.Parameter(parameter, requires_grad=parameter.is_floating_point())
 
 
+class LocalColwiseParallel(ColwiseParallel):
+    """
+    Colwise parallel with use_dtensor=False for local tensor operations.
+    """
+
+    def __init__(self, **kwargs):
+        super().__init__(use_dtensor=False, **kwargs)
+
+
 class RowwiseParallel(TensorParallelLayer):
     """
     Partition a compatible nn.Module in a row-wise fashion. Currently supports nn.Linear and nn.Embedding.
@@ -782,7 +791,7 @@ def __init__(
         input_layouts: Placement | None = None,
         output_layouts: Placement | None = None,
         use_local_output: bool = True,
-        use_dtensor=True,
+        use_dtensor: bool = True,
         **kwargs,
     ):
         super().__init__(**kwargs)
@@ -913,6 +922,24 @@ def partition_tensor(self, param, empty_param, param_type, param_casting_dtype,
         return nn.Parameter(parameter, requires_grad=parameter.is_floating_point())
 
 
+class LocalRowwiseParallel(RowwiseParallel):
+    """
+    Rowwise parallel with use_dtensor=False for local tensor operations.
+    """
+
+    def __init__(self, **kwargs):
+        super().__init__(use_dtensor=False, **kwargs)
+
+
+class LocalPackedRowwiseParallel(PackedRowwiseParallel):
+    """
+    Packed rowwise parallel with use_dtensor=False for local tensor operations.
+    """
+
+    def __init__(self, **kwargs):
+        super().__init__(use_dtensor=False, **kwargs)
+
+
 class SequenceParallel(TensorParallelLayer):
     """
     SequenceParallel replicates a compatible ``nn.Module`` parameters and runs the sharded computation with
@@ -1064,10 +1091,10 @@ class RouterParallel(TensorParallelLayer):
     Allows to reshape the router scores to support running expert parallel.
     """
 
-    def __init__(self, *args, **kwargs):
+    def __init__(self, use_dtensor: bool = False, *args, **kwargs):
         super().__init__(**kwargs)
         self.args = args
-        self.use_dtensor = False
+        self.use_dtensor = use_dtensor
 
     @staticmethod
     def _prepare_input_fn(input_layouts, desired_input_layouts, mod, inputs, device_mesh):
@@ -1118,7 +1145,7 @@ def _prepare_output_fn(output_layouts, use_local_output, mod, outputs, device_me
                 f"The number of experts must be divisible by number of ep_size: {mod.num_experts} % {ep_size} != 0"
             )
         num_local_experts = mod.num_experts // ep_size
-        router_scores, router_indices = outputs
+        router_logits, router_scores, router_indices = outputs
         router_scores = router_scores[:, ep_rank * num_local_experts : (ep_rank + 1) * num_local_experts]
         router_indices = router_indices.masked_fill((router_indices // num_local_experts) != ep_rank, -1)
         # As -1 % 1 is 0, we can only use mask fill when num_local_experts is 1
@@ -1129,7 +1156,7 @@ def _prepare_output_fn(output_layouts, use_local_output, mod, outputs, device_me
         router_indices = router_indices.masked_fill(
             router_indices == -1, num_local_experts
         )  # masking class for one hot
-        return router_scores, router_indices
+        return router_logits, router_scores, router_indices
 
     def shard_tensor(
         self,
@@ -1171,11 +1198,11 @@ class ParallelInterface(GeneralInterface):
             "rowwise": RowwiseParallel(),
             "colwise_rep": ColwiseParallel(output_layouts=Replicate()),
             "rowwise_rep": RowwiseParallel(input_layouts=Replicate()),
-            "local_colwise": ColwiseParallel(use_dtensor=False),
-            "local_rowwise": RowwiseParallel(use_dtensor=False),
+            "local_colwise": LocalColwiseParallel(),
+            "local_rowwise": LocalRowwiseParallel(),
             "local": IsolatedParallel(),
             "gather": GatherParallel(),
-            "local_packed_rowwise": PackedRowwiseParallel(use_dtensor=False),
+            "local_packed_rowwise": LocalPackedRowwiseParallel(),
             "sequence_parallel": SequenceParallel(),
             "replicate": ReplicateParallel(),
             "grouped_gemm": GroupedGemmParallel(),