YUA-MoE-9.45B

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

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

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. Pre-training on raw multilingual text (~126B tokens)
2. SFT on 5.8M instruction-response pairs (25+ sources)
3. GRPO/DPO alignment (planned)

Roadmap

• 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