QAT_to_Llama.cpp_GGUF_TQ2_0_JirackUltra_1b.md · JiRackUltra_1b

1

# Instructions for Ternarization and Export to GGUF/TQ2_0 — JiRack Ultra 1B

2

3

Model: `/mnt/nfs_share/JiRackUltra_1b/`, class `JiRackTransformer` in `JiRackTernaryUltra_1b.py`

4

(standard dense transformer, `model_type: qwen2`; base per [DS-4]:

5

[`CMSManhattan/JiRackUltra_1b`](https://huggingface.co/CMSManhattan/JiRackUltra_1b)).

6

7

> **Naming note:** the repo says "1b" but the config is a **~1.5B-class**

8

> architecture (~1.78B params, [DS1.5-1] values). Treat "1b" as the product

9

> label; the checkpoint to load is `CMSManhattan/JiRackUltra_1b`.

10

>

11

> The code is structurally identical to the 7B/14B/32B files — same `BitLinear`,

12

> lambda warmup, RoPE half-split, `load_hf_state_dict` / `export_ternary_state_dict`,

13

> same 6 smoke tests. Two things are unique to this size: a **different vocab**

14

> (151936, not the siblings' 152064) and a **corrected tie-embeddings fact**

15

> (see below).

16

17

## Config constants [DS1.5-1] — the source of truth

18

19

These match the published `config.json` in `CMSManhattan/JiRackUltra_1b` exactly

20

(`Qwen2ForCausalLM`, `tie_word_embeddings: false`, `bos 151646`, `eos 151643`).

21

22

| Constant | 1B value | 7B | 14B | 32B |

23

|---|---|---|---|---|

24

| `VOCAB_SIZE` | **151936** ⚠ | 152064 | 152064 | 152064 |

25

| `HIDDEN_SIZE` | 1536 | 3584 | 5120 | 5120 |

26

| `INTERMEDIATE_SIZE` | 8960 | 18944 | 13824 | 27648 |

27

| `NUM_LAYERS` | 28 (same as 7B!) | 28 | 48 | 64 |

28

| `NUM_HEADS` / `NUM_KV_HEADS` | 12 / 2 (6:1) | 28 / 4 | 40 / 8 | 40 / 8 |

29

| `HEAD_DIM` | 128 (q dim 1536, kv dim 256) | 128 | 128 | 128 |

30

| `ROPE_THETA` | 10,000 | 10,000 | 1,000,000 | 1,000,000 |

31

| `RMS_EPS` | 1e-6 | 1e-6 | 1e-5 | 1e-5 |

32

| `MAX_SEQ_LEN` | 4096 ([DS-6]: ckpt supports 131072) | 4096 | 4096 | 4096 |

33

34

35

**The vocab trap:** 151936 ≠ 152064. `token_emb` and `lm_head` are 151936 × 1536

36

here. If a sibling's conversion `config.json` (vocab 152064) is reused, you get a

37

shape mismatch at best, or a GGUF whose metadata disagrees with the tokenizer at

38

worst.

39

40

**Tokenizer to use:** [`CMSManhattan/JiRackPrecisionTokenizer`](https://huggingface.co/CMSManhattan/JiRackPrecisionTokenizer)

41

— the JiRack tokenizer for the JiRack Ultra family (qwen2-compatible; robotics,

42

routing, FIM and tool-call special tokens included). Its vocab (151,779) fits

43

inside the model's padded embedding matrix (151,936), so:

44

- **Do NOT call `model.resize_token_embeddings()`** — 151,779 < 151,936 would

45

*shrink* the matrix and corrupt the checkpoint. The extra special tokens land

46

in already-existing padded slots; no resize is needed for this model.

47

- Keep `VOCAB_SIZE = 151936` in the config and GGUF metadata (embedding size),

48

regardless of the tokenizer's 151,779 used entries.

49

- The JiRack special tokens have untrained embeddings until QAT/fine-tuning

50

includes them — they're reserved slots, not yet functional semantics.

51

52

**Tie-embeddings — corrected fact ([DS1.5-2]):** the older 7B/32B file comments

53

claimed this size ties embeddings. This file corrects that: the published

54

`CMSManhattan/JiRackUltra_1b` checkpoint **ships a real, untied `lm_head.weight`**

55

(`tie_word_embeddings: false` in its config.json); the fallback branch in

56

`load_hf_state_dict` exists only for third-party re-uploads that strip

57

`lm_head`. If the fallback triggers on the official checkpoint, something is

58

wrong with your shard merge — investigate, don't shrug.

59

60

## Status template (fill in as you go)

61

62

⏳ **To do:**

63

- HF → JiRack conversion via `load_hf_state_dict()` (per [DS-4]: load `CMSManhattan/JiRackUltra_1b` or merge its safetensors shards)

64

- Built-in smoke tests: `python JiRackTernaryUltra_1b.py` — all 6 pass

65

- GGUF conversion (pt → safetensors with HF names → GGUF, bf16)

66

- llama-cli smoke test ("What is the capital of France?" → "Paris")

67

- Baseline perplexity on wikitext-100k, c=2048, bf16, lambda=0.0

68

- QAT warmup lambda 0.0→1.0, materialization via `export_ternary_state_dict()`, TQ2_0 export, final perplexity

69

70

Checkpoint size expectation: ~3.5 GB at bf16 for ~1.78B params (~7 GB at fp32) —

71

matches the published `model.safetensors` (3.55 GB, bf16). At this size, fp32

72

loading is fine on any machine — no shard-by-shard gymnastics needed.

73

74

## Quantization mechanism (identical to siblings, [FIX-1..5])

75

76

`BitLinear` with per-module persistent `lambda_` buffer:

77

- **Weights:** per-tensor absmean gamma (fp32 stats, eps=1e-5 clamp, [FIX-5]);

78

`clamp(round(w/gamma), -1, 1) * gamma` → exactly `(-gamma, 0, +gamma)` →

79

lossless TQ2_0 round-trip.

80

- **Activations:** per-token absmax int8 fake-quant ([FIX-2],[FIX-3]) — QAT only.

81

- **STE via lambda:** `w + lam * (w_quant - w).detach()`; `lambda_` persistent

82

([FIX-4]) — survives resume, pollutes exports (Section 2).

83

- **Fast path:** eval with `lambda_ < 1e-6` → plain `F.linear`; the λ=0 baseline

84

PPL genuinely measures the unquantized model.

85

86

---

87

88

## 0. MTP Block — Not applicable

89

90

Standard `qwen2` attention; no MTP/Medusa draft layer exists or is needed.

91

92

---

93

94

## 1. What gets ternarized and what doesn't

95

96

**Ternarized — every `BitLinear`, by construction (7 per block × 28 blocks = 196).**

97

No `apply_bitlinear_patch` step; scope is fixed by the class:

98

99

| JiRack name | HF name (for GGUF export mapping) | Bias |

100

|---|---|---|

101

| `blocks.N.q_proj` | `model.layers.N.self_attn.q_proj` | yes ([DS-2]) |

102

| `blocks.N.k_proj` | `model.layers.N.self_attn.k_proj` | yes ([DS-2]) |

103

| `blocks.N.v_proj` | `model.layers.N.self_attn.v_proj` | yes ([DS-2]) |

104

| `blocks.N.out_proj` | `model.layers.N.self_attn.o_proj` | no |

105

| `blocks.N.ffn_w1` (gate) | `model.layers.N.mlp.gate_proj` | no |

106

| `blocks.N.ffn_w3` (up) | `model.layers.N.mlp.up_proj` | no |

107

| `blocks.N.ffn_w2` (down) | `model.layers.N.mlp.down_proj` | no |

108

109

Same renames as all siblings: `out_proj` ↔ `o_proj`, `ffn_w1/w3/w2` ↔

110

`gate/up/down`. `load_hf_state_dict()` is HF→JiRack; **GGUF export needs the

111

reverse**. [DS-4]: no RoPE permutation needed (half-split matches HF/llama.cpp).

112

113

**NOT ternarized — stays in full precision:**

114

- **QKV biases** ([DS-2]/[DS-5]): weight-only ternarization; preserved by

115

`export_ternary_state_dict()` ([FIX-9]). Verify `blk.*.attn_q.bias` /

116

`attn_k.bias` / `attn_v.bias` land in the GGUF as f32.

117

- All `RMSNorm` weights: `blocks.N.norm1`, `blocks.N.norm2`, `ln_f` — eps = **1e-6**.

118

- `token_emb` (151936 × 1536) and `lm_head` — **untied** ([DS1.5-2], verified).

119

- RoPE: `freqs_cos/sin` are non-persistent buffers — never exported.

120

121

**A size-specific caution on quality:** at this scale, `embed_tokens` + `lm_head`

122

(2 × 151936 × 1536 ≈ 0.47B params) are ~26% of the model — vs ~7% at 7B. Keeping

123

them full precision is therefore both more important for quality *and* a bigger

124

share of the final file size; don't expect the same compression ratio as the

125

larger siblings, and expect ternarization PPL degradation to be relatively larger

126

at this scale (small models have less redundancy to absorb quantization).

127

128

## 2. Conversion pitfalls — CHECK ALL FIVE

129

130

1. **`lambda_` buffers pollute the state_dict:** 196 persistent `*.lambda_`

131

scalars. Strip before safetensors/GGUF. Clean tensor count: 28 × 12

132

(7 weights + 3 biases + 2 norms) + `token_emb` + `ln_f` + `lm_head`

133

= **339 tensors** — identical to the 7B count! Layer/tensor counts cannot

134

distinguish this model from the 7B; only shapes can (hidden 1536 vs 3584,

135

vocab 151936 vs 152064).

136

2. **`vocab_size` = 151936** — unique among the family. The conversion

137

`config.json` and the GGUF metadata must say 151936 (embedding size), never a

138

sibling's 152064. Tokenizer files: `CMSManhattan/JiRackPrecisionTokenizer`

139

(151,779 used entries — fits the padded matrix, no resize; see vocab-trap

140

section).

141

3. **θ = 10,000 and eps = 1e-6** — same values as the 7B, so 1B↔7B config reuse

142

is safe *for these two values only*; a 14B/32B config (θ=1M, eps=1e-5)

143

silently degrades generation and shifts every norm.

144

4. **RoPE table is 4096 positions** (`MAX_SEQ_LEN`; [DS-6]: ckpt nominally

145

supports 131072). Advertise conservative context in GGUF metadata unless you

146

deliberately raise it and re-verify.

147

5. **Materialization path:** `export_ternary_state_dict()` emits `{name}.codes`

148

(int8 ∈ {-1,0,1}) + `{name}.gamma`, not dense weights. Reconstruct

149

`w = codes.to(dtype) * gamma`, rename to HF names, write safetensors, convert

150

to GGUF, `llama-quantize` to TQ2_0 — lossless round-trip.

151

152

Baseline verification (before any ternarization):

153

- `python JiRackTernaryUltra_1b.py` — all 6 smoke tests pass.

154

- GGUF bf16 smoke test in llama-cli with chat template.

155

- Record your own baseline PPL on wikitext-100k (11 × 2048-token chunks, bf16,

156

λ=0.0). Expect a noticeably higher absolute PPL than the big siblings — that's

157

the model size, not a bug.

158

159

**Note for the future:** an abnormal PPL spike under TQ2_0 means a pipeline bug —

160

suspect QAT warmup, un-stripped `lambda_`, a mis-pinned exception tensor, a

161

sibling config value leaking in (θ/eps/vocab), or the wrong tokenizer.

162

163

## 3. QAT training

164

165

1. Load the JiRack base weights and the JiRack tokenizer (per [DS-4]):

166

```python

167

from transformers import AutoModelForCausalLM, AutoTokenizer

168

from JiRackTernaryUltra_1b import JiRackTransformer, JiRackConfig

169

170

tokenizer = AutoTokenizer.from_pretrained("CMSManhattan/JiRackPrecisionTokenizer")

171

assert len(tokenizer) <= 151936, "tokenizer must fit the padded embedding matrix"

172

# NOTE: no model.resize_token_embeddings() -- see the vocab-trap section.

173

174

model = JiRackTransformer(JiRackConfig()) # grad ckpt optional at this size

175

hf = AutoModelForCausalLM.from_pretrained(

176

"CMSManhattan/JiRackUltra_1b", torch_dtype=torch.bfloat16)

177

model.load_hf_state_dict(hf.state_dict())

178

model = model.to(torch.bfloat16)

179

```

180

With `strict=True`, the official checkpoint must load with **zero** real

181

missing keys — if the tied-embedding fallback fires, your input is a stripped

182

re-upload, not the official `CMSManhattan/JiRackUltra_1b`.

183

184

**λ-start caveat:** if the published checkpoint is already post-QAT (weights

185

materialized in ternary form), a fresh 0→1 warmup partially "un-freezes" them

186

before re-clamping. That won't break training, but consider starting λ near

187

1.0 (or a short ramp) when continuing from an already-ternarized checkpoint;

188

a full 0→1 warmup is for a full-precision base.

189

2. Ramp lambda 0.0 → 1.0 with `model.set_lambda(x)` — one call covers all 196

190

BitLinears, no skip list. Continuity / STE gradients (weights *and* QKV

191

biases) / lambda serialization are covered by smoke tests #1/#3/#4.

192

This is also the size to prototype your λ schedule cheaply before spending

193

compute on 7B/14B/32B — the QAT dynamics transfer, the cost doesn't.

194

3. Save checkpoint (`lambda_` serializes with it).

195

4. Materialize: `export_ternary_state_dict()` → dense ternary weights →

196

safetensors (HF names) → GGUF → TQ2_0 (Section 2, item 5).

197

198

---

199

200

## Cross-model trap table — 1B / 7B / 14B / 32B (same code, different constants)

201

202

203

|---|---|---|---|---|

204

| Base checkpoint | `CMSManhattan/JiRackUltra_1b` | JiRack Ultra 7B base | JiRack Ultra 14B base | JiRack Ultra 32B base |

205

206

| `vocab_size` | **151936** | 152064 | 152064 | 152064 |

207

| `rope_theta` | 10,000 | 10,000 | 1,000,000 | 1,000,000 |

208

| `rms_norm_eps` | 1e-6 | 1e-6 | 1e-5 | 1e-5 |

209

| Layers | 28 | 28 | 48 | 64 |

210

| Hidden / FFN | 1536 / 8960 | 3584 / 18944 | 5120 / 13824 | 5120 / 27648 |

211

| GQA (q / kv) | 12 / 2 | 28 / 4 | 40 / 8 | 40 / 8 |

212

| BitLinear count | 196 | 196 | 336 | 448 |

213

| Clean tensor count | 339 | 339 | 579 | 771 |

214

| `lambda_` buffers to strip | 196 | 196 | 336 | 448 |

215

| bf16 checkpoint size | ~3.5 GB | ~15–16 GB | ~29–30 GB | ~65 GB |

216

| embed+lm_head share of params | ~26% | ~7% | ~5% | ~3% |

217

| File-specific quirks | "1b" naming vs ~1.5B-class config | — | — | — |

218

| Identical across all four | `BitLinear` math, λ warmup API, RoPE half-split, HF name mapping, QKV bias handling, `export_ternary_state_dict()`, smoke-test suite | ← | ← | ← |