selective 2d api/example added for fine-grained tp/pp demo #830
Conversation
examples/selective2d/2d_train.py
Outdated
| from pippy.microbatch import TensorChunkSpec, sum_reducer | ||
|
|
||
| pp_dim, tp_dim = 0, 1 | ||
| pp_rank, tp_rank = args.local_rank // args.tp_size, args.local_rank % args.tp_size |
There was a problem hiding this comment.
What if we are doing this on multiple host?
There was a problem hiding this comment.
in multiple host case, we should use args.rank instead of args.local_rank. i'll change the line to handle multiple hosts in the next commit
|
|
||
| def __init__(self, mesh, config): | ||
| super().__init__() | ||
| assert config.n_embd % config.n_head == 0 |
There was a problem hiding this comment.
Let's also add an assert for self.n_head % tp_size == 0?
examples/selective2d/2d_train.py
Outdated
|
|
||
| # PP | ||
| cut_fn(model, args, args.pp_size) | ||
| stage = compile_stage(model, pp_rank, args.world_size, args.pp_size, args.device, pp_groups, |
There was a problem hiding this comment.
For my learning purpose, does it work if we first do TP and then call compile_stage? So DTensor can already be traced by torch.fx?
There was a problem hiding this comment.
for now we apply TP first and then PP --> so Dtensor is traced by torch.fx
if we do PP first and TP, the result does not change but PP changes layer name after tracing
(e.g., transformer.block.i.attn --> transformer_block_i_attn) so we should change the name that we pass to TP.
There was a problem hiding this comment.
I see. Good to know that DTensor is traceable by torch.fx. I have a n00b question here, what's the difference between first applying TP vs first applying PP?
There was a problem hiding this comment.
performance-wise, there is no difference. i think only difference is api. pippy's compile stage breaks higher class (e.g., block/transformer/MLP/Attention) into low-level layers (linear) so we just need to be careful of layer name.
examples/selective2d/2d_train.py
Outdated
| from pippy.IR import annotate_split_points, PipeSplitWrapper | ||
| from pippy import split_into_equal_size | ||
| from pippy.compile import compile_stage | ||
| from pippy.microbatch import TensorChunkSpec, sum_reducer |
There was a problem hiding this comment.
Import should always live in the top of the file.
examples/selective2d/2d_train.py
Outdated
|
|
||
| return model, stage | ||
|
|
||
| def even_cut(model, args, pp_size, cut={}): |
There was a problem hiding this comment.
Should add a one line docstring to describe the function.
examples/selective2d/2d_train.py
Outdated
|
|
||
| annotate_split_points(model, cut) | ||
|
|
||
| def after_ar_cut(model, args, pp_size, cut={}): |
There was a problem hiding this comment.
Should add a one line docstring to describe the function.
examples/selective2d/2d_train.py
Outdated
|
|
||
| annotate_split_points(model, cut) | ||
|
|
||
| def pp_and_tp_fg(model, mesh, args, tp_attn_layers=None, tp_mlp_layers=None, cut_fn=even_cut): |
There was a problem hiding this comment.
The naming is pretty confusing. What is fg?
There was a problem hiding this comment.
Should add a one line docstring to describe the function.
examples/selective2d/2d_train.py
Outdated
|
|
||
| def pp(model, pp_device_mesh, args): | ||
| from pippy.IR import annotate_split_points, PipeSplitWrapper | ||
| from pippy import split_into_equal_size |
There was a problem hiding this comment.
Any reason why this is imported but not used?
There was a problem hiding this comment.
My bad. I used in my draft and moved on to my own cut. Will remove.
examples/selective2d/2d_train.py
Outdated
| return model, stage | ||
|
|
||
| def even_cut(model, args, pp_size, cut={}): | ||
| from pippy.IR import annotate_split_points, PipeSplitWrapper |
There was a problem hiding this comment.
Move the import to the top of the file.
examples/selective2d/2d_train.py
Outdated
|
|
||
| return model, stage | ||
|
|
||
| def even_cut(model, args, pp_size, cut={}): |
There was a problem hiding this comment.
Why cut is passed as an argument? And using {} as the default value is never good.
There was a problem hiding this comment.
I was thinking mixing two algorithms but at this point we may not need it. removed :)
examples/selective2d/2d_train.py
Outdated
| return local_iter_num, iter_time | ||
|
|
||
| def tp_train(): | ||
| local_iter_num = 0 |
There was a problem hiding this comment.
The inconsistent indentation is never good.
## Description added 2d parallelism (tp+pp) API and example for fine-grained tp/pp ## Checklist: - [v] Has code been commented, particularly in hard-to-understand areas? - [v] Have you made corresponding changes to the documentation?
## Description added 2d parallelism (tp+pp) API and example for fine-grained tp/pp ## Checklist: - [v] Has code been commented, particularly in hard-to-understand areas? - [v] Have you made corresponding changes to the documentation?
Description
added 2d parallelism (tp+pp) API/example for fine-grained tp/pp
Checklist: