suchirsalhan commited on
Commit
1f39d9b
·
verified ·
1 Parent(s): fa5de2d

merged checkpoint (qmd_guided/transport+aligned)

Browse files
zoo-language-beetle-monolingual-fineweb-100m-kor__beetle-monolingual-fineweb-100m-spa/qmd_guided__transport-aligned/.destination.json ADDED
@@ -0,0 +1,20 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "repo_id": "Mergeability/fineweb-100m-language-03",
3
+ "subfolder": "zoo-language-beetle-monolingual-fineweb-100m-kor__beetle-monolingual-fineweb-100m-spa/qmd_guided__transport-aligned",
4
+ "bytes": 778365041,
5
+ "meta": {
6
+ "pair_id": "zoo-language-beetle-monolingual-fineweb-100m-kor__beetle-monolingual-fineweb-100m-spa",
7
+ "parent_a": "Beetle-FineWeb-100M/beetle-monolingual-fineweb-100m-kor",
8
+ "parent_b": "Beetle-FineWeb-100M/beetle-monolingual-fineweb-100m-spa",
9
+ "ceiling": "(floor-relative)",
10
+ "operator": "qmd_guided",
11
+ "alignment": "transport+aligned",
12
+ "align_method": "permutation",
13
+ "regime": "beetle_language",
14
+ "eval_langs": "kor+spa",
15
+ "nll_merge": 3.4174,
16
+ "nll_floor": 2.6033,
17
+ "param_coverage": 1.0,
18
+ "MS": null
19
+ }
20
+ }
zoo-language-beetle-monolingual-fineweb-100m-kor__beetle-monolingual-fineweb-100m-spa/qmd_guided__transport-aligned/README.md ADDED
@@ -0,0 +1,43 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ ---
2
+ library_name: transformers
3
+ tags: [model-merging, mergeability, training-free, quotient-merge-distance]
4
+ ---
5
+
6
+ # qmd_guided__transport-aligned
7
+
8
+ Training-free merged checkpoint from the **Mergeability** sweep
9
+ (`benchmark/emit_lm.py --real`), produced by weight-space merging of two independently
10
+ trained parents. No gradient steps were taken.
11
+
12
+ | field | value |
13
+ |---|---|
14
+ | pair_id | `zoo-language-beetle-monolingual-fineweb-100m-kor__beetle-monolingual-fineweb-100m-spa` |
15
+ | parent_a | `Beetle-FineWeb-100M/beetle-monolingual-fineweb-100m-kor` |
16
+ | parent_b | `Beetle-FineWeb-100M/beetle-monolingual-fineweb-100m-spa` |
17
+ | ceiling | `(floor-relative)` |
18
+ | operator | `qmd_guided` |
19
+ | alignment | `transport+aligned` |
20
+ | align_method | `permutation` |
21
+ | regime | `beetle_language` |
22
+ | eval_langs | `kor+spa` |
23
+ | nll_merge | `3.4174` |
24
+ | nll_floor | `2.6033` |
25
+ | param_coverage | `1.0` |
26
+ | MS | `None` |
27
+
28
+ ## How it was made
29
+
30
+ Parents were loaded, activations extracted on a shared calibration corpus, and the merge applied
31
+ either **naive** (parents combined in their own coordinates) or **aligned** (parent B carried into
32
+ parent A's residual-stream basis via `common.alignment.residual_basis_map` before merging —
33
+ permutation for same-width pairs, orthogonal/rectangular for cross-width).
34
+
35
+ `MS` is the recovery score from `common.eval.mergeability_score` (merged vs. floor vs. ceiling), the
36
+ same normalisation used by Zhou et al., so it is comparable across rows of the sweep.
37
+
38
+ ## Caveats
39
+
40
+ Sub-1B merges are noisy; an aligned signal where the naive one is noise is the finding, not a bug.
41
+ Rows without a joint ceiling are floor-relative and must not be read as absolute recovery.
42
+
43
+ Generated automatically — see the [mergeschool repo](https://github.com/suchirsalhan/merge-school).
zoo-language-beetle-monolingual-fineweb-100m-kor__beetle-monolingual-fineweb-100m-spa/qmd_guided__transport-aligned/config.json ADDED
@@ -0,0 +1,24 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "activation_hidden_dim": 3072,
3
+ "architectures": [
4
+ "PicoDecoderHF"
5
+ ],
6
+ "attention_n_heads": 12,
7
+ "attention_n_kv_heads": 1,
8
+ "auto_map": {
9
+ "AutoConfig": "pico_decoder.PicoDecoderHFConfig",
10
+ "AutoModelForCausalLM": "pico_decoder.PicoDecoderHF"
11
+ },
12
+ "batch_size": 64,
13
+ "d_model": 768,
14
+ "dropout": 0.1,
15
+ "dtype": "float32",
16
+ "max_seq_len": 512,
17
+ "model_type": "pico_decoder",
18
+ "n_layers": 14,
19
+ "norm_eps": 1e-05,
20
+ "position_emb_theta": 10000.0,
21
+ "training_seed": 42,
22
+ "transformers_version": "5.14.1",
23
+ "vocab_size": 50000
24
+ }
zoo-language-beetle-monolingual-fineweb-100m-kor__beetle-monolingual-fineweb-100m-spa/qmd_guided__transport-aligned/generation_config.json ADDED
@@ -0,0 +1,6 @@
 
 
 
 
 
 
 
