SmolMoE From-Scratch Research Build

TLDR

I built a compact SmolLM-style Mixture-of-Experts pipeline end to end: a custom decoder/MoE backbone, dense-to-MoE sparse upcycling from SmolLM weights, small continued pretraining, and two research extensions. The Cosmopedia track studies whether 1000 carefully selected examples can improve continued pretraining efficiency, while the Nemotron track trains chat/code/math expert specialization with a staged curriculum.

This repository is a from-scratch implementation and experimental study of a small SmolLM-style Mixture-of-Experts language model. It was built for the Cohere Labs Scholars take-home assignment, but the project is organized as a complete engineering and research pipeline rather than as a notebook-only exercise.

The project starts with a custom Transformer/MoE implementation, initializes it from a dense SmolLM checkpoint through sparse upcycling, continues pretraining the upcycled model, and then uses the resulting checkpoint for two open-ended experiments:

1. Cosmopedia dataset intervention: choose better 1000-sample training sets from cosmopedia-100k and measure whether data selection improves continued pretraining efficiency.
2. Nemotron expert specialization: train the existing three MoE experts on the NVIDIA Llama-Nemotron post-training dataset so they specialize on chat, code, and math SFT data while preserving held-out language-modeling quality.

This repository is the cleaned public project root. Large model checkpoints, raw datasets, and generated per-row metric files are intentionally excluded; see ARTIFACTS.md for the expected local artifact locations when reproducing full runs.

High-Level Results

Base MoE Pipeline

I implemented a compact SmolLM-style causal language model with:

The dense SmolLM checkpoint is loaded from:

smolLM_Dense_weights/

The upcycled MoE checkpoint is:

upcycled_smolMoE/upcycled_smolMoE.pt

The continued-pretrained MoE checkpoint used by the later experiments is:

upcycled_smolMoE/continued_pretrained_smolMoE.pt

Cosmopedia Dataset Intervention

Detailed report:

efficient_training/README.md

Best heldout result from the downloaded Kaggle outputs:

The best overall Cosmopedia run reduced token-weighted CE by 20.80% and token-weighted perplexity by 37.84% relative to the CPT baseline.

The key finding: random in-domain training is already a strong baseline, but a diversity-aware, moderately hard selection becomes best after a few epochs. The best stopping point was epoch 3; more epochs over the same 1000 rows overfit.

Nemotron Expert Specialization

Detailed report:

nemotron_specialization/README.md

Final held-out summary from the downloaded Kaggle outputs:

The key finding: a staged curriculum with forced routing, router annealing, and fully learned routing significantly improved both language-modeling loss and domain-router alignment.

Assignment Framing

The assignment had four coding-challenge parts:

1. Implement and validate a custom SmolMoELM.
2. Upcycle a dense SmolLM checkpoint into the MoE architecture.
3. Continue pretraining the upcycled MoE.
4. Explore an open-ended research extension.

This project treats those parts as a single build:

custom MoE architecture
    -> dense checkpoint upcycling
    -> continued pretraining
    -> dataset intervention research
    -> expert specialization research

The original assignment described the first part as a debugging challenge. In this repository, the work is presented as a from-scratch model implementation and validation pipeline: the model architecture, routing behavior, upcycling logic, training loop, metrics, and open-ended experiments are all included as explicit code and artifacts.

Repository Layout

.
  README.md
  ARTIFACTS.md
  assignment.md

  moe/
    backbone.py

  utils/
    config.py
    dataset.py
    load_model.py
    utils.py

  smolLM_Dense_weights/
    ...

  upcycle.py
  train_router.py
  test_moe_backbone.py

  upcycled_smolMoE/
    README.md
    cpt_metrics.png
    gini_spec.png
    cpt.txt
    config.txt

  efficient_training/
    README.md
    ...

  nemotron_specialization/
    README.md
    ...

The two Part 4 experiments are intentionally isolated:

• efficient_training/ contains the Cosmopedia dataset-intervention work.
• nemotron_specialization/ contains the Nemotron expert-specialization work.

This keeps the base model/upcycling pipeline separate from the research extensions.

Part 1: Custom SmolMoE Language Model

The core model is implemented in:

moe/backbone.py

The model is a decoder-only causal language model inspired by SmolLM/Llama-style building blocks:

