非对齐线性层
背景与挑战
类Megatron-LM框架已成为大模型训练的主流方案之一，张量并行 (Tensor Parallelism，TP) 是大模型训练的基本并行范式，但该范式在部分场景仍存在不足，例如要求大模型的注意力头数、序列长度能够整除TP，如不满足条件将在参数校验中抛出异常。
解决方案
本特性提供了如下解决方案：
- 序列长度不能整除TP时：和pad方案（将序列长度pad到TP的整数倍）不同，该方案通过序列分配策略来解决，小于 (seq_len%tp_size) 的tp_rank分配 (seq_len//tp_size+1) 序列长度，其他分配 (seq_len//tp_size) 序列长度。
例如seq_len=1026,tp_size=4, tp0和tp1分配的序列长度为257，tp2和tp3分配的序列长度为256。
- 注意力头数不能整除TP时：MHA（多头注意力）结构的模型，小于 (num_attention_heads%tp_size) 的tp_rank分配 (num_attention_heads//tp_size+1) 个注意力头，其他卡分配 (num_attention_heads//tp_size) 注意力头。例如num_attention_heads=25,tp_size=4, tp0分配的注意力头为7个，tp1、tp2和tp3分配的注意力头均为6个。
需要注意的是，对模型的注意力相关权重进行TP切分与注意力头数有关。假设hidden_size=3200，qkv_weight大小为(9600,3200)[MHA]，dense_weight大小为（3200,3200)，tp0的qkv权重大小为(2688,3200)，dense权重大小为(3200,896)，tp1、tp2及tp3的qkv权重大小为(2304, 3200)，dense权重大小为（3200,768）；GQA（分组查询注意力）结构的模型权重切分和注意力头按num_query_groups比例分配。
使用场景
序列长度或注意力头数不能整除TP的场景可以使用该特性。
使用方法
在训练脚本中设置如下参数即启用该特性。
--unaligned-linear
- 非对齐的线性层会引起各TP的负载不均衡。
- 该特性不支持2D张量并行、CP特性（要求TP*CP能被注意力头数整除）等。
- 对特殊的模型结构需要进行特殊适配，该特性当前已适配MHA、GQA结构，暂不支持MoE、MLA等结构。
使用效果
- 补充功能场景：补充注意力头数、序列长度不能被TP整除的场景。
- 潜在性能影响：各TP处理的注意力头数、序列长度不一致，负载不均衡，建议模型结构设计时考虑该情况。
该特性是为了完善TP（张量并行）场景下的限制约束，特性本身会带来负载不均衡的性能影响，所以在模型设计和超参优化时注意这一影响。
父主题： 多模态能力支持