# answerdotai/modernbert architecture

> Architecture diagram of the GitHub repository answerdotai/modernbert, generated by GitDiagram from its file tree, README and a sample of source files (last updated 2026-08-13). It is an AI-made overview: check details against the source.

- Interactive diagram: https://gitdiagram.com/answerdotai/modernbert
- Repository: https://github.com/answerdotai/modernbert

## Overview

#### 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`](https://github.com/answerdotai/modernbert/blob/main/main.py) (Composer launcher), [`run_evals.py`](https://github.com/answerdotai/modernbert/blob/main/run_evals.py) (Evaluation dispatch), [`src/text_data.py`](https://github.com/answerdotai/modernbert/blob/main/src/text_data.py) (Text datasets), [`src/flex_bert.py`](https://github.com/answerdotai/modernbert/blob/main/src/flex_bert.py) (FlexBERT encoder)

## Diagram (Mermaid)

```mermaid
flowchart TD

subgraph group_config["Experiment configuration"]
  node_yaml_runs["Run YAMLs<br/>configuration"]
  node_model_yaml["Model YAMLs<br/>configuration<br/>[flex_bert.yaml]"]
end

subgraph group_training["Training runtime"]
  node_main["Composer launcher<br/>python entry point<br/>[main.py]"]
  node_composer{{"Composer trainer<br/>training runtime"}}
  node_flexbert["FlexBERT encoder<br/>model implementation<br/>[flex_bert.py]"]
  node_alternate_models["Alternate BERT models<br/>model implementations<br/>[mosaic_bert.py]"]
  node_bert_primitives["BERT primitives<br/>model layers<br/>[model.py]"]
  node_optimization["Optimizer and schedule<br/>training policy<br/>[optimizer.py]"]
  node_callbacks["Runtime callbacks<br/>instrumentation"]
end

subgraph group_data["Data pipeline"]
  node_source_corpora[("Text and tokenized corpora<br/>data source")]
  node_dataset_conversion["Dataset conversion<br/>data preparation<br/>[hf_to_mds.py]"]
  node_text_datasets["Text datasets<br/>dataset loader<br/>[text_data.py]"]
  node_sequence_preparation["Pack and pad batches<br/>batch preparation<br/>[sequence_packer.py]"]
end

subgraph group_evaluation["Evaluation and integrations"]
  node_checkpoints[("Trained checkpoints<br/>model artifact")]
  node_evaluation_dispatch["Evaluation dispatch<br/>python entry point<br/>[run_evals.py]"]
  node_task_jobs["Finetuning task jobs<br/>evaluation jobs<br/>[finetuning_jobs.py]"]
  node_classification["Classification flows<br/>python entry points<br/>[glue.py]"]
  node_retrieval["Retrieval workflows<br/>example training<br/>[train_st.py]"]
  node_hf_integration["Hugging Face conversion<br/>artifact integration<br/>[convert_to_hf.py]"]
  node_wandb(("Weights &amp; Biases<br/>experiment integration"))
end

node_yaml_runs -->|"drives"| node_main
node_model_yaml -->|"selects model"| node_main
node_main -->|"constructs"| node_flexbert
node_main -.->|"can construct"| node_alternate_models
node_flexbert -->|"composes"| node_bert_primitives
node_alternate_models -.->|"shares architecture"| node_bert_primitives
node_main -->|"configures"| node_optimization
node_main -->|"registers"| node_callbacks
node_main -->|"launches"| node_composer
node_source_corpora -->|"converted to MDS"| node_dataset_conversion
node_source_corpora -->|"loaded from"| node_text_datasets
node_dataset_conversion -->|"produces MDS"| node_text_datasets
node_text_datasets -->|"token sequences"| node_sequence_preparation
node_sequence_preparation -->|"Composer batches"| node_main
node_composer -->|"writes"| node_checkpoints
node_checkpoints -->|"evaluated by"| node_evaluation_dispatch
node_evaluation_dispatch -->|"dispatches"| node_task_jobs
node_checkpoints -->|"finetunes"| node_classification
node_checkpoints -->|"adapts"| node_retrieval
node_checkpoints -->|"exports"| node_hf_integration
node_evaluation_dispatch -->|"logs metrics"| node_wandb

click node_yaml_runs "https://github.com/answerdotai/modernbert/blob/main/yamls/main/flex-bert-base.yaml"
click node_model_yaml "https://github.com/answerdotai/modernbert/blob/main/yamls/models/flex_bert.yaml"
click node_main "https://github.com/answerdotai/modernbert/blob/main/main.py"
click node_flexbert "https://github.com/answerdotai/modernbert/blob/main/src/flex_bert.py"
click node_alternate_models "https://github.com/answerdotai/modernbert/blob/main/src/mosaic_bert.py"
click node_bert_primitives "https://github.com/answerdotai/modernbert/blob/main/src/bert_layers/model.py"
click node_optimization "https://github.com/answerdotai/modernbert/blob/main/src/optimizer.py"
click node_callbacks "https://github.com/answerdotai/modernbert/blob/main/src/callbacks/dataloader_speed.py"
click node_dataset_conversion "https://github.com/answerdotai/modernbert/blob/main/src/data/hf_to_mds.py"
click node_text_datasets "https://github.com/answerdotai/modernbert/blob/main/src/text_data.py"
click node_sequence_preparation "https://github.com/answerdotai/modernbert/blob/main/src/sequence_packer.py"
click node_evaluation_dispatch "https://github.com/answerdotai/modernbert/blob/main/run_evals.py"
click node_task_jobs "https://github.com/answerdotai/modernbert/blob/main/src/evals/finetuning_jobs.py"
click node_classification "https://github.com/answerdotai/modernbert/blob/main/glue.py"
click node_retrieval "https://github.com/answerdotai/modernbert/blob/main/examples/train_st.py"
click node_hf_integration "https://github.com/answerdotai/modernbert/blob/main/convert_to_hf.py"
click node_wandb "https://github.com/answerdotai/modernbert/blob/main/wandb_log_live_eval.py"

classDef toneNeutral fill:#f8fafc,stroke:#334155,stroke-width:1.5px,color:#0f172a
classDef toneBlue fill:#dbeafe,stroke:#2563eb,stroke-width:1.5px,color:#172554
classDef toneAmber fill:#fef3c7,stroke:#d97706,stroke-width:1.5px,color:#78350f
classDef toneMint fill:#dcfce7,stroke:#16a34a,stroke-width:1.5px,color:#14532d
classDef toneRose fill:#ffe4e6,stroke:#e11d48,stroke-width:1.5px,color:#881337
classDef toneIndigo fill:#e0e7ff,stroke:#4f46e5,stroke-width:1.5px,color:#312e81
classDef toneTeal fill:#ccfbf1,stroke:#0f766e,stroke-width:1.5px,color:#134e4a
class node_yaml_runs,node_model_yaml toneBlue
class node_main,node_composer,node_flexbert,node_alternate_models,node_bert_primitives,node_optimization,node_callbacks toneAmber
class node_source_corpora,node_dataset_conversion,node_text_datasets,node_sequence_preparation toneMint
class node_checkpoints,node_evaluation_dispatch,node_task_jobs,node_classification,node_retrieval,node_hf_integration,node_wandb toneRose
```

