大模型开发基本功:模型转换、推理优化与训练监控实战
除了 Transformer、RoPE、SwiGLU、RMSNorm,FlashAttention 原理,Megatron 并行策略,量化和推理加速技术,CUDA 编程等核心技术外,大模型的基本功还有哪些?
虽然这些追求性能极限优化的核心技术很重要,但在实际工作中,大多数场景并不需要完全理解其底层实现。很多时候我们的需求是使用它们,而不是去优化它们。因此,本文将介绍一些在实际工程开发中非常常用,但往往被忽视的'基本功'。这些技能对于提升开发效率、解决实际问题至关重要。
1. 模型结构转换能力
在开源社区,Llama 的网络结构已经占据了主导地位,modeling_llama.py 理论上可以加载任何一个兼容该结构的开源模型。掌握模型权重转换的能力是理解不同架构差异的关键。
基础任务:跨架构权重映射
尝试完成以下脚本,使得我们可以用 modeling_llama.py 加载其他开源模型的权重:
trans_qwen_to_llama.pytrans_llama_to_qwen.py
实现思路: 通过遍历源模型的 state_dict,将键名(key)和目标模型的键名进行映射。在此过程中,你会观察到每个开源模型的独特之处。例如:
- Qwen2 的 q、k、v 线性变换可能包含 bias。
- Baichuan 的 lm_head 之前有一个 normalize() 操作。
- 某些模型存在 attention 机制的魔改。
对着论文去理解作者为什么要做这些改动,是深入理解模型设计的重要环节。
进阶任务:Megatron 并行策略转换
并非所有同学都用 Megatron 训练代码,但如果使用,这两个脚本是基本功中的基本功:
trans_llama_to_megatron.py(给定参数 tp 和 pp)trans_megatron_to_llama.py
注意事项:
Megatron checkpoint 的 pp_size 实现 merge 和 split 相对简单,但在对 tp_size 进行 merge 和 split 的时候,一定要留意 Megatron 的 GQA(Grouped Query Attention)实现方式,否则会导致权重对齐错误。
2. 自定义 Modeling 文件
既然我们已经可以万物转 Llama,为什么还要实现一个属于自己的 modeling_XX.py?
modeling_llama.py 在通用性上存在局限,例如:
- 加载模型时没有
skip_build选项; - 缺少
stream_generate支持; - 不支持
sequence_parallel; - 默认使用 FlashAttention;
- 并没有一个可以作为 reward_model 的 lm_head。
解决方案:
实现一个集百家之长的 modeling_XX.py。收集 modeling_llama.py、modeling_qwen.py、modeling_baichuan.py、modeling_yi.py、modeling_deepseek.py、modeling_glm.py 等所有开源文件,把各家公司实现得比较好用的 def 加入到自己的文件中。
这样,当市面上出现任何新的开源模型,我们都可以通过 trans_newModel_to_myModel.py 快速对该模型进行微调操作,而无需修改训练代码。
调试增强功能
为了助力日常 debug,可以在 modeling 文件中加入以下有趣的功能:


