fix linter errors 3

KingsleyLiu-NV · KingsleyLiu-NV · commit b478235879eb · 2023-11-19T21:33:28.000-08:00
diff --git a/torch/distributed/_tensor/dispatch.py b/torch/distributed/_tensor/dispatch.py
@@ -396,7 +396,31 @@ def to_dt(res, spec):
                 else:
                     res_list.append(None)  # type: ignore[arg-type]
 
+<<<<<<< HEAD
             return tuple(res_list) if isinstance(res, tuple) else res_list
+=======
+        # run local op computation with potentially modified args/kwargs
+        local_tensor_args = cast(Tuple[object, ...], local_tensor_args)
+        if op_call in _random_ops and is_rng_supported_mesh(mesh):
+            if not random._rng_tracker:
+                raise RuntimeError(
+                    "A CudaRNGStateTracker instance must be instantiated "
+                    "before executing a random op over a DTensor. "
+                    "Try calling random.manual_seed() or distribute_tensor() "
+                    "before executing a DTensor random op."
+                )
+            # For DTensor random operator, run it within a distribute region
+            with random._rng_tracker._distribute_region(
+                cast(DTensorSpec, args_schema[0])
+            ):
+                local_results = op_call(*local_tensor_args, **local_kwargs)
+        elif op_call == torch.ops.aten.convolution.default:
+            local_results = tp_convolution(op_call, local_tensor_args, local_kwargs)
+        elif op_call == torch.ops.aten.convolution_backward.default:
+            local_results = tp_convolution_backward(
+                op_call, local_tensor_args, local_kwargs
+            )
+>>>>>>> fix linter errors 3
         else:
             # if the res contains only non tensor values, we simply return it without rewrapping
             return res
diff --git a/torch/distributed/_tensor/examples/convnext_example.py b/torch/distributed/_tensor/examples/convnext_example.py
@@ -105,7 +105,8 @@ def forward(self, x):
 def init_weights(m):
     if type(m) == nn.Conv2d or type(m) == nn.Linear:
         nn.init.ones_(m.weight)
-        nn.init.zeros_(m.bias)
+        if m.bias is not None:
+            nn.init.zeros_(m.bias)
 
 
 class ConvNeXt(nn.Module):
diff --git a/torch/distributed/_tensor/ops/conv_ops.py b/torch/distributed/_tensor/ops/conv_ops.py
@@ -1,5 +1,6 @@
 # Copyright (c) Meta Platforms, Inc. and affiliates
 # implement matrix related ops for distributed tensor
+from typing import List
 
 import torch
 from torch.distributed._tensor.op_schema import OpSchema, OutputSharding
@@ -26,8 +27,14 @@ def convolution_rules(op_schema: OpSchema) -> OutputSharding:
     assert isinstance(input_spec, DTensorSpec)
     assert isinstance(weight_spec, DTensorSpec)
     assert isinstance(bias_spec, DTensorSpec)
+    assert input_spec.tensor_meta is not None
+    assert weight_spec.tensor_meta is not None
     in_shape = input_spec.tensor_meta.shape
     weight_shape = weight_spec.tensor_meta.shape
+    assert isinstance(stride, List)
+    assert isinstance(padding, List)
+    assert isinstance(dilation, List)
+    assert isinstance(weight_shape, torch.Size)
     N, C_in, H_in, W_in = in_shape[0], in_shape[1], in_shape[2], in_shape[3]
     C_out = weight_shape[0]
     H_out = (H_in + 2 * padding[0] - dilation[0] * (weight_shape[2] - 1) - 1) // stride[
@@ -76,7 +83,8 @@ def convolution_backward_rules(op_schema: OpSchema) -> OutputSharding:
     assert isinstance(grad_output_spec, DTensorSpec)
     assert isinstance(input_spec, DTensorSpec)
     assert isinstance(weight_spec, DTensorSpec)
