# tx-llm-base

**Repository Path**: homecommunity/tx-llm-base

## Basic Information

- **Project Name**: tx-llm-base
- **Description**: 甜心大模型训练基础版，方便学习和理解，不用pytorch等框架去实现，而是自己实现方便熟悉大模型的预训练、后训练过程
- **Primary Language**: Unknown
- **License**: Not specified
- **Default Branch**: main
- **Homepage**: None
- **GVP Project**: No

## Statistics

- **Stars**: 1
- **Forks**: 0
- **Created**: 2026-09-23
- **Last Updated**: 2026-10-01

## Categories & Tags

**Categories**: Uncategorized

**Tags**: None

## README

# tx-llm-base

零深度学习框架依赖的 Llama-2 训练代码库：自动求导、优化器、训练循环全部手写，数组计算统一走 **CuPy**
（数据常驻显存，底层 CUDA / cuBLAS），不引入 PyTorch / TensorFlow / JAX，**不自研任何 CUDA 内核**。

## 模块划分

```text
tx_bpe/       分词器（CPU 侧纯 Python）：BPE 训练 / 编解码 / Llama-2 词表导入
tx_model/     唯一基础设施层：数组入口 / 设备 / 算子 / Llama-2 模型 / 优化器 / 序列化
tx_pretrain/  预训练：语料 token 化、批构造、训练循环、checkpoint
tx_sft/       有监督微调：指令数据、Qwen-2.5 ChatML 模板、label mask
tx_grpo/      GRPO：rollout 采样、奖励、优势、KL、策略损失
tx_opd/       OPD（On-Policy Distillation）本期仅占位，不实现功能
```

`tx_bpe` 属 CPU 侧逻辑（字符串操作为主），仅用 Python 标准库 + NumPy；`tx_model` 及三个训练模块的张量
计算 100% 在 GPU 上以 CuPy 完成，层与层之间只传 `dx`，无全局计算图。

## 环境要求

| 环境 | 能否运行 CuPy | 用途 |
| --- | --- | --- |
| 本机 macOS | 否（无 NVIDIA GPU / CUDA） | 写代码、静态检查、纯 CPU 逻辑单测（分词器 / 数据格式 / Chat 模板） |
| 远程 Linux + NVIDIA GPU | 是（`cupy-cuda12x` wheel） | 显存数值验证、gradcheck、真实单卡训练 |

CuPy wheel 必须与 CUDA 驱动匹配：CUDA 12.x 驱动装 `cupy-cuda12x`，CUDA 11.x 装 `cupy-cuda11x`。
**CuPy 14 起 CUDA 运行库（cudart / curand / cublas / nvrtc）与头文件不再随 wheel 内置**，需装 `[ctk]` extra
（或使用系统 CUDA 12 toolkit 并设 `CUDA_PATH`），否则随机数（cuRAND）与逐元素算子 JIT 会失败。
安装：

```bash
pip install -r requirements.txt          # 含 cupy-cuda12x[ctk]
# 可选（仅在需要导入官方 tokenizer / 转换官方权重时安装）
pip install sentencepiece safetensors
```

启动训练时会打印一次环境自检（CuPy 版本、CUDA 驱动版本、设备名、总显存）。
可用 `python -m tx_pretrain.check_env` 做更完整的环境自检（设备信息 + cuBLAS / JIT 头文件 / cuRAND 三类探针 + 显存峰值）。

### CuPy 版本兼容与精度限制

- `tx_model/dtype.py` 对 `bfloat16` 做**可选注册**：CuPy 不提供 `bfloat16` dtype，`resolve("bfloat16")` 会给出
  明确中文报错；`tx_pretrain` / `tx_sft` / `tx_grpo` 在 CuPy 12/13/14 上均可导入。
- `tx_model/device.py` 兼容 CuPy 14：设备名走 `cuda.runtime.getDeviceProperties`，显存池 getter 走顶层
  `cupy.get_default_memory_pool()`，`compute_capability` 兼容旧版 `tuple(6,1)` 与 CuPy 14 的字符串 `"61"`。
- **Pascal（CC 6.1，如 GTX 1080 Ti）无 bf16 硬件支持**，`dtype` 只能用 `float32`（或 fp16 存激活）。

## CPU / GPU 边界（强制约定）

1. 张量代码只写 `xp.*`（`tx_model/xp.py` 暴露 `xp = cupy`），不出现 `import numpy as np` 混用数组计算；
2. `numpy.ndarray` 仅允许出现在**分词、数据文件读取、日志、checkpoint 落盘**四处；
3. CPU ↔ GPU 转换只允许通过 `tx_model/device.py` 的 `to_device` / `to_host`；
4. 训练 step 内**禁止隐式同步**（不 `float(loss)`、不 `.tolist()`、不 `print(array)`），同步点只能在日志 / 保存周期处。

## 运行方式（远程单卡）

