Apex-2.5-Base

The pretrained model behind Apex-2.5: a test release of the new Apex-3 architecture with a SuperBPE tokenizer, ready for your continued pretraining or fine-tuning.

Thank you for your interest in this model.

Apex-2.5-Base is a mixture-of-experts language model with 21B total and 3.36B active parameters, trained from scratch with no external pretrained weights. It is a test version of a new architecture that also introduces a SuperBPE tokenizer. It was trained on 40B tokens and stopped at the end of a learning-rate decay, so you can continue training it on your own data. The guide below explains how.

It is pure next-token prediction, so give it text to continue rather than chat messages.

Here is what Apex-2.5-Base offers:

• A new architecture. The routed experts are fine-grained: 64 per layer, 4 active. Each token also passes through a shared expert. Attention alternates five sliding-window layers with one global layer, and uses grouped-query attention with QK-norm. A sigmoid router balances the experts mainly with a selection bias updated outside the gradient, helped by a small sequence-level balance loss.
• SuperBPE. Tokens may span several words, so the same text needs about 21% fewer tokens than with an ordinary BPE vocabulary of the same size. The 40B training tokens hold roughly as much text as 51B BPE tokens.
• Pretraining. Three stages, 8,502 optimizer steps in all. 4.0B tokens on one NVIDIA GH200; a 31.0B-token continuation across two GH200s through DiLoCo; then 5.0B tokens of linear learning-rate decay toward code and math. The global batch grew from 2M to 8M tokens. Optimization used Muon and AdamW with FP32 master weights. From 4B tokens on, the attention, shared-expert and dense-layer linears ran in FP8, about half of the compute.
• Ready to continue. The checkpoint ends right after the decay. See Continue training for data format, token budgets, learning rates and hardware.

Evaluation

Measured by us after pretraining: zero-shot log-likelihood on the full evaluation sets, no chat template. acc_norm divides each answer's log-likelihood by its length in bytes.

Bits per byte of the reference answer (lower is better), no generation: GSM8K 0.667, HumanEval 0.264, TriviaQA 1.155.

WinoGrande uses partial scoring: each option fills the blank, and the shared continuation is scored. LAMBADA counts an example as correct when every token of the last word is the greedy prediction. SuperBPE tokens can cross word boundaries, so treat LAMBADA as a check between checkpoints of this model rather than a number to compare with other tokenizers.

Documents overlapping these evaluation sets were removed from the training data by 13-gram matching. This removal is a lower bound, and the FLAN data in the decay stage may contain the training splits of ARC and HellaSwag.

Comparison with other open base models

Scores of the other models are copied from their published papers and model cards, not re-measured by us. Shots, normalization and evaluation tools differ between rows, so read this table as rough context rather than a ranking. Apex-2.5-Base uses acc_norm (byte-normalized) for HellaSwag, ARC and PIQA, and acc for WinoGrande.

¹ TinyLlama paper (arXiv 2401.02385, Table 3; EVAL.md), acc_norm, lm-evaluation-harness. ² OLMo paper (arXiv 2402.00838, Table 3 and 4.1): per-token normalization for HellaSwag, PIQA and WinoGrande, unconditional normalization for ARC. Training tokens from the model card; the paper's Table 1 lists 2T. ³ OLMoE paper (arXiv 2409.02060, Tables 4 and 12). 5-shot OLMES; ARC-C and PIQA use the multiple-choice format, which scores higher than cloze. Not comparable with the 0-shot rows.

Quick start

from transformers import AutoModelForCausalLM, AutoTokenizer

model_id = "YOON1v/Apex-2.5-Preview-Base"
tok = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(
    model_id, trust_remote_code=True, torch_dtype="auto", device_map="auto")

# A base model continues text: give it a prefix, not a chat message
ids = tok("def is_prime(n):", return_tensors="pt").to(model.device)
out = model.generate(**ids, max_new_tokens=128)
print(tok.decode(out[0], skip_special_tokens=True))

• trust_remote_code=True is required because this is a custom architecture.
• The BF16 weights take 39.1 GiB, so use a 48 GB or larger GPU. With device_map="auto", layers that do not fit are offloaded to CPU memory, which is slow.
• The Transformers code in this repository is a simple reference implementation, not an optimized one: it builds a dense attention mask and runs the experts in a loop. The routers are kept in FP32, as in training.
• vLLM and SGLang cannot load this architecture yet.

Continue training

1. What carries over from pretraining

Special tokens:

2. Prepare your data

Tokenize every document with the bundled tokenizer and add <|endoftext|> after it. Then concatenate the documents and cut them into rows of 8,192 tokens.

from datasets import load_dataset
from transformers import AutoTokenizer