## Components

### Experiment configuration

Declarative run and model definitions.

- [Run YAMLs](https://github.com/answerdotai/modernbert/blob/main/yamls/main/flex-bert-base.yaml): configuration · Training experiment definitions and shared defaults. · `yamls/main/flex-bert-base.yaml`
- [Model YAMLs](https://github.com/answerdotai/modernbert/blob/main/yamls/models/flex_bert.yaml): configuration · Model-family and architecture settings. · `yamls/models/flex_bert.yaml`

### Training runtime

Composer-based pretraining assembly.

- [Composer launcher](https://github.com/answerdotai/modernbert/blob/main/main.py): python entry point · Builds and launches configured training. · `main.py`
- **Composer trainer**: training runtime · Orchestrates multi-GPU training, algorithms, and callbacks.
- [FlexBERT encoder](https://github.com/answerdotai/modernbert/blob/main/src/flex_bert.py): model implementation · Modular configurable encoder family. · `src/flex_bert.py`
- [Alternate BERT models](https://github.com/answerdotai/modernbert/blob/main/src/mosaic_bert.py): model implementations · MosaicBERT and Hugging Face-compatible variants. · `src/mosaic_bert.py`
- [BERT primitives](https://github.com/answerdotai/modernbert/blob/main/src/bert_layers/model.py): model layers · Attention, embeddings, normalization, and BERT configuration. · `src/bert_layers/model.py`
- [Optimizer and schedule](https://github.com/answerdotai/modernbert/blob/main/src/optimizer.py): training policy · Optimizer, LR scheduler, and RoPE schedule construction. · `src/optimizer.py`
- [Runtime callbacks](https://github.com/answerdotai/modernbert/blob/main/src/callbacks/dataloader_speed.py): instrumentation · Data speed, packing efficiency, and gradient monitoring. · `src/callbacks/dataloader_speed.py`

### Data pipeline

Corpus conversion and batch preparation.

- **Text and tokenized corpora**: data source · Streaming MDS, local MDS, CSV, TSV, or JSONL inputs.
- [Dataset conversion](https://github.com/answerdotai/modernbert/blob/main/src/data/hf_to_mds.py): data preparation · Converts Hugging Face and source datasets to MDS. · `src/data/hf_to_mds.py`
- [Text datasets](https://github.com/answerdotai/modernbert/blob/main/src/text_data.py): dataset loader · Streaming and local MDS dataset paths. · `src/text_data.py`
- [Pack and pad batches](https://github.com/answerdotai/modernbert/blob/main/src/sequence_packer.py): batch preparation · Token sequence packing and BERT padding. · `src/sequence_packer.py`

### Evaluation and integrations

Checkpoint-consuming downstream workflows.

- **Trained checkpoints**: model artifact · Pretrained weights consumed by downstream jobs.
- [Evaluation dispatch](https://github.com/answerdotai/modernbert/blob/main/run_evals.py): python entry point · Runs configured checkpoint evaluation jobs. · `run_evals.py`
- [Finetuning task jobs](https://github.com/answerdotai/modernbert/blob/main/src/evals/finetuning_jobs.py): evaluation jobs · General, GLUE, and SuperGLUE task definitions. · `src/evals/finetuning_jobs.py`
- [Classification flows](https://github.com/answerdotai/modernbert/blob/main/glue.py): python entry points · Standalone GLUE and sequence-classification workflows. · `glue.py`
- [Retrieval workflows](https://github.com/answerdotai/modernbert/blob/main/examples/train_st.py): example training · Sentence Transformers, PyLate, and ColBERT adaptations. · `examples/train_st.py`
- [Hugging Face conversion](https://github.com/answerdotai/modernbert/blob/main/convert_to_hf.py): artifact integration · Exports checkpoints for Hugging Face use. · `convert_to_hf.py`
- [Weights & Biases](https://github.com/answerdotai/modernbert/blob/main/wandb_log_live_eval.py): experiment integration · Artifact retrieval and live evaluation logging. · `wandb_log_live_eval.py`

## Connections

- 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