BS-TDM: 位切片时分复用大语言模型

1 物理层 × 6 步位切片 × 4 bit = 24 bit/w (8× 压缩) — 30M 参数超越标准 6 层 Transformer, 部署推理 2.6× 加速, ppl 无损。

快速开始

推理 (两种模式)

# 方式 1: 部署态推理 (推荐, 2.3x 加速)
#   4-bit 量化权重打包存储, 加载时解包为 bf16 缓存, 免实时量化开销
python inference_deployed.py

# 方式 2: 训练态推理 (对照)
#   fp32 权重 + 每步实时 4-bit 量化
python inference.py

推理性能对比 (实测, batch=1, seq=128, RTX 4090)

三种模式 ppl 完全一致 (17.06),量化无损。

• bf16 缓存部署 (A): 训练后 4-bit 量化 → 一次性解包为 bf16 缓存 → 推理纯 GEMM, 2.6× 加速 + 显存最低。小模型首选。
• 原生 4-bit 打包 (B): 权重以 uint8 (2 值/字节) 存 GPU, 每步 forward 时解包。量化权重仅 4.9 MB (vs fp32 39.3MB = 8× 压缩)。每步解包有开销, 速度与训练态持平。大模型 (B 级) 首选 (配合 int4 GEMM 可加速)。
• 30M 模型差异小 (embedding/head fp32 占 60MB 主导); 模型越大, 量化权重占比越高, 原生 4-bit 优势越大。

训练复现

# 全流程: 预训练(3ep) + SFT(2ep) + 推理 benchmark + 生成样本
python run_full_pipeline.py

# 三模型对比 (标准6层 vs BS-TDM 2×3 vs 1×6)
python train_timeshare.py

文件清单

模型架构

输入 token IDs
  → nn.Embedding (fp)
  → 1 × TimeSharedBlock (6 步循环, 每步残差):
      for step in [0..5]:
          h = h + Attention(h, step)    ← 标准 MHA + 第 step 组 4-bit 位切片权重
          h = h + FFN(h, step)           ← SwiGLU + 4-bit 位切片权重
  → nn.Linear head (fp)

部署: 每权重 6 步 × 4 bit = 24 bit/w → 量化权重仅 5.0MB (vs fp32 39.3MB = 7.8× 压缩)

核心结果

训练对比 (同逻辑深度 6, minimind 120k, 3 epoch)

标准 Transformer:

frsmash3.7 SSM (直接循环 + 混合 + GLA 突破):

SSM 不能时分 (固定 A 矩阵), 但 GLA 可以 (数据相关门控 g_t)! frsmash3.8 纯 GLA 时分超越 v3.7 且 8× 压缩.

部署推理 (三种模式)

选型建议:

• 小模型 (≤100M): 用 bf16 缓存部署 (最快)
• 大模型 (≥1B): 用 原生 4-bit (最省显存, 配合 int4 GEMM 加速)

全流程 (预训练 → SFT → 推理)

采样参数

generate(model, voc, "你的提示", dev,
         max_new=120, top_k=50, top_p=0.95, rep_penalty=1.2, temp=0.8, seed=0)

环境

• Python 3.12, PyTorch 2.13+cu126, Triton 3.7
• CUDA GPU (sm_89 / RTX 4090 测试通过)
• Windows 需 PYTHONUTF8=1

引用

详见 paper.md / paper_en.md。