QAT_to_Llama.cpp_GGUF_TQ2_0_JirackUltra_14b.md · JiRackUltra_14b

1

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

2

3

Model: class `JiRackTransformer` in `JiRackTernaryUltra_14b.py`

4

(standard dense transformer, `model_type: qwen2`; base:

5

`CMSManhattan/JiRackUltra_14b`).

6

7

> **Base repo status:** verify `CMSManhattan/JiRackUltra_14b` exists on HF

8

> before relying on the HF load path. Until it's published, load base

9

> weights from a local `.pt` via `BASE_CHECKPOINT` in the training script;

10

> once the repo goes live, the HF path works without code changes.

11

>

12

> The code is structurally identical to the 1B/7B/32B files — same

13

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

14

> `export_ternary_state_dict`, the same 6 smoke tests. The pipeline carries

15

> over one-to-one; only the config constants change — **and at 14B two of

16

> them flip** (θ and eps, see below).

17

18

## 0. Config constants [DS14-1..3] — the source of truth

19

20

```python

21

VOCAB_SIZE = 152064 # same as 7B/32B; NOT 151936 like the 1B

22

HIDDEN_SIZE = 5120

23

INTERMEDIATE_SIZE = 13824 # NOTE: smaller than 32B's 27648 despite same hidden

24

NUM_LAYERS = 48

25

NUM_HEADS = 40

26

NUM_KV_HEADS = 8 # GQA 5:1

27

HEAD_DIM = 128 # 40 * 128 = 5120; kv dim = 1024

28

ROPE_THETA = 1000000.0 # [DS14-2] — NOT 10000 like the 1B/7B!

29

RMS_EPS = 1e-5 # [DS14-3] — NOT 1e-6 like the 1B/7B!

30

MAX_SEQ_LEN = 4096 # checkpoint nominally supports 131072

31

ATTN_QKV_BIAS = True # [DS14-5] bias on q/k/v only

32

tie_word_embeddings = False # separate lm_head.weight

33

```

34

35

**The θ/eps flip — the #1 trap at this size:** 14B sits on the *other*

36

branch of the family. Reusing a 1B/7B config here (θ=10,000, eps=1e-6)

37

does NOT crash — shapes are unrelated so you'd catch a 1B/7B config by

38

shape mismatch anyway, but a hand-copied θ or eps *value* slips through

39

silently: generation quality degrades (wrong rotary frequencies) and every

40

norm shifts slightly (wrong eps). Same in reverse: never copy 14B's θ/eps

41

down to the 1B/7B.

42

43

**Hidden-size collision with the 32B:** hidden 5120 and GQA 40/8 are

44

identical to the 32B — only `NUM_LAYERS` (48 vs 64) and

45

`INTERMEDIATE_SIZE` (13824 vs 27648) differ. A per-layer tensor from a 32B

46

export has the same attention shapes as a 14B one; only the FFN shapes and

47

the layer count distinguish the two models.

48

49

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

50

— the JiRack tokenizer for the family (robotics, routing, FIM and tool-call

51

special tokens). Its vocab (151,779) fits inside the 14B's padded matrix

52

(152,064):

53

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

54

would *shrink* the matrix and corrupt the checkpoint. The extra special

55

tokens land in already-existing padded slots; no resize is needed.

56

- Keep `VOCAB_SIZE = 152064` in the conversion config and GGUF metadata

57

(embedding size), regardless of the tokenizer's 151,779 used entries.

58

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

59

includes them — reserved slots, not yet functional semantics.

60

61

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

62

63

**Ternarized — every `BitLinear`, by construction (7 per block × 48 blocks

64

= 336).** Scope is fixed by the class; no `apply_bitlinear_patch` step:

65

66

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

67

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

68

| `blocks.N.q_proj` | `model.layers.N.self_attn.q_proj` | yes ([DS14-5]) |

69

| `blocks.N.k_proj` | `model.layers.N.self_attn.k_proj` | yes ([DS14-5]) |

70

| `blocks.N.v_proj` | `model.layers.N.self_attn.v_proj` | yes ([DS14-5]) |

71

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

72

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

73

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

74

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

75

76

`load_hf_state_dict()` is HF→JiRack; **GGUF export needs the reverse**. No

77

RoPE permutation is needed (half-split matches HF/llama.cpp).

78

79

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

80

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

81

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

82

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

83

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

84

- `token_emb` (152064 × 5120) and `lm_head` — **untied**.

85

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

86

87

Embed+lm_head share of params: ~5% — good compression ratio, and PPL

88

degradation from ternarization is smaller than at 1B/7B (more redundancy).

89

90

## 2. Conversion pitfalls — CHECK ALL FIVE

91

92

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

93

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

94

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

95

= **579 tensors**.

96

2. **`vocab_size` = 152064.** Tokenizer files:

97

`CMSManhattan/JiRackPrecisionTokenizer` (151,779 used entries — fits, no

98

resize; see vocab-trap section).

99

3. **θ = 1,000,000 and eps = 1e-5** — the OPPOSITE of the 1B/7B values.

100

Config reuse across the θ/eps boundary silently degrades generation and

101

shifts every norm; only 14B↔32B reuse is safe for these two values (but

102

watch layers/FFN, see the hidden-size collision note).

103

4. **RoPE table is 4096 positions** (`MAX_SEQ_LEN`; checkpoint nominally

104

supports 131072). Advertise conservative context in GGUF metadata unless

105

you deliberately raise it and re-verify.

106

5. **Materialization path:** `export_ternary_state_dict()` emits

107

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

108

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

109

safetensors, convert to GGUF, `llama-quantize` to TQ2_0 — lossless

110

round-trip.

111

112

Baseline verification (before any ternarization):

113

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

114

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

115

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

116

bf16, λ=0.0).

117

118

## 3. QAT training (lambda warmup 0 → 1)

119

120

```python

121

import torch

122

from transformers import AutoModelForCausalLM, AutoTokenizer

123

from JiRackTernaryUltra_14b import JiRackTransformer, JiRackConfig

124

125

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

126

assert len(tokenizer) <= 152064, "tokenizer must fit the 14B padded embedding matrix"

127

# Do NOT call model.resize_token_embeddings() -- see the vocab-trap section.

128

129

model = JiRackTransformer(JiRackConfig(), use_checkpoint=True) # grad ckpt mandatory

130

hf = AutoModelForCausalLM.from_pretrained(

131

"CMSManhattan/JiRackUltra_14b", torch_dtype=torch.bfloat16) # bf16: ~30 GB host RAM (fp32 would need ~59 GB)

132

model.load_hf_state_dict(hf.state_dict(), strict=True)

133

model = model.to(torch.bfloat16)

134

```

135

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

136

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

137

stripped re-upload, not the official checkpoint.

138

139

**Hardware reality check:** ~30 GB bf16 weights + ~30 GB grads + Adafactor

140

factored state + activations. A single 96GB card fits this at BATCH_SIZE=1,

141

GRAD_ACCUM=10, sequences ≤ 1024, gradient checkpointing on — and not much

142

more. If OOM: shorten sequences first; freezing embed+lm_head (~5% of

143

params, kept full precision anyway) is the next lever.

144

145

Then ramp `lambda` 0.0 → 1.0 via `model.set_lambda(x)` (one call covers all

146

336 BitLinears), save the checkpoint (`lambda_` is persistent and

147

serializes on its own). Prototype the λ schedule on the 1B first — QAT

148

dynamics transfer, the cost doesn't.

149

150

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

151

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

152

them first. Won't break training, but consider starting λ near 1.0 (or a

153

short ramp) when continuing from an already-ternarized checkpoint; a full

154

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

155

156

## 4. Export back to HF format (lambda = 1.0)

157

158

