inner cache implemented #858
Conversation
| # Find my submodule | ||
| self.split_gm = self.pipe.split_gm | ||
| named_children = list(self.split_gm.named_children()) | ||
|
|
There was a problem hiding this comment.
@fduwjj comment is left on blank line. if it indicate L91-L92 or L95, we will leave it for safety reason until we modify all functions to use 'list of submods' instead of one submod.
There was a problem hiding this comment.
I am indeed asking about L95.
…last inner node (self.nodes[-1]) instead of self.node
|
code rebased |
pippy/PipelineStage.py
Outdated
| if self.inner_depth > 1: | ||
| # Forward pass of all chunks | ||
| for chunk in range(self.chunks): | ||
| s = self.streams[chunk % self.nstreams] | ||
| with torch.cuda.stream(s): | ||
| output, send_reqs = self.forward_one_chunk_ipipe( | ||
| chunk, args_split, kwargs_split, fwd_cache | ||
| ) | ||
| all_send_reqs += send_reqs | ||
| # Prepare for final output merge or reduction | ||
| output_chunks[chunk] = output | ||
| else: | ||
| # Forward pass of all chunks | ||
| for chunk in range(self.chunks): | ||
| s = self.streams[chunk % self.nstreams] | ||
| with torch.cuda.stream(s): | ||
| output, send_reqs = self.forward_one_chunk( | ||
| chunk, args_split, kwargs_split, fwd_cache | ||
| ) | ||
| all_send_reqs += send_reqs | ||
| # Prepare for final output merge or reduction | ||
| output_chunks[chunk] = output |
There was a problem hiding this comment.
The code has lots of duplication here, can you kindly consolidate a little bit?
There was a problem hiding this comment.
cleaned in the subsequent commit
| kwargs_split, | ||
| fwd_cache: Dict[int, Any], | ||
| ): | ||
| if self.rank == self.nstages - 1: |
There was a problem hiding this comment.
If I understand correctly, self.nstages is the global stages? If so does this line of code ever gets hit?
There was a problem hiding this comment.
self.nstages has been used in original codebase so i added a new variable self.global_depth to store global stages. so self.nstages is same as pp_group_size. (# of ranks)
|
|
||
| try: | ||
| if self.rank == self.nstages - 1: # last stage | ||
| output = self.forward_maybe_with_nosync( |
There was a problem hiding this comment.
Do you need a inner_depth == 0 check here?
There was a problem hiding this comment.
no we don't. inner_depth == 1 is default (no inner pipelining) and when inner_depth == 1, this function is not called. Instead, it calls forward_one_chunk function.
Description
instead of direct passing, use internal cache
Type of change
Checklist: