sow-arith-compare / README.md
tritueviet's picture
sow-space: chạy backbone ternary bằng kernel số nguyên — 1,50× tok/s
a14bd1e verified
|
Raw
History Blame Contribute Delete
9.83 kB
metadata
title: sow-match vs sow-ternary
emoji: 🧮
colorFrom: gray
colorTo: indigo
sdk: docker
app_port: 7860
pinned: false

sow-match vs sow-ternary

https://huggingface.co/spaces/ml-intern-explorers/sow-arith-compare

Hai model 191M params, train từ đầu trên cùng một task số học cộng/trừ, phục vụ trong một Space duy nhất và chạy song song trên cùng một prompt — hai panel stream cạnh nhau.

sow-match sow-ternary
Trọng số fp32 đặc (nn.Linear) ternary {−1, 0, +1} + activation int8 (BitNet b1.58 QAT + STE)
Backbone lượng tử không 72 ma trận (12 layer × 6 Linear)
Checkpoint 764 MB (.pth) 233 MB (2-bit packed .pt, 3,27×)
Train data ADD_base_8M — vét hết (a,b) ∈ [−999,999]² cho +/, 1 epoch cùng ADD_base_8M

Cả hai đều là L12/D768, head 64, vocab 65536, chạy RNN-mode (token-by-token, không cần CUDA kernel) trên CPU.

Định dạng đầu vào

Cả hai checkpoint chỉ được train trên đúng một dạng câu:

<số nguyên> <+ hoặc -> <số nguyên> = ?

Đúng: 561 + 432 = ? · -605 + 264 = ? · 723 - 909 = ? · -757 + -567 = ? Sai: 561+432 (thiếu dấu cách) · 5 * 3 = ? (chỉ +/) · 561 + 432 (thiếu = ?) · five plus three = ? (phải là chữ số)

UI kiểm tra định dạng ngay khi gõ và cảnh báo nếu lệch, hoặc nếu toán hạng vượt ±999.

Đã đo được gì

Exact-match 500 câu ADD_base_test (greedy) — hai model giống nhau từng điểm một:

trong ±999 ngoài ±999 (extrapolation) tổng
sow-match (fp32, 764 MB) 118/118 = 100% 0/382 = 0% 118/500 = 23,6%
sow-ternary (1.58-bit, 233 MB) 118/118 = 100% 0/382 = 0% 118/500 = 23,6%

Vì cùng data và cùng shape, đây là so sánh sạch: ternary 1.58-bit không mất điểm nào so với fp32, trong khi checkpoint nhỏ hơn 3,27×. Con số tổng 23,6% thấp chỉ vì 382/500 câu của test set có toán hạng ngoài ±999 — ngoài không gian train, nên cả hai đều sai. Sai ở đó là giới hạn đã biết, không phải bug.

Tốc độ — vì sao bản ternary nhanh hơn

Ban đầu bản ternary chậm hơn bản fp32: load_any() nở 72 ma trận 2-bit ngược thành fp32 (764 MB RAM, đúng bằng bản dense), rồi mỗi bước vẫn nhân fp32 cộng thêm chi phí fake-quant activation. Tức là trả đủ giá của lượng tử hoá mà không thu được gì.

Nguyên tắc của engine hiện tại, cũng là nguyên tắc chung của mọi runtime low-bit chạy được: code không bao giờ nở ra fp32, và dot product chạy hẳn ở miền số nguyên. Lượng tử hoá chỉ trả công khi nó đi thẳng vào kernel; nở ngược ra float lúc load là vứt bỏ toàn bộ phần lời.

trước sau
Backbone trong RAM 340 MB fp32 85 MB int8 code
Đỉnh RSS lúc load 1440 MB 1074 MB
Matvec backbone fp32 GEMV + fake-quant int8 × int8 → int32 accumulate, nhân scale 1 lần
Recurrence state @ ab — batched matmul O(N³) (state @ -kk) ⊗ (kk*a) — rank-1, O(N²)

Ba thay đổi, theo thứ tự tác động:

  1. Backbone chạy kernel nguyên. Ma trận ternary = code int8 {−1,0,+1} × một scale alpha, và activation cũng đã là int8 × một scale 1/s, nên cả phép nhân rút gọn thành

    y = (alpha / s) · (T_int8 @ a_int8)        T ∈ {−1,0,+1}, a ∈ [−128,127]
    

    — một tích vô hướng nguyên, cộng dồn int32, chỉ nhân scale đúng một lần ở cuối. Đây không phải xấp xỉ của phép toán QAT mà chính là nó: đường fp32 cũ chỉ đang thổi cả hai vế thành float rồi nhân lại đúng những tích đó. Kernel SIMD lấy từ quantized.linear_dynamic (fbgemm/qnnpack). Decode ở batch=1 nghẽn băng thông bộ nhớ chứ không nghẽn FLOP, nên đọc 1 byte/trọng số thay vì 4 mới là phần lớn của khoản lời — đó cũng là lý do khoản lời này chỉ có được khi trọng số ở lại dạng code. Ghi chú: backend int8 trên x86 tiền-VNNI thường phải hạ activation xuống 7 bit (reduce_range) vì bước nhân-cộng trung gian tích luỹ ở int16 và có thể tràn. Với trọng số ternary thì tổng riêng bị chặn ở 2·255, không đời nào tràn, nên tắt được reduce_range và lấy lại đủ 8 bit activation.

  2. Recurrence dạng rank-1. ab là tích ngoài (−kk) ⊗ (kk*a), nên state @ ab bằng (state @ (−kk)) ⊗ (kk*a): hai phép O(N²) thay cho một batched matmul O(N³). Docstring của engine vốn đã mô tả đúng dạng này, chỉ có code là làm theo dạng đắt. Áp dụng cho cả hai model.

  3. Vi tối ưu. Sáu vector token-shift của mỗi block gộp thành một tensor (6, D) (1 phép nhân broadcast thay 12 phép elementwise); các tham số (1,1,D) được flatten một lần lúc load thay vì view(-1) ~100 lần mỗi token; unpack_codes giải mã bằng số học (v&1) − (v>>1) thay vì gather qua LUT (gather cần index int64 = 680 MB cấp phát vô ích cho 85M trọng số).

Đo trên 1 thread (i5-9400F, AVX2, không VNNI), 25 token decode:

ms/token tok/s
sow-ternary — engine cũ 38,6 25,9 chậm hơn cả bản dense
sow-ternary — engine mới 25,7 38,9 1,50×
sow-match (dense) — cũ → mới 34,0 → 31,8 29,4 → 31,4 1,07× (chỉ được rank-1 + vi tối ưu)

Kiểm chứng. Chạy lại toàn bộ 500 câu ADD_base_test bằng cả hai engine trên chính rwkv-final-packed.pt: 118/118 đáp án in-range giống nhau từng byte, và verdict đúng/sai khớp trên cả 500 câu. 30 câu có prediction khác nhau đều là câu extrapolation mà cả hai engine đều sai — ở đó output là rác và cực nhạy với sai số làm tròn. Bản dense cũng cho output y hệt bản cũ.

Trên phần cứng không có backend số nguyên, engine tự động quay về đường fp32 cũ. Ép tắt để so sánh A/B: SOW_TERNARY_INT8=0.

Layout

sow-space/
├── Dockerfile             # HF Spaces (Docker SDK) build, listen :7860
├── requirements.txt
├── sow_space_app.py       # Gradio Blocks: 1 prompt -> 2 panel stream song song (entrypoint)
├── sow_space_model.py     # engine inference RNN-mode dùng chung cho cả hai checkpoint
├── tokenizer/
│   ├── sow_tokenizer.py   # TRIE World tokenizer
│   └── sow_vocab.txt      # vocab 65536
├── weights/
│   ├── sow-match-dense.pth   # dense fp32, 764 MB (git-lfs)
│   └── sow-ternary-packed.pt # ternary 2-bit packed, 233 MB (git-lfs)
└── NOTICE

Thiếu file weights nào thì panel tương ứng báo "checkpoint chưa có" — Space vẫn build và chạy.

Chạy local

pip install -r requirements.txt
pip install torch --index-url https://download.pytorch.org/whl/cpu
python3 sow_space_app.py          # http://localhost:7860

Hoặc build đúng image của Space:

docker build -t sow-space .
docker run --rm -p 7860:7860 sow-space

Push lên Hugging Face Spaces

git lfs install
hf auth login
git remote add hf https://huggingface.co/spaces/ml-intern-explorers/sow-arith-compare
git push hf main

Khi tạo Space mới, HF tự sinh sẵn một commit "initial commit" ⇒ push đầu tiên bị rejected. git fetch hf main && git rebase hf/main, giữ README của mình (có app_port), lấy .gitattributes của HF (danh sách LFS đầy đủ).

Ghi chú kỹ thuật

  • Mỗi request dựng lại state RNN từ prompt (stateless giữa các request) ⇒ nhiều người dùng cùng lúc không giẫm lên state của nhau.
  • Hai model decode trong hai thread riêng, mỗi thread 1 intra-op thread của torch (torch.set_num_threads(1)) — trên cpu-basic 2 vCPU thì như vậy nhanh hơn để cả hai cùng spawn N thread rồi tranh core.
  • Checkpoint packed đã chứa giá trị ternary chính xác, nên loader bỏ qua weight_quant lúc load (chạy lại sẽ rescale theo mật độ của chính pattern ternary thay vì theo biên độ đã train).
  • Ở đường int8, kernel tự lượng tử hoá activation nên activation_quant() trong Python bị tắt — giữ lại chỉ là chi phí thừa: đo thử thì nó tốn 4 ms/token mà không đổi được câu nào trong 500 câu test.
  • Còn tối ưu thêm được — mặc định tắt: SowRNN(..., head_int8=True) lượng tử hoá nốt head.weight (65536×768, fp32, chiếm ~10 ms/token) sang int8 per-row. Đo được: RSS 965 → 836 MB25,7 → 21,7 ms/token (1,78× so với engine cũ) — tức là vừa nhẹ RAM hơn vừa nhanh hơn, không có đánh đổi về tốc độ. Đánh đổi nằm ở chỗ khác: QAT không ternary hoá ma trận này, nên bật cờ là thêm một bước post-training quantization mà checkpoint chưa từng được train cho. Chạy đủ 500 câu ADD_base_test thì toàn bộ đáp án in-range vẫn byte-identical với engine cũ, nhưng vì nó đổi numerics của model đang phục vụ (và làm nhãn "ternary 1.58-bit" cần thêm chú thích) nên để tắt làm mặc định — bật khi nào cần RAM/tốc độ hơn là cần mô tả sạch.