Enable DualPipeV by adding a multiplexed graph #258
Conversation
stack-info: PR: #258, branch: sanketpurandare/stack/1
sanketpurandare
force-pushed
the
sanketpurandare/stack/1
branch
from
26237a1 to
ab091ce
Compare
meta-cla
bot
added
the
CLA Signed
sanketpurandare
changed the title
| insert_point = n | ||
| if n.op == "output": | ||
| multiplexed_graph_op_node = n | ||
| multiplxed_gm_placeholders = multiplexed_gm.graph.find_nodes(op="placeholder") |
| # Collect output arguments from forward graph | ||
| # Collect output arguments from multiplexed graph (will contain only fwd_outs) | ||
| multiplexed_graph_outputs = multiplexed_gm.graph.find_nodes(op="output") | ||
| assert len(multiplexed_graph_outputs) == 1 |
There was a problem hiding this comment.
why this assert? is this a tuple of forward outs or are you hardcoding the model arch into the infra?
There was a problem hiding this comment.
graph.find_nodes() returns a list of nodes with the filter criteria, this assert is for sanity check that there is exactly one output node. It says nothing about the args of the output node
| ) -> tuple[list[Any], list[Any], Any, tuple[tuple[Any], tuple[Any]]]: | ||
| multiplexed_outs = fx.Interpreter(multiplexed_fw_bw_module).boxed_run(bw_fw_args) | ||
|
|
||
| num_params_buffers = bw_graph_meta.num_params + bw_graph_meta.num_buffers |
There was a problem hiding this comment.
maybe in a later pr, it would be nice if the decoding of the multiplexed_outs could be done by a module that was generated in the same place that generated the multiplexed graph. i mostly don't like how this function hardcodes a specific 'API' that's defined somewhere else and there isn't any enforcement that it stays consistent.
| numerics_logs: Optional[list[str]] = None, | ||
| ) -> None: | ||
| ) -> tuple[ | ||
| _PipelineScheduleRuntime, |
There was a problem hiding this comment.
its a bit odd to me that this util function is (a) not a class method on the runtime, and (b) returns a runtime object among other things. (from the description, mostly doing setup stuff and waiting ops, it seems this function could plausibly have a null return type)
| mb_index, | ||
| is_next_stage_on_this_rank, | ||
| is_prev_stage_on_this_rank, | ||
| ) = _prepare_fwd_common(action, ctx) |
There was a problem hiding this comment.
i see why its this way now. but maybe you can
(a) get any widely used objects off ctx into local vars inside stage_forward and pass those into the util functions
(b) just access the less widely used objects directly off ctx inside the utils to reduce the number of passed around locals / API surface
There was a problem hiding this comment.
Will try some this refactoring in the next PR
| ), f"num_params_buffers: {num_params_buffers}, num_params: {num_params}, num_buffers: {num_buffers}" | ||
| num_input_grads = ( | ||
| len(bw_module.graph.find_nodes(op="output")[0].args[0]) - num_params_buffers | ||
| ) |
There was a problem hiding this comment.
i'm probably getting tripped up by the naming convention here, num_input_grads is generally filtered from num_fwd_outputs right?
There was a problem hiding this comment.
You can get the number of input_grads either from the placeholders of the fwd_graph or the outputs of the bwd_graph. num_fwd_outputs will give you the num_out_grads (tangents) that will be passed in as bwd_graph inputs.
Stacked PRs:
Enable DualPipeV by adding a multiplxed graph that currently doesn't do overlapping but executes the backward and forward of different stages and microbatches in a single multiplexed graph.