[pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

Author: pre-commit-ci[bot]

transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py

@@ -555,7 +555,7 @@ def set_device_mesh(
         """
         Set DeviceMesh(s) used for sharding weights and convert main weights into DTensor
         depending on the TransformerEngine class to support FSDP-TP sharding with FSDP2.
-        
+
         TransformerEngine manages tensor parallel mechanics, while DTensor offers seamless
         integration with Torch DCP checkpointing. This method should only be invoked when
         using DTensor parameters, e.g. when using FSDP2 or DCP.
@@ -579,19 +579,20 @@ def set_device_mesh(
             # Validate TP DeviceMesh / Group. Must be consistent with tp_size.
             assert (
                 tp_mesh.ndim == 1 and self.tp_size == tp_mesh.size(),
-                f"TransformerEngine {self.__class__.__name__} TP init size ({self.tp_size}) "
-                f"does not match the size of the provided TP DeviceMesh ({tp_mesh.size()})."
+                (
+                    f"TransformerEngine {self.__class__.__name__} TP init size ({self.tp_size}) "
+                    f"does not match the size of the provided TP DeviceMesh ({tp_mesh.size()})."
+                ),
             )
             # Set the tensor parallel group from the mesh.
             self.set_tensor_parallel_group(tp_mesh.get_group())
 
             # Construct TP-sharded DTensors.
             if self.softmax_type == "learnable":
                 from torch.distributed.tensor.placement_types import Shard
+
                 self.softmax_offset = _convert_param_to_dtensor_param(
-                    self.softmax_offset,
-                    tp_mesh,
-                    placements=(Shard(dim=0),)
+                    self.softmax_offset, tp_mesh, placements=(Shard(dim=0),)
                 )
 
     def set_context_parallel_group(
@@ -862,7 +863,7 @@ def set_meta_tensor(self, fwd: bool, recipe: Union[Recipe, List[Recipe]]) -> Non
         self.quantizers[fp8_meta_tensor_key] = []
         for recipe_state in recipe_states:
             self.quantizers[fp8_meta_tensor_key].extend(recipe_state.make_quantizers())
-    
+
     def _get_softmax_offset(self) -> torch.Tensor:
         """Get the softmax offset."""
         softmax_offset = (

transformer_engine/pytorch/attention/multi_head_attention.py

@@ -611,7 +611,7 @@ def set_device_mesh(
         """
         Set DeviceMesh(s) used for sharding weights and convert main weights into DTensor
         depending on the TransformerEngine class to support FSDP-TP sharding with FSDP2.
-        
+
         TransformerEngine manages tensor parallel mechanics, while DTensor offers seamless
         integration with Torch DCP checkpointing. This method should only be invoked when
         using DTensor parameters, e.g. when using FSDP2 or DCP.
@@ -643,8 +643,10 @@ def set_device_mesh(
             # Validate TP DeviceMesh / Group. Must be consistent with tp_size.
             assert (
                 tp_mesh.ndim == 1 and self.tp_size == tp_mesh.size(),
-                f"TransformerEngine {self.__class__.__name__} TP init size ({self.tp_size}) "
-                f"does not match the size of the provided TP DeviceMesh ({tp_mesh.size()})."
+                (
+                    f"TransformerEngine {self.__class__.__name__} TP init size ({self.tp_size}) "
+                    f"does not match the size of the provided TP DeviceMesh ({tp_mesh.size()})."
+                ),
             )
             # Set the tensor parallel group from the mesh.
             self.set_tensor_parallel_group(tp_mesh.get_group())

transformer_engine/pytorch/distributed.py

@@ -1943,6 +1943,7 @@ def _convert_param_to_dtensor_param(
 ):
     """Convert the parameter into a DTensor."""
     from torch.distributed.tensor import DTensor
+
     # If the parameter is already a DTensor, extract local Tensor.
     # We overwrite the original DTensor's distributed configuration.
     param_tensor = param

transformer_engine/pytorch/module/base.py

@@ -1300,7 +1300,7 @@ def reset_parameters(self, defer_init: Optional[bool] = False) -> None:
             #       re-applying the nn.Parameter() wrap is a no-op when the input is already
             #       a parameter so we always re-apply it just for extra safety.
             if is_dtensor:
-                # Recreate the DTensor from the Parameter, inheriting 
+                # Recreate the DTensor from the Parameter, inheriting
                 # all attributes originally set on the Parameter.
                 dtensor_param = _convert_param_to_dtensor_param(
                     param,

transformer_engine/pytorch/module/grouped_linear.py

@@ -828,7 +828,7 @@ def set_device_mesh(
         """
         Set DeviceMesh(s) used for sharding weights and convert main weights into DTensor
         depending on the TransformerEngine class to support FSDP-TP sharding with FSDP2.
-        
+
         TransformerEngine manages tensor parallel mechanics, while DTensor offers seamless
         integration with Torch DCP checkpointing. This method should only be invoked when
         using DTensor parameters, e.g. when using FSDP2 or DCP.
@@ -860,36 +860,39 @@ def set_device_mesh(
             # Validate TP DeviceMesh / Group. Must be consistent with tp_size.
             assert (
                 tp_mesh.ndim == 1 and self.tp_size == tp_mesh.size(),
-                f"TransformerEngine {self.__class__.__name__} TP init size ({self.tp_size}) "
-                f"does not match the size of the provided TP DeviceMesh ({tp_mesh.size()})."
+                (
+                    f"TransformerEngine {self.__class__.__name__} TP init size ({self.tp_size}) "
+                    f"does not match the size of the provided TP DeviceMesh ({tp_mesh.size()})."
+                ),
             )
             # Set the tensor parallel group from the mesh.
             self.set_tensor_parallel_group(tp_mesh.get_group())
 
             # Construct TP-sharded DTensors.
             from torch.distributed.tensor.placement_types import Replicate, Shard
+
             for weight in self.weight_names:
                 param = getattr(self, weight)
                 placements = (Replicate(),)
                 if self.parallel_mode == "column":
                     placements = (Shard(dim=0),)
                 elif self.parallel_mode == "row":
                     placements = (Shard(dim=1),)
-                setattr(self, weight, _convert_param_to_dtensor_param(
-                    param,
-                    tp_mesh,
-                    placements=placements
-                ))
+                setattr(
+                    self,
+                    weight,
+                    _convert_param_to_dtensor_param(param, tp_mesh, placements=placements),
+                )
             for bias in self.bias_names:
                 param = getattr(self, bias)
                 placements = (Replicate(),)
                 if self.parallel_mode == "column":
                     placements = (Shard(dim=0),)
-                setattr(self, bias, _convert_param_to_dtensor_param(
-                    param,
-                    tp_mesh,
-                    placements=placements
-                ))
+                setattr(
+                    self,
+                    bias,
+                    _convert_param_to_dtensor_param(param, tp_mesh, placements=placements),
+                )
 
         # Set amax_reduction_group to the FSDP and/or TP sharding mesh
         # for per-tensor scaling recipes. Parameters must be registered.
@@ -1105,7 +1108,7 @@ def _get_weight_tensors(self) -> List[Union[torch.Tensor, QuantizedTensorStorage
                 for w in weight_tensors
             ]
         return weight_tensors
-    
+
     def _get_bias_tensors(self) -> List[torch.Tensor]:
         """Get the bias tensors of the module."""
         bias_tensors = []

transformer_engine/pytorch/module/layernorm.py

@@ -147,7 +147,7 @@ def set_device_mesh(
         """
         Set DeviceMesh(s) used for sharding weights and convert main weights into DTensor
         depending on the TransformerEngine class to support FSDP-TP sharding with FSDP2.
-        
+
         TransformerEngine manages tensor parallel mechanics, while DTensor offers seamless
         integration with Torch DCP checkpointing. This method should only be invoked when
         using DTensor parameters, e.g. when using FSDP2 or DCP.
@@ -173,15 +173,12 @@ def set_device_mesh(
             # with DTensor parameters in TP layers to support DTensor operations.
             from transformer_engine.pytorch.distributed import _convert_param_to_dtensor_param
             from torch.distributed.tensor.placement_types import Replicate
+
             self.weight = _convert_param_to_dtensor_param(
-                self.weight,
-                tp_mesh,
-                placements=(Replicate(),)
+                self.weight, tp_mesh, placements=(Replicate(),)
             )
             self.bias = _convert_param_to_dtensor_param(
-                self.bias,
-                tp_mesh,
-                placements=(Replicate(),)
+                self.bias, tp_mesh, placements=(Replicate(),)
             )
 
     @property

transformer_engine/pytorch/module/layernorm_linear.py

@@ -1462,7 +1462,7 @@ def set_device_mesh(
         """
         Set DeviceMesh(s) used for sharding weights and convert main weights into DTensor
         depending on the TransformerEngine class to support FSDP-TP sharding with FSDP2.
-        
+
         TransformerEngine manages tensor parallel mechanics, while DTensor offers seamless
         integration with Torch DCP checkpointing. This method should only be invoked when
         using DTensor parameters, e.g. when using FSDP2 or DCP.
@@ -1494,14 +1494,17 @@ def set_device_mesh(
             # Validate TP DeviceMesh / Group. Must be consistent with tp_size.
             assert (
                 tp_mesh.ndim == 1 and self.tp_size == tp_mesh.size(),
-                f"TransformerEngine {self.__class__.__name__} TP init size ({self.tp_size}) "
-                f"does not match the size of the provided TP DeviceMesh ({tp_mesh.size()})."
+                (
+                    f"TransformerEngine {self.__class__.__name__} TP init size ({self.tp_size}) "
+                    f"does not match the size of the provided TP DeviceMesh ({tp_mesh.size()})."
+                ),
             )
             # Set the tensor parallel group from the mesh.
             self.set_tensor_parallel_group(tp_mesh.get_group())
 
             # Construct TP-sharded DTensors.
             from torch.distributed.tensor.placement_types import Replicate, Shard
+
             # Linear
             for weight in self.weight_names:
                 param = getattr(self, weight)
@@ -1510,35 +1513,31 @@ def set_device_mesh(
                     placements = (Shard(dim=0),)
                 elif self.parallel_mode == "row":
                     placements = (Shard(dim=1),)
-                setattr(self, weight, _convert_param_to_dtensor_param(
-                    param,
-                    tp_mesh,
-                    placements=placements
-                ))
+                setattr(
+                    self,
+                    weight,
+                    _convert_param_to_dtensor_param(param, tp_mesh, placements=placements),
+                )
             for bias in self.bias_names:
                 param = getattr(self, bias)
                 placements = (Replicate(),)
                 if self.parallel_mode == "column":
                     placements = (Shard(dim=0),)
-                setattr(self, bias, _convert_param_to_dtensor_param(
-                    param,
-                    tp_mesh,
-                    placements=placements
-                ))
+                setattr(
+                    self,
+                    bias,
+                    _convert_param_to_dtensor_param(param, tp_mesh, placements=placements),
+                )
             # LayerNorm
             placements = (Replicate(),)
             if self.parallel_mode == "row":
                 placements = (Shard(dim=0),)
             self.layer_norm_weight = _convert_param_to_dtensor_param(
-                self.layer_norm_weight,
-                tp_mesh,
-                placements=placements
+                self.layer_norm_weight, tp_mesh, placements=placements
             )
             if self.layer_norm_bias is not None:
                 self.layer_norm_bias = _convert_param_to_dtensor_param(
-                    self.layer_norm_bias,
-                    tp_mesh,
-                    placements=placements
+                    self.layer_norm_bias, tp_mesh, placements=placements
                 )
 
         # Set amax_reduction_group to the FSDP and/or TP sharding mesh
@@ -1941,7 +1940,7 @@ def _get_weight_quantizers(self) -> List[Quantizer]:
         weight_quantizer.internal = True
         return [weight_quantizer]
 
-    def _set_tensor_parallel_attributes(self, defer_init = False) -> None:
+    def _set_tensor_parallel_attributes(self, defer_init=False) -> None:
         """Set tensor and sequence parallelism attributes."""
         if not defer_init:
             # Set parallelism attributes for layer norm parameters

transformer_engine/pytorch/module/layernorm_mlp.py

@@ -2028,7 +2028,7 @@ def set_device_mesh(
         """
         Set DeviceMesh(s) used for sharding weights and convert main weights into DTensor
         depending on the TransformerEngine class to support FSDP-TP sharding with FSDP2.
-        
+
         TransformerEngine manages tensor parallel mechanics, while DTensor offers seamless
         integration with Torch DCP checkpointing. This method should only be invoked when
         using DTensor parameters, e.g. when using FSDP2 or DCP.
@@ -2060,47 +2060,38 @@ def set_device_mesh(
             # Validate TP DeviceMesh / Group. Must be consistent with tp_size.
             assert (
                 tp_mesh.ndim == 1 and self.tp_size == tp_mesh.size(),
-                f"TransformerEngine {self.__class__.__name__} TP init size ({self.tp_size}) "
-                f"does not match the size of the provided TP DeviceMesh ({tp_mesh.size()})."
+                (
+                    f"TransformerEngine {self.__class__.__name__} TP init size ({self.tp_size}) "
+                    f"does not match the size of the provided TP DeviceMesh ({tp_mesh.size()})."
+                ),
             )
             # Set the tensor parallel group from the mesh.
             self.set_tensor_parallel_group(tp_mesh.get_group())
 
             # Construct TP-sharded DTensors.
             from torch.distributed.tensor.placement_types import Replicate, Shard
+
             # FC1 -> Column-Parallel -> Shard(dim=0)
             self.fc1_weight = _convert_param_to_dtensor_param(
-                self.fc1_weight,
-                tp_mesh,
-                placements=(Shard(dim=0),)
+                self.fc1_weight, tp_mesh, placements=(Shard(dim=0),)
             )
             self.fc1_bias = _convert_param_to_dtensor_param(
-                self.fc1_bias,
-                tp_mesh,
-                placements=(Shard(dim=0),)
+                self.fc1_bias, tp_mesh, placements=(Shard(dim=0),)
             )
             # FC2 Weight -> Row-Parallel -> Shard(dim=1)
             self.fc2_weight = _convert_param_to_dtensor_param(
-                self.fc2_weight,
-                tp_mesh,
-                placements=(Shard(dim=1),)
+                self.fc2_weight, tp_mesh, placements=(Shard(dim=1),)
             )
             # LN & FC2 Bias -> Replicate()
             self.fc2_bias = _convert_param_to_dtensor_param(
-                self.fc2_bias,
-                tp_mesh,
-                placements=(Replicate(),)
+                self.fc2_bias, tp_mesh, placements=(Replicate(),)
             )
             self.layer_norm_weight = _convert_param_to_dtensor_param(
-                self.layer_norm_weight,
-                tp_mesh,
-                placements=(Replicate(),)
+                self.layer_norm_weight, tp_mesh, placements=(Replicate(),)
             )
             if self.layer_norm_bias is not None:
                 self.layer_norm_bias = _convert_param_to_dtensor_param(
-                    self.layer_norm_bias,
-                    tp_mesh,
-                    placements=(Replicate(),)
+                    self.layer_norm_bias, tp_mesh, placements=(Replicate(),)
                 )
 
         # Set amax_reduction_group to the FSDP and/or TP sharding mesh
@@ -2583,7 +2574,7 @@ def _get_weight_tensors(self) -> List[Union[torch.Tensor, QuantizedTensorStorage
         if isinstance(fc2_weight, DTensor):
             fc2_weight = fc2_weight.to_local()
         return [fc1_weight, fc2_weight]
-    
+
     def _get_bias_tensors(self) -> List[torch.Tensor]:
         """Get the bias tensors of the module."""
         fc1_bias = self.fc1_bias if self.use_bias else None
@@ -2593,7 +2584,7 @@ def _get_bias_tensors(self) -> List[torch.Tensor]:
         if isinstance(fc2_bias, DTensor):
             fc2_bias = fc2_bias.to_local()
         return [fc1_bias, fc2_bias]
-    
+
     def _get_layernorm_weight_and_bias(self) -> List[Optional[torch.Tensor]]:
         """Get the weight and bias of the layer norm."""
         ln_weight = self.layer_norm_weight
@@ -2651,7 +2642,7 @@ def backward_dw(self):
             del fc1_bias_grad
             self._trigger_wgrad_accumulation_and_reduce_hooks()
 
-    def _set_tensor_parallel_attributes(self, defer_init = False) -> None:
+    def _set_tensor_parallel_attributes(self, defer_init=False) -> None:
         """Set tensor and sequence parallelism attributes."""
         if not defer_init:
             # Set parallel attributes for layer norm parameters
@@ -2665,4 +2656,4 @@ def _set_tensor_parallel_attributes(self, defer_init = False) -> None:
             if self.use_bias:
                 set_tensor_model_parallel_attributes(self.fc1_bias, True, 0, 1)
                 if self.set_parallel_mode:
-                    setattr(self.fc2_bias, "sequence_parallel", self.sequence_parallel)
+                    setattr(self.fc2_bias, "sequence_parallel", self.sequence_parallel)