35 · Content revision 5
Save decoder state, replay one specified update
Learn how a versioned checkpoint stores the tokenizer, decoder, trainer-paired AdamW state, and a separate sampling RNG, rejects corrupted bytes, and matches one update whose inputs, targets, and learning rate the caller supplies.
Save the declared component state, not only weights
Chapter 33 selected decoder step . Its trainer returns one sealed selected training state containing the model snapshot, the matching AdamW snapshot, and their shared completed step. In this fixed fixture the selection is also the last training step, but Chapter 35 does not infer pairing from that coincidence. An untrusted file that declares selected step while its AdamW counter is is rejected.
Checkpoint creation requires that trainer-issued capture. Version stores the model snapshot, AdamW snapshot, and both step values, but no independent model-lineage proof. Loading checks that the selected-step value equals the AdamW counter; equal counters in arbitrary bytes do not by themselves prove a common training trajectory.
The checkpoint stores:
- schema, tokenizer-layout, and sampling-RNG versions;
- five ordered literal token pieces and every decoder configuration field;
- stable named parameter tensors;
- named AdamW moment tensors, parameter groups, step, and exact accumulated beta powers; and
- the selected step plus one separate raw SplitMix64 sampling state.
At this required clean post-update boundary there are no gradients to save. The stored AdamW configuration includes its base learning rate, but not the next caller-supplied override or the Chapter 33 schedule.
Keep the ownership boundary explicit:
- Stored: tokenizer, decoder configuration and parameter bits, trainer-paired AdamW state, their shared recorded step, and the sampling RNG.
- Supplied by the caller for the demonstrated update: inputs
[0,1], targets[1,2], and learning rate . - Still required to continue training: corpus and split identity, tokenized data, batch order and cursor, training RNG, learning-rate schedule, gradient clipping, and validation policy.
- Outside this checkpoint: the Chapter 34 evaluation report and test provenance, gradients at this clean boundary, and the attention cache that Chapter 38 will own.
The saved SplitMix64 state is used for later token sampling. It is not the Chapter 33 batch-shuffle seed or another training RNG.
The fixture’s five token IDs, through , map to the one-byte program
labels 0 through 4. Together with the model and optimizer records they make
ordered payload records. Encoding the same state twice produces
the same bytes and the same checksum label
fnv1a64:2b8b6097eaed6a91.
The promise is deliberately narrow: under the same arithmetic environment, loading reproduces exact bits. Changing hardware or arithmetic kernels can change floating-point behavior even when a checkpoint is valid. FNV-1a detects accidental corruption; it does not authenticate a file against an attacker.
Advance each offset by shape times byte width
For canonical payload record , the next absolute byte offset is
The header is bytes. Five one-byte token records occupy
, so the first parameter begins at offset . It is
token_embedding.weight, has shape , and stores each element as an
eight-byte f64. Its exclusive end is therefore
No alignment padding is inserted. An on-disk f64 need not start at a native
memory alignment boundary because the loader copies eight bytes and applies
little-endian decoding; it never casts a file pointer to an f64 pointer.
Every shape product, byte-width multiplication, offset addition, and conversion
from on-disk u64 to process usize is checked. The loader requires each start
to equal the preceding end and the final end to equal byte . A gap,
overlap, arithmetic overflow, truncation, or trailing byte is an error.
Separate record order, shape, and representation
- is the absolute start offset of record .
- is the absolute start offset of the following record.
- is the byte width implied by record ‘s stored dtype: for
u8, foru32, or forf64. - is the length of axis in record .
- is the record’s element count.
- selects one shape axis.
- follows one stable tokenizer, parameter, and optimizer-state order.
- is the complete header size and first payload offset, so .
A descriptor stores role, name, dtype, shape, absolute offset, and byte length. Those fields have different jobs. Shape says how values form a tensor; dtype determines representation and byte width; role and name identify the stored state and its component; canonical order fixes the deterministic sequence. A matching file length cannot compensate for a changed name or shape.
The five grouped spans below cover all descriptors without gaps. Aggregated rows summarize adjacent descriptors; the loader still validates every descriptor individually.
| Record or group | Role | Dtype | Shape or count | Bytes per element | Half-open byte range |
|---|---|---|---|---|---|
literal-token-0 | literal token | u8 | |||
literal-token-1 through literal-token-3 | literal tokens | u8 | one-byte records | ||
literal-token-4 | literal token | u8 | |||
token_embedding.weight | model parameter | f64 | |||
| remaining parameters and moments | model and AdamW state | f64 | tensor records / scalar values |
The first bytes are short enough to inspect byte by byte:
4c 4c 4d 43 50 33 35 00 01 00 04 03 02 01 35 0b
They begin with the magic bytes for LLMCP35, followed by schema version ,
the little-endian marker, and the start of the encoded header length. The table,
not machine alignment, determines where every payload record lives.
From artifact bundles to validated LLM checkpoints
This progression follows the state needed to preserve a language model’s meaning and continue its computation.
An isolated parameter blob does not say which tokenizer, model configuration, tensor shapes, optimizer moments, or random stream gives those bytes their meaning, while separately coordinated artifacts can drift apart.
OpenAI’s GPT-2 downloader is a concrete 2019 language-model artifact boundary. OpenAI’s downloader retrieves GPT-2 checkpoint data, index, and metadata together with a checkpoint pointer, hyperparameters, encoder data, and BPE vocabulary, showing that the released language model was a coordinated artifact bundle rather than one isolated weight file. The script does not establish a single self-contained file, exact training resumption, optimizer or RNG restoration, checksums, or this chapter’s wire format.
ZeRO exposes the scale pressure behind model-and-optimizer state families. Rajbhandari and colleagues classify large-model training state as parameters, gradients, and optimizer state such as Adam momentum and variance, then quantify and partition those state families as language models scale. From that accounting we can infer that weights alone do not determine the next adaptive update. ZeRO studies distributed memory; its activation-checkpoint discussion is about recomputation. It does not define durable serialization or require gradients at this chapter’s clean post-update boundary.
The later safetensors format specification is a self-describing tensor boundary. The safetensors format describes a little-endian header length, per-tensor dtype, shape, and half-open byte offsets, and a completely indexed row-major data buffer without holes. It does not by itself define tokenizer, decoder configuration, optimizer, RNG, checksum, or atomic-publication semantics. Its safety goal must not be confused with this chapter’s accidental-corruption check.
Released neural language models coordinated tokenizer, configuration, and checkpoint artifacts; large-model training made optimizer state a major state family, and later tensor containers exposed dtype, shape, and byte offsets before loading values.
A reproducible LLM checkpoint combines self-describing tensor storage with an application schema that names both stored component state and caller-owned continuation inputs; this chapter stores tokenizer, decoder, paired AdamW, and sampling-RNG state without claiming a full trainer restart.
The road to reproducible modern LLMs joins coordinated model artifacts,
optimizer state, and inspectable tensor metadata with an application-level
schema that states what is stored and what a caller must still supply. The local
Rust contrast derives bytes from the selected model’s
f64 parameter values. It then records that the complete -byte file
also has five tokenizer records, named parameter records, optimizer
moment records, and the sampling RNG state. This measured comparison does
not describe GPT-2 or safetensors as a raw-memory dump, and the course format is
not safetensors-compatible or a proposed universal standard.
rust/demos/ch35-checkpoints/src/lib.rs#historical-checkpoint-contrast #[derive(Clone, Debug, PartialEq, Eq)]
pub struct HistoricalCheckpointContrast {
pub isolated_parameter_bytes: Vec<u8>,
pub isolated_parameter_tensors: usize,
pub isolated_parameter_scalars: usize,
pub checkpoint_records: usize,
pub tokenizer_records: usize,
pub optimizer_moment_records: usize,
pub checkpoint_file_bytes: usize,
pub checkpoint_sampling_rng_state: u64,
}
/// Contrasts model-derived value bytes with the complete local LLM checkpoint.
pub fn historical_checkpoint_contrast(
checkpoint: &Checkpoint,
encoded: &llm_from_scratch::checkpoint::EncodedCheckpoint,
) -> HistoricalCheckpointContrast {
let isolated_parameter_bytes = checkpoint
.model_state()
.bit_pattern()
.into_iter()
.flat_map(u64::to_le_bytes)
.collect();
let tokenizer_records = encoded
.tensors()
.iter()
.filter(|descriptor| {
matches!(
descriptor.role(),
CheckpointTensorRole::LiteralToken | CheckpointTensorRole::BpePairs
)
})
.count();
let optimizer_moment_records = encoded
.tensors()
.iter()
.filter(|descriptor| {
matches!(
descriptor.role(),
CheckpointTensorRole::OptimizerFirstMoment
| CheckpointTensorRole::OptimizerSecondMoment
)
})
.count();
HistoricalCheckpointContrast {
isolated_parameter_bytes,
isolated_parameter_tensors: checkpoint.model_state().parameter_names().len(),
isolated_parameter_scalars: checkpoint.model_state().scalar_count(),
checkpoint_records: encoded.tensors().len(),
tokenizer_records,
optimizer_moment_records,
checkpoint_file_bytes: encoded.bytes().len(),
checkpoint_sampling_rng_state: checkpoint.sampling_rng_state(),
}
} Encode, validate, and replace atomically
AdamWStateEntry owns one parameter name, shape, finite first moment, and
finite non-negative second moment. AdamWState adds configuration, optional
decay groups, step count, exact repeatedly multiplied beta powers, and a stable
name-keyed map. Step requires powers exactly and no moments; later steps
require powers in and nonempty moments. Loading restores those powers
directly instead of recomputing them through a potentially different arithmetic
path.
rust/crates/llm-from-scratch/src/training/adamw.rs#adamw-persistence-state /// One validated name-keyed moment pair prepared for persistence.
#[derive(Clone, Debug, PartialEq)]
pub struct AdamWStateEntry {
name: String,
moments: AdamWMomentState,
}
impl AdamWStateEntry {
pub fn new(
name: impl Into<String>,
shape: Vec<usize>,
first_moment: Vec<f64>,
second_moment: Vec<f64>,
) -> Result<Self, AdamWStateError> {
let name = name.into();
if name.is_empty() {
return Err(AdamWStateError::EmptyParameterName);
}
let elements = shape.iter().try_fold(1_usize, |product, &dimension| {
product.checked_mul(dimension)
});
let Some(elements) = elements else {
return Err(AdamWStateError::ShapeProductOverflow { name });
};
if first_moment.len() != elements || second_moment.len() != elements {
return Err(AdamWStateError::MomentLengthMismatch {
name,
shape,
expected: elements,
first: first_moment.len(),
second: second_moment.len(),
});
}
for (kind, values) in [
("first", first_moment.as_slice()),
("second", second_moment.as_slice()),
] {
if let Some((index, &value)) = values
.iter()
.enumerate()
.find(|(_, value)| !value.is_finite())
{
return Err(AdamWStateError::NonFiniteMoment {
name,
kind,
index,
value,
});
}
}
if let Some((index, &value)) = second_moment
.iter()
.enumerate()
.find(|(_, value)| **value < 0.0)
{
return Err(AdamWStateError::NegativeSecondMoment { name, index, value });
}
Ok(Self {
name,
moments: AdamWMomentState {
shape,
first: first_moment,
second: second_moment,
},
})
}
pub fn name(&self) -> &str {
&self.name
}
pub const fn moments(&self) -> &AdamWMomentState {
&self.moments
}
}
/// A complete graph-free AdamW continuation snapshot.
#[derive(Clone, Debug, PartialEq)]
pub struct AdamWState {
config: AdamWConfig,
groups: Option<AdamWParameterGroups>,
step: u64,
beta1_power: f64,
beta2_power: f64,
states: BTreeMap<String, AdamWMomentState>,
}
impl AdamWState {
pub fn new(
config: AdamWConfig,
groups: Option<AdamWParameterGroups>,
step: u64,
beta1_power: f64,
beta2_power: f64,
entries: Vec<AdamWStateEntry>,
) -> Result<Self, AdamWStateError> {
validate_beta_power("beta1", beta1_power, step)?;
validate_beta_power("beta2", beta2_power, step)?;
if (step == 0) != entries.is_empty() {
return Err(AdamWStateError::StatePresence {
step,
entries: entries.len(),
});
}
let mut states = BTreeMap::new();
for entry in entries {
if states.insert(entry.name.clone(), entry.moments).is_some() {
return Err(AdamWStateError::DuplicateParameterName { name: entry.name });
}
}
if let Some(groups) = &groups {
let expected = groups.parameter_names();
let actual = states.keys().cloned().collect::<Vec<_>>();
if step > 0 && expected != actual {
return Err(AdamWStateError::ParameterGroupsMismatch { expected, actual });
}
}
Ok(Self {
config,
groups,
step,
beta1_power,
beta2_power,
states,
})
}
pub const fn config(&self) -> AdamWConfig {
self.config
}
pub const fn parameter_groups(&self) -> Option<&AdamWParameterGroups> {
self.groups.as_ref()
}
pub const fn step_count(&self) -> u64 {
self.step
}
pub const fn beta1_power(&self) -> f64 {
self.beta1_power
}
pub const fn beta2_power(&self) -> f64 {
self.beta2_power
}
pub fn parameter_names(&self) -> impl ExactSizeIterator<Item = &str> {
self.states.keys().map(String::as_str)
}
pub fn state(&self, name: &str) -> Option<&AdamWMomentState> {
self.states.get(name)
}
}
fn validate_beta_power(name: &'static str, value: f64, step: u64) -> Result<(), AdamWStateError> {
let valid = if step == 0 {
value.to_bits() == 1.0_f64.to_bits()
} else {
value.is_finite() && (0.0..1.0).contains(&value)
};
if valid {
Ok(())
} else {
Err(AdamWStateError::InvalidBetaPower { name, value, step })
}
}
/// A malformed optimizer snapshot rejected before an AdamW instance is built.
#[derive(Clone, Debug, PartialEq)]
pub enum AdamWStateError {
InvalidBetaPower {
name: &'static str,
value: f64,
step: u64,
},
StatePresence {
step: u64,
entries: usize,
},
EmptyParameterName,
DuplicateParameterName {
name: String,
},
ShapeProductOverflow {
name: String,
},
MomentLengthMismatch {
name: String,
shape: Vec<usize>,
expected: usize,
first: usize,
second: usize,
},
NonFiniteMoment {
name: String,
kind: &'static str,
index: usize,
value: f64,
},
NegativeSecondMoment {
name: String,
index: usize,
value: f64,
},
ParameterGroupsMismatch {
expected: Vec<String>,
actual: Vec<String>,
},
}
impl fmt::Display for AdamWStateError {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
match self {
Self::InvalidBetaPower { name, value, step } => write!(
formatter,
"{name} power must be exactly 1 at step zero or finite in [0,1) later, got {value} at step {step}"
),
Self::StatePresence { step, entries } => write!(
formatter,
"AdamW step {step} is incompatible with {entries} persisted moment entries"
),
Self::EmptyParameterName => {
formatter.write_str("an AdamW persistence entry has an empty parameter name")
}
Self::DuplicateParameterName { name } => write!(
formatter,
"AdamW persistence parameter name {name:?} appears more than once"
),
Self::ShapeProductOverflow { name } => write!(
formatter,
"AdamW persistence shape for parameter {name:?} overflows usize"
),
Self::MomentLengthMismatch {
name,
shape,
expected,
first,
second,
} => write!(
formatter,
"AdamW persistence parameter {name:?} with shape {shape:?} needs {expected} moments, got {first} first and {second} second"
),
Self::NonFiniteMoment {
name,
kind,
index,
value,
} => write!(
formatter,
"AdamW persistence parameter {name:?} has non-finite {kind} moment at flat index {index}: {value}"
),
Self::NegativeSecondMoment { name, index, value } => write!(
formatter,
"AdamW persistence parameter {name:?} has negative second moment at flat index {index}: {value}"
),
Self::ParameterGroupsMismatch { expected, actual } => write!(
formatter,
"AdamW persistence groups name {expected:?}, but moments name {actual:?}"
),
}
}
}
impl Error for AdamWStateError {} CheckpointTokenizer is sealed after construction. Its literal constructor
rejects an empty vocabulary, an empty piece, and repeated byte spellings. Its
byte-BPE constructor accepts an already validated tokenizer; untrusted decoded
pairs are accepted only after the existing BPE constructor validates them during
loading. Each valid path records the resulting vocabulary size, and the private
representation has no mutator that could make that size stale. The file also
records the tokenizer-layout and SplitMix64 algorithm versions. The saved
sampling-RNG state continues only the later sampling stream, not the Chapter 33
batch-shuffle seed.
Checkpoint creation and loading start from different ownership situations.
Checkpoint::from_snapshot borrows the trainer-issued selected training state,
which seals a model snapshot, its matching AdamW snapshot, and their shared step.
Because the trainer-owned capture remains available afterward, checkpoint
construction copies those buffers into an independently owned checkpoint.
Checkpoint::from_bytes instead returns a checkpoint that
owns the tensor buffers decoded from the file.
restore_independent_model copies because the checkpoint must remain available;
into_model consumes the checkpoint and moves its model buffers. Both produce
the same parameter values. Either model-restoration boundary creates the live
tied embedding/output relationship: restore_independent_model first copies the
graph-free state, whereas into_model moves its buffers. The checkpoint state
itself has one token_embedding.weight value slot, no separate output-head
parameter, and no live component alias.
rust/crates/llm-from-scratch/src/checkpoint.rs#checkpoint-state-transfer /// Copies one trainer-issued selected model/AdamW/step bundle.
///
/// The sealed bundle prevents callers from attaching a free step label or an
/// independently chosen optimizer to the selected model state.
pub fn from_snapshot(
tokenizer: CheckpointTokenizer,
selected: &SelectedTrainingState,
sampling_rng_state: u64,
) -> Result<Self, CheckpointError> {
let selected_step =
u64::try_from(selected.step()).map_err(|_| CheckpointError::SizeOverflow {
context: "selected training step",
})?;
Self::from_owned_parts(
tokenizer,
selected.model_state().independent_snapshot(),
selected.optimizer_state().clone(),
selected_step,
sampling_rng_state,
)
}
fn from_owned_parts(
tokenizer: CheckpointTokenizer,
model_state: DecoderModelState,
optimizer_state: AdamWState,
selected_step: u64,
sampling_rng_state: u64,
) -> Result<Self, CheckpointError> {
let checkpoint = Self {
tokenizer,
model_state,
optimizer_state,
selected_step,
sampling_rng_state,
};
checkpoint.validate_parts()?;
Ok(checkpoint)
}
pub const fn tokenizer(&self) -> &CheckpointTokenizer {
&self.tokenizer
}
pub const fn model_state(&self) -> &DecoderModelState {
&self.model_state
}
pub const fn optimizer_state(&self) -> &AdamWState {
&self.optimizer_state
}
pub const fn selected_step(&self) -> u64 {
self.selected_step
}
pub const fn sampling_rng_state(&self) -> u64 {
self.sampling_rng_state
}
/// Compatibility accessor for the version-1 sampling RNG state.
pub const fn rng_state(&self) -> u64 {
self.sampling_rng_state()
}
/// Rebuilds an independent decoder while retaining the checkpoint.
pub fn restore_independent_model(&self) -> Result<DecoderModel, CheckpointError> {
self.model_state
.restore_independent_model()
.map_err(Into::into)
}
/// Consumes the checkpoint and moves its model buffers into one decoder.
pub fn into_model(self) -> Result<DecoderModel, CheckpointError> {
self.model_state.into_model().map_err(Into::into)
}
pub fn restore_optimizer(&self) -> AdamW {
AdamW::from_persistence_state(&self.optimizer_state)
} Checkpoint construction and untrusted loading establish the tokenizer/model/
optimizer relationships before exposing a checkpoint. Because the checkpoint’s
fields are private and expose no mutating access, encode does not repeat that
semantic validation. It creates a record plan whose descriptors own their names,
shapes, roles, dtypes, and eventual offsets. Each TensorPayload borrows either
literal-token bytes, a BPE-pair slice, or an f64 slice; model tensors and AdamW
moments both use the f64 variant. The plan creates no encoded byte vector per
record.
rust/crates/llm-from-scratch/src/checkpoint.rs#checkpoint-record-planning #[derive(Clone, Copy, Debug)]
enum TensorPayload<'a> {
Bytes(&'a [u8]),
BpePairs(&'a [TokenPair]),
Float64(&'a [f64]),
}
impl TensorPayload<'_> {
const fn dtype(self) -> CheckpointDType {
match self {
Self::Bytes(_) => CheckpointDType::U8,
Self::BpePairs(_) => CheckpointDType::U32,
Self::Float64(_) => CheckpointDType::F64,
}
}
fn byte_len(self) -> Result<usize, CheckpointError> {
match self {
Self::Bytes(values) => Ok(values.len()),
Self::BpePairs(pairs) => pairs
.len()
.checked_mul(2)
.and_then(|values| values.checked_mul(CheckpointDType::U32.byte_width()))
.ok_or(CheckpointError::SizeOverflow {
context: "BPE pair payload",
}),
Self::Float64(values) => values
.len()
.checked_mul(CheckpointDType::F64.byte_width())
.ok_or(CheckpointError::SizeOverflow {
context: "f64 tensor payload",
}),
}
}
fn write_to(self, bytes: &mut Vec<u8>) {
match self {
Self::Bytes(values) => bytes.extend_from_slice(values),
Self::BpePairs(pairs) => {
for pair in pairs {
put_u32(bytes, pair.left());
put_u32(bytes, pair.right());
}
}
Self::Float64(values) => {
for &value in values {
put_f64(bytes, value);
}
}
}
}
}
/// Owns record metadata while borrowing the values that will become file bytes.
#[derive(Debug)]
struct TensorRecordPlan<'a> {
descriptor: CheckpointTensorDescriptor,
payload: TensorPayload<'a>,
} The encoder measures a provisional header, assigns every checked absolute
offset, and computes the complete byte count. It then reserves one final
Vec<u8> sized for the complete header-plus-payload file, writes the header and
descriptor table, and converts each referenced payload value directly into that
final buffer in canonical little-endian order. FNV-1a covers the complete file
while the checksum field itself is treated as zero.
rust/crates/llm-from-scratch/src/checkpoint.rs#versioned-checkpoint-encoding /// Encodes one canonical little-endian file without native-memory casts.
pub fn encode(&self) -> Result<EncodedCheckpoint, CheckpointError> {
let mut records = self.tensor_record_plan()?;
let model_parameter_count = self.model_state.parameter_names().len();
let optimizer_state_count = self.optimizer_state.parameter_names().len();
let mut provisional_header = Vec::new();
write_fixed_header(&mut provisional_header, 0, 0, 0);
write_metadata(
&mut provisional_header,
self,
model_parameter_count,
optimizer_state_count,
&records,
)?;
let header_bytes = usize_to_u64(provisional_header.len(), "header length")?;
let mut next_offset = header_bytes;
let mut payload_bytes = 0_u64;
for record in &mut records {
record.descriptor.offset = next_offset;
next_offset = next_offset.checked_add(record.descriptor.byte_len).ok_or(
CheckpointError::SizeOverflow {
context: "tensor end offset",
},
)?;
payload_bytes = payload_bytes
.checked_add(record.descriptor.byte_len)
.ok_or(CheckpointError::SizeOverflow {
context: "payload length",
})?;
}
let total_bytes =
header_bytes
.checked_add(payload_bytes)
.ok_or(CheckpointError::SizeOverflow {
context: "complete file length",
})?;
let total_capacity = u64_to_usize(total_bytes, "complete file length")?;
let mut bytes = Vec::new();
bytes
.try_reserve_exact(total_capacity)
.map_err(|_| CheckpointError::Allocation {
context: "complete checkpoint file",
})?;
write_fixed_header(&mut bytes, header_bytes, payload_bytes, 0);
write_metadata(
&mut bytes,
self,
model_parameter_count,
optimizer_state_count,
&records,
)?;
if bytes.len() != u64_to_usize(header_bytes, "header length")? {
return Err(CheckpointError::Layout(
"two-pass header length changed".to_owned(),
));
}
for record in &records {
record.payload.write_to(&mut bytes);
}
if bytes.len() != total_capacity {
return Err(CheckpointError::Layout(
"encoded file length changed after layout".to_owned(),
));
}
let checksum = checkpoint_checksum(&bytes);
bytes[CHECKSUM_OFFSET..CHECKSUM_OFFSET + CHECKSUM_WIDTH]
.copy_from_slice(&checksum.to_le_bytes());
Ok(EncodedCheckpoint {
bytes,
header_bytes,
checksum,
tensors: records
.into_iter()
.map(|record| record.descriptor)
.collect(),
})
} This is not allocation-free or zero-copy serialization. Separate descriptor and
provisional-header allocations remain, and the complete encoded file needs its
final byte buffer. Numeric values are converted to explicit little-endian bytes
rather than reinterpreted in native memory. It is also not streaming to disk:
save_atomic first obtains the complete encoded buffer, then writes and
synchronizes that buffer in the destination directory. Encoding no longer
materializes a separate encoded payload buffer for every record; collectively,
those removed record buffers previously retained one extra copy of all payload
bytes before the final file was assembled.
The reader checks the fixed header, complete extent, checksum, known roles and
dtypes, and descriptor ranges before decoding owned state. After tokenizer
decoding, every model descriptor must have the model-parameter role and f64
dtype. Its bytes become one Tensor, and the ordinary leaf check validates the
parameter name and finite scalar values at this loading stage.
DecoderModelState retains those buffers and exposes its ordered list through
DecoderParameterSource. The shared decoder-layout validator reads each name
and tensor through a scoped reference and checks the configuration, parameter
count, order, names, shapes, and one token_embedding.weight slot. It does not
copy a tensor, construct NamedParameter handles, build decoder components, or
create a live embedding/output alias. This layout check finishes before the
reader decodes optimizer tensors.
rust/crates/llm-from-scratch/src/training/trainer.rs#decoder-state-layout-validation impl DecoderParameterSource for DecoderModelState {
fn len(&self) -> usize {
self.parameters.len()
}
fn name(&self, index: usize) -> &str {
&self.parameters[index].name
}
fn with_tensor<R>(&self, index: usize, inspect: impl FnOnce(&Tensor) -> R) -> R {
inspect(&self.parameters[index].value)
}
}
impl DecoderModelState {
/// Builds graph-free state after the caller has validated every parameter leaf.
pub(crate) fn try_from_leaf_validated_parameters(
config: DecoderModelConfig,
parameters: Vec<(String, Tensor)>,
) -> Result<Self, TrainerError> {
let state = Self {
config,
parameters: parameters
.into_iter()
.map(|(name, value)| StateParameter { name, value })
.collect(),
};
validate_parameter_layout(config, &state)?;
Ok(state)
}
} The reader then validates the optimizer, the relationships among tokenizer,
model, and optimizer, and exact canonical re-encoding. Only after every check
does it return the Checkpoint that owns the decoded model and optimizer
buffers. A failure at any stage returns no Checkpoint.
rust/crates/llm-from-scratch/src/checkpoint.rs#validated-checkpoint-loading /// Rejects the complete file before exposing any partially restored state.
pub fn from_bytes(bytes: &[u8]) -> Result<Self, CheckpointError> {
if bytes.len() < FIXED_HEADER_BYTES {
return Err(CheckpointError::Truncated {
context: "fixed header",
});
}
if bytes[..CHECKPOINT_MAGIC.len()] != CHECKPOINT_MAGIC {
return Err(CheckpointError::InvalidMagic);
}
let version = u16::from_le_bytes(
bytes[8..10]
.try_into()
.expect("the fixed header length was checked"),
);
if version != CHECKPOINT_VERSION {
return Err(CheckpointError::UnsupportedVersion { found: version });
}
let endian = u32::from_le_bytes(
bytes[10..14]
.try_into()
.expect("the fixed header length was checked"),
);
if endian != LITTLE_ENDIAN_MARKER {
return Err(CheckpointError::UnsupportedEndianness { found: endian });
}
let header_bytes = read_fixed_u64(bytes, 14);
let payload_bytes = read_fixed_u64(bytes, 22);
let stored_checksum = read_fixed_u64(bytes, CHECKSUM_OFFSET);
let declared_total = header_bytes.checked_add(payload_bytes);
if header_bytes < usize_to_u64(FIXED_HEADER_BYTES, "fixed header length")?
|| declared_total != Some(usize_to_u64(bytes.len(), "file length")?)
{
return Err(CheckpointError::FileExtent {
header: header_bytes,
payload: payload_bytes,
actual: bytes.len(),
});
}
let header_end = u64_to_usize(header_bytes, "header length")?;
if header_end > bytes.len() {
return Err(CheckpointError::FileExtent {
header: header_bytes,
payload: payload_bytes,
actual: bytes.len(),
});
}
let actual_checksum = checkpoint_checksum(bytes);
if stored_checksum != actual_checksum {
return Err(CheckpointError::ChecksumMismatch {
expected: stored_checksum,
actual: actual_checksum,
});
}
let header = decode_variable_header(&bytes[FIXED_HEADER_BYTES..header_end])?;
validate_descriptors(&header, header_bytes, bytes.len())?;
let tokenizer_count = match header.tokenizer_kind {
1 => header.tokenizer_vocabulary,
2 => 1,
tag => return Err(CheckpointError::UnknownTokenizerKind { tag }),
};
let tokenizer_end = tokenizer_count;
let model_end = tokenizer_end
.checked_add(header.model_parameter_count)
.ok_or(CheckpointError::SizeOverflow {
context: "model descriptor boundary",
})?;
let optimizer_descriptor_count =
header
.optimizer_state_count
.checked_mul(2)
.ok_or(CheckpointError::SizeOverflow {
context: "optimizer descriptor count",
})?;
let expected_total = model_end.checked_add(optimizer_descriptor_count).ok_or(
CheckpointError::SizeOverflow {
context: "optimizer descriptor boundary",
},
)?;
if header.descriptors.len() != expected_total {
return Err(CheckpointError::Layout(format!(
"descriptor table has {} records, expected {expected_total}",
header.descriptors.len()
)));
}
let tokenizer = decode_tokenizer(
header.tokenizer_kind,
header.tokenizer_vocabulary,
&header.descriptors[..tokenizer_end],
bytes,
)?;
let model_state = decode_model(
header.config,
&header.descriptors[tokenizer_end..model_end],
bytes,
)?;
let optimizer_state = decode_optimizer(
header.optimizer_config,
header.optimizer_groups,
header.optimizer_step,
header.beta1_power,
header.beta2_power,
&header.descriptors[model_end..],
bytes,
)?;
let checkpoint = Self::from_owned_parts(
tokenizer,
model_state,
optimizer_state,
header.selected_step,
header.sampling_rng_state,
)?;
let canonical = checkpoint.encode()?;
if canonical.bytes() != bytes {
return Err(CheckpointError::NonCanonical(
"decoded state does not reproduce the original bytes".to_owned(),
));
}
Ok(checkpoint)
} Saving creates a unique file in the destination directory with create_new,
writes and synchronizes the complete bytes, renames it over the destination,
then synchronizes the directory. That is atomic replacement under the supported
Unix same-filesystem rename semantics. Other targets return an explicit
unsupported error instead of deleting the prior checkpoint.
rust/crates/llm-from-scratch/src/checkpoint.rs#atomic-checkpoint-save /// Replaces one file through a synchronized same-directory temporary name.
///
/// The course's supported Unix workflow relies on same-filesystem rename
/// semantics. Other targets return an explicit error instead of deleting an
/// existing destination before a non-atomic move.
pub fn save_atomic(
&self,
path: impl AsRef<Path>,
) -> Result<EncodedCheckpoint, CheckpointError> {
#[cfg(not(unix))]
{
let _ = path;
return Err(CheckpointError::UnsupportedAtomicReplacement);
}
#[cfg(unix)]
{
let path = path.as_ref();
let file_name = path.file_name().ok_or_else(|| {
CheckpointError::Layout("checkpoint destination has no file name".to_owned())
})?;
let parent = path
.parent()
.filter(|parent| !parent.as_os_str().is_empty())
.unwrap_or_else(|| Path::new("."));
let encoded = self.encode()?;
let (temporary_path, mut temporary) =
create_temporary(parent, file_name.to_string_lossy().as_ref())?;
let publication = (|| -> Result<(), CheckpointError> {
temporary
.write_all(encoded.bytes())
.map_err(|source| CheckpointError::Io {
operation: "write temporary",
path: temporary_path.clone(),
source,
})?;
temporary.sync_all().map_err(|source| CheckpointError::Io {
operation: "synchronize temporary",
path: temporary_path.clone(),
source,
})?;
drop(temporary);
fs::rename(&temporary_path, path).map_err(|source| CheckpointError::Io {
operation: "atomically replace",
path: path.to_owned(),
source,
})?;
File::open(parent)
.and_then(|directory| directory.sync_all())
.map_err(|source| CheckpointError::Io {
operation: "synchronize parent directory",
path: parent.to_owned(),
source,
})?;
Ok(())
})();
if publication.is_err() {
let _ = fs::remove_file(&temporary_path);
}
publication?;
Ok(encoded)
}
} The executable fixture needs two independent component-replay branches. It retains
the original checkpoint for later corruption and atomic-save evidence, so that
branch uses restore_independent_model. The loaded checkpoint has no later
owner: after obtaining the AdamW state and saved sampling-RNG value needed by
the loaded branch, the fixture consumes it with into_model and moves its
tensors into the second decoder. The caller gives both branches inputs ,
targets , and learning rate . The fixture directly invokes loss,
backward, and one AdamW update, then compares every parameter bit, exact optimizer
state, post-update logits, and the next sampling SplitMix64 draw. It does not
call train_decoder, restore a corpus or batch cursor, or apply the Chapter 33
learning-rate schedule, clipping, or validation policy.
rust/demos/ch35-checkpoints/src/lib.rs#learner-evidence pub fn learner_evidence() -> Result<LearnerEvidence, FixtureError> {
let selected = selection_evidence()?;
let selected_training_state = selected.result.selected_training_state();
let selected_step = u64::try_from(selected.result.selected_step())
.map_err(|_| FixtureError::Invariant("selected step does not fit u64"))?;
require(
selected_step == selected_training_state.optimizer_state().step_count(),
"trainer-issued selected model and optimizer no longer share one step",
)?;
let mut sampling_rng = SplitMix64::from_seed(SAMPLING_RNG_SEED);
let _ = sampling_rng.next_u64();
let saved_sampling_rng_state = sampling_rng.state();
let expected_sampling_rng_next = sampling_rng.next_u64();
let checkpoint = Checkpoint::from_snapshot(
literal_tokenizer()?,
selected_training_state,
saved_sampling_rng_state,
)?;
let encoded = checkpoint.encode()?;
let repeated = checkpoint.encode()?;
let loaded = Checkpoint::from_bytes(encoded.bytes())?;
let loaded_encoded = loaded.encode()?;
let original_model = checkpoint.restore_independent_model()?;
let loaded_optimizer = loaded.restore_optimizer();
let loaded_sampling_rng_state = loaded.sampling_rng_state();
let loaded_model = loaded.into_model()?;
let original_logits = logits_bits(&original_model)?;
let loaded_logits = logits_bits(&loaded_model)?;
let (updated_original, updated_original_optimizer) = apply_component_update(
original_model,
checkpoint.restore_optimizer(),
&LOGIT_INPUTS,
&LOGIT_TARGETS,
NEXT_LEARNING_RATE,
)?;
let (updated_loaded, updated_loaded_optimizer) = apply_component_update(
loaded_model,
loaded_optimizer,
&LOGIT_INPUTS,
&LOGIT_TARGETS,
NEXT_LEARNING_RATE,
)?;
let updated_original_bits = parameter_bits(&updated_original);
let updated_loaded_bits = parameter_bits(&updated_loaded);
let updated_original_logits = logits_bits(&updated_original)?;
let updated_loaded_logits = logits_bits(&updated_loaded)?;
let (changed_batch_model, changed_batch_optimizer) = apply_component_update(
checkpoint.restore_independent_model()?,
checkpoint.restore_optimizer(),
&[1, 2],
&[2, 3],
NEXT_LEARNING_RATE,
)?;
let changed_batch_diverges = parameter_bits(&changed_batch_model) != updated_original_bits
|| changed_batch_optimizer != updated_original_optimizer;
let (changed_learning_rate_model, changed_learning_rate_optimizer) = apply_component_update(
checkpoint.restore_independent_model()?,
checkpoint.restore_optimizer(),
&LOGIT_INPUTS,
&LOGIT_TARGETS,
NEXT_LEARNING_RATE * 2.0,
)?;
let changed_learning_rate_diverges = parameter_bits(&changed_learning_rate_model)
!= updated_original_bits
|| changed_learning_rate_optimizer != updated_original_optimizer;
let loaded_sampling_rng_next = SplitMix64::from_state(loaded_sampling_rng_state).next_u64();
let corruption = corruption_evidence(selected_training_state, &checkpoint, encoded.bytes())?;
let atomic = atomic_evidence(selected_training_state, &checkpoint)?;
require(
encoded.tensors().len() == 38,
"checkpoint record count changed",
)?;
require(
checkpoint.model_state().parameter_names().len() == 11,
"model parameter tensor count changed",
)?;
require(
checkpoint.model_state().scalar_count() == 144,
"model scalar count changed",
)?;
require(
updated_original_optimizer.step_count() == selected_step + 1,
"component-replay optimizer step count changed",
)?;
require(
changed_batch_diverges && changed_learning_rate_diverges,
"caller-supplied batch or learning-rate adversary no longer diverges",
)?;
require(
corruption.version_rejected
&& corruption.vocabulary_mismatch_rejected
&& corruption.step_mismatch_rejected
&& corruption.truncation_rejected
&& corruption.checksum_rejected,
"one corruption fixture was accepted",
)?;
require(
atomic.replaced_complete_file && atomic.temporary_files == 0,
"atomic replacement evidence changed",
)?;
let history = historical_checkpoint_contrast(&checkpoint, &encoded);
require(
history.isolated_parameter_bytes.len()
== history.isolated_parameter_scalars * std::mem::size_of::<f64>(),
"isolated parameter byte length changed",
)?;
require(
history.tokenizer_records
+ history.isolated_parameter_tensors
+ history.optimizer_moment_records
== history.checkpoint_records,
"checkpoint record-family counts changed",
)?;
Ok(LearnerEvidence {
checkpoint,
bytes_deterministic: encoded == repeated,
round_trip_identical: encoded == loaded_encoded,
logits_before_bits_identical: original_logits == loaded_logits,
logits_before_fingerprint: bits_fingerprint(&loaded_logits),
parameter_bits_after_identical: updated_original_bits == updated_loaded_bits,
optimizer_after_identical: updated_original_optimizer == updated_loaded_optimizer,
logits_after_bits_identical: updated_original_logits == updated_loaded_logits,
logits_after_fingerprint: bits_fingerprint(&updated_loaded_logits),
changed_batch_diverges,
changed_learning_rate_diverges,
sampling_rng_next_identical: expected_sampling_rng_next == loaded_sampling_rng_next,
sampling_rng_next: loaded_sampling_rng_next,
corruption,
atomic,
history,
encoded,
})
} The executable report records these five measured boundaries:
roundtrip=bytes_deterministic:true loaded_bytes_identical:true logits_bits_identical:true logits_fingerprint:fnv1a64:6029064fe7cd162d sampling_rng_next_identical:true sampling_rng_next:0x9a8c505971939232
component_replay=caller_inputs:[0,1] caller_targets:[1,2] caller_learning_rate:0.006000 next_step:9 parameter_bits_identical:true optimizer_state_identical:true logits_bits_identical:true logits_fingerprint:fnv1a64:0b875a0c9f380d8f changed_batch_diverges:true changed_learning_rate_diverges:true
scope=tokenizer:stored model:stored optimizer:stored selected_step:stored optimizer_step:stored optimizer_base_learning_rate:stored sampling_rng:stored step_equality:validated model_lineage:not_stored corpus_identity:not_stored split_identity:not_stored epoch_materialization:not_stored epoch_cursor:not_stored batch_order:not_stored batch_cursor:not_stored shuffle_rng:not_stored training_rng:not_stored learning_rate_schedule:not_stored next_learning_rate:not_stored clipping_policy:not_stored validation_policy:not_stored gradients:not_stored trainer_capture:creation_required caller_next_batch:required caller_next_learning_rate:required clean_post_update:required whole_job_resume:false
reject=version:true vocabulary_mismatch:true step_mismatch:true truncation:true checksum:true
atomic=replaced_complete_file:true loaded_sampling_rng_state:0x9e3779b97f4a7c39 temporary_files:0 unix_same_directory:true
The divergence fields make the condition observable: changing either the batch or the learning rate changes the result. Equality belongs to this one specified update, not to an unstored trainer trajectory.
The complete executable example also checks the byte-BPE variant, mixed element widths, version and endian drift, metadata corruption, trailing bytes, tokenizer and model mismatches, negative second moments, and deterministic re-encoding.
rust/demos/ch35-checkpoints/src/main.rs fn main() -> Result<(), Box<dyn std::error::Error>> {
print!("{}", ch35_checkpoints::learner_report()?);
Ok(())
} Audit the byte layout before trusting the payload
A checkpoint loader needs exact record order, element width, shape, and byte range. A flow diagram would discard those details, while the descriptor table keeps them available for direct inspection.
Audit the table from top to bottom. Each row must begin exactly where the previous row ends, its width must agree with its dtype and shape, and the final exclusive end must equal the declared file length. The fixed -byte prefix then lets you cross-check the magic, schema version, endian marker, and start of the header length against the same bytes the loader receives.
The rejection line tests complementary failure boundaries: an unknown version,
an incompatible vocabulary, a mismatched recorded step, a truncated range, and
a changed checked byte. Together with roundtrip, component_replay, and
scope, the result separates stored-state replay, one caller-supplied update,
and the trainer state that remains outside the file.
Read the evidence conservatively. logits_bits_identical:true establishes exact
replay for this fixture. checksum:true shows that one modified checked byte is
rejected; it does not turn FNV-1a into a cryptographic signature.
Predict offsets before loading corruptions
- Start at offset with a
f64tensor. Predict its exclusive end. - A BPE-pair tensor stores
u32. How many bytes does it occupy? - One literal token grows from one byte to four. Which later offsets change?
- Parameter values stay intact but a descriptor shape changes. Should loading succeed when the checksum covers metadata and values?
- Weights reload exactly but AdamW moments reset. Will the next update match?
- An untrusted file declares selected step while its AdamW counter remains at step . Which validation invariant rejects it?
- A process stops after the temporary file is synchronized but before rename. Which complete destination remains visible?
- Does a matching FNV-1a value prove that no adversary changed the file?
Check your reasoning
- , because .
- bytes.
- The changed token’s end and every subsequent absolute offset move by three bytes; canonical re-encoding also changes the checked metadata.
- No. The metadata checksum changes, and semantic shape validation would reject a freshly checksummed inconsistent descriptor too.
- No. AdamW’s next adaptive direction depends on its first and second moments and accumulated beta powers.
- Checkpoint validation requires the recorded selected step and AdamW counter to agree.
- The prior complete destination remains visible; the unpublished temporary file can be cleaned up.
- No. FNV-1a detects accidental corruption but provides no authentication.
The central misconception is that a checkpoint is either only weights or a
complete trainer restart. This file stores enough to interpret the decoder,
restore its trainer-paired AdamW state, and continue a separate sampling stream.
Full training continuation would also need the caller-owned corpus and split,
batch order and cursor, training RNG, update inputs and targets, next learning
rate and schedule, and clipping and validation policy. The demo supplies one
batch and learning rate manually; it does not resume train_decoder.
Load the same state before choosing a token
The cumulative decoder can now leave memory with the same tokenizer meaning, architecture, parameter bits, trainer-paired AdamW state and shared step, and a separate sampling stream. The loader restores model meaning before any consumer can use the values. Replaying one specified update still requires the caller to supply the same inputs, targets, and learning rate.
Chapter 36 will load this checkpoint before converting logits into token choices with temperature and top- sampling. That next boundary can vary sampling policy without silently changing tokenizer, model configuration, weights, trainer-paired AdamW state, or the sampling stream being continued. It does not turn this file into a Chapter 33 trainer restart.