Qwen2.5-Coder-3B Fine-Tuned for Triton Kernel Gen
10 months ago
- #AI
- #Triton
- #PyTorch
- Qwen2.5-Coder-3B-KernelBook 是一个专精模型，用于将 PyTorch nn.Module 代码转译为 Triton 内核代码。
- 基于 GPUMODE/KernelBook 数据集微调，包含由 torch.compile 生成的 18,162 个 PyTorch-Triton 代码对。
- 采用低秩自适应（LoRA）技术，结合 PyTorch 2.5.0、Transformers、PEFT 和 TRL 框架进行微调。
- 在 NVIDIA H100 80GB 显卡上训练 1 小时 37 分钟，最终训练损失 0.0922，平均 token 准确率达 98.34%。
- 关键超参数：学习率 2e-4、批量大小 1、梯度累积步数 8、最大序列长度 4096。
- 提供示例用法，展示如何通过 Hugging Face Transformers 库从 PyTorch 代码生成 Triton 内核。
- 数据集详情：18,162 个 PyTorch-Triton 代码对，建议与 torch==2.5.0 版本配套使用。
- 基础模型 Qwen2.5-Coder-3B 含 30.9 亿参数，支持 32,768 token 上下文长度，采用 RoPE、SwiGLU 和 RMSNorm 技术。
- 同时引用了 KernelBook 数据集和 Qwen2.5-Coder 基础模型的文献来源。