picoLM

≤100M parameter language model for the CSE 251B NanoGPT contest.

Two architectures:

• picoLM Dense — dense transformer (--model_variant dense)
• picoLM MoE — mixture-of-experts transformer (--model_variant moe)

Config

Model defaults in config.yaml:

• picolm_dense for picoLM Dense
• picolm_moe for picoLM MoE

Dataset defaults in dataset_config.yaml.

Training

python3 train.py                              # uses config defaults (picoLM Dense)
python3 train.py --model_variant moe          # picoLM MoE
python3 train.py --config /path/to/config.yaml

Hyperparameter Tuning (Optuna)

pip install optuna

# picoLM Dense sweep (90-95M param budget)
python3 optuna_sweep.py --arch dense --n_trials 50 --proxy_steps 300

# picoLM MoE sweep
python3 optuna_sweep.py --arch moe --n_trials 50 --proxy_steps 300

# Quick smoke test
python3 optuna_sweep.py --arch dense --n_trials 2 --proxy_steps 10 \
    --dataset roneneldan/TinyStories --max_train_samples 128

Results are saved to optuna_results/ with best configs exported as YAML.

Evaluation

python3 evaluate.py --model_dir ./checkpoints --data val.bin