• token embeddings,
• rotary position embeddings,
• grouped-query self-attention,
• RMSNorm,
• residual decoder blocks,
• MoE feed-forward layers,
• final RMSNorm,
• tied-style causal LM output head shape, implemented as a separate lm_head.

The implementation centers on three classes:

MoE Layer

The MoE layer replaces the dense feed-forward MLP with expert banks:

gate_bank: E x D x H
up_bank:   E x D x H
down_bank: E x H x D

where:

E = number of experts
D = hidden size
H = intermediate size

For each token, the router produces expert logits:

router_logits = gate(hidden_state)

The layer selects the top k experts per token, applies SwiGLU-style expert MLPs, gathers the selected expert outputs, and combines them with the selected router probabilities.

In this assignment configuration:

num_experts = 3
num_experts_per_tok = 1

so each token routes to one expert.

Attention

The attention block uses:

num_heads = 9
kv_heads = 3

This gives grouped-query attention: key/value heads are repeated to match the number of query heads. Rotary embeddings are applied to query and key states.

Load-Balancing Metric

The MoE layer computes a Switch Transformer-style load-balancing auxiliary term. For each layer:

load_i       = fraction of tokens assigned to expert i
importance_i = mean router probability mass for expert i
L_lb         = E * sum_i(load_i * importance_i)

The model exposes:

expert_utilization_per_layer, lb_loss = model.get_expert_utilization()

This is used in the base continued-pretraining loop to monitor whether routing is collapsing or staying balanced.

Validation Script

Basic model validation lives in:

test_moe_backbone.py

It instantiates the model, loads provided trial weights, generates from a test prompt, and checks that the load-balancer loss matches the expected value.

Part 2: Dense-to-MoE Upcycling

Upcycling is implemented in:

upcycle.py

The goal is to initialize the custom MoE from an already useful dense SmolLM checkpoint instead of training the MoE from random weights.

This follows the sparse-upcycling idea from Sparse Upcycling: Training Mixture-of-Experts from Dense Checkpoints: reuse a dense checkpoint, copy useful dense weights into a sparse MoE layout, and continue training from that stronger initialization.

The dense checkpoint is:

smolLM_Dense_weights/

The upcycling process copies:

The important design choice is that every expert initially receives the same dense MLP weights:

expert 0 = dense MLP
expert 1 = dense MLP
expert 2 = dense MLP

That makes the upcycled MoE behave like the dense model at initialization, provided the router chooses one of the identical experts. The router weights are zeroed so the top-1 selection deterministically chooses expert 0 in eval mode.

This gives a stable starting point:

dense model knowledge
    -> copied into shared blocks
    -> copied into all expert MLPs
    -> router starts simple

The saved output is:

upcycled_smolMoE/upcycled_smolMoE.pt

Part 3: Continued Pretraining

Continued pretraining is implemented in:

train_router.py

This stage starts from the upcycled checkpoint:

upcycled_smolMoE/upcycled_smolMoE.pt

and produces:

upcycled_smolMoE/continued_pretrained_smolMoE.pt

Dataset

The base CPT run uses HuggingFaceTB/cosmopedia-100k through:

utils/dataset.py

The dataset utility:

1. loads the dataset from Hugging Face,
2. tokenizes the text column,
3. appends EOS,
4. groups each document into fixed blocks,
5. optionally splits into train/validation sets.

For the base CPT run:

max_samples = 1000
block_size = 256
val_fraction = 0.2
batch_size = 4
steps = 100

This stage is deliberately small. It is meant to adapt the upcycled MoE and exercise the router/expert path, not to perform large-scale pretraining.

Training Objective

The base CPT objective is:

loss = causal_lm_ce + LB_ALPHA * load_balancing_loss

with:

LB_ALPHA = 0.01

The causal LM loss predicts the next token over the whole fixed-length text block. The load-balancing loss nudges expert usage away from pathological collapse.

MoE-Specific Metric

For Part 3, I implemented a Gini-style specialization metric:

specialization = (sum_i p_i^2 - 1/E) / (1 - 1/E)

where p_i is the expert load distribution for a layer and E is the number of experts.

Interpretation:

0%   = uniform expert usage
100% = all tokens routed to a single expert

This is not a "higher is always better" metric. It is a compact way to track whether routing is uniform, specialized, or collapsed. In the later Nemotron experiment, specialization is evaluated more directly with domain-by-expert routing matrices and router accuracy.

Base CPT Artifacts

upcycled_smolMoE/cpt_metrics.png
upcycled_smolMoE/gini_spec.png
upcycled_smolMoE/cpt.txt
upcycled_smolMoE/continued_pretrained_smolMoE.pt

The saved CPT log shows validation CE around 2.136 by step 100 and confirms that the model still generates coherent text after continued pretraining.

Part 4A: Efficient Training Through Dataset Intervention

Detailed report:

efficient_training/README.md

This subproject answers:

If only 1000 training examples can be used, can we process the full
[cosmopedia-100k](https://huggingface.co/datasets/HuggingFaceTB/cosmopedia-100k)
dataset and choose those examples better?

What Was Built

The efficient_training/ directory includes:

Training Sets

The tested 1000-sample sets were:

1. random seed 42,
2. top dense CE,
3. random rows from dense CE percentile 15-35,
4. highest empirical token entropy,
5. sqrt-diverse moderate dense CE,
6. random seed 123,
7. random seed 777.

The strongest selection idea was the sqrt-diverse set:

quota_c proportional to sqrt(N_c)

where N_c is the size of embedding cluster c.

Within each cluster, rows were filtered to:

token_count in [512, 1536]
dense CE between cluster-local p50 and p85

This protects smaller semantic clusters while avoiding extreme hard examples and extreme sequence lengths.

Main Cosmopedia Result

The detailed README includes plots:

efficient_training/report_figures/

and compact result tables:

efficient_training/report_tables/

Interpretation

The biggest lesson is that random in-domain continued training is a strong baseline. The curated selector did not beat random after one epoch, but it did become the best checkpoint after three epochs.

The result is not "hardest rows win." In fact:

• top dense CE underperformed random,
• top token entropy was the weakest trained intervention,
• low/mid CE random was strong after one epoch,
• sqrt-diverse became best with early stopping.

This suggests that small-data continued pretraining needs a balance of:

domain relevance + semantic coverage + moderate difficulty + length control

Part 4B: Nemotron Expert Specialization

Detailed report:

nemotron_specialization/README.md

This subproject answers:

Can the existing three experts specialize on chat, code, and math SFT data
rather than routing uniformly?

What Was Built

The nemotron_specialization/ directory is a self-contained training pipeline:

The original backbone outside the directory is not edited. The experiment uses a local model copy so the specialization work stays isolated.

Dataset

Source:

nvidia/Llama-Nemotron-Post-Training-Dataset

Domains:

chat
code
math

Balanced selected data:

Split:

The completed Kaggle run used smaller disjoint phase slices for time reasons, as documented in the subdirectory README.

Expert Mapping

The three experts are assigned fixed intended domains:

DOMAIN_TO_EXPERT = {
    "chat": 0,
    "code": 1,
    "math": 2,
}

Three-Phase Curriculum

The specialization method has three phases:

Phase 3 includes:

• language-modeling loss,
• small domain-router CE,
• router z-loss,
• router-distribution orthogonality,
• router variance/confidence loss.

The orthogonality term is behavioral: it pushes average router distributions for chat/code/math away from each other. It does not directly penalize expert weight similarity.

Main Nemotron Result

The important qualitative shift is that phase 1 improves LM loss through forced expert practice, while phase 2 is where the router learns the domain mapping. Phase 3 improves both CE/PPL and router accuracy further under fully learned routing.

The detailed README links the analysis plots:

nemotron_specialization/analysis/

Engineering Choices

Keep Experiments Isolated

The base MoE lives in:

moe/backbone.py

The Cosmopedia CUDA trainer uses a local patched model copy under:

efficient_training/cuda_modeling/

The Nemotron curriculum uses another local copy under:

nemotron_specialization/modeling/

This is intentional. The open-ended experiments required CUDA dtype fixes, forced routing, router logits, and extra routing metadata. Instead of mutating the base assignment implementation repeatedly, the experimental model copies keep each research track self-contained.

Prefer Simple, Auditable Metrics

The project uses:

• causal LM CE and PPL,
• token-weighted CE/PPL for variable-length heldout sets,
• load-balancing loss,
• Gini-style expert specialization,
• router accuracy,
• expert usage,
• domain-by-expert routing matrices.

The metrics are simple enough to inspect manually and specific enough to answer the assignment questions.

Work Around Real Hardware Limits

The project was developed across a local Mac with MPS and Kaggle T4 GPUs.

Important constraints:

• full 100k CE/PPL scoring was too slow locally,
• T4 15 GB memory was too tight for stable 2048-token MoE training,
• fp16 CUDA exposed dtype issues in attention and logits,
• multi-epoch experiments had to be scoped to avoid runaway Kaggle time.

The final stable training setup for Cosmopedia used:

2 x T4
block_size = 1024
batch_size_per_gpu = 1
grad_accum_steps = 4
effective_batch = 8

The Nemotron completed run similarly used smaller phase slices than the full planned 37k/domain training split.

Known Limitations

Cosmopedia Heldout Usage

The 3000-row heldout set was used for both validation and model selection in the follow-up runs. A stricter final report should split it into:

1000 validation rows
2000 final test rows

and only use the 2000-row test set once at the end.

Weights And Data Excluded From Git

Model checkpoints, exported model binaries, raw dataset shards, selected training parquets, caches, and large per-row eval files are intentionally not tracked. The no-weights result artifacts are present, including logs, configs, metrics, plots, and compact CSV/JSON summaries. The README reports are based on those downloaded no-weights outputs.

Custom Model Is Educational, Not Production-Optimized

The custom backbone is intentionally direct and readable. It does not use FlashAttention, sequence packing, activation checkpointing, fused kernels, or other production training optimizations.

No Broad Downstream Benchmarking

The project evaluates CE/PPL and router behavior. It does not run broad downstream tasks, human evaluation, or generation-quality benchmarks.

Reproducing The Main Flow

1. Validate the custom MoE

python test_moe_backbone.py

2. Upcycle the dense model

python upcycle.py

This writes:

upcycled_smolMoE/upcycled_smolMoE.pt

3. Continue pretraining

python train_router.py

This writes:

upcycled_smolMoE/continued_pretrained_smolMoE.pt

4. Cosmopedia dataset intervention

See:

efficient_training/README.md

Key scripts:

python efficient_training/compute_cosmopedia_metrics.py
python efficient_training/embed_cosmopedia_diversity.py
python efficient_training/select_cosmopedia_diverse_ce.py
python efficient_training/train_cosmopedia_cpt_cuda.py
python efficient_training/eval_cosmopedia_cpt_cuda.py

The final report artifacts are:

efficient_training/report_figures/
efficient_training/report_tables/
efficient_training/experimentResults/

5. Nemotron expert specialization

See:

nemotron_specialization/README.md

Main launcher:

python nemotron_specialization/train.py

Analysis artifacts:

nemotron_specialization/analysis/

How To Read This Repository

Recommended order:

1. Read this README for the full project story.
2. Inspect moe/backbone.py to understand the custom architecture.
3. Inspect upcycle.py to see how dense weights become MoE weights.
4. Inspect train_router.py for the base CPT loop and MoE metric.
5. Read efficient_training/README.md for the dataset-intervention research.
6. Read nemotron_specialization/README.md for the expert-specialization research.

The two subdirectory READMEs are intentionally detailed. This root README is the project map and high-level technical narrative; the sub-READMEs are the experiment logs and analysis reports.

Final Takeaway

This project builds a small MoE language-model pipeline end to end:

custom architecture
    -> dense weight transfer
    -> continued pretraining
    -> data selection research
    -> expert specialization research

The Cosmopedia work shows that data selection matters, but simple random continued training is a serious baseline. The best intervention was not the hardest examples or the highest-entropy examples, but a diversity-aware, moderately hard subset with early stopping.

The Nemotron work shows that expert specialization can be trained directly with a staged curriculum. Forced expert warmup improves the experts, router annealing teaches domain routing, and fully learned routing with light regularization produces strong router/domain alignment.

Together, the two experiments show the same broader lesson: with a small upcycled MoE, the useful question is not only "does the loss go down?" It is "what data or routing signal made the model improve, and can we show that improvement with the right heldout and MoE-specific metrics?"