-
+    assert isinstance(bias_shape_opt, List)
+    assert input_spec.tensor_meta is not None
     weight_tensor_meta = weight_spec.tensor_meta
     bias_tensor_meta = TensorMeta(
         torch.Size(bias_shape_opt),
diff --git a/torch/distributed/_tensor/ops/other_ops.py b/torch/distributed/_tensor/ops/other_ops.py
@@ -16,6 +16,8 @@
 def slice_backward_rules(op_schema: OpSchema) -> OutputSharding:
     grad_output_spec, input_sizes, dim, start, end, step = op_schema.args_schema
     assert isinstance(grad_output_spec, DTensorSpec)
+    assert isinstance(input_sizes, List)
+    assert grad_output_spec.tensor_meta is not None
     grad_input_stride = list(np.cumprod(input_sizes[::-1])[:-1][::-1])
     grad_input_stride.append(1)
     dim_map = grad_output_spec.dim_map
@@ -48,6 +50,7 @@ def bernoulli_rules(op_schema: OpSchema) -> OutputSharding:
 def nll_loss_forward_rules(op_schema: OpSchema) -> OutputSharding:
     input_spec = op_schema.args_schema[0]
     assert isinstance(input_spec, DTensorSpec)
+    assert input_spec.tensor_meta is not None
     result_shape: List[int] = []
     result_stride: List[int] = []
     result_dim = 0
diff --git a/torch/distributed/_tensor/tp_conv.py b/torch/distributed/_tensor/tp_conv.py
@@ -1,6 +1,6 @@
 # Copyright (c) Meta Platforms, Inc. and affiliates
 # implement matrix related ops for distributed tensor
-from typing import cast, Dict, Tuple
+from typing import cast, Dict, List, Tuple
 
 import torch
 import torch.distributed as dist
@@ -102,20 +102,20 @@ def _ring_send_recv_aggregate(grad_in_tensor, d1, d2, left, right, rank, size):
 
 def tp_convolution(
     op_call: torch._ops.OpOverload,
-    *local_tensor_args: Tuple[object, ...],
-    **local_tensor_kwargs: Dict[str, object],
+    local_tensor_args: Tuple[object, ...],
+    local_tensor_kwargs: Dict[str, object],
 ) -> object:
     assert op_call == aten.convolution.default
     assert len(local_tensor_args) == 9
 
     rank = dist.get_rank()
     size = dist.get_world_size()
-    in_tensor = local_tensor_args[0]
-    weight = local_tensor_args[1]
-    bias = local_tensor_args[2]
+    in_tensor = cast(torch.Tensor, local_tensor_args[0])
+    weight = cast(torch.Tensor, local_tensor_args[1])
     stride, padding, dilation = local_tensor_args[3:6]
 
     assert _is_supported(in_tensor.shape, weight.shape, stride, padding, dilation)
+    assert isinstance(padding, List)
 
     if not _requires_data_exchange(padding):
         local_results = op_call(*local_tensor_args, **local_tensor_kwargs)
@@ -141,35 +141,35 @@ def tp_convolution(
         local_results = op_call(*local_tensor_args, **local_tensor_kwargs)
 
         # step3 remove extra ouputs from the results
-        padding = local_tensor_args[4][1]
+        padding_w = padding[1]
         w = local_results.size(3)
         if rank == 0:
-            local_results = local_results[:, :, :, : w - padding]
+            local_results = local_results[:, :, :, : w - padding_w]
         elif rank == size - 1:
-            local_results = local_results[:, :, :, padding:]
+            local_results = local_results[:, :, :, padding_w:]
         else:
-            local_results = local_results[:, :, :, padding : w - padding]
+            local_results = local_results[:, :, :, padding_w : w - padding_w]
 
         return local_results
 
 
 def tp_convolution_backward(
     op_call: torch._ops.OpOverload,
-    *local_tensor_args: Tuple[object, ...],
-    **local_tensor_kwargs: Dict[str, object],
+    local_tensor_args: Tuple[object, ...],
+    local_tensor_kwargs: Dict[str, object],
 ) -> object:
     assert op_call == aten.convolution_backward.default
     assert len(local_tensor_args) == 11
 
     rank = dist.get_rank()
     size = dist.get_world_size()
-    grad_out_tensor = local_tensor_args[0]
-    in_tensor = local_tensor_args[1]
-    weight = local_tensor_args[2]
-    bias = local_tensor_args[3]
+    grad_out_tensor = cast(torch.Tensor, local_tensor_args[0])
+    in_tensor = cast(torch.Tensor, local_tensor_args[1])
+    weight = cast(torch.Tensor, local_tensor_args[2])
     stride, padding, dilation = local_tensor_args[4:7]
 
     assert _is_supported(in_tensor.shape, weight.shape, stride, padding, dilation)
+    assert isinstance(padding, List)
 
     if not _requires_data_exchange(padding):
         local_results = op_call(*local_tensor_args, **local_tensor_kwargs)
@@ -190,18 +190,18 @@ def tp_convolution_backward(
 
         # step2 reconstruct local gradient output tensor
         N, C_out, H_out, _ = grad_out_tensor.shape
-        padding = local_tensor_args[5][1]
+        padding_w = padding[1]
         if rank == 0:
             grad_out_tensor = torch.nn.functional.pad(
-                grad_out_tensor, (0, padding), "constant", 0
+                grad_out_tensor, (0, padding_w), "constant", 0
             )
         elif rank == size - 1:
             grad_out_tensor = torch.nn.functional.pad(
-                grad_out_tensor, (padding, 0), "constant", 0
+                grad_out_tensor, (padding_w, 0), "constant", 0
             )
         else:
             grad_out_tensor = torch.nn.functional.pad(
-                grad_out_tensor, (padding, padding), "constant", 0
+                grad_out_tensor, (padding_w, padding_w), "constant", 0
             )
 
         # step3 feed local input tensor to op_call
