ML building blocks for training experiments.
# Mac:
# # Required for quick install.
# brew install uv
# Ubuntu/Debian:
# # Required for quick install.
# sudo apt-get install -y curl
# curl -LsSf https://astral.sh/uv/install.sh | sh
uv add priml
# Alternatively: python -m pip install priml- model -- composable model definitions and building blocks.
- optimizers -- optimizer implementations for training.
- loss -- loss functions.
- metrics -- evaluation metrics.
- math -- numerical and math utilities.
- train -- the training loop and experiment scaffolding.
- data -- the data pipeline and dataset utilities.
- inference -- inference helpers.
Modules, costs, and metrics pass open keyword arguments down the tree; see the keyword bus for the rule and why it is sound.
from priml.baselines.nanochat import experiments as nanochat
nanochat.exp000().pprint(indent=3)NanoChatLoop.Config(
study_name='nanochat',
experiment_name='exp000',
working_dir=PosixPath('/opt/scratch/runs/nanochat/exp000'),
step=NanoChatTrainStep.Config(
│ model=NanoChatLM.Config(
│ │ block=[
│ │ │ TransformerBlock.Config(
│ │ │ channels_in=512,
│ │ │ channels_out=512,
│ │ │ attn=ValueGatedAttention.Config(
│ │ │ channels_in=512,
│ │ │ channels_out=512,
│ │ │ num_heads=4,
│ │ │ norm_qk=RMSNorm.Config(
│ │ │ channels_in=128,
│ │ │ channels_out=128,
│ │ │ eps=1.1920928955078125e-07
│ │ │ ),
│ │ │ kernel=Flash4Attention.Config(),
│ │ │ window=1_024,
│ │ │ max_seq_len=2_048,
│ │ │ gated=False,
│ │ │ depth_index=((0, 8),)
│ │ │ ),
│ │ │ ffn=SwiGLUReluSquared.Config(
│ │ │ channels_in=512,
│ │ │ channels_out=512,
│ │ │ channels_hidden=2_048,
│ │ │ expansion=4.0,
│ │ │ round_to=1,
│ │ │ depth_index=((0, 8),),
│ │ │ shard='colwise'
│ │ │ ),
│ │ │ norm1=RMSNorm.Config(
│ │ │ channels_in=512,
│ │ │ channels_out=512,
│ │ │ eps=1.1920928955078125e-07
│ │ │ ),
│ │ │ norm2=RMSNorm.Config(
│ │ │ channels_in=512,
│ │ │ channels_out=512,
│ │ │ eps=1.1920928955078125e-07
│ │ │ ),
│ │ │ depth_index=((0, 8),)
│ │ │ ),
...
│ │ │ TransformerBlock.Config(
│ │ │ channels_in=512,
│ │ │ channels_out=512,
│ │ │ attn=ValueGatedAttention.Config(
│ │ │ channels_in=512,
│ │ │ channels_out=512,
│ │ │ num_heads=4,
│ │ │ norm_qk=RMSNorm.Config(
│ │ │ channels_in=128,
│ │ │ channels_out=128,
│ │ │ eps=1.1920928955078125e-07
│ │ │ ),
│ │ │ kernel=Flash4Attention.Config(),
│ │ │ window=2_048,
│ │ │ max_seq_len=2_048,
│ │ │ depth_index=((7, 8),)
│ │ │ ),
│ │ │ ffn=SwiGLUReluSquared.Config(
│ │ │ channels_in=512,
│ │ │ channels_out=512,
│ │ │ channels_hidden=2_048,
│ │ │ expansion=4.0,
│ │ │ round_to=1,
│ │ │ depth_index=((7, 8),),
│ │ │ shard='colwise'
│ │ │ ),
│ │ │ norm1=RMSNorm.Config(
│ │ │ channels_in=512,
│ │ │ channels_out=512,
│ │ │ eps=1.1920928955078125e-07
│ │ │ ),
│ │ │ norm2=RMSNorm.Config(
│ │ │ channels_in=512,
│ │ │ channels_out=512,
│ │ │ eps=1.1920928955078125e-07
│ │ │ ),
│ │ │ depth_index=((7, 8),)
│ │ │ )
│ │ ],
│ │ value_embedding_stride=2,
│ │ embedding=NarrowEmbedding.Config(
│ │ channels_in=8_192,
│ │ channels_out=512,
│ │ dtype=torch.bfloat16,
│ │ inner=Embedding.Config(
│ │ channels_in=8_192,
│ │ channels_out=512,
│ │ init_weight=functools.partial(<function priml.model.init.normal at 0xdefacedeface>, std=1.0)
│ │ )
│ │ ),
│ │ norm=RMSNorm.Config(
│ │ channels_in=512,
│ │ channels_out=512,
│ │ eps=None
│ │ ),
│ │ lm_head=SoftCap.Config(
│ │ channels_in=512,
│ │ channels_out=8_192,
│ │ inner=Linear.Config(
│ │ channels_in=512,
│ │ channels_out=8_192,
│ │ init_weight=functools.partial(<function priml.model.init.normal at 0xdefacedeface>, std=0.001)
│ │ )
│ │ ),
│ │ rope=RoPE.Config(
│ │ channels_head=128,
│ │ frequencies=HuggingFaceFrequencies.Config(base=10000.0),
│ │ dtype=torch.bfloat16
│ │ ),
│ │ mix=ResidualMix.Config(
│ │ num_layers=8,
│ │ channels_in=512
│ │ )
│ ),
│ optimizer=CompositeOptimizer.Config(
│ optimizers=[
│ PartialConfig(<class 'functools.partial'>, <class 'priml.optimizers.fused_adamw.FusedAdamW'>, betas=(0.8, 0.95), eps=1e-10, weight_decay=0.0, compile=True, lr=0.004898979485566357),
│ PartialConfig(<class 'functools.partial'>, <class 'priml.optimizers.fused_adamw.FusedAdamW'>, betas=(0.8, 0.95), eps=1e-10, weight_decay=0.0, compile=True, lr=0.7348469228349535),
│ PartialConfig(<class 'functools.partial'>, <class 'priml.optimizers.fused_adamw.FusedAdamW'>, betas=(0.8, 0.95), eps=1e-10, weight_decay=0.0, compile=True, lr=0.7348469228349535),
│ PartialConfig(<class 'functools.partial'>, <class 'priml.optimizers.fused_adamw.FusedAdamW'>, betas=(0.8, 0.95), eps=1e-10, weight_decay=0.0, compile=True, lr=0.005),
│ PartialConfig(<class 'functools.partial'>, <class 'priml.optimizers.fused_adamw.FusedAdamW'>, betas=(0.96, 0.95), eps=1e-10, weight_decay=0.0, compile=True, lr=0.5),
│ NorMuon.Config()
│ ],
│ select=[
│ matching('lm_head'),
│ excluding(matching('embed'), 'value_embeds'),
│ matching('value_embeds'),
│ matching('mix.running'),
│ matching('mix.original'),
│ excluding(NorMuon.eligible_tensor, 'embed', 'lm_head')
│ ],
│ drop_empty=True
│ ),
│ loss=TokenCrossEntropy.Config(ignore_index=-1),
│ lr_schedule=PartialConfig(<class 'functools.partial'>, <function priml.math.schedules.constant at 0xdefacedeface>),
│ parallelism=NoParallel.Config(device='cuda'),
│ model_quantization=NoModelQuantization.Config(),
│ activation_memoization=DefaultActivationStorage.Config(),
│ compile=PartialConfig(<class 'functools.partial'>, <function torch.compile at 0xdefacedeface>),
│ ema=NoEMA.Config(),
│ schedule=PartialConfig(<class 'functools.partial'>, <function priml.math.schedules.trapezoidal at 0xdefacedeface>, flat=0.5),
│ adam_lr_tuned_at_channels=512
),
dataset=NanoChatData.Config(
base_dir='/opt/scratch',
working_dir=PosixPath('/opt/scratch/datasets/nanochat'),
tokenizer_dir=PosixPath('/opt/scratch/datasets/nanochat/tokenizer'),
device='cuda',
vocab_size=8_192
),
metrics_train={},
metrics_eval={'val': BitsPerByte.Config()},
checkpointer=Checkpointer.Config(
base_dir=PosixPath('/opt/scratch/runs/nanochat/exp000'),
working_dir=PosixPath('/opt/scratch/runs/nanochat/exp000/checkpoints'),
storer=SyncLocalStateDictStorer.Config(),
resume=False
),
phase_timer=PhaseTimer.Config(
base_dir=PosixPath('/opt/scratch/runs/nanochat/exp000'),
working_dir=PosixPath('/opt/scratch/runs/nanochat/exp000/profiling')
),
tracker=TrackerList.Config(
trackers={ 'metrics': FileTracker.Config(
base_dir=PosixPath('/opt/scratch/runs/nanochat/exp000'),
working_dir=PosixPath('/opt/scratch/runs/nanochat/exp000/metrics.json')
),
'wandb': AsyncTracker.Config(
tracker=WandbTracker.Config(
project='nanochat',
base_dir=PosixPath('/opt/scratch/runs/nanochat/exp000'),
working_dir=PosixPath('/opt/scratch/runs/nanochat/exp000/wandb'),
ingestion=WandbIngestion(),
metric_step_metrics={},
run_config={}
)
)},
base_dir=PosixPath('/opt/scratch/runs/nanochat/exp000'),
working_dir=PosixPath('/opt/scratch/runs/nanochat/exp000')
),
max_time=300.0,
max_time_kind='train',
num_steps_eval=-1,
num_steps_log=1,
early_train_log_steps=0,
eval_every_epoch=False,
seed=42,
runtime=SingleProcess.Config(
device='cuda',
float32_matmul_precision='high'
)
)
Priml has rich support for understanding where compute is (theoretically) being spent.
For example,
from priml import cost as theoretical
from priml.baselines.sudoku import experiments as sudoku
sudoku.exp000().finalize().step.model.cost(batch_size=64, dtype=torch.bfloat16)
theoretical.peak()["rtx5090", torch.bfloat16]
theoretical.peak()["h100", torch.bfloat16]prints,
Cost(params=6833666, params_active=6823426, bytes_state=20480)
flops flops bytes bytes bytes intensity intensity
bf16 f32 bf16 f32 int64 bf16 f32
primal matmul 1.307T - 4.807G - - 271.8 -
primal elementwise 3.506G 524.3K 12.24G 4.194M - 0.2864 0.125
primal reduction 762.3M - 1.539G - - 0.4954 -
primal selection - - 22.61M - 84.48K - -
adjoint matmul 871G - 3.205G - - 271.8 -
adjoint elementwise 1.585G 524.3K 7.525G 4.194M - 0.2106 0.125
adjoint reduction 206.4M - 415.6M - - 0.4967 -
adjoint selection 5.308M - 31.87M - 82.94K 0.1665 -
total 2.184T 1.049M 29.79G 8.389M 167.4K 73.31 0.125
flops bytes intensity
matmul 209.5T 1.792T 116.9
elementwise 104.8T 1.792T 58.48
reduction 104.8T 1.792T 58.48
selection 104.8T 1.792T 58.48
sort 104.8T 1.792T 58.48
flops bytes intensity
matmul 989T 3.35T 295.2
elementwise 67T 3.35T 20
reduction 67T 3.35T 20
selection 67T 3.35T 20
sort 67T 3.35T 20
The H100's BF16 matmul roofline knee is 295.2 FLOPs/byte. The
transformer's matmul intensity is about 8% below this knee
(271.8 / 295.2 - 1 = -0.07927). The knee is the crossover from
memory-bound to compute-bound operation.
Without a duration, utilization returns intensity divided by the device's
roofline knee for each cell. With duration_sec, it returns achieved FLOP/s
divided by min(peak FLOP/s, intensity * peak bytes/s). For a compute-bound
matmul cell this equals MFU; for a memory-bound cell the denominator is lower
than peak compute. The device argument accepts a name such as "h100"
or a single-device report such as theoretical.peak()["h100"].
theoretical.utilization(
sudoku.exp000().finalize().step.model.cost(batch_size=64, dtype=torch.bfloat16),
device="h100",
)prints,
bf16 f32
primal matmul 0.9206 -
primal elementwise 0.01432 0.00625
primal reduction 0.02477 -
adjoint matmul 0.9206 -
adjoint elementwise 0.01053 0.00625
adjoint reduction 0.02484 -
adjoint selection 0.008327 -
total 1.924 0.0125
The displayed utilization total sums per-cell ratios; it is not a
whole-model utilization and should not be interpreted as one.
These bytes describe logical operand I/O for the analytical unfused algorithm. They are not measured HBM traffic: fusion and cache reuse can change physical memory transfers. The roofline comparisons above use that analytical traffic model; they do not establish achieved hardware utilization.
See CONTRIBUTING.md for local validation and the public contribution flow.
Sibling projects in the rekursiv-ai family:
- sagent — The self-mutating multi-provider coding-agent CLI and typed Python library.
- trackinizer — Centralized agent database for tracking inquiries, work, and the evidence behind conclusions.
- wesearch — Web search, resilient page fetch, and scholarly-paper lookup without a browser stack.
- madcatter — Rich-based Markdown renderer for the terminal; ships the
mdcatCLI. - configgle — Hierarchical experiment configuration in typed pure-Python dataclasses instead of YAML.
- copybarista — Bidirectional source sync for publishing OSS-ready trees from a monorepo.
If you find our work useful, please consider citing:
@misc{rekursivai2026priml,
title={Priml - Strongly typed configurable building blocks for A/B testing ML research.}
author={Joshua V. Dillon and Dan Kondratyuk},
year={2026},
howpublished={Github},
url={https://github.com/rekursiv-ai/priml},
}