fix linter errors 2

KingsleyLiu-NV · KingsleyLiu-NV · commit b97cb7c03c19 · 2023-11-19T21:33:27.000-08:00
diff --git a/torch/distributed/_tensor/dispatch.py b/torch/distributed/_tensor/dispatch.py
@@ -38,6 +38,7 @@ def decompose_handler(
     args: Tuple[object, ...],
     kwargs: Dict[str, object],
 ) -> object:
+<<<<<<< HEAD
     """
     Decomposes a op to core ATen op, this handler is mostly here
     for inference mode usage where the ops are not core aten ops.
@@ -47,6 +48,18 @@ def decompose_handler(
         return r
     else:
         raise RuntimeError("Decomposition failed")
+=======
+    # Redistribute grad_output tensor to the same placement as input tensor
+    args = list(args)
+    if op_call == torch.ops.aten.convolution_backward.default:
+        assert isinstance(args[0], dtensor.DTensor) and isinstance(
+            args[1], dtensor.DTensor
+        )
+        args[0] = args[0].redistribute(args[1].device_mesh, args[1].placements)
+    args = tuple(args)
+    out, _, _ = _operator_dispatch(op_call, args, kwargs, sharding_propagator)
+    return out
+>>>>>>> fix linter errors 2
 
 
 def is_same_size_handler(
diff --git a/torch/distributed/_tensor/examples/convnext_example.py b/torch/distributed/_tensor/examples/convnext_example.py
@@ -119,8 +119,8 @@ def __init__(
         self,
         in_chans=3,
         num_classes=10,
-        depths=[1, 1],
-        dims=[2, 4],
+        depths=[1, 1],  # noqa: B006
+        dims=[2, 4],  # noqa: B006
         drop_path_rate=0.0,
         layer_scale_init_value=1e-6,
         head_init_scale=1.0,
@@ -221,9 +221,8 @@ def test_tp_convnext_train(rank, world_size):
     optimizer.step()
     torch.cuda.synchronize(device)
 
-    forward_time = 0
-    backward_time = 0
-    optimizer_step_time = 0
+    forward_time = 0.0
+    backward_time = 0.0
     start = time.time()
     for i in range(ITER_TIME):
         t1 = time.time()
diff --git a/torch/distributed/_tensor/ops/__init__.py b/torch/distributed/_tensor/ops/__init__.py
@@ -6,5 +6,5 @@
 from .pointwise_ops import *  # noqa: F403
 from .random_ops import *  # noqa: F403
 from .view_ops import *  # noqa: F403
-from .conv_ops import *
-from .other_ops import *
+from .conv_ops import *  # noqa: F403
+from .other_ops import *  # noqa: F403
diff --git a/torch/distributed/_tensor/ops/conv_ops.py b/torch/distributed/_tensor/ops/conv_ops.py
@@ -23,6 +23,9 @@ def convolution_rules(op_schema: OpSchema) -> OutputSharding:
         groups,
     ) = op_schema.args_schema
 
+    assert isinstance(input_spec, DTensorSpec)
+    assert isinstance(weight_spec, DTensorSpec)
+    assert isinstance(bias_spec, DTensorSpec)
     in_shape = input_spec.tensor_meta.shape
     weight_shape = weight_spec.tensor_meta.shape
     N, C_in, H_in, W_in = in_shape[0], in_shape[1], in_shape[2], in_shape[3]
@@ -70,6 +73,10 @@ def convolution_backward_rules(op_schema: OpSchema) -> OutputSharding:
         output_mask,
     ) = op_schema.args_schema
 
+    assert isinstance(grad_output_spec, DTensorSpec)
+    assert isinstance(input_spec, DTensorSpec)
+    assert isinstance(weight_spec, DTensorSpec)
+
     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
@@ -1,5 +1,7 @@
 # Copyright (c) Meta Platforms, Inc. and affiliates
 # implement matrix related ops for distributed tensor
+from typing import List
+
 import numpy as np
 
 import torch
@@ -13,6 +15,7 @@
 @register_prop_rule(aten.slice_backward.default)
 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)
     grad_input_stride = list(np.cumprod(input_sizes[::-1])[:-1][::-1])
     grad_input_stride.append(1)
     dim_map = grad_output_spec.dim_map
@@ -36,7 +39,7 @@ def slice_backward_rules(op_schema: OpSchema) -> OutputSharding:
 @register_prop_rule(aten.native_layer_norm.default)
 def layer_norm_rules(op_schema: OpSchema) -> OutputSharding:
     input_spec, normalized_shape = op_schema.args_schema[0:2]
-
+    assert isinstance(input_spec, DTensorSpec)
     input_shape = input_spec.tensor_meta.shape
     input_dim = len(input_shape)
     norm_dim = len(normalized_shape)
@@ -88,7 +91,9 @@ def layer_norm_backward_rules(op_schema: OpSchema) -> OutputSharding:
     input_spec = op_schema.args_schema[1]
     weight_spec = op_schema.args_schema[5]
     bias_spec = op_schema.args_schema[6]
-
+    assert isinstance(input_spec, DTensorSpec)
+    assert isinstance(weight_spec, DTensorSpec)
+    assert isinstance(bias_spec, DTensorSpec)
     grad_input_spec = DTensorSpec.from_dim_map(
         input_spec.mesh,
         [-1 for _ in range(len(input_spec.tensor_meta.shape))],
@@ -104,18 +109,19 @@ def layer_norm_backward_rules(op_schema: OpSchema) -> OutputSharding:
 @register_prop_rule(aten.bernoulli_.float)
 def bernoulli_rules(op_schema: OpSchema) -> OutputSharding:
     input_spec = op_schema.args_schema[0]
+    assert isinstance(input_spec, DTensorSpec)
     return OutputSharding(input_spec)
 
 
 @register_prop_rule(aten.nll_loss_forward.default)
 def nll_loss_forward_rules(op_schema: OpSchema) -> OutputSharding:
     input_spec = op_schema.args_schema[0]
-
-    result_shape = []
-    result_stride = []
+    assert isinstance(input_spec, DTensorSpec)
+    result_shape: List[int] = []
+    result_stride: List[int] = []
     result_dim = 0
-    total_weight_shape = []
-    total_weight_stride = []
+    total_weight_shape: List[int] = []
+    total_weight_stride: List[int] = []
     total_weight_dim = 0
 
     result_tensor_meta = TensorMeta(
@@ -146,4 +152,5 @@ def nll_loss_forward_rules(op_schema: OpSchema) -> OutputSharding:
 @register_prop_rule(aten.nll_loss_backward.default)
 def nll_loss_backward_rules(op_schema: OpSchema) -> OutputSharding:
     input_spec = op_schema.args_schema[1]
+    assert isinstance(input_spec, DTensorSpec)
     return OutputSharding(input_spec)
