Hugging Face's logo

Tensor-compression
/
TR_main

tensor-decomposition
svd
low-rank
model-compression
knowledge-distillation
Model card Files
xet
 Community
Rhett1 commited on
Commit
ae89081
·
verified ·
1 Parent(s): deb7d02

Add asvd_llm family: ASVD baselines @ matched budgets (heal 60M + unheal)

Browse files
Files changed (25) hide show
  1. README.md +5 -1
  2. asvd_llm/heal/llama2-13b/README.md +16 -0
  3. asvd_llm/heal/llama2-13b/config.json +26 -0
  4. asvd_llm/heal/llama2-13b/lowrank_linear.py +83 -0
  5. asvd_llm/heal/llama2-13b/provenance.json +199 -0
  6. asvd_llm/heal/llama2-13b/ranks.json +253 -0
  7. asvd_llm/heal/llama2-13b/state_dict_fp16.pt +3 -0
  8. asvd_llm/heal/llama2-7b/README.md +16 -0
  9. asvd_llm/heal/llama2-7b/config.json +25 -0
  10. asvd_llm/heal/llama2-7b/lowrank_linear.py +83 -0
  11. asvd_llm/heal/llama2-7b/provenance.json +199 -0
  12. asvd_llm/heal/llama2-7b/ranks.json +206 -0
  13. asvd_llm/heal/llama2-7b/state_dict_fp16.pt +3 -0
  14. asvd_llm/unheal/llama2-13b/README.md +11 -0
  15. asvd_llm/unheal/llama2-13b/config.json +26 -0
  16. asvd_llm/unheal/llama2-13b/lowrank_linear.py +83 -0
  17. asvd_llm/unheal/llama2-13b/provenance.json +15 -0
  18. asvd_llm/unheal/llama2-13b/ranks.json +253 -0
  19. asvd_llm/unheal/llama2-13b/state_dict_fp16.pt +3 -0
  20. asvd_llm/unheal/llama2-7b/README.md +11 -0
  21. asvd_llm/unheal/llama2-7b/config.json +25 -0
  22. asvd_llm/unheal/llama2-7b/lowrank_linear.py +83 -0
  23. asvd_llm/unheal/llama2-7b/provenance.json +15 -0
  24. asvd_llm/unheal/llama2-7b/ranks.json +206 -0
  25. asvd_llm/unheal/llama2-7b/state_dict_fp16.pt +3 -0
