Oguz-MoE-354M-Base

Oğuz-MoE-354M-Base

Oğuz-MoE-354M-Base is a Turkish mixture-of-experts (MoE) causal language model with 354.2M total parameters and ~205.5M active parameters per token. It was produced by a four-stage progressive-scaling pipeline — from-scratch dense pretraining, two rounds of depth upscaling, and a final FFN→MoE upcycle — executed entirely on a single consumer-grade GPU (NVIDIA GeForce RTX 4060, 8 GB VRAM).

The model is a base (pretrained-only) checkpoint: it completes text, it is not instruction-tuned.

⚠️ This model uses a custom architecture (MoeLlamaForCausalLM). Load it with trust_remote_code=True.

---

Model Summary

---

Architecture

• Routing: Each layer's FFN is replaced by a MoEBlock containing a linear router (576 → 4) and 4 SwiGLU experts (intermediate size 2048). Per token, the top-2 experts are selected; their softmax-normalized gates sum to 1.0.
• Dispatch: Tokens are flattened across top-k, sorted by expert index, split into contiguous per-expert tensors, processed, and restored via inverse permutation (fully differentiable, autograd-safe).
• Load balancing: Switch-transformer-style auxiliary loss (num_experts × Σ frac_i × prob_i, ideal value 1.0) with coefficient 0.01.
• Upcycling initialization: The trained 1536-wide dense FFN weights were copied exactly into the first 1536 intermediate dimensions of every expert; new gate/up dimensions ~ N(0, 0.01); new down_proj columns ~ N(0, 0.001) so the added capacity starts near-silent (near-function-preserving initialization).

---

Training Pipeline (Lineage)

Total token exposure across all stages: ~2.2B

Stage details

• Stage 0 (pretrain): AdamW (fused), lr 5e-4, weight decay 0.1, cosine schedule, bf16 autocast. Two-phase context curriculum: 80% of tokens at context 512, 20% at context 1024.
• Stages 1–2 (depth healing): Duplicated layers inserted interleaved next to their source layers to preserve activation-distribution continuity; small Gaussian noise (σ = 0.02) for symmetry breaking. Fresh AdamW (lr 5e-4, wd 0.1, cosine) per stage, bf16.
• Stage 3 (MoE upcycle): Optimizer: 8-bit AdamW (bitsandbytes), lr 3e-4, wd 0.1, cosine decay over 9,155 optimizer steps (73,240 microsteps, gradient accumulation 8). Constant 16,384 tokens per optimizer update in both context phases. Gradient clipping 1.0. No gradient checkpointing (aux-loss graph compatibility); VRAM verified empirically instead.

---

Datasets & Tokenizer

• Tokenizer: ytu-ce-cosmos/turkish-gpt2 (50,257 tokens).
• Training corpus: streamed interleaved mixture (seed 4242361):
  • 30% — ytu-ce-cosmos/Cosmos-Turkish-Corpus-v1.0
  • 70% — HuggingFaceFW/fineweb-2 (tur_Latn split)
• Tokenized offline once into a flat uint16 binary (1.6B tokens) and streamed via memory-mapping during all stages.

---

Usage

from transformers import AutoModelForCausalLM, AutoTokenizer
import torch

repo = "Yigit-Karaman/Oguz-MoE-354M-Base"

tokenizer = AutoTokenizer.from_pretrained(repo)
model = AutoModelForCausalLM.from_pretrained(
    repo,
    trust_remote_code=True,      # required: custom MoE architecture
    torch_dtype=torch.bfloat16,
    device_map="auto",
)

prompt = "Türkiye'nin başkenti"
inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
with torch.no_grad():
    out = model.generate(**inputs, max_new_tokens=64, do_sample=True, temperature=0.8, top_p=0.9)
print(tokenizer.decode(out[0], skip_special_tokens=True))

---

Training Monitoring (Stage 3, Final Run)

Metrics below are taken from the actual completed healing run (73,240 microsteps / 9,155 optimizer steps, ~149.996M tokens):