```bash
# 1) 训练 BPE 词表（CPU，可在本机执行）；--out 即 HuggingFace 规范目录
#    语料在 /data/llm/tx-llm-base/corpus/ 下，产物落仓库内 tx_data/tokenizer/
python -m tx_bpe.train_bpe \
    --corpus /data/llm/tx-llm-base/corpus/smoke_5m.txt \
    --out tx_data/tokenizer --vocab-size 6144

#    100MiB 正式档耗时数小时，务必后台长跑 + 日志重定向（-u 保证日志秒级落盘）；
#    阶段/心跳/ETA 等日志参数与 tail 观察见 tx_bpe/README.md「日志与后台运行」：
nohup python -u -m tx_bpe.train_bpe \
    --corpus /data/llm/tx-llm-base/corpus/general_100mi.txt \
    --out tx_data/tokenizer --vocab-size 6144 \
    --log-every 100 --log-every-seconds 5 \
    > /data/llm/tx-llm-base/logs/train_bpe_100mi.log 2>&1 &

# 2) 语料 token 化 -> 二进制分片（CPU，可在本机执行）
python -m tx_pretrain.tokenize_corpus \
    --corpus /data/llm/tx-llm-base/corpus/smoke_5m.txt \
    --out /data/llm/tx-llm-base/tokens_100mi \
    --tokenizer tx_data/tokenizer

# 3) 预训练（GPU）；tx_pretrain/config.json 默认已是单卡 11GB 可用的小配置
#    （d_model 768 / 12 层 / n_head 16 / n_kv_head 8 / ffn 2048 / max_seq_len 512 / train.seq_len 512，约 87.31M 参数）
#    长序列建议加 --chunked_attention --chunk_size 256 降低 score 显存峰值
CUDA_VISIBLE_DEVICES=0 python -m tx_pretrain.train --config tx_pretrain/config.json \
    --chunked_attention --chunk_size 256

# 4) SFT 数据准备（CPU / IO）：BelleGroup 多轮数据集 → data/sft/{train,val}.jsonl
#    离线推荐先用 huggingface-cli 下载再转换；不给 --input 则经 datasets 从 Hub 拉取
huggingface-cli download BelleGroup/multiturn_chat_0.8M \
    --repo-type dataset --local-dir /data/llm/tx-llm-base/sft/multiturn_chat_0.8M
bash tx_sft/prepare_belle.sh -f        # 内部：python -m tx_sft.prepare_belle --input ...

# 5) SFT / GRPO（GPU）
CUDA_VISIBLE_DEVICES=0 python -m tx_sft.train --config tx_sft/config.json
CUDA_VISIBLE_DEVICES=0 python -m tx_grpo.train --config tx_grpo/config.json

# 冒烟测试：缩小同构配置、少量 step，仍必须在有 CUDA 的机器上运行
# （缺语料 / 词表时自动退化为随机 token 与仅字节级最小词表，仅验证通路）
CUDA_VISIBLE_DEVICES=0 python -m tx_pretrain.train --config tx_pretrain/config.json --smoke-test
CUDA_VISIBLE_DEVICES=0 python -m tx_grpo.train --config tx_grpo/config.json --smoke-test
```

三个训练入口都支持 `--override a.b.c=value`（点号路径，可重复）覆盖任意配置项，例如
`--override sft.seq_len=512 --override grpo.inner_steps=4`。

SFT 使用 **Qwen-2.5 ChatML** 模板（角色字面量 `system` / `human` / `assistant`，`user` 按数据集改为 `human`）：
多轮对话顺序拼接进同一条序列（不做 BOS/EOS 重置），`<|im_start|>` / `<|im_end|>` 不扩词表（作为普通文本由 BPE 编码），
仅 **assistant 回复段（含结尾 `<|im_end|>`）** 计入 loss。模板见 `tx_sft/chat_template.py`，
数据准备见 `tx_sft/prepare_belle.py`，基座与词表须与 `tx_pretrain/config.json` 一致（`vocab_size=6144`，`seq_len ≤ 512`）。
SFT 使用**自实现训练循环**（`tx_sft/train_loop.py:SFTTrainer`，算子仍复用 `tx_model`），不复用预训练的 `Trainer`。

SFT 的序列打包（`sft.pack=true`）会给每条样本分配 `segment_ids`，模型据此构造**块对角因果掩码**
（`tx_model/ops_attn.py: build_block_diag_mask`），打包后不同样本之间互不可见，padding 位置同样不可见。

显存吃紧时用命令行覆盖参数调节，不修改配置文件：

```bash
CUDA_VISIBLE_DEVICES=0 python -m tx_pretrain.train --config tx_pretrain/config.json \
    --micro-batch 1 --grad-accum 8
```

远程长跑建议 `tmux` / `nohup` + 日志重定向，本机只做代码同步与日志查看。

## 本机可运行的纯 CPU 单测

不依赖 CuPy 的逻辑（分词器 / 数据格式 / Chat 模板 / label mask / 奖励函数）：

```bash
python -m pytest tests_bpe tests_sft tests_grpo -q
```

## 一致性验证

- `tx_model/gradcheck.py`：GPU 上对 RMSNorm / RoPE / GQA Attention / SwiGLU 做有限差分梯度校验（fp32 阈值 ~1e-2）；
- `tx_model/convert_hf.py`：官方 Llama-2 `safetensors` 权重转 `npz`，加载后与 HF 前向逐 token 比对（误差 < 1e-3）；
- `tx_model/metrics.py`：参数量自检（7B 配置应 ≈ 6.7B），用于快速发现结构错误。

`tx_bpe` 训练产物即 HuggingFace 规范单目录，可被 `transformers` 原生加载（与本地编解码逐样本比对，
见 `tests_bpe/test_hf_format.py`，未装 `transformers` 时该用例自动跳过）：

```python
from transformers import AutoTokenizer
tok = AutoTokenizer.from_pretrained("tx_data/tokenizer")   # PreTrainedTokenizerFast
assert tok.decode(tok.encode("hello 世界")) == "hello 世界"
```