CoolFace
Modelpublic

jungwon-ai/YUA-MoE-9.45B

sourceHugging Faceapache-2.0updated 6mo agoView on Hugging Face
0likes14downloads
Model Card

YUA-MoE-9.45B

Multilingual MoE language model, pre-trained from scratch on TPU.

9.45B total parameters with 8 experts (top-2 gating), giving ~1.93B active parameters per forward pass. Trained on Korean, English, Japanese, Chinese, and code data using the MaxText framework on TPU v4-32.

Note: Pre-training is still in progress. Weights will be uploaded after training completes. The model code (configuration + architecture) is available now.

GitHub: https://github.com/yuaone/yua

Model Details

Parameters9.45B total, ~1.93B active
ArchitectureTransformer MoE (8 experts, top-2)
Hidden dim2048
Layers32
AttentionGQA (32 heads, 8 KV heads), RoPE, QK-norm
FFNSwiGLU, 5461 per expert
Vocab128K (custom SentencePiece, multilingual)
Context4096 (YaRN extension to 256K planned)
Training data~126B tokens (multilingual + code)
Precisionbfloat16

The model was initialized via Drop-Upcycling (ICLR 2025) from a 1.93B dense checkpoint, converting the single FFN into 8 specialized experts.

How to Use

python
from transformers import AutoModelForCausalLM
import torch

model = AutoModelForCausalLM.from_pretrained(
    "jungwon-ai/YUA-MoE-9.45B",
    trust_remote_code=True,
    torch_dtype=torch.bfloat16,
    device_map="auto",
)

Tokenizer is SentencePiece-based (128K vocab). Will be included with weights.

Architecture Notes

  • —Routing: Standard top-2 softmax gating. Future versions will use ReLU routing (ReMoE).
  • —QK-norm: Applied to query and key projections for training stability.
  • —No shared expert in current version. Planned for Gen2.
  • —Untied embeddings: Input and output embeddings are separate.

Training

Trained using MaxText (JAX) on Google Cloud TPU v4-32 through the TRC program. The training pipeline includes:

  1. 1.Pre-training on raw multilingual text (~126B tokens)
  2. 2.SFT on 5.8M instruction-response pairs (25+ sources)
  3. 3.GRPO/DPO alignment (planned)

Roadmap

  • —[x] Architecture design + HF model class
  • —[ ] Pre-training completion (~May 2026)
  • —[ ] Context extension (YaRN, 4K → 256K)
  • —[ ] SFT + alignment
  • —[ ] Weight upload to this repo
  • —[ ] vLLM serving support
  • —[ ] Gen2: d=10240, shared experts, ReLU routing

Citation

If you use this model, please cite:

@misc{yua2026,
  title={YUA-MoE-9.45B: A Multilingual Mixture-of-Experts Language Model},
  author={Jungwon Eom},
  year={2026},
  url={https://huggingface.co/jungwon-ai/YUA-MoE-9.45B}
}

License

Apache 2.0