pipe partitioning

Author: Shaden Smith

deepspeed/runtime/pipe/engine.py

@@ -111,8 +111,8 @@ def __init__(self, *super_args, **super_kwargs):
 
         # Partition input/output buffers
         # XXX temporarily disable while I revert some partition hacks.
-        self.is_pipe_partitioned = False #self.is_model_parallel
-        self.is_grad_partitioned = False
+        self.is_pipe_partitioned = self.is_model_parallel
+        self.is_grad_partitioned = False #self.is_model_parallel
 
         model_parameters = filter(lambda p: p.requires_grad, self.module.parameters())
         num_params = sum([p.numel() for p in model_parameters])
@@ -554,12 +554,18 @@ def _exec_forward_pass(self, buffer_id):
                 local_part=inputs[1],
                 group=self.grid.get_slice_parallel_group())
 
+            inputs = part_input.full()
+            inputs.requires_grad = True
+            part_input = None
+            self.pipe_buffers['inputs'][buffer_id] = inputs
+            '''
             inputs = tuple([part_input.full(), inputs[2]])
             inputs[0].requires_grad = True
             # skip mask
             #inputs[1].requires_grad = True
             part_input = None
             self.pipe_buffers['inputs'][buffer_id] = inputs
+            '''
 
         # Zero out the gradients each time we use the tensor because only the data in
         # tensor changes across batches
@@ -569,13 +575,14 @@ def _exec_forward_pass(self, buffer_id):
 
         # Partition the outputs if we are not the last stage
         if self.is_pipe_partitioned and not self.is_last_stage():
-            part = PartitionedTensor(tensor=outputs[0],
+            assert torch.is_tensor(outputs)
+            part = PartitionedTensor(tensor=outputs,
                                      group=self.grid.get_slice_parallel_group())
             # Clear the large output data, but save the computation graph
-            outputs[0].data = torch.zeros(1)
-            self.pipe_buffers['output_tensors'][buffer_id] = outputs[0]
+            outputs.data = torch.zeros(1)
+            self.pipe_buffers['output_tensors'][buffer_id] = outputs
             # Inject the partitioned tensor into the output before sending
-            outputs = tuple([part.to_meta(), part.data(), outputs[1]])
+            outputs = tuple([part.to_meta(), part.data()])
             part = None
 
         self.pipe_buffers['outputs'][buffer_id] = outputs
@@ -633,15 +640,11 @@ def _exec_backward_pass(self, buffer_id):
                     local_part=outputs[1],
                     group=self.grid.get_slice_parallel_group())
                 self.pipe_buffers['output_tensors'][buffer_id].data = part_output.full()
-                outputs = tuple(
-                    [self.pipe_buffers['output_tensors'][buffer_id],
-                     outputs[2]])
+                outputs = self.pipe_buffers['output_tensors'][buffer_id]
             else:
                 # Already restored from partition
-                self.pipe_buffers['output_tensors'][buffer_id].data = outputs[0]
-                outputs = tuple(
-                    [self.pipe_buffers['output_tensors'][buffer_id],
-                     outputs[1]])
+                self.pipe_buffers['output_tensors'][buffer_id].data = outputs
+                outputs = self.pipe_buffers['output_tensors'][buffer_id]
 
         grad_tensors = self.grad_layer
         if self.is_grad_partitioned:
@@ -650,7 +653,7 @@ def _exec_backward_pass(self, buffer_id):
                 meta=self.grad_layer[0],
                 local_part=self.grad_layer[1],
                 group=self.grid.get_slice_parallel_group())
-            grad_tensors = tuple([part_grad.full(), self.grad_layer[2]])
+            grad_tensors = part_grad.full()
             part_grad = None
             #print(f'RANK={self.global_rank} BEFORE-BWD restored grad={self.grad_layer[0].size()} {self.grad_layer[1].size()}')
 
@@ -873,13 +876,10 @@ def _exec_send_grads(self, buffer_id):
 
         # Partition the gradient
         if self.is_grad_partitioned:
-            part = PartitionedTensor(tensor=inputs[0].grad,
+            assert torch.is_tensor(inputs)
+            part = PartitionedTensor(tensor=inputs.grad,
                                      group=self.grid.get_slice_parallel_group())
-            # Clear the large output data, but save the computation graph
-            # Inject the partitoned tensor into the output before sending
-
-            # XXX Hack
-            inputs = tuple([part.to_meta(), part.data(), inputs[1]])
+            inputs = tuple([part.to_meta(), part.data()])
 
         # XXX Terrible hack
         # Drop the attention mask from the input buffer here. It does not have
@@ -900,8 +900,6 @@ def _exec_send_grads(self, buffer_id):
                 # First two sends are partitioned gradient
                 p2p.send(inputs[0], self.prev_stage)
                 p2p.send(inputs[1], self.prev_stage)
-                # XXX hack hack hack
-                #p2p.send(inputs[2].grad, self.prev_stage)
             else:
                 for idx, buffer in enumerate(inputs):
                     # Skip tensors that will not produce a grad
@@ -975,7 +973,7 @@ def _exec_recv_grads(self, buffer_id):
                 local_part=outputs[1],
                 group=self.grid.get_slice_parallel_group())
             outputs[0].data = part_output.full()
-            outputs = tuple([outputs[0], outputs[2]])
+            outputs = outputs[0]
             # save for backward
             self.pipe_buffers['outputs'][buffer_id] = outputs
 
@@ -985,7 +983,7 @@ def _exec_recv_grads(self, buffer_id):
                 s = list(outputs.size())
                 self.grad_layer = self._allocate_buffer(s, num_buffers=1)[0]
             else:
-                sizes = [list(t.size()) for t in outputs if t.is_floating_point()]
+                sizes = [list(t.size()) for t in outputs]# if t.is_floating_point()]
                 self.grad_layer = self._allocate_buffers(sizes, num_buffers=1)[0]
 
         if isinstance(self.grad_layer, torch.Tensor):
@@ -999,7 +997,7 @@ def _exec_recv_grads(self, buffer_id):
                                               dtype=torch.long,
                                               device=self.device)
                 p2p.recv(buffer, self.next_stage)
-
+            
         if self.wall_clock_breakdown():
             self.timers('pipe_recv_grad').stop()