README.md CHANGED
@@ -7,7 +7,7 @@ tags: [tensor-decomposition, svd, low-rank, model-compression, knowledge-distill
7
  # TR-3 study models: main compressed models and matched-budget baselines
8
 
9
  Layout: **<method> / <stage> / <model>**. Method families are `tr3` (our 3-site tensor ring +
10
- sparse residual), `tt3` (3-site tensor train ablation), `svd_llm` and `svd_llm_v2`
11
  (matched-budget baselines). Stages are `heal`
12
  (KD-distilled, deployable) and `unheal` (raw post-compression, for measuring the healing delta).
13
 
@@ -31,6 +31,10 @@ sparse residual), `tt3` (3-site tensor train ablation), `svd_llm` and `svd_llm_v
31
  | svd_llm_v2 ‡ | heal | [`llama2-13b`](svd_llm2/heal/llama2-13b) | 7.740 | llama2 | MMLU 44.96 |
32
  | svd_llm_v2 ‡ | unheal | [`llama2-7b`](svd_llm2/unheal/llama2-7b) | 4.023 | llama2 | — |
33
  | svd_llm_v2 ‡ | unheal | [`llama2-13b`](svd_llm2/unheal/llama2-13b) | 7.740 | llama2 | — |
 
 
 
 
34
 
35
  \* full-benchmark lm-eval 0.4.12 result; SVD probe columns are limit-40 healing probes here —
36
  each dir's README/provenance carries the precise protocol.
 
7
  # TR-3 study models: main compressed models and matched-budget baselines
8
 
9
  Layout: **<method> / <stage> / <model>**. Method families are `tr3` (our 3-site tensor ring +
10
+ sparse residual), `tt3` (3-site tensor train ablation), `svd_llm`, `svd_llm_v2` and `asvd_llm`
11
  (matched-budget baselines). Stages are `heal`
12
  (KD-distilled, deployable) and `unheal` (raw post-compression, for measuring the healing delta).
13
 
 
31
  | svd_llm_v2 ‡ | heal | [`llama2-13b`](svd_llm2/heal/llama2-13b) | 7.740 | llama2 | MMLU 44.96 |
32
  | svd_llm_v2 ‡ | unheal | [`llama2-7b`](svd_llm2/unheal/llama2-7b) | 4.023 | llama2 | — |
33
  | svd_llm_v2 ‡ | unheal | [`llama2-13b`](svd_llm2/unheal/llama2-13b) | 7.740 | llama2 | — |
34
+ | asvd_llm | heal | [`llama2-7b`](asvd_llm/heal/llama2-7b) | 4.023 | llama2 | MMLU 35.09 |
35
+ | asvd_llm | heal | [`llama2-13b`](asvd_llm/heal/llama2-13b) | 7.736 | llama2 | MMLU 42.46 |
36
+ | asvd_llm | unheal | [`llama2-7b`](asvd_llm/unheal/llama2-7b) | 4.023 | llama2 | — |
37
+ | asvd_llm | unheal | [`llama2-13b`](asvd_llm/unheal/llama2-13b) | 7.736 | llama2 | — |
38
 
39
  \* full-benchmark lm-eval 0.4.12 result; SVD probe columns are limit-40 healing probes here —
40
  each dir's README/provenance carries the precise protocol.
asvd_llm/heal/llama2-13b/README.md ADDED
@@ -0,0 +1,16 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ ---
2
+ license: llama2
3
+ base_model: NousResearch/Llama-2-13b-chat-hf
4
+ tags: [svd, asvd, low-rank, model-compression, knowledge-distillation, llama2]
5
+ ---
6
+
7
+ # ASVD (arXiv:2312.05821) LLaMA-2-13B-chat, KD-healed (60M tokens)
8
+
9
+ - **HEALED checkpoint** — the UNHEALED sibling is at `../../unheal/llama2-13b/`. Same `ranks.json`,
10
+ different weights (this dir's SHA256 is in `provenance.json`).
11
+ - Compression: official ASVD (abs_mean activation-aware scaling, per-layer rank assigned by
12
+ ppl-sensitivity binary search, 256 wikitext2 calib samples, seed 233), global kept ratio
13
+ retuned to a factored budget of **7,736,414,208 (-0.050% vs 7,740,282,880 target)** (exact budget match to the TR-3 main row).
14
+ - Healing: pure forward-KL distillation from the frozen Llama-2-13b teacher, 60M tokens;
15
+ held-out KL **3.8606 -> 0.2780**, probe MMLU (5-shot, limit 40) **25.66% -> 42.46%**.
16
+ - Load with `lowrank_linear.py` in this dir (same recipe as the svd_llm dirs).
asvd_llm/heal/llama2-13b/config.json ADDED
@@ -0,0 +1,26 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "_name_or_path": null,
3
+ "architectures": [
4
+ "LlamaForCausalLM"
5
+ ],
6
+ "bos_token_id": 1,
7
+ "eos_token_id": 2,
8
+ "hidden_act": "silu",
9
+ "hidden_size": 5120,
10
+ "initializer_range": 0.02,
11
+ "intermediate_size": 13824,
12
+ "max_position_embeddings": 4096,
13
+ "model_type": "llama",
14
+ "num_attention_heads": 40,
15
+ "num_hidden_layers": 40,
16
+ "num_key_value_heads": 40,
17
+ "pad_token_id": 0,
18
+ "pretraining_tp": 1,
19
+ "rms_norm_eps": 1e-05,
20
+ "rope_scaling": null,
21
+ "tie_word_embeddings": false,
22
+ "torch_dtype": "float16",
23
+ "transformers_version": "4.31.0.dev0",
24
+ "use_cache": true,
25
+ "vocab_size": 32000
26
+ }
asvd_llm/heal/llama2-13b/lowrank_linear.py ADDED
@@ -0,0 +1,83 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """Minimal standalone loader for the SVD-baseline compressed LLaMA-2-7B checkpoints.
2
+
3
+ No dependency beyond torch + transformers. The state_dict stores every compressed
4
+ projection as two factors of a LowRankLinear (A: out x k, B: k x in); ranks.json
5
+ holds the per-layer integer rank of each projection. Untouched parameters
6
+ (embeddings, norms) keep their stock transformers names.
7
+
8
+ Usage:
9
+ import json, torch
10
+ from lowrank_linear import build
11
+ model = build(
12
+ base_id="NousResearch/Llama-2-7b-chat-hf",
13
+ ranks=json.load(open("ranks.json")),
14
+ state_dict_path="state_dict_fp16.pt",
15
+ device="cuda", dtype=torch.float16,
16
+ )
17
+ tokenizer = AutoTokenizer.from_pretrained("NousResearch/Llama-2-7b-chat-hf")
18
+ out = model.generate(tokenizer("2+2=", return_tensors="pt").to("cuda").input_ids,
19
+ max_new_tokens=16, do_sample=False)
20
+ print(tokenizer.decode(out[0]))
21
+ """
22
+
23
+ from __future__ import annotations
24
+
25
+ import torch
26
+ import torch.nn as nn
27
+ import torch.nn.functional as F
28
+
29
+
30
+ class LowRankLinear(nn.Module):
31
+ """y = A (B x). A: out x r, B: r x in. (LLaMA projections carry no bias.)"""
32
+
33
+ def __init__(self, A: torch.Tensor, B: torch.Tensor):
34
+ super().__init__()
35
+ self.A = nn.Parameter(A)
36
+ self.B = nn.Parameter(B)
37
+
38
+ def forward(self, x):
39
+ return F.linear(F.linear(x, self.B), self.A)
40
+
41
+ @torch.no_grad()
42
+ def materialize_(self) -> nn.Linear:
43
+ """Replace this module with the dense W = A @ B (inference only; never call
44
+ while training -- it cuts the gradient path)."""
45
+ W = (self.A @ self.B).to(self.A.dtype)
46
+ dense = nn.Linear(self.B.shape[1], self.A.shape[0], bias=False,
47
+ device=W.device, dtype=W.dtype)
48
+ dense.weight.data = W
49
+ return dense
50
+
51
+
52
+ def build(base_id: str, ranks: dict, state_dict_path: str,
53
+ device: str = "cuda", dtype=torch.float16):
54
+ from transformers import AutoModelForCausalLM
55
+
56
+ model = AutoModelForCausalLM.from_pretrained(base_id, torch_dtype=dtype)
57
+ search = {}
58
+ for name, mod in model.named_modules():
59
+ if isinstance(mod, nn.Linear):
60
+ search[name] = mod
61
+ for path, k in ranks.items():
62
+ out_f, in_f = search[path].weight.shape
63
+ model_mod = model
64
+ for p in path.split(".")[:-1]:
65
+ model_mod = getattr(model_mod, p)
66
+ setattr(model_mod, path.split(".")[-1],
67
+ LowRankLinear(torch.zeros(out_f, int(k), dtype=dtype),
68
+ torch.zeros(int(k), in_f, dtype=dtype)))
69
+ sd = torch.load(state_dict_path, map_location="cpu", weights_only=True)
70
+ model.load_state_dict(sd, strict=True)
71
+ del sd
72
+ return model.to(device)
73
+
74
+
75
+ def materialize_all(model):
76
+ """Fuse every LowRankLinear back into a dense nn.Linear in place."""
77
+ for name, mod in list(model.named_modules()):
78
+ if isinstance(mod, LowRankLinear):
79
+ parent = model
80
+ for p in name.split(".")[:-1]:
81
+ parent = getattr(parent, p)
82
+ setattr(parent, name.split(".")[-1], mod.materialize_())
83
+ return model
asvd_llm/heal/llama2-13b/provenance.json ADDED
@@ -0,0 +1,199 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "method": "KD healing: alpha*T^2*KL(teacher||student) + (1-alpha)*CE",
3
+ "student_init": "/data1/rkp00022/zixuan/projects/TR3_main/checkpoints/asvd_llama2_13b_chat_r05947/tr3fmt/state_dict_fp16.pt",
4
+ "teacher": "NousResearch/Llama-2-13b-chat-hf",
5
+ "alpha": 1.0,
6
+ "temperature": 1.0,
7
+ "lr": 2e-05,
8
+ "tokens_seen": 59998208,
9
+ "epochs": 0.228,
10
+ "global_batch_tokens": 32768,
11
+ "dataset_mixture_pre_exclusion": {
12
+ "alpaca": 50000,
13
+ "openhermes": 200000,
14
+ "ultrachat": 100000
15
+ },
16
+ "dataset_revisions": {
17
+ "tatsu-lab/alpaca": "dce01c9b08f87459cf36a430d809084718273017",
18
+ "teknium/OpenHermes-2.5": "b82037821055c377bed0d495e72e46de3bc72e84",
19
+ "HuggingFaceH4/ultrachat_200k": "8049631c405ae6576f93f445c6b8166f76f5505a"
20
+ },
21
+ "canonical_pool_excluded_texts": 10000,
22
+ "train_texts_after_exclusion": 356000,
23
+ "chat_texts": 340000,
24
+ "mmlu_aux_texts_kept": 16000,
25
+ "mmlu_aux_split": "cais/mmlu auxiliary_train (official TRAIN split, disjoint from the evaluated test split; audit verified 0/99842 exact overlap)",
26
+ "mmlu_aux_stats": {
27
+ "requested_groups": 16000,
28
+ "rendered_groups": 16000,
29
+ "shots": 5,
30
+ "rows_total": 99842,
31
+ "rows_dropped_by_decontam": 71,
32
+ "rows_clean": 99771,
33
+ "revision": "c30699e8356da336a370243923dbaf21066bb9fe",
34
+ "format": "MMLU-shaped (NOT byte-identical to lm-eval; auxiliary_train has no subject)"
35
+ },
36
+ "gold_label_supervision": false,
37
+ "attribution_note": "alpha=1.0 => pure KD, no gold-label CE, so any MMLU change is attributable to distillation on in-domain text. alpha<1 adds CE on text containing answers = labeled MCQ training, which is a DIFFERENT and weaker claim.",
38
+ "train_tokens_total": 263112704,
39
+ "decontamination": {
40
+ "removed": {
41
+ "alpaca": 11,
42
+ "openhermes": 120,
43
+ "ultrachat": 33
44
+ }
45
+ },
46
+ "start_kl": 3.8606380820274353,
47
+ "final_kl": 0.2779867332428694,
48
+ "start_mmlu": 0.2565789473684211,
49
+ "final_mmlu": 0.4245614035087719,
50
+ "history": [
51
+ {
52
+ "tokens_m": 5.01,
53
+ "step": 153,
54
+ "held_out_kl": 0.3917487971484661,
55
+ "mmlu": 0.37236842105263157
56
+ },
57
+ {
58
+ "tokens_m": 10.03,
59
+ "step": 306,
60
+ "held_out_kl": 0.34532562270760536,
61
+ "mmlu": 0.3881578947368421
62
+ },
63
+ {
64
+ "tokens_m": 15.01,
65
+ "step": 458,
66
+ "held_out_kl": 0.32420871779322624,
67
+ "mmlu": 0.4083333333333333
68
+ },
69
+ {
70
+ "tokens_m": 20.02,
71
+ "step": 611,
72
+ "held_out_kl": 0.30829225666821003,
73
+ "mmlu": 0.42280701754385963
74
+ },
75
+ {
76
+ "tokens_m": 25.0,
77
+ "step": 763,
78
+ "held_out_kl": 0.2985613960772753,
79
+ "mmlu": 0.4192982456140351
80
+ },
81
+ {
82
+ "tokens_m": 30.02,
83
+ "step": 916,
84
+ "held_out_kl": 0.2917508389800787,
85
+ "mmlu": 0.4236842105263158
86
+ },
87
+ {
88
+ "tokens_m": 35.03,
89
+ "step": 1069,
90
+ "held_out_kl": 0.28504760190844536,
91
+ "mmlu": 0.4267543859649123
92
+ },
93
+ {
94
+ "tokens_m": 40.01,
95
+ "step": 1221,
96
+ "held_out_kl": 0.28255406580865383,
97
+ "mmlu": 0.42280701754385963
98
+ },
99
+ {
100
+ "tokens_m": 45.02,
101
+ "step": 1374,
102
+ "held_out_kl": 0.2797098886221647,
103
+ "mmlu": 0.42543859649122806
104
+ },
105
+ {
106
+ "tokens_m": 50.0,
107
+ "step": 1526,
108
+ "held_out_kl": 0.2789408154785633,
109
+ "mmlu": 0.42412280701754385
110
+ },
111
+ {
112
+ "tokens_m": 55.02,
113
+ "step": 1679,
114
+ "held_out_kl": 0.278350081294775,
115
+ "mmlu": 0.425
116
+ },
117
+ {
118
+ "tokens_m": 60.0,
119
+ "step": 1831,
120
+ "held_out_kl": 0.2779867332428694,
121
+ "mmlu": 0.4245614035087719,
122
+ "final": true
123
+ }
124
+ ],
125
+ "teacher_mmlu_same_limit": 0.5442982456140351,
126
+ "checkpoint": "/data1/rkp00022/zixuan/projects/TR3_main/checkpoints/kdheal_asvd13b_r05947_h1/state_dict_fp16.pt",
127
+ "sha256": "5f19311e6529f6d2d21484fe9fa2af4a238fe3f1c84ac6f3b146020032b61077",
128
+ "seed": 1234,
129
+ "mmlu_limit_per_subtask": 40,
130
+ "teacher_revision": "d73f5fa9c4bc135502e04c27b39660747172d76b",
131
+ "student_init_sha256": "49b4e7679fb063510a8d64c114c2a532708fbad8f822b51baa2969d5f1e20211",
132
+ "ranks_file": "/data1/rkp00022/zixuan/projects/TR3_main/checkpoints/asvd_llama2_13b_chat_r05947/tr3fmt/ranks.json",
133
+ "ranks_sha256": "145f1c36080f3e7d53725ca5c35583d3f64396c0cfb2f0b5342af12829cc2627",
134
+ "input_identity_gate": {
135
+ "student_ckpt": {
136
+ "key": "tr3fmt",
137
+ "sha256": "49b4e7679fb063510a8d64c114c2a532708fbad8f822b51baa2969d5f1e20211",
138
+ "status": "unregistered_allowed"
139
+ },
140
+ "ranks": {
141
+ "key": "tr3fmt",
142
+ "sha256": "145f1c36080f3e7d53725ca5c35583d3f64396c0cfb2f0b5342af12829cc2627",
143
+ "status": "unregistered_allowed"
144
+ }
145
+ },
146
+ "train_seed": 1234,
147
+ "mmlu_seed": 1234,
148
+ "dynamic_batch": false,
149
+ "dbl": null,
150
+ "snapshot_milestones": true,
151
+ "snapshots": [
152
+ "(PosixPath('/data1/rkp00022/zixuan/projects/TR3_main/checkpoints/kdheal_asvd13b_r05947_h1_snap15m'), 15.007744)",
153
+ "(PosixPath('/data1/rkp00022/zixuan/projects/TR3_main/checkpoints/kdheal_asvd13b_r05947_h1_snap30m'), 30.015488)",
154
+ "(PosixPath('/data1/rkp00022/zixuan/projects/TR3_main/checkpoints/kdheal_asvd13b_r05947_h1_snap45m'), 45.023232)",
155
+ "(PosixPath('/data1/rkp00022/zixuan/projects/TR3_main/checkpoints/kdheal_asvd13b_r05947_h1_snap50m'), 50.003968)",
156
+ "(PosixPath('/data1/rkp00022/zixuan/projects/TR3_main/checkpoints/kdheal_asvd13b_r05947_h1_snap55m'), 55.017472)"
157
+ ],
158
+ "stream_skip_tokens_m": 0.0,
159
+ "stream_blocks_consumed": [
160
+ 0,
161
+ 58592
162
+ ],
163
+ "stream_note": "the block order is a seeded permutation consumed from index 0; stages sharing a --seed with stream_skip=0 re-train on the SAME blocks. [start, end] block indices are recorded so any continuation can be checked for overlap with what earlier stages consumed.",
164
+ "student_master_dtype": "bf16",
165
+ "fake_quant": null,
166
+ "rebalance_report": null,
167
+ "student_param_count": 7736414208,
168
+ "micro": 8,
169
+ "grad_accum": 4,
170
+ "seq_len": 1024,
171
+ "held_out_source": "canonical: mixture 2000/4000/4000, seed 1234, seq 1024, blocks 256:264 (identical to the compression gates)",
172
+ "decontam_stats": {
173
+ "mmlu": {
174
+ "items": 14042,
175
+ "new_ngrams": 764261,
176
+ "revision": "c30699e8356da336a370243923dbaf21066bb9fe"
177
+ },
178
+ "hellaswag": {
179
+ "items": 10042,
180
+ "new_ngrams": 1334302,
181
+ "revision": "218ec52e09a7e7462a5400043bb9a69a41d06b76"
182
+ },
183
+ "boolq": {
184
+ "items": 3270,
185
+ "new_ngrams": 276806,
186
+ "revision": "3de24cf8022e94f4ee4b9d55a6f539891524d646"
187
+ },
188
+ "triviaqa": {
189
+ "items": 17944,
190
+ "new_ngrams": 32608,
191
+ "revision": "0f7faf33a3908546c6fd5b73a660e0f8ff173c2f"
192
+ },
193
+ "gsm8k": {
194
+ "items": 1319,
195
+ "new_ngrams": 46281,
196
+ "revision": "740312add88f781978c0658806c59bc2815b9866"
197
+ }
198
+ }
199
+ }
asvd_llm/heal/llama2-13b/ranks.json ADDED
@@ -0,0 +1,253 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "model.layers.0.mlp.down_proj": 1494,
3
+ "model.layers.0.mlp.gate_proj": 1494,
4
+ "model.layers.0.mlp.up_proj": 1494,
5
+ "model.layers.0.self_attn.k_proj": 1024,
6
+ "model.layers.0.self_attn.o_proj": 1024,
7
+ "model.layers.0.self_attn.q_proj": 1024,
8
+ "model.layers.0.self_attn.v_proj": 1024,
9
+ "model.layers.1.mlp.down_proj": 1494,
10
+ "model.layers.1.mlp.gate_proj": 1494,
11
+ "model.layers.1.mlp.up_proj": 1494,
12
+ "model.layers.1.self_attn.k_proj": 1024,
13
+ "model.layers.1.self_attn.o_proj": 1024,
14
+ "model.layers.1.self_attn.q_proj": 1024,
15
+ "model.layers.1.self_attn.v_proj": 1024,
16
+ "model.layers.10.mlp.gate_proj": 2615,
17
+ "model.layers.10.mlp.up_proj": 2988,
18
+ "model.layers.10.self_attn.k_proj": 1024,
19
+ "model.layers.10.self_attn.o_proj": 1024,
20
+ "model.layers.10.self_attn.q_proj": 1024,
21
+ "model.layers.10.self_attn.v_proj": 1024,
22
+ "model.layers.11.mlp.gate_proj": 3362,
23
+ "model.layers.11.mlp.up_proj": 2241,
24
+ "model.layers.11.self_attn.k_proj": 1024,
25
+ "model.layers.11.self_attn.o_proj": 1024,
26
+ "model.layers.11.self_attn.q_proj": 1024,
27
+ "model.layers.11.self_attn.v_proj": 1024,
28
+ "model.layers.12.mlp.gate_proj": 2988,
29
+ "model.layers.12.mlp.up_proj": 2241,
30
+ "model.layers.12.self_attn.k_proj": 1024,
31
+ "model.layers.12.self_attn.o_proj": 1024,
32
+ "model.layers.12.self_attn.q_proj": 1024,
33
+ "model.layers.12.self_attn.v_proj": 1024,
34
+ "model.layers.13.mlp.gate_proj": 1494,
35
+ "model.layers.13.mlp.up_proj": 1494,
36
+ "model.layers.13.self_attn.k_proj": 1024,
37
+ "model.layers.13.self_attn.o_proj": 1024,
38
+ "model.layers.13.self_attn.q_proj": 1024,
39
+ "model.layers.13.self_attn.v_proj": 1024,
40
+ "model.layers.14.mlp.gate_proj": 1494,
41
+ "model.layers.14.self_attn.k_proj": 1536,
42
+ "model.layers.14.self_attn.o_proj": 1024,
43
+ "model.layers.14.self_attn.q_proj": 1024,
44
+ "model.layers.14.self_attn.v_proj": 1024,
45
+ "model.layers.2.mlp.down_proj": 1868,
46
+ "model.layers.2.mlp.gate_proj": 1494,
47
+ "model.layers.2.mlp.up_proj": 1868,
48
+ "model.layers.2.self_attn.k_proj": 1024,
49
+ "model.layers.2.self_attn.o_proj": 1024,
50
+ "model.layers.2.self_attn.q_proj": 1024,
51
+ "model.layers.2.self_attn.v_proj": 1024,
52
+ "model.layers.3.mlp.down_proj": 1494,
53
+ "model.layers.3.mlp.gate_proj": 1494,
54
+ "model.layers.3.mlp.up_proj": 1868,
55
+ "model.layers.3.self_attn.k_proj": 1024,
56
+ "model.layers.3.self_attn.o_proj": 1024,
57
+ "model.layers.3.self_attn.q_proj": 1024,
58
+ "model.layers.3.self_attn.v_proj": 1024,
59
+ "model.layers.4.mlp.down_proj": 1494,
60
+ "model.layers.4.mlp.gate_proj": 1494,
61
+ "model.layers.4.mlp.up_proj": 1494,
62
+ "model.layers.4.self_attn.k_proj": 1024,
63
+ "model.layers.4.self_attn.o_proj": 1280,
64
+ "model.layers.4.self_attn.q_proj": 1024,
65
+ "model.layers.4.self_attn.v_proj": 1024,
66
+ "model.layers.5.mlp.down_proj": 1494,
67
+ "model.layers.5.mlp.gate_proj": 1494,
68
+ "model.layers.5.mlp.up_proj": 1494,
69
+ "model.layers.5.self_attn.k_proj": 1024,
70
+ "model.layers.5.self_attn.o_proj": 1024,
71
+ "model.layers.5.self_attn.q_proj": 1024,
72
+ "model.layers.5.self_attn.v_proj": 1280,
73
+ "model.layers.6.mlp.down_proj": 1494,
74
+ "model.layers.6.mlp.gate_proj": 1494,
75
+ "model.layers.6.mlp.up_proj": 1868,
76
+ "model.layers.6.self_attn.k_proj": 1024,
77
+ "model.layers.6.self_attn.o_proj": 1024,
78
+ "model.layers.6.self_attn.q_proj": 1024,
79
+ "model.layers.6.self_attn.v_proj": 1280,
80
+ "model.layers.7.mlp.down_proj": 1494,
81
+ "model.layers.7.mlp.gate_proj": 1494,
82
+ "model.layers.7.mlp.up_proj": 1494,
83
+ "model.layers.7.self_attn.k_proj": 1024,
84
+ "model.layers.7.self_attn.o_proj": 1024,
85
+ "model.layers.7.self_attn.q_proj": 1024,
86
+ "model.layers.7.self_attn.v_proj": 1024,
87
+ "model.layers.8.mlp.down_proj": 2241,
88
+ "model.layers.8.mlp.gate_proj": 1868,
89
+ "model.layers.8.mlp.up_proj": 2615,
90
+ "model.layers.8.self_attn.k_proj": 1024,
91
+ "model.layers.8.self_attn.o_proj": 1024,
92
+ "model.layers.8.self_attn.q_proj": 1024,
93
+ "model.layers.8.self_attn.v_proj": 1024,
94
+ "model.layers.9.mlp.gate_proj": 1494,
95
+ "model.layers.9.mlp.up_proj": 2988,
96
+ "model.layers.9.self_attn.k_proj": 1024,
97
+ "model.layers.9.self_attn.o_proj": 1280,
98
+ "model.layers.9.self_attn.q_proj": 1024,
99
+ "model.layers.9.self_attn.v_proj": 1536,
100
+ "model.layers.14.mlp.up_proj": 1494,
101
+ "model.layers.15.mlp.gate_proj": 1494,
102
+ "model.layers.15.mlp.up_proj": 1494,
103
+ "model.layers.15.self_attn.k_proj": 1280,
104
+ "model.layers.15.self_attn.o_proj": 1024,
105
+ "model.layers.15.self_attn.q_proj": 1024,
106
+ "model.layers.15.self_attn.v_proj": 1024,
107
+ "model.layers.16.mlp.gate_proj": 1494,
108
+ "model.layers.16.mlp.up_proj": 1494,
109
+ "model.layers.16.self_attn.k_proj": 1024,
110
+ "model.layers.16.self_attn.o_proj": 1024,
111
+ "model.layers.16.self_attn.q_proj": 1024,
112
+ "model.layers.16.self_attn.v_proj": 1024,
113
+ "model.layers.17.mlp.gate_proj": 1494,
114
+ "model.layers.17.mlp.up_proj": 1494,
115
+ "model.layers.17.self_attn.k_proj": 1024,
116
+ "model.layers.17.self_attn.o_proj": 1024,
117
+ "model.layers.17.self_attn.q_proj": 1024,
118
+ "model.layers.17.self_attn.v_proj": 1024,
119
+ "model.layers.18.mlp.gate_proj": 1494,
120
+ "model.layers.18.mlp.up_proj": 1494,
121
+ "model.layers.18.self_attn.k_proj": 1024,
122
+ "model.layers.18.self_attn.o_proj": 1024,
123
+ "model.layers.18.self_attn.q_proj": 1024,
124
+ "model.layers.18.self_attn.v_proj": 1792,
125
+ "model.layers.19.mlp.gate_proj": 1494,
126
+ "model.layers.19.mlp.up_proj": 1494,
127
+ "model.layers.19.self_attn.k_proj": 1280,
128
+ "model.layers.19.self_attn.o_proj": 1024,
129
+ "model.layers.19.self_attn.q_proj": 1024,
130
+ "model.layers.19.self_attn.v_proj": 1024,
131
+ "model.layers.20.mlp.gate_proj": 1494,
132
+ "model.layers.20.mlp.up_proj": 1494,
133
+ "model.layers.20.self_attn.k_proj": 1024,
134
+ "model.layers.20.self_attn.o_proj": 1024,
135
+ "model.layers.20.self_attn.q_proj": 1024,
136
+ "model.layers.20.self_attn.v_proj": 1536,
137
+ "model.layers.21.mlp.gate_proj": 1494,
138
+ "model.layers.21.mlp.up_proj": 2241,
139
+ "model.layers.21.self_attn.k_proj": 1024,
140
+ "model.layers.21.self_attn.o_proj": 1024,
141
+ "model.layers.21.self_attn.q_proj": 1024,
142
+ "model.layers.21.self_attn.v_proj": 1024,
143
+ "model.layers.22.mlp.gate_proj": 2615,
144
+ "model.layers.22.mlp.up_proj": 3362,
145
+ "model.layers.22.self_attn.k_proj": 1024,
146
+ "model.layers.22.self_attn.o_proj": 1024,
147
+ "model.layers.22.self_attn.q_proj": 1024,
148
+ "model.layers.22.self_attn.v_proj": 1024,
149
+ "model.layers.23.mlp.gate_proj": 2615,
150
+ "model.layers.23.mlp.up_proj": 2615,
151
+ "model.layers.23.self_attn.k_proj": 1280,
152
+ "model.layers.23.self_attn.o_proj": 1280,
153
+ "model.layers.23.self_attn.q_proj": 1536,
154
+ "model.layers.23.self_attn.v_proj": 1280,
155
+ "model.layers.24.mlp.down_proj": 3362,
156
+ "model.layers.24.mlp.gate_proj": 1868,
157
+ "model.layers.24.mlp.up_proj": 2241,
158
+ "model.layers.24.self_attn.k_proj": 1024,
159
+ "model.layers.24.self_attn.o_proj": 1024,
160
+ "model.layers.24.self_attn.q_proj": 1024,
161
+ "model.layers.24.self_attn.v_proj": 1792,
162
+ "model.layers.25.mlp.down_proj": 3362,
163
+ "model.layers.25.mlp.gate_proj": 2241,
164
+ "model.layers.25.mlp.up_proj": 2988,
165
+ "model.layers.25.self_attn.k_proj": 1024,
166
+ "model.layers.25.self_attn.o_proj": 1280,
167
+ "model.layers.25.self_attn.q_proj": 1024,
168
+ "model.layers.25.self_attn.v_proj": 1536,
169
+ "model.layers.26.mlp.down_proj": 2988,
170
+ "model.layers.26.mlp.gate_proj": 1868,
171
+ "model.layers.26.mlp.up_proj": 1868,
172
+ "model.layers.26.self_attn.k_proj": 1024,
173
+ "model.layers.26.self_attn.o_proj": 1024,
174
+ "model.layers.26.self_attn.q_proj": 1024,
175
+ "model.layers.26.self_attn.v_proj": 1280,
176
+ "model.layers.27.mlp.down_proj": 2988,
177
+ "model.layers.27.mlp.gate_proj": 1494,
178
+ "model.layers.27.mlp.up_proj": 1494,
179
+ "model.layers.27.self_attn.k_proj": 1024,
180
+ "model.layers.27.self_attn.o_proj": 1024,
181
+ "model.layers.27.self_attn.q_proj": 1024,
182
+ "model.layers.27.self_attn.v_proj": 1536,
183
+ "model.layers.28.mlp.down_proj": 2988,
184
+ "model.layers.28.mlp.gate_proj": 2615,
185
+ "model.layers.28.mlp.up_proj": 2615,
186
+ "model.layers.28.self_attn.k_proj": 1024,
187
+ "model.layers.28.self_attn.o_proj": 1024,
188
+ "model.layers.28.self_attn.q_proj": 1024,
189
+ "model.layers.28.self_attn.v_proj": 1024,
190
+ "model.layers.29.mlp.down_proj": 2988,
191
+ "model.layers.29.mlp.gate_proj": 2615,
192
+ "model.layers.29.mlp.up_proj": 2241,
193
+ "model.layers.29.self_attn.k_proj": 1024,
194
+ "model.layers.29.self_attn.o_proj": 1024,
195
+ "model.layers.29.self_attn.q_proj": 1024,
196
+ "model.layers.29.self_attn.v_proj": 1024,
197
+ "model.layers.30.mlp.down_proj": 1868,
198
+ "model.layers.30.mlp.gate_proj": 2988,
199
+ "model.layers.30.mlp.up_proj": 2615,
200
+ "model.layers.30.self_attn.k_proj": 1024,
201
+ "model.layers.30.self_attn.o_proj": 1024,
202
+ "model.layers.30.self_attn.q_proj": 1024,
203
+ "model.layers.30.self_attn.v_proj": 1024,
204
+ "model.layers.31.mlp.down_proj": 1868,
205
+ "model.layers.31.mlp.gate_proj": 2988,
206
+ "model.layers.31.mlp.up_proj": 2241,
207
+ "model.layers.31.self_attn.k_proj": 1024,
208
+ "model.layers.31.self_attn.o_proj": 1024,
209
+ "model.layers.31.self_attn.q_proj": 1024,
210
+ "model.layers.31.self_attn.v_proj": 1024,
211
+ "model.layers.32.mlp.down_proj": 2241,
212
+ "model.layers.32.mlp.gate_proj": 3362,
213
+ "model.layers.32.mlp.up_proj": 2988,
214
+ "model.layers.32.self_attn.k_proj": 1024,
215
+ "model.layers.32.self_attn.o_proj": 1024,
216
+ "model.layers.32.self_attn.q_proj": 1024,
217
+ "model.layers.32.self_attn.v_proj": 1024,
218
+ "model.layers.33.mlp.down_proj": 1494,
219
+ "model.layers.33.mlp.up_proj": 3362,
220
+ "model.layers.33.self_attn.k_proj": 1024,
221
+ "model.layers.33.self_attn.o_proj": 1024,
222
+ "model.layers.33.self_attn.q_proj": 1024,
223
+ "model.layers.33.self_attn.v_proj": 1536,
224
+ "model.layers.34.mlp.down_proj": 1868,
225
+ "model.layers.34.self_attn.k_proj": 1024,
226
+ "model.layers.34.self_attn.o_proj": 1024,
227
+ "model.layers.34.self_attn.q_proj": 1024,
228
+ "model.layers.34.self_attn.v_proj": 1024,
229
+ "model.layers.35.mlp.down_proj": 2988,
230
+ "model.layers.35.self_attn.k_proj": 1024,
231
+ "model.layers.35.self_attn.o_proj": 1024,
232
+ "model.layers.35.self_attn.q_proj": 1024,
233
+ "model.layers.35.self_attn.v_proj": 1536,
234
+ "model.layers.36.mlp.down_proj": 1868,
235
+ "model.layers.36.mlp.up_proj": 2615,
236
+ "model.layers.36.self_attn.k_proj": 1024,
237
+ "model.layers.36.self_attn.o_proj": 1280,
238
+ "model.layers.36.self_attn.q_proj": 1024,
239
+ "model.layers.37.mlp.down_proj": 2615,
240
+ "model.layers.37.self_attn.o_proj": 1024,
241
+ "model.layers.37.self_attn.q_proj": 1024,
242
+ "model.layers.37.self_attn.v_proj": 1536,
243
+ "model.layers.38.mlp.down_proj": 2988,
244
+ "model.layers.38.mlp.gate_proj": 3362,
245
+ "model.layers.38.self_attn.k_proj": 1024,
246
+ "model.layers.38.self_attn.o_proj": 1024,
247
+ "model.layers.38.self_attn.q_proj": 1024,
248
+ "model.layers.38.self_attn.v_proj": 1024,
249
+ "model.layers.39.self_attn.k_proj": 1024,
250
+ "model.layers.39.self_attn.o_proj": 1024,
251
+ "model.layers.39.self_attn.q_proj": 1024,
252
+ "model.layers.39.self_attn.v_proj": 1024
253
+ }
asvd_llm/heal/llama2-13b/state_dict_fp16.pt ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:5f19311e6529f6d2d21484fe9fa2af4a238fe3f1c84ac6f3b146020032b61077
3
+ size 15473033811
asvd_llm/heal/llama2-7b/README.md ADDED
@@ -0,0 +1,16 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ ---
2
+ license: llama2
3
+ base_model: NousResearch/Llama-2-7b-chat-hf
4
+ tags: [svd, asvd, low-rank, model-compression, knowledge-distillation, llama2]
5
+ ---
6
+
7
+ # ASVD (arXiv:2312.05821) LLaMA-2-7B-chat, KD-healed (60M tokens)
8
+
9
+ - **HEALED checkpoint** — the UNHEALED sibling is at `../../unheal/llama2-7b/`. Same `ranks.json`,
10
+ different weights (this dir's SHA256 is in `provenance.json`).
11
+ - Compression: official ASVD (abs_mean activation-aware scaling, per-layer rank assigned by
12
+ ppl-sensitivity binary search, 256 wikitext2 calib samples, seed 233), global kept ratio
13
+ retuned to a factored budget of **4,023,459,584 (+0.008% vs 4,023,119,424 target)** (exact budget match to the TR-3 main row).
14
+ - Healing: pure forward-KL distillation from the frozen Llama-2-7b teacher, 60M tokens;
15
+ held-out KL **7.3180 -> 0.3822**, probe MMLU (5-shot, limit 40) **23.86% -> 35.09%**.
16
+ - Load with `lowrank_linear.py` in this dir (same recipe as the svd_llm dirs).
asvd_llm/heal/llama2-7b/config.json ADDED
@@ -0,0 +1,25 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "architectures": [
3
+ "LlamaForCausalLM"
4
+ ],
5
+ "bos_token_id": 1,
6
+ "eos_token_id": 2,
7
+ "hidden_act": "silu",
8
+ "hidden_size": 4096,
9
+ "initializer_range": 0.02,
10
+ "intermediate_size": 11008,
11
+ "max_position_embeddings": 4096,
12
+ "model_type": "llama",
13
+ "num_attention_heads": 32,
14
+ "num_hidden_layers": 32,
15
+ "num_key_value_heads": 32,
16
+ "pad_token_id": 0,
17
+ "pretraining_tp": 1,
18
+ "rms_norm_eps": 1e-05,
19
+ "rope_scaling": null,
20
+ "tie_word_embeddings": false,
21
+ "torch_dtype": "float16",
22
+ "transformers_version": "4.31.0.dev0",
23
+ "use_cache": true,
24
+ "vocab_size": 32000
25
+ }
asvd_llm/heal/llama2-7b/lowrank_linear.py ADDED
@@ -0,0 +1,83 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """Minimal standalone loader for the SVD-baseline compressed LLaMA-2-7B checkpoints.
2
+
3
+ No dependency beyond torch + transformers. The state_dict stores every compressed
4
+ projection as two factors of a LowRankLinear (A: out x k, B: k x in); ranks.json
5
+ holds the per-layer integer rank of each projection. Untouched parameters
6
+ (embeddings, norms) keep their stock transformers names.
7
+
8
+ Usage:
9
+ import json, torch
10
+ from lowrank_linear import build
11
+ model = build(
12
+ base_id="NousResearch/Llama-2-7b-chat-hf",
13
+ ranks=json.load(open("ranks.json")),
14
+ state_dict_path="state_dict_fp16.pt",
15
+ device="cuda", dtype=torch.float16,
16
+ )
17
+ tokenizer = AutoTokenizer.from_pretrained("NousResearch/Llama-2-7b-chat-hf")
18
+ out = model.generate(tokenizer("2+2=", return_tensors="pt").to("cuda").input_ids,
19
+ max_new_tokens=16, do_sample=False)
20
+ print(tokenizer.decode(out[0]))
21
+ """
22
+
23
+ from __future__ import annotations
24
+
25
+ import torch
26
+ import torch.nn as nn
27
+ import torch.nn.functional as F
28
+
29
+
30
+ class LowRankLinear(nn.Module):
31
+ """y = A (B x). A: out x r, B: r x in. (LLaMA projections carry no bias.)"""
32
+
33
+ def __init__(self, A: torch.Tensor, B: torch.Tensor):
34
+ super().__init__()
35
+ self.A = nn.Parameter(A)
36
+ self.B = nn.Parameter(B)
37
+
38
+ def forward(self, x):
39
+ return F.linear(F.linear(x, self.B), self.A)
40
+
41
+ @torch.no_grad()
42
+ def materialize_(self) -> nn.Linear:
43
+ """Replace this module with the dense W = A @ B (inference only; never call
44
+ while training -- it cuts the gradient path)."""
45
+ W = (self.A @ self.B).to(self.A.dtype)
46
+ dense = nn.Linear(self.B.shape[1], self.A.shape[0], bias=False,
47
+ device=W.device, dtype=W.dtype)
48
+ dense.weight.data = W
49
+ return dense
50
+
51
+
52
+ def build(base_id: str, ranks: dict, state_dict_path: str,
53
+ device: str = "cuda", dtype=torch.float16):
54
+ from transformers import AutoModelForCausalLM
55
+
56
+ model = AutoModelForCausalLM.from_pretrained(base_id, torch_dtype=dtype)
57
+ search = {}
58
+ for name, mod in model.named_modules():
59
+ if isinstance(mod, nn.Linear):
60
+ search[name] = mod
61
+ for path, k in ranks.items():
62
+ out_f, in_f = search[path].weight.shape
63
+ model_mod = model
64
+ for p in path.split(".")[:-1]:
65
+ model_mod = getattr(model_mod, p)
66
+ setattr(model_mod, path.split(".")[-1],
67
+ LowRankLinear(torch.zeros(out_f, int(k), dtype=dtype),
68
+ torch.zeros(int(k), in_f, dtype=dtype)))
69
+ sd = torch.load(state_dict_path, map_location="cpu", weights_only=True)
70
+ model.load_state_dict(sd, strict=True)
71
+ del sd
72
+ return model.to(device)
73
+
74
+
75
+ def materialize_all(model):
76
+ """Fuse every LowRankLinear back into a dense nn.Linear in place."""
77
+ for name, mod in list(model.named_modules()):
78
+ if isinstance(mod, LowRankLinear):
79
+ parent = model
80
+ for p in name.split(".")[:-1]:
81
+ parent = getattr(parent, p)
82
+ setattr(parent, name.split(".")[-1], mod.materialize_())
83
+ return model
asvd_llm/heal/llama2-7b/provenance.json ADDED
@@ -0,0 +1,199 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "method": "KD healing: alpha*T^2*KL(teacher||student) + (1-alpha)*CE",
3
+ "student_init": "/data1/rkp00022/zixuan/projects/TR3_main/checkpoints/asvd_llama2_7b_chat_r0597/tr3fmt/state_dict_fp16.pt",
4
+ "teacher": "NousResearch/Llama-2-7b-chat-hf",
5
+ "alpha": 1.0,
6
+ "temperature": 1.0,
7
+ "lr": 2e-05,
8
+ "tokens_seen": 59998208,
9
+ "epochs": 0.228,
10
+ "global_batch_tokens": 32768,
11
+ "dataset_mixture_pre_exclusion": {
12
+ "alpaca": 50000,
13
+ "openhermes": 200000,
14
+ "ultrachat": 100000
15
+ },
16
+ "dataset_revisions": {
17
+ "tatsu-lab/alpaca": "dce01c9b08f87459cf36a430d809084718273017",
18
+ "teknium/OpenHermes-2.5": "b82037821055c377bed0d495e72e46de3bc72e84",
19
+ "HuggingFaceH4/ultrachat_200k": "8049631c405ae6576f93f445c6b8166f76f5505a"
20
+ },
21
+ "canonical_pool_excluded_texts": 10000,
22
+ "train_texts_after_exclusion": 356000,
23
+ "chat_texts": 340000,
24
+ "mmlu_aux_texts_kept": 16000,
25
+ "mmlu_aux_split": "cais/mmlu auxiliary_train (official TRAIN split, disjoint from the evaluated test split; audit verified 0/99842 exact overlap)",
26
+ "mmlu_aux_stats": {
27
+ "requested_groups": 16000,
28
+ "rendered_groups": 16000,
29
+ "shots": 5,
30
+ "rows_total": 99842,
31
+ "rows_dropped_by_decontam": 71,
32
+ "rows_clean": 99771,
33
+ "revision": "c30699e8356da336a370243923dbaf21066bb9fe",
34
+ "format": "MMLU-shaped (NOT byte-identical to lm-eval; auxiliary_train has no subject)"
35
+ },
36
+ "gold_label_supervision": false,
37
+ "attribution_note": "alpha=1.0 => pure KD, no gold-label CE, so any MMLU change is attributable to distillation on in-domain text. alpha<1 adds CE on text containing answers = labeled MCQ training, which is a DIFFERENT and weaker claim.",
38
+ "train_tokens_total": 263112704,
39
+ "decontamination": {
40
+ "removed": {
41
+ "alpaca": 11,
42
+ "openhermes": 120,
43
+ "ultrachat": 33
44
+ }
45
+ },
46
+ "start_kl": 7.317993998527527,
47
+ "final_kl": 0.38215016573667526,
48
+ "start_mmlu": 0.23859649122807017,
49
+ "final_mmlu": 0.3508771929824561,
50
+ "history": [
51
+ {
52
+ "tokens_m": 5.01,
53
+ "step": 153,
54
+ "held_out_kl": 0.5403434894979,
55
+ "mmlu": 0.29078947368421054
56
+ },
57
+ {
58
+ "tokens_m": 10.03,
59
+ "step": 306,
60
+ "held_out_kl": 0.470409594476223,
61
+ "mmlu": 0.2964912280701754
62
+ },
63
+ {
64
+ "tokens_m": 15.01,
65
+ "step": 458,
66
+ "held_out_kl": 0.4413424991071224,
67
+ "mmlu": 0.32280701754385965
68
+ },
69
+ {
70
+ "tokens_m": 20.02,
71
+ "step": 611,
72
+ "held_out_kl": 0.42158643528819084,
73
+ "mmlu": 0.33903508771929824
74
+ },
75
+ {
76
+ "tokens_m": 25.0,
77
+ "step": 763,
78
+ "held_out_kl": 0.4071510024368763,
79
+ "mmlu": 0.34342105263157896
80
+ },
81
+ {
82
+ "tokens_m": 30.02,
83
+ "step": 916,
84
+ "held_out_kl": 0.39785025641322136,
85
+ "mmlu": 0.34035087719298246
86
+ },
87
+ {
88
+ "tokens_m": 35.03,
89
+ "step": 1069,
90
+ "held_out_kl": 0.3906019814312458,
91
+ "mmlu": 0.3447368421052632
92
+ },
93
+ {
94
+ "tokens_m": 40.01,
95
+ "step": 1221,
96
+ "held_out_kl": 0.3873727582395077,
97
+ "mmlu": 0.3447368421052632
98
+ },
99
+ {
100
+ "tokens_m": 45.02,
101
+ "step": 1374,
102
+ "held_out_kl": 0.384848989546299,
103
+ "mmlu": 0.34649122807017546
104
+ },
105
+ {
106
+ "tokens_m": 50.0,
107
+ "step": 1526,
108
+ "held_out_kl": 0.38350963220000267,
109
+ "mmlu": 0.3460526315789474
110
+ },
111
+ {
112
+ "tokens_m": 55.02,
113
+ "step": 1679,
114
+ "held_out_kl": 0.382674265652895,
115
+ "mmlu": 0.35043859649122805
116
+ },
117
+ {
118
+ "tokens_m": 60.0,
119
+ "step": 1831,
120
+ "held_out_kl": 0.38215016573667526,
121
+ "mmlu": 0.3508771929824561,
122
+ "final": true
123
+ }
124
+ ],
125
+ "teacher_mmlu_same_limit": 0.47543859649122805,
126
+ "checkpoint": "/data1/rkp00022/zixuan/projects/TR3_main/checkpoints/kdheal_asvd7b_r0597_h1/state_dict_fp16.pt",
127
+ "sha256": "b1d633866caf2e818cfad660c86ab4c008582590ad75dfbe956d892d60b9d357",
128
+ "seed": 1234,
129
+ "mmlu_limit_per_subtask": 40,
130
+ "teacher_revision": "351844e75ed0bcbbe3f10671b3c808d2b83894ee",
131
+ "student_init_sha256": "237129f39105d18bd8f4e428476368fb79f3d9bf38d03b84070492044af69b57",
132
+ "ranks_file": "/data1/rkp00022/zixuan/projects/TR3_main/checkpoints/asvd_llama2_7b_chat_r0597/tr3fmt/ranks.json",
133
+ "ranks_sha256": "0462e6d17a528107a17701f1910b202d1aa231412cadfbcf54de00e8a0d0bba5",
134
+ "input_identity_gate": {
135
+ "student_ckpt": {
136
+ "key": "tr3fmt",
137
+ "sha256": "237129f39105d18bd8f4e428476368fb79f3d9bf38d03b84070492044af69b57",
138
+ "status": "unregistered_allowed"
139
+ },
140
+ "ranks": {
141
+ "key": "tr3fmt",
142
+ "sha256": "0462e6d17a528107a17701f1910b202d1aa231412cadfbcf54de00e8a0d0bba5",
143
+ "status": "unregistered_allowed"
144
+ }
145
+ },
146
+ "train_seed": 1234,
147
+ "mmlu_seed": 1234,
148
+ "dynamic_batch": false,
149
+ "dbl": null,
150
+ "snapshot_milestones": true,
151
+ "snapshots": [
152
+ "(PosixPath('/data1/rkp00022/zixuan/projects/TR3_main/checkpoints/kdheal_asvd7b_r0597_h1_snap15m'), 15.007744)",
153
+ "(PosixPath('/data1/rkp00022/zixuan/projects/TR3_main/checkpoints/kdheal_asvd7b_r0597_h1_snap30m'), 30.015488)",
154
+ "(PosixPath('/data1/rkp00022/zixuan/projects/TR3_main/checkpoints/kdheal_asvd7b_r0597_h1_snap45m'), 45.023232)",
155
+ "(PosixPath('/data1/rkp00022/zixuan/projects/TR3_main/checkpoints/kdheal_asvd7b_r0597_h1_snap50m'), 50.003968)",
156
+ "(PosixPath('/data1/rkp00022/zixuan/projects/TR3_main/checkpoints/kdheal_asvd7b_r0597_h1_snap55m'), 55.017472)"
157
+ ],
158
+ "stream_skip_tokens_m": 0.0,
159
+ "stream_blocks_consumed": [
160
+ 0,
161
+ 58592
162
+ ],
163
+ "stream_note": "the block order is a seeded permutation consumed from index 0; stages sharing a --seed with stream_skip=0 re-train on the SAME blocks. [start, end] block indices are recorded so any continuation can be checked for overlap with what earlier stages consumed.",
164
+ "student_master_dtype": "bf16",
165
+ "fake_quant": null,
166
+ "rebalance_report": null,
167
+ "student_param_count": 4023459584,
168
+ "micro": 16,
169
+ "grad_accum": 2,
170
+ "seq_len": 1024,
171
+ "held_out_source": "canonical: mixture 2000/4000/4000, seed 1234, seq 1024, blocks 256:264 (identical to the compression gates)",
172
+ "decontam_stats": {
173
+ "mmlu": {
174
+ "items": 14042,
175
+ "new_ngrams": 764261,
176
+ "revision": "c30699e8356da336a370243923dbaf21066bb9fe"
177
+ },
178
+ "hellaswag": {
179
+ "items": 10042,
180
+ "new_ngrams": 1334302,
181
+ "revision": "218ec52e09a7e7462a5400043bb9a69a41d06b76"
182
+ },
183
+ "boolq": {
184
+ "items": 3270,
185
+ "new_ngrams": 276806,
186
+ "revision": "3de24cf8022e94f4ee4b9d55a6f539891524d646"
187
+ },
188
+ "triviaqa": {
189
+ "items": 17944,
190
+ "new_ngrams": 32608,
191
+ "revision": "0f7faf33a3908546c6fd5b73a660e0f8ff173c2f"
192
+ },
193
+ "gsm8k": {
194
+ "items": 1319,
195
+ "new_ngrams": 46281,
196
+ "revision": "740312add88f781978c0658806c59bc2815b9866"
197
+ }
198
+ }
199
+ }
asvd_llm/heal/llama2-7b/ranks.json ADDED
@@ -0,0 +1,206 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "model.layers.0.mlp.down_proj": 1194,
3
+ "model.layers.0.mlp.gate_proj": 1194,
4
+ "model.layers.0.mlp.up_proj": 1194,
5
+ "model.layers.0.self_attn.k_proj": 819,
6
+ "model.layers.0.self_attn.o_proj": 819,
7
+ "model.layers.0.self_attn.q_proj": 819,
8
+ "model.layers.0.self_attn.v_proj": 819,
9
+ "model.layers.1.mlp.down_proj": 1194,
10
+ "model.layers.1.mlp.gate_proj": 1194,
11
+ "model.layers.1.mlp.up_proj": 1194,
12
+ "model.layers.1.self_attn.k_proj": 819,
13
+ "model.layers.1.self_attn.o_proj": 819,
14
+ "model.layers.1.self_attn.q_proj": 819,
15
+ "model.layers.1.self_attn.v_proj": 819,
16
+ "model.layers.10.mlp.down_proj": 2388,
17
+ "model.layers.10.mlp.up_proj": 2686,
18
+ "model.layers.10.self_attn.k_proj": 819,
19
+ "model.layers.10.self_attn.o_proj": 819,
20
+ "model.layers.10.self_attn.q_proj": 819,
21
+ "model.layers.10.self_attn.v_proj": 819,
22
+ "model.layers.11.mlp.gate_proj": 1194,
23
+ "model.layers.11.mlp.up_proj": 2388,
24
+ "model.layers.11.self_attn.k_proj": 819,
25
+ "model.layers.11.self_attn.o_proj": 819,
26
+ "model.layers.11.self_attn.q_proj": 819,
27
+ "model.layers.11.self_attn.v_proj": 819,
28
+ "model.layers.12.mlp.gate_proj": 1194,
29
+ "model.layers.12.mlp.up_proj": 1194,
30
+ "model.layers.12.self_attn.k_proj": 819,
31
+ "model.layers.12.self_attn.o_proj": 1228,
32
+ "model.layers.12.self_attn.q_proj": 819,
33
+ "model.layers.12.self_attn.v_proj": 1228,
34
+ "model.layers.13.mlp.gate_proj": 1194,
35
+ "model.layers.13.mlp.up_proj": 1194,
36
+ "model.layers.13.self_attn.k_proj": 819,
37
+ "model.layers.13.self_attn.o_proj": 819,
38
+ "model.layers.13.self_attn.q_proj": 819,
39
+ "model.layers.13.self_attn.v_proj": 819,
40
+ "model.layers.14.mlp.gate_proj": 1194,
41
+ "model.layers.14.mlp.up_proj": 1194,
42
+ "model.layers.14.self_attn.k_proj": 819,
43
+ "model.layers.14.self_attn.o_proj": 819,
44
+ "model.layers.14.self_attn.q_proj": 819,
45
+ "model.layers.14.self_attn.v_proj": 819,
46
+ "model.layers.15.mlp.gate_proj": 1194,
47
+ "model.layers.15.mlp.up_proj": 1194,
48
+ "model.layers.15.self_attn.k_proj": 819,
49
+ "model.layers.15.self_attn.o_proj": 819,
50
+ "model.layers.15.self_attn.q_proj": 819,
51
+ "model.layers.15.self_attn.v_proj": 819,
52
+ "model.layers.16.mlp.gate_proj": 1492,
53
+ "model.layers.16.mlp.up_proj": 1791,
54
+ "model.layers.16.self_attn.k_proj": 819,
55
+ "model.layers.16.self_attn.o_proj": 819,
56
+ "model.layers.16.self_attn.q_proj": 819,
57
+ "model.layers.16.self_attn.v_proj": 1843,
58
+ "model.layers.17.mlp.gate_proj": 1492,
59
+ "model.layers.17.mlp.up_proj": 1791,
60
+ "model.layers.17.self_attn.k_proj": 819,
61
+ "model.layers.17.self_attn.o_proj": 819,
62
+ "model.layers.17.self_attn.q_proj": 819,
63
+ "model.layers.17.self_attn.v_proj": 819,
64
+ "model.layers.18.mlp.gate_proj": 1194,
65
+ "model.layers.18.mlp.up_proj": 1194,
66
+ "model.layers.18.self_attn.k_proj": 819,
67
+ "model.layers.18.self_attn.o_proj": 819,
68
+ "model.layers.18.self_attn.q_proj": 819,
69
+ "model.layers.18.self_attn.v_proj": 819,
70
+ "model.layers.19.mlp.gate_proj": 2089,
71
+ "model.layers.19.mlp.up_proj": 2388,
72
+ "model.layers.19.self_attn.k_proj": 819,
73
+ "model.layers.19.self_attn.o_proj": 1024,
74
+ "model.layers.19.self_attn.q_proj": 819,
75
+ "model.layers.19.self_attn.v_proj": 1228,
76
+ "model.layers.2.mlp.down_proj": 1492,
77
+ "model.layers.2.mlp.gate_proj": 1194,
78
+ "model.layers.2.mlp.up_proj": 1194,
79
+ "model.layers.2.self_attn.k_proj": 819,
80
+ "model.layers.2.self_attn.o_proj": 819,
81
+ "model.layers.2.self_attn.q_proj": 819,
82
+ "model.layers.2.self_attn.v_proj": 819,
83
+ "model.layers.20.mlp.down_proj": 2686,
84
+ "model.layers.20.mlp.gate_proj": 1492,
85
+ "model.layers.20.mlp.up_proj": 1492,
86
+ "model.layers.20.self_attn.k_proj": 819,
87
+ "model.layers.20.self_attn.o_proj": 819,
88
+ "model.layers.20.self_attn.q_proj": 819,
89
+ "model.layers.20.self_attn.v_proj": 819,
90
+ "model.layers.21.mlp.gate_proj": 2089,
91
+ "model.layers.21.self_attn.k_proj": 819,
92
+ "model.layers.21.self_attn.o_proj": 819,
93
+ "model.layers.21.self_attn.q_proj": 819,
94
+ "model.layers.21.self_attn.v_proj": 819,
95
+ "model.layers.3.mlp.down_proj": 1194,
96
+ "model.layers.3.mlp.gate_proj": 1194,
97
+ "model.layers.3.mlp.up_proj": 1492,
98
+ "model.layers.3.self_attn.k_proj": 819,
99
+ "model.layers.3.self_attn.o_proj": 819,
100
+ "model.layers.3.self_attn.q_proj": 819,
101
+ "model.layers.3.self_attn.v_proj": 1024,
102
+ "model.layers.4.mlp.down_proj": 1194,
103
+ "model.layers.4.mlp.gate_proj": 1194,
104
+ "model.layers.4.mlp.up_proj": 1791,
105
+ "model.layers.4.self_attn.k_proj": 819,
106
+ "model.layers.4.self_attn.o_proj": 819,
107
+ "model.layers.4.self_attn.q_proj": 819,
108
+ "model.layers.4.self_attn.v_proj": 819,
109
+ "model.layers.5.mlp.down_proj": 1492,
110
+ "model.layers.5.mlp.gate_proj": 1194,
111
+ "model.layers.5.mlp.up_proj": 1492,
112
+ "model.layers.5.self_attn.k_proj": 819,
113
+ "model.layers.5.self_attn.o_proj": 819,
114
+ "model.layers.5.self_attn.q_proj": 819,
115
+ "model.layers.5.self_attn.v_proj": 1228,
116
+ "model.layers.6.mlp.down_proj": 2388,
117
+ "model.layers.6.mlp.gate_proj": 1194,
118
+ "model.layers.6.mlp.up_proj": 1791,
119
+ "model.layers.6.self_attn.k_proj": 819,
120
+ "model.layers.6.self_attn.o_proj": 819,
121
+ "model.layers.6.self_attn.q_proj": 819,
122
+ "model.layers.6.self_attn.v_proj": 1228,
123
+ "model.layers.7.mlp.down_proj": 1194,
124
+ "model.layers.7.mlp.gate_proj": 1194,
125
+ "model.layers.7.mlp.up_proj": 1791,
126
+ "model.layers.7.self_attn.k_proj": 819,
127
+ "model.layers.7.self_attn.o_proj": 819,
128
+ "model.layers.7.self_attn.q_proj": 819,
129
+ "model.layers.7.self_attn.v_proj": 819,
130
+ "model.layers.8.mlp.down_proj": 1492,
131
+ "model.layers.8.mlp.gate_proj": 1194,
132
+ "model.layers.8.mlp.up_proj": 2089,
133
+ "model.layers.8.self_attn.k_proj": 819,
134
+ "model.layers.8.self_attn.o_proj": 1228,
135
+ "model.layers.8.self_attn.q_proj": 819,
136
+ "model.layers.8.self_attn.v_proj": 1638,
137
+ "model.layers.9.mlp.down_proj": 2686,
138
+ "model.layers.9.mlp.gate_proj": 2089,
139
+ "model.layers.9.self_attn.k_proj": 819,
140
+ "model.layers.9.self_attn.o_proj": 1228,
141
+ "model.layers.9.self_attn.q_proj": 819,
142
+ "model.layers.9.self_attn.v_proj": 819,
143
+ "model.layers.21.mlp.up_proj": 2089,
144
+ "model.layers.22.mlp.down_proj": 2388,
145
+ "model.layers.22.mlp.gate_proj": 1492,
146
+ "model.layers.22.mlp.up_proj": 1791,
147
+ "model.layers.22.self_attn.k_proj": 819,
148
+ "model.layers.22.self_attn.o_proj": 819,
149
+ "model.layers.22.self_attn.q_proj": 819,
150
+ "model.layers.22.self_attn.v_proj": 819,
151
+ "model.layers.23.mlp.down_proj": 1791,
152
+ "model.layers.23.mlp.gate_proj": 1194,
153
+ "model.layers.23.mlp.up_proj": 1194,
154
+ "model.layers.23.self_attn.k_proj": 819,
155
+ "model.layers.23.self_attn.o_proj": 819,
156
+ "model.layers.23.self_attn.q_proj": 819,
157
+ "model.layers.23.self_attn.v_proj": 819,
158
+ "model.layers.24.mlp.down_proj": 2089,
159
+ "model.layers.24.mlp.gate_proj": 2089,
160
+ "model.layers.24.mlp.up_proj": 1791,
161
+ "model.layers.24.self_attn.k_proj": 819,
162
+ "model.layers.24.self_attn.o_proj": 819,
163
+ "model.layers.24.self_attn.q_proj": 819,
164
+ "model.layers.24.self_attn.v_proj": 1024,
165
+ "model.layers.25.mlp.down_proj": 1791,
166
+ "model.layers.25.mlp.gate_proj": 1791,
167
+ "model.layers.25.mlp.up_proj": 2089,
168
+ "model.layers.25.self_attn.k_proj": 819,
169
+ "model.layers.25.self_attn.o_proj": 819,
170
+ "model.layers.25.self_attn.q_proj": 819,
171
+ "model.layers.25.self_attn.v_proj": 819,
172
+ "model.layers.26.mlp.gate_proj": 2089,
173
+ "model.layers.26.mlp.up_proj": 1492,
174
+ "model.layers.26.self_attn.k_proj": 819,
175
+ "model.layers.26.self_attn.o_proj": 819,
176
+ "model.layers.26.self_attn.q_proj": 819,
177
+ "model.layers.26.self_attn.v_proj": 1843,
178
+ "model.layers.27.mlp.down_proj": 2089,
179
+ "model.layers.27.mlp.gate_proj": 2089,
180
+ "model.layers.27.mlp.up_proj": 2686,
181
+ "model.layers.27.self_attn.k_proj": 819,
182
+ "model.layers.27.self_attn.o_proj": 819,
183
+ "model.layers.27.self_attn.q_proj": 819,
184
+ "model.layers.27.self_attn.v_proj": 819,
185
+ "model.layers.28.mlp.down_proj": 2686,
186
+ "model.layers.28.mlp.up_proj": 2686,
187
+ "model.layers.28.self_attn.k_proj": 819,
188
+ "model.layers.28.self_attn.o_proj": 819,
189
+ "model.layers.28.self_attn.q_proj": 819,
190
+ "model.layers.29.mlp.down_proj": 2388,
191
+ "model.layers.29.mlp.gate_proj": 2686,
192
+ "model.layers.29.self_attn.k_proj": 819,
193
+ "model.layers.29.self_attn.o_proj": 819,
194
+ "model.layers.29.self_attn.q_proj": 819,
195
+ "model.layers.29.self_attn.v_proj": 819,
196
+ "model.layers.30.mlp.gate_proj": 2388,
197
+ "model.layers.30.self_attn.k_proj": 819,
198
+ "model.layers.30.self_attn.o_proj": 819,
199
+ "model.layers.30.self_attn.q_proj": 819,
200
+ "model.layers.30.self_attn.v_proj": 819,
201
+ "model.layers.31.mlp.down_proj": 2686,
202
+ "model.layers.31.self_attn.k_proj": 819,
203
+ "model.layers.31.self_attn.o_proj": 819,
204
+ "model.layers.31.self_attn.q_proj": 819,
205
+ "model.layers.31.self_attn.v_proj": 819
206
+ }
asvd_llm/heal/llama2-7b/state_dict_fp16.pt ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:b1d633866caf2e818cfad660c86ab4c008582590ad75dfbe956d892d60b9d357
3
+ size 8047083200
asvd_llm/unheal/llama2-13b/README.md ADDED
@@ -0,0 +1,11 @@
 
 
 
 
 
 
 
 
 
 
 
 
1
+ ---
2
+ license: llama2
3
+ base_model: NousResearch/Llama-2-13b-chat-hf
4
+ tags: [svd, asvd, low-rank, model-compression, llama2]
5
+ ---
6
+
7
+ # ASVD (arXiv:2312.05821) LLaMA-2-13B-chat, UNHEALED
8
+
9
+ - **UNHEALED checkpoint**: raw post-training compression weights — no distillation. The
10
+ KD-healed counterpart is at `../../heal/llama2-13b/`.
11
+ - Total factored params: **7,736,414,208** (our exact accounting; see `provenance.json`).
asvd_llm/unheal/llama2-13b/config.json ADDED
@@ -0,0 +1,26 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "_name_or_path": null,
3
+ "architectures": [
4
+ "LlamaForCausalLM"
5
+ ],
6
+ "bos_token_id": 1,
7
+ "eos_token_id": 2,
8
+ "hidden_act": "silu",
9
+ "hidden_size": 5120,
10
+ "initializer_range": 0.02,
11
+ "intermediate_size": 13824,
12
+ "max_position_embeddings": 4096,
13
+ "model_type": "llama",
14
+ "num_attention_heads": 40,
15
+ "num_hidden_layers": 40,
16
+ "num_key_value_heads": 40,
17
+ "pad_token_id": 0,
18
+ "pretraining_tp": 1,
19
+ "rms_norm_eps": 1e-05,
20
+ "rope_scaling": null,
21
+ "tie_word_embeddings": false,
22
+ "torch_dtype": "float16",
23
+ "transformers_version": "4.31.0.dev0",
24
+ "use_cache": true,
25
+ "vocab_size": 32000
26
+ }
asvd_llm/unheal/llama2-13b/lowrank_linear.py ADDED
@@ -0,0 +1,83 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """Minimal standalone loader for the SVD-baseline compressed LLaMA-2-7B checkpoints.
2
+
3
+ No dependency beyond torch + transformers. The state_dict stores every compressed
4
+ projection as two factors of a LowRankLinear (A: out x k, B: k x in); ranks.json
5
+ holds the per-layer integer rank of each projection. Untouched parameters
6
+ (embeddings, norms) keep their stock transformers names.
7
+
8
+ Usage:
9
+ import json, torch
10
+ from lowrank_linear import build
11
+ model = build(
12
+ base_id="NousResearch/Llama-2-7b-chat-hf",
13
+ ranks=json.load(open("ranks.json")),
14
+ state_dict_path="state_dict_fp16.pt",
15
+ device="cuda", dtype=torch.float16,
16
+ )
17
+ tokenizer = AutoTokenizer.from_pretrained("NousResearch/Llama-2-7b-chat-hf")
18
+ out = model.generate(tokenizer("2+2=", return_tensors="pt").to("cuda").input_ids,
19
+ max_new_tokens=16, do_sample=False)
20
+ print(tokenizer.decode(out[0]))
21
+ """
22
+
23
+ from __future__ import annotations
24
+
25
+ import torch
26
+ import torch.nn as nn
27
+ import torch.nn.functional as F
28
+
29
+
30
+ class LowRankLinear(nn.Module):
31
+ """y = A (B x). A: out x r, B: r x in. (LLaMA projections carry no bias.)"""
32
+
33
+ def __init__(self, A: torch.Tensor, B: torch.Tensor):
34
+ super().__init__()
35
+ self.A = nn.Parameter(A)
36
+ self.B = nn.Parameter(B)
37
+
38
+ def forward(self, x):
39
+ return F.linear(F.linear(x, self.B), self.A)
40
+
41
+ @torch.no_grad()
42
+ def materialize_(self) -> nn.Linear:
43
+ """Replace this module with the dense W = A @ B (inference only; never call
44
+ while training -- it cuts the gradient path)."""
45
+ W = (self.A @ self.B).to(self.A.dtype)
46
+ dense = nn.Linear(self.B.shape[1], self.A.shape[0], bias=False,
47
+ device=W.device, dtype=W.dtype)
48
+ dense.weight.data = W
49
+ return dense
50
+
51
+
52
+ def build(base_id: str, ranks: dict, state_dict_path: str,
53
+ device: str = "cuda", dtype=torch.float16):
54
+ from transformers import AutoModelForCausalLM
55
+
56
+ model = AutoModelForCausalLM.from_pretrained(base_id, torch_dtype=dtype)
57
+ search = {}
58
+ for name, mod in model.named_modules():
59
+ if isinstance(mod, nn.Linear):
60
+ search[name] = mod
61
+ for path, k in ranks.items():
62
+ out_f, in_f = search[path].weight.shape
63
+ model_mod = model
64
+ for p in path.split(".")[:-1]:
65
+ model_mod = getattr(model_mod, p)
66
+ setattr(model_mod, path.split(".")[-1],
67
+ LowRankLinear(torch.zeros(out_f, int(k), dtype=dtype),
68
+ torch.zeros(int(k), in_f, dtype=dtype)))
69
+ sd = torch.load(state_dict_path, map_location="cpu", weights_only=True)
70
+ model.load_state_dict(sd, strict=True)
71
+ del sd
72
+ return model.to(device)
73
+
74
+
75
+ def materialize_all(model):
76
+ """Fuse every LowRankLinear back into a dense nn.Linear in place."""
77
+ for name, mod in list(model.named_modules()):
78
+ if isinstance(mod, LowRankLinear):
79
+ parent = model
80
+ for p in name.split(".")[:-1]:
81
+ parent = getattr(parent, p)
82
+ setattr(parent, name.split(".")[-1], mod.materialize_())
83
+ return model
asvd_llm/unheal/llama2-13b/provenance.json ADDED
@@ -0,0 +1,15 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "method": "ASVD official (abs_mean scaling, ppl sensitivity, param_ratio target)",
3
+ "model_id": "NousResearch/Llama-2-13b-chat-hf",
4
+ "param_ratio_target": 0.5897,
5
+ "factored_params": 5444860928,
6
+ "untouched_params": 2291553280,
7
+ "total_params": 7736414208,
8
+ "alpha": 0.5,
9
+ "seed": 233,
10
+ "n_calib": 256,
11
+ "sensitivity_metric": "ppl",
12
+ "state_dict_sha256": "49b4e7679fb063510a8d64c114c2a532708fbad8f822b51baa2969d5f1e20211",
13
+ "healed_counterpart": "../../heal/llama2-13b/state_dict_fp16.pt",
14
+ "note": "converted to TR3 LowRank format (pure re-keying, meta counts verified)"
15
+ }
asvd_llm/unheal/llama2-13b/ranks.json ADDED
@@ -0,0 +1,253 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "model.layers.0.mlp.down_proj": 1494,
3
+ "model.layers.0.mlp.gate_proj": 1494,
4
+ "model.layers.0.mlp.up_proj": 1494,
5
+ "model.layers.0.self_attn.k_proj": 1024,
6
+ "model.layers.0.self_attn.o_proj": 1024,
7
+ "model.layers.0.self_attn.q_proj": 1024,
8
+ "model.layers.0.self_attn.v_proj": 1024,
9
+ "model.layers.1.mlp.down_proj": 1494,
10
+ "model.layers.1.mlp.gate_proj": 1494,
11
+ "model.layers.1.mlp.up_proj": 1494,
12
+ "model.layers.1.self_attn.k_proj": 1024,
13
+ "model.layers.1.self_attn.o_proj": 1024,
14
+ "model.layers.1.self_attn.q_proj": 1024,
15
+ "model.layers.1.self_attn.v_proj": 1024,
16
+ "model.layers.10.mlp.gate_proj": 2615,
17
+ "model.layers.10.mlp.up_proj": 2988,
18
+ "model.layers.10.self_attn.k_proj": 1024,
19
+ "model.layers.10.self_attn.o_proj": 1024,
20
+ "model.layers.10.self_attn.q_proj": 1024,
21
+ "model.layers.10.self_attn.v_proj": 1024,
22
+ "model.layers.11.mlp.gate_proj": 3362,
23
+ "model.layers.11.mlp.up_proj": 2241,
24
+ "model.layers.11.self_attn.k_proj": 1024,
25
+ "model.layers.11.self_attn.o_proj": 1024,
26
+ "model.layers.11.self_attn.q_proj": 1024,
27
+ "model.layers.11.self_attn.v_proj": 1024,
28
+ "model.layers.12.mlp.gate_proj": 2988,
29
+ "model.layers.12.mlp.up_proj": 2241,
30
+ "model.layers.12.self_attn.k_proj": 1024,
31
+ "model.layers.12.self_attn.o_proj": 1024,
32
+ "model.layers.12.self_attn.q_proj": 1024,
33
+ "model.layers.12.self_attn.v_proj": 1024,
34
+ "model.layers.13.mlp.gate_proj": 1494,
35
+ "model.layers.13.mlp.up_proj": 1494,
36
+ "model.layers.13.self_attn.k_proj": 1024,
37
+ "model.layers.13.self_attn.o_proj": 1024,
38
+ "model.layers.13.self_attn.q_proj": 1024,
39
+ "model.layers.13.self_attn.v_proj": 1024,
40
+ "model.layers.14.mlp.gate_proj": 1494,
41
+ "model.layers.14.self_attn.k_proj": 1536,
42
+ "model.layers.14.self_attn.o_proj": 1024,
43
+ "model.layers.14.self_attn.q_proj": 1024,
44
+ "model.layers.14.self_attn.v_proj": 1024,
45
+ "model.layers.2.mlp.down_proj": 1868,
46
+ "model.layers.2.mlp.gate_proj": 1494,
47
+ "model.layers.2.mlp.up_proj": 1868,
48
+ "model.layers.2.self_attn.k_proj": 1024,
49
+ "model.layers.2.self_attn.o_proj": 1024,
50
+ "model.layers.2.self_attn.q_proj": 1024,
51
+ "model.layers.2.self_attn.v_proj": 1024,
52
+ "model.layers.3.mlp.down_proj": 1494,
53
+ "model.layers.3.mlp.gate_proj": 1494,
54
+ "model.layers.3.mlp.up_proj": 1868,
55
+ "model.layers.3.self_attn.k_proj": 1024,
56
+ "model.layers.3.self_attn.o_proj": 1024,
57
+ "model.layers.3.self_attn.q_proj": 1024,
58
+ "model.layers.3.self_attn.v_proj": 1024,
59
+ "model.layers.4.mlp.down_proj": 1494,
60
+ "model.layers.4.mlp.gate_proj": 1494,
61
+ "model.layers.4.mlp.up_proj": 1494,
62
+ "model.layers.4.self_attn.k_proj": 1024,
63
+ "model.layers.4.self_attn.o_proj": 1280,
64
+ "model.layers.4.self_attn.q_proj": 1024,
65
+ "model.layers.4.self_attn.v_proj": 1024,
66
+ "model.layers.5.mlp.down_proj": 1494,
67
+ "model.layers.5.mlp.gate_proj": 1494,
68
+ "model.layers.5.mlp.up_proj": 1494,
69
+ "model.layers.5.self_attn.k_proj": 1024,
70
+ "model.layers.5.self_attn.o_proj": 1024,
71
+ "model.layers.5.self_attn.q_proj": 1024,
72
+ "model.layers.5.self_attn.v_proj": 1280,
73
+ "model.layers.6.mlp.down_proj": 1494,
74
+ "model.layers.6.mlp.gate_proj": 1494,
75
+ "model.layers.6.mlp.up_proj": 1868,
76
+ "model.layers.6.self_attn.k_proj": 1024,
77
+ "model.layers.6.self_attn.o_proj": 1024,
78
+ "model.layers.6.self_attn.q_proj": 1024,
79
+ "model.layers.6.self_attn.v_proj": 1280,
80
+ "model.layers.7.mlp.down_proj": 1494,
81
+ "model.layers.7.mlp.gate_proj": 1494,
82
+ "model.layers.7.mlp.up_proj": 1494,
83
+ "model.layers.7.self_attn.k_proj": 1024,
84
+ "model.layers.7.self_attn.o_proj": 1024,
85
+ "model.layers.7.self_attn.q_proj": 1024,
86
+ "model.layers.7.self_attn.v_proj": 1024,
87
+ "model.layers.8.mlp.down_proj": 2241,
88
+ "model.layers.8.mlp.gate_proj": 1868,
89
+ "model.layers.8.mlp.up_proj": 2615,
90
+ "model.layers.8.self_attn.k_proj": 1024,
91
+ "model.layers.8.self_attn.o_proj": 1024,
92
+ "model.layers.8.self_attn.q_proj": 1024,
93
+ "model.layers.8.self_attn.v_proj": 1024,
94
+ "model.layers.9.mlp.gate_proj": 1494,
95
+ "model.layers.9.mlp.up_proj": 2988,
96
+ "model.layers.9.self_attn.k_proj": 1024,
97
+ "model.layers.9.self_attn.o_proj": 1280,
98
+ "model.layers.9.self_attn.q_proj": 1024,
99
+ "model.layers.9.self_attn.v_proj": 1536,
100
+ "model.layers.14.mlp.up_proj": 1494,
101
+ "model.layers.15.mlp.gate_proj": 1494,
102
+ "model.layers.15.mlp.up_proj": 1494,
103
+ "model.layers.15.self_attn.k_proj": 1280,
104
+ "model.layers.15.self_attn.o_proj": 1024,
105
+ "model.layers.15.self_attn.q_proj": 1024,
106
+ "model.layers.15.self_attn.v_proj": 1024,
107
+ "model.layers.16.mlp.gate_proj": 1494,
108
+ "model.layers.16.mlp.up_proj": 1494,
109
+ "model.layers.16.self_attn.k_proj": 1024,
110
+ "model.layers.16.self_attn.o_proj": 1024,
111
+ "model.layers.16.self_attn.q_proj": 1024,
112
+ "model.layers.16.self_attn.v_proj": 1024,
113
+ "model.layers.17.mlp.gate_proj": 1494,
114
+ "model.layers.17.mlp.up_proj": 1494,
115
+ "model.layers.17.self_attn.k_proj": 1024,
116
+ "model.layers.17.self_attn.o_proj": 1024,
117
+ "model.layers.17.self_attn.q_proj": 1024,
118
+ "model.layers.17.self_attn.v_proj": 1024,
119
+ "model.layers.18.mlp.gate_proj": 1494,
120
+ "model.layers.18.mlp.up_proj": 1494,
121
+ "model.layers.18.self_attn.k_proj": 1024,
122
+ "model.layers.18.self_attn.o_proj": 1024,
123
+ "model.layers.18.self_attn.q_proj": 1024,
124
+ "model.layers.18.self_attn.v_proj": 1792,
125
+ "model.layers.19.mlp.gate_proj": 1494,
126
+ "model.layers.19.mlp.up_proj": 1494,
127
+ "model.layers.19.self_attn.k_proj": 1280,
128
+ "model.layers.19.self_attn.o_proj": 1024,
129
+ "model.layers.19.self_attn.q_proj": 1024,
130
+ "model.layers.19.self_attn.v_proj": 1024,
131
+ "model.layers.20.mlp.gate_proj": 1494,
132
+ "model.layers.20.mlp.up_proj": 1494,
133
+ "model.layers.20.self_attn.k_proj": 1024,
134
+ "model.layers.20.self_attn.o_proj": 1024,
135
+ "model.layers.20.self_attn.q_proj": 1024,
136
+ "model.layers.20.self_attn.v_proj": 1536,
137
+ "model.layers.21.mlp.gate_proj": 1494,
138
+ "model.layers.21.mlp.up_proj": 2241,
139
+ "model.layers.21.self_attn.k_proj": 1024,
140
+ "model.layers.21.self_attn.o_proj": 1024,
141
+ "model.layers.21.self_attn.q_proj": 1024,
142
+ "model.layers.21.self_attn.v_proj": 1024,
143
+ "model.layers.22.mlp.gate_proj": 2615,
144
+ "model.layers.22.mlp.up_proj": 3362,
145
+ "model.layers.22.self_attn.k_proj": 1024,
146
+ "model.layers.22.self_attn.o_proj": 1024,
147
+ "model.layers.22.self_attn.q_proj": 1024,
148
+ "model.layers.22.self_attn.v_proj": 1024,
149
+ "model.layers.23.mlp.gate_proj": 2615,
150
+ "model.layers.23.mlp.up_proj": 2615,
151
+ "model.layers.23.self_attn.k_proj": 1280,
152
+ "model.layers.23.self_attn.o_proj": 1280,
153
+ "model.layers.23.self_attn.q_proj": 1536,
154
+ "model.layers.23.self_attn.v_proj": 1280,
155
+ "model.layers.24.mlp.down_proj": 3362,
156
+ "model.layers.24.mlp.gate_proj": 1868,
157
+ "model.layers.24.mlp.up_proj": 2241,
158
+ "model.layers.24.self_attn.k_proj": 1024,
159
+ "model.layers.24.self_attn.o_proj": 1024,
160
+ "model.layers.24.self_attn.q_proj": 1024,
161
+ "model.layers.24.self_attn.v_proj": 1792,
162
+ "model.layers.25.mlp.down_proj": 3362,
163
+ "model.layers.25.mlp.gate_proj": 2241,
164
+ "model.layers.25.mlp.up_proj": 2988,
165
+ "model.layers.25.self_attn.k_proj": 1024,
166
+ "model.layers.25.self_attn.o_proj": 1280,
167
+ "model.layers.25.self_attn.q_proj": 1024,
168
+ "model.layers.25.self_attn.v_proj": 1536,
169
+ "model.layers.26.mlp.down_proj": 2988,
170
+ "model.layers.26.mlp.gate_proj": 1868,
171
+ "model.layers.26.mlp.up_proj": 1868,
172
+ "model.layers.26.self_attn.k_proj": 1024,
173
+ "model.layers.26.self_attn.o_proj": 1024,
174
+ "model.layers.26.self_attn.q_proj": 1024,
175
+ "model.layers.26.self_attn.v_proj": 1280,
176
+ "model.layers.27.mlp.down_proj": 2988,
177
+ "model.layers.27.mlp.gate_proj": 1494,
178
+ "model.layers.27.mlp.up_proj": 1494,
179
+ "model.layers.27.self_attn.k_proj": 1024,
180
+ "model.layers.27.self_attn.o_proj": 1024,
181
+ "model.layers.27.self_attn.q_proj": 1024,
182
+ "model.layers.27.self_attn.v_proj": 1536,
183
+ "model.layers.28.mlp.down_proj": 2988,
184
+ "model.layers.28.mlp.gate_proj": 2615,
185
+ "model.layers.28.mlp.up_proj": 2615,
186
+ "model.layers.28.self_attn.k_proj": 1024,
187
+ "model.layers.28.self_attn.o_proj": 1024,
188
+ "model.layers.28.self_attn.q_proj": 1024,
189
+ "model.layers.28.self_attn.v_proj": 1024,
190
+ "model.layers.29.mlp.down_proj": 2988,
191
+ "model.layers.29.mlp.gate_proj": 2615,
192
+ "model.layers.29.mlp.up_proj": 2241,
193
+ "model.layers.29.self_attn.k_proj": 1024,
194
+ "model.layers.29.self_attn.o_proj": 1024,
195
+ "model.layers.29.self_attn.q_proj": 1024,
196
+ "model.layers.29.self_attn.v_proj": 1024,
197
+ "model.layers.30.mlp.down_proj": 1868,
198
+ "model.layers.30.mlp.gate_proj": 2988,
199
+ "model.layers.30.mlp.up_proj": 2615,
200
+ "model.layers.30.self_attn.k_proj": 1024,
201
+ "model.layers.30.self_attn.o_proj": 1024,
202
+ "model.layers.30.self_attn.q_proj": 1024,
203
+ "model.layers.30.self_attn.v_proj": 1024,
204
+ "model.layers.31.mlp.down_proj": 1868,
205
+ "model.layers.31.mlp.gate_proj": 2988,
206
+ "model.layers.31.mlp.up_proj": 2241,
207
+ "model.layers.31.self_attn.k_proj": 1024,
208
+ "model.layers.31.self_attn.o_proj": 1024,
209
+ "model.layers.31.self_attn.q_proj": 1024,
210
+ "model.layers.31.self_attn.v_proj": 1024,
211
+ "model.layers.32.mlp.down_proj": 2241,
212
+ "model.layers.32.mlp.gate_proj": 3362,
213
+ "model.layers.32.mlp.up_proj": 2988,
214
+ "model.layers.32.self_attn.k_proj": 1024,
215
+ "model.layers.32.self_attn.o_proj": 1024,
216
+ "model.layers.32.self_attn.q_proj": 1024,
217
+ "model.layers.32.self_attn.v_proj": 1024,
218
+ "model.layers.33.mlp.down_proj": 1494,
219
+ "model.layers.33.mlp.up_proj": 3362,
220
+ "model.layers.33.self_attn.k_proj": 1024,
221
+ "model.layers.33.self_attn.o_proj": 1024,
222
+ "model.layers.33.self_attn.q_proj": 1024,
223
+ "model.layers.33.self_attn.v_proj": 1536,
224
+ "model.layers.34.mlp.down_proj": 1868,
225
+ "model.layers.34.self_attn.k_proj": 1024,
226
+ "model.layers.34.self_attn.o_proj": 1024,
227
+ "model.layers.34.self_attn.q_proj": 1024,
228
+ "model.layers.34.self_attn.v_proj": 1024,
229
+ "model.layers.35.mlp.down_proj": 2988,
230
+ "model.layers.35.self_attn.k_proj": 1024,
231
+ "model.layers.35.self_attn.o_proj": 1024,
232
+ "model.layers.35.self_attn.q_proj": 1024,
233
+ "model.layers.35.self_attn.v_proj": 1536,
234
+ "model.layers.36.mlp.down_proj": 1868,
235
+ "model.layers.36.mlp.up_proj": 2615,
236
+ "model.layers.36.self_attn.k_proj": 1024,
237
+ "model.layers.36.self_attn.o_proj": 1280,
238
+ "model.layers.36.self_attn.q_proj": 1024,
239
+ "model.layers.37.mlp.down_proj": 2615,
240
+ "model.layers.37.self_attn.o_proj": 1024,
241
+ "model.layers.37.self_attn.q_proj": 1024,
242
+ "model.layers.37.self_attn.v_proj": 1536,
243
+ "model.layers.38.mlp.down_proj": 2988,
244
+ "model.layers.38.mlp.gate_proj": 3362,
245
+ "model.layers.38.self_attn.k_proj": 1024,
246
+ "model.layers.38.self_attn.o_proj": 1024,
247
+ "model.layers.38.self_attn.q_proj": 1024,
248
+ "model.layers.38.self_attn.v_proj": 1024,
249
+ "model.layers.39.self_attn.k_proj": 1024,
250
+ "model.layers.39.self_attn.o_proj": 1024,
251
+ "model.layers.39.self_attn.q_proj": 1024,
252
+ "model.layers.39.self_attn.v_proj": 1024
253
+ }
asvd_llm/unheal/llama2-13b/state_dict_fp16.pt ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:49b4e7679fb063510a8d64c114c2a532708fbad8f822b51baa2969d5f1e20211
3
+ size 15473031795
asvd_llm/unheal/llama2-7b/README.md ADDED
@@ -0,0 +1,11 @@
 
 
 
 
 
 
 
 
 
 
 
 
1
+ ---
2
+ license: llama2
3
+ base_model: NousResearch/Llama-2-7b-chat-hf
4
+ tags: [svd, asvd, low-rank, model-compression, llama2]
5
+ ---
6
+
7
+ # ASVD (arXiv:2312.05821) LLaMA-2-7B-chat, UNHEALED
8
+
9
+ - **UNHEALED checkpoint**: raw post-training compression weights — no distillation. The
10
+ KD-healed counterpart is at `../../heal/llama2-7b/`.
11
+ - Total factored params: **4,023,459,584** (our exact accounting; see `provenance.json`).
asvd_llm/unheal/llama2-7b/config.json ADDED
@@ -0,0 +1,25 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "architectures": [
3
+ "LlamaForCausalLM"
4
+ ],
5
+ "bos_token_id": 1,
6
+ "eos_token_id": 2,
7
+ "hidden_act": "silu",
8
+ "hidden_size": 4096,
9
+ "initializer_range": 0.02,
10
+ "intermediate_size": 11008,
11
+ "max_position_embeddings": 4096,
12
+ "model_type": "llama",
13
+ "num_attention_heads": 32,
14
+ "num_hidden_layers": 32,
15
+ "num_key_value_heads": 32,
16
+ "pad_token_id": 0,
17
+ "pretraining_tp": 1,
18
+ "rms_norm_eps": 1e-05,
19
+ "rope_scaling": null,
20
+ "tie_word_embeddings": false,
21
+ "torch_dtype": "float16",
22
+ "transformers_version": "4.31.0.dev0",
23
+ "use_cache": true,
24
+ "vocab_size": 32000
25
+ }
asvd_llm/unheal/llama2-7b/lowrank_linear.py ADDED
@@ -0,0 +1,83 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """Minimal standalone loader for the SVD-baseline compressed LLaMA-2-7B checkpoints.
2
+
3
+ No dependency beyond torch + transformers. The state_dict stores every compressed
4
+ projection as two factors of a LowRankLinear (A: out x k, B: k x in); ranks.json
5
+ holds the per-layer integer rank of each projection. Untouched parameters
6
+ (embeddings, norms) keep their stock transformers names.
7
+
8
+ Usage:
9
+ import json, torch
10
+ from lowrank_linear import build
11
+ model = build(
12
+ base_id="NousResearch/Llama-2-7b-chat-hf",
13
+ ranks=json.load(open("ranks.json")),
14
+ state_dict_path="state_dict_fp16.pt",
15
+ device="cuda", dtype=torch.float16,
16
+ )
17
+ tokenizer = AutoTokenizer.from_pretrained("NousResearch/Llama-2-7b-chat-hf")
18
+ out = model.generate(tokenizer("2+2=", return_tensors="pt").to("cuda").input_ids,
19
+ max_new_tokens=16, do_sample=False)
20
+ print(tokenizer.decode(out[0]))
21
+ """
22
+
23
+ from __future__ import annotations
24
+
25
+ import torch
26
+ import torch.nn as nn
27
+ import torch.nn.functional as F
28
+
29
+
30
+ class LowRankLinear(nn.Module):
31
+ """y = A (B x). A: out x r, B: r x in. (LLaMA projections carry no bias.)"""
32
+
33
+ def __init__(self, A: torch.Tensor, B: torch.Tensor):
34
+ super().__init__()
35
+ self.A = nn.Parameter(A)
36
+ self.B = nn.Parameter(B)
37
+
38
+ def forward(self, x):
39
+ return F.linear(F.linear(x, self.B), self.A)
40
+
41
+ @torch.no_grad()
42
+ def materialize_(self) -> nn.Linear:
43
+ """Replace this module with the dense W = A @ B (inference only; never call
44
+ while training -- it cuts the gradient path)."""
45
+ W = (self.A @ self.B).to(self.A.dtype)
46
+ dense = nn.Linear(self.B.shape[1], self.A.shape[0], bias=False,
47
+ device=W.device, dtype=W.dtype)
48
+ dense.weight.data = W
49
+ return dense
50
+
51
+
52
+ def build(base_id: str, ranks: dict, state_dict_path: str,
53
+ device: str = "cuda", dtype=torch.float16):
54
+ from transformers import AutoModelForCausalLM
55
+
56
+ model = AutoModelForCausalLM.from_pretrained(base_id, torch_dtype=dtype)
57
+ search = {}
58
+ for name, mod in model.named_modules():
59
+ if isinstance(mod, nn.Linear):
60
+ search[name] = mod
61
+ for path, k in ranks.items():
62
+ out_f, in_f = search[path].weight.shape
63
+ model_mod = model
64
+ for p in path.split(".")[:-1]:
65
+ model_mod = getattr(model_mod, p)
66
+ setattr(model_mod, path.split(".")[-1],
67
+ LowRankLinear(torch.zeros(out_f, int(k), dtype=dtype),
68
+ torch.zeros(int(k), in_f, dtype=dtype)))
69
+ sd = torch.load(state_dict_path, map_location="cpu", weights_only=True)
70
+ model.load_state_dict(sd, strict=True)
71
+ del sd
72
+ return model.to(device)
73
+
74
+
75
+ def materialize_all(model):
76
+ """Fuse every LowRankLinear back into a dense nn.Linear in place."""
77
+ for name, mod in list(model.named_modules()):
78
+ if isinstance(mod, LowRankLinear):
79
+ parent = model
80
+ for p in name.split(".")[:-1]:
81
+ parent = getattr(parent, p)
82
+ setattr(parent, name.split(".")[-1], mod.materialize_())
83
+ return model
asvd_llm/unheal/llama2-7b/provenance.json ADDED
@@ -0,0 +1,15 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "method": "ASVD official (abs_mean scaling, ppl sensitivity, param_ratio target)",
3
+ "model_id": "NousResearch/Llama-2-7b-chat-hf",
4
+ "param_ratio_target": 0.589,
5
+ "factored_params": 2887585536,
6
+ "untouched_params": 1135874048,
7
+ "total_params": 4023459584,
8
+ "alpha": 0.5,
9
+ "seed": 233,
10
+ "n_calib": 256,
11
+ "sensitivity_metric": "ppl",
12
+ "state_dict_sha256": "237129f39105d18bd8f4e428476368fb79f3d9bf38d03b84070492044af69b57",
13
+ "healed_counterpart": "../../heal/llama2-7b/state_dict_fp16.pt",
14
+ "note": "converted to TR3 LowRank format (pure re-keying, meta counts verified)"
15
+ }
asvd_llm/unheal/llama2-7b/ranks.json ADDED
@@ -0,0 +1,206 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "model.layers.0.mlp.down_proj": 1194,
3
+ "model.layers.0.mlp.gate_proj": 1194,
4
+ "model.layers.0.mlp.up_proj": 1194,
5
+ "model.layers.0.self_attn.k_proj": 819,
6
+ "model.layers.0.self_attn.o_proj": 819,
7
+ "model.layers.0.self_attn.q_proj": 819,
8
+ "model.layers.0.self_attn.v_proj": 819,
9
+ "model.layers.1.mlp.down_proj": 1194,
10
+ "model.layers.1.mlp.gate_proj": 1194,
11
+ "model.layers.1.mlp.up_proj": 1194,
12
+ "model.layers.1.self_attn.k_proj": 819,
13
+ "model.layers.1.self_attn.o_proj": 819,
14
+ "model.layers.1.self_attn.q_proj": 819,
15
+ "model.layers.1.self_attn.v_proj": 819,
16
+ "model.layers.10.mlp.down_proj": 2388,
17
+ "model.layers.10.mlp.up_proj": 2686,
18
+ "model.layers.10.self_attn.k_proj": 819,
19
+ "model.layers.10.self_attn.o_proj": 819,
20
+ "model.layers.10.self_attn.q_proj": 819,
21
+ "model.layers.10.self_attn.v_proj": 819,
22
+ "model.layers.11.mlp.gate_proj": 1194,
23
+ "model.layers.11.mlp.up_proj": 2388,
24
+ "model.layers.11.self_attn.k_proj": 819,
25
+ "model.layers.11.self_attn.o_proj": 819,
26
+ "model.layers.11.self_attn.q_proj": 819,
27
+ "model.layers.11.self_attn.v_proj": 819,
28
+ "model.layers.12.mlp.gate_proj": 1194,
29
+ "model.layers.12.mlp.up_proj": 1194,
30
+ "model.layers.12.self_attn.k_proj": 819,
31
+ "model.layers.12.self_attn.o_proj": 1228,
32
+ "model.layers.12.self_attn.q_proj": 819,
33
+ "model.layers.12.self_attn.v_proj": 1228,
34
+ "model.layers.13.mlp.gate_proj": 1194,
35
+ "model.layers.13.mlp.up_proj": 1194,
36
+ "model.layers.13.self_attn.k_proj": 819,
37
+ "model.layers.13.self_attn.o_proj": 819,
38
+ "model.layers.13.self_attn.q_proj": 819,
39
+ "model.layers.13.self_attn.v_proj": 819,
40
+ "model.layers.14.mlp.gate_proj": 1194,
41
+ "model.layers.14.mlp.up_proj": 1194,
42
+ "model.layers.14.self_attn.k_proj": 819,
43
+ "model.layers.14.self_attn.o_proj": 819,
44
+ "model.layers.14.self_attn.q_proj": 819,
45
+ "model.layers.14.self_attn.v_proj": 819,
46
+ "model.layers.15.mlp.gate_proj": 1194,
47
+ "model.layers.15.mlp.up_proj": 1194,
48
+ "model.layers.15.self_attn.k_proj": 819,
49
+ "model.layers.15.self_attn.o_proj": 819,
50
+ "model.layers.15.self_attn.q_proj": 819,
51
+ "model.layers.15.self_attn.v_proj": 819,
52
+ "model.layers.16.mlp.gate_proj": 1492,
53
+ "model.layers.16.mlp.up_proj": 1791,
54
+ "model.layers.16.self_attn.k_proj": 819,
55
+ "model.layers.16.self_attn.o_proj": 819,
56
+ "model.layers.16.self_attn.q_proj": 819,
57
+ "model.layers.16.self_attn.v_proj": 1843,
58
+ "model.layers.17.mlp.gate_proj": 1492,
59
+ "model.layers.17.mlp.up_proj": 1791,
60
+ "model.layers.17.self_attn.k_proj": 819,
61
+ "model.layers.17.self_attn.o_proj": 819,
62
+ "model.layers.17.self_attn.q_proj": 819,
63
+ "model.layers.17.self_attn.v_proj": 819,
64
+ "model.layers.18.mlp.gate_proj": 1194,
65
+ "model.layers.18.mlp.up_proj": 1194,
66
+ "model.layers.18.self_attn.k_proj": 819,
67
+ "model.layers.18.self_attn.o_proj": 819,
68
+ "model.layers.18.self_attn.q_proj": 819,
69
+ "model.layers.18.self_attn.v_proj": 819,
70
+ "model.layers.19.mlp.gate_proj": 2089,
71
+ "model.layers.19.mlp.up_proj": 2388,
72
+ "model.layers.19.self_attn.k_proj": 819,
73
+ "model.layers.19.self_attn.o_proj": 1024,
74
+ "model.layers.19.self_attn.q_proj": 819,
75
+ "model.layers.19.self_attn.v_proj": 1228,
76
+ "model.layers.2.mlp.down_proj": 1492,
77
+ "model.layers.2.mlp.gate_proj": 1194,
78
+ "model.layers.2.mlp.up_proj": 1194,
79
+ "model.layers.2.self_attn.k_proj": 819,
80
+ "model.layers.2.self_attn.o_proj": 819,
81
+ "model.layers.2.self_attn.q_proj": 819,
82
+ "model.layers.2.self_attn.v_proj": 819,
83
+ "model.layers.20.mlp.down_proj": 2686,
84
+ "model.layers.20.mlp.gate_proj": 1492,
85
+ "model.layers.20.mlp.up_proj": 1492,
86
+ "model.layers.20.self_attn.k_proj": 819,
87
+ "model.layers.20.self_attn.o_proj": 819,
88
+ "model.layers.20.self_attn.q_proj": 819,
89
+ "model.layers.20.self_attn.v_proj": 819,
90
+ "model.layers.21.mlp.gate_proj": 2089,
91
+ "model.layers.21.self_attn.k_proj": 819,
92
+ "model.layers.21.self_attn.o_proj": 819,
93
+ "model.layers.21.self_attn.q_proj": 819,
94
+ "model.layers.21.self_attn.v_proj": 819,
95
+ "model.layers.3.mlp.down_proj": 1194,
96
+ "model.layers.3.mlp.gate_proj": 1194,
97
+ "model.layers.3.mlp.up_proj": 1492,
98
+ "model.layers.3.self_attn.k_proj": 819,
99
+ "model.layers.3.self_attn.o_proj": 819,
100
+ "model.layers.3.self_attn.q_proj": 819,
101
+ "model.layers.3.self_attn.v_proj": 1024,
102
+ "model.layers.4.mlp.down_proj": 1194,
103
+ "model.layers.4.mlp.gate_proj": 1194,
104
+ "model.layers.4.mlp.up_proj": 1791,
105
+ "model.layers.4.self_attn.k_proj": 819,
106
+ "model.layers.4.self_attn.o_proj": 819,
107
+ "model.layers.4.self_attn.q_proj": 819,
108
+ "model.layers.4.self_attn.v_proj": 819,
109
+ "model.layers.5.mlp.down_proj": 1492,
110
+ "model.layers.5.mlp.gate_proj": 1194,
111
+ "model.layers.5.mlp.up_proj": 1492,
112
+ "model.layers.5.self_attn.k_proj": 819,
113
+ "model.layers.5.self_attn.o_proj": 819,
114
+ "model.layers.5.self_attn.q_proj": 819,
115
+ "model.layers.5.self_attn.v_proj": 1228,
116
+ "model.layers.6.mlp.down_proj": 2388,
117
+ "model.layers.6.mlp.gate_proj": 1194,
118
+ "model.layers.6.mlp.up_proj": 1791,
119
+ "model.layers.6.self_attn.k_proj": 819,
120
+ "model.layers.6.self_attn.o_proj": 819,
121
+ "model.layers.6.self_attn.q_proj": 819,
122
+ "model.layers.6.self_attn.v_proj": 1228,
123
+ "model.layers.7.mlp.down_proj": 1194,
124
+ "model.layers.7.mlp.gate_proj": 1194,
125
+ "model.layers.7.mlp.up_proj": 1791,
126
+ "model.layers.7.self_attn.k_proj": 819,
127
+ "model.layers.7.self_attn.o_proj": 819,
128
+ "model.layers.7.self_attn.q_proj": 819,
129
+ "model.layers.7.self_attn.v_proj": 819,
130
+ "model.layers.8.mlp.down_proj": 1492,
131
+ "model.layers.8.mlp.gate_proj": 1194,
132
+ "model.layers.8.mlp.up_proj": 2089,
133
+ "model.layers.8.self_attn.k_proj": 819,
134
+ "model.layers.8.self_attn.o_proj": 1228,
135
+ "model.layers.8.self_attn.q_proj": 819,
136
+ "model.layers.8.self_attn.v_proj": 1638,
137
+ "model.layers.9.mlp.down_proj": 2686,
138
+ "model.layers.9.mlp.gate_proj": 2089,
139
+ "model.layers.9.self_attn.k_proj": 819,
140
+ "model.layers.9.self_attn.o_proj": 1228,
141
+ "model.layers.9.self_attn.q_proj": 819,
142
+ "model.layers.9.self_attn.v_proj": 819,
143
+ "model.layers.21.mlp.up_proj": 2089,
144
+ "model.layers.22.mlp.down_proj": 2388,
145
+ "model.layers.22.mlp.gate_proj": 1492,
146
+ "model.layers.22.mlp.up_proj": 1791,
147
+ "model.layers.22.self_attn.k_proj": 819,
148
+ "model.layers.22.self_attn.o_proj": 819,
149
+ "model.layers.22.self_attn.q_proj": 819,
150
+ "model.layers.22.self_attn.v_proj": 819,
151
+ "model.layers.23.mlp.down_proj": 1791,
152
+ "model.layers.23.mlp.gate_proj": 1194,
153
+ "model.layers.23.mlp.up_proj": 1194,
154
+ "model.layers.23.self_attn.k_proj": 819,
155
+ "model.layers.23.self_attn.o_proj": 819,
156
+ "model.layers.23.self_attn.q_proj": 819,
157
+ "model.layers.23.self_attn.v_proj": 819,
158
+ "model.layers.24.mlp.down_proj": 2089,
159
+ "model.layers.24.mlp.gate_proj": 2089,
160
+ "model.layers.24.mlp.up_proj": 1791,
161
+ "model.layers.24.self_attn.k_proj": 819,
162
+ "model.layers.24.self_attn.o_proj": 819,
163
+ "model.layers.24.self_attn.q_proj": 819,
164
+ "model.layers.24.self_attn.v_proj": 1024,
165
+ "model.layers.25.mlp.down_proj": 1791,
166
+ "model.layers.25.mlp.gate_proj": 1791,
167
+ "model.layers.25.mlp.up_proj": 2089,
168
+ "model.layers.25.self_attn.k_proj": 819,
169
+ "model.layers.25.self_attn.o_proj": 819,
170
+ "model.layers.25.self_attn.q_proj": 819,
171
+ "model.layers.25.self_attn.v_proj": 819,
172
+ "model.layers.26.mlp.gate_proj": 2089,
173
+ "model.layers.26.mlp.up_proj": 1492,
174
+ "model.layers.26.self_attn.k_proj": 819,
175
+ "model.layers.26.self_attn.o_proj": 819,
176
+ "model.layers.26.self_attn.q_proj": 819,
177
+ "model.layers.26.self_attn.v_proj": 1843,
178
+ "model.layers.27.mlp.down_proj": 2089,
179
+ "model.layers.27.mlp.gate_proj": 2089,
180
+ "model.layers.27.mlp.up_proj": 2686,
181
+ "model.layers.27.self_attn.k_proj": 819,
182
+ "model.layers.27.self_attn.o_proj": 819,
183
+ "model.layers.27.self_attn.q_proj": 819,
184
+ "model.layers.27.self_attn.v_proj": 819,
185
+ "model.layers.28.mlp.down_proj": 2686,
186
+ "model.layers.28.mlp.up_proj": 2686,
187
+ "model.layers.28.self_attn.k_proj": 819,
188
+ "model.layers.28.self_attn.o_proj": 819,
189
+ "model.layers.28.self_attn.q_proj": 819,
190
+ "model.layers.29.mlp.down_proj": 2388,
191
+ "model.layers.29.mlp.gate_proj": 2686,
192
+ "model.layers.29.self_attn.k_proj": 819,
193
+ "model.layers.29.self_attn.o_proj": 819,
194
+ "model.layers.29.self_attn.q_proj": 819,
195
+ "model.layers.29.self_attn.v_proj": 819,
196
+ "model.layers.30.mlp.gate_proj": 2388,
197
+ "model.layers.30.self_attn.k_proj": 819,
198
+ "model.layers.30.self_attn.o_proj": 819,
199
+ "model.layers.30.self_attn.q_proj": 819,
200
+ "model.layers.30.self_attn.v_proj": 819,
201
+ "model.layers.31.mlp.down_proj": 2686,
202
+ "model.layers.31.self_attn.k_proj": 819,
203
+ "model.layers.31.self_attn.o_proj": 819,
204
+ "model.layers.31.self_attn.q_proj": 819,
205
+ "model.layers.31.self_attn.v_proj": 819
206
+ }
asvd_llm/unheal/llama2-7b/state_dict_fp16.pt ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:237129f39105d18bd8f4e428476368fb79f3d9bf38d03b84070492044af69b57
3
+ size 8047081433