Model
本目录存放模型定义。迁移目标是使用 NeMo/Megatron 配置表达模型,而不是继续维护旧的手写 PyTorch 模型定义。
当前目标 preset 来自旧工程的 2b-moe-active0.9b-balanced:
- 12 transformer layers
- hidden size 1536
- attention heads 24
- grouped-query attention kv heads 4
- dense FFN hidden size 4608
- MoE layers: 2, 4, 6, 8, 10
- experts: 12
- active experts/top-k: 4
- expert FFN hidden size 6144
- tokenizer: laoyaomodel bytes BPE, vocab size约 65K
关键改动是强制专家负载均衡:router aux loss、z-loss、capacity factor 和 router bias update 都必须显式配置,避免旧实验中专家负载失衡。