Eclipse-Senpai commited on
Commit
91941e4
·
verified ·
1 Parent(s): 7edbdc4

scrub internal project references from vendored model source

Browse files
Files changed (1) hide show
  1. meiosis.py +10 -11
meiosis.py CHANGED
@@ -1,10 +1,9 @@
1
- """Meiosis: PICO release 1 (2026-07). Tied-embedding looped decoder-only LM.
2
 
3
- Spec: research/2026-07-first-release/final-spec.md (approved 2026-07-02).
4
  embed -> prelude x1 -> [body of `body_blocks` distinct blocks xK loops,
5
  per-loop LoRA + loop embed, Deep Delta vdim1 residuals] -> coda x1
6
  -> RMSNorm -> tied unembed. Attention is MHA by default, GQA when
7
- `n_kv_heads` < `n_heads` (2026-07-05 overhaul knobs, ADR-0009).
8
  """
9
 
10
  import math
@@ -20,7 +19,7 @@ LOOP_EMBED_STD = 0.02
20
 
21
  @dataclass
22
  class MeiosisConfig:
23
- # defaults = release shape per ADR-0009 (B'-GQA overhaul, 2026-07-05):
24
  # 3-block GQA body x3 loops, vocab 4096, ~5.76M total under the <6M cap
25
  vocab_size: int = 4096
26
  dim: int = 288
@@ -37,12 +36,12 @@ class MeiosisConfig:
37
  max_seq_len: int = 512
38
  ddl_beta_init: float = 1.0
39
  # rms_norm backward amplifies grads by 1/sqrt(eps_rms) when k_in ~ 0 — which is
40
- # exactly the zero-init state. 1e-5 gave a 1.7e6x amplifier (1e5-magnitude grad
41
- # spikes; > fp16 max at ANY loss scale — the 2026-07-06 fp16 divergence, ADR-0012).
42
- # 1e-2 caps it at 1.7e3: fp16-safe, and identical bf16 training curves.
43
  ddl_k_eps: float = 1e-2
44
  ddl_v_sigmoid_scale: float = 4.0
45
- # intra-document attention (ADR-0019): tokens attend only within their own
46
  # EOS-delimited document. None = plain causal (pre-mask checkpoints).
47
  doc_mask_eos: int | None = 2
48
 
@@ -79,7 +78,7 @@ def apply_rope(x: Tensor, cos: Tensor, sin: Tensor) -> Tensor:
79
  def build_doc_mask(tokens: Tensor, eos_id: int) -> Tensor:
80
  """(B,T) tokens -> (B,1,T,T) bool, True where attention is allowed:
81
  causal AND same document. Exclusive EOS scan, so an EOS token is the
82
- last token of its document (FSX-1 convention)."""
83
  is_eos = tokens == eos_id
84
  doc_id = torch.cumsum(is_eos, dim=1) - is_eos.to(torch.long)
85
  same = doc_id.unsqueeze(2) == doc_id.unsqueeze(1)
@@ -301,7 +300,7 @@ class Meiosis(nn.Module):
301
 
302
 
303
  def init_meiosis(model: Meiosis) -> None:
304
- """Mandatory MythosMini-validated init. Never mu-center the tied embedding."""
305
  with torch.no_grad():
306
  model.embed.weight.normal_(mean=0.0, std=EMBED_STD)
307
  model.loop_embed.weight.normal_(mean=0.0, std=LOOP_EMBED_STD)
@@ -315,7 +314,7 @@ def count_parameters(model: nn.Module) -> int:
315
 
316
 
317
  def muon_param_split(model: Meiosis) -> tuple[list[nn.Parameter], list[nn.Parameter]]:
318
- """Explicit Muon/aux split (ADR-0005). Muon gets the block and LoRA
319
  matrices; the tied embedding, loop embeddings, norm gains, and 1-row DDL
320
  heads stay on NAdamW. Listed explicitly - no shape heuristics, so a
321
  rank-8 pilot LoRA cannot silently fall out of the Muon group.
 
1
+ """Tied-embedding looped decoder-only language model.
2
 
 
3
  embed -> prelude x1 -> [body of `body_blocks` distinct blocks xK loops,
4
  per-loop LoRA + loop embed, Deep Delta vdim1 residuals] -> coda x1
5
  -> RMSNorm -> tied unembed. Attention is MHA by default, GQA when
6
+ `n_kv_heads` < `n_heads`.
7
  """
8
 
9
  import math
 
19
 
20
  @dataclass
21
  class MeiosisConfig:
22
+ # defaults = min-spark release shape:
23
  # 3-block GQA body x3 loops, vocab 4096, ~5.76M total under the <6M cap
24
  vocab_size: int = 4096
25
  dim: int = 288
 
36
  max_seq_len: int = 512
37
  ddl_beta_init: float = 1.0
38
  # rms_norm backward amplifies grads by 1/sqrt(eps_rms) when k_in ~ 0 — which is
39
+ # exactly the zero-init state. 1e-5 gives a 1.7e6x amplifier (1e5-magnitude grad
40
+ # spikes, above fp16 max at any loss scale).
41
+ # 1e-2 caps it at 1.7e3: fp16-safe, with identical bf16 training curves.
42
  ddl_k_eps: float = 1e-2
43
  ddl_v_sigmoid_scale: float = 4.0
44
+ # intra-document attention: tokens attend only within their own
45
  # EOS-delimited document. None = plain causal (pre-mask checkpoints).
46
  doc_mask_eos: int | None = 2
47
 
 
78
  def build_doc_mask(tokens: Tensor, eos_id: int) -> Tensor:
79
  """(B,T) tokens -> (B,1,T,T) bool, True where attention is allowed:
80
  causal AND same document. Exclusive EOS scan, so an EOS token is the
81
+ last token of its document."""
82
  is_eos = tokens == eos_id
83
  doc_id = torch.cumsum(is_eos, dim=1) - is_eos.to(torch.long)
84
  same = doc_id.unsqueeze(2) == doc_id.unsqueeze(1)
 
300
 
301
 
302
  def init_meiosis(model: Meiosis) -> None:
303
+ """Mandatory init. Never mu-center the tied embedding."""
304
  with torch.no_grad():
305
  model.embed.weight.normal_(mean=0.0, std=EMBED_STD)
306
  model.loop_embed.weight.normal_(mean=0.0, std=LOOP_EMBED_STD)
 
314
 
315
 
316
  def muon_param_split(model: Meiosis) -> tuple[list[nn.Parameter], list[nn.Parameter]]:
317
+ """Explicit Muon/aux split. Muon gets the block and LoRA
318
  matrices; the tied embedding, loop embeddings, norm gains, and 1-row DDL
319
  heads stay on NAdamW. Listed explicitly - no shape heuristics, so a
320
  rank-8 pilot LoRA cannot silently fall out of the Muon group.