Add Sequence Parallelism to llama

[wanchaol]

Stack from ghstack (oldest at bottom):

• -> Add Sequence Parallelism to llama #32

Somehow the torch.compile not working although eager sequence
parallelism working, so currently don't turn it on by default

[wconstab]

I wonder if its useful to log more info about the SP plan. I was thinking about it for PP, what info do we want to print. Should each parallelism print its own summary, or should we have one overall function that prints overall parallel info in a unified way?

[wanchaol]

🤔 That's a good point. I think yeah we should probably log the parallelize plan for SP. This would require some changes in PyTorch to add __str__ to our ParallelStyles, I can add the log once the PyTorch PR is merged.

Should each parallelism print its own summary, or should we have one overall function that prints overall parallel info in a unified way

My two cents: It's a bit tricky to give overall summary. I think we can figure out how to even print the intended summary for each parallelism first, i.e. when transformerblock stacked too many, we can't log/print every layer parallel plan, so I think maybe we print pp degree of transformerblock, and we might not want to print the SP plan for each PP transformerblock.

[tianyu-l]

looks great to me! one inline question

[tianyu-l]

shall we also apply it on the final norm after all transformer blocks?

[wanchaol]

not sth currently enabled, but I think we can explore this in real training and see if shard the final norm would give additional memory/perf benefits :)