VLLM适配新模型

Written by with ♥ on in 默认分类

前言

最近在尝试预训练大模型,想将我们的模型适配到VLLM框架上,所以也对VLLM框架的代码进行了阅读总结,下面给出一点经验.

将Transformers库中Modeling转换为VLLM模型

如VLLM官方文档所言,大部分Model的适配需要参考Transformer库中的modeling文件,大部分结构是类似的,但是需要额外注意以下几点:

  • Transformers库会使用模型对应的Config类来反序列化JSON文件,但是VLLM会直接将JSON配置反序列化成一个字典,以固定的Key来获取配置值.所以如果你要支持的模型并不是以标准的Huggingface约定的风格来规定config.json,你可能要在如get_head_size()等函数中打上若干补丁. 👇🏻
def get_head_size(self) -> int:
        # TODO remove hard code
        if hasattr(self.hf_text_config, "model_type"
                   ) and self.hf_text_config.model_type == 'deepseek_v2':
            # FlashAttention supports only head_size 32, 64, 128, 256,
            # we need to pad head_size 192 to 256
            return 256
  • VLLM只关注推理,因此在适配模型时只需要关注两件事情: Forward和Load Weights. 大部分Modeling函数中定义的Class是可以很好的复用的,但是你可能要修改一部分函数签名. 每个Model的组成部分都需要继承nn.Module并记得调用父类构造函数.一个经典的例子如下:
class MyModelForCausalLM(nn.Module):
    def __init__(
        self,
        config,
        cache_config: Optional[CacheConfig] = None,
        quant_config: Optional[QuantizationConfig] = None,
    ):
        super().__init__()
        self.config = config
        self.quant_config = quant_config
        self.model = MyModel(config, cache_config, quant_config)
        if self.config.tie_word_embeddings:
            self.lm_head = self.model.decoder.embed_tokens
        else:
            self.lm_head = ParallelLMHead(config.vocab_size,
                                          config.word_embed_proj_dim)
        self.logits_processor = LogitsProcessor(config.vocab_size)
        self.sampler = Sampler()

    def forward(
        self,
        input_ids: torch.Tensor,
        positions: torch.Tensor,
        kv_caches: List[torch.Tensor],
        attn_metadata: AttentionMetadata,
        intermediate_tensors: Optional[IntermediateTensors] = None,
    ) -> torch.Tensor:
        hidden_states = self.model(input_ids, positions, kv_caches,
                                   attn_metadata)
        return hidden_states
  • 替换Transformer Blocks中的模块.VLLM始终是为多卡并行而设计的,而且基于算子融合的理念把大部分常用算子进行了整合.一些常用的替代方案是:
LMHead => ParallelLMHead
Embedding => VocabParallelEmbedding
RMSNorm => vllm.model_executor.layers.layernorm.RMSNorm
ROPE => get_rope()

对于激活函数,VLLM提供了Silu+MatrixMul和GELU+MatrixMul两种融合算子(用CUDA实现),另外也用CUDA写了单个的激活函数算子.如果你的模型用到了比较奇怪的自定义激活函数,那记得写好Forward函数之后加上@torch.jit.script 来优化下.

VLLM中的各种Linear层

在适配模型的时候,需要将Modelling代码中的Linear更换为各种可并行的Linear.VLLM提供了各种Linear,下面分析下他们的实现.

LinearBase和LinearMethod

LinearBase是一个nn.Module ,它通过quant_conf 参数来决定是否调用量化相关的矩阵乘算子.当不需要量化时,则调用 UnquantizedLinearMethod 来执行计算. LinearMethod类实际上是Torch Fx函数式的封装,UnquantizedLinearMethod会调用F.linear来进行矩阵乘计算.之所以使用函数式来进行实际计算,是因为当涉及到分布式并行时,Weights和Input Tensor通常是被切分的向量,自行维护Weights的加载和切分可能更有效率.

一个LinearMethod 类包含create_weights (依据Linear层的Input/Output Size和Parameter Type等参数创建Weight向量并注册到相关的LinearBase Module中,并设置相关的Weights Loader.

ReplicatedLinear

朴实无华的线性层,每个分布式节点执行相同的操作,存在大量重复计算.

ColumnParallelLinear

如图所示,对于常见的 形如 Y=AX+b类型的矩阵乘法,ColumnParallelLinear将权重A按照第二维(如果LinearMethod涉及量化或者Lora合并,维度数目可能不同)进行切分,每个卡只执行一部分形如 $A_1X+b$的操作.但是注意,在执行完成之后,每个卡上的$Y_i$的Shape和内容均只有原Y的一部分,并不会自动Gather.这是考虑到TransformerBlock中的并行操作不需要完整的Tensor,如MLP层等.

Image.png

实现 : 先根据PP或TP的WorldSize得出分片i的数目并将output_size切分为output_size_per_partition来设置Weights Tensor和输出. 同样因为切分了权重向量,因此重写Weights Loader,计算当前卡对应的Weights Tensor的子向量Offeset.

适用: 一般适用于MLP的前半个Linear.

MergedColumnParallelLinear

和上面提到的ColumnParallelLinear的计算流程是完全相同的,但是它用于将多个Linear融合在一起进行运算(比如MLP中Gate和Up两个Linear可以进行融合).因此它在Weights加载过程中不仅需要分加载每一行权重,更要把权重在Output这个维度进行concat,来保证算子融合的正确性.此外在Weights的加载过程中加入了对用于量化的PackedParameter的支持.

适用: 一般适用于MLP的Gate-Up 融合.

QKVParallelLinear

见名知意,是将三个Linear融合在一起,而且针对于QKV这三个Linear进行特化,输出是将qkv拼接在一起的Tensor.相比于MergedColumnParallelLinear的区别是shared的offset可以根据QKV HeadNum以及HeadSize来计算得出.

tag(s): none
show comments · back · home
Edit with markdown