AstrAI/README-zh-CN.md at 2c2697390db3490db850f8a8cc7c1e382394ed45

6.6 KiB

Raw Blame History

English • 中文

轻量级 Transformer 训练与推理框架

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

📖 目录

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

中文

特性

🚀 高性能: 训练与推理双向优化，高效并行。
🔧 灵活: 支持 seq/sft/dpo/grpo 多种训练方式，可定制模型架构。
💡 易用: 简洁的 API 与丰富的示例、演示。
📦 轻量: 依赖少，部署简单。
🔬 研究友好: 模块化设计，便于实验新想法。
🤗 HuggingFace 风格 API: 类 HuggingFace 的 AutoModel/AutoTokenizer 接口，方便加载模型和分词器。
🔌 双 API 兼容: 同时支持 OpenAI 和 Anthropic 聊天补全 API，开箱即用。

快速开始

安装

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

安装开发依赖：

pip install -e ".[dev]"

下载预训练模型

下载预训练模型权重（1B 双语检查点）到 params/ 目录：

python scripts/demo/download.py

或从 HuggingFace 手动下载放入 params/。

训练模型

export CUDA_VISIBLE_DEVICES=0,1,2,3

nohup python scripts/tools/train.py \
    --nprocs=4 \
    --train_type=pt \
    --data_root_path=/path/to/dataset \
    --param_path=/path/to/model \
    --batch_per_device=4 \
    --grad_accum_steps=8 \
    --warmup_ratio=0.05 \
    --max_lr=1e-4 \
    --max_grad_norm=1.0 \
    --adamw_beta1=0.9 \
    --adamw_beta2=0.95 \
    --adamw_weight_decay=0.01 \
    --window_size=2048 \
    --ckpt_interval=10000 \
    --ckpt_dir=./checkpoint \
    --random_seed=3407 \
    --label_smoothing=0.05 \
    > out.log 2> err.log &

完整参数列表见参数说明。

文本生成

python scripts/tools/generate.py \
    --param_path /path/to/model \
    --input_json_file /path/to/input.json \
    --output_json_file /path/to/output.json

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

# Docker Compose（GPU，默认）
docker compose up -d

# Docker Compose（仅 CPU）
docker compose --profile cpu up -d

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

启动 HTTP 服务

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

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

发起请求：

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

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

# Anthropic 兼容
curl -X POST http://localhost:8000/v1/messages \
  -H "Content-Type: application/json" \
  -d '{
    "model": "astrai",
    "system": "你是一个乐于助人的助手。",
    "messages": [{"role": "user", "content": "你好"}],
    "max_tokens": 512
  }'

# Anthropic 兼容流式并设置停止序列
curl -X POST http://localhost:8000/v1/messages \
  -H "Content-Type: application/json" \
  -d '{
    "model": "astrai",
    "messages": [{"role": "user", "content": "写个故事"}],
    "max_tokens": 500,
    "stream": true,
    "stop_sequences": ["结束"]
  }'

# 健康检查
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 上的视频演示。

文档

文档	说明
参数说明	训练与推理参数配置
架构文档	系统架构、类图与设计模式
训练文档	训练循环、策略与公式
推理文档	KVCache、连续批处理、采样与 HTTP API
数据流程	数据管道、存储后端与数据集架构

贡献

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

Fork 本仓库。
创建功能分支。
提交更改。
发起 Pull Request。

重大更改请先开 issue 讨论。

社区

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

许可证

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

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

6.6 KiB Raw Blame History Unescape Escape