File size: 4,567 Bytes
9239f91
 
 
 
9941aae
9239f91
 
 
 
 
 
 
 
 
9941aae
9239f91
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
9941aae
9239f91
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
---
language: en
library_name: transformers
license: apache-2.0
model_name: project-828-gpt-base
pipeline_tag: text-generation
tags:
- code
- moe
- mixture-of-experts
- flash-attention
- custom-architecture
---

# project-828-gpt-base

A **398.7M parameter Mixture-of-Experts (MoE)** causal language model optimized for code generation and technical reasoning, pretrained on ~60B tokens.

> **398.7M total parameters, ~286M active per token** β€” 4 routed experts + 1 shared expert with top-2 routing per token.

## Model Architecture

| Property | Value |
|:---|:---|
| **Architecture** | Decoder-only Transformer with MoE FFN layers |
| **Total Parameters** | **398.7M** |
| **Active Parameters / Token** | **~286M** |
| **Hidden Dimension ($d_{\text{model}}$)** | 768 |
| **Intermediate Size ($d_{\text{ff}}$)** | 760 |
| **Hidden Layers** | 24 |
| **Attention Heads / KV Heads** | 12 / 6 (GQA 2:1) |
| **Head Dimension** | 64 |
| **Routed Experts** | 4 |
| **Active Experts / Token** | 2 (Top-2) |
| **Shared Experts** | 1 |
| **Context Length** | 2048 (extensible to 8192 via YaRN) |
| **Vocabulary Size** | 49152 |
| **Precision** | BFloat16 Mixed Precision |
| **Positional Encoding** | RoPE with YaRN scaling support |

### Parameter Breakdown

| Component | Parameters |
|:---|---:|
| Embeddings | ~37.7M |
| Unembedding | ~37.8M |
| Attention (Γ—24 layers) | ~1.8M each |
| MoE FFN (Γ—24 layers) | ~11.7M each |
| Layer Norms + Misc | ~0.07M |
| **Total** | **398.7M** |

### Key Architectural Features

- **Grouped Query Attention (GQA)**: 12 query heads with 6 key-value heads (2:1 ratio) for memory-efficient attention
- **QK-Norm**: RMSNorm applied to query and key projections before RoPE for attention stability
- **Auxiliary-Loss-Free MoE Routing**: Sigmoid gating with dynamic bias adjustment ([DeepSeek-V3 paper](https://arxiv.org/abs/2408.15664)) β€” achieves near-perfect ~25% utilization per expert without auxiliary losses
- **SwiGLU Activation with Soft-Clamping**: `limit=7.0` prevents activation explosions during long training runs
- **Batched Expert Dispatch**: Sort-and-slice dispatch with `searchsorted` boundaries for contiguous memory access
- **RoPE with YaRN Extension**: Base context of 2048 tokens, extensible to 8192 via YaRN scaling

## Training Details

### Phase 1 β€” Pretraining (~60B tokens)

| Training Config | Value |
|:---|:---|
| **Hardware** | H200 GPU |
| **Peak Learning Rate** | 3e-4 |
| **Min Learning Rate** | 3e-5 |
| **Scheduler** | WSD (Warmup-Stable-Decay) |
| **Warmup Steps** | 500 |
| **Total Steps** | 101,726 |
| **Effective Batch Size** | 37 Γ— 8 = 296 sequences |
| **Tokens per Step** | ~0.61M |
| **Gradient Clipping** | 1.0 |

### Training Data Mix

| Dataset | Weight | Category |
|:---|:---:|:---|
| `starcoderdata` β€” Python | 14 | Source Code |
| `starcoderdata` β€” JavaScript | 8 | Source Code |
| `starcoderdata` β€” Java | 6 | Source Code |
| `starcoderdata` β€” TypeScript | 4 | Source Code |
| `starcoderdata` β€” C++ | 6 | Source Code |
| `starcoderdata` β€” C | 4 | Source Code |
| `starcoderdata` β€” C# | 3 | Source Code |
| `starcoderdata` β€” Go | 4 | Source Code |
| `starcoderdata` β€” Rust | 3 | Source Code |
| `starcoderdata` β€” PHP | 3 | Source Code |
| `fineweb-edu-dedup` | 20 | General Knowledge |
| `cosmopedia-v2` | 7 | General Knowledge |
| `wikipedia-en` | 3 | General Knowledge |
| `finemath-4plus` | 8 | Math / Reasoning |
| `stackexchange` (programming/CS) | 7 | CS / Engineering |

**Category breakdown**: Source Code 55% Β· General Knowledge 30% Β· Math/Reasoning 8% Β· CS/Engineering 7%

## Quickstart

```python
from transformers import AutoModelForCausalLM, AutoTokenizer

repo_id = "AkshithAI/project-828-gpt-base"

tokenizer = AutoTokenizer.from_pretrained(repo_id, trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained(
    repo_id,
    torch_dtype="bfloat16",
    device_map="auto",
    trust_remote_code=True
)

prompt = "def binary_search(arr, target):"
inputs = tokenizer(prompt, return_tensors="pt").to(model.device)

outputs = model.generate(
    **inputs,
    max_new_tokens=200,
    temperature=0.7,
    top_p=0.9,
    do_sample=True,
    use_cache=False,
    eos_token_id=tokenizer.eos_token_id,
    pad_token_id=tokenizer.pad_token_id
)
print(tokenizer.decode(outputs[0], skip_special_tokens=True))
```

## Citation

```bibtex
@misc{project828,
  author = {AkshithAI},
  title = {Project 828: MoE Transformer with Training Pipeline},
  year = {2025},
  publisher = {GitHub},
  url = {https://github.com/AkshithAI/project_828}
}
```