LiBai下分布式推理说明

[CPFLAME]

前言

LiBai发布新模块, 分布式推理

• 适用场景:

  • 当一张GPU卡无法容纳一个大模型的权重时, 需要把权重切分到多个卡上, 来完成模型的推理

• 支持的功能:

  • 模型并行和流水并行的推理
  • 流水并行下可以手动设置每张卡下的layer层数, 方便进行负载均衡
  • 同时支持libai的model, 和torch的model进行分布式推理(需要重载加载模型的模块)

• demo的演示

  一份简单的多机多卡分布式(模型并行2X流水并行2)运行代码

  # test_inference.py
  from libai.inference.text_generation import TextGenerationPipeline
  from libai.utils import distributed as dist
  
  if __name__ == "__main__":
    pipeline = TextGenerationPipeline(
        "projects/MT5/configs/t5_inference.py",
        data_parallel=1,
        tensor_parallel=2,
        pipeline_parallel=2,
        pipeline_stage_id=[0] * 12 + [1] * 12,
        pipeline_num_layers=12 * 2,
        model_path="data_test/t5_inference_model",
        mode="huggingface",
    )
  
    text = ["summarize: She is a student, She is tall, She loves study"]
    dict1 = pipeline(text)
    if dist.is_main_process():
        print(dict1)

  多机多卡的分布式推理脚本

  在node0上输入指令:

  NODE=2 NODE_RANK=0 ADDR=192.168.0.1 PORT=12345 bash tools/infer.sh test_inference.py 2

  在node1上输入指令:

  NODE=2 NODE_RANK=1 ADDR=192.168.0.1 PORT=12345 bash tools/infer.sh test_inference.py 2

如何使用pytorch的模型进行分布式推理?

在LiBai下面加载torch的模型进行分布式推理之前, 我们需要有一些预备知识

预备知识

一个完整的模型由两个部分构成:

• 模型结构, 换种说法就是model.py
• 模型权重, 换种说法就是model_best.pth

那么假设我们在框架A下面, 有modelA.py和model_best_A.pth, 我们想在框架B上面跑起来这个框架A下面的模型, 应该怎么做呢?

1. 在框架B下面, 用框架B的算子搭建出一个modelB.py, 该modelB的参数名字可以和modelA的不一致, 但是前向推理的逻辑运算最好一致
2. 然后去加载model_best_A.pth得到model_A_state_dict(), 把model_A_state_dict()里面的参数格式全部转换成框架B下面支持的格式, 其中可以运用中间格式进行转换. 举个例子, 比如torch.tensor()->np.numpy()(中间格式)->oneflow.tensor()
3. 之前提到了modelB中的参数名字可以和modelA中的不一致, 如果不一致的话, 那么我们需要把model_A_state_dict()中的key值改一下和modelB的一致.
4. 做完了以后, 直接加载我们转换好的参数modelB.load_state_dict(model_A_state_dict, ), 就可以在框架B下面进行推理了.
5. 为了保证模型转换好以后的准确性, 可以喂给modelA以及modelB相同的输入, 检查一下是否能得到相同的输出.

在LiBai下面跑pytorch模型的分布式推理

在有了预备知识了以后, 再来看看怎么在LiBai下面跑pytorch模型的分布式推理.

主要分为以下的几步:

1. 把torch的算子替换为oneflow: 把torch_model.py下面的torch全部替换为oneflow, 得到oneflow_model.py.
2. 把oneflow_model.py中的layer尽可能的替换成LiBai中支持的layer.
3. 写好转换权重的脚本, 可以参考ModelLoaderHuggerFace.
4. 编写model_config.py, 这个部分比较简单, 如果没有完整训练的train_config.py也没有关系, 只用写一份调用model的config就可以了, 可以参考dalle2_config.py
5. 继承libai/inference/basic.py, 写好预处理和后处理, 把第三步写好的转换权重的脚本重载到方法load_pretrain_weight中.
6. 进行分布式推理.

