Repository: huggingface/transformers
Pull request: Fix Qwen3 tp plan with FP8

# What does this PR do?
This PR : 
- Enables the device context for the FP8 activation kernel (if disabled, it causes a triton memory error)
- Updates the tp plan for the qwen3 fp8 models (since they use finegrained fp8, the scales will follow the sharding strategy of the weights)

Changed files:
filename: src/transformers/integrations/finegrained_fp8.py
patch: @@ -333,12 +333,8 @@ def forward(self, input: torch.Tensor) -> torch.Tensor:
             return F.linear(input, self.weight, self.bias)
         else:
             # Context manager used to switch among the available cuda devices
-            # with torch.cuda.device(input.device):
-            qinput, scale = act_quant(input, self.block_size[1])
-            # Blocks the CPU until all CUDA operations on the specified device are complete. It is used to ensure that the results of the
-            # preceding operations are ready before proceeding
-            # torch.cuda.synchronize(device=self.weight.device)
             with torch.cuda.device(input.device):
+                qinput, scale = act_quant(input, self.block_size[1])
                 output = w8a8_block_fp8_matmul_triton(
                     qinput,
                     self.weight,
@@ -347,6 +343,8 @@ def forward(self, input: torch.Tensor) -> torch.Tensor:
                     self.block_size,
                     output_dtype=input.dtype,
                 )
+            # Blocks the CPU until all CUDA operations on the specified device are complete. It is used to ensure that the results of the
+            # preceding operations are ready before proceeding
             torch.cuda.synchronize()
             if self.bias is not None:
                 output = output + self.bias

filename: src/transformers/quantizers/quantizer_finegrained_fp8.py
patch: @@ -194,6 +194,31 @@ def update_missing_keys(self, model, missing_keys: List[str], prefix: str) -> Li
                         not_missing_keys.append(missing)
         return [k for k in missing_keys if k not in not_missing_keys]
 
+    def update_tp_plan(self, config):
+        if "Qwen3" in config.__class__.__name__:
+            text_plan = {
+                "layers.*.self_attn.q_proj.weight": "local_colwise",
+                "layers.*.self_attn.q_proj.weight_scale_inv": "local_colwise",
+                "layers.*.self_attn.k_proj.weight": "local_colwise",
+                "layers.*.self_attn.k_proj.weight_scale_inv": "local_colwise",
+                "layers.*.self_attn.v_proj.weight": "local_colwise",
+                "layers.*.self_attn.v_proj.weight_scale_inv": "local_colwise",
+                "layers.*.self_attn.o_proj.weight": "local_rowwise",
+                "layers.*.self_attn.o_proj.weight_scale_inv": "local_rowwise",
+                "layers.*.self_attn": "gather",
+                "layers.*.mlp.gate_proj.weight": "local_colwise",
+                "layers.*.mlp.gate_proj.weight_scale_inv": "local_colwise",
+                "layers.*.mlp.up_proj.weight": "local_colwise",
+                "layers.*.mlp.up_proj.weight_scale_inv": "local_colwise",
+                "layers.*.mlp.down_proj.weight": "local_rowwise",
+                "layers.*.mlp.down_proj.weight_scale_inv": "local_rowwise",
+                "layers.*.mlp": "gather",
+            }
+
+            config.base_model_tp_plan = text_plan
+
+            return config
+
     def is_serializable(self, safe_serialization=None):
         return True