```python

159

model.eval()

160

model.set_lambda(1.0)

161

hf_out = {}

162

hf_out["model.embed_tokens.weight"] = model.token_emb.weight.detach().clone()

163

hf_out["model.norm.weight"] = model.ln_f.weight.detach().clone()

164

hf_out["lm_head.weight"] = model.lm_head.weight.detach().clone()

165

166

for i, blk in enumerate(model.blocks):

167

p = f"model.layers.{i}"

168

hf_out[f"{p}.input_layernorm.weight"] = blk.norm1.weight.detach().clone()

169

hf_out[f"{p}.post_attention_layernorm.weight"] = blk.norm2.weight.detach().clone()

170

for name, lin in [("self_attn.q_proj", blk.q_proj),

171

("self_attn.k_proj", blk.k_proj),

172

("self_attn.v_proj", blk.v_proj),

173

("self_attn.o_proj", blk.out_proj),

174

("mlp.gate_proj", blk.ffn_w1),

175

("mlp.up_proj", blk.ffn_w3),

176

("mlp.down_proj", blk.ffn_w2)]:

177

w = lin.weight.float()

178

gamma = w.abs().mean().clamp(min=lin.eps)

179

w_ternary = torch.clamp(torch.round(w / gamma), -1, 1) * gamma

180

hf_out[f"{p}.{name}.weight"] = w_ternary.to(lin.weight.dtype)

181

if lin.bias is not None:

182

hf_out[f"{p}.{name}.bias"] = lin.bias.detach().clone()

183

184

torch.save(hf_out, "jirack_ultra14b_ternary_hf.bin")

185

# config.json: architectures=["Qwen2ForCausalLM"], model_type="qwen2",

186

# + constants from Section 0 (θ=1e6, eps=1e-5 — do not copy from 1B/7B!)

187

```

188

189

## 5. GGUF conversion

190

191

```bash

192

python convert_hf_to_gguf.py jirack_ultra14b_ternary_hf/ \

193

--outfile jirack_ultra14b_ternary_f16.gguf --outtype f16

194

195

./llama-quantize jirack_ultra14b_ternary_f16.gguf \

196

jirack_ultra14b_ternary_tq2_0.gguf TQ2_0

197

```

198

Plain Qwen2 — long-standing support, no selective `--tensor-type` needed,

199

the whole file quantizes in one pass.

200

201

## 6. Verification

202

203

- TQ2_0 GGUF perplexity vs f16 GGUF of the same model.

204

- The only expected source of residual error: activations are quantized

205

differently than in the original `BitLinear` (per-token int8 here vs

206

block-wise q8_K in llama.cpp) — not a bug, an unavoidable difference.

207

- An abnormal PPL spike under TQ2_0 means a pipeline bug — suspect QAT

208

warmup, un-stripped `lambda_`, a sibling config value leaking in

209

(**θ=10,000 or eps=1e-6 from the 1B/7B is the classic one here**), or the

210

wrong tokenizer.

211

212

## 7. Pre-release checklist

213

214

- [ ] Constants `VOCAB_SIZE=152064`, `ROPE_THETA=1000000`, `RMS_EPS=1e-5` set correctly

215

- [ ] Tokenizer `CMSManhattan/JiRackPrecisionTokenizer`, no resize

216

- [ ] Lambda warmup 0→1 complete

217

- [ ] Export back done, 579 clean tensors, 336 `lambda_` stripped

218

- [ ] `convert_hf_to_gguf.py` ran with zero missing/unexpected keys

219

- [ ] `llama-quantize` to TQ2_0 succeeded

220

- [ ] Perplexity compared against the f16 GGUF

221

222

---

223

224

## Cross-model trap table — 1B / 7B / 14B / 32B

225

226

227

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

228

| Base checkpoint | `CMSManhattan/JiRackUltra_1b` | `CMSManhattan/JiRackUltra_7b` (not published) | `CMSManhattan/JiRackUltra_14b` (verify) | JiRack Ultra 32B base |

229

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

230

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

231

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

232

| Layers | 28 | 28 | **48** | 64 |

233

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

234

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

235

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

236

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

237

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

238

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

239

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

240

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