To train long-context models, or in some cases simply as an alternative to adding TP, we need to have an implementation of context parallelism. For MLP, this is simply sharding on the sequence dim. For attn, this is some variant of ring attention.
Scope:
- Gemma and Llama models for now.
- Some MFU degradation expected, but needs to work with compile and do reasonable overlapping.
Acceptance criteria:
- Allows training models with many times longer sequence length on LUMI.
- Relevant correctness test suite that passes.
Decision: previously discussed and already started.
To train long-context models, or in some cases simply as an alternative to adding TP, we need to have an implementation of context parallelism. For MLP, this is simply sharding on the sequence dim. For attn, this is some variant of ring attention.
Scope:
Acceptance criteria:
Decision: previously discussed and already started.