Genesis 深度学习框架¶
轻量级深度学习框架 | 从零构建 | Python + Triton + CUDA
🚀 项目概述¶
Genesis 是一个基于 Python 从零构建的轻量级深度学习框架。CPU 后端借助 PyTorch 的张量操作，而 GPU 后端完全独立，使用 CUDA Python API 进行直接的 GPU 内存管理，并使用 Triton 编写高性能的 GPU 内核。项目旨在提供清晰的架构设计和教育价值，同时保持代码的可读性和可扩展性。
✨ 核心特性¶
- 🎯 轻量级设计 - 清晰的API设计，易于理解和使用
- ⚡ 高性能 - Triton优化的GPU内核，接近主流框架性能
- 🔄 自动微分 - 完整的反向传播和梯度计算系统
- 🧠 神经网络 - 丰富的神经网络层和优化器实现
- 🔧 混合精度 - 支持FP16/BF16混合精度训练 (AMP)
- 📊 分布式训练 - 多GPU并行训练支持
- 🎨 模型库 - 内置主流LLM模型如Qwen的实现
- 💾 模型管理 - 完整的检查点保存/加载系统
- 📈 学习率调度 - 多种学习率调度器和梯度裁剪
- 🚀 性能优化 - 内核缓存、内存池化和自适应配置
🏗️ 架构亮点¶
graph TB
A[用户API] --> B[自动微分引擎]
A --> C[神经网络模块]
B --> D[张量系统]
C --> D
D --> E[后端抽象层]
E --> F[CPU Backend]
E --> G[CUDA Backend]
G --> H[Triton Kernels]
style A fill:#e1f5fe
style B fill:#f3e5f5
style C fill:#e8f5e8
style D fill:#fff3e0
style E fill:#fce4ec
style F fill:#f1f8e9
style G fill:#e3f2fd
style H fill:#fff8e1🎯 设计目标¶
教育价值¶
- 清晰的代码结构 - 每个模块职责明确
- 完整的文档 - 从设计理念到实现细节的完整文档
- 渐进式学习 - 从基础概念到高级特性的学习路径
工程实践¶
- 现代化架构 - 借鉴PyTorch等主流框架的优秀设计
- 高效实现 - 使用Triton等现代工具进行性能优化
- 可扩展性 - 模块化设计便于添加新功能
实用性¶
- 功能完整 - 支持从模型定义到训练部署的完整工作流
- 性能优化 - 多种优化策略确保实际训练性能
- 生态兼容 - 与现有深度学习生态良好兼容
📊 性能状态¶
内存分配器性能（最新优化）¶
| 场景 | Genesis vs PyTorch | 状态 |
|---|---|---|
| 同尺寸分配 | 1.43x | ✅ 优秀 |
| 大内存(>1MB) | 3.92x | ✅ 杰出 |
| Transformer训练 | 1.89x | ✅ 优秀 |
| 内存压力 | 4.83x | ✅ 杰出 |
| 变化尺寸 | 0.83x | 🔄 良好 |
算子性能¶
| 操作 | Genesis vs PyTorch | 状态 |
|---|---|---|
| 矩阵乘法 | 0.95x | ✅ 良好 |
| 元素级操作 | 1.02x | ✅ 优秀 |
| 归约操作 | 0.87x | 🔄 优化中 |
| Softmax | 1.15x | ✅ 优秀 |
| LayerNorm | 1.08x | ✅ 优秀 |
| Cat操作 | 0.02x | ❌ 修复中 |
| LogSumExp | 0.02x | ❌ 修复中 |
| 广播操作 | 0.04x | ❌ 修复中 |
近期性能改进¶
- ✅ 块分配器: Transformer训练场景38倍性能提升
- ✅ 内存管理: 消除cudaMalloc/cudaFree同步开销
- ✅ Fill操作: GPU原生内核36倍性能提升
- 🔄 Cat操作: GPU原生实现进行中（修复0.02x问题）
- 🔄 归约操作: Triton内核优化进行中
🛠️ 技术栈¶
核心依赖¶
- Python 3.8+ - 主要开发语言
- PyTorch - 内存管理和部分操作
- Triton 2.0+ - GPU内核优化
- CUDA 11.0+ - GPU计算支持
- NumPy - CPU数值计算
- cuda-python - 直接CUDA API访问
开发工具¶
- pytest - 单元测试框架
- black - 代码格式化
- mypy - 类型检查
- MkDocs - 文档生成
- Material for MkDocs - 文档主题
🎓 学习路径¶
初学者¶
高级用户¶
贡献者¶
🌟 项目亮点¶
代码质量¶
- 类型注解 - 完整的类型提示，IDE友好
- 单元测试 - 95%+测试覆盖率
- 完整文档 - 从API到设计的全面文档
- 代码规范 - 统一的代码风格和最佳实践
近期更新 (2025-01)¶
- ✅ 内存分配器优化 - 达到PyTorch级性能
- ✅ Qwen模型支持 - 完整Qwen LLM架构实现
- ✅ 混合精度训练 - FP16/BF16自动混合精度(AMP)
- ✅ 梯度裁剪 - 支持梯度范数和值裁剪
- ✅ 学习率调度器 - StepLR, ExponentialLR, CosineAnnealingLR
- ✅ 检查点系统 - 模型保存/加载与优化器状态保存
- 🔄 算子性能 - 修复关键算子(cat, logsumexp, broadcast)
- 🔄 内核优化 - Triton内核持续改进中
🤝 社区与贡献¶
我们欢迎各种形式的贡献：
- 🐛 错误报告 - 请及时报告发现的bug
- 💡 功能建议 - 欢迎新功能想法
- 📝 文档改进 - 帮助改善文档质量
- 💻 代码贡献 - 直接参与代码开发
详情请参考 贡献指南。
📞 联系方式¶
- GitHub Issues - bug报告和功能请求
- Discussions - 技术讨论和使用交流
- 邮箱 - genesis-dev@example.com