We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
Runtime error when enable memory-centric-tiled-linear
Hi team,

I'm running the [example](https://github.com/microsoft/DeepSpeedExamples/blob/25d73cf73fb3dc66faefa141b7319526555be9fc/Megatron-LM-v1.1.5-ZeRO3/examples/ds_pretrain_gpt2-zero3.sh) on 8 * 40GB A100 Node,

the configurations are 

<img width="653" alt="Screen Shot 2021-07-28 at 8 48 35 PM" src="https://user-images.githubusercontent.com/3942083/127324796-14439be1-f272-4901-9d8e-3d4fbe48bf06.png">

and the stacktrace is

<img width="1226" alt="Screen Shot 2021-07-28 at 8 41 53 PM" src="https://user-images.githubusercontent.com/3942083/127324847-b7e5cb23-9b5e-4ebe-bc95-1cf976abee07.png">

it seems that the args --memory-centric-tiled-linear and --tile-factor are the root cause. So Anyone can help me to fix this problems? Thanks.


--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `deepspeed/runtime/zero/tiling.py`
Content:
```
1 import torch
2 import deepspeed
3 from deepspeed.runtime.utils import partition_uniform as partition
4 
5 
6 def split_tensor_along_last_dim(tensor, partitions, contiguous_split_chunks=False):
7     """Split a tensor along its last dimension. Adapted from Megatron-LM.
8 
9     Arguments:
10         tensor: input tensor.
11         partitions: list of partition sizes to supply to torch.split
12         contiguous_split_chunks: If True, make each chunk contiguous
13                                  in memory.
14     """
15     # Get the size and dimension.
16     last_dim = tensor.dim() - 1
17     # Split.
18     tensor_list = torch.split(tensor, partitions, dim=last_dim)
19     # Note: torch.split does not create contiguous tensors by default.
20     if contiguous_split_chunks:
21         return tuple(chunk.contiguous() for chunk in tensor_list)
22 
23     return tensor_list
24 
25 
26 class TiledLinear(torch.nn.Module):
27     def __init__(self,
28                  in_features,
29                  out_features,
30                  bias=True,
31                  in_splits=1,
32                  out_splits=1,
33                  input_is_already_split=False,
34                  combine_out_splits=True,
35                  linear_cls=torch.nn.Linear,
36                  init_linear=None,
37                  **kwargs):
38         """A replacement for ``torch.nn.Linear`` that works with ZeRO-3 to reduce
39         memory requirements via tiling.
40 
41         TiledLinear breaks the input and output dimensions of a linear layer
42         into tiles that are processed in sequence. This class enables huge
43         linear layers when combined with ZeRO-3 because inactive tiles can be
44         partitioned and offloaded.
45 
46         .. note::
47             We recommend using as few tiles as necessary. Tiling
48             significantly reduces memory usage, but can reduce throughput
49             for inexpensive layers. This due to the smaller kernels having
50             less parallelism and lower arithmetic intensity, while
51             introducing more frequent synchronization and communication.
52 
53         Args:
54             in_features (int): See ``torch.nn.Linear``
55             out_features (int): See ``torch.nn.Linear``
56             bias (bool, optional): See ``torch.nn.Linear``
57             in_splits (int, optional): The number of tiles along the input dimension. Defaults to 1.
58             out_splits (int, optional): The number of tiles along the output dimension. Defaults to 1.
59             input_is_already_split (bool, optional): If set to ``True``, assume that the ``input_`` in
60                 to ``forward()`` is already split into ``in_splits`` chunks. Defaults to ``False``.
61             combine_out_splits (bool, optional): If set to ``False``, do not combine the ``out_splits`` outputs
62                 into a single tensor. Defaults to ``True``.
63             linear_cls (class, optional): The underlying class to build individual tiles.
64                 Defaults to ``torch.nn.Linear``.
65             init_linear (``torch.nn.Linear``, optional): If set, copy the parameters of
66                 ``init_linear``. Useful for debugging. Defaults to ``None``.
67             kwargs (dict, optional): additional keyword arguments to provide to ``linear_cls()``.
68 
69         Raises:
70             RuntimeError: ``in_splits`` must be within the range [1, in_features).
71             RuntimeError: ``out_splits`` must be within the range of [1, out_features).
72         """
73 
74         super().__init__()
75 
76         if (in_splits < 1) or (in_splits > in_features):
77             raise RuntimeError('in splits must be in range [1, in_features].')
78         if (out_splits < 1) or (out_splits > out_features):
79             raise RuntimeError('out splits must be in range [1, out_features].')
80 
81         # global, not necessarily local
82         self.in_features = in_features
83         self.out_features = out_features
84         self.use_bias = bias
85 
86         self.out_splits = out_splits
87         self.in_splits = in_splits
88         self.input_is_already_split = input_is_already_split
89         self.combine_out_splits = combine_out_splits
90 
91         # Build partition-lists. These are CSR-style splits [0, part0, part1, ..., features]
92         # For example, row_parts[p] gives the start of partition p and row_parts[p+1]
93         # is the exclusive end.
94         self.in_parts = partition(num_items=in_features, num_parts=in_splits)
95         self.out_parts = partition(num_items=out_features, num_parts=out_splits)
96 
97         assert len(self.out_parts) == out_splits + 1
98         assert len(self.in_parts) == in_splits + 1
99         assert self.out_parts[0] == 0
100         assert self.out_parts[out_splits] == out_features
101         assert self.in_parts[in_splits] == in_features
102 
103         self.linears = torch.nn.ModuleList()
104         for out_id in range(out_splits):
105             self.linears.append(torch.nn.ModuleList())
106 
107             local_out_dim = self.out_parts[out_id + 1] - self.out_parts[out_id]
108 
109             for in_id in range(in_splits):
110                 #if input_size is split, we only need one bias
111                 local_bias = bias if in_id == (in_splits - 1) else False
112 
113                 local_in_dim = self.in_parts[in_id + 1] - self.in_parts[in_id]
114                 local = linear_cls(local_in_dim,
115                                    local_out_dim,
116                                    bias=local_bias,
117                                    **kwargs)
118                 self.linears[out_id].append(local)
119 
120         # Optionally initialize with a known tensor
121         if init_linear is not None:
122             self.copy_params_from(init_linear)
123 
124     def forward(self, input_):
125         if self.in_splits > 1 and not self.input_is_already_split:
126             split_sizes = [
127                 self.in_parts[p + 1] - self.in_parts[p] for p in range(self.in_splits)
128             ]
129             inputs = self._split_global_input(input_, split_sizes)
130         elif self.in_splits > 1:
131             inputs = input_
132             assert len(inputs) == self.in_splits, f"Col splits {self.in_splits} does not match input splits {len(inputs)}"
133         else:
134             # no splits
135             inputs = [input_]
136 
137         outputs = [None] * self.out_splits
138         for out_id in range(self.out_splits):
139             for in_id in range(self.in_splits):
140                 local_output = self.linears[out_id][in_id](inputs[in_id])
141 
142                 outputs[out_id] = self._reduce_local_output(in_id=in_id,
143                                                             out_id=out_id,
144                                                             current_out=outputs[out_id],
145                                                             new_out=local_output)
146 
147         if self.combine_out_splits:
148             return self._combine_output_splits(outputs)
149 
150         return outputs
151 
152     def _split_global_input(self, input, split_sizes):
153         """Partition an input tensor along the last dimension, aligned with given splits.
154 
155         Subclasses should override this method to account for new input types.
156 
157         Args:
158             input (List[Tensor]): The tensor to partition along the last dimension.
159             split_sizes (List[int]): The size of each partition.
160 
161         Returns:
162             List[Any]: A list of the chunks of ``input``.
163         """
164         return split_tensor_along_last_dim(input, split_sizes)
165 
166     def _reduce_local_output(self, in_id, out_id, current_out, new_out):
167         """Reduce (sum) a new local result into the existing local results.
168 
169         Subclasses should override this method.
170 
171         For a given ``out_id``, this method is called ``in_id-1`` times. The first input
172         split is a simple assignment.
173 
174         Args:
175             in_id (int): The input split that produced ``new_out``.
176             out_id (int): The output split that produced ``new_out``.
177             current_out (Any): The reduced form of all previous ``out_id`` results.
178             new_out (Any): The local result from forward (``in_id``, ``out_id``)e
179 
180         Returns:
181             Any: The combined result of ``current_out`` and ``new_out``.
182         """
183 
184         if current_out is None:
185             #this clone is necessary to preserve auto grad
186             #there is some issue with inplace update for outputs that are views
187             return new_out.clone()
188         else:
189             return current_out + new_out
190 
191     def _combine_output_splits(self, outputs):
192         """Join the splits of the output into a single result.
193 
194         Args:
195             outputs (List[Any]): The reduced outputs for each output split.
196 
197         Returns:
198             Any: The combined outputs.
199         """
200         assert len(outputs) == self.out_splits
201         return torch.cat(outputs, dim=-1)
202 
203     @torch.no_grad()
204     def copy_params_from(self, other):
205         """Copy the weight and bias data from ``other``.
206 
207         This is especially useful for reproducible initialization and testing.
208 
209         Equivalent to:
210 
211         .. code-block:: python
212 
213             with torch.no_grad():
214                 self.weight.copy_(other.weight)
215                 if self.bias is not None:
216                     self.bias.copy_(other.bias)
217 
218         .. note::
219             If ZeRO-3 is enabled, this is a collective operation and the
220             updated parameters of data-parallel rank 0 will be visible on all
221             ranks. See :class:`deepspeed.zero.GatheredParameters` for more
222             information.
223 
224 
225         Args:
226             other (``torch.nn.Linear``): the linear layer to copy from.
227         """
228         assert hasattr(other, 'weight')
229         assert other.weight.size() == (self.out_features, self.in_features)
230         if self.use_bias:
231             assert hasattr(other, 'bias')
232             assert other.bias is not None
233             assert other.bias.size() == (self.out_features, )
234         else:
235             assert other.bias is None
236 
237         for row in range(self.out_splits):
238             rstart = self.out_parts[row]
239             rstop = self.out_parts[row + 1]
240 
241             for col in range(self.in_splits):
242                 cstart = self.in_parts[col]
243                 cstop = self.in_parts[col + 1]
244 
245                 local = self.linears[row][col]
246                 global_weight = other.weight[rstart:rstop, cstart:cstop]
247                 with deepspeed.zero.GatheredParameters(local.weight, modifier_rank=0):
248                     local.weight.copy_(global_weight)
249 
250             if local.bias is not None:
251                 with deepspeed.zero.GatheredParameters(local.bias, modifier_rank=0):
252                     local.bias.data.copy_(other.bias[rstart:rstop].data)
253 
254 
255 class TiledLinearReturnBias(TiledLinear):
256     """Wrapper for a Linear class that returns its own bias parameter, such as
257     used by Megatron-LM.
258     """
259     def _reduce_local_output(self, in_id, out_id, current_out, new_out):
260         """Reduces output tensors, but not the returned bias. """
261         if current_out is not None:
262             old_tensor, old_bias = current_out
263         else:
264             old_tensor, old_bias = None, None
265 
266         assert isinstance(new_out, tuple)
267         assert len(new_out) == 2
268 
269         tensor, bias = new_out
270         assert tensor is not None
271 
272         tensor = super()._reduce_local_output(in_id=in_id,
273                                               out_id=out_id,
274                                               current_out=old_tensor,
275                                               new_out=tensor)
276 
277         if bias is None:
278             bias = old_bias
279 
280         return tensor, bias
281 
282     def _combine_output_splits(self, outputs):
283         # stack output tensors
284         tensors = [o[0] for o in outputs]
285         tensor = super()._combine_output_splits(tensors)
286 
287         # stack biases if applicable
288         biases = [o[1] for o in outputs if o[1] is not None]
289         if len(biases) > 0:
290             bias = super()._combine_output_splits(biases)
291         else:
292             bias = None
293 
294         return tensor, bias
295
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

diff --git a/deepspeed/runtime/zero/tiling.py b/deepspeed/runtime/zero/tiling.py
--- a/deepspeed/runtime/zero/tiling.py
+++ b/deepspeed/runtime/zero/tiling.py
@@ -14,6 +14,7 @@
     """
     # Get the size and dimension.
     last_dim = tensor.dim() - 1
+
     # Split.
     tensor_list = torch.split(tensor, partitions, dim=last_dim)
     # Note: torch.split does not create contiguous tensors by default.
@@ -123,8 +124,9 @@
 
     def forward(self, input_):
         if self.in_splits > 1 and not self.input_is_already_split:
+            input_parts = partition(input_.shape[-1], self.in_splits)
             split_sizes = [
-                self.in_parts[p + 1] - self.in_parts[p] for p in range(self.in_splits)
+                input_parts[p + 1] - input_parts[p] for p in range(self.in_splits)
             ]
             inputs = self._split_global_input(input_, split_sizes)
         elif self.in_splits > 1:

{"golden_diff": "diff --git a/deepspeed/runtime/zero/tiling.py b/deepspeed/runtime/zero/tiling.py\n--- a/deepspeed/runtime/zero/tiling.py\n+++ b/deepspeed/runtime/zero/tiling.py\n@@ -14,6 +14,7 @@\n     \"\"\"\n     # Get the size and dimension.\n     last_dim = tensor.dim() - 1\n+\n     # Split.\n     tensor_list = torch.split(tensor, partitions, dim=last_dim)\n     # Note: torch.split does not create contiguous tensors by default.\n@@ -123,8 +124,9 @@\n \n     def forward(self, input_):\n         if self.in_splits > 1 and not self.input_is_already_split:\n+            input_parts = partition(input_.shape[-1], self.in_splits)\n             split_sizes = [\n-                self.in_parts[p + 1] - self.in_parts[p] for p in range(self.in_splits)\n+                input_parts[p + 1] - input_parts[p] for p in range(self.in_splits)\n             ]\n             inputs = self._split_global_input(input_, split_sizes)\n         elif self.in_splits > 1:\n", "issue": "Runtime error when enable memory-centric-tiled-linear\nHi team,\r\n\r\nI'm running the [example](https://github.com/microsoft/DeepSpeedExamples/blob/25d73cf73fb3dc66faefa141b7319526555be9fc/Megatron-LM-v1.1.5-ZeRO3/examples/ds_pretrain_gpt2-zero3.sh) on 8 * 40GB A100 Node,\r\n\r\nthe configurations are \r\n\r\n<img width=\"653\" alt=\"Screen Shot 2021-07-28 at 8 48 35 PM\" src=\"https://user-images.githubusercontent.com/3942083/127324796-14439be1-f272-4901-9d8e-3d4fbe48bf06.png\">\r\n\r\nand the stacktrace is\r\n\r\n<img width=\"1226\" alt=\"Screen Shot 2021-07-28 at 8 41 53 PM\" src=\"https://user-images.githubusercontent.com/3942083/127324847-b7e5cb23-9b5e-4ebe-bc95-1cf976abee07.png\">\r\n\r\nit seems that the args --memory-centric-tiled-linear and --tile-factor are the root cause. So Anyone can help me to fix this problems? Thanks.\r\n\n", "before_files": [{"content": "import torch\nimport deepspeed\nfrom deepspeed.runtime.utils import partition_uniform as partition\n\n\ndef split_tensor_along_last_dim(tensor, partitions, contiguous_split_chunks=False):\n    \"\"\"Split a tensor along its last dimension. Adapted from Megatron-LM.\n\n    Arguments:\n        tensor: input tensor.\n        partitions: list of partition sizes to supply to torch.split\n        contiguous_split_chunks: If True, make each chunk contiguous\n                                 in memory.\n    \"\"\"\n    # Get the size and dimension.\n    last_dim = tensor.dim() - 1\n    # Split.\n    tensor_list = torch.split(tensor, partitions, dim=last_dim)\n    # Note: torch.split does not create contiguous tensors by default.\n    if contiguous_split_chunks:\n        return tuple(chunk.contiguous() for chunk in tensor_list)\n\n    return tensor_list\n\n\nclass TiledLinear(torch.nn.Module):\n    def __init__(self,\n                 in_features,\n                 out_features,\n                 bias=True,\n                 in_splits=1,\n                 out_splits=1,\n                 input_is_already_split=False,\n                 combine_out_splits=True,\n                 linear_cls=torch.nn.Linear,\n                 init_linear=None,\n                 **kwargs):\n        \"\"\"A replacement for ``torch.nn.Linear`` that works with ZeRO-3 to reduce\n        memory requirements via tiling.\n\n        TiledLinear breaks the input and output dimensions of a linear layer\n        into tiles that are processed in sequence. This class enables huge\n        linear layers when combined with ZeRO-3 because inactive tiles can be\n        partitioned and offloaded.\n\n        .. note::\n            We recommend using as few tiles as necessary. Tiling\n            significantly reduces memory usage, but can reduce throughput\n            for inexpensive layers. This due to the smaller kernels having\n            less parallelism and lower arithmetic intensity, while\n            introducing more frequent synchronization and communication.\n\n        Args:\n            in_features (int): See ``torch.nn.Linear``\n            out_features (int): See ``torch.nn.Linear``\n            bias (bool, optional): See ``torch.nn.Linear``\n            in_splits (int, optional): The number of tiles along the input dimension. Defaults to 1.\n            out_splits (int, optional): The number of tiles along the output dimension. Defaults to 1.\n            input_is_already_split (bool, optional): If set to ``True``, assume that the ``input_`` in\n                to ``forward()`` is already split into ``in_splits`` chunks. Defaults to ``False``.\n            combine_out_splits (bool, optional): If set to ``False``, do not combine the ``out_splits`` outputs\n                into a single tensor. Defaults to ``True``.\n            linear_cls (class, optional): The underlying class to build individual tiles.\n                Defaults to ``torch.nn.Linear``.\n            init_linear (``torch.nn.Linear``, optional): If set, copy the parameters of\n                ``init_linear``. Useful for debugging. Defaults to ``None``.\n            kwargs (dict, optional): additional keyword arguments to provide to ``linear_cls()``.\n\n        Raises:\n            RuntimeError: ``in_splits`` must be within the range [1, in_features).\n            RuntimeError: ``out_splits`` must be within the range of [1, out_features).\n        \"\"\"\n\n        super().__init__()\n\n        if (in_splits < 1) or (in_splits > in_features):\n            raise RuntimeError('in splits must be in range [1, in_features].')\n        if (out_splits < 1) or (out_splits > out_features):\n            raise RuntimeError('out splits must be in range [1, out_features].')\n\n        # global, not necessarily local\n        self.in_features = in_features\n        self.out_features = out_features\n        self.use_bias = bias\n\n        self.out_splits = out_splits\n        self.in_splits = in_splits\n        self.input_is_already_split = input_is_already_split\n        self.combine_out_splits = combine_out_splits\n\n        # Build partition-lists. These are CSR-style splits [0, part0, part1, ..., features]\n        # For example, row_parts[p] gives the start of partition p and row_parts[p+1]\n        # is the exclusive end.\n        self.in_parts = partition(num_items=in_features, num_parts=in_splits)\n        self.out_parts = partition(num_items=out_features, num_parts=out_splits)\n\n        assert len(self.out_parts) == out_splits + 1\n        assert len(self.in_parts) == in_splits + 1\n        assert self.out_parts[0] == 0\n        assert self.out_parts[out_splits] == out_features\n        assert self.in_parts[in_splits] == in_features\n\n        self.linears = torch.nn.ModuleList()\n        for out_id in range(out_splits):\n            self.linears.append(torch.nn.ModuleList())\n\n            local_out_dim = self.out_parts[out_id + 1] - self.out_parts[out_id]\n\n            for in_id in range(in_splits):\n                #if input_size is split, we only need one bias\n                local_bias = bias if in_id == (in_splits - 1) else False\n\n                local_in_dim = self.in_parts[in_id + 1] - self.in_parts[in_id]\n                local = linear_cls(local_in_dim,\n                                   local_out_dim,\n                                   bias=local_bias,\n                                   **kwargs)\n                self.linears[out_id].append(local)\n\n        # Optionally initialize with a known tensor\n        if init_linear is not None:\n            self.copy_params_from(init_linear)\n\n    def forward(self, input_):\n        if self.in_splits > 1 and not self.input_is_already_split:\n            split_sizes = [\n                self.in_parts[p + 1] - self.in_parts[p] for p in range(self.in_splits)\n            ]\n            inputs = self._split_global_input(input_, split_sizes)\n        elif self.in_splits > 1:\n            inputs = input_\n            assert len(inputs) == self.in_splits, f\"Col splits {self.in_splits} does not match input splits {len(inputs)}\"\n        else:\n            # no splits\n            inputs = [input_]\n\n        outputs = [None] * self.out_splits\n        for out_id in range(self.out_splits):\n            for in_id in range(self.in_splits):\n                local_output = self.linears[out_id][in_id](inputs[in_id])\n\n                outputs[out_id] = self._reduce_local_output(in_id=in_id,\n                                                            out_id=out_id,\n                                                            current_out=outputs[out_id],\n                                                            new_out=local_output)\n\n        if self.combine_out_splits:\n            return self._combine_output_splits(outputs)\n\n        return outputs\n\n    def _split_global_input(self, input, split_sizes):\n        \"\"\"Partition an input tensor along the last dimension, aligned with given splits.\n\n        Subclasses should override this method to account for new input types.\n\n        Args:\n            input (List[Tensor]): The tensor to partition along the last dimension.\n            split_sizes (List[int]): The size of each partition.\n\n        Returns:\n            List[Any]: A list of the chunks of ``input``.\n        \"\"\"\n        return split_tensor_along_last_dim(input, split_sizes)\n\n    def _reduce_local_output(self, in_id, out_id, current_out, new_out):\n        \"\"\"Reduce (sum) a new local result into the existing local results.\n\n        Subclasses should override this method.\n\n        For a given ``out_id``, this method is called ``in_id-1`` times. The first input\n        split is a simple assignment.\n\n        Args:\n            in_id (int): The input split that produced ``new_out``.\n            out_id (int): The output split that produced ``new_out``.\n            current_out (Any): The reduced form of all previous ``out_id`` results.\n            new_out (Any): The local result from forward (``in_id``, ``out_id``)e\n\n        Returns:\n            Any: The combined result of ``current_out`` and ``new_out``.\n        \"\"\"\n\n        if current_out is None:\n            #this clone is necessary to preserve auto grad\n            #there is some issue with inplace update for outputs that are views\n            return new_out.clone()\n        else:\n            return current_out + new_out\n\n    def _combine_output_splits(self, outputs):\n        \"\"\"Join the splits of the output into a single result.\n\n        Args:\n            outputs (List[Any]): The reduced outputs for each output split.\n\n        Returns:\n            Any: The combined outputs.\n        \"\"\"\n        assert len(outputs) == self.out_splits\n        return torch.cat(outputs, dim=-1)\n\n    @torch.no_grad()\n    def copy_params_from(self, other):\n        \"\"\"Copy the weight and bias data from ``other``.\n\n        This is especially useful for reproducible initialization and testing.\n\n        Equivalent to:\n\n        .. code-block:: python\n\n            with torch.no_grad():\n                self.weight.copy_(other.weight)\n                if self.bias is not None:\n                    self.bias.copy_(other.bias)\n\n        .. note::\n            If ZeRO-3 is enabled, this is a collective operation and the\n            updated parameters of data-parallel rank 0 will be visible on all\n            ranks. See :class:`deepspeed.zero.GatheredParameters` for more\n            information.\n\n\n        Args:\n            other (``torch.nn.Linear``): the linear layer to copy from.\n        \"\"\"\n        assert hasattr(other, 'weight')\n        assert other.weight.size() == (self.out_features, self.in_features)\n        if self.use_bias:\n            assert hasattr(other, 'bias')\n            assert other.bias is not None\n            assert other.bias.size() == (self.out_features, )\n        else:\n            assert other.bias is None\n\n        for row in range(self.out_splits):\n            rstart = self.out_parts[row]\n            rstop = self.out_parts[row + 1]\n\n            for col in range(self.in_splits):\n                cstart = self.in_parts[col]\n                cstop = self.in_parts[col + 1]\n\n                local = self.linears[row][col]\n                global_weight = other.weight[rstart:rstop, cstart:cstop]\n                with deepspeed.zero.GatheredParameters(local.weight, modifier_rank=0):\n                    local.weight.copy_(global_weight)\n\n            if local.bias is not None:\n                with deepspeed.zero.GatheredParameters(local.bias, modifier_rank=0):\n                    local.bias.data.copy_(other.bias[rstart:rstop].data)\n\n\nclass TiledLinearReturnBias(TiledLinear):\n    \"\"\"Wrapper for a Linear class that returns its own bias parameter, such as\n    used by Megatron-LM.\n    \"\"\"\n    def _reduce_local_output(self, in_id, out_id, current_out, new_out):\n        \"\"\"Reduces output tensors, but not the returned bias. \"\"\"\n        if current_out is not None:\n            old_tensor, old_bias = current_out\n        else:\n            old_tensor, old_bias = None, None\n\n        assert isinstance(new_out, tuple)\n        assert len(new_out) == 2\n\n        tensor, bias = new_out\n        assert tensor is not None\n\n        tensor = super()._reduce_local_output(in_id=in_id,\n                                              out_id=out_id,\n                                              current_out=old_tensor,\n                                              new_out=tensor)\n\n        if bias is None:\n            bias = old_bias\n\n        return tensor, bias\n\n    def _combine_output_splits(self, outputs):\n        # stack output tensors\n        tensors = [o[0] for o in outputs]\n        tensor = super()._combine_output_splits(tensors)\n\n        # stack biases if applicable\n        biases = [o[1] for o in outputs if o[1] is not None]\n        if len(biases) > 0:\n            bias = super()._combine_output_splits(biases)\n        else:\n            bias = None\n\n        return tensor, bias\n", "path": "deepspeed/runtime/zero/tiling.py"}], "after_files": [{"content": "import torch\nimport deepspeed\nfrom deepspeed.runtime.utils import partition_uniform as partition\n\n\ndef split_tensor_along_last_dim(tensor, partitions, contiguous_split_chunks=False):\n    \"\"\"Split a tensor along its last dimension. Adapted from Megatron-LM.\n\n    Arguments:\n        tensor: input tensor.\n        partitions: list of partition sizes to supply to torch.split\n        contiguous_split_chunks: If True, make each chunk contiguous\n                                 in memory.\n    \"\"\"\n    # Get the size and dimension.\n    last_dim = tensor.dim() - 1\n\n    # Split.\n    tensor_list = torch.split(tensor, partitions, dim=last_dim)\n    # Note: torch.split does not create contiguous tensors by default.\n    if contiguous_split_chunks:\n        return tuple(chunk.contiguous() for chunk in tensor_list)\n\n    return tensor_list\n\n\nclass TiledLinear(torch.nn.Module):\n    def __init__(self,\n                 in_features,\n                 out_features,\n                 bias=True,\n                 in_splits=1,\n                 out_splits=1,\n                 input_is_already_split=False,\n                 combine_out_splits=True,\n                 linear_cls=torch.nn.Linear,\n                 init_linear=None,\n                 **kwargs):\n        \"\"\"A replacement for ``torch.nn.Linear`` that works with ZeRO-3 to reduce\n        memory requirements via tiling.\n\n        TiledLinear breaks the input and output dimensions of a linear layer\n        into tiles that are processed in sequence. This class enables huge\n        linear layers when combined with ZeRO-3 because inactive tiles can be\n        partitioned and offloaded.\n\n        .. note::\n            We recommend using as few tiles as necessary. Tiling\n            significantly reduces memory usage, but can reduce throughput\n            for inexpensive layers. This due to the smaller kernels having\n            less parallelism and lower arithmetic intensity, while\n            introducing more frequent synchronization and communication.\n\n        Args:\n            in_features (int): See ``torch.nn.Linear``\n            out_features (int): See ``torch.nn.Linear``\n            bias (bool, optional): See ``torch.nn.Linear``\n            in_splits (int, optional): The number of tiles along the input dimension. Defaults to 1.\n            out_splits (int, optional): The number of tiles along the output dimension. Defaults to 1.\n            input_is_already_split (bool, optional): If set to ``True``, assume that the ``input_`` in\n                to ``forward()`` is already split into ``in_splits`` chunks. Defaults to ``False``.\n            combine_out_splits (bool, optional): If set to ``False``, do not combine the ``out_splits`` outputs\n                into a single tensor. Defaults to ``True``.\n            linear_cls (class, optional): The underlying class to build individual tiles.\n                Defaults to ``torch.nn.Linear``.\n            init_linear (``torch.nn.Linear``, optional): If set, copy the parameters of\n                ``init_linear``. Useful for debugging. Defaults to ``None``.\n            kwargs (dict, optional): additional keyword arguments to provide to ``linear_cls()``.\n\n        Raises:\n            RuntimeError: ``in_splits`` must be within the range [1, in_features).\n            RuntimeError: ``out_splits`` must be within the range of [1, out_features).\n        \"\"\"\n\n        super().__init__()\n\n        if (in_splits < 1) or (in_splits > in_features):\n            raise RuntimeError('in splits must be in range [1, in_features].')\n        if (out_splits < 1) or (out_splits > out_features):\n            raise RuntimeError('out splits must be in range [1, out_features].')\n\n        # global, not necessarily local\n        self.in_features = in_features\n        self.out_features = out_features\n        self.use_bias = bias\n\n        self.out_splits = out_splits\n        self.in_splits = in_splits\n        self.input_is_already_split = input_is_already_split\n        self.combine_out_splits = combine_out_splits\n\n        # Build partition-lists. These are CSR-style splits [0, part0, part1, ..., features]\n        # For example, row_parts[p] gives the start of partition p and row_parts[p+1]\n        # is the exclusive end.\n        self.in_parts = partition(num_items=in_features, num_parts=in_splits)\n        self.out_parts = partition(num_items=out_features, num_parts=out_splits)\n\n        assert len(self.out_parts) == out_splits + 1\n        assert len(self.in_parts) == in_splits + 1\n        assert self.out_parts[0] == 0\n        assert self.out_parts[out_splits] == out_features\n        assert self.in_parts[in_splits] == in_features\n\n        self.linears = torch.nn.ModuleList()\n        for out_id in range(out_splits):\n            self.linears.append(torch.nn.ModuleList())\n\n            local_out_dim = self.out_parts[out_id + 1] - self.out_parts[out_id]\n\n            for in_id in range(in_splits):\n                #if input_size is split, we only need one bias\n                local_bias = bias if in_id == (in_splits - 1) else False\n\n                local_in_dim = self.in_parts[in_id + 1] - self.in_parts[in_id]\n                local = linear_cls(local_in_dim,\n                                   local_out_dim,\n                                   bias=local_bias,\n                                   **kwargs)\n                self.linears[out_id].append(local)\n\n        # Optionally initialize with a known tensor\n        if init_linear is not None:\n            self.copy_params_from(init_linear)\n\n    def forward(self, input_):\n        if self.in_splits > 1 and not self.input_is_already_split:\n            input_parts = partition(input_.shape[-1], self.in_splits)\n            split_sizes = [\n                input_parts[p + 1] - input_parts[p] for p in range(self.in_splits)\n            ]\n            inputs = self._split_global_input(input_, split_sizes)\n        elif self.in_splits > 1:\n            inputs = input_\n            assert len(inputs) == self.in_splits, f\"Col splits {self.in_splits} does not match input splits {len(inputs)}\"\n        else:\n            # no splits\n            inputs = [input_]\n\n        outputs = [None] * self.out_splits\n        for out_id in range(self.out_splits):\n            for in_id in range(self.in_splits):\n                local_output = self.linears[out_id][in_id](inputs[in_id])\n\n                outputs[out_id] = self._reduce_local_output(in_id=in_id,\n                                                            out_id=out_id,\n                                                            current_out=outputs[out_id],\n                                                            new_out=local_output)\n\n        if self.combine_out_splits:\n            return self._combine_output_splits(outputs)\n\n        return outputs\n\n    def _split_global_input(self, input, split_sizes):\n        \"\"\"Partition an input tensor along the last dimension, aligned with given splits.\n\n        Subclasses should override this method to account for new input types.\n\n        Args:\n            input (List[Tensor]): The tensor to partition along the last dimension.\n            split_sizes (List[int]): The size of each partition.\n\n        Returns:\n            List[Any]: A list of the chunks of ``input``.\n        \"\"\"\n        return split_tensor_along_last_dim(input, split_sizes)\n\n    def _reduce_local_output(self, in_id, out_id, current_out, new_out):\n        \"\"\"Reduce (sum) a new local result into the existing local results.\n\n        Subclasses should override this method.\n\n        For a given ``out_id``, this method is called ``in_id-1`` times. The first input\n        split is a simple assignment.\n\n        Args:\n            in_id (int): The input split that produced ``new_out``.\n            out_id (int): The output split that produced ``new_out``.\n            current_out (Any): The reduced form of all previous ``out_id`` results.\n            new_out (Any): The local result from forward (``in_id``, ``out_id``)e\n\n        Returns:\n            Any: The combined result of ``current_out`` and ``new_out``.\n        \"\"\"\n\n        if current_out is None:\n            #this clone is necessary to preserve auto grad\n            #there is some issue with inplace update for outputs that are views\n            return new_out.clone()\n        else:\n            return current_out + new_out\n\n    def _combine_output_splits(self, outputs):\n        \"\"\"Join the splits of the output into a single result.\n\n        Args:\n            outputs (List[Any]): The reduced outputs for each output split.\n\n        Returns:\n            Any: The combined outputs.\n        \"\"\"\n        assert len(outputs) == self.out_splits\n        return torch.cat(outputs, dim=-1)\n\n    @torch.no_grad()\n    def copy_params_from(self, other):\n        \"\"\"Copy the weight and bias data from ``other``.\n\n        This is especially useful for reproducible initialization and testing.\n\n        Equivalent to:\n\n        .. code-block:: python\n\n            with torch.no_grad():\n                self.weight.copy_(other.weight)\n                if self.bias is not None:\n                    self.bias.copy_(other.bias)\n\n        .. note::\n            If ZeRO-3 is enabled, this is a collective operation and the\n            updated parameters of data-parallel rank 0 will be visible on all\n            ranks. See :class:`deepspeed.zero.GatheredParameters` for more\n            information.\n\n\n        Args:\n            other (``torch.nn.Linear``): the linear layer to copy from.\n        \"\"\"\n        assert hasattr(other, 'weight')\n        assert other.weight.size() == (self.out_features, self.in_features)\n        if self.use_bias:\n            assert hasattr(other, 'bias')\n            assert other.bias is not None\n            assert other.bias.size() == (self.out_features, )\n        else:\n            assert other.bias is None\n\n        for row in range(self.out_splits):\n            rstart = self.out_parts[row]\n            rstop = self.out_parts[row + 1]\n\n            for col in range(self.in_splits):\n                cstart = self.in_parts[col]\n                cstop = self.in_parts[col + 1]\n\n                local = self.linears[row][col]\n                global_weight = other.weight[rstart:rstop, cstart:cstop]\n                with deepspeed.zero.GatheredParameters(local.weight, modifier_rank=0):\n                    local.weight.copy_(global_weight)\n\n            if local.bias is not None:\n                with deepspeed.zero.GatheredParameters(local.bias, modifier_rank=0):\n                    local.bias.data.copy_(other.bias[rstart:rstop].data)\n\n\nclass TiledLinearReturnBias(TiledLinear):\n    \"\"\"Wrapper for a Linear class that returns its own bias parameter, such as\n    used by Megatron-LM.\n    \"\"\"\n    def _reduce_local_output(self, in_id, out_id, current_out, new_out):\n        \"\"\"Reduces output tensors, but not the returned bias. \"\"\"\n        if current_out is not None:\n            old_tensor, old_bias = current_out\n        else:\n            old_tensor, old_bias = None, None\n\n        assert isinstance(new_out, tuple)\n        assert len(new_out) == 2\n\n        tensor, bias = new_out\n        assert tensor is not None\n\n        tensor = super()._reduce_local_output(in_id=in_id,\n                                              out_id=out_id,\n                                              current_out=old_tensor,\n                                              new_out=tensor)\n\n        if bias is None:\n            bias = old_bias\n\n        return tensor, bias\n\n    def _combine_output_splits(self, outputs):\n        # stack output tensors\n        tensors = [o[0] for o in outputs]\n        tensor = super()._combine_output_splits(tensors)\n\n        # stack biases if applicable\n        biases = [o[1] for o in outputs if o[1] is not None]\n        if len(biases) > 0:\n            bias = super()._combine_output_splits(biases)\n        else:\n            bias = None\n\n        return tensor, bias\n", "path": "deepspeed/runtime/zero/tiling.py"}]}