[Parallelism] Implement vocabulary parallelism

[casper-hansen]

Balancing Pipeline Parallelism with Vocabulary Parallelism introduces a way to handle vocabulary scaling together with PP. While context parallelism splits the sequence dimension across devices, Vocabulary Parallelism specifically splits the vocabulary dimension of embedding layers.

Results:

[gnadathur]

cc: @H-Huang

[ParamThakkar123]

@casper-hansen @gnadathur @H-Huang may I take on this issue ??

[ParamThakkar123]

This is how I am planning to implement this:

1. make a new file vocab_parallel.py
2. Implement Vocabulary Parallel Layers
3. Modify Pipeline Schedule to add a custom schedule that supports vocabulary passes
4. Update LLaMA Pipeline Logic to add vocab_parallelism support
5. Update JobConfig to add new field to enable or disable vocab_parallelism

[tianyu-l]

@ParamThakkar123
Thanks for your interest!

Instead of working directly on top of current code, it's better to have a dedicated folder for your contribution, where you can create a new TrainSpec. In the case of Vocab Parallel, a simple & ideal way to do it is to substitute the pipelining_fn while keeping other parts.

In fact we'll be encouraging such explorations soon, if contributors agree to maintain the code. The guideline will come out shortly. Please stay tuned.

[tianyu-l]

@ParamThakkar123 The experiments folder (torchtitan/experiments) is almost ready to accept contributions (we still need to support flexible custom JobConfig, but it'll come shortly). Please consider starting submitting PRs following the contributing guidelines. Please let us know for any concerns you may have!