# triton-test

**Repository Path**: AD_Andy/triton-test

## Basic Information

- **Project Name**: triton-test
- **Description**: No description available
- **Primary Language**: Unknown
- **License**: Not specified
- **Default Branch**: main
- **Homepage**: None
- **GVP Project**: No

## Statistics

- **Stars**: 0
- **Forks**: 0
- **Created**: 2025-10-16
- **Last Updated**: 2025-10-30

## Categories & Tags

**Categories**: Uncategorized

**Tags**: None

## README

# Triton性能调研项目

> 在ROCm/DCU环境下对比PyTorch原生实现与Triton优化kernel的性能差异

[![Python](https://img.shields.io/badge/Python-3.10-blue.svg)](https://www.python.org/)
[![PyTorch](https://img.shields.io/badge/PyTorch-2.4.1-orange.svg)](https://pytorch.org/)
[![Triton](https://img.shields.io/badge/Triton-3.0.0-green.svg)](https://triton-lang.org/)

## 📋 目录

- [快速开始](#快速开始)
- [项目概述](#项目概述)
- [项目结构](#项目结构)
- [环境配置](#环境配置)
- [使用指南](#使用指南)
- [技术实现](#技术实现)
- [测试计划](#测试计划)
- [项目进度](#项目进度)
- [文档说明](#文档说明)

## 🚀 快速开始

### 本地开发

```bash
# 1. 克隆代码
git clone <your-gitee-repo> triton-test
cd triton-test

# 2. 安装依赖
pip install -r requirements.txt

# 3. 查看可用命令
make help
```

### 服务器部署

```bash
# 1. 在登录节点下载模型（有网络）
make download

# 2. 在计算节点运行训练
make train-pytorch    # PyTorch基线
make train-triton     # Triton优化

# 多卡训练
NUM_GPUS=4 make train-pytorch-multi
NUM_GPUS=4 make train-triton-multi

# 3. 性能对比
make compare

# 4. 一键完整测试
bash scripts/run_full_test.sh
```

## 📖 项目概述

### 目标

对比PyTorch原生实现与Triton优化kernel在**训练和推理**中的性能差异：

- **训练性能**: 吞吐量、延迟、内存使用
- **推理性能**: 延迟、吞吐量（次要）
- **多卡扩展性**: 1/2/4/8卡DDP性能

### 测试环境

| 类别               | 配置                    |
| ------------------ | ----------------------- |
| **硬件**     | 8×海光DCU卡 (800W功耗) |
| **操作系统** | Ubuntu 22.04            |
| **ROCm**     | DTK 25.04.1             |
| **PyTorch**  | 2.4.1 (ROCm版本)        |
| **Triton**   | 3.0.0                   |
| **Python**   | 3.10                    |

### 测试模型

**阶段一：GPT-2 Medium**（当前）

- 参数量: 355M (1.32GB)
- 数据集: WikiText-2 (36,718样本)
- 目标: 快速验证Triton优化效果

**阶段二：Qwen2.5-7B**（后续）

- 参数量: 7B (约14GB)
- 目标: 在更大模型上验证

**阶段三：Qwen3-235B**（可选）

- 参数量: 235B (MoE架构)
- 目标: 超大模型推理测试

## 📁 项目结构

```
triton-test/
├── README.md                      # 项目主文档
├── requirements.txt               # Python依赖
├── Makefile                       # 快捷命令
│
├── docs/                          # 📚 文档目录
│   ├── USAGE.md                  # 详细使用指南
│   ├── DEVELOPMENT.md            # 开发者文档
│   ├── TESTING_GUIDE.md          # 测试流程
│   ├── OFFLINE_INSTALL.md        # 离线安装
│   └── TROUBLESHOOTING.md        # 故障排查
│
├── configs/                       # ⚙️ 配置文件
│   ├── pytorch.json              # PyTorch训练配置
│   └── triton.json               # Triton训练配置
│
├── scripts/                       # 🔧 脚本目录
│   ├── run_full_test.sh         # 完整测试流程
│   ├── train/                    # 训练脚本
│   │   ├── train_pytorch.py     # PyTorch基线训练 (385行)
│   │   ├── train_triton.py      # Triton优化训练 (290行)
│   │   ├── run_pytorch.sh       # PyTorch启动脚本
│   │   └── run_triton.sh        # Triton启动脚本
│   └── tools/                    # 工具脚本
│       ├── download_models.py   # 模型下载
│       ├── compare_results.py   # 性能对比
│       └── benchmark.py         # 基准测试
│
├── triton_kernels/               # 🚀 Triton Kernel实现
│   ├── attention.py              # Flash Attention (270行)
│   ├── layer_norm.py             # LayerNorm融合 (200行)
│   ├── mlp.py                    # MLP+GELU优化 (172行)
│   └── utils.py                  # 工具函数 (69行)
│
├── models/                       # 模型文件（需下载）
├── data/                         # 数据集（需下载）
├── results/                      # 测试结果（自动生成）
└── logs/                         # 训练日志（自动生成）
```

### 核心组件说明

| 组件 | 文件 | 功能亮点 |
|------|------|---------|
| **训练脚本** | `scripts/train/train_pytorch.py` | 完整性能监控、DDP支持 |
|  | `scripts/train/train_triton.py` | 自动模块替换、fallback机制 |
| **Triton Kernels** | `triton_kernels/attention.py` | Flash Attention风格、减少HBM访问 |
|  | `triton_kernels/layer_norm.py` | 算子融合、减少kernel启动 |
|  | `triton_kernels/mlp.py` | 快速GELU、激活函数融合 |
| **工具脚本** | `scripts/tools/compare_results.py` | 自动生成图表和报告 |
|  | `scripts/tools/benchmark.py` | 多配置基准测试 |

## ⚙️ 环境配置

### 硬件环境

```
┌─────────────────┐        ┌──────────────────┐
│  登录节点        │  共享   │  计算节点        │
│  (有网络)        │◄──────►│  (node109)       │
│  - 模型下载      │  存储   │  - 8×DCU卡       │
│  - 数据准备      │        │  - 训练执行      │
└─────────────────┘        └──────────────────┘
```

### 软件依赖

核心依赖（见 `requirements.txt`）：

```
torch==2.4.1
triton==3.0.0
transformers==4.39.2
datasets==3.6.0
```

### 容器配置

训练容器需要以下配置：

```bash
docker run -itd \
  --name ad-triton-test \
  --device=/dev/kfd --device=/dev/mkfd --device=/dev/dri \
  --group-add video --group-add 39 \
  --shm-size=32g \
  -v /workspace/triton-test:/workspace/triton-test \
  -v /opt/hyhal:/opt/hyhal:ro \
  -v /opt/dtk-25.04-rc4:/opt/dtk-25.04-rc4:ro \
  -e HIP_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 \
  <base-image>
```

## 📚 使用指南

### Makefile命令

```bash
make install          # 安装依赖
make download         # 下载模型和数据
make train-pytorch    # PyTorch单卡训练
make train-triton     # Triton单卡训练
make compare          # 生成性能对比
make benchmark        # 运行基准测试
make clean            # 清理临时文件
```

### 配置调整

编辑 `configs/pytorch.json` 或 `configs/triton.json`：

```json
{
  "data": {
    "block_size": 512              // 序列长度
  },
  "training": {
    "per_device_train_batch_size": 8,  // batch size
    "learning_rate": 5e-5,         // 学习率
    "bf16": true,                  // 混合精度
    "num_train_epochs": 3          // 训练轮数
  },
  "triton": {
    "enable_attention": true,      // 启用Triton Attention
    "enable_layer_norm": true,     // 启用Triton LayerNorm
    "enable_mlp": true             // 启用Triton MLP
  }
}
```

### 多配置测试

```bash
make benchmark

# 或自定义配置
cd scripts/tools
python benchmark.py \
  --batch-sizes 1 2 4 8 \
  --seq-lengths 128 256 512 1024 \
  --num-gpus 1 2 4 8
```

### 查看结果

```bash
# 性能指标
cat results/pytorch_baseline/metrics.json
cat results/triton_optimized/metrics.json

# 对比报告
cat results/comparison/comparison_report.md

# 可视化图表
ls results/comparison/*.png
```

## 🔧 技术实现

### 1. 性能监控系统

实现了完整的训练性能监控：

```python
class PerformanceMonitor:
    metrics = {
        "loss": [],                 # 训练损失
        "throughput_samples": [],   # 样本吞吐量 (samples/sec)
        "throughput_tokens": [],    # token吞吐量 (tokens/sec)
        "memory_allocated": [],     # GPU内存使用 (GB)
        "step_time": [],           # 单步时间 (sec)
    }
```

**功能**：

- ✅ 实时记录训练指标
- ✅ 定期打印性能统计
- ✅ 保存JSON格式详细数据
- ✅ 计算平均值和趋势

### 2. Triton Kernel优化

#### Flash Attention

```python
# 关键优化技术
- Tiling策略: 减少HBM访问
- 在线Softmax: 融合计算，避免多次遍历
- Causal Masking: 支持自回归生成
```

#### LayerNorm融合

```python
# 优化策略
- 算子融合: 归一化+仿射变换一次完成
- 减少内存访问: 单kernel调用
- 数值稳定性: 与PyTorch一致
```

#### MLP优化

```python
# 快速GELU近似
GELU(x) ≈ 0.5 * x * (1 + tanh(√(2/π) * (x + 0.044715 * x³)))
```

### 3. 自动模块替换

```python
def replace_modules_with_triton(model, config):
    """自动替换PyTorch模块为Triton版本"""
    for name, module in model.named_modules():
        if "attn" in name:
            # 替换Attention
            triton_attn = TritonGPT2Attention(module.config)
            triton_attn.load_state_dict(module.state_dict())
            # ... 执行替换
```

**特点**：

- 🔍 自动识别可替换模块
- 💾 保留原始权重
- 🔄 支持渐进式替换
- 🛡️ 自动fallback到PyTorch

### 4. 分布式训练

- **DDP集成**: 多卡数据并行
- **自动设备分配**: 基于环境变量
- **混合精度**: BF16减少通信开销
- **梯度同步**: 高效的all-reduce

## 📊 测试计划

### 训练性能测试（优先级高）

| 测试项   | PyTorch                 | Triton | 对比指标 |
| -------- | ----------------------- | ------ | -------- |
| 吞吐量   | samples/sec, tokens/sec | ✓     | 加速比   |
| 内存使用 | 峰值内存 (GB)           | ✓     | 优化率   |
| 延迟     | 单步时间 (ms)           | ✓     | 时间减少 |
| 多卡扩展 | 1/2/4/8卡 DDP           | ✓     | 扩展效率 |

**配置矩阵**：

  - Batch size: 1, 2, 4, 8, 16
  - 序列长度: 128, 256, 512, 1024
  - 学习率: 1e-5, 5e-5, 1e-4

### 推理性能测试（次要）

- 首token延迟
- 后续token延迟
- 批处理吞吐量
- 内存占用对比

## 📈 项目进度

### ✅ 已完成（本地）

**代码实现** (约2000行Python代码)

- ✅ PyTorch基线训练 (385行)

  - 完整训练流程
  - PerformanceMonitor性能监控
  - DDP多卡支持
  - FP16/BF16混合精度
- ✅ Triton优化训练 (290行)

  - 自动模块替换
  - PyTorch fallback
  - 配置化kernel启用
- ✅ Triton Kernels (711行)

  - attention.py (270行): Flash Attention
  - layer_norm.py (200行): 融合LayerNorm
  - mlp.py (172行): 优化MLP+GELU
  - utils.py (69行): 工具函数
- ✅ 工具脚本

  - download_models.py: 支持镜像下载
  - compare_results.py: 性能对比+可视化
  - benchmark.py: 多配置测试
- ✅ 文档系统 (约1000行)

  - README.md: 项目总览
  - USAGE.md: 使用指南 (320行)
  - DEVELOPMENT.md: 开发文档 (390行)

### ⏳ 待执行（服务器）

**环境准备**

- ⏳ 代码推送到Gitee
- ⏳ 服务器克隆代码
- ⏳ 安装依赖和配置

**数据准备**

- ⏳ 下载GPT-2 Medium模型
- ⏳ 下载WikiText-2数据集
- ⏳ 验证数据可用性

**性能测试**（1-2天）

- ⏳ PyTorch基线训练
- ⏳ Triton优化训练
- ⏳ 生成性能对比
- ⏳ 多卡扩展性测试

**结果分析**（3-5天）

- ⏳ 分析性能提升
- ⏳ 识别优化瓶颈
- ⏳ 生成详细报告
- ⏳ ROCm/DCU优化建议

**扩展验证**（1-2周）

- ⏳ Qwen2.5-7B测试
- ⏳ 多配置benchmark
- ⏳ 推理性能测试（可选）

## 📄 文档说明

| 文档 | 用途 |
|------|------|
| **README.md** | 项目概览和快速开始（本文档） |
| **[docs/USAGE.md](docs/USAGE.md)** | 详细使用指南和配置说明 |
| **[docs/DEVELOPMENT.md](docs/DEVELOPMENT.md)** | 开发文档和kernel开发教程 |
| **[docs/TESTING_GUIDE.md](docs/TESTING_GUIDE.md)** | 完整测试流程指南 |
| **[docs/OFFLINE_INSTALL.md](docs/OFFLINE_INSTALL.md)** | 离线环境安装方案 |
| **[docs/TROUBLESHOOTING.md](docs/TROUBLESHOOTING.md)** | 常见问题和故障排查 |

**快速导航**：
- 🚀 运行训练 → [快速开始](#快速开始)
- ⚙️ 调整配置 → [docs/USAGE.md](docs/USAGE.md)
- 🧪 完整测试 → [docs/TESTING_GUIDE.md](docs/TESTING_GUIDE.md)
- 🔧 开发扩展 → [docs/DEVELOPMENT.md](docs/DEVELOPMENT.md)
- 🐛 解决问题 → [docs/TROUBLESHOOTING.md](docs/TROUBLESHOOTING.md)

## 🎯 预期成果

### 代码交付物 ✅

- **完整代码库**: 约2000行高质量代码
- **文档系统**: 3个文档，约1000行
- **配置管理**: JSON配置+Shell脚本
- **自动化工具**: 下载、训练、对比、测试

### 性能数据 ⏳

待测试后生成：

- PyTorch vs Triton性能对比
- 多卡扩展性数据
- 不同配置下的性能表现
- 可视化图表和报告

### 优化建议 ⏳

待分析后提供：

- ROCm/DCU环境特定优化
- Kernel参数调优建议
- 内存管理策略
- 多卡通信优化

## 📊 项目统计

```
代码统计:
├── Python代码: ~2000行
├── 配置文件: 2个 (configs/)
├── Shell脚本: 3个
├── 文档: 6个 (README + docs/)
└── 总文件数: ~30个

目录结构:
├── docs/: 5个专业文档
├── configs/: 2个配置文件
├── scripts/train/: 训练脚本
├── scripts/tools/: 工具脚本
└── triton_kernels/: 3个Triton kernel实现
```

## 🤝 贡献和反馈

### 获取帮助

1. **使用问题**: 查看 [docs/USAGE.md](docs/USAGE.md)
2. **测试流程**: 查看 [docs/TESTING_GUIDE.md](docs/TESTING_GUIDE.md)
3. **开发问题**: 查看 [docs/DEVELOPMENT.md](docs/DEVELOPMENT.md)
4. **故障排查**: 查看 [docs/TROUBLESHOOTING.md](docs/TROUBLESHOOTING.md)

### 项目特点

✨ **完整性**: 从训练到对比的全流程
🔧 **可配置**: 灵活的JSON配置系统
📊 **可视化**: 自动生成性能图表
📖 **文档全**: 使用和开发文档齐全
🧪 **可测试**: 完整的基准测试套件
🚀 **易部署**: 一键运行完整测试

---

**项目作者**: Triton Performance Testing Team
**更新时间**: 2024年10月
**许可协议**: MIT License