Apex-2-Base

The base checkpoint of a Mixture-of-Experts LLM pretrained from scratch: 3.87B total / 1.45B active parameters, trained on 86.5B tokens. Released as a starting point for continued pretraining (CPT) and fine-tuning.

[!TIP] Ready for continued pretraining

• Standard Qwen3MoeForCausalLM format. It loads and trains in transformers, vLLM and other tools that support Qwen3-MoE, with no trust_remote_code.
• Full-precision fp32 weights, exactly as in the training checkpoint (no bf16 rounding). Checked against the training checkpoint: logits relative difference 9.3e-7, 100% next-token agreement.
• All training settings are public: optimizer, LR schedule and the exact data mix. See the Continued pretraining guide below.
• A pre-decay checkpoint at peak LR (stable-44500) is also available for WSD-style continuation.

• Chat (SFT) version: YOON1v/Apex-2
• Code, architecture writeup and training diary: github.com/DW-dev-UE/LLM-from-scratch

Model details

Tokens trained per dataset (86.5B total, click to expand)

In billions of tokens, computed from each stage's mixture weights (see Training data).

Usage (inference)

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained("YOON1v/Apex-2-Base")
model = AutoModelForCausalLM.from_pretrained("YOON1v/Apex-2-Base", dtype=torch.bfloat16, device_map="cuda")

inputs = tokenizer("def is_prime(n: int) -> bool:\n", return_tensors="pt").to("cuda")
out = model.generate(**inputs, max_new_tokens=128, do_sample=False)
print(tokenizer.decode(out[0], skip_special_tokens=True))

vLLM: LLM(model="YOON1v/Apex-2-Base", dtype="bfloat16", max_model_len=4096)

This is a base model: it continues text and does not follow instructions or chat. For chat, use Apex-2.

Continued pretraining guide

Loading

import torch
from transformers import AutoModelForCausalLM

model = AutoModelForCausalLM.from_pretrained(
    "YOON1v/Apex-2-Base",
    dtype=torch.float32,        # keep the fp32 weights; train with bf16 autocast (mixed precision)
    output_router_logits=True,  # adds the expert load-balancing loss (off by default)
)

• With output_router_logits=True, the load-balancing loss (router_aux_loss_coef=0.01) is added to the loss whenever labels are passed. Without it, routing can collapse onto a few experts.
• The router z-loss (1e-3) used in pretraining is not part of the transformers implementation. Add it yourself if you need it.

Original training setup

Recommendations

• Re-warm the LR. This checkpoint has finished its decay (LR 0). Start CPT with a warmup and tune the peak LR at or below the original 4e-4.
• Replay old data. Training only on new data can erode existing skills, code especially. Mixing in some of the original data (listed below) reduces forgetting.
• Packing. If you pack several documents into one sequence, block attention across document boundaries (e.g. transformers DataCollatorWithFlattening + FlashAttention).
• Context length. Trained at 4,096 tokens. Going longer needs changes such as a larger rope_theta.
• Tokenizer. Use the default settings. transformers may suggest fix_mistral_regex=True; do not enable it, because it changes tokenization from what training used.
• Special tokens. The embeddings of <|im_start|>, <|im_end|> and the other special tokens were barely trained during pretraining. Chat formatting is learned in SFT.

Intermediate checkpoint

For WSD-style continuation, stable-44500 is the natural starting point: keep training on new data at the peak LR, then run your own decay whenever you choose.

model = AutoModelForCausalLM.from_pretrained("YOON1v/Apex-2-Base", revision="stable-44500", dtype=torch.float32)

Benchmarks

Measured by us with lm-eval 0.4.13 (HF, bf16) and EvalPlus 0.3.1 (vLLM, base prompt, greedy).

Compared with similar-size base models (official numbers; protocols differ, so this is a rough comparison):

With about 1/200 of the pretraining data, Apex-2-Base matches Qwen2.5-1.5B (18T tokens) on HumanEval+ (32.9). Knowledge (MMLU) and math still show the gap in pretraining scale.

Training data

Public datasets only. 86.5B tokens in three stages:

Tokens trained per dataset: see the expandable table at the top of Model details.

StarCoderData by language (18.65B total)

• Token counts are computed from each stage's mixture weights. Every 4,096-token sequence was drawn from a source at random according to those weights, so actual amounts can differ by about 1%.
• Documents that also appear in the stage 1 data were removed from the later corpus, and no source was used for more than one epoch, so no data was repeated.

Limitations

• English-centric. Other languages, including Korean, barely work.
• Limited world knowledge (MMLU ~28%).
• As a base model it has no instruction following, chat ability or safety tuning.
• 4,096-token context.

License

Apache 2.0. Each training dataset has its own license and terms (ODC-By, CC-BY, CC-BY-SA, the StarCoderData terms of use, NVIDIA data licenses, etc.); see the dataset pages.