1
+ {
2
+ "_from_model_config": true,
3
+ "output_attentions": false,
4
+ "output_hidden_states": false,
5
+ "transformers_version": "5.14.1"
6
+ }
zoo-language-beetle-monolingual-fineweb-100m-kor__beetle-monolingual-fineweb-100m-spa/qmd_guided__transport-aligned/model.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:49e72956f1cd345755caf92fa98dbaecfa9d0d53bc18f55731473ec3a4894dab
3
+ size 775229392
zoo-language-beetle-monolingual-fineweb-100m-kor__beetle-monolingual-fineweb-100m-spa/qmd_guided__transport-aligned/pico_decoder.py ADDED
@@ -0,0 +1,410 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """
2
+ Pico Decoder: A Lightweight Causal Transformer Language Model
3
+ Implementation from https://github.com/pico-lm/pico-train/blob/main/src/model/pico_decoder.py
4
+
5
+ Key features:
6
+ - RMSNorm for layer normalization
7
+ - Rotary Positional Embeddings (RoPE)
8
+ - Multi-head attention with KV-cache support
9
+ - SwiGLU activation function
10
+ - Residual connections throughout
11
+ - KV-cache for faster autoregressive generation
12
+
13
+ References:
14
+ - RoPE: https://arxiv.org/abs/2104.09864
15
+ - SwiGLU: https://arxiv.org/abs/2002.05202
16
+ - LLAMA: https://arxiv.org/abs/2302.13971
17
+
18
+ HuggingFace compatibility notes
19
+ ---------------------------------
20
+ PicoDecoderHF stores weights at the TOP LEVEL (embedding_proj, layers,
21
+ output_norm, de_embedding_proj) so that state dict keys match the raw
22
+ PicoDecoder checkpoint format exactly. Do NOT add a self.pico_decoder
23
+ wrapper — it would prepend a key prefix that does not exist in any saved
24
+ checkpoint, causing every weight to be MISSING on load.
25
+
26
+ vocab_size in config.json must be the BASE BPE vocabulary size
27
+ (tokenizer.vocab_size), NOT the padded len(tokenizer) which includes
28
+ <unusedN> padding tokens. The embedding table was built with the base size.
29
+ """
30
+
31
+ from dataclasses import asdict, is_dataclass
32
+ from typing import TYPE_CHECKING, Any, Dict, Optional, Tuple, Union
33
+ import torch
34
+ import torch.nn as nn
35
+ import torch.nn.functional as F
36
+ from torch.nn.attention import SDPBackend, sdpa_kernel
37
+ from transformers import GenerationMixin, PretrainedConfig, PreTrainedModel
38
+ from transformers.modeling_outputs import CausalLMOutput, CausalLMOutputWithPast
39
+ try:
40
+ if TYPE_CHECKING:
41
+ from src.config import ModelConfig
42
+ except ImportError:
43
+ pass
44
+
45
+
46
+ class RMSNorm(torch.nn.Module):
47
+ def __init__(self, config):
48
+ super().__init__()
49
+ self.eps = config.norm_eps
50
+ self.weight = nn.Parameter(torch.ones(config.d_model))
51
+ def _norm(self, x):
52
+ return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
53
+ def forward(self, x):
54
+ return self._norm(x.float()).type_as(x) * self.weight
55
+
56
+
57
+ class RoPE(nn.Module):
58
+ """
59
+ Rotary Position Embedding.
60
+ freqs_cis is computed lazily on first use and cached per-device,
61
+ avoiding meta-tensor issues when HF loads with low_cpu_mem_usage=True.
62
+ The cache auto-extends if a call needs a longer sequence than the
63
+ current cached length (grows geometrically).
64
+ """
65
+ def __init__(self, config):
66
+ super().__init__()
67
+ self.theta = config.position_emb_theta
68
+ self.dim = config.d_model // config.attention_n_heads
69
+ self.max_seq = config.max_seq_len
70
+ # NOT a buffer — plain dict so it never touches the meta device
71
+ self._cache: Dict[torch.device, torch.Tensor] = {}
72
+ self._cache_max_seq: Dict[torch.device, int] = {}
73
+
74
+ def _build_cache(self, device: torch.device, length: int) -> torch.Tensor:
75
+ freqs = 1.0 / (
76
+ self.theta ** (
77
+ torch.arange(0, self.dim, 2, device=device).float() / self.dim
78
+ )
79
+ )
80
+ t = torch.arange(length, device=device)
81
+ freqs = torch.outer(t, freqs)
82
+ return torch.polar(torch.ones_like(freqs), freqs)
83
+
84
+ def _get_freqs_cis(self, device: torch.device, min_length: int = 0) -> torch.Tensor:
85
+ cached_len = self._cache_max_seq.get(device, 0)
86
+ needed = max(min_length, self.max_seq)
87
+ if device not in self._cache or cached_len < needed:
88
+ new_len = max(needed, cached_len * 2)
89
+ self._cache[device] = self._build_cache(device, new_len)
90
+ self._cache_max_seq[device] = new_len
91
+ return self._cache[device]
92
+
93
+ def get_freqs_cis(self, input_shape, start_pos, end_pos, device):
94
+ _f = self._get_freqs_cis(device, min_length=end_pos)[start_pos:end_pos]
95
+ ndim = len(input_shape)
96
+ if ndim < 2:
97
+ raise ValueError(
98
+ f"RoPE expects input with ndim >= 2, got shape {tuple(input_shape)}"
99
+ )
100
+ expected = (input_shape[1], input_shape[-1])
101
+ if tuple(_f.shape) != expected:
102
+ raise ValueError(
103
+ f"RoPE freqs_cis shape mismatch: got {tuple(_f.shape)}, "
104
+ f"expected {expected} (input_shape={tuple(input_shape)}, "
105
+ f"start_pos={start_pos}, end_pos={end_pos}, "
106
+ f"cached_max_seq={self._cache_max_seq.get(device)}, dim={self.dim})"
107
+ )
108
+ return _f.view(*[d if i == 1 or i == ndim - 1 else 1
109
+ for i, d in enumerate(input_shape)])
110
+
111
+ def forward(self, queries, keys, start_pos=0):
112
+ device = queries.device
113
+ q_ = torch.view_as_complex(queries.float().reshape(*queries.shape[:-1], -1, 2))
114
+ k_ = torch.view_as_complex(keys.float().reshape(*keys.shape[:-1], -1, 2))
115
+ fc = self.get_freqs_cis(q_.shape, start_pos, start_pos + q_.shape[1], device)
116
+ return (torch.view_as_real(q_ * fc).flatten(3).type_as(queries),
117
+ torch.view_as_real(k_ * fc).flatten(3).type_as(keys))
118
+
119
+
120
+ class Attention(nn.Module):
121
+ def __init__(self, config):
122
+ super().__init__()
123
+ self.n_heads = config.attention_n_heads
124
+ self.n_kv_heads = config.attention_n_kv_heads
125
+ self.batch_size = config.batch_size
126
+ self.max_seq_len = config.max_seq_len
127
+ d = config.d_model
128
+ self.head_dim = d // self.n_heads
129
+ self.n_rep = self.n_heads // self.n_kv_heads
130
+ self.q_proj = nn.Linear(d, self.n_heads * self.head_dim, bias=False)
131
+ self.k_proj = nn.Linear(d, self.n_kv_heads * self.head_dim, bias=False)
132
+ self.v_proj = nn.Linear(d, self.n_kv_heads * self.head_dim, bias=False)
133
+ self.o_proj = nn.Linear(self.n_heads * self.head_dim, d, bias=False)
134
+ self.rope = RoPE(config)
135
+ # Optional working-memory hook. Set by PicoDecoder.set_working_memory().
136
+ # When present, it must be a callable returning a float tensor of shape
137
+ # (1, n_heads, seq_len, seq_len) to be added to the attention mask.
138
+ self._working_memory_hook = None
139
+ def forward(self, input, mask=None, past_key_values=None, use_cache=False):
140
+ bsz, seq_len, _ = input.shape
141
+ queries = self.q_proj(input).view(bsz, seq_len, self.n_heads, self.head_dim)
142
+ keys = self.k_proj(input).view(bsz, seq_len, self.n_kv_heads, self.head_dim)
143
+ values = self.v_proj(input).view(bsz, seq_len, self.n_kv_heads, self.head_dim)
144
+ start_pos = past_key_values[0].shape[1] if past_key_values is not None else 0
145
+ queries, keys = self.rope(queries, keys, start_pos)
146
+ if past_key_values is not None:
147
+ keys = torch.cat([past_key_values[0], keys], dim=1)
148
+ values = torch.cat([past_key_values[1], values], dim=1)
149
+ cached_keys = keys if use_cache else None
150
+ cached_values = values if use_cache else None
151
+ queries = queries.transpose(1, 2)
152
+ keys = keys.transpose(1, 2)
153
+ values = values.transpose(1, 2)
154
+ apply_gqa = self.n_rep > 1
155
+ if apply_gqa and queries.device.type == "mps":
156
+ keys = keys.repeat_interleave(self.n_rep, dim=-3)
157
+ values = values.repeat_interleave(self.n_rep, dim=-3)
158
+ apply_gqa = False
159
+ attn_mask = mask.to(queries.dtype) if mask is not None else None
160
+ if self._working_memory_hook is not None:
161
+ key_len = keys.shape[-2]
162
+ alibi_bias = self._working_memory_hook(
163
+ key_len, queries.device, queries.dtype
164
+ )
165
+ # Crop / broadcast to (1, n_heads, q_len, k_len).
166
+ alibi_bias = alibi_bias[:, :, key_len - seq_len:key_len, :key_len]
167
+ if attn_mask is None:
168
+ attn_mask = alibi_bias.expand(bsz, -1, -1, -1)
169
+ else:
170
+ # attn_mask may be (q_len, k_len) — broadcast to 4D first.
171
+ while attn_mask.dim() < 4:
172
+ attn_mask = attn_mask.unsqueeze(0)
173
+ attn_mask = attn_mask + alibi_bias
174
+ with sdpa_kernel(backends=[SDPBackend.CUDNN_ATTENTION, SDPBackend.MATH]):
175
+ attn_output = F.scaled_dot_product_attention(
176
+ queries.contiguous(), keys.contiguous(), values.contiguous(),
177
+ attn_mask=attn_mask, enable_gqa=apply_gqa,
178
+ )
179
+ attn_output = attn_output.transpose(1, 2).contiguous().view(bsz, seq_len, -1)
180
+ return self.o_proj(attn_output), (cached_keys, cached_values)
181
+
182
+
183
+ class SwiGLU(nn.Module):
184
+ def __init__(self, config):
185
+ super().__init__()
186
+ self.w_0 = nn.Linear(config.d_model, config.activation_hidden_dim, bias=False)
187
+ self.w_1 = nn.Linear(config.d_model, config.activation_hidden_dim, bias=False)
188
+ self.w_2 = nn.Linear(config.activation_hidden_dim, config.d_model, bias=False)
189
+ def forward(self, x):
190
+ return self.w_2(F.silu(self.w_0(x)) * self.w_1(x))
191
+
192
+
193
+ class PicoDecoderBlock(nn.Module):
194
+ def __init__(self, config):
195
+ super().__init__()
196
+ self.attention = Attention(config)
197
+ self.swiglu = SwiGLU(config)
198
+ self.attention_norm = RMSNorm(config)
199
+ self.swiglu_norm = RMSNorm(config)
200
+ def forward(self, input, mask=None, past_key_values=None, use_cache=False):
201
+ attention_output, cached_key_values = self.attention(
202
+ self.attention_norm(input), mask=mask,
203
+ past_key_values=past_key_values, use_cache=use_cache)
204
+ h = input + attention_output
205
+ return h + self.swiglu(self.swiglu_norm(h)), cached_key_values
206
+
207
+
208
+ class PicoDecoder(nn.Module):
209
+ def __init__(self, model_config):
210
+ super().__init__()
211
+ self.config = model_config
212
+ self.embedding_proj = nn.Embedding(model_config.vocab_size, model_config.d_model)
213
+ self.layers = nn.ModuleList(
214
+ [PicoDecoderBlock(model_config) for _ in range(model_config.n_layers)])
215
+ self.output_norm = RMSNorm(model_config)
216
+ self.de_embedding_proj = nn.Linear(
217
+ model_config.d_model, model_config.vocab_size, bias=False)
218
+ # Optional working-memory scheduler (see src/working_memory). When
219
+ # absent, the forward pass is identical to the vanilla PicoDecoder.
220
+ self._working_memory = None
221
+
222
+ # ──────────────────────────────────────────────────────────────────────
223
+ # Optional working-memory ALiBi hook (see src/working_memory)
224
+ # ──────────────────────────────────────────────────────────────────────
225
+ def set_working_memory(self, scheduler) -> None:
226
+ """Attach a ``WorkingMemoryScheduler``; pass ``None`` to detach."""
227
+ self._working_memory = scheduler
228
+ hook = (lambda n, d, t: scheduler.bias(n, d, t)) if scheduler is not None else None
229
+ for layer in self.layers:
230
+ layer.attention._working_memory_hook = hook
231
+
232
+ def update_progress(self, step: int, max_steps: int) -> None:
233
+ """Tick the working-memory scheduler (no-op if not attached)."""
234
+ if self._working_memory is not None:
235
+ self._working_memory.step(step, max_steps)
236
+
237
+ def convert_to_hf_model(self):
238
+ # Fabric/DDP wrappers expose the inner module via .module; unwrap so
239
+ # state_dict() never carries a "module." prefix that PicoDecoderHF
240
+ # cannot match. Without this, strict=False below silently dropped
241
+ # every weight, save_pretrained wrote an empty checkpoint without a
242
+ # valid model_type key, and the HF-smoke round-trip raised
243
+ # "Unrecognized model in step_0".
244
+ raw = getattr(self, "module", self)
245
+ hf = PicoDecoderHF(PicoDecoderHFConfig.from_dataclass(raw.config))
246
+ # Strip auxiliary heads added by continual-learning modules (MAML's
247
+ # classifier_head, etc.). They are not part of the published HF
248
+ # architecture; loading them would raise "Unexpected key(s)" in
249
+ # PicoDecoderHF.load_state_dict and crash the per-step save.
250
+ backbone_state = {
251
+ k: v for k, v in raw.state_dict().items()
252
+ if not k.startswith("classifier_head")
253
+ }
254
+ missing, unexpected = hf.load_state_dict(backbone_state, strict=False)
255
+ if unexpected:
256
+ raise RuntimeError(
257
+ f"PicoDecoder.convert_to_hf_model: {len(unexpected)} unexpected "
258
+ f"keys in state_dict: {sorted(unexpected)[:5]}…"
259
+ )
260
+ return hf
261
+ def forward(self, input_ids, past_key_values=None, use_cache=False,
262
+ return_hidden=False):
263
+ """
264
+ return_hidden=True returns (hidden_states, cached_key_values) where
265
+ hidden_states is the final-layer post-norm representation (pre-
266
+ de_embedding_proj), used by MAML's classifier head and any other
267
+ consumer that needs a feature view of the model rather than logits.
268
+ Default behaviour is unchanged.
269
+ """
270
+ seq_len = input_ids.shape[-1]
271
+ h = self.embedding_proj(input_ids)
272
+ start_pos = 0 if past_key_values is None else past_key_values[0][0].shape[1]
273
+ mask = None
274
+ if seq_len > 1:
275
+ mask = torch.full((seq_len, seq_len), float("-inf"))
276
+ mask = torch.triu(mask, diagonal=1)
277
+ if past_key_values is not None:
278
+ mask = torch.hstack([torch.zeros((seq_len, start_pos)), mask])
279
+ mask = mask.to(h.device)
280
+ cached_key_values = () if use_cache else None
281
+ for idx, layer in enumerate(self.layers):
282
+ layer_past = past_key_values[idx] if past_key_values is not None else None
283
+ h, layer_cached = layer(
284
+ h, mask=mask, past_key_values=layer_past, use_cache=use_cache)
285
+ if use_cache:
286
+ cached_key_values += (layer_cached,)
287
+ hidden = self.output_norm(h)
288
+ if return_hidden:
289
+ return hidden, cached_key_values
290
+ return self.de_embedding_proj(hidden).float(), cached_key_values
291
+
292
+
293
+ class PicoDecoderHFConfig(PretrainedConfig):
294
+ model_type = "pico_decoder"
295
+ def __init__(self,
296
+ n_layers=14, d_model=768, vocab_size=32768,
297
+ attention_n_heads=12, attention_n_kv_heads=1,
298
+ max_seq_len=512, batch_size=64, position_emb_theta=10000.0,
299
+ activation_hidden_dim=3072, norm_eps=1e-5, dropout=0.1,
300
+ **kwargs):
301
+ if not attention_n_kv_heads:
302
+ attention_n_kv_heads = attention_n_heads
303
+ super().__init__(**kwargs)
304
+ self.n_layers = n_layers
305
+ self.d_model = d_model
306
+ self.vocab_size = vocab_size
307
+ self.attention_n_heads = attention_n_heads
308
+ self.attention_n_kv_heads = attention_n_kv_heads
309
+ self.max_seq_len = max_seq_len
310
+ self.batch_size = batch_size
311
+ self.position_emb_theta = position_emb_theta
312
+ self.activation_hidden_dim = activation_hidden_dim
313
+ self.norm_eps = norm_eps
314
+ self.dropout = dropout
315
+ @classmethod
316
+ def from_dict(cls, config_dict: Dict[str, Any], **kwargs) -> "PicoDecoderHFConfig":
317
+ pico_config = cls(**config_dict)
318
+ return_unused_kwargs = kwargs.pop("return_unused_kwargs", False)
319
+ unused_kwargs = {k: v for k, v in kwargs.items() if not hasattr(pico_config, k)}
320
+ if return_unused_kwargs:
321
+ return pico_config, unused_kwargs
322
+ return pico_config
323
+ @classmethod
324
+ def from_dataclass(cls, model_config):
325
+ if is_dataclass(model_config) and not isinstance(model_config, type):
326
+ d = asdict(model_config)
327
+ elif isinstance(model_config, dict):
328
+ d = dict(model_config)
329
+ elif hasattr(model_config, "__dict__"):
330
+ d = dict(vars(model_config))
331
+ else:
332
+ raise TypeError(
333
+ f"Cannot build PicoDecoderHFConfig from {type(model_config).__name__}"
334
+ )
335
+ return cls.from_dict(d)
336
+
337
+
338
+ class PicoDecoderHF(PreTrainedModel, GenerationMixin):
339
+ """
340
+ HuggingFace wrapper for BeetleLM PicoDecoder.
341
+ Usage: AutoModelForCausalLM.from_pretrained(repo, trust_remote_code=True)
342
+ Works with CPU, CUDA (A100, etc.), and MPS out of the box.
343
+ """
344
+ config_class = PicoDecoderHFConfig
345
+ _no_split_modules = ["PicoDecoderBlock"]
346
+ _tied_weights_keys = []
347
+
348
+ def __init__(self, config: PicoDecoderHFConfig):
349
+ super().__init__(config)
350
+ self.embedding_proj = nn.Embedding(config.vocab_size, config.d_model)
351
+ self.layers = nn.ModuleList(
352
+ [PicoDecoderBlock(config) for _ in range(config.n_layers)])
353
+ self.output_norm = RMSNorm(config)
354
+ self.de_embedding_proj = nn.Linear(config.d_model, config.vocab_size, bias=False)
355
+ # Required: lets HF finalize weight init and meta-device materialization
356
+ self.post_init()
357
+
358
+ # Required for low_cpu_mem_usage / Accelerate device-dispatch to work
359
+ def _init_weights(self, module):
360
+ if isinstance(module, nn.Linear):
361
+ nn.init.normal_(module.weight, mean=0.0, std=0.02)
362
+ if module.bias is not None:
363
+ nn.init.zeros_(module.bias)
364
+ elif isinstance(module, nn.Embedding):
365
+ nn.init.normal_(module.weight, mean=0.0, std=0.02)
366
+ elif isinstance(module, RMSNorm):
367
+ nn.init.ones_(module.weight)
368
+
369
+ def get_input_embeddings(self): return self.embedding_proj
370
+ def set_input_embeddings(self, value): self.embedding_proj = value
371
+
372
+ def forward(self, input_ids=None, past_key_values=None,
373
+ use_cache=False, labels=None, **kwargs):
374
+ seq_len = input_ids.shape[-1]
375
+ h = self.embedding_proj(input_ids)
376
+ start_pos = 0 if past_key_values is None else past_key_values[0][0].shape[1]
377
+ mask = None
378
+ if seq_len > 1:
379
+ mask = torch.full((seq_len, seq_len), float("-inf"), device=h.device)
380
+ mask = torch.triu(mask, diagonal=1)
381
+ if past_key_values is not None:
382
+ mask = torch.hstack([torch.zeros((seq_len, start_pos), device=h.device), mask])
383
+ cached_key_values = () if use_cache else None
384
+ for idx, layer in enumerate(self.layers):
385
+ layer_past = past_key_values[idx] if past_key_values is not None else None
386
+ h, layer_cached = layer(
387
+ h, mask=mask, past_key_values=layer_past, use_cache=use_cache)
388
+ if use_cache:
389
+ cached_key_values += (layer_cached,)
390
+ logits = self.de_embedding_proj(self.output_norm(h)).float()
391
+ loss = None
392
+ if labels is not None:
393
+ loss = F.cross_entropy(
394
+ logits[:, :-1].contiguous().view(-1, self.config.vocab_size),
395
+ labels[:, 1:].contiguous().clamp(0, self.config.vocab_size - 1).view(-1),
396
+ )
397
+ if use_cache:
398
+ return CausalLMOutputWithPast(
399
+ loss=loss, logits=logits, past_key_values=cached_key_values)
400
+ return CausalLMOutput(loss=loss, logits=logits)
401
+
402
+ def prepare_inputs_for_generation(self, input_ids, past_key_values=None, **kwargs):
403
+ return {"input_ids": input_ids,
404
+ "past_key_values": past_key_values,
405
+ "use_cache": True}
406
+
407
+
408
+ PicoDecoderHFConfig.register_for_auto_class()
409
+ PicoDecoderHF.register_for_auto_class("AutoModel")
410
+ PicoDecoderHF.register_for_auto_class("AutoModelForCausalLM")
zoo-language-beetle-monolingual-fineweb-100m-kor__beetle-monolingual-fineweb-100m-spa/qmd_guided__transport-aligned/tokenizer.json ADDED
The diff for this file is too large to render. See raw diff
 
zoo-language-beetle-monolingual-fineweb-100m-kor__beetle-monolingual-fineweb-100m-spa/qmd_guided__transport-aligned/tokenizer_config.json ADDED
@@ -0,0 +1,12 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "backend": "tokenizers",
3
+ "bos_token": "<s>",
4
+ "clean_up_tokenization_spaces": false,
5
+ "eos_token": "</s>",
6
+ "is_local": false,
7
+ "local_files_only": false,
8
+ "model_max_length": 1000000000000000019884624838656,
9
+ "pad_token": "<pad>",
10
+ "tokenizer_class": "TokenizersBackend",
11
+ "unk_token": "<unk>"
12
+ }