load value from pipe_cache

pytorch · Aug 8, 2023 · 6fdbc19 · 6fdbc19
1 parent 84f765b
commit 6fdbc19
Showing 1 changed file with 4 additions and 4 deletions.
diff --git a/pippy/PipelineStage.py b/pippy/PipelineStage.py
@@ -699,13 +699,13 @@ def forward_one_chunk_ipipe(
 
                 for i in range(1, self.inner_depth-1):
                     output = self.forward_maybe_with_nosync(
-                        i, targets, output, **composite_kwargs
+                        i, targets, self.pipe_cache[chunk][i-1], **composite_kwargs
                     )
                     self.pipe_cache[chunk][i] = output
 
                 output = self.forward_maybe_with_nosync(
-                    self.inner_depth-1, targets, output, targets, **composite_kwargs
-                )
+                    self.inner_depth-1, targets, self.pipe_cache[chunk][self.inner_depth-2], targets, **composite_kwargs
+                ) # self.inner_pipe >= 2 is asserted
                 self.pipe_cache[chunk][self.inner_depth-1] = output
             else:
                 # 0th (first) inner node
@@ -717,7 +717,7 @@ def forward_one_chunk_ipipe(
                 # other inner nodes uses 'output' of previous inner node
                 for i in range(1, self.inner_depth):
                     output = self.forward_maybe_with_nosync(
-                        i, targets, output, **composite_kwargs
+                        i, targets, self.pipe_cache[chunk][i-1], **composite_kwargs
                     )
                     self.pipe_cache[chunk][i] = output