OpenASH

OpenASH 是一个基于 PyTorch 的轻量级语言模型项目，使用自研的 累积最大值注意力机制（Cumulative Max Attention） 替代传统 Softmax 注意力，支持预训练、SFT 微调、流式推理和 WebUI 对话。

模型架构

• MaxStateSuper: 核心注意力层，使用 torch.cummax 操作替代 softmax，通过可学习的 alpha 参数混合多种交互项，支持有状态推理（stateful inference）
• DecoderLayer: 解码器层，包含 MaxStateSuper 注意力 + 门控前馈网络（Gated FFN）+ 残差连接 + LayerNorm
• OpenASH: 完整模型，由 Embedding + 多层 DecoderLayer + 输出投影层组成

默认配置: hidden_size=640, num_layers=8, num_heads=8（约数百万参数）。

项目结构

demo/
├── open_ash.py            # 模型定义（MaxStateSuper, DecoderLayer, OpenASH）
├── open_ash_voc.py         # 分词器（基于 jieba + 代理词表编码）
├── open_ash_infer.py       # 推理引擎（采样、流式生成、交互式聊天）
├── open_ash_webui.py       # Streamlit WebUI（支持工具调用、思考模式）
├── open_ash_dataset.py     # 数据集（Pretrain/SFT/DPO/RLAIF/AgentRL）
├── train_pretrain.py       # 预训练脚本
├── train_full_sft.py       # SFT 全参数微调脚本
├── trainer_utils.py        # 训练工具（分布式、检查点、学习率调度）
├── config.py               # 配置文件
├── open_ash_voc_agent.json # 代理词表映射
└── vocabulary_nnn.json     # 基础词表

特性

• 自定义注意力机制: 使用累积最大值（cummax）替代 softmax，计算效率更高
• 有状态推理: 支持跨 chunk 的状态传递，可处理超长序列
• 代理词表编码: 通过二维/三维索引编码扩展词表，在有限 token 空间内覆盖更多词汇
• 多种采样策略: 支持 Temperature、Top-K、Top-P、重复惩罚
• 流式输出: 支持逐 token 流式解码
• 工具调用: 支持 function calling / tool use
• 思考模式: 支持 <|think|>...<|end_think|> 链式思考
• 分布式训练: 支持多卡 DDP、混合精度（BFloat16）、梯度累积、断点续训
• Streamlit WebUI: 提供可视化对话界面，支持中英文切换

依赖

torch
jieba
streamlit
numpy
tqdm

快速开始

推理

python open_ash_infer.py

WebUI 对话

streamlit run demo/open_ash_webui.py

启动后会自动在浏览器中打开 WebUI 界面（默认 http://localhost:8501），支持以下功能：

• 模型选择: 支持加载不同配置的模型权重（自动扫描 out/ 目录）
• 参数调节: 侧边栏可调节历史对话轮次、最大生成长度、温度、Top-K、Top-P、重复惩罚
• 思考模式: 开启后模型会在 <|think|>...<|end_think|> 中展示推理过程
• 工具调用: 支持选择并使用工具（数学计算、时间查询、随机数、文本长度、单位转换、天气、汇率、翻译等）
• 多语言: 支持中英文界面切换
• 流式输出: 实时显示生成过程
• 对话历史: 保存当前会话的对话记录

配置说明：

• WEIGHT_DIR: 模型权重目录（默认 F:\OpenASH\out）
• HIDDEN_SIZE: 模型隐藏层维度（默认 640）
• NUM_LAYERS: 模型层数（默认 8）
• NUM_HEADS: 注意力头数（默认 8）
• MAX_SEQ_LEN: 最大序列长度（默认 2048）

预训练

python train_pretrain.py \
    --data_path pretrain_t2t.jsonl \
    --epochs 6 \
    --batch_size 40 \
    --learning_rate 1.25e-4 \
    --max_seq_len 512 \
    --save_dir ../out

SFT 微调

python train_full_sft.py \
    --data_path sft_t2t.jsonl \
    --epochs 6 \
    --batch_size 40 \
    --learning_rate 1.25e-5 \
    --from_weight pretrain \
    --max_seq_len 512 \
    --save_dir ../out \
    --save_weight full_sft

多卡训练

torchrun --nproc_per_node=4 train_pretrain.py \
    --data_path pretrain_t2t.jsonl \
    --batch_size 10

训练参数说明

数据格式

预训练数据

每行一个 JSON 对象，包含 text 字段：

{"text": "这是一段训练文本..."}

SFT 数据

每行一个 JSON 对象，包含 conversations 字段：

{"conversations": [{"role": "user", "content": "你好"}, {"role": "assistant", "content": "你好！有什么可以帮你的？"}]}

支持 system、user、assistant、tool 角色，以及 reasoning_content（思考内容）、tool_calls（工具调用）、tools（工具定义）等字段。

License

MIT