步骤1:把torch的算子替换为oneflow

得益于oneflow和pytorch的api高度一致, 通常情况下, 可以直接把torch代码中的import torch, 直接替换成import oneflow, 正常情况下面, 可以把torch_model.py下面所有的torch关键字全部替换为oneflow.

在成功跑通了以后, 我们就有了一份以oneflow为底层算子的单机代码oneflow_model.py了.

如果没有办法跑通, 可以在报错的地方查看一下, 可能是算子没有对齐, 这个时候可以去oneflow下提issue, 或者直接换另外一种不影响运行结果的写法直接绕过去.

步骤2: 替换为LiBai中的layer

在LiBai下面提供了封装好的layers可以进行调用, 这些layers是基于oneflow的再次封装, 目的是在于用户可以直接调用这些封装好的模块, 而不用自己去定义sbp和placement, 这些layer是提供分布式推理的关键.

libai/layers
├── activation.py
├── attention.py
├── cross_entropy.py
├── droppath.py
├── embedding.py
├── __init__.py
├── layer_norm.py
├── linear.py
├── lm_logits.py
├── mlp.py

拿一个最常用也最实用的例子: 在transformer的模型中, 一般Linear层是运用最多的层, 也是最适合切分其权重到不同的gpu上面, 以完成分布式训练或者推理的层.

得益于LiBai的设计, 用户可以把模型中的layer部分替换为LiBai的layer, 其他的代码可以不动. 所以对于用户来说, 如果Linear层是单卡GPU放不下的瓶颈, 那么用户可以只把代码中的Linear层用LiBai的Linear层进行替换.

其中LiBai的Linear层提供了两种并行方式parallel="row"和parallel="col", 分别代表着模型并行下, 参数是按照按行切分以及以及按列切分. 通常来说, 多个连续的Linear层, 按照col->row->col->row->col->....切分方式的运算效率是比较高的.

# torch_model.py
from torch import nn
from torch.nn import Linear

class MultiheadAttention(nn.Module):
    def __init__(hidden_size):
        self.query_key_value = Linear(
            self.hidden_size,
            self.hidden_size * 3,
        )

        self.dense = Linear(
            self.hidden_size,
            self.hidden_size,
        )
    ...

###########替换为oneflow的算子##########
# oneflow_model.py
from oneflow import nn
from oneflow.nn import Linear

class MultiheadAttention(nn.Module):
    def __init__(hidden_size):
        self.query_key_value = Linear(
            self.hidden_size,
            self.hidden_size * 3,
        )

        self.dense = Linear(
            self.hidden_size,
            self.hidden_size,
        )
    ...

### 替换为`LiBai`中的`layer`
# libai_model.py

from oneflow import nn
from libai.layers import Linear

class MultiheadAttention(nn.Module):
    def __init__(hidden_size):
        self.query_key_value = Linear(
            self.hidden_size,
            self.hidden_size * 3,
            parallel="col",
        )

        self.dense = Linear(
            self.hidden_size,
            self.hidden_size,
            parallel="row",
        )
    ...

步骤3: 实现自己的ModelLoader用于加载模型权重

利用训练好的模型权重是推理的关键之一，一般来说模型能够正确加载权重需要满足以下几点：

• 训练好的模型权重文件：huggingface中有丰富的模型权重资源，可以获取到本地。
• 确保权重文件中权重的state_dict.keys()与model.state_dict().keys()一致。
• 确保权重文件中权重的state_dict.values()与model.state_dict().values()一致。

一般满足上述3点的pytorch模型即可直接加载模型权重，但是LiBai中支持分布式推理功能，可以在多机多卡的形式下正确加载模型权重并且其中的tensor都以分布式的形式存在，但这些处理细节不需要用户考虑，只需确保上述3点正确。

确保权重文件中权重的state_dict.keys()与model.state_dict().keys()一致：

由于模型的定义风格不同，LiBai中的模型参数命名与Transformers仓库有所不同，例如以下是LiBai中和Transfomers中Bert的Embedding层的命名差异，其中的vocab_embeddings与word_embeddings是等价的：

"""Bert Embedding in LiBai"""
class BertEmbeddings(nn.Module):
    def __init__(self, vocab_size, hidden_size, 
                 max_sequence_length, embedding_dropout_prob, num_tokentypes=0, 
                 init_method=nn.init.xavier_normal_, amp_e):
        super().__init__()
        self.vocab_embeddings = VocabEmbedding(
            vocab_size, hidden_size, init_method=init_method, amp_enabled=amp_enabled
        )
        self.position_embeddings = Embedding(
            max_sequence_length, hidden_size, init_method=init_method, amp_enabled=amp_enabled
        )
        ...

"""Bert Embedding in Transformers"""
class BertEmbeddings(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.word_embeddings = nn.Embedding(config.vocab_size, config.hidden_size, padding_idx=config.pad_token_id)
        self.position_embeddings = nn.Embedding(config.max_position_embeddings, config.hidden_size)
        ...

"""所以BertEmbedding.state_dict()会出现以下差异："""
# print(BertEmbedding.named_parameters().keys()) in LiBai
embeddings.vocab_embeddings.weight
embeddings.position_embeddings.weight
embeddings.tokentype_embeddings.weight

# print(BertEmbedding.named_parameters().keys()) in Transformers
embeddings.word_embeddings.weight
embeddings.position_embeddings.weight
embeddings.token_type_embeddings.weight

所以在ModelLoader中需要实现_convert_state_dict，用于将huggingface获取的权重文件中的key转为LiBai的形式，可以参考BertLoader._convert_state_dict.
另外在ModelLoader中还需要实现_load_config_from_json，用于将huggingface获取的config.py中的模型配置加载到LiBai中，可以参考BertLoader._load_config_from_json.

确保权重文件中权重的state_dict.values()与model.state_dict().values()一致

经过上述的步骤后已经实现了key值正确匹配，最后还需确认values值是一致的，在LiBai的模型实现中有两点需要注意：

• query，key，value的计算方式：LiBai中的query，key，value的计算方式是固定的，Transformers中query，key，value的计算方式在不同的模型有所区别，所以当qkv的计算方式不同时需要使用_fix_qkv_ordering方法改变LiBai中query，key，valuetensor的排列顺序来与huggingface保持一致，可以参考BertLoader中的做法：here。

"""LiBai中query，key，value的计算方式"""
query_key_value = query_key_value(hidden_states)
query_key_value = query_key_value.view(bsz, -1, self.num_heads, 3 * self.head_size)
query_key_value = query_key_value.permute(0, 2, 1, 3)  # [bsz, num_heads, src_len, 3 * head_size]
query, key, value = flow.chunk(query_key_value, chunks=3, dim=-1)

"""Transformers中BertModel的query，key，value的计算方式"""
query, key, value = flow.chunk(query_key_value, chunks=3, dim=-1)
query = query.view(query.size(0), query.size(1), self.num_heads, -1).permute(0, 2, 1, 3)
key = key.view(key.size(0), key.size(1), self.num_heads, -1).permute(0, 2, 1, 3)
value = value.view(value.size(0), value.size(1), self.num_heads, -1).permute(0, 2, 1, 3)

• layernorm的放置位置
  layernorm层的放置位置可以通过apply_residual_post_layernorm参数来改变，其原理可以参考LiBai的文档How to use Huggingface’s pretrained weights in LiBai.

目前LiBai当中支持了Bert, Roberta, GPT2, T5, MT5, Swin, Swin2, Vit模型加载huggingface权重的方法，使用方式可以参考LiBai文档，如果需要实现暂未支持的模型的权重加载也可以参考该文档实现：ModelLoaderHuggerFace.

步骤4: 编写config.py

如果使用的是LiBai训练好的model进行分布式推理, 那么直接采用训练时的train_config.py即可, 参考Couplets下的distribute_infer.py

如果是迁移的pytorch的model, 在这种情况下面, 没有完整训练的train_config.py也没有关系, 只用写一份调用model的model_config.py就可以了(记得要包含语句model=LazyCall(...)(...)), 可以参考dalle2_config.py

步骤5: 编写自己的PipelineInference.py

在LiBai中已经提供好了基类libai/inference/basic.py, 用户只需要重载自己需要的函数就可以了. 可以参考text_generation.py

如果是采用LiBai训练出来的model, 那么形式就比较简单了, 可以直接参考Couplets下的distribute_infer.py

简单说明一下在LiBai中需要重载的函数:

from libai.inference.basic import BasePipeline
from libai.utils import distributed as dist

class MyPipeline(BasePipeline):
    def _parse_parameters(self, **pipeline_parameters):
        ## 通过重载本函数, 把MyPipeline.__call__()中的入参, 分发到推理的各个阶段中.
        preprocess_params = {
            "preprocess_param1": pipeline_parameters["preprocess_param1"],
            "preprocess_param2": pipeline_parameters["preprocess_param2"],
        }
        forward_params = {
            "forward_param": pipeline_parameters["forward_param"]
        }
        postprocess_params = {
            "postprocess_param": pipeline_parameters["postprocess_param"]
        }
        return preprocess_params, forward_params, postprocess_params

    def load_pretrain_weight(self, libai_cfg_model, model_path, mode="myloader"):
        # 如果跑的是pytorch的model, 那么需要先完成步骤3, 然后把步骤3写好的`MyLoader`载入到下面的函数中
        # 如果跑的是libai训练好的model, 那么不需要重载本函数, 直接使用mode="libai"即可.
        if mode == "myloader":
            import MyLoader

            model_loader = MyLoader(
                libai_cfg_model,
                libai_cfg_model.cfg,
                model_path,
                ...,
            )
            return model_loader.load()
        else:
            return super().load_pretrain_weight(
                libai_cfg_model,
                model_path,
                mode=mode,
            )

    def preprocess(self, inputs, preprocess_param1, preprocess_param2, **kwargs):
        ...
        # model_input_dict: {"key1": flow.Tensor1, ...}
        return model_input_dict
    
    def forward(self, model_input_dict, forward_param, **kwargs):
        ...
        model_output_dict = self.model(**model_input_dict)
        return model_output_dict

    def postprocess(self, model_output_dict, postprocess_param, **kwargs):
        ...
        return out_dict

if __name__ == "__main__":
    pipeline = MyPipeline(...)
    out = pipeline(
        input_text=..., 
        preprocess_param1=..., 
        preprocess_param2=...,
        forward_param=...,
        postprocess_param=...,
    )
    if dist.is_main_process():
        print(out)

步骤5: 进行分布式推理

在步骤4中编写好了自己的PipelineInference.py以后, 就可以进行分布式的推理了. 如前言的demo演示中提到的

    # test_inference.py
    import MyPipeline
    from libai.utils import distributed as dist

    if __name__ == "__main__":
        pipeline = MyPipeline(
            "path/to/model_config.py",
            data_parallel=1,
            tensor_parallel=2,
            pipeline_parallel=2,
            pipeline_stage_id=None,
            pipeline_num_layers=...,
            model_path=...",
            mode="huggingface",
        )

但是值得注意的是, 如果是load_pytorch的model, 由于我们只在模型初始化的时候替换了LiBai的layers, 所以在这种情况下面, 只支持tensor_parallel的并行方式.

这是因为如果想要支持pipeline_parallel的话, 还需要在model.forward()中加入语句x.to_global()(参考LiBai_code.py), 把上一个pipeline_stage的中间结果, 同步到本pipeline_stage中来, 这个和hugging_face中的x.to(cuda_divce)(参考hugging_face_code.py)原理是一致的.

至于具体在哪个语句上面加上to_global()语句, 可以打开pipeline_parallel运行一遍, 定位到报错的语句添加即可.