Core framework development

Changes to xlm-core live under src/xlm/ — the training harness, datamodule, metrics, shared modules, tasks, Hydra configs, and CLI. This is separate from model families in xlm-models/.

What counts as "core"

Area	Location	Examples
Training harness	harness.py	Harness, LossFunction, Predictor protocols
Data pipeline	datamodule.py	TextDataModule, DatasetManager, collator protocols
Metrics	metrics.py	MetricWrapper, custom metrics
Shared modules	modules/	Rotary transformer, DDiT backbones
Tasks / datasets	tasks/	Preprocessing, post-hoc evaluators
CLI	commands/	xlm, xlm-scaffold, job dispatch
Hydra configs	configs/	Shared dataset YAMLs, trainer defaults

Prerequisites

Set up your environment first: Install from source in the Contributing hub.

For core work, install the same stacks CI uses so optional task paths and model families stay importable:

pip install -e ".[all]"
pip install -e ./xlm-models
pip install -r requirements/test_requirements.txt \
            -r requirements/dev_requirements.txt \
            -r requirements/docs_requirements.txt \
            -r requirements/lint_requirements.txt

Protocol or harness changes can break any model family — keep xlm-models editable while you develop.

Before you start

1. Search existing Issues and Discussions for overlap.
2. Open an issue for non-trivial work. Use the feature request template and note which component is affected (core, CLI, metrics, datamodule, etc.).
3. Assess impact radius:
4. Protocol changes (LossFunction, Predictor, Collator, or Harness in harness.py / datamodule.py) affect every family under xlm-models/ (arlm, ilm, mlm, mdlm, flexmdm, dream) and external models using the same contracts.
5. Hydra config changes (renamed keys, moved defaults under configs/) may require updating experiment and datamodule YAMLs across xlm-models/ and downstream repos.
6. Skim the closest existing code and the Related guides below before proposing API shape.

Implementation conventions

Follow repo conventions in the files you touch:

• Types: use jaxtyping for tensor shapes where applicable (same pattern as model families).
• Docstrings: Google style (see .darglint and .flake8).
• Formatting: line length 79 — black src xlm-models tests (see pyproject.toml).
• Protocols: if you add or change a protocol in harness.py or datamodule.py, update docstrings and consider backward compatibility for existing model packages.
• Configs: if you add Hydra groups or change shared YAMLs, update affected experiment configs and mention the migration in your PR.

There is no single “core feature” template — use neighboring modules and tests as reference.

Related guides

Depending on what you are changing, these existing pages may help:

• Data pipeline — how DatasetManager and dataloader names connect to metrics
• Metrics — defining and wiring metrics
• Custom commands — extending job_type via external model commands.yaml
• Adding a task or dataset — if your core change includes a new benchmark
• Debugging utilities — debug modes and utilities

Testing

• Unit tests: add or extend tests under tests/core/.
• Datamodule / collator changes: update matrix coverage in test_datamodule.py; consider scenarios under tests/integration/ — see Integration tests.
• Harness or protocol changes: run model-family tests to catch regressions:

pytest tests/models/ -v

• Fast loop (while developing):

pytest -m "not slow and not cli"

• CI-aligned run (matches tests.yml):

pytest -m "not gpu and not cli and not integration"

• Style: format and lint changed paths, e.g. black src xlm-models tests and flake8 src/xlm/ on files you modified.

See Running tests for markers, coverage, and smoke tests.

Documentation

• Public API: modules under src/xlm/ are largely covered by api-autonav for src/xlm in mkdocs.yml; add docstrings on new public symbols.
• Conceptual changes: if you introduce a new user-facing workflow, consider a page under docs/guide/ and a nav: entry in mkdocs.yml.
• Verify the site builds:

mkdocs build

More detail: Docs and issues.

Opening a PR

1. Push your branch to your fork and open a PR against main.
2. Use the pull request template and check Core framework (src/xlm/) under contribution type.
3. Summarize what behavior or protocol changed and why.
4. Call out breaking changes for model families or external packages.
5. Link the related issue and note tests/docs you ran or updated.

General workflow: Contributing hub.