AstrAI/assets/docs/README-zh-CN.md

English • 中文

轻量级 Transformer 训练与推理框架

English • 中文 • 问题追踪 • 讨论区 • HuggingFace

📖 目录

• 特性
• 快速开始
• 文档
• 贡献
• 社区
• 许可证

---

中文

特性

• 🚀 高性能: 训练与推理双向优化，高效并行。
• 🔧 灵活: 支持 seq/sft/dpo/grpo 多种训练方式，可定制模型架构。
• 💡 易用: 简洁的 API 与丰富的示例、演示。
• 📦 轻量: 依赖少，部署简单。
• 🔬 研究友好: 模块化设计，便于实验新想法。
• 🤗 HuggingFace 集成: 兼容 HuggingFace 模型与数据集。

快速开始

安装

git clone https://github.com/ViperEkura/AstrAI.git
cd AstrAI
pip install -e .

安装开发依赖：

pip install -e ".[dev]"

训练模型

python scripts/tools/train.py \
  --train_type=seq \
  --data_root_path=/path/to/dataset \
  --param_path=/path/to/param_path

文本生成

python scripts/tools/generate.py --param_path=/path/to/param_path

Docker

使用 Docker 构建和运行（推荐用于 GPU 环境）：

# 构建镜像
docker build -t astrai:latest .

# 启用 GPU 运行
docker run --gpus all -it astrai:latest

# 指定特定 GPU
docker run --gpus '"device=0,1"' -it astrai:latest

# 运行推理服务
docker run --gpus all -p 8000:8000 astrai:latest \
  python -m scripts.tools.server --port 8000 --device cuda

# 挂载数据卷
docker run --gpus all -v /path/to/data:/data -it astrai:latest

注意: 必须使用 --gpus all 才能启用 CUDA 支持，否则 torch.cuda.is_available() 将返回 False。

启动 HTTP 服务

启动推理服务器，支持 OpenAI 兼容的 HTTP API：

python -m scripts.tools.server --port 8000 --device cuda

发起请求：

# Chat API（OpenAI 兼容）
curl -X POST http://localhost:8000/v1/chat/completions \
  -H "Content-Type: application/json" \
  -d '{
    "messages": [{"role": "user", "content": "你好"}],
    "max_tokens": 512
  }'

# 流式响应
curl -X POST http://localhost:8000/v1/chat/completions \
  -H "Content-Type: application/json" \
  -d '{
    "messages": [{"role": "user", "content": "讲个故事"}],
    "stream": true,
    "max_tokens": 500
  }'

# 健康检查
curl http://localhost:8000/health

演示

查看 scripts/demo/ 文件夹中的演示：

# 下载预处理数据（运行演示前必需）
python scripts/demo/download.py

# 交互式流式聊天
python scripts/demo/stream_chat.py

# 批量生成
python scripts/demo/generate_batch.py

# 自回归生成
python scripts/demo/generate_ar.py

观看 bilibili 上的视频演示。

文档

文档	说明
参数说明	训练与推理参数配置
设计文档	系统架构与模块设计
数据流程	数据处理管道详解
模型介绍	模型架构与技术细节

贡献

我们欢迎贡献！请参阅贡献指南了解详情。

1. Fork 本仓库。
2. 创建功能分支。
3. 提交更改。
4. 发起 Pull Request。

重大更改请先开 issue 讨论。

社区

• GitHub Issues: 问题追踪
• Discussions: GitHub 讨论区
• HuggingFace: 模型中心

许可证

本项目采用 GPL-3.0 许可证。

---

专为高性能与易用性设计的轻量级 Transformer 框架。