working

Author: jerryzh168

tutorials/developer_api_guide/my_dtype_tensor_subclass.py

@@ -24,6 +24,32 @@
 
 aten = torch.ops.aten
 
+# TODO: move to torchao/utils.py
+def fill_defaults(args, n, defaults_tail):
+    """
+    __torch_dispatch__ doesn't guarantee the number of arguments you are
+    passed (e.g., defaulted arguments are not passed); but usually it is
+    convenient to pad out the arguments list with defaults.  This function
+    helps you do that.
+    Args:
+        args: the list of positional arguments passed to __torch_dispatch__
+        n: the number of arguments you are expecting to get
+        defaults_tail: default values for the arguments, starting from the
+            end of the list
+    Example:
+        >>> fill_defaults([1, 2, 3], 5, [3, 4, 5])
+        [1, 2, 3, 4, 5]
+        >>> fill_defaults([1, 2, 3], 5, [None, None, None])
+        [1, 2, 3, None, None]]
+    """
+    if n - len(defaults_tail) > len(args):
+        raise RuntimeError("not enough defaults to fill arguments")
+    r = list(args)
+    for i in range(len(args), n):
+        r.append(defaults_tail[i - n + len(defaults_tail)])
+    return r
+
+
 ###############################
 # Base Layout Tensor Subclass #
 ###############################
@@ -204,6 +230,7 @@ def __new__(
         cls,
         int_data: torch.Tensor,
         scale: torch.Tensor,
+        transposed: bool,
         layout_type: LayoutType,
     ):
         kwargs = {}
@@ -220,22 +247,24 @@ def __init__(
         self,
         int_data: torch.Tensor,
         scale: torch.Tensor,
+        transposed: bool,
         layout_type: LayoutType,
     ):
         self.int_data = int_data
         self.scale = scale
+        self.transposed = transposed
         self.layout_type = layout_type
 
     def __tensor_flatten__(self):
-        return ["int_data", "scale"], [self.layout_type]
+        return ["int_data", "scale"], [self.transposed, self.layout_type]
 
     @classmethod
     def __tensor_unflatten__(
         cls, tensor_data_dict, tensor_attributes, outer_size, outer_stride
     ):
         int_data, scale = tensor_data_dict["int_data"], tensor_data_dict["scale"]
