Stage 1 initial: step 1000, loss=0.29, cka=0.60

README.md

@@ -0,0 +1,80 @@
+---
+language:
+  - multilingual
+  - en
+  - es
+  - hi
+  - zh
+  - ar
+  - sw
+  - tr
+  - ja
+  - id
+  - te
+license: apache-2.0
+tags:
+  - mamba
+  - moe
+  - ssm
+  - multilingual
+  - distillation
+  - aya
+library_name: aetheris
+pipeline_tag: text-generation
+---
+
+# Aetheris — Hybrid Mamba-MoE Multilingual Model
+
+**Aetheris** is a ~800M parameter hybrid SSM/MoE language model distilled from
+[CohereLabs/tiny-aya-global](https://huggingface.co/CohereLabs/tiny-aya-global) (3.35B).
+
+Built by [Wayy Research](https://github.com/Wayy-Research).
+
+## Architecture
+
+- **Type**: Hybrid Mamba (SSM) + Mixture of Experts (MoE)
+- **Layers**: 24 (interleaved: even=SSM, odd=MoE)
+- **Hidden dim**: 1024
+- **Experts**: 4 per MoE layer, top-1 routing
+- **SSM state dim**: 16
+- **Vocab size**: 256,000 (shared with tiny-aya-global)
+- **Parameters**: ~800M
+
+## Training
+
+3-stage MambaInLlama distillation pipeline:
+
+| Stage | Method | Data | Steps |
+|-------|--------|------|-------|
+| 1 | CKA-guided Layer Alignment | ClimbMix | 10,000 |
+| 2 | KL Distillation (T=2.0, alpha=0.7) | ClimbMix | 20,000 |
+| 3 | Supervised Fine-Tuning | aya_collection | 5,000 |
+
+Key research findings applied:
+- SSM 10x LR boost (compensates 27x gradient imbalance)
+- SVD split for MoE expert initialization (CKA=0.097 diversity)
+- Per-language KL tracking for multilingual equity
+
+## Current Checkpoint
+
+- **Stage**: 1 (layer-alignment)
+- **Step**: 1000
+- **Loss**: 0.4459
+- **Updated**: 2026-03-12T21:27:28.377034+00:00
+
+## Languages
+
+Supports 70+ languages inherited from tiny-aya-global. Core evaluation
+languages: English, Spanish, Hindi, Chinese, Arabic, Swahili, Turkish,
+Japanese, Indonesian, Telugu.
+
+## Citation
+
+```bibtex
+@misc{aetheris2026,
+  title={Aetheris: Hybrid Mamba-MoE Multilingual Model via Knowledge Distillation},
+  author={Wayy Research},
+  year={2026},
+  url={https://huggingface.co/wayyresearch/aetheris}
+}
+```

stage1_checkpoint.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:2f58f5ce1d22f5e3a4144c29157a695ff88e9a1e07767a52b59a20d774560fa0
+size 1644712894

stage1_metadata.json

@@ -0,0 +1,21 @@
+{
+  "step": 1000,
+  "checkpoint_file": "best.pt",
+  "timestamp": "2026-03-12T21:27:28.377034+00:00",
+  "cka_history_keys": [
+    "0_0",
+    "1_1",
+    "3_2",
+    "4_3",
+    "6_4"
+  ],
+  "final_loss": 0.4459029510617256,
+  "best_loss": 0.4196140021085739,
+  "training_config": {
+    "loss_type": "mse+cosine",
+    "cka_threshold": 0.75,
+    "seed": 42
+  },
+  "stage": 1,
+  "stage_name": "layer-alignment"
+}

student_config.yaml

@@ -0,0 +1,31 @@
+## Aetheris Student Model Configuration
+## Target: ~500-800M parameters (HybridMambaMoE)
+##
+## Architecture: 24 layers alternating SSM (even) and MoE (odd)
+## Vocab sized to match Aya tokenizer (256k)
+##
+## Wayy Research, 2024-2026
+
+vocab_size: 256000
+d_model: 1024
+n_layer: 24
+num_experts: 4
+top_k: 1
+d_ff: 3072  # d_model * 3
+
+# SSM parameters
+ssm_d_state: 16
+ssm_expand: 2
+# d_inner: null  # defaults to d_model * ssm_expand = 2048
+
+# Training parameters
+load_balancing_coef: 0.01
+router_z_loss_coef: 0.001
+max_seq_len: 2048
+dtype: "float16"
+
+# Optimization
+use_cpu_offload: false
+gradient_checkpointing: true
+checkpoint_ssm_layers: true
+use_flash_attention: false

training_config.yaml

@@ -0,0 +1,108 @@
+## Production Distillation Config for RunPod (RTX A6000 48GB)
+##
+## 3-stage MambaInLlama pipeline with research-validated hyperparameters:
+##   - SSM 10x LR boost (NB11)
+##   - T=2.0 for KL distillation (NB09)
+##   - alpha=0.7 for KL/CE balance (NB09)
+##   - SVD split for MoE diversity (NB03)
+##
+## Teacher: CohereLabs/tiny-aya-global (3.35B, validated in NB01-NB11)
+## Student: Aetheris HybridMambaMoE (~800M params)
+## Data: ClimbMix (Stage 1-2), multilingual chat (Stage 3)
+##
+## Wayy Research, 2024-2026
+
+# Teacher model (tiny-aya-global: 3.35B, 70+ langs, no gated access required)
+teacher:
+  name: "CohereLabs/tiny-aya-global"
+  dtype: "bfloat16"
+  device_map: "auto"
+
+# Student model
+student:
+  config_path: "configs/student.yaml"
+  dtype: "bfloat16"
+  checkpoint: null  # Set to Stage 1 checkpoint for Stage 2, etc.
+
+# Languages (10 core for multilingual equity tracking)
+languages: [en, es, hi, zh, ar, sw, tr, ja, id, te]
+
+# Seed
+seed: 42
+
+# --- Stage 0: Block Conversion ---
+conversion:
+  strategy: "weight_map"
+  a_init: "exponential_decay"
+  delta_init: "uniform"
+  ffn_to_moe: "svd_split"    # Best diversity (CKA=0.097 vs replicate=0.88)
+
+# --- Stage 1: Layer Alignment ---
+stage1:
+  enabled: true
+  total_steps: 10000
+  lr: 1.0e-4
+  warmup_steps: 500
+  batch_size: 4
+  gradient_accumulation: 8
+  gradient_checkpointing: true
+  max_seq_len: 512
+  loss_type: "mse+cosine"
+  cka_threshold: 0.75
+  cka_check_every: 500
+  save_every: 1000
+  log_every: 50
+  output_dir: "checkpoints/stage1_alignment"
+
+# --- Stage 2: KL Distillation ---
+stage2:
+  enabled: true
+  total_steps: 20000
+  lr: 5.0e-5                   # Base LR
+  ssm_lr_multiplier: 10.0      # SSM blocks get 10x (NB11: KL -26%, agreement +12x)
+  warmup_steps: 500
+  batch_size: 4
+  gradient_accumulation: 8
+  gradient_checkpointing: true
+  max_seq_len: 512
+  temperature: 2.0              # NB09: T=2.0 good balance
+  alpha: 0.7                    # NB09: alpha=0.7
+  save_every: 2000
+  log_every: 50
+  output_dir: "checkpoints/stage2_kl"
+
+# --- Stage 3: SFT ---
+stage3:
+  enabled: true
+  total_steps: 5000
+  lr: 2.0e-5
+  warmup_steps: 200
+  batch_size: 4
+  gradient_accumulation: 4
+  gradient_checkpointing: true
+  max_seq_len: 1024
+  save_every: 500
+  log_every: 25
+  output_dir: "checkpoints/stage3_sft"
+
+# --- Data ---
+data:
+  # Stage 1 & 2: ClimbMix (retokenized with Aya vocab)
+  climbmix:
+    dataset: "nvidia/ClimbMix"
+    mode: "retokenize"
+    streaming: true
+    buffer_size: 500
+    min_tokens: 32
+
+  # Stage 3: Multilingual chat data (aya_collection is non-gated)
+  sft:
+    dataset_name: "CohereForAI/aya_collection"
+    streaming: true
+
+# --- Evaluation ---
+eval:
+  max_new_tokens: 128
+  temperature: 0.7
+  top_p: 0.9
+  output_dir: "results/runpod"