Upload 6 files

README.md

@@ -4,51 +4,101 @@ license: cc-by-nc-4.0
 tags:
   - complexity-deep
   - token-routed
-  - checkpoint
-  - resumable
-  - chinchilla
+  - moe
+  - deterministic-routing
+  - zipf-routing
+  - mu-guidance
+library_name: complexity-framework
+pipeline_tag: text-generation
 ---
 
-# COMPLEXITY-DEEP Token-Routed MoE (383.5M) — Training Checkpoint (Step 15,259)
+# COMPLEXITY-DEEP Token-Routed MoE (383.5M)
 
-Resumable training checkpoint with full optimizer state at the end of 8B tokens training.
+## Model Details
 
-**Note**: This model was trained with a Chinchilla-like token budget (8B tokens for 383.5M parameters, ~21 tokens/param). The model may benefit from continued training beyond this point.
+- **Architecture**: Token-Routed MLP + Mu-Guidance + Shared Lexical Expert
+- **Parameters**: 383.5M total, ~105M active per token
+- **Hidden size**: 1024
+- **Layers**: 20
+- **Attention heads**: 16 (GQA, 4 KV heads)
+- **Intermediate size**: 3200 (800 per expert)
+- **Experts**: 4 (deterministic Zipf-balanced routing)
+- **Shared expert**: 800 intermediate
+- **Vocabulary**: 32,000
+- **Max context**: 4,096
 
-## Contents
+## Training
 
-- `checkpoint.pt` - Model weights + training state
-- `model.safetensors` - Model weights (safetensors format)
-- `optimizer_rank0.pt` - AdamW optimizer state (GPU 0)
-- `optimizer_rank1.pt` - AdamW optimizer state (GPU 1)
-- `training_state.json` - Step counter, LR, etc.
+- **Dataset**: FineWeb-Edu (streaming)
+- **Tokens**: 8B (15,259 steps)
+- **Batch size**: 128 per GPU x 2 GPUs = 256 effective
+- **Optimizer**: AdamW (lr=2.1e-4, auto-scaled to 4.2e-4)
+- **Scheduler**: Cosine with 5% warmup (762 steps)
+- **Precision**: BF16
+- **Hardware**: 2x NVIDIA RTX PRO 6000 (96GB each)
+- **Training time**: ~30 hours
 
-## Model Config
+## Results
 
-- **Parameters**: 383.5M total, ~105M active per token
-- **Hidden**: 1024, Layers: 20, Heads: 16, KV Heads: 4
-- **Experts**: 4, Intermediate: 3200 (800/expert), Shared: 800
-- **Training**: 8B tokens (15,259 steps), AdamW lr=2.1e-4, cosine 5% warmup
+### Loss
+- **Final loss**: ~2.96
+- **vs Dense baseline (384.5M)**: +0.09 gap (stable from step 5K)
+- **Gap trend**: 0.28 (step 1K) -> 0.09 (step 5K+)
 
-## Resume Training
+### Zero-Shot Benchmarks
+| Benchmark | MoE (383.5M) | Dense (384.5M) |
+|-----------|-------------|----------------|
+| ARC-Easy  | 43.6%       | 45.9%          |
+| HellaSwag | 28.7%       | 30.1%          |
+| MMLU      | 23.0%       | 23.1%          |
 
