Skip to content

增加可选的MLA支持、修复模型内部精度一致,优化代码add mla, fix model dtype, improve codes#240

Open
Zephor5 wants to merge 2 commits intojingyaogong:masterfrom
Zephor5:master
Open

增加可选的MLA支持、修复模型内部精度一致,优化代码add mla, fix model dtype, improve codes#240
Zephor5 wants to merge 2 commits intojingyaogong:masterfrom
Zephor5:master

Conversation

@Zephor5
Copy link
Copy Markdown

@Zephor5 Zephor5 commented Feb 28, 2025

针对预训练pretrain稍作了一些更新和优化

  1. 移除__pychache__多余文件
  2. 添加一个可选的基于单个字符(汉字)的tokenizer
  3. 添加Deepseek的MLA(多头潜在注意力)支持(参数控制可选)
  4. pretrain dataset将数据准备移到批数据处理,增加可支持单个和多个数据文件 (另有额外未使用验证的映射训练数据读取机制,避免一次性加载全部训练数据)
  5. 更新了model内部的数据精度,保持跟参数dtype一致。
  6. 另有pretrain中注释了的显存调试的代码
  7. pretrain的学习率更新修改为累积步数时

TODO:
除了预训练、eval_model 其他的训练和api代码未做针对性适配,不影响原逻辑默认参数使用,若要启用MLA可能需要对应修改

@fangzhangmnm
Copy link
Copy Markdown

fangzhangmnm commented Mar 9, 2025

你试过训练吗。我自己写了一个mla结果loss收敛的很差
image
Disclaimer:
这个是我用自己写的代码在自己的数据集上训练的. 不是这个repo里的代码和数据
Edit:
后来发现是一个低级错误,导致这个attention层被短路了. 修复之后loss曲线和没有MLA时差不多,虽然巨慢(2x time) (因为没有MLA版的fast attention)

@Zephor5
Copy link
Copy Markdown
Author

Zephor5 commented Mar 11, 2025

@fangzhangmnm 相同参数情况下内存消耗多一点,因为中间参数多一些;loss曲线差不太多。感觉因为小模型用mla位置信息没有gqa多,所以效果在小模型下可能不太如gqa。

@LeoWootsi
Copy link
Copy Markdown

你试过训练吗。我自己写了一个mla结果loss收敛的很差 image

请问这个是在仓库里公开的那个pretrain_hq.jsonl 上训练的loss吗?

@fangzhangmnm
Copy link
Copy Markdown

你试过训练吗。我自己写了一个mla结果loss收敛的很差 image

请问这个是在仓库里公开的那个pretrain_hq.jsonl 上训练的loss吗?

这个是我用自己写的代码在自己的数据集上训练的. 不是这个repo里的代码和数据

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants