LLM优化：开源星火13B显卡及内存占用优化

1. 背景

本qiang~这两天接了一个任务，部署几个开源的模型，并且将本地经过全量微调的模型与开源模型做一个效果对比。

部署的开源模型包括：星火13B，Baichuan2-13B, ChatGLM6B等

其他两个模型基于transformers架构封装，因此推理服务启动还是十分丝滑，但星火13B是基于Megatron-DeepSpeed框架实现，地址是：https://gitee.com/iflytekopensource/iFlytekSpark-13B，启动推理服务的过程中发现启动13B的显卡占用71G-78G，有些反直觉。

此文就是整理开源星火13B的显存及内存排查并优化的整理过程，至于哪家开源模型效果好，不在此文的讨论范围内。

2. 原因分析

直观上来说，13B的模型，数据类型为bf16，显卡占用大概在26G左右，但星火13B直接占用70G+，不可思议，怪不得网上关于星火开源模型的讨论少之又少，原因显而易见，这么大的显存占用只能用多卡或者A800等80G显卡才能适配。穷人家的孩子，哪有这么多余粮。

排查原因的过程中，少不了源码的调试与分析。在排查的过程中，启动推理服务的文件run_iFlytekSpark_text_generation.py中，model_provider方法是初始化模型并加载模型文件的方法。

def model_provider(pre_process=True, post_process=True):
    """Build the model."""
    print_rank_0('building iFlytekSpark model ...')
    args = get_args()
    config = core_transformer_config_from_args(args)

    ### 初始化星火模型
    model = iFlytekSparkModel(
        config,
        num_tokentypes=0,
        parallel_output=False,
        pre_process=pre_process,
        post_process=post_process,
        return_moe_loss=False
    )

    if args.from_pretrained is not None:
        assert os.path.exists(args.from_pretrained)
        ckpt_path = get_checkpoint_name(args.from_pretrained)
        print_rank_0('Loading from {} '.format(
                args.from_pretrained))
        # 模型加载权重文件
        state_dict = torch.load(ckpt_path, map_location=f"cuda:{torch.cuda.current_device()}")
        if 'module' in state_dict:
            state_dict = state_dict['module']
        model.load_state_dict(state_dict)
    return model

其中，加载权重文件可以看到，加载state_dict时，直接将权重文件加载到显卡中，而非加载至CPU，然后再执行to方法，转移到GPU。因此该处是一个潜在的优化点。

再打入iFlytekSparkModel内部，词表Embedding层，线性转换层，等初始化weight时，也是直接将weight分配在GPU上运行。例如下例：

class RowParallelLinear(torch.nn.Module):
    def __init__(self, input_size: int, output_size: int, *,
                 config: ModelParallelConfig,
                 init_method: Callable,
                 bias: bool = True,
                 input_is_parallel: bool = False,
                 stride: int = 1,
                 keep_master_weight_for_test: bool = False,
                 skip_bias_add: bool = False,
                 moe=False, enable_expert_tensor_parallelism=False):
        super(RowParallelLinear, self).__init__()

        # .........

        if config.use_cpu_initialization:
            self.weight = Parameter(torch.empty(self.output_size,
                                                self.input_size_per_partition,
                                                dtype=config.params_dtype))
            if config.perform_initialization:
                self.master_weight = _initialize_affine_weight_cpu(
                    self.weight, self.output_size, self.input_size,
                    self.input_size_per_partition, 1, init_method,
                    stride=stride, return_master_weight=keep_master_weight_for_test,
                    params_dtype=config.params_dtype)
        else:
            # 默认按照启动sh命令，会走该分支
            self.weight = Parameter(torch.empty(
                self.output_size, self.input_size_per_partition,
                device=get_accelerator().current_device_name(), dtype=config.params_dtype))
            if config.perform_initialization:
                _initialize_affine_weight_gpu(self.weight, init_method,
                                              partition_dim=1, stride=stride)
        if bias:
            if config.use_cpu_initialization:
                self.bias = Parameter(torch.empty(self.output_size,
                                                  dtype=config.params_dtype))
            else:
                # 默认按照启动sh命令，会走该分支
                self.bias = Parameter(torch.empty(
                    self.output_size, device=get_accelerator().current_device_name(),
                    dtype=config.params_dtype))
            setattr(self.bias, 'sequence_parallel', self.sequence_parallel)

            if config.perform_initialization:
                # Always initialize bias to zero.
                with torch.no_grad():
                    self.bias.zero_()
        else:
            self.register_parameter('bias', None) 

3. 优化方案

1. 模型初始化时，模型的Embedding，线性层的权重weight均直接加载至GPU，因此可以优化为先将这些weight加载至CPU。

改进的方式也很简单，从上面的源码层面，可以看到，当增加参数” use_cpu_initialization”，将使用CPU进行初始化权重，因此只需要在启动推理服务的脚本中增加” --use-cpu-initialization”参数即可。

2. 加载模型文件时，直接加载至GPU，然后run_iFlytekSpark_text_generation.py中的get_model方法中，当模型加载完成后，会进行分配至GPU以及FP16的转换的操作。如下代码所示。

def get_model(model_provider_func, model_type=ModelType.encoder_or_decoder, wrap_with_ddp=True):
    """Build the model."""
    args = get_args()
    args.model_type = model_type

    # ..........

    # GPU allocation.
    for model_module in model:
        model_module.to(get_accelerator().current_device_name())

    # Fp16 conversion.
    if args.fp16 or args.bf16:
        model = [Float16Module(model_module, args) for model_module in model]

    # .......

    return model

因此，优化的方式也很简单，可以优化为先加载至CPU，再运行get_model中的默认分配至GPU，加载完后，再使用垃圾回收机制清除CPU占用的内存即可。

话不多说，优化后的代码如下：

def model_provider(pre_process=True, post_process=True):
    """Build the model."""
    print_rank_0('building iFlytekSpark model ...')
    args = get_args()
    config = core_transformer_config_from_args(args)
    model = iFlytekSparkModel(
        config,
        num_tokentypes=0,
        parallel_output=False,
        pre_process=pre_process,
        post_process=post_process,
        return_moe_loss=False
    )

    if args.from_pretrained is not None:
        print(args.from_pretrained)
        assert os.path.exists(args.from_pretrained)
        ckpt_path = get_checkpoint_name(args.from_pretrained)
        print_rank_0('Loading from {} '.format(
                args.from_pretrained))

        # state_dict = torch.load(ckpt_path, map_location=f"cuda:{torch.cuda.current_device()}")
        # CPU进行加载
        state_dict = torch.load(ckpt_path, map_location=f"cpu")
        if 'module' in state_dict:
            state_dict = state_dict['module']
        model.load_state_dict(state_dict)

        # 加载完成，删除state_dict，并垃圾回收
        del state_dict
        gc.collect()
        torch.cuda.empty_cache()

    return model

4. 效果对比

(1) 优化前的显卡占用: 71.5G

(2) 优化前的内存占用: 虚拟内存占用94.5G

(3) 优化后的显卡占用: 26G

(4) 优化后的内存占用: 43.1G

5. 总结

一句话足矣~

本文主要是针对开源星火13B的显存及内存占用过大的一个代码优化。核心思想是使用CPU预加载模型，再转换至GPU。

后期如有遇到此类问题，可以借鉴之~