GPUs, memory, checkpoints, storage, and reproducible experiments
Learn what occupies GPU memory, how music representation and sequence length affect training cost, when memory-saving and distributed techniques help, and how to build a recoverable local, cloud, or hybrid workflow.
About 16 minutes
Guiding question
What actually occupies GPU memory during training?
By the end, you’ll be able to
Distinguish computational throughput, GPU memory capacity, system memory, and storage.
Identify the main sources of training memory use, including parameters, gradients, optimiser state, activations, and temporary buffers.
Explain how audio duration, token rate, sequence length, batch size, model width, and precision affect resource use.
Distinguish data parallelism, fully sharded data parallelism, and model or tensor parallelism.
Design checkpoints that support inference, comparison, and exact or approximate training resumption.
Compare local, cloud, and hybrid experimentation workflows.
Create an experiment-tracking and reproducibility record containing configurations, data versions, code, metrics, samples, and environment information.
A model architecture can be mathematically valid and still be impossible to train with the available hardware.
Infrastructure determines which model sizes, sequence lengths, batch sizes, datasets, and experiments can be run in practice. It also determines whether a failed job can resume, whether two runs can be compared, and whether the result can be reproduced later.
The central question is therefore not simply Which GPU is fastest? It is:
Which combination of hardware, memory strategy, storage, software, and experiment controls can run this specific training plan reliably?
Four different resource constraints
Swipe sideways to view the full comparison
Resource
What it limits
Typical symptom
GPU compute
How quickly tensor operations are completed
Low step throughput despite fitting in memory
GPU memory
How many parameters, activations, gradients, and states fit at once
Out-of-memory error
CPU and data pipeline
How quickly examples are decoded, transformed, and transferred
The GPU waits between batches
Storage and network
How quickly datasets and checkpoints can be read, written, or transferred
Slow startup, checkpoint stalls, or repeated data downloads
What occupies training memory?
Loading model weights is only the beginning. A conventional training step may require memory for:
Model parameters
Parameter gradients
Optimiser states
Forward-pass activations saved for backpropagation
Input batches, targets, masks, and conditions
Temporary operation buffers
Mixed-precision scaling or master values
Distributed communication buffers
Framework allocator reserves
The largest component depends on the model, optimiser, sequence length, batch size, precision, and implementation.
How memory changes during one training step
1
Load parameters
Model values occupy device memory before the batch begins.
2
Load the batch
Inputs, targets, conditions, lengths, and masks are transferred or created.
3
Run the forward pass
Activation memory grows as intermediate tensors are retained.
4
Run the backward pass
Activations are progressively released while gradients are created.
5
Run the optimiser
Optimiser states and temporary update tensors contribute to peak memory.
6
Clear gradients
Gradient storage can be released or reused before the next ordinary step.
Allocated and reserved memory
PyTorch uses a caching allocator for CUDA memory. Memory released by tensors may remain reserved by PyTorch so it can be reused quickly.
As a result, a system monitor can show memory as occupied even when that memory is not currently assigned to live tensors.
Useful measurements include:
Memory allocated by tensors
Memory reserved by the caching allocator
Peak allocated memory
Peak reserved memory
This distinction is important when debugging apparent memory leaks.
Audio duration becomes sequence length
A music model does not usually receive thirty seconds as an abstract unit. The representation converts those seconds into samples, frames, latent positions, or tokens.
Approximate sequence length depends on:
Duration
Sample or frame rate
Codec token rate
Number of codebooks or token streams
Channel handling
Interleaving pattern
Padding and special tokens
Doubling duration usually increases the amount of represented data, but the exact memory and compute effect depends on the architecture.
In practice
Real-world example: MusicGen token sequences
MusicGen uses a 32 kHz EnCodec representation with four codebooks sampled at 50 frames per second. Its interleaving pattern determines how the four token streams become the sequence processed by the language model. Duration, codebook count, and interleaving therefore affect training sequence length and cost.
Why long sequences are expensive
Longer sequences increase the number of token representations and layer activations that must be processed.
For conventional dense self-attention, the attention relationship between sequence positions forms a square matrix. Materialising this matrix creates memory and computational costs that grow strongly as sequence length increases.
Efficient attention implementations such as FlashAttention reduce memory traffic and avoid storing some large intermediate matrices. They improve feasibility and throughput, but they do not make long-context training free.
What raises training cost?
Swipe sideways to view the full comparison
Setting
Primary effect
Important qualification
More parameters
More parameter, gradient, and optimiser state
Frozen parameters reduce some training state but still require model storage
Longer sequence
More activations and more sequence operations
Attention implementation changes the memory profile
Larger microbatch
More examples and activations processed together
Can improve throughput if memory remains available
More audio codebooks
More tokens or prediction streams
Interleaving and architecture determine the final sequence
Higher sample or frame rate
More representation positions per second
Compression can reduce the downstream sequence
Full fine-tuning
Gradients and optimiser state for the complete trainable model
LoRA or adapters can train substantially fewer values
More generated evaluations
Additional decoding and audio storage
Evaluation cost can become significant even without backpropagation
A larger batch processes more examples together. This commonly increases activation and input memory because more examples pass through the network at once.
The microbatch size is the number of examples processed in one forward and backward pass on one worker. The effective batch size can be larger through gradient accumulation and multiple workers.
Reducing the microbatch is often the first response to an out-of-memory error, but very small microbatches can reduce hardware utilisation or change optimisation behaviour.
Gradient accumulation
Gradient accumulation divides a large effective batch into several smaller microbatches.
Each microbatch runs a forward and backward pass. Gradients are accumulated, and the optimiser updates the parameters only after the selected number of microbatches.
This lowers the activation memory required at one moment compared with processing the full effective batch together. It also adds more sequential computation and requires correct loss scaling and gradient clearing.
In practice
Effective batch size
A common approximation is: microbatch size multiplied by accumulation steps multiplied by the number of data-parallel workers. The actual optimisation behaviour can still depend on loss reduction, incomplete batches, masking, and distributed implementation.
Numerical precision
Numerical precision determines how tensor values are represented.
Common training formats include:
FP32
FP16
BF16
Mixed-precision combinations
Lower-precision operations can reduce tensor memory and improve throughput on compatible hardware. Some calculations remain in higher precision for numerical stability. FP16 training may use gradient scaling to reduce underflow risk.
Total memory does not necessarily fall in direct proportion to the bit width because optimiser states, master values, temporary tensors, and unsupported operations may remain in higher precision.
Activation checkpointing
Ordinary backpropagation keeps many forward activations in memory. Activation checkpointing saves only selected tensors and recomputes omitted forward operations during the backward pass.
This trades additional computation for lower activation memory. It is especially useful when activations, rather than parameters or optimiser state, are the main bottleneck.
This technique is unrelated to saving a model checkpoint to disk.
Two different meanings of checkpoint
Swipe sideways to view the full comparison
Term
Purpose
Where it exists
Activation checkpointing
Reduce GPU memory through recomputation
Inside the forward and backward computation
Model checkpoint
Save model or training state for later use
In persistent storage
Choosing a memory-saving technique
Swipe sideways to view the full comparison
Technique
Reduces
Trade-off
Smaller microbatch
Batch-dependent activations and inputs
Lower parallel utilisation or different batch behaviour
Gradient accumulation
Peak activation memory for a target effective batch
More sequential forward and backward passes
Mixed precision
Memory and time for supported tensor operations
Numerical-stability requirements
Activation checkpointing
Saved activation memory
Forward operations are recomputed
Freeze the base model
Gradients and optimiser state for frozen parameters
Lower adaptation capacity
LoRA or adapters
Trainable gradient and optimiser state
May not express every required model change
CPU offload
Selected GPU-resident state
Transfers can reduce speed
FSDP or ZeRO-style sharding
Replicated parameters, gradients, or optimiser states
Communication and distributed complexity
The GPU can be fast while training remains slow
Music data may require file reading, decompression, resampling, channel conversion, tokenisation, caption parsing, and random cropping before it reaches the model.
If this pipeline is slower than the GPU, the accelerator waits between batches.
Potential responses include:
Additional data-loader workers
Persistent workers
Pinned host memory where appropriate
Prefetching
Cached or precomputed codec tokens
Faster local dataset storage
Sharding large datasets
Moving expensive preprocessing outside the training loop
Each change should be measured because excessive workers or caching can create new CPU, memory, or storage bottlenecks.
From one GPU to distributed training
A single GPU is usually the simplest environment for debugging and small experiments. Distributed training becomes useful when:
The model or training state does not fit on one device
A larger global batch is required
One-device throughput makes the experiment impractically slow
Evaluation or generation needs to be parallelised
Adding GPUs introduces communication, synchronisation, process management, failure handling, and distributed checkpointing. More devices do not guarantee proportional speedup.
DDP and FSDP
Swipe sideways to view the full comparison
Property
DDP
FSDP
Model parameters
Normally replicated per worker
Sharded outside relevant computation
Gradients
Normally present for each replica before synchronisation
Sharded through reduce-scatter behaviour
Optimiser state
Normally replicated unless another technique is added
Sharded across workers
Primary benefit
Relatively direct multi-GPU data parallelism
Lower per-worker persistent model-state memory
Primary cost
Full replica must fit on each worker
Additional sharding, communication, and checkpoint complexity
Saving training progress
A checkpoint is a persistent record created at a particular training state.
Different checkpoints serve different purposes:
Generation or inference
Comparing earlier and later models
Selecting the best validation checkpoint
Resuming an interrupted training run
Preserving milestones for an audit
Rolling back after a software or data error
A weights-only checkpoint is smaller and simpler, but it generally does not contain everything needed to resume the exact optimiser trajectory.
Checkpoint types
Swipe sideways to view the full comparison
Checkpoint
Typical contents
Primary use
Weights only
Model state dictionary
Inference or later initialisation
Resumable training
Model, optimiser, scheduler, scaler, counters, and configuration
Continue an interrupted run
Best validation
State selected by a held-out metric
Model comparison and final evaluation
Periodic milestone
State at fixed steps or epochs
Audit learning progression or recover from later failure
Distributed checkpoint
Sharded state stored across distributed workers
Resume large sharded training jobs
What a resumable checkpoint should preserve
1
Model state
Store the trainable and required non-trainable tensors.
2
Optimiser state
Preserve moments and other update history.
3
Scheduler state
Resume the correct learning-rate position.
4
Precision scaler state
Preserve mixed-precision gradient-scaling progress when used.
5
Progress counters
Record epoch, global step, tokens, examples, or audio hours processed.
6
Model and data configuration
Record architecture, tokenizer, codec, conditions, and dataset version.
7
Randomness state
Preserve relevant random-generator states when a closer continuation is required.
8
Code and environment identifiers
Record the commit, package versions, platform, and hardware.
How often should checkpoints be saved?
Frequent checkpoints reduce the amount of work lost after interruption. They also increase storage use and can pause training while large states are written.
The appropriate interval depends on:
Cost per training step
Failure or interruption risk
Checkpoint size
Storage bandwidth
Validation frequency
Whether asynchronous checkpointing is available
How many historical states must be retained
Test checkpoint loading before depending on the recovery plan.
Storage can become a larger problem than the model
A music project may need to store:
Original audio and MIDI
Cleaned and resampled versions
Captions and metadata
Precomputed codec tokens or embeddings
Dataset manifests and hashes
Model checkpoints
Optimiser states
Generated evaluation samples
Similarity indexes
Logs, plots, and profiler traces
Keeping every artifact from every run can consume substantial space. Deleting everything except the final model removes evidence needed for debugging and reproducibility.
A practical storage plan
Swipe sideways to view the full comparison
Artifact
Access pattern
Suggested treatment
Current training shards
Read repeatedly during the run
Keep on fast storage close to the workers
Original source data
Read during preparation and audits
Preserve immutably with provenance records
Token or embedding cache
Read frequently but can be rebuilt
Version against the representation and preprocessing code
Latest resumable checkpoint
Needed quickly after interruption
Keep in reliable persistent storage
Best validation checkpoint
Needed for final evaluation and release decisions
Protect from automatic deletion
Old periodic checkpoints
Needed occasionally for analysis
Apply a documented retention policy
Generated samples and logs
Needed for comparison and audit
Store representative outputs with run identifiers
Local, cloud, and hybrid workflows
A local workstation offers direct access to fixed hardware and data. It is well suited to pipeline development, small models, short runs, and repeated debugging.
Cloud or cluster infrastructure provides access to hardware that may not be available locally and can support larger or temporary experiments. It also introduces remote storage, job scheduling, environment setup, data transfer, access control, and interruption handling.
A hybrid workflow develops and validates the smallest version locally, then runs a versioned configuration on larger remote infrastructure.
Local and cloud training
Swipe sideways to view the full comparison
Question
Local hardware
Cloud or managed cluster
Capacity
Limited to owned hardware
Can access larger temporary configurations
Setup
Direct but requires local drivers and maintenance
Requires remote environment and job configuration
Data access
Fast when data already resides locally
May require upload, replication, or remote storage
Cost behaviour
Hardware is paid for independently of each run
Resources and storage are commonly metered while allocated or used
Interruption
Affected by local crashes and power availability
Affected by job limits, preemption, quota, or remote failures
Scaling
Constrained by the physical machine
Can support multi-GPU and multi-node experiments
Reproducibility
Environment may drift through local updates
Images and job specifications can standardise repeated launches
A reliable hybrid workflow
1
Build locally
Test parsing, tokenisation, one forward pass, one backward pass, and checkpoint loading.
2
Run a tiny overfit test
Confirm that the model can fit a deliberately small sample when debugging the pipeline.
3
Run a short validation experiment
Use the real split and evaluation code for a limited number of steps.
4
Freeze the configuration
Record code commit, environment, data manifest, model settings, and seed.
5
Launch remotely
Use the same versioned configuration on larger hardware.
6
Track and checkpoint
Stream metrics and save recoverable states to persistent storage.
7
Reproduce locally
Load selected outputs or checkpoints into the evaluation environment.
Meta's AudioCraft training pipelines are built on PyTorch, use Flashy for training-pipeline design, and use Dora as an experiment manager.
The framework supports named configurations, experiment grids, team-specific environment settings, cluster configuration, and model-specific solvers. This separates experiment definitions from the physical machine or cluster used to run them.
For MusicGen, AudioCraft combines the compression model, language model, conditions, optimiser, metrics, validation, and generation stages within the training system.
Reproducibility has several levels
Exact bit-for-bit repetition is not always available across PyTorch releases, devices, platforms, kernels, and distributed schedules.
A project can still aim for:
Configuration reproducibility: the same settings can be reconstructed.
Data reproducibility: the same dataset and splits can be identified.
Procedural reproducibility: the same pipeline can be executed.
Statistical reproducibility: repeated runs produce comparable distributions of results.
Artifact reproducibility: the selected checkpoint and outputs can be recovered.
Record seeds, but do not treat one seed as a complete reproducibility system.
Minimum reproducibility record
1
Identify the code
Store the repository commit and note uncommitted changes.
2
Identify the environment
Record package, framework, driver, and platform versions.
3
Identify the data
Store manifests, split group IDs, preprocessing versions, and hashes.
4
Identify the model
Record architecture, parameter count, tokenizer, codec, and conditions.
5
Identify the optimisation
Record optimiser, scheduler, precision, batch, accumulation, and clipping.
6
Identify the infrastructure
Record GPU type, worker count, distributed strategy, and storage layout.
7
Identify randomness
Record seeds and relevant random-generator states.
8
Preserve artifacts
Retain selected checkpoints, logs, generated samples, and evaluation reports.
Plan infrastructure before the full run
1
Define the model and representation
Estimate parameters, token rate, codebooks, context, and trainable fraction.
2
Measure one real step
Run a representative batch and record peak memory, step time, and data wait.
3
Locate the bottleneck
Separate activation, gradient, optimiser, data, storage, and communication limits.
4
Apply the smallest suitable technique
Reduce the microbatch, change precision, checkpoint activations, freeze parameters, or shard state.
5
Estimate the complete experiment
Include training, validation, generation, repeated seeds, memorisation audits, and failed-run allowance.
6
Test recovery
Save, transfer, load, and resume a checkpoint before the long run.
7
Freeze tracking requirements
Ensure configuration, metrics, samples, and environment metadata are recorded automatically.
8
Set resource limits
Define maximum steps, wall time, storage, cloud spend, and checkpoint retention.
Interactive lesson
Estimate a Music Training Run
Try it: Begin with the small symbolic preset and inspect each memory category. Switch to codec-token music generation, increase duration and model size, then recover from an out-of-memory state using one technique at a time. Finish by selecting a checkpoint and tracking plan.
Training memory includes model parameters, gradients, optimiser states, activations, batches, temporary buffers, and allocator reserves. Duration affects sequence length through the representation. Mixed precision, smaller microbatches, accumulation, activation checkpointing, parameter-efficient tuning, and sharding solve different bottlenecks.
When the system still fails
A model that fits in memory can still produce NaN losses, silent data errors, stalled workers, broken checkpoints, invalid audio, or gradients that never reach part of the network.
The next section, Debugging Model Training, turns logs, gradients, memory traces, generated samples, and controlled tests into a systematic debugging process.
Check your understanding
Ready for a quick check?
Test whether you can identify memory components, choose infrastructure techniques, design checkpoints, and create a reproducible training workflow.
Ready to continue?
Save this section to your account and continue from any device.