lobotomy

A PyTorch library for language-model training and experimentation, derived from OLMo-core. It includes transformer models, data pipelines, distributed training, evaluation, and checkpoint tools under the lobotomy Python namespace.

The Fizz pretrainer provides an end-to-end training workflow with configurable model architectures, optimizers, data mixtures, and single-device, DDP, or FSDP execution. Experiments include looped transformers, attention residuals, multi-token prediction, Engram memory, and patch pretraining.

Setup

Requires Python 3.12 or newer and uv. From the repository root:

uv sync

This installs all workspace packages in editable mode, along with the development dependencies. Optional integrations are declared on individual packages. For example, install Hugging Face dataset tooling and W&B support with:

uv sync --package cli --extra hf --extra wandb

Other extras include eval and comet on cli, and fa4, fla, and torchao on modeling. Install the dependencies required by your chosen attention backend and training configuration; the base installation does not include every backend.

Upgrading an existing environment

Older builds used hatch-una to bundle internal dependencies into multiple packages. This could mix an editable CLI with stale modeling/data code copied into site-packages. Packages now declare normal dependencies and ship only their own namespace directory; uv sync installs workspace members in editable mode.

After updating the checkout, stop running training/tokenization processes and run this once from the repository root:

uv pip uninstall cli tokenize config data eval io modeling runtime train utils
uv sync --locked --inexact

Uninstall the old workspace packages together before reinstalling so their shared file records cannot delete or shadow another package's files. The uninstall targets only workspace packages; --inexact retains extra installed packages, while sync still enforces locked dependency versions. Checkpoints and datasets are unaffected. Keep any optional extras you normally specify when syncing. Remove the earlier PYTHONPATH workaround and restart Python processes or notebook kernels; already-imported modules stay in memory until restarted. Future source updates require no manual path overrides.

Building packages

uv build --all-packages --wheel