-```python
-import torch
+### Inference (vLLM 0.18, RTX PRO 6000 96GB)
+- **Sustained throughput**: 4,900 tok/s
+- **Peak throughput**: 5,700 tok/s
+- **Median TTFT**: 39.6 ms
+- **Median ITL**: 16.0 ms
+- **CUDA graph**: natively compatible (deterministic routing)
+
+## Expert Specialization (3D t-SNE)
+
+Interactive visualization of expert activations across layers. Each point is an expert at a given layer; proximity = similar activation patterns.
+
+**[▶ Open Interactive 3D t-SNE Visualization](https://www.complexity-ai.fr/expert_tsne_3d.html)**
 
-checkpoint = torch.load("checkpoint.pt", map_location="cpu")
-model.load_state_dict(checkpoint["model"])
+## Generation Example
 
-# Load optimizer for your GPU rank (0 or 1)
-rank = torch.distributed.get_rank()
-optimizer_state = torch.load(f"optimizer_rank{rank}.pt", map_location="cpu")
-optimizer.load_state_dict(optimizer_state)
+No supervised fine-tuning. Raw base model output:
 
-# Resume from step 15,259
+**Prompt:** "The meaning of life is"
+
+**Output:** "very much the same. The same thing happens to all living things. They live in a constant
+state of flux. The single cell of a living cell, in this case a cell nucleus, constantly changes to become
+an organism, and that organism is the organism. The human body is a system of interconnected
+cells. Each cell is made up of a set of parts, which are connected by a network of specialized cells."
+
+## Files
+
+- `model.safetensors` - Model weights
+- `model_config.yaml` - Architecture configuration
+- `config.json` - HuggingFace-compatible config
+
+## Usage
+
+```python
+from complexity.config import ModelConfig
+from complexity.models import ComplexityModel
+from safetensors.torch import load_file
+
+config = ModelConfig.load("model_config.yaml")
+model = ComplexityModel(config)
+state = load_file("model.safetensors", device="cpu")
+model.load_state_dict(state, strict=False)
+model.eval().cuda()
 ```
 
-## Pretrained Weights (inference)
+## Paper
 
-For inference use the safetensors checkpoint in `../final/` instead.
+Under review at TMLR: https://openreview.net/forum?id=jZq6EVboC6
 
 ## License
 

config.json

@@ -0,0 +1,32 @@
+{
+  "model_type": "deep",
+  "architectures": ["DeepForCausalLM"],
+  "hidden_size": 1024,
+  "num_hidden_layers": 20,
+  "intermediate_size": 3200,
+  "vocab_size": 32000,
+  "num_attention_heads": 16,
+  "num_key_value_heads": 4,
+  "attention_type": "gqa",
+  "attention_dropout": 0.0,
+  "use_qk_norm": true,
+  "sliding_window": null,
+  "max_position_embeddings": 4096,
+  "rope_theta": 10000.0,
+  "rope_type": "standard",
+  "mlp_type": "token_routed",
+  "hidden_act": "silu",
+  "num_experts": 4,
+  "shared_expert": true,
+  "shared_intermediate_size": 800,
+  "use_mu_guidance": true,
+  "use_mu_projection": false,
+  "disable_mu_guidance": false,
+  "norm_type": "rmsnorm",
+  "norm_eps": 1e-06,
+  "tie_word_embeddings": true,
+  "use_sdpa": true,
+  "use_cache": true,
+  "initializer_range": 0.02,
+  "extra_config": {}
+}

model_config.yaml

@@ -0,0 +1,28 @@
+attention_dropout: 0.0
+attention_type: gqa
+disable_mu_guidance: false
+extra_config: {}
+hidden_act: silu
+hidden_size: 1024
+initializer_range: 0.02
+intermediate_size: 3200
+max_position_embeddings: 4096
+mlp_type: token_routed
+norm_eps: 1.0e-06
+norm_type: rmsnorm
+num_attention_heads: 16
+num_experts: 4
+num_hidden_layers: 20
+num_key_value_heads: 4
+rope_theta: 10000.0
+rope_type: standard
+shared_expert: true
+shared_intermediate_size: null
+sliding_window: null
+tie_word_embeddings: true
+use_cache: true
+use_mu_guidance: true
+use_mu_projection: false
+use_qk_norm: true
+use_sdpa: true
+vocab_size: 32000

special_tokens_map.json

@@ -0,0 +1,6 @@
+{
+  "bos_token": "<s>",
+  "eos_token": "</s>",
+  "pad_token": "<pad>",
+  "unk_token": "<unk>"
+}

tokenizer_config.json

@@ -0,0 +1,47 @@
+{
+  "add_bos_token": true,
+  "add_eos_token": true,
+  "chat_template": "{% if messages[0]['role'] == 'system' %}{{ messages[0]['content'] }}\n\n{% set ns = namespace(messages=messages[1:]) %}{% else %}{% set ns = namespace(messages=messages) %}{% endif %}{% for message in ns.messages %}{% if message['role'] == 'user' %}User: {{ message['content'] }}\n\n{% elif message['role'] == 'assistant' %}Assistant: {{ message['content'] }}{{ eos_token }}\n{% endif %}{% endfor %}{% if add_generation_prompt %}Assistant:{% endif %}",
+  "added_tokens_decoder": {
+    "0": {
+      "content": "</s>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "1": {
+      "content": "<pad>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "2": {
+      "content": "<s>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "3": {
+      "content": "<unk>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    }
+  },
+  "bos_token": "<s>",
+  "clean_up_tokenization_spaces": false,
+  "eos_token": "</s>",
+  "extra_special_tokens": {},
+  "model_max_length": 2048,
+  "pad_token": "<pad>",
+  "tokenizer_class": "PreTrainedTokenizerFast",
+  "unk_token": "<unk>"
+}