We support sequence parallelism (SP) via the +ring-flash-attention project. This +allows one to split up sequences across GPUs, which is useful in the event that a +single sequence causes OOM errors during model training.
+First, install ring-flash-attn, recommended via pip install axolotl[ring-flash-attn],
+or from source with pip install .[ring-flash-attn].
Your Axolotl YAML config should contain the following lines:
+sequence_parallel_degree: 4 # Split each sequence into 4 parts, one per GPU
+flash_attention: true # Required with sequence parallelism
+
+# Optional; strides across the key dimension. Larger values use more memory but will make training faster.
+heads_k_stride: 1See our dedicated guide for more details.
+For combining FSDP with QLoRA, see our dedicated guide.
Please see docs for more info.
For NCCL-related problems, see our NCCL troubleshooting guide.
To enable sequence parallelism, add the following to your configuration file:
# Set to a divisor (> 1) of the number of GPUs available
-sequence_parallel_degree: 4 # Split sequences across 4 GPUsThe sequence_parallel_degree should be a divisor of the total number of GPUs. For example:
# Example config with sequence parallelism
-base_model: meta-llama/Llama-3-8B-Instruct
-sequence_len: 8192
-sequence_parallel_degree: 2 # Split each sequence into 4 parts
-flash_attention: true # Required with sequence parallelism
-...base_model: meta-llama/Llama-3-8B-Instruct
+sequence_len: 8192
+
+...
+
+sequence_parallel_degree: 4 # Split each sequence into 4 parts, one per GPU
+flash_attention: true # Required with sequence parallelism
+# Optional; strides across the key dimension. Larger values use more memory but should make training faster.
+heads_k_stride: 1
+
+...This will train the Llama 3 8B model with 8K context length, with each sequence split into 2 subsequences of length 4096 across 2 GPUs.