WAR the fact that it is not possible to set __setattr__ dynamically

Signed-off-by: Przemek Tredak <ptredak@nvidia.com>

Author: ptrendx

transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py

@@ -482,7 +482,7 @@ def remove_extra_states_check(self, incompatible_keys):  # pylint: disable=unuse
 
         self.register_load_state_dict_post_hook(remove_extra_states_check)
 
-        self.__setattr__ = self.default_setattr
+        self._default_setattr = self._warning_setattr
 
     def _load_from_state_dict(
         self, state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs

transformer_engine/pytorch/module/base.py

@@ -639,7 +639,7 @@ def fast_setattr(self, name: str, value: Any) -> None:
     def module_setattr(self, name: str, value: Any) -> None:
         super().__setattr__(name, value)
 
-    def default_setattr(self, name: str, value: Any) -> None:
+    def _warning_setattr(self, name: str, value: Any) -> None:
         warnings.warn(
             """The default implementation of torch.nn.Module introduces significant CPU overhead
             when setting attributes and is therefore not recommended. Please use the explicit calls
@@ -649,6 +649,12 @@ def default_setattr(self, name: str, value: Any) -> None:
         )
         self.module_setattr(name, value)
 
+    def _default_setattr(self, name: str, value: Any) -> None:
+        return self.module_setattr(name, value)
+
+    def __setattr__(self, name: str, value: Any) -> None:
+        return self._default_setattr(name, value)
+
     def adjust_amax_history_length(self, length: int, fwd: Optional[bool] = None) -> None:
         """
         Delayed scaling only.

transformer_engine/pytorch/module/grouped_linear.py

@@ -716,7 +716,7 @@ def __init__(
                     if name in (f"weight{i}", f"bias{i}"):
                         param.skip_backward_post_hook = True
 
-        self.__setattr__ = self.default_setattr
+        self._default_setattr = self._warning_setattr
 
     def set_meta_tensor(self, fwd: bool, recipe: Recipe) -> None:
         """Init scales and amaxes for fwd | bwd."""

transformer_engine/pytorch/module/layernorm_linear.py

@@ -1405,7 +1405,7 @@ def __init__(
                 if name in self.weight_names or name in self.bias_names:
                     param.skip_backward_post_hook = True
 
-        self.__setattr__ = self.default_setattr
+        self._default_setattr = self._warning_setattr
 
     def set_meta_tensor(self, fwd: bool, recipe: Recipe) -> None:
         """Init scales and amaxes for fwd | bwd."""

transformer_engine/pytorch/module/layernorm_mlp.py

@@ -1960,7 +1960,7 @@ def __init__(
         self.bwd_ln_sm_margin = int(os.getenv("NVTE_BWD_LAYERNORM_SM_MARGIN", "0"))
         self.inf_ln_sm_margin = int(os.getenv("NVTE_INF_LAYERNORM_SM_MARGIN", "0"))
 
-        self.__setattr__ = self.default_setattr
+        self._default_setattr = self._warning_setattr
 
     def set_meta_tensor(self, fwd: bool, recipe: Recipe) -> None:
         """Init scales and amaxes for fwd | bwd."""

transformer_engine/pytorch/module/linear.py

@@ -1309,7 +1309,7 @@ def __init__(
                 if name in self.weight_names or name in self.bias_names:
                     param.skip_backward_post_hook = True
 
-        self.__setattr__ = self.default_setattr
+        self._default_setattr = self._warning_setattr
 
     def set_meta_tensor(self, fwd: bool, recipe: Recipe) -> None:
         """Init scales and amaxes for fwd | bwd."""