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。
后期如有遇到此类问题，可以借鉴之~
