Training and Inference Systems › Data, Checkpoint, and Experiment Infrastructure
World Models HandbookAll chapters Training and Inference Systems › Data, Checkpoint, and Experiment InfrastructureOne section · Sections in this chapter3 / 8 ↓
Reading tools Definitions & inline links
Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.
Article links
Make inline references clickable
Distributed training supplies compute; data and artifact infrastructure supply the intended inputs and evidence of what happened. This section connects data loading, checkpointing and experiment tracking. Exact continuation requires the causally relevant training state to be captured or reconstructible. A stochastic loader can be replayable if its random streams, ordering and pending work are recorded. An incomplete checkpoint may still support some comparisons, but cannot justify exact continuation. Code identity can come from a commit, a content hash or an archived executable snapshot; losing every usable copy prevents faithful reconstruction. These records also support audits, without establishing legal compliance by themselves.
Data Loading for World-Model Trajectories
For action-conditioned reinforcement-learning setups, training trajectories can contain temporally aligned sequences. Two properties distinguish such trajectory data from typical image classification pipelines:
- Temporal association. Dynamics learning needs evidence linking states, actions and subsequent outcomes. Sampling isolated frames while discarding their episode and time associations removes those observed transition pairs, even though individual images can still support representation learning or prior knowledge. A sampling unit can be a contiguous window or an episode; its length is a workload choice, not a fixed rule.
- Heterogeneous modalities. A sample may combine RGB, depth, proprioception, torques, actions, language and metadata. Their decoding and transforms can differ. Whether implemented as separate paths or one dispatching loader, the requirement is to preserve temporal alignment and sample identity across the streams.
A loader can be organized as an assembly line: retrieve and decode data, sample contiguous windows, transform aligned modalities, collate tensors and masks, then prefetch the next batches. Record augmentation randomness when replay matters. Choose prefetch depth from service times, memory and recovery requirements rather than always buffering the same number of batches. Any stage can become the bottleneck; measure where the consumer waits.
The main failure patterns here:
- Shard hotspots. Assigning entire unequal-work shards to workers can create imbalance; unequal file sizes alone do not prove it if workers process equal-cost windows. Balance assignments using measured decode and sample cost, and distinguish load-balancing weights from statistical sampling weights. Exposed synchronization waits for required ranks, so an overloaded rank can limit throughput. Overlap and other serial work determine the magnitude.
- Repeated retrieval overhead. Many small requests and repeated decoding can cost more than fetching and reusing a shard. Whole-shard reads can amortize overhead when many windows reuse the bytes; indexed range reads can be preferable when little of a large shard is needed. Benchmark request concurrency, locality, transferred bytes and decode reuse rather than assume one access pattern is always cheaper.
- Decode bottlenecks. When the decoder cannot supply batches at the model's consumption rate, it can starve the accelerator. Pre-decoding, supported hardware decoders or parallel CPU workers are possible remedies, with storage and format tradeoffs. A large slow model may instead be the bottleneck, so buffering and parallelism should follow measurements.
Determinism needs more than a dataset seed. Record the dataset manifest and version, sampler algorithm, epoch and cursor, augmentation RNG states, rank count and work assignment, plus any streaming offsets or prefetched items that affect the next batch. A different shuffle order is a different realized optimization sequence, not necessarily a different underlying data distribution. Framework kernels, releases and hardware can also affect numerical reproducibility, as PyTorch's reproducibility notes explain. The restricted example below tests a fixed CPU setting; it does not promise replay across arbitrary worker counts or platforms.
Checkpointing
A checkpoint serializes the state needed to continue training: parameters, optimizer moments, relevant schedules, random streams, data position and any auxiliary network state the algorithm uses. Under a fixed replayable configuration, restoring complete state can reproduce subsequent updates. Missing state can repeat data or change updates without an immediate error. That motivates a continuation test near the save/load code rather than relying on a plausible-looking loss curve.
Three concepts are essential:
- Atomic publication and durability. Publish a checkpoint as current only after every required snapshot file is complete and validated. A partial file may fail to load, or may pass inadequate checks. On a supporting local filesystem, a same-filesystem temporary-file replacement can provide atomic visibility; it does not by itself prove persistence after power loss. Durable recovery additionally depends on file and directory synchronization or the storage service's committed-write guarantees. Distributed shards need a protocol that publishes one complete snapshot identity, such as an atomically committed manifest, not a promise that every storage system supports rename.
- Consistency across ranks. Snapshot one agreed logical optimizer-step boundary, including all sharded states. A backward all-reduce does not establish that every subsequent optimizer update or asynchronous write has completed. Complete the relevant updates, coordinate ranks and capture an immutable consistent snapshot before committing the checkpoint manifest. Replicated parameters and sharded state require different save protocols.
- Tiering and pruning. One retention policy saves every steps, retains the last , and promotes selected snapshots to longer-term storage. Choose retention from recovery, reproducibility and audit requirements; keeping every snapshot can be justified. Provenance comes from an explicit deployed checkpoint identity and traceable records, not from deleting alternatives. Storage tier and retained recovery points affect retrieval cost and lost work.
Two gotchas specific to world models:
- Data loader state. Preserve enough state to reconstruct the next data selection, not just the model tensors. Replaying or skipping windows changes the realized update sequence, although it need not introduce statistical bias in every sampling scheme. Some pipelines use a deterministic global sampler keyed on step so the loader state is reconstructible from a step number and unchanged dataset and sampler configuration. In a streaming or prefetched loader, reconstructing the next selection may also require pending work and random-stream state. Model and data state together determine continuation.
- Auxiliary learning state. Save whichever target networks, slow critics, EMA copies, normalization statistics and update counters the specific algorithm actually uses. Dreamer, TD-MPC and MuZero variants do not all maintain the same EMA copy of the entire model. A missing auxiliary tensor or schedule position can change learning after resumption even if the primary parameters load correctly. Audit the implemented algorithm rather than infer checkpoint contents from its family name.
Checkpoint size depends on the actual saved tensors. As an illustration, suppose a checkpoint stores fp16 weights, a separate fp32 master copy and two fp32 Adam moments. That is bytes per parameter, excluding auxiliary tensors and serialization overhead. The separate master-weight approach is described in Mixed Precision Training. It is not universal: ordinary PyTorch autocast can retain fp32 parameters without storing an additional fp16 copy. The tutorial below uses fp32 parameters and AdamW. Activations and gradients normally are not saved in an optimizer-step checkpoint. For the stated illustrative layout:
where:
- : total core training state in bytes (weights, master copy, and optimizer moments)
- : the parameter count, here
- : bytes per parameter for fp16 weights
- : bytes per parameter for the fp32 master copy
- : bytes per parameter for the two fp32 Adam moments
The stated tensors occupy 14 GB decimal, about 13.04 GiB, before auxiliary buffers or file overhead; this conversion does not round buffers into the total. Repeated writes create an I/O budget. Options include asynchronous checkpointing with an immutable snapshot and a durable completion marker, sharded checkpoints with a committed manifest, and delta checkpointing when actual encoded deltas reduce bytes. Parameter values changing slowly does not by itself make dense checkpoints smaller.
Experiment Infrastructure
Experiment infrastructure is the layer that makes a training run an experiment rather than merely a computation. Minimally, it must record:
- Configuration: model architecture, hyperparameters, seeds, code version (git commit or content hash), container image digest.
- Data: dataset version, shard manifest, preprocessing script version, licensing/consent tags for any human data.
- Metrics: training loss curves, wall-clock per step, throughput (tokens or frames per second), gradient norm, effective learning rate, memory headroom, and validation metrics captured per checkpoint.
- Artifacts: checkpoints, evaluation rollouts, any visualization dumps (a critical debugging aid for latent world models, see Interpretability and World-Model Debugging).
- Environment: accelerator type, interconnect, driver/CUDA versions, and whether mixed precision, sequence parallelism, or activation checkpointing were enabled.
These records make comparisons interpretable. A difference between two evaluation scores does not establish that one intended change caused it. Record potential confounders, use controlled comparisons and replication where appropriate, and report uncertainty. Different realized sample orders can contribute run-to-run variability even when the dataset is unchanged. Missing provenance may require rerunning an experiment, but logging is evidence for a comparison, not proof that every other factor was identical.
A Minimal Checkpoint and Data Pipeline
The following code demonstrates a scheduler-free, single-device checkpoint loop that captures the essential state, plus a sharded, deterministic data sampler. The goal is clarity, not performance, and it runs comfortably on CPU. Read it with an eye to which pieces of state are captured and why each one is necessary for a resumed run to match an uninterrupted one.
class TinyLatentDynamics(nn.Module):
"""Minimal latent world model: predicts next latent from (z, a)."""
def __init__(self, latent_dim=16, action_dim=2, hidden=64):
super().__init__()
self.latent_dim = latent_dim
self.action_dim = action_dim
self.net = nn.Sequential(
nn.Linear(latent_dim + action_dim, hidden),
nn.ReLU(),
nn.Linear(hidden, hidden),
nn.ReLU(),
nn.Linear(hidden, latent_dim),
)
def forward(self, z, a):
return self.net(torch.cat([z, a], dim=-1))
class DeterministicWindowSampler:
"""Each worker gets a deterministic, non-overlapping shard of windows."""
def __init__(
self,
n_episodes,
episode_len,
window_len,
stride,
rank=0,
world_size=1,
seed=1234,
):
for name, value in (
("n_episodes", n_episodes),
("episode_len", episode_len),
("window_len", window_len),
("stride", stride),
("world_size", world_size),
):
if (
isinstance(value, bool)
or not isinstance(value, (int, np.integer))
or value <= 0
):
raise ValueError(f"{name} must be a positive integer")
if (
isinstance(rank, bool)
or not isinstance(rank, (int, np.integer))
or not 0 <= rank < world_size
):
raise ValueError("rank must be an integer in [0, world_size)")
if window_len > episode_len:
raise ValueError("window_len cannot exceed episode_len")
windows = []
for ep in range(n_episodes):
for start in range(0, episode_len - window_len + 1, stride):
windows.append((ep, start))
rng = np.random.default_rng(seed)
rng.shuffle(windows)
self.windows = windows[rank::world_size]
def __len__(self):
return len(self.windows)
def __getitem__(self, i):
return self.windows[i]We can now define checkpoint state for this restricted CPU example. The step count identifies the next window only while the dataset, sampler seed, rank count and sampler algorithm remain unchanged. It is not sufficient for an arbitrary streaming loader. The helper assumes one writer, no concurrent training update, and a local filesystem supporting same-directory atomic replacement. Its fixed temporary name is not a concurrent-writer protocol, and it calls neither file nor directory fsync, so it is not a crash-durability test. Distributed storage requires its own consistency, commit and durability protocol. Load only a trusted checkpoint written by this example, because weights_only=False enables pickle deserialization.
def save_checkpoint(path, model, optimizer, step, sampler_seed):
checkpoint_path = Path(path)
temporary_path = checkpoint_path.with_suffix(
checkpoint_path.suffix + ".tmp"
)
torch.save(
{
"step": step,
"weights": model.state_dict(),
"optimizer": optimizer.state_dict(),
"rng_torch": torch.get_rng_state(),
"rng_numpy": np.random.get_state(),
"sampler_seed": sampler_seed,
},
temporary_path,
)
temporary_path.replace(checkpoint_path)
def load_checkpoint(path, model, optimizer):
state = torch.load(path, weights_only=False)
model.load_state_dict(state["weights"])
optimizer.load_state_dict(state["optimizer"])
torch.set_rng_state(state["rng_torch"])
np.random.set_state(state["rng_numpy"])
return state["step"], state["sampler_seed"]Let's generate a controlled synthetic state sequence, with 16 state dimensions and two action dimensions. Each episode follows , where has shape and independent process noise has standard deviation 0.02. The observations here equal the synthetic states; no encoder or learned state representation is being tested. The model fits one-step transitions, not a validated real-world simulator.
model = TinyLatentDynamics()
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3)
# Controlled synthetic dynamics, not independently sampled next-state targets.
T_gen, n_ep = 32, 64
sim_rng = np.random.default_rng(0)
actions = sim_rng.normal(size=(n_ep, T_gen, 2)).astype("float32")
action_map = sim_rng.normal(scale=0.5, size=(2, 16)).astype("float32")
latent_obs = np.zeros((n_ep, T_gen + 1, 16), dtype="float32")
latent_obs[:, 0] = sim_rng.normal(size=(n_ep, 16))
process_noise = sim_rng.normal(scale=0.02, size=(n_ep, T_gen, 16))
for t in range(T_gen):
latent_obs[:, t + 1] = (
0.8 * latent_obs[:, t]
+ 0.2 * actions[:, t] @ action_map
+ process_noise[:, t]
)
sampler = DeterministicWindowSampler(
n_episodes=n_ep, episode_len=T_gen + 1, window_len=8, stride=1, seed=7
)
def train_window(model, optimizer, sampler, window_idx):
ep, start = sampler[window_idx]
z = torch.from_numpy(latent_obs[ep, start : start + 7])
a = torch.from_numpy(actions[ep, start : start + 7])
z_next = torch.from_numpy(latent_obs[ep, start + 1 : start + 8])
pred = model(z, a)
loss = F.mse_loss(pred, z_next)
optimizer.zero_grad()
loss.backward()
optimizer.step()
return float(loss.detach())
checkpoint_step = 512
losses = [
train_window(model, optimizer, sampler, i) for i in range(checkpoint_step)
]
save_checkpoint(
"tiny_wm.ckpt.pt", model, optimizer, checkpoint_step, sampler_seed=7
)Now compare the next 16 uninterrupted updates with a fresh model and optimizer loaded from that checkpoint. Both paths use the same CPU runtime and dataset. Equal losses alone are insufficient, so we also compare all resulting model parameters. We additionally sample and compare the next Torch CPU and legacy NumPy random draws; the deterministic training continuation alone would not exercise those saved random streams.
continuation_steps = 16
expected_torch_draw = torch.rand(4)
expected_numpy_draw = np.random.standard_normal(4)
expected_losses = [
train_window(model, optimizer, sampler, i)
for i in range(checkpoint_step, checkpoint_step + continuation_steps)
]
expected_weights = {
name: value.clone() for name, value in model.state_dict().items()
}
resumed_model = TinyLatentDynamics()
resumed_optimizer = torch.optim.AdamW(resumed_model.parameters(), lr=1e-3)
resumed_step, resumed_seed = load_checkpoint(
"tiny_wm.ckpt.pt", resumed_model, resumed_optimizer
)
assert torch.equal(expected_torch_draw, torch.rand(4))
assert np.array_equal(expected_numpy_draw, np.random.standard_normal(4))
resumed_sampler = DeterministicWindowSampler(
n_ep, T_gen + 1, 8, 1, seed=resumed_seed
)
resumed_losses = [
train_window(resumed_model, resumed_optimizer, resumed_sampler, i)
for i in range(resumed_step, resumed_step + continuation_steps)
]
assert np.array_equal(expected_losses, resumed_losses)
assert all(
torch.equal(expected_weights[name], value)
for name, value in resumed_model.state_dict().items()
)
model = resumed_model
optimizer = resumed_optimizerAvailable windows: 1664; checkpoint after 512 updates Final loss: 0.0038 Mean loss over first 10 steps: 0.0889 Mean loss over last 10 steps: 0.0552 Exact loss and weight match for 16 CPU continuation updates
Compare the printed first-ten and last-ten means rather than assuming every noisy window loss decreases. Matching future losses and weights exercises the restored optimizer and sampler on the next 16 updates; it is not a direct equality assertion over every optimizer field. The next-draw assertions cover the two random streams saved here, not an independent default_rng generator or a device RNG. This is not a distributed-recovery proof: topology changes, prefetch state, schedules, scalers, interrupted writes and different software or hardware require additional tests. PyTorch's reproducibility guidance limits cross-platform and cross-release reproducibility.
The pre-checkpoint loss curve shows optimization on a controlled synthetic process, including variation between sampled windows. The separate continuation assertions, not the shape of this plot, test resumption.
