Serialization

An important aspect of a deep learning framework is the ability to save and load training state to and from disk. Burn serializes its records with the burnpack format, a compact binary container implemented by the burn-pack crate.

Constraints

1. Users should be able to add any field to a module, even fields that are not serializable.

   This can include constants, database connections, other module references, or any other information. Only the parameters (tensors) should be serialized; the structure of the module itself is encapsulated by its configuration (hyperparameters).

2. Records should be decoupled from the backend in use.

   A record holds plain tensor data (TensorData), so weights saved with one backend can be loaded on another. Parameter initialization is lazy, so loading a record does not require eagerly materializing the module first.

3. The format should be fast to load and embeddable.

   Tensor data is stored contiguously and aligned, so it can be read back with zero-copy / memory-mapped loading, and a record can be saved straight to bytes for no-std environments.

The burnpack format

The burn-pack crate is intentionally minimal and tensor-library-agnostic: it knows how to read and write the container format but has no notion of Burn modules. A burnpack file has three parts:

┌────────────────────────────────────────────────────────────┐
│ Header (fixed size)                                          │
│   magic "BURN", format version, metadata byte length         │
├────────────────────────────────────────────────────────────┤
│ Metadata (CBOR)                                              │
│   tensors : map<name, descriptor>                            │
│     dtype, shape, data_offsets, optional param_id            │
│   scalars  : map<name, typed scalar>                         │
│   metadata : map<string, string>  user key/value pairs       │
├────────────────────────────────────────────────────────────┤
│ Tensor data section                                          │
│   each tensor's bytes start on a 256-byte boundary so the    │
│   data can be sliced zero-copy / memory-mapped from a file   │
└────────────────────────────────────────────────────────────┘

All multi-byte integers are little-endian. Tensor entries carry an optional param_id used to preserve a parameter's identity across save/load. Besides tensors, a pack can store named typed scalars (integers, floats, booleans), which the optimizer and learning rate scheduler records use to persist their non-tensor state.

A pack is read back with burn_pack::Reader, which yields burn_pack::Tensor entries plus a scalar map. Tensor bytes stay lazy: a file-backed reader only touches the disk when an entry's data is actually used.

A pack is written with burn_pack::Writer, which is symmetric about this, because laziness is a property of Tensor rather than of the writer's input. Tensor::new carries bytes that are already resident; Tensor::deferred carries a byte length plus a provider that yields the data on demand. The writer reads only the metadata while planning, to compute every descriptor and offset before any I/O, then calls each provider once in write order, dropping each tensor's bytes before requesting the next. burn-store collects a module's parameters as deferred tensors, so each device readback waits until the writer reaches that tensor. Paired with Writer::write_to_file, which streams to disk, saving a large module costs one tensor of host memory at a time rather than the whole set. The in-memory sinks (into_bytes, write_into) still build the container as a whole.

Because the offset table is committed from the declared length before the bytes exist, a length that turns out to be wrong would misplace every later tensor. Two checks prevent that: planning rejects a length that disagrees with the tensor's own shape and dtype, and the write pass rejects bytes whose length differs from what was reserved. Quantized tensors are exempt from the first, since their packed values and inline scales are not a product of shape and dtype; that exception is why Tensor::deferred takes an explicit length rather than deriving one.

Writer::write_to_file builds the container in a scratch file beside the destination and renames it into place only once it is complete, so a save that fails partway leaves whatever was already at that path untouched. That covers disk failure (a full disk, a quota, an I/O error) for every caller, and provider failure for deferred tensors, whose bytes are produced mid-write. Writer::write_to_file_in_place truncates and rewrites the destination instead: it skips the fsync and the transient second copy, but a failed save leaves a truncated file.

The three record types

Higher layers bridge their state to and from burnpack through three record types. Each one can be serialized to a file (save / load, appending the .bpk extension when the path has none) or to an in-memory byte buffer (into_bytes / from_bytes).

ModuleRecord (burn-core, burn::store)

Holds a module's parameters: a flat list of (path, ParamId, TensorData) entries keyed by module path. It is produced and applied through the Module trait itself rather than a separate codegen type:

• module.into_record() walks the module with a ModuleVisitor (the Collector), recording each float/int/bool parameter under its dotted path.
• module.load_record(record) (or the fallible try_load_record) walks the module with a ModuleMapper that looks each parameter up by path and loads the matching tensor.

Load-time behavior is configured with builder methods on the record, ignored when saving:

• allow_partial(bool) — tolerate module parameters absent from the record.
• validate(bool) — toggle shape-mismatch / missing-tensor validation.
• with_dtype_policy(..) / cast_to_module_dtype() — choose whether a parameter adopts the record's dtype (DTypePolicy::FromRecord, the default) or casts the data to the module parameter's current dtype (DTypePolicy::CastToModule).

The save-side dtype is not configurable: the record stores whatever dtype the module currently holds. The dtype applied on load is controlled by the record's DTypePolicy (.cast_to_module_dtype() / .with_dtype_policy(..)).

This module in burn-core is intentionally tiny — no filtering, key remapping, or adapters. The richer snapshot/import tooling (filtering, key remapping, PyTorch/SafeTensors cross-framework stores) lives in the burn-store crate.

OptimizerRecord (burn-optim)

Holds an optimizer's per-parameter state. Unlike a module record (keyed by module path), it is keyed per parameter: each parameter's state is decomposed into tensors named "{param_id}.{field}" plus a few typed scalar entries kept in the burnpack scalar map (including a __rank scalar so the state can be reconstructed without inferring rank from tensor shapes).

• optimizer.to_record() flattens each parameter's DynState into tensors and scalars.
• optimizer.load_record(record) reconstructs the states (no device argument: tensors load on the default device and migrate to each parameter's device on the next step).

LrSchedulerRecord (burn-optim)

Holds a learning rate scheduler's state, which is just a handful of scalars (step counters, current learning rate) and no tensors. Produced/applied through the LrScheduler trait's to_record() / load_record(). Composed schedulers nest their children's records under an index prefix (with_record / record).

Checkpointing in burn-train

During training, burn-train defines a Checkpoint trait (save(path) / load(path)) implemented for all three record types — ModuleRecord, OptimizerRecord, LrSchedulerRecord — and for () (a stateless no-op). The Checkpointer<R: Checkpoint> trait drives periodic saves; the FileCheckpointer writes each record to {name}-{epoch}.bpk under the experiment directory. This is how the model, optimizer, and scheduler are persisted and restored across epochs.

Notes

• There is no Recorder, PrecisionSettings, Module::Record associated type, or #[derive(Record)] any more. All of those were part of the previous serde-based record system, which has been removed.
• Cross-framework import/export (PyTorch .pt, SafeTensors) still lives in burn-store (PytorchStore, SafetensorsStore, BurnpackStore).