From 57dee5a6467c2f5dd13009213497904693934aa4 Mon Sep 17 00:00:00 2001 From: Conglong Li Date: Thu, 21 Oct 2021 10:56:03 -0700 Subject: [PATCH 1/3] fix pp --- deepspeed/runtime/pipe/engine.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/deepspeed/runtime/pipe/engine.py b/deepspeed/runtime/pipe/engine.py index 6b56f8c7ad18..edec2b9fbc6a 100644 --- a/deepspeed/runtime/pipe/engine.py +++ b/deepspeed/runtime/pipe/engine.py @@ -1060,7 +1060,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 = ([outputs[0], *outputs[2:]]) + outputs = (outputs[0], *outputs[2:]) # save for backward self.pipe_buffers['outputs'][buffer_id] = outputs @@ -1072,8 +1072,7 @@ def _exec_recv_grads(self, buffer_id): dtype=outputs.dtype, num_buffers=1)[0] else: - sizes_and_dtypes = [(list(t.size()), - t.dtype) for t in outputs if t.is_floating_point()] + sizes_and_dtypes = [(list(t.size()), t.dtype) for t in outputs] self.grad_layer = self._allocate_buffers(sizes_and_dtypes, num_buffers=1)[0] From 285ea953beef171188057445db9999a081b17796 Mon Sep 17 00:00:00 2001 From: Conglong Li Date: Thu, 21 Oct 2021 12:15:08 -0700 Subject: [PATCH 2/3] better fix --- deepspeed/runtime/pipe/engine.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/deepspeed/runtime/pipe/engine.py b/deepspeed/runtime/pipe/engine.py index edec2b9fbc6a..9395ce1a687e 100644 --- a/deepspeed/runtime/pipe/engine.py +++ b/deepspeed/runtime/pipe/engine.py @@ -1072,7 +1072,12 @@ def _exec_recv_grads(self, buffer_id): dtype=outputs.dtype, num_buffers=1)[0] else: - sizes_and_dtypes = [(list(t.size()), t.dtype) for t in outputs] + if self.is_grad_partitioned: + sizes_and_dtypes = [(list(t.size()), t.dtype) for t in outputs] + else: + sizes_and_dtypes = [(list(t.size()), + t.dtype) for t in outputs + if t.is_floating_point()] self.grad_layer = self._allocate_buffers(sizes_and_dtypes, num_buffers=1)[0] From d40833a502039dbbd7af6f4b653721be9a49f361 Mon Sep 17 00:00:00 2001 From: Conglong Li Date: Thu, 21 Oct 2021 14:36:31 -0700 Subject: [PATCH 3/3] add hack comment --- deepspeed/runtime/pipe/engine.py | 14 ++++++++++++++ 1 file changed, 14 insertions(+) diff --git a/deepspeed/runtime/pipe/engine.py b/deepspeed/runtime/pipe/engine.py index 9395ce1a687e..8ecc0663fe95 100644 --- a/deepspeed/runtime/pipe/engine.py +++ b/deepspeed/runtime/pipe/engine.py @@ -1072,6 +1072,20 @@ def _exec_recv_grads(self, buffer_id): dtype=outputs.dtype, num_buffers=1)[0] else: + # XXX This is a HACK + # When we exchange activations/gradients, the two pipe stages + # need to issue the send/recv with the same buffer sizes or + # else there is a deadlock. The is_floating_point() filter is + # used to avoid sending gradients for tensors that do not + # produce gradients. When TP>1, we partition the first + # activations/gradients across TP ranks to save communication + # volume and memory. That partitioned tensor is represented as + # two tensors: a 1/TPth chunk of the original data and also a + # small LongTensor storing the metadata used to reconstruct on + # the other side. When combined, the floating point filter also + # filtered out the metadata tensor. A quick (hacky) fix just + # branches on is_grad_partitioned so we don't filter out the + # metadata tensor. if self.is_grad_partitioned: sizes_and_dtypes = [(list(t.size()), t.dtype) for t in outputs] else: