1. 模型定义与注册
1.1 在 Transformers 库中新建一个名为 SailorCausalLM 的模型,并完成注册,整体架构继承自 Qwen2CausalLM。
2. 模块替换
2.1 将 SailorCausalLM 的 Attention 替换为 fla-org/flame 仓库中 RetNet 的实现。
2.2 将 MLP 替换为 Qwen2MOE 模块中的实现。
2.3 将 Attention 中的 RMSNorm 替换为 LayerNorm(MLP 中保持 RMSNorm)。
2.4 将残差连接机制改为 DenseNet 风格的残差结构。
3. 推理框架搭建
3.1 为本地路径 ~/models/qwen2-0.5B 编写支持 batch 推理与 streaming 输出的 generator 函数。
4. 推理优化技术集成 4.1 添加量化推理支持。
4.2 实现张量并行(Tensor Parallel, TP)推理。
4.3 对 Attention 使用数据并行,对 MLP 使用张量并行。
4.4 使用 FlashAttention 替换 Huggingface 原生 KV Cache 以提升推理速度。
4.5 接入 FlashInfer 加速推理。
4.6 实现带有 window attention 的 Streaming LLM(支持轮转窗口)。
4.7 集成 SnapKV KV 缓存机制。
4.8 使用 Flash-MLA 提升推理效率。
4.9 实现专家并行(Expert Parallel)支持。
5. 基础训练框架开发
5.1 不依赖 Huggingface Transformers,编写独立的 train.py,支持数据加载、自动验证等基础功能。
6. 训练功能增强与并行策略 6.1 集成 DeepSpeed 以提升训练性能。
6.2 使用 TorchTitan 进行模型训练。
6.3 训练中使用 FlashAttention(注意处理 padding 问题)。
6.4 支持数据并行(Data Parallel, DP)训练。
6.5 支持张量并行(TP)训练。
6.6 支持流水线并行(Pipeline Parallel, PP)训练。
6.7 支持使用 Ulysses 和 Flash-Ring-Attention 实现序列并行与三维混合并行(Sequence / 3D Parallel)。
6.8 使用 Muon 优化器。
6.9 使用 AdamW 优化器。
6.10 配置学习率调度器(LR Scheduler)。
6.11 支持多机训练。
6.12 在 WMT 数据集上训练 Transformer 编码器-解码器,并评估 BLEU 分数。
6.13 替换 Transformer 中的关键模块为 Liger Kernel 的实现(不采用官方给得 Monkey Patching 脚本)。
7. 模型训练任务 7.1 模块化训练 CNN 于 FashionMNIST 数据集上。
7.2 训练一个标准 ResNet 模型。
7.3 训练一个 DenseNet 模型。
7.4 上述模型均集成 WandB 监控与 Early Stopping 策略。
8. 核心功能实现 8.1 使用 PyTorch 实现 matmul 的前向与反向传播。 8.2 使用 NumPy 实现一个自动微分框架(挑战性较高)。