背景

XTuner近期开源了长序列文本训练技术，该技术可以在多张显卡并行训练的条件下，将tokens均分到不同的显卡上，达到减少训练时显存占用的效果。魔搭社区正在和XTuner在训练推理技术层面进行深入合作，目前已经将该技术引入魔搭社区训练框架SWIFT，并支持社区LLM、MLLM（多模态大模型）及各类数据集的训练支持。

技术原理

该序列并行设计思路参考了 DeepSpeed 的工作 DeepSpeed Ulysses，并加以优化，以达到直接基于 transformers 算法库的开源模型训练 1M 以上超长序列的目标。

上图展示了序列并行策略的实现方案。由于 Transformer 结构较为规整，除 attention 计算外，其他计算过程中 token 之间不会互相影响（即每个 token 的计算是独立的），这一条件为序列并行提供了有利条件。上图展示了序列并行的核心设计。设由 P 个 GPUs 共同计算一个长度为 N 的长序列，在 Attention 计算的第一阶段，长度为 N / P 的子序列会通过线性层投影为 Query、Key、Value。接下来， QKV Tensor 会在参与序列并行计算的多个 GPUs 之间通过高度优化的 all-to-all 通信算子汇聚，得到序列长度为 N ，但更少注意力头的子序列。注意力计算后，通过另一个 all-to-all 通信算子将其转换为长度为 N / P 的子序列，进行后续计算。伪代码如下所示。

# Pseudo code for an Attention Layer
# Input: hidden_states with shape (bs, seq_len, dim)
# Output: attn_out with shape (bs, seq_len, dim)
def attn_forward(hidden_states):
    q, k, v = qkv_proj(hidden_states)
    q, k, v = reshape(q, k, v)  # (bs, q_len, dim) -> (bs, q_len, nhead, hdim)
    q, k = apply_rotary_pos_emb(q, k, cos, sin)
    sp_size = get_sequence_parallel_world_size()
    # (bs, q_len, nhead, hdim) -> (bs, q_len * sp_size, nhead / sp_size, hdim)
    q, k, v = all_to_all(q, k, v, sp_size)
    attn_out = local_attn(q, k, v)
    # (bs, q_len * sp_size, nhead / sp_size, hdim) -> (bs, q_len, nhead, hdim)
    attn_out = all_to_all(attn_out)
    attn_out = reshape(attn_out)  # (bs, q_len, nhead, hdim) -> (bs, q_len, dim)
    attn_out = o_proj(attn_out)
    return attn_out

通过上述技术方案，每个GPU上只会计算Sequence中的一部分token，在多进程条件下可以显著减少显存使用。

框架介绍

SWIFT

SWIFT是魔搭社区开发的一站式大模型训练框架，其主要特点包括：

1. 支持包含Llama3、Qwen、DeepSeek、Llava、Qwen-vl、DeepSeek-vl等一系列自然语言与多模态开源大模型
   
2. 支持近百个预训练、微调、对齐数据集的直接训练使用
   
3. 支持多种最新的技术，如Galore、LISA、Unsloth、LoRA、QLoRA等轻量级训练方式，以及VLLM部署方式
   
4. 支持使用界面直接训练和部署，支持评测能力
   

SWIFT官方REPO：

https://github.com/modelscope/swift

XTuner

XTuner是一站式大模型训练框架，其主要特点包括：

1. 支持包含Llava、Llama3、Qwen、DeepSeek、Gemma等一系列自然语言与多模态开源大模型
   
2. 支持LoRA、QLoRA、DeepSpeed等一系列轻量级训练方式
   
3. 支持XTuner自主开发的LMDeploy部署技术，在很多场景下可以超越VLLM的性能
   
4. 支持评测能力
   

XTuner的官方REPO：

https://github.com/InternLM/xtuner

在魔搭中使用序列并行技术

在本次的合作中，魔搭社区引入了XTuner底层的超长序列切分训练能力到SWIFT框架中，开发者可以使用魔搭社区的训练框架添加一个参数来实现超长序列的切分训练。

安装依赖：

pip install 'swift[seq_parallel]'

魔搭为序列并行技术提供了一个新的参数：

--sequence_parallel_size n 
# n设置为显卡数量，或者可以被显卡数量整除的正整数， 举个例子，比如显卡数量是4， 则n可以是4/2/1，当n是1时序列并行不生效

下面给出了一个完整的训练命令：

nproc_per_node=2
NPROC_PER_NODE=$nproc_per_node \
MASTER_PORT=29500 \
CUDA_VISIBLE_DEVICES=0,1 \
swift sft \
    --model_type chatglm3-6b-32k \
    --model_revision master \
    --sft_type lora \
    --tuner_backend peft \
    --dtype AUTO \
    --output_dir output \
    --ddp_backend nccl \
    --dataset AI-ModelScope/LongAlpaca-12k \
    --train_dataset_sample -1 \
    --num_train_epochs 1 \
    --max_length 16000 \
    --check_dataset_strategy warning \
    --lora_rank 8 \
    --lora_alpha 32 \
    --lora_dropout_p 0.05 \
    --lora_target_modules ALL \
    --gradient_checkpointing true \
    --batch_size 1 \
    --weight_decay 0.1 \
    --learning_rate 1e-4 \
    --gradient_accumulation_steps 1 \
    --max_grad_norm 0.5 \
    --warmup_ratio 0.03 \
    --eval_steps 10000 \
    --save_steps 10000 \
    --save_total_limit 2 \
    --logging_steps 10 \
    --save_only_model true \
    --sequence_parallel_size 1

上面的命令中我们针对AI-ModelScope/LongAlpaca-12k数据集做了特殊处理，使其每个语料输出统一长度，这样更方便统计训练时长和显存占用。

注意，序列并行需要在多卡环境下运行。

也可以在web-ui界面中使用这项功能或者进行实验：

swift web-ui

实验

下面针对不同的情况给出了显存占用和训练速度：

硬件环境：A100*4

理论上分析，由于序列长度折半，因此在训练时可以承载sequence_parallel_size*n尺寸的序列。我们对此进行了压力测试, 在脚本中添加如下参数：

--packing true # packing到max-length
--gpu_memory_fraction 0.3 # 限制显存在24G

之后动态调节max_length参数，可以观察到，使用DDP训练，chatglm3-6b模型可以最大支持9000的序列长度，使用sequence_parallel_size=2的双卡序列并行后，可以达到18000的序列长度，可训练序列增加到200%，这与我们的假设是相符的。

结论

通过实验结果可以看到：

• sequence_parallel_size在>1时显存占用比DDP低, 但会牺牲一部分的训练时间(实验中可以看到，约为20%的时间增加)
  
• 在sequence_parallel_size持续增加时，显存降低的的收益比会有所降低（sequence_parallel_size=4比sequence_parallel_size=2的单卡显存占用下降2G）
  

序列并行技术可以在长序列条件下减缓显存压力，一般在开发者有多张显卡，但显卡规格不高或序列长度过长显存吃紧时推荐使用，但需要注意序列并行的拆分数量，过多的数量的边际效益会有一定降低，建议开发者在使用时根据自己的硬件条件实际实验和调节。

对这项技术有兴趣的开发者欢迎进入我们的开发者群共同讨论，扫描下方二维码即可进入：