Fix type-safety of torch.nn.Module instances

Summary:
X-link: ctrl-labs/src2#38809

As laid out in pytorch/pytorch#81462 (comment) the change in pytorch/pytorch#104321 was not necessary and largely destroys the type-safety of `torch.nn.Module` instances.

As far as I can see, the underlying issue of pytorch/pytorch#81462 in `torch.nn.parallel.DistributedDataParallel` has been fixed in the meantime by actually typing `register_comm_hook` correctly.

The proper solution to issues like pytorch/pytorch#81462 is to give the underlying field/method a proper type annotation, then there should be no need to go for a "type system disabling `__getattr__`".

(I'll probably be offline for a while, not able to react here...)

cc H-Huang awgu kwen2501 wanchaol fegin fduwjj wz337 wconstab d4l3k c-p-i-o voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx ipiszy yf225 chenyang78 kadeng muchulee8 ColinPeppler amjames desertfire chauhang aakhundov avikchaudhuri gmagogsfm zhxchen17 tugsbayasgalan angelayi suo ydwu4 XilunWu rec mrshenli pritamdamania87 zhaojuanmao satgera rohan-varma gqchen aazzolini osalpekar jiayisuse tianyu-l kiukchung lucasllc

X-link: pytorch/pytorch#115074

Differential Revision: D52890934

Pulled By: ezyang

Author: bluenote10

backends/cadence/aot/quantizer/patterns.py

@@ -75,6 +75,7 @@ def partition_types(self) -> List[OpOverload]:
     def get_anchors(
         self, gm: fx.GraphModule, fused_partition: List[fx.GraphModule]
     ) -> PartitionAnchors:
+        # pyre-fixme[29]: `Union[BoundMethod[typing.Callable(torch._C.TensorBase.__ge...
         addmm_node = fused_partition[0].nodes[-1]
 
         bias_qspec = DerivedQuantizationSpec(
@@ -107,6 +108,7 @@ def partition_types(self) -> List[OpOverload]:
     def get_anchors(
         self, gm: fx.GraphModule, fused_partition: List[fx.GraphModule]
     ) -> PartitionAnchors:
+        # pyre-fixme[29]: `Union[BoundMethod[typing.Callable(torch._C.TensorBase.__ge...
         bmm_node = fused_partition[0].nodes[-1]
 
         return PartitionAnchors(
@@ -127,6 +129,7 @@ def partition_types(self) -> List[OpOverload]:
     def get_anchors(
         self, gm: fx.GraphModule, fused_partition: List[fx.GraphModule]
     ) -> PartitionAnchors:
+        # pyre-fixme[29]: `Union[BoundMethod[typing.Callable(torch._C.TensorBase.__ge...
         conv1d_node = fused_partition[0].nodes[-1]
 
         bias_qspec = DerivedQuantizationSpec(
@@ -165,6 +168,7 @@ def partition_types(self) -> List[OpOverload]:
     def get_anchors(
         self, gm: fx.GraphModule, fused_partition: List[fx.GraphModule]
     ) -> PartitionAnchors:
+        # pyre-fixme[29]: `Union[BoundMethod[typing.Callable(torch._C.TensorBase.__ge...
         conv2d_node = fused_partition[0].nodes[-1]
 
         bias_qspec = DerivedQuantizationSpec(
@@ -203,6 +207,7 @@ def partition_types(self) -> List[OpOverload]:
     def get_anchors(
         self, gm: fx.GraphModule, fused_partition: List[fx.GraphModule]
     ) -> PartitionAnchors:
+        # pyre-fixme[29]: `Union[BoundMethod[typing.Callable(torch._C.TensorBase.__ge...
         layer_norm_node = fused_partition[0].nodes[-1]
 
         others = [(layer_norm_node, 1)]
@@ -237,6 +242,7 @@ def partition_types(self) -> List[OpOverload]:
     def get_anchors(
         self, gm: fx.GraphModule, fused_partition: List[fx.GraphModule]
     ) -> PartitionAnchors:
+        # pyre-fixme[29]: `Union[BoundMethod[typing.Callable(torch._C.TensorBase.__ge...
         linear_node = fused_partition[0].nodes[-1]
 
         bias_qspec = DerivedQuantizationSpec(
@@ -275,6 +281,7 @@ def partition_types(self) -> List[OpOverload]:
     def get_anchors(
         self, gm: fx.GraphModule, fused_partition: List[fx.GraphModule]
     ) -> PartitionAnchors:
+        # pyre-fixme[29]: `Union[BoundMethod[typing.Callable(torch._C.TensorBase.__ge...
         matmul_node = fused_partition[0].nodes[-1]
 
         return PartitionAnchors(
@@ -297,6 +304,7 @@ def partition_types(self) -> List[OpOverload]:
     def get_anchors(
         self, gm: fx.GraphModule, fused_partition: List[fx.GraphModule]
     ) -> PartitionAnchors:
+        # pyre-fixme[29]: `Union[BoundMethod[typing.Callable(torch._C.TensorBase.__ge...
         relu_node = fused_partition[0].nodes[-1]
 
         return PartitionAnchors(

backends/example/example_backend_delegate_passes/merge_to_dim_pass.py

@@ -73,7 +73,11 @@ def call(self, graph_module: torch.fx.GraphModule) -> PassResult:
                         )
                         out_to_dim_node_user.args = tuple(out_to_dim_node_user_new_args)
 
+                        # pyre-fixme[29]: `Union[torch._tensor.Tensor,
+                        #  torch.nn.modules.module.Module]` is not a function.
                         graph_module.erase_node(out_to_dim_node)
+                        # pyre-fixme[29]: `Union[torch._tensor.Tensor,
+                        #  torch.nn.modules.module.Module]` is not a function.
                         graph_module.erase_node(node)
             # TODO: Handle other merging rules, including 1->N, N->1, N->N
         return PassResult(graph_module, True)

backends/vulkan/_passes/int4_weight_only_quantizer.py

@@ -136,6 +136,7 @@ def _vk_replace_linear_int4(
                     scales_precision=scales_precision,
                 )
                 if copy_weights and child.weight.device != torch.device("meta"):
+                    # pyre-fixme[16]: `Module` has no attribute `weight`.
                     new_linear.weight = child.weight
                 setattr(module, name, new_linear)
         else:

backends/xnnpack/test/ops/conv2d.py

@@ -155,6 +155,8 @@ def _test(
         conv_count=1,
         dtype: torch.dtype = torch.float,
     ):
+        # pyre-fixme[29]: `Union[torch._tensor.Tensor,
+        #  torch.nn.modules.module.Module]` is not a function.
         tester = Tester(m.eval(), m.get_inputs())
 
         if quant_config is not None:

examples/models/llama/export_llama_lib.py

@@ -960,8 +960,14 @@ def _load_llama_model(
             use_kv_cache,
             use_sdpa_with_kv_cache,
             enable_dynamic_shape,
+            # pyre-fixme[6]: For 5th argument expected `ModelArgs` but got
+            #  `Union[Tensor, Module]`.
             model.max_seq_len,
+            # pyre-fixme[6]: For 6th argument expected `int` but got `Union[Tensor,
+            #  Module]`.
             model.n_layers,
+            # pyre-fixme[6]: For 7th argument expected `int` but got `Union[Tensor,
+            #  Module]`.
             model.vocab_size,
             metadata_str,
         ),
@@ -1045,6 +1051,7 @@ def _get_source_transforms(  # noqa
                 transforms.append(replace_attention_to_attention_sha)
                 transforms.append(replace_causal_mask)
                 transforms.append(replace_rms_norm_with_native_rms_norm)
+                # pyre-fixme[16]: Module `backends` has no attribute `qualcomm`.
                 transforms.append(convert_linear_to_conv2d)
             else:
                 transforms.append(replace_kv_cache_with_simple_kv_cache)
@@ -1056,6 +1063,7 @@ def _get_source_transforms(  # noqa
                     transforms.append(
                         get_model_with_r1_r2(args.optimized_rotation_path)
                     )
+                # pyre-fixme[16]: Module `backends` has no attribute `qualcomm`.
                 transforms.append(convert_linear_to_conv2d)
 
         elif args.mps:

examples/models/llama/source_transformation/apply_spin_quant_r1_r2.py

@@ -101,13 +101,17 @@ def apply_spin_quant_r1_r2(model: torch.nn.Module, optimized_rotation_path: str)
     optimized_rotation = torch.load(optimized_rotation_path, weights_only=True)
     R1 = optimized_rotation["R1"].to(torch.float32)
     config = model.params
+    # pyre-fixme[16]: Item `Tensor` of `Union[Tensor, Module]` has no attribute
+    #  `n_heads`.
     num_heads = config.n_heads
     head_dim = config.dim // num_heads
 
     rotate_embeddings(model, R1)
     rotate_head(model, R1)
     cleanup_memory()
 
+    # pyre-fixme[6]: For 1st argument expected `Iterable[Variable[_T]]` but got
+    #  `Union[Tensor, Module]`.
     for idx, layer in enumerate(model.layers):
         key = f"model.layers.{idx}.self_attn.R2"
         R2 = optimized_rotation[key].to(torch.float32)
@@ -130,6 +134,8 @@ def fuse_ln_linear(
 
         # Calculating new weight and bias
         W_ = linear.weight.data.to(dtype=torch.float32)
+        # pyre-fixme[58]: `*` is not supported for operand types `Tensor` and
+        #  `Union[torch._tensor.Tensor, torch.nn.modules.module.Module]`.
         linear.weight.data = (W_ * layernorm.weight.to(dtype=torch.float32)).to(
             linear_dtype
         )
@@ -140,6 +146,8 @@ def fuse_ln_linear(
                     torch.zeros(linear.out_features, dtype=torch.float32)
                 )
             linear.bias.data = linear.bias.data.to(dtype=torch.float32) + torch.matmul(
+                # pyre-fixme[6]: For 2nd argument expected `Tensor` but got
+                #  `Union[Tensor, Module]`.
                 W_, layernorm.bias.to(dtype=torch.float32)
             )
             linear.bias.data = linear.bias.data.to(linear_dtype)
@@ -148,10 +156,18 @@ def fuse_ln_linear(
 def fuse_layer_norms(model: torch.nn.Module):
     # Embedding fusion
     for W in [model.tok_embeddings]:
+        # pyre-fixme[16]: Item `Tensor` of `Union[Tensor, Module]` has no attribute
+        #  `weight`.
         W_ = W.weight.data.to(dtype=torch.float32)
+        # pyre-fixme[16]: Item `Tensor` of `Union[Tensor, Module]` has no attribute
+        #  `weight`.
         W.weight.data = (W_ - W_.mean(dim=-1, keepdim=True)).to(W.weight.data.dtype)
 
     # Fuse the linear operations in Layernorm into the adjacent linear blocks.
+    # pyre-fixme[29]:
+    #  `Union[BoundMethod[typing.Callable(torch._tensor.Tensor.__iter__)[[Named(self,
+    #  torch._tensor.Tensor)], typing.Any], torch._tensor.Tensor],
+    #  torch._tensor.Tensor, torch.nn.modules.module.Module]` is not a function.
     for layer in model.layers:
         # fuse the input layernorms into the linear layers
         fuse_ln_linear(layer.ffn_norm, [layer.feed_forward.w3, layer.feed_forward.w1])
@@ -170,9 +186,15 @@ def fuse_layer_norms(model: torch.nn.Module):
         layer.attention_norm.weight.data = torch.ones_like(W_norm, dtype=torch.float32)
 
     fuse_ln_linear(
+        # pyre-fixme[6]: For 1st argument expected `Module` but got `Union[Tensor,
+        #  Module]`.
         model.norm,
+        # pyre-fixme[6]: For 2nd argument expected `Iterable[Linear]` but got
+        #  `Iterable[Union[Tensor, Module]]`.
         [model.output],
     )
+    # pyre-fixme[16]: Item `Tensor` of `Union[Tensor, Module]` has no attribute
+    #  `weight`.
     W_norm = model.norm.weight.data
     model.norm.weight.data = torch.ones_like(W_norm, dtype=torch.float32)
 

examples/models/llama/source_transformation/attention.py

@@ -165,13 +165,19 @@ def __init__(self, attention_mha: nn.Module):
 
         for i in range(self.n_heads):
             self.wq[i].weight.data.copy_(
+                # pyre-fixme[16]: Item `Tensor` of `Union[Tensor, Module]` has no
+                #  attribute `weight`.
                 attention_mha.wq.weight[i * self.head_dim : (i + 1) * self.head_dim]
             )
         for i in range(self.n_kv_heads):
             self.wk[i].weight.data.copy_(
+                # pyre-fixme[16]: Item `Tensor` of `Union[Tensor, Module]` has no
+                #  attribute `weight`.
                 attention_mha.wk.weight[i * self.head_dim : (i + 1) * self.head_dim]
             )
             self.wv[i].weight.data.copy_(
+                # pyre-fixme[16]: Item `Tensor` of `Union[Tensor, Module]` has no
+                #  attribute `weight`.
                 attention_mha.wv.weight[i * self.head_dim : (i + 1) * self.head_dim]
             )
         self.wo = attention_mha.wo

examples/models/llama/source_transformation/lora.py

@@ -110,13 +110,23 @@ def filter_fn(child: torch.nn.Module, cur_fqn: str) -> bool:
 
     def replacement_fn(child: torch.nn.Module) -> torch.nn.Module:
         new_linear = Int8DynActInt4WeightLinearLoRA(
+            # pyre-fixme[6]: For 1st argument expected `int` but got `Union[Module,
+            #  Tensor]`.
             child.in_features,
+            # pyre-fixme[6]: For 2nd argument expected `int` but got `Union[Module,
+            #  Tensor]`.
             child.out_features,
             lora_rank=lora_rank,
             bias=False,
             device=child.weight.device,
+            # pyre-fixme[6]: For 6th argument expected `int` but got `Union[Module,
+            #  Tensor]`.
             groupsize=child.groupsize,
+            # pyre-fixme[6]: For 7th argument expected `dtype` but got
+            #  `Union[Module, Tensor]`.
             precision=child.precision,
+            # pyre-fixme[6]: For 8th argument expected `dtype` but got
+            #  `Union[Module, dtype, Tensor]`.
             scales_precision=child.scales.dtype,
         )
         return new_linear

examples/models/llama/source_transformation/pre_quantization.py

@@ -38,7 +38,11 @@ def filter_fn(child: torch.nn.Module, cur_fqn: str) -> bool:
 
     def replacement_fn(child: torch.nn.Module) -> torch.nn.Module:
         new_linear = Int8DynActInt4WeightLinear(
+            # pyre-fixme[6]: For 1st argument expected `int` but got `Union[Module,
+            #  Tensor]`.
             child.in_features,
+            # pyre-fixme[6]: For 2nd argument expected `int` but got `Union[Module,
+            #  Tensor]`.
             child.out_features,
             bias=False,
             device=child.weight.device,
@@ -99,7 +103,11 @@ def filter_fn(child: torch.nn.Module, cur_fqn: str) -> bool:
     def replacement_fn(child: torch.nn.Module) -> torch.nn.Module:
         new_linear = Int8DynActInt8WeightLinear(
             device=child.weight.device,
+            # pyre-fixme[6]: For 2nd argument expected `int` but got `Union[Module,
+            #  Tensor]`.
             in_features=child.in_features,
+            # pyre-fixme[6]: For 3rd argument expected `int` but got `Union[Module,
+            #  Tensor]`.
             out_features=child.out_features,
             precision=dtype,
             bias=False,

exir/capture/_capture.py

@@ -120,7 +120,11 @@ def _capture_legacy_do_not_use(f, args) -> ExirExportedProgram:
                 signature=ModuleCallSignature(
                     inputs=[],
                     outputs=[],
+                    # pyre-fixme[6]: For 3rd argument expected `TreeSpec` but got
+                    #  `Union[Tensor, Module]`.
                     in_spec=in_spec,
+                    # pyre-fixme[6]: For 4th argument expected `TreeSpec` but got
+                    #  `Union[Tensor, Module]`.
                     out_spec=out_spec,
                 ),
             )
@@ -350,6 +354,8 @@ def convert_to_fake(x):
                     inputs=[],
                     outputs=[],
                     in_spec=in_spec,
+                    # pyre-fixme[6]: For 4th argument expected `TreeSpec` but got
+                    #  `Union[None, TreeSpec, Tensor, Module]`.
                     out_spec=out_spec,
                 ),
             )