Each wheel contains only its own lobotomy.<package> files and retains its internal dependencies in package metadata. For deployment outside the workspace, provide the required workspace wheels together, for example with uv pip install dist/*.whl. A single application wheel no longer bundles all internal library code. Source distributions can also be built and rebuilt outside the workspace with standard Hatchling.

Training

Inspect the available options and resolve a configuration before training:

uv run lobotomy-fizz-pretrain --help
uv run lobotomy-fizz-pretrain \
  --recipe.model_shape=10m \
  --trainer.save_folder=./checkpoints \
  --dry-run

--dry-run prints the resolved configuration and parameter counts without loading data, allocating the model, or initializing distributed training. Settings use --section.field=value, with explicit true and false values for booleans.

Before starting a run, configure the tokenized data source, attention backend, batch sizes, and evaluation settings for your environment. The training guide covers data preparation, model presets, optimizers, distributed launches, checkpoint loading, and experimental features.

For streaming tokenization of Hugging Face datasets:

uv run lobotomy-fizz-tokenize --help

See the tokenization guide for a complete example. Both commands are also available as Python modules: lobotomy.cli.scripts.fizz.fizz_training.pretrain and lobotomy.tokenize.tokenize_hf_stream.

Repository layout

Directory	Python namespace	Contents
libs/utils	lobotomy.utils	Dependency-free shared primitives: base exceptions and environment errors
libs/io	lobotomy.io	Local and remote file operations, storage backends, and filesystem caching
libs/config	lobotomy.config	Configuration dataclasses, enums, registries, and configuration errors
libs/runtime	lobotomy.runtime	Devices, tensors, distributed helpers, RNG state, logging, and environment setup
libs/modeling	lobotomy.modeling	Model layers, attention backends, generation, distributed primitives, and tokenizers
libs/data	lobotomy.data	Datasets, data loaders, packing, and data mixtures
libs/eval	lobotomy.eval	In-loop evaluation tasks, metrics, tokenizers, and bundled datasets
libs/train	lobotomy.train	Training modules, optimizers, and learning-rate schedules
apps/tokenize	lobotomy.tokenize	Streaming Hugging Face dataset tokenization and shard publication
apps/cli	lobotomy.cli	Trainer, callbacks, evaluation, RL, checkpoint tools, and training scripts

utils, io, and config form the dependency base. runtime adds execution-time primitives on top of config. modeling depends on all four. data depends on modeling; train depends on both; cli brings everything together.

Import path changes

In-loop evaluation now lives in libs/eval. Replace olmo_eval imports with lobotomy.eval; the CLI eval extra uses this workspace package. The imported source, bundled data, and provenance are documented in the eval guide.

lobotomy.config and lobotomy.runtime were extracted out of lobotomy.modeling. The old import paths were removed rather than re-exported, so existing code has to be updated. Mapping for the moved APIs:

Old	New
lobotomy.modeling.config.Config	lobotomy.config.Config
lobotomy.modeling.config.StrEnum	lobotomy.config.StrEnum
lobotomy.modeling.config.Registrable	lobotomy.config.Registrable
lobotomy.modeling.config.UNSET	lobotomy.config.UNSET
lobotomy.modeling.config.DType	lobotomy.runtime.dtypes.DType
lobotomy.modeling.exceptions.OLMoConfigurationError	lobotomy.config.exceptions.OLMoConfigurationError
lobotomy.modeling.exceptions.OLMoCLIError	lobotomy.runtime.exceptions.OLMoCLIError
lobotomy.modeling.exceptions.OLMoThreadError	lobotomy.runtime.exceptions.OLMoThreadError
lobotomy.modeling.distributed.utils.*	lobotomy.runtime.distributed.*
lobotomy.modeling.utils.apply_to_tensors, move_to_device, mark_dynamic, same_storage, serialize_to_tensor, deserialize_from_tensor, get_element_size, info_value_of_dtype, min_value_of_dtype	lobotomy.runtime.tensor.*
lobotomy.modeling.utils.get_default_device, has_compute_capability, gc_cuda, cuda_sync_debug_mode, get_or_init_stream	lobotomy.runtime.device.*
lobotomy.modeling.utils.get_default_thread_count, wait_for, threaded_generator	lobotomy.runtime.concurrency.*
lobotomy.modeling.utils.set_env_var, set_env_variables, filter_warnings, prepare_cli_environment	lobotomy.runtime.environment.*
lobotomy.modeling.utils.LogFilterType, setup_logging, log_extra_field, logging_configured, excepthook, install_excepthook, log_or_print, log_once, warn_once, format_float, format_int, format_timedelta	lobotomy.runtime.logging.*
lobotomy.modeling.utils.seed_all, EnvRngStates	lobotomy.runtime.random.*

The following stay in lobotomy.modeling.utils because they are modeling concerns: generate_uuid, roundrobin, powers_of_2, capped_powers_of_2, flatten_dict, ensure_multiple_of, has_flash_attn, and record_flops.

This is an intentional compatibility break. Saved configuration that names the old class paths through _CLASS_, and pickled objects referencing moved symbols, are not guaranteed to load. Newly written configuration and state use the canonical new paths; no automatic fallback resolution is performed.

Tests

Run each package's suite separately from the repository root. The suites share module names and should not be collected in a single pytest session.

uv run pytest libs/utils/tests
uv run pytest libs/io/tests
uv run pytest libs/config/tests
uv run pytest libs/runtime/tests
uv run pytest libs/modeling/tests
uv run pytest libs/data/tests
uv run pytest libs/eval/tests -m "not slow"
uv run pytest libs/train/tests
uv run pytest apps/cli/tests

Use -m "not gpu" to exclude GPU-marked tests. Depending on the suite, tests may also require optional dependencies, distributed execution, or authenticated access to S3, GCS, or Hugging Face. Some cloud fixtures create and delete remote objects; excluding GPU tests does not exclude those fixtures.

Embedding training

Train a native lobotomy transformer checkpoint on streamed query/document pairs from Nomic's unsupervised dataset:

uv run --package embedding lobotomy-train-embeddings CHECKPOINT_DIR OUTPUT_DIR

By default, training streams every split once and runs until exhaustion; use --split reddit_title_body to select a split or --steps 1000 to limit the run. At each optimizer update the trainer chooses uniformly among the source splits that still have data. Every microbatch in that update comes from that split, including when gradient accumulation is enabled. Splits are streamed until exhausted; no source is restarted. Smaller splits can run out earlier, leaving only larger splits near the end. --seed controls both source selection and each split's shuffle. --shuffle-buffer is divided across the selected splits (with at least one row per split), bounding the total shuffle buffer approximately. Progress uses published split row counts to estimate total steps; filtering invalid pairs can reduce the actual total. The final batch from each split is trained if it contains at least two pairs; a lone final pair is skipped because in-batch InfoNCE needs a negative.

To train on Nomic's supervised dataset, select it explicitly:

uv run --package embedding lobotomy-train-embeddings CHECKPOINT_DIR OUTPUT_DIR \
  --dataset supervised --hard-negatives 1

The trainer samples up to one hard negative per query by default; change that with --hard-negatives (use all for every valid negative, or 0 for in-batch negatives only). Sampling is seeded by --seed. All positive documents and sampled hard negatives in a microbatch are candidates for each query; only the query's paired document is labeled positive. Rows with fewer valid hard negatives contribute those available. Each additional negative requires another encoded document and increases memory use; all can use very large amounts of memory because all negatives are encoded with gradients in each microbatch. Adjust the microbatch size for your device. The dataset defaults to unsupervised, and --split, accumulation, and checkpointing work in both modes.

Microbatch sequence lengths are padded up to a multiple of 64 by default, capped at --max-length. This limits the number of allocation shapes for variable-length data, at the cost of extra padding and computation. Use --pad-to-multiple-of 1 to disable rounding or select a different multiple.

Use --gradient-accumulation-steps 4 to accumulate four microbatches before an optimizer update. --batch-size is the number of pairs per microbatch, so this example processes up to four times that many pairs per update. InfoNCE negatives remain within each microbatch; accumulation does not enlarge the negative pool. Partial accumulation groups are weighted by pair count and still produce an update. --steps, the estimated total, metrics, and checkpoint intervals all count optimizer updates.

The checkpoint must contain config.json and model_and_optim, including the original tokenizer configuration. The trainer appends EOS after truncation and pools its final hidden state after the language model head's normalization. It normalizes the embedding and uses query-to-document InfoNCE with in-batch negatives. The saved native checkpoint retains the language model head; use the same EOS pooling and normalization for embedding inference. This trainer runs on one device and saves model weights, not optimizer or data stream state. Use --help for training options.

Execution settings can be overridden using the same --model.field=value syntax as Fizz. For example, --model.block.sequence_mixer.backend=torch changes the attention backend (the exact path depends on the checkpoint's block layout). Overrides are applied to the checkpoint configuration before model construction and saved in the output config. Weight-incompatible architecture changes will fail when loading the checkpoint.

Use --compile-model and optionally --compile-mode, --autocast bfloat16, or --activation-checkpointing full to control execution. Activation checkpointing also supports selected_blocks with --checkpoint-block-interval, selected_ops, and budget with --activation-memory-budget and compilation. Kernel support and compilation depend on the device and installed dependencies. Fizz's distributed trainer, optimizer and scheduler options are not used by this single-device app.

Training logs use the Fizz console and W&B callbacks. Console metrics are reported on step 1 and every 10 steps by default (--log-interval changes the interval). To enable W&B and log InfoNCE loss and learning rate on every step:

uv run --package embedding --extra wandb lobotomy-train-embeddings CHECKPOINT_DIR OUTPUT_DIR \
  --wandb --wandb-project PROJECT --wandb-name RUN_NAME

On CUDA, logged metrics also include current and peak allocated and reserved memory, inactive split memory, and cumulative allocator retries and OOMs. Peak values cover the period since the previous metrics report. The same diagnostics are logged if a CUDA out-of-memory exception occurs.

The W&B callback requires WANDB_API_KEY; --wandb-entity and --wandb-group are also available. W&B logging is disabled unless --wandb is passed.

To save periodic model checkpoints, set an interval and an optional local or remote destination. For example, to save every 1,000 steps to an HF bucket:

uv run --package embedding lobotomy-train-embeddings CHECKPOINT_DIR OUTPUT_DIR \
  --checkpoint-interval 1000 --checkpoint-dir hf://buckets/ORG/BUCKET/embeddings

Checkpoints are written synchronously to step1000, step2000, and so on. Without --checkpoint-dir, they go under OUTPUT_DIR/checkpoints. The final model is still saved to OUTPUT_DIR. Existing checkpoints are not overwritten. Each checkpoint includes model weights and configuration; loading one starts a new run with a fresh optimizer and data stream, rather than resuming training.

Export embeddings

Export a native embedding checkpoint to a local Sentence Transformers model:

uv run --package embedding --extra export lobotomy-export-embeddings CHECKPOINT_DIR OUTPUT_DIR --max-length 512

The export uses the checkpoint tokenizer and the native Hugging Face converter, then adds last-token pooling and L2 normalization. It preserves tokenizer EOS handling when it matches training; otherwise it configures a fast tokenizer to append EOS after truncation. Slow tokenizers that do not already match this behavior cannot be exported. Set --max-length to the training sequence length (the default is 512). Conversion validates the native-to-Hugging Face model by default; --no-validate skips that check. --overwrite replaces an existing output directory. For exports with custom model code, load with:

from sentence_transformers import SentenceTransformer
model = SentenceTransformer("OUTPUT_DIR", trust_remote_code=True)
embeddings = model.encode(["example text"])