Open-source code for training and scaling Energy-based Foundation Models

Keywords: Energy-based Foundation Models, language models, pretraining, scaling laws

🎉 Updates • 🚀 Quick Start • 📈 Scaling Law • ⚙️ Pretrained Checkpoints • 💬 Chat Demo

English · 简体中文

---

🎉 Updates

• 2026-10-23 — ✨✨ Technical report released. See [report].
• 2026-10-07 — ✨✨ Full codebase released.

🚀 Quick Start

Install the GPU, Hugging Face, and web-chat dependencies once:

uv sync --extra gpu --extra hf --extra web
source .venv/bin/activate

Download and tokenize the training data (run once before the first training job; requires network). Prepares climbmix, dclm, and fineweb by default; use --datasets to prepare only a subset:

bash runs/prepare_data.sh                  # default: climbmix + dclm + fineweb
bash runs/prepare_data.sh --datasets dclm  # only the named datasets

Run a complete training job by setting the token budget:

TARGET_TOTAL_TOKENS=100000000 bash runs/train.sh

Data preparation needs network access, so run it once on a networked cpu-worker. Training itself is offline, but must run where the data is ready.

📈 Scaling Law

On DCLM, validation BPB falls as training tokens and model size increase. Zero-shot accuracy scales with model size across ClimbMix, DCLM, and FineWeb, and improves with training tokens.

DCLM pretraining	Parameter scaling
Zero-shot model-size scaling	Zero-shot token scaling

Use configs/train.yaml as the default configuration and set TARGET_TOTAL_TOKENS to choose the training budget.

⚙️ Pretrained Checkpoints

Pretrained models are published in the EBFM Hugging Face collection. The collection includes 160M, 520M, and 1B models trained on DCLM, FineWeb, and ClimbMix; one 160M OWT model; and a FineWeb SFT model named EBFM-FineWeb-1B-CHAT.

Load a published model directly by its Hugging Face ID. The first call downloads and caches the weights and tokenizer:

import torch
from ebfm.modeling_ebfm import load_checkpoint

model_id = "guan-wang/EBFM-FineWeb-1B-CHAT"
model, tokenizer, hparams, device = load_checkpoint(model_id, device="cuda")
tokens = tokenizer.encode("Hello, EBFM.", append=tokenizer.get_bos_token_id())
input_ids = torch.tensor([tokens], dtype=torch.long, device=device)
with torch.no_grad():
    logits = model(input_ids)
print(logits.shape)

Replace the model ID with any repository listed in the collection.

💬 Chat Demo

Run the browser chat with the published SFT model:

bash runs/chat.sh --web \
  --checkpoint guan-wang/EBFM-FineWeb-1B-CHAT \
  --host 0.0.0.0 \
  --port 8000

Open http://localhost:8000 in a browser on the server machine. If the server runs on a remote GPU node, use your cluster's port forwarding to expose port 8000 on your local machine.

The PyTorch CPU and CUDA indexes are configured in pyproject.toml. Choose one extra per environment.

Run

The shell launchers are deliberately thin. They always run from the project root, so relative paths are stable in local shells and distributed jobs.

# Pretraining: training length comes from TARGET_TOTAL_TOKENS.
TARGET_TOTAL_TOKENS=100000000 bash runs/train.sh

# Fine-tuning: initialize from a published Hugging Face model.
TARGET_TOTAL_TOKENS=100000000 bash runs/sft.sh \
  --dataset_name ebfm_sft \
  --finetuning_model_ckpt guan-wang/EBFM-DCLM-1B

# Position-wise BPB evaluation.
MODEL_DIR="$(python -c 'from huggingface_hub import snapshot_download; print(snapshot_download(repo_id="guan-wang/EBFM-DCLM-1B"))')"
bash runs/eval.sh \
  --checkpoint guan-wang/EBFM-DCLM-1B \
  --dataset dclm \
  --data-dir data/dclm \
  --tokenizer-path "${MODEL_DIR}" \
  --output_root . \
  --run_name eval-dclm

# QA evaluation.
bash runs/qa.sh \
  --checkpoint guan-wang/EBFM-DCLM-1B \
  --tokenizer-path "${MODEL_DIR}" \
  --eval-bundle data/eval_bundle \
  --output_root . \
  --run_name qa-dclm

# Interactive generation.
bash runs/chat.sh --checkpoint guan-wang/EBFM-FineWeb-1B-CHAT

Chat, validation evaluation, QA evaluation, and SFT initialization accept published Hugging Face model IDs. BPB evaluation also needs the local snapshot directory for token_bytes.pt.

Training reads configs/train.yaml; evaluation reads configs/eval.yaml. Explicit command-line arguments take precedence over YAML values. Training length is controlled by TARGET_TOTAL_TOKENS. The effective gradient accumulation is derived from global_batch_size, device_batch_size, and the distributed world size, so max_steps and accumulate_grad_batches are not independent training controls.

Every run uses the same output layout:

outputs/<run-name>/
├── checkpoints/
├── checkpoints_hf/
├── logs/
│   └── wandb/
└── results/

When training finishes, checkpoints_hf/ holds a clean Hugging Face model (config.json + model.safetensors) exported from the final checkpoint. It contains no tokenizer and no custom modeling code; read it back with load_checkpoint("<run-dir>/checkpoints_hf", tokenizer_path=...). The export needs the hf extra (uv sync --extra gpu --extra hf); if Transformers is missing it is skipped with a log message and the .ckpt files are unaffected.

Data and tokenizer assets are not included in the repository. Pass their locations through the configuration or command line. Supported pretraining and BPB datasets include DCLM, FineWeb, ClimbMix, and OWT.

Load from Hugging Face

load_checkpoint() accepts a Hugging Face model ID and downloads the Transformers weights and tokenizer automatically. For example:

import torch
from ebfm.modeling_ebfm import load_checkpoint

model, tokenizer, hparams, device = load_checkpoint(
    "guan-wang/EBFM-FineWeb-1B", device="cuda"
)
tokens = tokenizer.encode("hello", append=tokenizer.get_bos_token_id())
input_ids = torch.tensor([tokens], dtype=torch.long, device=device)
with torch.no_grad():
    logits = model(input_ids)
print(logits.shape)

For BPB evaluation, download the same model repository locally and pass that directory as --tokenizer-path; it contains the tokenizer byte table used by the metric. The loader also supports local Transformers directories and Lightning checkpoints produced by this training code.

Repository layout

ebfm/
├── configs/                                 # Default run configurations
│   ├── eval.yaml                 # Evaluation defaults
│   └── train.yaml                # Pretraining defaults
├── ebfm/                                    # Model, data, and training library
│   ├── __init__.py                           # Package initialization
│   ├── common.py                             # Shared constants and utilities
│   ├── config.py                             # Training and evaluation config parsing
│   ├── configuration_ebfm.py                  # Transformers model configuration
│   ├── core_eval.py                          # Core benchmark evaluation
│   ├── dataloader.py                         # Streaming data loaders
│   ├── dataset.py                             # Shared dataset utilities
│   ├── dataset_sft.py                        # Supervised fine-tuning datasets
│   ├── disk_aware_checkpoint.py              # Checkpoint save/load utilities
│   ├── logger.py                              # Training metric and artifact logging
│   ├── metrics.py                             # Loss and BPB metrics
│   ├── modeling_ebfm.py                        # EBFM architecture, loading, and inference
│   ├── optim.py                               # Muon and AdamW optimizers and schedules
│   ├── pretrain_dataset.py                    # Pretraining data pipeline
│   ├── tokenizer.py                           # EBFM tokenizer and tokenizer loading
│   └── trainer.py                             # Lightning training and evaluation module
├── figs/                                    # README figures and demo images
│   ├── OpenEBFM_logo.png                        # Project logo / title banner
│   ├── acc_size.png                             # Zero-shot accuracy by model size
│   ├── acc_tokens.png                           # Zero-shot accuracy by tokens
│   ├── chat.jpg                                 # Web chat interface screenshot
│   ├── dclm_pretrain.png                        # DCLM validation BPB by tokens
│   └── pretrain_size.png                        # DCLM parameter scaling plot
├── runs/                                    # Shell entry points for common workflows
│   ├── chat.sh                   # Launch interactive chat
│   ├── eval.sh                   # Run validation evaluation
│   ├── prepare_data.sh           # Prepare and tokenize datasets
│   ├── qa.sh                     # Run task-based QA evaluation
│   ├── sft.sh                    # Fine-tune from a checkpoint
│   └── train.sh                  # Pretrain an EBFM model
├── scripts/                                 # Python command-line entry points
│   ├── chat.py                   # Serve terminal or web chat
│   ├── eval.py                   # Evaluate validation loss and BPB
│   ├── export_hf.py              # Export a checkpoint to Transformers format
│   ├── prepare_data.py           # Build training and evaluation data
│   ├── qa.py                     # Evaluate QA task suites
│   ├── sft.py                    # Run supervised fine-tuning
│   └── train.py                  # Run pretraining
├── tasks/                                   # Dataset/task adapters for evaluation
│   ├── common.py                             # Shared task interfaces and helpers
│   ├── customjson.py                         # Load custom JSON evaluation tasks
│   ├── gsm8k.py                              # GSM8K math reasoning task
│   ├── mmlu.py                               # MMLU multiple-choice task
│   ├── smoltalk.py                           # SmolTalk conversation task
│   └── spellingbee.py                        # SpellingBee word-generation task
├── tests/                                   # Lightweight unit and compatibility tests
│   ├── test_config_and_boundaries.py       # Validate config and data boundaries
│   ├── test_dataloader_lightweight_state.py # Check lightweight loader state
│   ├── test_dataset_sft_mask.py            # Check SFT loss masking
│   ├── test_hf_configuration.py            # Check Transformers configuration
│   └── test_modeling_ebfm_rope.py           # Check rotary embeddings
├── .gitignore                                # Excludes local data, outputs, and caches
├── pyproject.toml                            # Package metadata and dependency groups
├── README.md                                 # English project guide
├── README_zh.md                              # Simplified Chinese project guide
└── uv.lock                                   # Locked dependency versions

Development

uv sync --extra cpu --group dev
pytest -q
python -m scripts.train --help
python -m scripts.eval --help
python -m scripts.qa --help
python -m scripts.chat --help

Citation

If you found this work helpful, please star and cite this work:

@misc{pu-openesm-energy-steered-models,
  author = {Pu, Yuan and Lu, Yudong and Wang, Ting and Hong, Zhicong and Li, Qinuo and Wu, Xuansheng and Zou, Na and Hu, Xia and Wang, Guanchu},
  keywords = {Computation and Language (cs.CL), Machine Learning (cs.LG), Energy-Based Models, Foundation Models, Language Modeling, Energy Steering, Scaling},
  title = {OpenESM: Energy-Steered Models Are Scalable Foundation Learners},
  publisher = {arXiv}
}

Acknowledgement

Acknowledgement to Nanochat and EBT for their open-source contributions to GPT and energy-based Transformers.