-        layout_type, = tensor_attributes
-        return cls(int_data, scale, layout_type)
+        transposed, layout_type, = tensor_attributes
+        return cls(int_data, scale, transposed, layout_type)
 
     @classmethod
     def from_plain(
@@ -248,12 +277,13 @@ def from_plain(
         extra metadata for packing etc.
         """
         assert isinstance(layout_type, PlainLayoutType)
-        return cls(int_data, scale, layout_type)
+        return cls(int_data, scale, False, layout_type)
 
     def _apply_fn_to_data(self, fn):
         return self.__class__(
             fn(self.int_data),
             fn(self.scale),
+            self.transposed,
             self.layout_type,
         )
 
@@ -273,11 +303,22 @@ def __torch_dispatch__(cls, func, types, args, kwargs):
         elif func is aten.split.Tensor:
             int_data_list = func(args[0].int_data, *args[1:], **kwargs)
             scale_list = func(args[0].scale, *args[1:], **kwargs)
-            out = [PlainMyDTypeLayout(int_data, scale, args[0].layout_type) for int_data, scale in zip(int_data_list, scale_list)]
+            out = [PlainMyDTypeLayout(int_data, scale, args[0].transposed, args[0].layout_type) for int_data, scale in zip(int_data_list, scale_list)]
             return out
         elif func is aten.empty_like.default:
             int_data_empty_like = func(args[0].int_data, *args[1:], **kwargs)
-            return PlainMyDTypeLayout(int_data_empty_like, args[0].scale, args[0].layout_type)
+            return PlainMyDTypeLayout(int_data_empty_like, args[0].scale, args[0].transposed, args[0].layout_type)
+        elif func is aten.slice.Tensor:
+            self, dim, start, end, step = fill_defaults(args, 5, [0, None, None, 1])
+            if dim == 0:
+                return return_and_correct_aliasing(
+                    func, args, kwargs, args[0]._apply_fn_to_data(lambda x: aten.slice.Tensor(x, dim, start, end, step))
+                )
+            elif dim == 1:
+                return PlainMyDTypeLayout(aten.slice.Tensor(self.int_data, dim, start, end, step), self.scale.view(-1, 1), self.transposed, self.layout_type)
+        elif func is aten.t.default:
+            args[0].transposed = not args[0].transposed
+            return return_and_correct_aliasing(func, args, kwargs, args[0])
 
         raise NotImplementedError(
             f"PlainMyDTypeLayout dispatch: attempting to run {func}, this is not supported"

tutorials/developer_api_guide/tensor_parallel.py

@@ -1,5 +1,5 @@
 import torch
-from my_dtype_tensor_subclass import MyDTypeTensor
+from my_dtype_tensor_subclass import MyDTypeTensor, fill_defaults
 from torch.utils._python_dispatch import return_and_correct_aliasing
 
 # a tensor subclass that supports tensor parallelism with DTensor
@@ -10,30 +10,6 @@ class MyDTypeTensorTP(MyDTypeTensor):
 
 aten = torch.ops.aten
 
-def fill_defaults(args, n, defaults_tail):
-    """
-    __torch_dispatch__ doesn't guarantee the number of arguments you are
-    passed (e.g., defaulted arguments are not passed); but usually it is
-    convenient to pad out the arguments list with defaults.  This function
-    helps you do that.
-    Args:
-        args: the list of positional arguments passed to __torch_dispatch__
-        n: the number of arguments you are expecting to get
-        defaults_tail: default values for the arguments, starting from the
-            end of the list
-    Example:
-        >>> fill_defaults([1, 2, 3], 5, [3, 4, 5])
-        [1, 2, 3, 4, 5]
-        >>> fill_defaults([1, 2, 3], 5, [None, None, None])
-        [1, 2, 3, None, None]]
-    """
-    if n - len(defaults_tail) > len(args):
-        raise RuntimeError("not enough defaults to fill arguments")
-    r = list(args)
-    for i in range(len(args), n):
-        r.append(defaults_tail[i - n + len(defaults_tail)])
-    return r
-
 @implements([aten._to_copy.default, aten.clone.default])
 def _(func, types, args, kwargs):
     return return_and_correct_aliasing(
@@ -51,20 +27,67 @@ def _(func, types, args, kwargs):
     empty_like_layout_tensor = func(args[0].layout_tensor, *args[1:], **kwargs)
     return MyDTypeTensorTP(empty_like_layout_tensor, empty_like_layout_tensor.shape)
 
-@implements([aten.slice.Tensor])
+@implements(aten.slice.Tensor)
 def _(func, types, args, kwargs):
     self, dim, start, end, step = fill_defaults(args, 5, [0, None, None, 1])
-    print("slice:", dim, start, end, step)
-    if dim == 0:
-        assert step == 1
-        return self.__class__(aten.slice.Tensor(self.layout_tensor), (end - start + 1,) + self.shape[1:], self.dtype)
-    return
+    assert step == 1
+    if end >= self.shape[dim]:
+        end = self.shape[dim]
+    print("dim:", dim, "start:", start, " end:", end, " shape:", end - start)
+    print("manual shape:", (end - start,) + self.shape[1:])
+    return self.__class__(aten.slice.Tensor(self.layout_tensor, dim, start, end, step), (end - start,) + self.shape[1:], self.dtype)
+
+# this is needed for DTensor.from_local() and for flattening tensor
+@implements(aten.view.default)
+def _(func, types, args, kwargs):
+    x, shape = args
+
+    if tuple(x.shape) == tuple(shape):
+        return x.__class__(x.layout_tensor, x.shape, x.dtype)
+
+    if len(shape) == 1 and shape[0] == -1:
+        return x.__class__(x.layout_tensor, (x.numel(),), x.dtype)
+
+    raise ValueError(f"{x.__class__.__name__} only supports .view() with same shape or shape=[-1]")
+
+@implements(aten.t.default)
+def _(func, types, args, kwargs):
+    tensor = args[0]
+    shape = tensor.shape[::-1]
+    new = tensor.__class__(tensor.layout_tensor.t(), shape, tensor.dtype)
+    return return_and_correct_aliasing(func, args, kwargs, new)
+
+@implements(aten.addmm.default)
+def _(func, types, args, kwargs):
+    input_tensor, weight_tensor, bias = (
+        args[1],
+        args[2],
+        args[0],
+    )
+    transposed = weight_tensor.layout_tensor.transposed
+    weight_tensor = weight_tensor.dequantize()
+    if transposed:
+        weight_tensor = weight_tensor.t()
+    return aten.addmm(input_tensor, weight_tensor, bias)
+
+@implements(aten.mm.default)
+def _(func, types, args, kwargs):
+    input_tensor, weight_tensor, bias = (
+        args[0],
+        args[1],
+        None
+    )
+    transposed = weight_tensor.layout_tensor.transposed
+    weight_tensor = weight_tensor.dequantize()
+    if transposed:
+        weight_tensor = weight_tensor.t()
+    return aten.mm(input_tensor, weight_tensor)
 
 
 class M(torch.nn.Module):
     def __init__(self, *args, **kwargs) -> None:
         super().__init__(*args, **kwargs)
-        self.linear = torch.nn.Linear(1024, 1024)
+        self.linear = torch.nn.Linear(1024, 1024, bias=False, device="cuda")
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
         return self.linear(x)
@@ -79,7 +102,7 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
     torch.manual_seed(5)
 
     m = M()
-    example_input = 100 * torch.randn(128, 1024)
+    example_input = 100 * torch.randn(128, 1024, device="cuda")
     m(example_input)
 
 
@@ -103,7 +126,7 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
     quantized_shard = quantized_weight[rank * n_local_rows : (rank + 1) * n_local_rows, :]
     print("quantized shard:", quantized_shard)
     # Construct DTensor from local shard
-    quantized_dtensor = DTensor.from_local(quantized_shard, device_mesh, [Shard(0)])
+    quantized_dtensor = DTensor.from_local(quantized_shard, mesh, [Shard(0)])
     print("quantized dtensor:", quantized_dtensor)
 
     # Replace parameter in module
@@ -117,4 +140,4 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
     )
     print("input dtensor:", input_dtensor)
 
-    m(input_dtensor)
+    print("result:", m(input_dtensor))