How answerdotai/modernbert works
Purpose and execution surface
This is the research training/evaluation repository for ModernBERT and FlexBERT encoder experiments: pretraining, scaling/architecture ablations, GLUE/SuperGLUE finetuning, and retrieval experiments. The primary Composer launch point is main.py, driven by YAML experiment definitions such as yamls/main/flex-bert-base.yaml and shared defaults in yamls/defaults.yaml. Classification and benchmark-oriented entry points include sequence_classification.py, glue.py, eval.py, and run_evals.py.
Configuration-driven training assembly
main.py interprets a YAML training configuration to construct the model, data loaders, optimizer/scheduler, callbacks, algorithms, and Composer trainer. Model-family configurations are separated from run configurations:
• yamls/models/flex_bert.yaml
• yamls/main/flex-bert-rope-base.yaml
• yamls/modernbert/modernbert-base-pretrain.yaml
This is the main control boundary: experiment behavior should generally be introduced declaratively before adding hard-coded launch logic.
Encoder/model layer
FlexBERT is the modular encoder implementation, while MosaicBERT and Hugging Face-compatible variants provide alternate architectures or interoperability:
• src/flex_bert.py
• src/mosaic_bert.py
• src/hf_bert.py
The architecture is decomposed into configurable BERT primitives—embeddings, attention, MLP, normalization, rotary position encodings, padding, initialization, and loss—under:
• src/bert_layers/model.py
• src/bert_layers/attention.py
• src/bert_layers/configuration_bert.py
Attention/runtime choices, including FlashAttention-oriented paths, are architectural performance dependencies rather than separate services.
Data ingestion and sequence preparation
Pretraining consumes text or tokenized corpora through streaming and non-streaming dataset paths:
• src/text_data.py
• src/sequence_packer.py
• src/bert_padding.py
StreamingTextDataset supports Mosaic Streaming-backed MDS plus CSV/TSV/JSONL inputs; NoStreamingDataset targets decompressed local MDS for throughput. Dataset preparation/conversion is isolated in:
• src/data/hf_to_mds.py
• src/data/mds_conversion.py
• src/convert_dataset.py
The boundary is MDS/local source data → tokenization/packing/padding → Composer-compatible batches.
Training optimization and runtime instrumentation
Training uses MosaicML Composer as the orchestration runtime, with multi-GPU execution implied by Composer launch. Optimizer and learning-rate policy construction are centralized in:
• src/optimizer.py
• src/scheduler.py
• src/algorithms/rope_schedule.py
Operational feedback and performance management are implemented as Composer callbacks:
• src/callbacks/dataloader_speed.py
• src/callbacks/packing_efficiency.py
• src/callbacks/log_grad_norm.py
The environment expects GPU acceleration and FlashAttention 2 (optionally FlashAttention 3 on H100), specified through environment.yaml and dependency files.
Evaluation and finetuning pipelines
Evaluation jobs consume trained checkpoints plus training/model configs. General evaluation dispatch and task definitions reside in:
• run_evals.py
• src/evals/finetuning_jobs.py
• src/evals/glue_jobs.py
GLUE/SuperGLUE-oriented standalone flows use:
• glue.py
• sequence_classification.py
• src/evals/superglue_jobs.py
These flows are distinct from pretraining: checkpoint/config → task-specific finetuning/evaluation → reported metrics.
Retrieval and artifact integrations
Retrieval examples adapt ModernBERT to external Sentence Transformers or PyLate/ColBERT workflows:
• examples/train_st.py
• examples/train_pylate.py
• src/colbert_beir/train.py
Hugging Face is used for model/dataset ecosystem integration and checkpoint conversion:
• convert_to_hf.py
• src/data/get_counts_from_hf.py
Weights & Biases supports artifact retrieval and live evaluation logging:
• download_artifacts_from_wandb.py
• wandb_log_live_eval.py
Main entry points: main.py Composer launcher · run_evals.py Evaluation dispatch · text_data.py Text datasets · flex_bert.py FlexBERT encoder
Components 20
Experiment configuration
Declarative run and model definitions.
- Run YAMLs · Training experiment definitions and shared defaults.
yamls/main/flex-bert-base.yaml - Model YAMLs · Model-family and architecture settings.
yamls/models/flex_bert.yaml
Training runtime
Composer-based pretraining assembly.
- Composer launcher · Builds and launches configured training.
main.py - Composer trainer · Orchestrates multi-GPU training, algorithms, and callbacks.
- FlexBERT encoder · Modular configurable encoder family.
src/flex_bert.py - Alternate BERT models · MosaicBERT and Hugging Face-compatible variants.
src/mosaic_bert.py - BERT primitives · Attention, embeddings, normalization, and BERT configuration.
src/bert_layers/model.py - Optimizer and schedule · Optimizer, LR scheduler, and RoPE schedule construction.
src/optimizer.py - Runtime callbacks · Data speed, packing efficiency, and gradient monitoring.
src/callbacks/dataloader_speed.py
Data pipeline
Corpus conversion and batch preparation.
- Text and tokenized corpora · Streaming MDS, local MDS, CSV, TSV, or JSONL inputs.
- Dataset conversion · Converts Hugging Face and source datasets to MDS.
src/data/hf_to_mds.py - Text datasets · Streaming and local MDS dataset paths.
src/text_data.py - Pack and pad batches · Token sequence packing and BERT padding.
src/sequence_packer.py
Evaluation and integrations
Checkpoint-consuming downstream workflows.
- Trained checkpoints · Pretrained weights consumed by downstream jobs.
- Evaluation dispatch · Runs configured checkpoint evaluation jobs.
run_evals.py - Finetuning task jobs · General, GLUE, and SuperGLUE task definitions.
src/evals/finetuning_jobs.py - Classification flows · Standalone GLUE and sequence-classification workflows.
glue.py - Retrieval workflows · Sentence Transformers, PyLate, and ColBERT adaptations.
examples/train_st.py - Hugging Face conversion · Exports checkpoints for Hugging Face use.
convert_to_hf.py - Weights & Biases · Artifact retrieval and live evaluation logging.
wandb_log_live_eval.py
Connections 21
- Run YAMLs → Composer launcher: drives
- Model YAMLs → Composer launcher: selects model
- Composer launcher → FlexBERT encoder: constructs
- Composer launcher → Alternate BERT models: can construct
- FlexBERT encoder → BERT primitives: composes
- Alternate BERT models → BERT primitives: shares architecture
- Composer launcher → Optimizer and schedule: configures
- Composer launcher → Runtime callbacks: registers
- Composer launcher → Composer trainer: launches
- Text and tokenized corpora → Dataset conversion: converted to MDS
- Text and tokenized corpora → Text datasets: loaded from
- Dataset conversion → Text datasets: produces MDS
- Text datasets → Pack and pad batches: token sequences
- Pack and pad batches → Composer launcher: Composer batches
- Composer trainer → Trained checkpoints: writes
- Trained checkpoints → Evaluation dispatch: evaluated by
- Evaluation dispatch → Finetuning task jobs: dispatches
- Trained checkpoints → Classification flows: finetunes
- Trained checkpoints → Retrieval workflows: adapts
- Trained checkpoints → Hugging Face conversion: exports
- Evaluation dispatch → Weights & Biases: logs metrics