Repository: microsoft/DeepSpeed

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.



Discussion:
Can anyone help me look at this problem? Thanks in advance!
@weberxie, apologies for the delayed response. I can take a look at this. Can you please confirm that the issue also occurs with 1-gpu run? 
@tjruwase Thanks for your response. 

After change the MP_SIZE to 1 and DLTS_NUM_GPU_PER_WORKER to 1, this issue will not occur by it will OOM at optimizer stage at the 2nd step.
@weberxie, thanks for trying that out. Can you please reduce your model size to get past the OOM. You can cut the number of layers in half for example. 
@tjruwase thanks for your advice, it passed the OOM after reduced the model size.
@weberxie, so does that mean the runtime error does not show with 1-gpu? If so, can you please try 2-GPU and 4-GPU and share the runtime error message? Thanks.
> so does that mean the runtime error does not show with 1-gpu? 

yes.

I'll try 2-GPU and 4-GPU, and shall i set the MP_SIZE according to the number of GPU?
@tjruwase When MP_SIZE is 1, it will run successfully Whether using 2 GPUs or 4 GPUs. but it will throw the same Runtime error when MP_SIZE is 2 and 2 GPUs or MP_SIZE is 4 and 4 GPUs.
@weberxie, thanks for the providing these data points. It does confirm that the tensor splitting of these flags are responsible. I will take a closer look. 
@weberxie, I am able to repro this issue on my side. Thanks for providing the repro guidance.

Resolving patch:
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 @@ def split_tensor_along_last_dim(tensor, partitions, contiguous_split_chunks=Fals
     """
     # 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 __init__(self,
 
     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:


Test patch: