Add files using upload-large-folder tool

README.md

@@ -1,390 +1,264 @@
----
-library_name: transformers
-pipeline_tag: text-generation
-language:
-  - zh
-  - en
-license: other
-license_name: baihu-custom-license
-license_link: https://huggingface.co/NovaAI6868/BaiHu-V1-Flash/blob/main/LICENSE.custom.md
-base_model: Qwen/Qwen3-0.6B-Base
-tags:
-  - sparse-attention
-  - subq
-  - ssa
-  - long-context
-  - commercial-license-required
-  - text-generation
----
-
-# BaiHu-V1-Flash
-
-**BaiHu-V1-Flash** is a retrofit of the dense-attention model `Qwen/Qwen3-0.6B-Base` into an
-**SSA (Sparse-attention + SubQ)** architecture, obtained by continued pretraining.
-
-- Base model: `Qwen/Qwen3-0.6B-Base` (28 layers / 16 Q heads / 8 KV heads / head_dim 128 / 32K context / tied embeddings)
-- Parameters: 598.8M
-- Training data: mixed Chinese + English (Fineweb-Edu-Chinese-V2.1 + fineweb-edu, 50/50)
-- License: **free for personal use; a paid license is required for commercial use** (see "License" below)
-
----
-
-## 1. Architecture: SSA (three paths, each with its own softmax, then summed)
-
-Every layer keeps the base model's MLP / RMSNorm weights and replaces full attention with an
-SSA layer built from three parallel paths:
-
-| Path | Role | Complexity |
-|---|---|---|
-| `shared` | every query sees all **completed** blocks through one compressed vector per block | `O(T·T/B)` |
-| `local` | dense causal attention over the most recent window | `O(T·w)` |
-| `sparse` (**SubQ**) | only 4 of 16 query heads produce block scores, shared across the whole head group; real attention is computed only for the selected top-k blocks | `O(T·k·B)` |
-
-### Hyperparameters
-
-| Parameter | Value | Meaning |
-|---|---|---|
-| `ssa_block_size` | 64 | block size B |
-| `ssa_top_k` | 8 | number of blocks selected by the sparse path |
-| `ssa_local_blocks` | 2 | local window = 3 × 64 = 192 tokens |
-| `ssa_num_subq_heads` | 4 | SubQ heads, r = 16 / 4 = 4 |
-| `ssa_router_dim` | 128 | router subspace dimension |
-| `ssa_compress_dim` | 128 | block compression dimension |
-
-Only **2.75M** parameters are new (≈0.46% of the model); all other weights are inherited
-from the base model.
-
----
-
-## 2. Training
-
-| Item | Setting |
-|---|---|
-| Starting point | a conversion checkpoint that is **bit-exact** with the base model (`max\|Δlogit\| = 0.000e+00`) |
-| Tokens seen | 5.0M (≈0.25 epoch of the corpus) |
-| Sequence length | 512 (must be a multiple of `ssa_block_size = 64`) |
-| Effective batch | 8192 tokens (batch 2 × grad_accum 8) |
-| Precision | fp32 |
-| Optimizer | SGD with momentum 0.9 |
-| Learning rate | 5e-4, 50-step warmup, cosine decay to 10% |
-| Hardware | single NVIDIA GTX TITAN X (Maxwell, sm_52, 12.9 GB) |
-| Throughput | ≈274 tokens/s |
-
-### Critical issues found and fixed during this retrofit
-
-Several defects silently break training and are worth documenting:
-
-1. **Both new branches had identically zero gradients (blocking).** To make the converted
-   model bit-exact with the base model, `compress_out` and `router_out` were initialized to
-   exactly zero, and the branches were skipped entirely by a gate. The branch output was
-   therefore always zero, so the back-propagated gradient was also always zero: all 2.75M
-   SSA parameters **stayed frozen for the entire run** and the sparse attention was dead
-   code. The fix is a small non-zero initialization.
-2. **Routing was non-differentiable.** `top-k` produces hard indices, and indexing is not
-   differentiable. If the routing scores are used only to decide *which* blocks to read and
-   never enter the softmax, the gradients of `router_q` / `router_k` are **exactly zero** —
-   the router can never learn to route. The fix is to feed the selected blocks' scores,
-   squashed through `tanh` and gently scaled, into the attention logits as an additive bias.
-3. **The shared summary was a sum, not a mean.** Its magnitude grew linearly with the
-   prefix, and because that branch is injected at full weight (a single-element softmax has
-   probability exactly 1), it swamped the residual stream: hidden states grew from 0.2 to
-   about 7 in layer 0 and to about 1900 by layer 27, and validation loss went 3.56 → 11.38.
-   The fix is to divide by the token count.
-4. **The shared branch needs an explicit gate.** With a single-element softmax the
-   probability is always 1, so the initialization scale of `compress_out` cannot control the
-   injection strength at all (measured: scales from 1e-4 to 0.03 all left the loss at
-   exactly 7.3526). A learnable scalar gate, initialized to a small positive value, lets the
-   optimizer decide how far to open it.
-
----
-
-## 3. How to Run Inference
-
-### 3.1 Important: this is a custom architecture
-
-`BaiHu-V1-Flash` uses `model_type: baihu_ssa`, which is **not** in the Transformers
-registry. Loading it with a plain `AutoModelForCausalLM.from_pretrained(...)` fails with:
-
-```
-ValueError: The checkpoint you are trying to load has model type `baihu_ssa`
-but Transformers does not recognize this architecture.
-```
-
-You must register the config and model classes first. This is a one-time, three-line step
-(see below).
-
-Also note: **this model does not go through `transformers.GenerationMixin`.** SSA owns its
-own KV-cache layout (`BaiHuSSACache`) because the cache stores per-block compressed prefix
-sums rather than a plain growing key/value tensor. Call the model's own `generate()` method;
-beam search is not supported.
-
-### 3.2 Setup
-
-```bash
-# The SSA implementation lives in the project repository (not on the Hub),
-# because it is a custom architecture.
-git clone <this-project-repo> ssa_model
-cd ssa_model
-uv venv --python 3.11 .venv
-
-# Ampere or newer (RTX 30xx/40xx, A100, ...): any recent torch works.
-uv pip install --python .venv/Scripts/python.exe \
-    "numpy>=1.26" "transformers>=4.51" "safetensors>=0.4" "torch>=2.6"
-
-# Maxwell / Pascal / Volta (GTX 9xx/10xx, TITAN X, V100, ...): pin torch 2.7.1 —
-# see the GPU note below for why `torch>=2.6` would resolve to a build without your kernels.
-uv pip install --python .venv/Scripts/python.exe \
-    "numpy>=1.26" "transformers>=4.51" "safetensors>=0.4" "torch==2.7.1"
-```
-
-> **GPU note (Maxwell and older).** If you are on an NVIDIA Maxwell card such as the
-> GTX TITAN X (compute capability sm_52), recent PyTorch wheels no longer ship kernels for
-> it: PyTorch **2.8 removed sm_50/sm_60**, and 2.8.0+cu126 only contains `sm_61…sm_90`.
-> You will get `CUDA error: no kernel image is available for execution on the device`.
-> Use **torch 2.7.1+cu126**, which still contains `sm_50`. Also train/infer in **fp32** on
-> Maxwell: measured fp32 5.30 / fp16 4.46 / bf16 3.20 TFLOPS, so half precision is a loss.
-
-### 3.3 Minimal working example
-
-```python
-import os
-import sys
-
-import torch
-from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer
-
-# Make the custom implementation importable (path to the cloned repo's src/)
-sys.path.insert(0, os.path.join("ssa_model", "src"))
-
-from baihu_ssa.configuration_baihu_ssa import BaiHuSSAConfig
-from baihu_ssa.model_baihu_ssa import BaiHuSSAForCausalLM
-
-# ---- register the custom architecture (REQUIRED) ----
-AutoConfig.register("baihu_ssa", BaiHuSSAConfig, exist_ok=True)
-AutoModelForCausalLM.register(BaiHuSSAConfig, BaiHuSSAForCausalLM, exist_ok=True)
-
-REPO = "NovaAI6868/BaiHu-V1-Flash"
-device = "cuda" if torch.cuda.is_available() else "cpu"
-dtype = torch.float32 if device == "cpu" else torch.float32  # see GPU note: fp32
-
-tokenizer = AutoTokenizer.from_pretrained(REPO)
-model = AutoModelForCausalLM.from_pretrained(REPO, dtype=dtype).to(device).eval()
-
-prompt = "人工智能的未来是"
-input_ids = tokenizer(prompt, return_tensors="pt").input_ids.to(device)
-
-with torch.no_grad():
-    out = model.generate(
-        input_ids,
-        max_new_tokens=32,
-        do_sample=False,        # greedy; supported
-        # do_sample=True, temperature=0.8, top_k=50,   # sampling also supported
-    )
-print(tokenizer.decode(out[0], skip_special_tokens=True))
-```
-
-### 3.4 `generate()` parameters
-
-The built-in `generate()` is a self-contained decoding loop (greedy or
-temperature/top-k sampling). Supported arguments:
-
-| Argument | Default | Meaning |
-|---|---|---|
-| `max_new_tokens` | 32 | number of tokens to generate |
-| `do_sample` | `False` | `False` = greedy, `True` = sample |
-| `temperature` | 1.0 | sampling temperature (used when `do_sample=True`) |
-| `top_k` | `None` | top-k sampling cutoff (used when `do_sample=True`) |
-| `eos_token_id` | `None` | stop early when all sequences emit this token |
-| `use_cache` | `True` | keep the SSA cache between steps; **leave on**, decoding is much slower without it |
-
-**Not supported:** beam search (the SSA cache has no `reorder_cache`), and
-`transformers` generation utilities such as `logits_processor` / `stopping_criteria`.
-
-### 3.5 Computing perplexity / loss
-
-```python
-with torch.no_grad():
-    inputs = tokenizer("Your text here", return_tensors="pt").to(device)
-    loss = model(**inputs, labels=inputs["input_ids"]).loss
-    print("loss", loss.item(), "ppl", torch.exp(loss).item())
-```
-
-### 3.6 Throughput you should expect
-
-Measured on a GTX TITAN X (sm_52, 12.9 GB), fp32, prompt 1024 tokens, 64 generated tokens:
-
-| | Qwen3-0.6B-Base | BaiHu-V1-Flash |
-|---|---|---|
-| Prefill / TTFT | 503 ms | 1431 ms |
-| Per-token decode (TPOT) | 36.6 ms | 80.5 ms |
-| Decode throughput | 27.3 tok/s | 12.4 tok/s |
-| Peak memory (generation) | 9.91 GB | 9.01 GB |
-
-**BaiHu-V1-Flash is currently ~2.2× slower to decode than the base model**, even though the
-sparse path reads fewer keys. The reason is implementation, not architecture: the attention
-layer loops over query blocks in Python and launches many small kernels, so kernel-launch
-overhead dominates the FLOPs saved. Memory use is lower; speed is not yet a win.
-See "Known Limitations" below.
-
-### 3.7 Expected output quality (honest warning)
-
-This checkpoint has seen only **5.0M tokens**, so the new branches are far from converged.
-Greedy decoding tends to **repeat itself** and it is **worse than the base model on
-perplexity** (ppl 82.3 vs 39.4 at 512 tokens). A real greedy sample:
-
-```
-prompt: 人工智能的未来是
-output: 人工智能的未来是怎样的？人工智能的未来是怎样的？人工智能的未来是怎样的？…
-```
-
-Use it to study or continue the SSA retrofit — **do not** expect it to match
-`Qwen3-0.6B-Base` as a general-purpose model yet.
-
-### 3.8 Troubleshooting
-
-| Error | Cause | Fix |
-|---|---|---|
-| `does not recognize this architecture` / `KeyError: 'baihu_ssa'` | custom architecture not registered | call `AutoConfig.register` + `AutoModelForCausalLM.register` as in 3.3 |
-| `CUDA error: no kernel image is available for execution on the device` | installed PyTorch has no kernels for your GPU (Maxwell/sm_52) | install torch 2.7.1+cu126 or older (2.8 dropped sm_50/sm_60) |
-| `AttributeError: ... has no attribute 'tie_weights'` | an external tool treating the model as `PreTrainedModel` | use a version of the implementation that defines `tie_weights()` (already fixed upstream) |
-| `BaiHuSSACache` errors when calling `model.generate(...)` from GenerationMixin | this model bypasses `GenerationMixin` | call `model.generate(...)` on the BaiHu model itself |
-
----
-
-## 4. Evaluation
-
-### 4.1 Language modeling perplexity (validation set, identical windows)
-
-| Sequence length | Qwen3-0.6B-Base | BaiHu-V1-Flash |
-|---|---|---|
-| 512 | 3.6726 / ppl 39.354 | 4.4108 / ppl 82.335 |
-| 1024 | 3.4691 / ppl 32.109 | 4.2575 / ppl 70.631 |
-| 2048 | 3.0801 / ppl 21.761 | 3.9213 / ppl 50.467 |
-
-### 4.2 Standard benchmarks (lm-evaluation-harness)
-
-| Task | Qwen3-0.6B-Base | BaiHu-V1-Flash | Delta |
-|---|---|---|---|
-| arc_easy | 0.5550 | 0.6250 | +0.0700 |
-| hellaswag | 0.5350 | 0.5200 | -0.0150 |
-| piqa | 0.7050 | 0.6950 | -0.0100 |
-| winogrande | 0.6300 | 0.6300 | +0.0000 |
-
-### 4.3 Inference compute and resource usage
-
-| Metric | Qwen3-0.6B-Base | BaiHu-V1-Flash |
-|---|---|---|
-| Parameters (M) | 596.0500 | 598.8000 |
-| Prefill peak memory (GB) | 9.6200 | 8.0210 |
-| Generation peak memory (GB) | 9.9120 | 9.0050 |
-| Prefill latency (s) | 0.5030 | 1.4310 |
-| TTFT (ms) | 502.7 | 1431.2 |
-| TPOT (ms) | 36.6 | 80.5 |
-| Decode throughput (tok/s) | 27.3300 | 12.4200 |
-| Attention FLOPs/token (GFLOPs) | 0.1176 | 0.1057 |
-| Attention key accesses vs full attention | 1.0000 | 0.8993 |
-| GPU utilization mean/max (%) | 93.9 | 43.3 |
-| Power mean/max (W) | 179.1 | 124.3 |
-
-Positive findings: peak inference memory is lower (generation 9.005 vs 9.912 GB, −9.2%), and the
-model draws less power because it is not compute-bound.
-
-Negative findings, stated plainly:
-
-- **Decode is 2.2× slower** (12.42 vs 27.33 tok/s) and **prefill is 2.8× slower**
-  (TTFT 1431 vs 503 ms), despite the sparse path reading fewer keys. The current
-  implementation loops over query blocks in Python and issues many small kernels, so
-  launch overhead dominates the FLOPs saved. **The sparse attention does not yet pay off
-  on this hardware.**
-- **Attention key accesses are still 89.9% of full attention** at this sequence length.
-  The reason is structural: the local window already covers 3 blocks (192 tokens) and the
-  sparse path reads up to `top_k + 1 = 9` blocks from a grid that only has 16 blocks at
-  1024 tokens, so the selected set is almost the whole grid. Sparsity only becomes a real
-  saving once the sequence is long relative to `top_k × block_size` (i.e. well beyond
-  10k tokens).
-- **Language modeling perplexity is clearly worse than the base model** at every length
-  tested (ppl 82.3 vs 39.4 at 512; 50.5 vs 21.8 at 2048). This is the honest cost of
-  shrinking the dense local window from the full prefix to 192 tokens while the new
-  long-range branches are still very weakly trained.
-- **Standard benchmarks are roughly neutral but not better**: arc_easy improves
-  (+0.070 acc_norm), winogrande is unchanged, while hellaswag (−0.015) and piqa (−0.010)
-  regress slightly.
-
-### 4.4 Sparsity
-
-Two measurements are reported because they answer different questions and are **not**
-interchangeable:
-
-| Scenario | Average blocks selected per query | Key access ratio vs full attention |
-|---|---|---|
-| Chunked forward over a 512-token validation window | 0.38 | 0.0938 |
-| Prefill of 1024 tokens (steady state) | up to `top_k` | 0.8993 |
-
-The first number averages over *all* query blocks including the early ones, which have no
-completed blocks available to select and therefore read nothing through the sparse path.
-The second is the steady-state ratio for later queries, and it is the one that matters for
-efficiency — see the note in section 4.3: at these sequence lengths the sparse path is not
-yet saving meaningful work.
-
----
-
-## 5. Known Limitations
-
-1. **Trained for very little.** Only 5.0M tokens (≈0.25 epoch). The new branches have not
-   converged; ppl is well above the base model and decode is slower (see 3.6 and 4.1). Continuing
-   to 200M tokens or more is required before the SSA layers can genuinely take over
-   long-range modeling.
-2. **The shared branch does not appear to help and was actively suppressed by the
-   optimizer.** Its learnable gate *decreased* over training (0.0100 → 0.0129 at step 100 →
-   0.0122 at step 610) instead of growing, meaning the optimizer found the single
-   prefix-mean summary not worth injecting. This is the single most important thing to
-   change next: replace it with **one compressed vector per block** (same `O(T·T/B)` cost,
-   far more information retained).
-3. **Sparse attention is not yet a net win on this hardware.** It reads fewer keys but runs
-   2.2× slower because the implementation loops over query blocks in Python and issues many
-   small kernels. It needs kernel-level batching (or a fused implementation) before the
-   sparsity can translate into speed.
-4. **Sparsity only pays off at long sequences.** With `top_k=8` and `block_size=64`, the
-   sparse path can read up to 9 blocks = 576 tokens; at 1024 tokens the grid only has 16
-   blocks, so the selected set covers most of the context and the local window already
-   covers the rest. Real savings require sequences well beyond 10k tokens.
-5. **Routing quality is not fully validated.** The distribution of selected top-k blocks
-   should be checked for degeneration (e.g. always selecting the same blocks). The
-   non-differentiable-routing bug that would have made this *impossible* to learn has been
-   fixed (see section 2), but the learned policy has not been analyzed in detail.
-6. **Block size B = 64 was not ablated.** Limited by memory and the Windows WDDM watchdog on
-   this machine; B ∈ {32, 64, 128} should be swept on a larger GPU.
-7. **Document-boundary packing.** The current packing strategy places multiple documents in
-   one sequence (separated by EOS).
-
----
-
-## 6. License
-
-**Free for personal use; a paid license is required for commercial use.**
-
-- ✅ Personal study, research, teaching, hobby projects: **free**, no application required
-- ✅ Academic research with public publication: **free** (please cite the source)
-- 💰 Internal company/studio use, paid API/SaaS, product integration, client deliverables:
-  **commercial license required**
-
-**Commercial licensing contact: novaweb6868@outlook.com**
-
-Full terms: [LICENSE.custom.md](./LICENSE.custom.md).
-
-This model is an architectural retrofit of `Qwen/Qwen3-0.6B-Base` (Apache License 2.0).
-This license governs only the newly added portions and does not alter the upstream
-component's original license.
-
----
-
-## 7. Citation
-
-```bibtex
-@misc{baihu-v1-flash,
-  title  = {BaiHu-V1-Flash: An SSA (Sparse-attention + SubQ) Retrofit of Qwen3-0.6B-Base},
-  author = {NovaAI6868},
-  year   = {2026},
-  url    = {https://huggingface.co/NovaAI6868/BaiHu-V1-Flash}
-}
-```
+---
+library_name: transformers
+pipeline_tag: text-generation
+language:
+  - zh
+  - en
+license: other
+license_name: baihu-custom-license
+license_link: https://huggingface.co/ZichenAI/BaiHu-V1-Flash/blob/main/LICENSE.custom.md
+base_model: Qwen/Qwen3-0.6B-Base
+base_model_relation: finetune
+tags:
+  - sparse-attention
+  - subq
+  - ssa
+  - long-context
+  - supervised-fine-tuning
+  - transfer-learning
+  - commercial-license-required
+  - text-generation
+---
+
+# BaiHu-V1-Flash
+
+**BaiHu-V1-Flash** is an SSA (Sparse-attention + SubQ) retrofit of `Qwen/Qwen3-0.6B-Base`,
+fine-tuned on **1,294 bilingual multi-turn dialogues** whose purpose is **举一反三** — the
+model is shown one worked rule / format / method in the first turn, and later turns ask it to
+reuse that rule on a *new* case it has never seen.
+
+- Base model: `Qwen/Qwen3-0.6B-Base` (28 layers / 16 Q heads / 8 KV heads / head_dim 128 / 32K context / tied embeddings)
+- Parameters: **598.8M** (2.75M are SSA-only modules)
+- Architecture: SSA — every layer runs three attention paths (shared / local / sparse-SubQ), each with its own softmax, then summed
+- Training data: 1,294 synthetic dialogues (English 693 / Chinese 601), six transfer types
+- License: **free for personal use; a paid license is required for commercial use** (see [License](#license))
+
+> **Revision note.** This repository previously hosted the base-pretrain checkpoint of the
+> same name (5.0M tokens of continued pretraining, no instruction tuning, no P0 revision).
+> It has been **replaced** by the checkpoint described here. The two are not interchangeable:
+> this one uses the revised shared branch (§2) and is a supervised fine-tune.
+
+---
+
+## 1. Revisions in this release
+
+| | earlier release (now replaced) | **this release** |
+|---|---|---|
+| Shared path | one mean vector for the **entire prefix** | **one compressed vector per completed block** (P0 revision) |
+| Training | continued pretraining, 5.0M tokens of web text | supervised fine-tune, 1,294 transfer dialogues (assistant-token loss) |
+| Behaviour | base LM | follows a rule / format established earlier in the dialogue |
+
+Both revisions share the same SSA architecture, base weights and license.
+
+## 2. Architecture
+
+Every layer keeps the base model's MLP / RMSNorm weights and replaces full attention with
+an SSA layer built from three parallel paths:
+
+| Path | Role | Complexity |
+|---|---|---|
+| `shared` | every query attends over **one compressed summary vector per completed block** | `O(T·T/B)` |
+| `local` | dense causal attention over the most recent window | `O(T·w)` |
+| `sparse` (**SubQ**) | only 4 of 16 query heads produce block scores, shared across the head group; real attention is computed only for the selected top-k blocks | `O(T·k·B)` |
+
+### The P0 revision (this release)
+
+The previous revision compressed the **whole prefix into a single mean vector**. A query
+could only see one undifferentiated global average, the single-element softmax made that
+average inject at full weight, and during 5.0M tokens of continued pretraining the learned
+gate **shrank instead of growing** (0.0100 → 0.0129 → 0.0122) — i.e. the optimizer actively
+suppressed the branch.
+
+This release implements the design the project always documented: **one compressed key/value
+per completed 64-token block**, so queries do a genuine softmax over blocks and "which block
+matters" becomes learnable. Cost is unchanged (`O(T·T/B)`; block summaries are still computed
+once per layer).
+
+### Hyperparameters
+
+| Parameter | Value | Meaning |
+|---|---|---|
+| `ssa_block_size` | 64 | block size B |
+| `ssa_top_k` | 8 | blocks selected by the sparse path |
+| `ssa_local_blocks` | 2 | local window = 3 × 64 = 192 tokens |
+| `ssa_num_subq_heads` | 4 | SubQ heads, r = 16 / 4 = 4 |
+| `ssa_router_dim` / `ssa_compress_dim` | 128 / 128 | router subspace / summary width |
+
+## 3. Training
+
+**Data.** 1,294 synthetic multi-turn dialogues (en 693 / zh 601), each 6–12 messages: the
+first user turn gives a worked case or a rule, a later user turn introduces a new instance
+that can only be handled by reusing it, and the last turn pushes generalization one step
+further. Six `transfer_kind`s are covered: `rule_induction`, `analogy_transfer`,
+`format_transfer`, `counterfactual`, `cross_domain`, `teaching_loop`.
+
+Split (stratified by language × kind, seed 0): **train 1,165 sessions / 345.7K tokens**
+(assistant 188.8K), **val 129 sessions / 38.2K tokens** (assistant 20.5K).
+
+**Recipe.** Loss is computed on assistant tokens only; sequences are padded to a multiple of
+the SSA block size (64); rendering uses the tokenizer's own chat template
+(`apply_chat_template`, `add_generation_prompt=False`, i.e. the final assistant turn carries
+Qwen3's empty `<think></think>` block).
+
+| Item | Value |
+|---|---|
+| Epochs / steps | 3 / 219 |
+| Tokens seen | 1.15M |
+| Batch | 4 × grad-accum 4 (effective 16) |
+| Optimizer | AdamW, lr 2e-5, cosine to 10%, 20-step warmup, wd 0.1 |
+| Precision | bfloat16 |
+| Hardware / time | RTX 4090D 24GB — **6.9 min**, ~2950 tok/s, 19.6GB peak |
+| Val loss / ppl | 2.534 → 1.646 → 1.539 → **1.537 / 4.65** (converged after ~2 epochs) |
+
+To make the SSA-only modules trainable, the two output projections and the shared gate are
+re-initialised to a small non-zero scale (0.01) before training — starting from the exact
+donor checkpoint would freeze all 2.75M of them at zero gradient.
+
+## 4. Evaluation
+
+### 4.1 Held-out transfer loss (129 unseen sessions, 20,520 assistant tokens)
+
+Both models evaluated with the identical script, mask and batching.
+
+| Metric | Qwen3-0.6B-Base | **BaiHu-V1-Flash** | Change |
+|---|---|---|---|
+| loss | 2.5339 | **1.5376** | −39.3% |
+| **ppl** | **12.603** | **4.654** | **−63.1%** |
+| ppl (en) | 14.836 | 5.118 | −65.5% |
+| ppl (zh) | 10.542 | 4.187 | −60.3% |
+
+Per transfer type (ppl):
+
+| Kind | Base | BaiHu-V1-Flash |
+|---|---|---|
+| rule_induction | 7.27 | 2.69 |
+| analogy_transfer | 23.69 | 9.08 |
+| format_transfer | 18.58 | 5.72 |
+| counterfactual | 9.13 | 3.67 |
+| cross_domain | 16.40 | 6.51 |
+| teaching_loop | 8.26 | 2.95 |
+
+The improvement holds in **all 8 slices** (2 languages × 6 kinds + overall).
+
+### 4.2 Generation samples (12 prompts, one per language × kind, greedy)
+
+- **Base**: **7/12 outputs collapse into repeated symbols** (`⚇⚇⚇`, `ацион`,
+  `.TRAILING`), the rest are off-topic or hallucinated — the base model has never seen this
+  dialogue format.
+- **BaiHu-V1-Flash**: **0/12 degenerate**; it reuses the rule/format from earlier turns and
+  is correct on most samples (`B12A`, `100, 50, 25, 12.5, …`, format rewrites). Arithmetic
+  errors remain — a 0.6B model with 1.3K training dialogues still slips on multi-step math.
+
+### 4.3 Effect of the P0 revision
+
+Measured with the same data, recipe and seed, P0 vs the old single-vector shared branch:
+
+| | old shared branch | P0 (per-block) |
+|---|---|---|
+| Held-out loss | 1.5376 | 1.5376 |
+| Shared gate after SFT | 0.01001 | 0.01001 |
+
+Under this regime (short dialogues, full causal window in the local path) the shared branch
+carries almost no load, so the revision is quality-neutral and the gate still does not grow.
+P0's intended benefit is long-range: it should only show up once the local window is allowed
+to shrink / contexts are far longer than the 192-token window. That experiment has **not**
+been run yet — see Limitations.
+
+## 5. How to Run Inference
+
+### 5.1 This is a custom architecture
+
+`model_type: baihu_ssa` is **not** in the Transformers registry; a plain
+`AutoModelForCausalLM.from_pretrained(...)` fails. Register the config/model classes first
+(three lines below). Also note the model **does not** use `GenerationMixin` — call its own
+`generate()` (greedy or temperature/top-k; no beam search).
+
+The implementation lives in the project repository (not on the Hub). **This checkpoint
+requires the P0 revision** of `modeling_baihu_ssa.py` — an older copy loads without error but
+computes a different shared branch.
+
+### 5.2 Minimal working example
+
+```python
+import os
+import sys
+
+import torch
+from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer
+
+sys.path.insert(0, os.path.join("ssa_model", "src"))  # path to the cloned repo's src/
+
+from baihu_ssa.configuration_baihu_ssa import BaiHuSSAConfig
+from baihu_ssa.model_baihu_ssa import BaiHuSSAForCausalLM
+
+# ---- register the custom architecture (REQUIRED) ----
+AutoConfig.register("baihu_ssa", BaiHuSSAConfig, exist_ok=True)
+AutoModelForCausalLM.register(BaiHuSSAConfig, BaiHuSSAForCausalLM, exist_ok=True)
+
+REPO = "ZichenAI/BaiHu-V1-Flash"
+device = "cuda" if torch.cuda.is_available() else "cpu"
+dtype = torch.bfloat16 if device == "cuda" else torch.float32
+
+tok = AutoTokenizer.from_pretrained(REPO)
+model = AutoModelForCausalLM.from_pretrained(REPO, dtype=dtype).to(device).eval()
+
+messages = [
+    {"role": "user", "content": "这个数列的规律是什么？1、4、9、16、25"},
+    {"role": "assistant", "content": "相邻两项的差是 3、5、7、9，每次加 2，所以第 n 项是 n 的平方。"},
+    {"role": "user", "content": "按同样的规律，下一项是多少？36、49 之后呢？"},
+]
+ids = tok.apply_chat_template(messages, add_generation_prompt=True, tokenize=True,
+                              enable_thinking=False)
+ids = ids["input_ids"] if not isinstance(ids, list) else ids   # transformers 5.x returns a dict here
+with torch.no_grad():
+    out = model.generate(torch.tensor([ids], device=device), max_new_tokens=64,
+                         do_sample=False, eos_token_id=tok.convert_tokens_to_ids("<|im_end|>"))
+text = tok.decode(out[0, len(ids):], skip_special_tokens=True)
+if "<think>" in text:            # non-thinking mode emits an empty think block
+    text = text.split("</think>")[-1]
+print(text.strip())
+```
+
+Training used `add_generation_prompt=True, enable_thinking=False` (the generation prompt then
+ends with `<|im_start|>assistant\n<think>\n\n</think>\n\n`). Keep that format for best results.
+
+### 5.3 Troubleshooting
+
+| Error | Cause | Fix |
+|---|---|---|
+| `does not recognize this architecture` | custom architecture not registered | `AutoConfig.register` + `AutoModelForCausalLM.register` as in §5.2 |
+| output is repetitive garbage | wrong chat format (e.g. a raw prompt with no ChatML wrapper) | render with the tokenizer's chat template, `enable_thinking=False` |
+| `CUDA error: no kernel image is available` | PyTorch without kernels for an old GPU (Maxwell, sm_52) | pin torch 2.7.1+cu126 (2.8 dropped sm_50/sm_60) |
+
+## 6. Known Limitations
+
+1. **0.6B scale.** Multi-step arithmetic is still unreliable; the SFT teaches the *behaviour*
+   (reuse the earlier rule/format, answer directly), not new reasoning ability.
+2. **The shared branch is still barely used** (gate ≈ 0.0101 after SFT, same as the old
+   revision). Whether P0 pays off can only be decided in a long-context / window-shrink regime.
+3. **Trained on synthetic dialogues.** The 1,294 conversations are model-generated; style and
+   coverage are limited to the six transfer kinds and the topics they cover. Evaluation above
+   is on a held-out slice of the same distribution — not a general capability claim.
+4. **Custom architecture, no llama.cpp/GGUF support** (SSA attention is not implemented there).
+5. **Not an instruct model at large.** It is a 0.6B research model; expect terse answers and
+   occasional arithmetic slips.
+
+## 7. License
+
+Custom license (`LICENSE.custom.md` in this repository):
+
+- **Personal / non-commercial use: free**, including running, modifying and publicly
+  distributing derivative models under the same license.
+- **Commercial use requires a paid license** — contact **novaweb6868@outlook.com**.
+- Metadata: `license: other`, tag `commercial-license-required`.
+
+## 8. Citation
+
+```bibtex
+@misc{baihu_v1_flash,
+  title  = {BaiHu-V1-Flash: an SSA (sparse-attention + SubQ) retrofit of Qwen3-0.6B-Base, fine-tuned on bilingual transfer dialogues},
+  author = {ZichenAI},
+  year   = {2026},
+  url    = {https://huggingface.co/ZichenAI/BaiHu-V1-Flash}
+}
+```

config.json

@@ -1,79 +1,79 @@
-{
-  "architectures": [
-    "BaiHuSSAForCausalLM"
-  ],
-  "attention_bias": false,
-  "attention_dropout": 0.0,
-  "bos_token_id": 151643,
-  "dtype": "bfloat16",
-  "eos_token_id": 151643,
-  "head_dim": 128,
-  "hidden_act": "silu",
-  "hidden_size": 1024,
-  "initializer_range": 0.02,
-  "intermediate_size": 3072,
-  "layer_types": [
-    "full_attention",
-    "full_attention",
-    "full_attention",
-    "full_attention",
-    "full_attention",
-    "full_attention",
-    "full_attention",
-    "full_attention",
-    "full_attention",
-    "full_attention",
-    "full_attention",
-    "full_attention",
-    "full_attention",
-    "full_attention",
-    "full_attention",
-    "full_attention",
-    "full_attention",
-    "full_attention",
-    "full_attention",
-    "full_attention",
-    "full_attention",
-    "full_attention",
-    "full_attention",
-    "full_attention",
-    "full_attention",
-    "full_attention",
-    "full_attention",
-    "full_attention"
-  ],
-  "max_position_embeddings": 32768,
-  "max_window_layers": 28,
-  "model_type": "baihu_ssa",
-  "num_attention_heads": 16,
-  "num_hidden_layers": 28,
-  "num_key_value_heads": 8,
-  "pad_token_id": null,
-  "rms_norm_eps": 1e-06,
-  "rope_parameters": {
-    "rope_theta": 1000000,
-    "rope_type": "default"
-  },
-  "sliding_window": null,
-  "ssa_block_size": 64,
-  "ssa_component_init": 0.01,
-  "ssa_component_seed": 1234,
-  "ssa_compress_dim": 128,
-  "ssa_donor_model": "Qwen3-0.6B-Base",
-  "ssa_donor_revision": null,
-  "ssa_force_full_window": false,
-  "ssa_local_blocks": 2,
-  "ssa_num_subq_heads": 4,
-  "ssa_router_bias_scale": 0.1,
-  "ssa_router_dim": 128,
-  "ssa_router_init_scale": 0.01,
-  "ssa_shared_init_scale": 0.01,
-  "ssa_shared_kv": true,
-  "ssa_top_k": 8,
-  "ssa_use_router_key": true,
-  "tie_word_embeddings": true,
-  "transformers_version": "5.17.0",
-  "use_cache": true,
-  "use_sliding_window": false,
-  "vocab_size": 151936
-}
+{
+  "architectures": [
+    "BaiHuSSAForCausalLM"
+  ],
+  "attention_bias": false,
+  "attention_dropout": 0.0,
+  "bos_token_id": 151643,
+  "dtype": "bfloat16",
+  "eos_token_id": 151643,
+  "head_dim": 128,
+  "hidden_act": "silu",
+  "hidden_size": 1024,
+  "initializer_range": 0.02,
+  "intermediate_size": 3072,
+  "layer_types": [
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention"
+  ],
+  "max_position_embeddings": 32768,
+  "max_window_layers": 28,
+  "model_type": "baihu_ssa",
+  "num_attention_heads": 16,
+  "num_hidden_layers": 28,
+  "num_key_value_heads": 8,
+  "pad_token_id": null,
+  "rms_norm_eps": 1e-06,
+  "rope_parameters": {
+    "rope_theta": 1000000,
+    "rope_type": "default"
+  },
+  "sliding_window": null,
+  "ssa_block_size": 64,
+  "ssa_component_init": 0.01,
+  "ssa_component_seed": 1234,
+  "ssa_compress_dim": 128,
+  "ssa_donor_model": "Qwen3-0.6B-Base",
+  "ssa_donor_revision": null,
+  "ssa_force_full_window": true,
+  "ssa_local_blocks": 2,
+  "ssa_num_subq_heads": 4,
+  "ssa_router_bias_scale": 0.1,
+  "ssa_router_dim": 128,
+  "ssa_router_init_scale": 0.01,
+  "ssa_shared_init_scale": 0.01,
+  "ssa_shared_kv": true,
+  "ssa_top_k": 8,
+  "ssa_use_router_key": true,
+  "tie_word_embeddings": true,
+  "transformers_version": "5.17.0",
+  "use_cache": true,
+  "use_sliding_window": false,
+  "vocab_size": 151936
+}

model.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:40f5a9f413e8ce888f0e44e8bf3c18acd9d711cf2d0c76482096d43ec473c4c6
+oid sha256:e869217a7b844cf394980908f70312b97634fcd8f3fac5f264967e65da323fae
 size 2395267832