Additive-Rand-Transformer: TinyGPT Arithmetic Toolkit & Manual

Complete Pipeline for Training, Probing, Evaluating, and Quantizing Compact Autoregressive Models (CPU / Colab GPU)

Open In Colab License: MIT PyTorch Hugging Face


Project Overview & Implemented Capabilities

This project provides an end-to-end research framework for training, evaluating, and quantizing compact autoregressive Transformers (TinyGPT, full ASCII-128 vocabulary where token id == byte value, parameter count ~62K–1.9M) on dynamically generated multi-digit addition and subtraction tasks. Sequences are wrapped in STX(0x02)/ETX(0x03) delimiters; subtraction supports negative results.

Core Modules & Architectures:

  1. Dynamic Data Generator: Generates 1–4+ digit arithmetic as STX..ETX plain expressions a op b = c and carry-resolution chains a op b = a1 op b1 = ... = c (one carry/borrow resolved per step, low digit upward); subtraction no longer swaps a/b, so negative results appear directly (a-b=-k). Counter-examples (gen_negative) with a rule verifier are available for positive/negative mixing (data.py, chain_data.py).
  2. Model Architectures: Configurable decoder-only architectures supporting standard Causal Self-Attention, Dynamic Sparse Attention (DSA), ALiBi relative positional biases, Mixture of Experts (MoE), and Low-Rank Adaptation (LoRA) (model.py).
  3. Comprehensive Sweep Matrix: Over 200 systematic configurations evaluating depth (=1\dots10$), channel width (=32\dots512$), training steps (00\dots1024K$), and dropout rates.
  4. Post-Training Quantization: Dynamic PyTorch INT8 and simulated INT4 quantization pipelines assessing size compression and latency speedups (quantize.py).
  5. Reinforcement Learning Pipelines: Self-play RL and Group Relative Policy Optimization (GRPO) with novelty and memory tracking (train.py).
  6. Unified Experiment Ledgers: Master matrices archived in EXPERIMENTS_ALL.xlsx.

Quickstart & Navigation


📚 Documentation:

  • docs/FORMAT.md — 数据格式权威说明(ASCII-128 码表 / STX·ETX / 普通式 / 进位消解链 / 反例 / 答案提取规则)
  • docs/EXPERIMENTS.md — 实验系列汇总(格式换代 → 链式×10 → 负结果 → 配比探索 → 长度泛化 → 炸炉 → 训练对比)

Step 1: Quickstart on Google Colab

Train, evaluate, and quantize models using free cloud GPUs/CPUs without local resource overhead:

  1. Download Colab_Run_Additive_Transformer.ipynb.
  2. Visit Google Colab -> Click Upload -> Select the .ipynb file.
  3. Click Runtime -> Run All (or execute cells interactively).

Step 2: Local Installation & Environment Setup

# 1. Clone repository
git clone https://huggingface.co/Hana-ame/additive-rand-transformer
cd additive-rand-transformer

# 2. Install dependencies (Python >= 3.8, PyTorch >= 2.0)
pip install torch openpyxl huggingface_hub pandas matplotlib
pip install -e .

On this dev machine (Hana-ame local): dependencies are ALREADY installed — use /mnt/d/WorkPlace/my_rl_project/venv/bin/python, do NOT pip install again (see knowledge-base simple-ai-runtime node). System python3 has no torch.


Step 3: Interactive Inference (REPL)

Option A: Shell Script

# Solve single problem
./use_model.sh -s "1234 + 5678"
./use_model.sh -s "9999 - 4321"

# Launch interactive REPL session
./use_model.sh -i

Option B: Python API

import torch
from huggingface_hub import hf_hub_download
from additive_rand_transformer.model import TinyGPT, TinyGPTConfig, TOK_TO_ID
from additive_rand_transformer.data import _int_to_tokens, decode

# 1. Download pretrained weights (best model: chain + negative results, 40k steps)
ckpt_path = hf_hub_download(
    repo_id="Hana-ame/additive-rand-transformer",
    filename="checkpoints/rope_l4_d128_bias05_final_chain_neg_s40000.pt",
    local_dir="."
)

# 2. Instantiate model
ck = torch.load(ckpt_path, map_location="cpu", weights_only=False)
cfg = TinyGPTConfig(**{k: v for k, v in ck["config"].items() if k in TinyGPTConfig.__dataclass_fields__})
model = TinyGPT(cfg)
model.load_state_dict(ck["model"])
model.eval()

STX, ETX, EQ = TOK_TO_ID["\x02"], TOK_TO_ID["\x03"], TOK_TO_ID["="]

# 3. Autoregressive generation: prompt is `STX a op b =`; model completes
#    the answer directly (plain models) or a carry-resolution chain ending
#    in the answer (chain models).
def solve(a: int, b: int, op_char: str = "+"):
    prompt = [STX] + _int_to_tokens(a) + [TOK_TO_ID[op_char]] + _int_to_tokens(b) + [EQ]
    ids = list(prompt)
    with torch.no_grad():
        for _ in range(120):
            x = torch.tensor([ids], dtype=torch.long)
            logits, _ = model(x, None)
            nxt = int(logits[0, -1].argmax())
            ids.append(nxt)
            if nxt == ETX:
                break
    out = decode(ids)
    ans = out[1:-1].rsplit("=", 1)[-1].strip()          # 最后一个 '=' 之后(去空格)
    print(f"Equation: {a} {op_char} {b}  ->  {out}  (answer {ans})")

solve(1234, 5678, "+")
solve(52, 194, "-")        # 负结果

也可直接用 ./use_model.sh -s "37 + 85"(单题/批量/REPL/judge 都支持)。


Step 4: Custom Training via config.json

1. Configure config.json

{
  "layers": 4,
  "d": 128,
  "heads": 4,
  "steps": 4000,
  "batch_size": 32,
  "lr": 3e-4,
  "wd": 0.1,
  "warmup": 200,
  "datasource": {
    "type": "cot",
    "max_digits": 4,
    "bias": 0.5,
    "max_spaces": 3,
    "single": true
  }
}

vocab_size defaults to the full ASCII-128 codebook. datasource.type: "cot" is a legacy wrapper that now produces plain STX..ETX expressions; use --chain for carry-resolution chains.

2. Launch Trainer

# Train from config (plain format via the legacy "cot" wrapper)
python -m additive_rand_transformer.train --config config.json

# Carry-resolution chain data source (recommended; supports negative results)
python -m additive_rand_transformer.train --chain --max_digits 4 --batch_size 32 \
       --lr 3e-4 --wd 0.1 --warmup 200 --seed 0 --steps 40000 \
       --n_layer 4 --n_embd 128 --n_head 4 --pos_enc rope --four_digit_bias 0.5

# Mix in counter-examples (0.01 = 100:1 positives:negatives) / Engram memory
python -m additive_rand_transformer.train --chain --neg_rate 0.01 ...
python -m additive_rand_transformer.train --chain --engram 16 --engram_table 4096 ...

# 50-step smoke verification
python -m additive_rand_transformer.train --quick

Step 5: Dynamic INT8 Quantization Benchmark

Quantize linear layers via PyTorch dynamic INT8 quantization to evaluate footprint compression and throughput:

python -m additive_rand_transformer.quantize --checkpoint checkpoints/l4_d128_cot_bias05_final.pt
  • Empirical Findings: 3.8x compression (1.7MB -> 0.45MB), 1.4x throughput acceleration, with 1–4 digit accuracy 100% preserved (zero degradation).

Step 6: Self-Play RL Training

Launch autonomous self-prompting and verification self-play reinforcement learning:

python -m additive_rand_transformer.rl_selfplay     --checkpoint checkpoints/l4_d128_cot_bias05_final.pt     --reward_mode both     --memory_bonus 1.5     --min_digits_reward 2     --runs_dir runs/selfplay

Pretrained Model Checkpoints

Checkpoint Identifier Architecture Parameters Key Characteristics
rope_l4_d128_bias05_final_chain_neg_s40000.pt 4L · 128d · RoPE · chain · neg · 40k steps 940,800 Best model: add1-4 ≈ 100/98/98/100, sub1-4 ≈ 100/100/92/92 (incl. negative results)
l4_d128_cot_bias05_final_chain_neg_s40000.pt 4L · 128d · chain · neg · 40k steps 940,800 Chain source, negative-result support, 40k-step SFT
l4_d128_cot_bias05_final.pt 4L · 128d · plain STX/ETX 940,800 Plain-format SFT baseline (4000 steps)
l4_d128_grpo_final.pt 4L · 128d · GRPO RL 940,800 GRPO fine-tune from the SFT final
l4_d128_reinforce_conservative.pt 4L · 128d · REINFORCE 940,800 Standard policy-gradient comparative baseline
l4_d128_selfplay_anchor.pt 4L · 128d · Self-Play 940,800 Self-play with difficulty anchoring
l2_d64_attn_dsa.pt 2L · 64d · DSA 173,824 2-layer dynamic sparse attention baseline
l2_d64_attn_causal.pt 2L · 64d · Causal 173,824 2-layer causal attention baseline

还有按实验系列命名的变体:*_chain.pt(链式源 4000 步)、*_s40000.pt(40k 步)、*_negrate*(正/反例配比 100:1 / 1000:1 探索)、epoch1/epoch2、overlap_a*(数据集重叠度探索)等,全部为 ASCII-128 词表。旧 vocab-16/32 的 checkpoint 已删除。


Experiment Excel Workbooks

All experimental configurations and quantitative evaluations are compiled in:
👉 EXPERIMENTS_ALL.xlsx


License

This repository is distributed under the MIT License.

Downloads last month

-

Downloads are not tracked for this model. How to track
Video Preview
loading