• VRAM preflight (measured with a real optimizer step): 4.42 GB (batch 4 × ctx 512) and 5.15 GB (batch 2 × ctx 1024) on an 8 GB RTX 4060 — no system-memory offloading during training.
• Checkpoint/resume validated in production: the run was interrupted and successfully resumed at microstep 52,824 from a 30-minute checkpoint, including full 8-bit AdamW state restoration.
• Phase transition (512 → 1024 context) occurred at microstep 58,550 with a phase-1 snapshot checkpoint; no loss discontinuity beyond normal microbatch noise.
• Expert load balance (EMA): remained within 23–28% per expert for the entire run; final balance [24%, 27%, 24%, 25%] — no expert collapse at any point.
• Auxiliary (load-balance) loss: 1.00 – 1.03 throughout (ideal ≈ 1.0).
• LM loss behavior: per-microbatch loss (2,048 tokens of heterogeneous web text) fluctuates in the ~1.3–5.4 range. The run shows the classic MoE-upcycling disruption–recovery pattern: loss rises to ~4.6 around step 250 while the router learns, then recovers below the dense baseline by step ~400. Occasional near-zero-loss microbatches (< 0.1) correspond to highly repetitive web content (memorization), not training instability.
• Learning rate: cosine decay from 3e-4 to 2.2e-10 (full schedule completed).
• Throughput: 7,200 – 9,400 tokens/s on a single RTX 4060 (bf16, no gradient checkpointing).

No formal downstream benchmarks have been run yet; this checkpoint is released as a research base model. Evaluation contributions from the community are welcome.

---

Reproducibility

• Global seed: 4242361 (Python / NumPy / PyTorch / CUDA), also used for dataset interleaving.
• Hardware: 1× NVIDIA GeForce RTX 4060 (8 GB), Windows.
• Software: PyTorch (bf16 autocast, SDPA attention), Hugging Face Transformers (custom-code loading), bitsandbytes 0.50.2 (AdamW8bit), Safetensors serialization.
• Data pipeline: offline uint16 token binary (1.6B tokens), memory-mapped streaming with tail+head wraparound (no tokens discarded at epoch boundaries).
• Optimizer (Stage 3): AdamW8bit, lr 3e-4, wd 0.1, cosine over 9,155 steps; gradient accumulation 8; 16,384 tokens per update in both phases; grad-norm clip 1.0.
• Checkpointing: time-based (30 min), saved only immediately after optimizer steps; loaded with map_location="cpu" and moved to GPU via bitsandbytes move_to_device to avoid VRAM spikes.
• MoE specifics: 4 experts / top-2; expert intermediate 2048; router Linear(576→4) init std 0.01; softmax-normalized top-2 gates; Switch-style aux loss coefficient 0.01; near-function-preserving expert initialization (exact copy of dense FFN into first 1536 dims).

---

Limitations & Biases

• Turkish (Latin script) only. Performance on other languages, or on Turkish in non-Latin scripts (e.g., Ottoman Turkish), is expected to be poor.
• Web-text biases. The model inherits stereotypes, errors, and topical biases present in public Turkish web corpora.
• Memorization. Occasional near-zero-loss microbatches during training indicate some memorization of repetitive web content.
• Small model. Prone to factual errors, hallucination, and weak long-range reasoning; effective trained context is 1024 tokens.
• Base model. Not aligned or instruction-tuned; unsuitable for conversational use without further fine-tuning.
• Custom code. Loading requires trust_remote_code=True; review the modeling code (modeling_moe_llama.py) before execution, per standard practice for custom-code repositories.

---

Intended Use

Research and prototyping for Turkish NLP: language-modeling experiments, fine-tuning/adapter studies, and as a worked example of progressive depth expansion + MoE upcycling on consumer hardware. Not intended for production deployment or safety-critical applications.

---

License

This model is released under the Apache 2.0 License (see the LICENSE file in this repository).

As required by the upstream datasets, this model is provided with the following attributions:

• Training data includes the FineWeb-2 dataset, licensed under ODC-BY 1.0.
• Training data includes the Cosmos-Turkish-Corpus-v1.0, courtesy of the YTÜ-Cosmos project.
• Tokenizer and base architecture configuration inspired by ytu-ce-cosmos/turkish-gpt2.

---

Citation

@misc{Oguz-MoE-354M-Base,
  title        = {Oğuz-MoE-354M-Base: A Turkish Mixture-of-Experts Language Model via Progressive Depth Expansion and FFN Upcycling},
  author       = {Yiğit Karaman},
  year         = {2026},
  howpublished = {\url{https://huggingface.co/Yigit-Karaman/Oguz-MoE-354M-Base}}
}

---

Acknowledgements

Tokenizer and part of the training corpus courtesy of the YTÜ-Cosmos project; Turkish web text from HuggingFace FineWeb-2. Depth-upscaling and MoE upcycling techniques follow the public literature on layer duplication and expert upcycling of dense checkpoints. Built with Hugging Face Transformers, bitsandbytes, and PyTorch.