Enabling ZeroBubbleV schedule in Graph PP #250
Conversation
[ghstack-poisoned]
meta-cla
bot
added
the
CLA Signed
autoparallel/graph_pp_runner.py
Outdated
| assert not any( | ||
| grad is None for grad in grads_to_accumulate | ||
| ), "All grads are None" |
There was a problem hiding this comment.
| assert not any( | |
| grad is None for grad in grads_to_accumulate | |
| ), "All grads are None" | |
| assert not all( | |
| grad is None for grad in grads_to_accumulate | |
| ), "All grads are None" |
There was a problem hiding this comment.
actually shouldn't we keep it as 'any' and change the string to match? or do we not care if some grads are none
There was a problem hiding this comment.
Yeah we should keep as any and change the string.
There was a problem hiding this comment.
I think some grads can be None, say forward outputs that don't require gradient, the usual torch.compile fw/bw would have None in the backward graph outputs as required by the custom autograd function API.
It's up to the partitioner/graph pass splitters and the runtime wrapper implementation though. As long as we know which grad belongs to which param.
There was a problem hiding this comment.
Switched it back to any
autoparallel/graph_pp_runner.py
Outdated
| assert num_placeholders == len( | ||
| fw_args | ||
| ), f"Mismatched number of inputs to fwd, {len([n for n in fw_module.graph.nodes if n.op == 'placeholder'])}, {len(fw_args)}" | ||
| ), f"Mismatched number of inputs to fwd: expected {num_placeholders}, got {len(fw_args)}" |
There was a problem hiding this comment.
no longer needed, we made the change upstream
There was a problem hiding this comment.
Code pointer?
There was a problem hiding this comment.
Removed the args len check
[ghstack-poisoned]
sanketpurandare
changed the base branch from
gh/sanketpurandare/2/base
to
main
sanketpurandare
changed the base branch from
main
to
gh/sanketpurandare/2/base
sanketpurandare
changed the base branch from
gh/sanketpurandare/2/base
to
main
sanketpurandare
changed the base branch from
main
to
gh/sanketpurandare/2/base
sanketpurandare
changed the title
Stack from ghstack (oldest at bottom):