remove registering layernorm in other_ops.py

KingsleyLiu-NV · KingsleyLiu-NV · commit 3d5ee229d5b2 · 2023-11-19T21:33:27.000-08:00
diff --git a/torch/distributed/_tensor/ops/other_ops.py b/torch/distributed/_tensor/ops/other_ops.py
@@ -36,76 +36,6 @@ def slice_backward_rules(op_schema: OpSchema) -> OutputSharding:
     return OutputSharding(grad_input_spec)
 
 
-@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)
-    mean_dim = input_dim - norm_dim
-    mean_shape = list(input_shape[0:mean_dim]) + [1 for _ in range(norm_dim)]
-    mean_stride = list(np.cumprod(mean_shape[::-1])[:-1][::-1])
-    mean_stride.append(1)
-    assert len(input_shape) > len(normalized_shape)
-    assert list(input_shape[-norm_dim:]) == normalized_shape
-
-    output_tensor_meta = input_spec.tensor_meta
-    mean_tensor_meta = TensorMeta(
-        torch.Size(mean_shape),
-        tuple(mean_stride),
-        input_spec.tensor_meta.dtype,
-    )
-    rstd_tensor_meta = TensorMeta(
-        torch.Size(mean_shape),
-        tuple(mean_stride),
-        input_spec.tensor_meta.dtype,
-    )
-
-    output_spec = DTensorSpec.from_dim_map(
-        input_spec.mesh,
-        [-1 for _ in range(input_dim)],
-        [],
-        tensor_meta=output_tensor_meta,
-    )
-
-    mean_spec = DTensorSpec.from_dim_map(
-        input_spec.mesh,
-        [-1 for _ in range(mean_dim)],
-        [],
-        tensor_meta=mean_tensor_meta,
-    )
-
-    rstd_spec = DTensorSpec.from_dim_map(
-        input_spec.mesh,
-        [-1 for _ in range(mean_dim)],
-        [],
-        tensor_meta=rstd_tensor_meta,
-    )
-
-    return OutputSharding([output_spec, mean_spec, rstd_spec])
-
-
-@register_prop_rule(aten.native_layer_norm_backward.default)
-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))],
-        [],
-        tensor_meta=input_spec.tensor_meta,
-    )
-
-    grad_weight_spec = weight_spec
-    grad_bias_spec = bias_spec
-    return OutputSharding([grad_input_spec, grad_weight_spec, grad_bias_spec])
-
-
 @register_prop_rule(aten.bernoulli_.float)
 def bernoulli_rules(op_schema: OpSchema) -> OutputSharding:
     input_spec = op_schema.args_schema[0]