tok = AutoTokenizer.from_pretrained("YOON1v/Apex-2.5-Preview-Base")
EOT, SEQ = 203776, 8192   # <|endoftext|>, training length

# sample_by="document": each file is one document
ds = load_dataset("text", data_files={"train": "my_corpus/*.txt"}, sample_by="document")["train"]

def tokenize(batch):
    # split_special_tokens: a literal "<|endoftext|>" in your text stays text, as in pretraining
    ids = tok(batch["text"], add_special_tokens=False, split_special_tokens=True)["input_ids"]
    return {"input_ids": [x + [EOT] for x in ids]}

def pack(batch):
    flat = [t for ids in batch["input_ids"] for t in ids]
    n = len(flat) // SEQ * SEQ
    rows = [flat[i:i + SEQ] for i in range(0, n, SEQ)]
    return {"input_ids": rows, "labels": [r[:] for r in rows]}

ds = ds.map(tokenize, batched=True, remove_columns=ds.column_names, num_proc=16)
ds = ds.map(pack, batched=True, remove_columns=ds.column_names, num_proc=16)

• SuperBPE tokenization is slower than ordinary BPE, so tokenize once and save the result (ds.save_to_disk).
• To mirror pretraining more closely, reset position_ids after each <|endoftext|>; the model accepts position_ids. Attention can still see earlier documents in the same row. For short documents you can skip packing and pad with <|pad|> instead, labelling the padding −100.
• To keep the model's general ability while you specialize it, mix some general text (web, code) into your data.

3. Turn a token budget into steps

steps = training tokens ÷ (8,192 × sequences per optimizer step)

Sequences per step = per-GPU batch × gradient accumulation × number of GPUs. Pretraining used 8M-token batches; for continued training, 0.5M to 2M tokens is a reasonable range.

4. Learning rate and optimizer

The values below are starting points, not tuned results.

• AdamW settings: β (0.9, 0.95), ε 1e-8, weight decay 0.1, no weight decay on norms and routers, gradient clipping 1.0. The Trainer already skips norms by name; the script in section 6 also skips the routers.
• Pretraining used Muon for the weight matrices. AdamW works too; a Muon implementation can stay closer to how the model was trained.
• Keep expert_bias frozen, as it already is. For runs longer than a few billion tokens, watch the expert load. If it drifts, update the bias outside the gradient after every step: bias += 1e-3 * clip((mean_load - load) / mean_load, -1, 1).

5. Hardware and memory

Without gradient checkpointing, an 8,192-token sequence keeps tens of GB of activations in the reference code, so turn it on (gradient_checkpointing=True). With LoRA on Transformers 4.x, also call model.enable_input_require_grads() or pass gradient_checkpointing_kwargs={"use_reentrant": False}.

• For LoRA, use target_modules=["wq", "wk", "wv", "wo", "w_gate", "w_up", "w_down"]. This adapts attention, the dense layer and the shared experts. The 64 routed experts are stored as stacked parameters, not nn.Linear, so they stay frozen. They hold about 18B of the 21B parameters.
• 4-bit quantization (QLoRA) saves little here, for the same reason: it quantizes nn.Linear layers only.
• The reference code is slow at scale because of the dense mask and the expert loop. For budgets of billions of tokens, the original training code is at https://github.com/DW-dev-UE/LLM-from-scratch. It reads its own checkpoint format, not these safetensors, so loading this release there needs a conversion step that is not included.

6. A minimal training script

import torch
from transformers import AutoModelForCausalLM, Trainer, TrainingArguments, default_data_collator

# load in FP32 (the master copy); bf16=True below runs the compute in BF16, and the routers stay FP32
model = AutoModelForCausalLM.from_pretrained(
    "YOON1v/Apex-2.5-Preview-Base", trust_remote_code=True, torch_dtype=torch.float32)

class ApexTrainer(Trainer):
    # no weight decay on norms and routers; by name, so it also works on FSDP-flattened parameters
    def get_decay_parameter_names(self, model):
        return [n for n, _ in model.named_parameters() if "norm" not in n and ".router." not in n]

args = TrainingArguments(
    output_dir="apex25-continued",
    per_device_train_batch_size=1,
    gradient_accumulation_steps=16,      # x GPUs x 8,192 tokens = tokens per step
    gradient_checkpointing=True,
    learning_rate=1e-4,
    lr_scheduler_type="warmup_stable_decay",
    lr_scheduler_kwargs={"num_decay_steps": 1000, "decay_type": "linear"},
    warmup_ratio=0.02,
    max_steps=9537,                      # from the table in section 3
    weight_decay=0.1, adam_beta1=0.9, adam_beta2=0.95, adam_epsilon=1e-8, max_grad_norm=1.0,
    bf16=True,
    logging_steps=10, save_steps=1000,
    fsdp="full_shard auto_wrap",
    fsdp_config={"transformer_layer_cls_to_wrap": ["Apex3Block"], "cpu_ram_efficient_loading": True},
)
ApexTrainer(model=model, args=args, train_dataset=ds, data_collator=default_data_collator).train()

• The model shifts labels by one position inside forward, so pass labels equal to input_ids. Use −100 for tokens that should not be learned.
• warmup_stable_decay with decay_type needs a recent Transformers release. The Transformers code of this repository was tested with Transformers 5.20.2.
• Do not pass a pre-built optimizer (optimizers=) together with FSDP; the Trainer rejects it.

7. Supervised fine-tuning

The tokenizer includes the chat template apex3-chat-v1:

<|system|>...<|end|><|user|>...<|end|><|assistant|><|think|>...<|/think|>answer<|end|>

• Build conversations with tok.apply_chat_template(messages). The template also breaks special-token strings inside message text, so user text cannot forge control tokens.
• Learn only the assistant turns: set labels to −100 everywhere else.
• <|think|>…<|/think|> is optional. Include it if you want the model to reason before answering.
• At inference, stop at <|end|>. generation_config.json already lists the stop tokens.

Training details

• Batch and length: 8,192-token sequences; the global batch grew from 2M to 8M tokens over the first 33B tokens.
• DiLoCo: the two workers exchanged weight differences every 1B tokens, using an outer Nesterov SGD with learning rate 0.7 and momentum 0.9.
• FP8 precision: tensor-wise FP8 for the attention, shared-expert and dense-layer linears, about 47% of the FLOPs. Routed experts, routers and the output layer stayed in higher precision.
• Memory layout: FP32 master weights, gradient accumulation and optimizer state were kept in the GH200's CPU memory, with BF16 working weights on the GPU.
• Throughput: the decay stage ran at about 23.9k tokens/s, 52% MFU as logged. Counting attention FLOPs per document, which matches the document masking, the MFU is about 1.6 points lower.

Data, stages 1–2 (0 → 35B)

Data, stage 3 (35B → 40B, decay)

• Synthetic data: SwallowCode-v2 and SwallowMath-v2 (rewritten by Qwen3-235B-A22B, Apache-2.0), Cosmopedia v2 (written by Mixtral-8x7B-Instruct), and the reasoning sets (generated by DeepSeek-R1 and QwQ-32B).
• Excluded data: no outputs of closed API models or Llama-licensed models, and no non-commercial data.
• Filtering: code was filtered for permissive licenses; secrets and personal data were removed.
• Decontamination: documents overlapping 24 evaluation sets were dropped by 13-gram matching.

Before you download

Apex-2.5-Base or Apex-2.5-SFT? Apex-2.5-Base continues text and does not follow instructions. Use it for your own continued pretraining or post-training. Apex-2.5-SFT is the instruction-tuned version.

What are its limits?

• It is a test release of a new architecture, trained on far fewer tokens than production base models, so its knowledge is limited.
• It is English-centric, with code and math.
• Context is 8,192 tokens.
• Because SuperBPE tokens can cross word boundaries, results can degrade when a prompt ends in the middle of a word (code completion, fill-in-the-middle, tool arguments). Token healing at serving time helps.

Is it free to use? The weights are free to download. License: Apache 2.0.

License

The model weights and code are released under the Apache License 2.0. Some training sources carry terms of their own, for example CC BY-SA (StackExchange, Wikipedia), ODC-By (FineWeb-Edu, Cosmopedia, peS2o) and the licenses of the models that generated the synthetic data. Review them before choosing a license for derivatives.

---

한국어 안내

저의 모델에 관심을 가져주셔서 감사합니다.

Apex-2.5-Base는 총 21B·활성 3.36B 파라미터의 MoE입니다. 해당 버전은 새로운 아키텍처에 SuperBPE를 사용한 테스트용 버전입니다.

40B 토큰을 사용하여 학습하였으며 이어서 학습을 진행하실 수 있습니다.

이어서 학습할 때 지킬 것:

• 동봉된 SuperBPE 토크나이저를 그대로 씁니다.
• 문서마다 끝에 <|endoftext|>(id 203776)를 붙여 8,192토큰 단위로 묶습니다.
• 학습률이 0 근처까지 내려간 상태라 다시 워밍업합니다(위 4절).
• 토큰 예산은 "학습 토큰 ÷ (8,192 × 스텝당 시퀀스 수)"로 스텝 수를 정합니다(위 3절).