Debugging Model Training

A failed training run does not immediately reveal which part failed.

A flat loss might come from the optimiser, but it might also come from incorrect targets, an empty mask, frozen parameters, detached tensors, stale cached tokens, or a data loader that repeatedly returns the same example.

Repetitive audio might reflect overfitting, but it might instead come from low-diversity sampling settings. Silent output might come from the model, the codec decoder, amplitude scaling, file writing, or playback.

Debugging therefore means locating the earliest point where the observed values stop matching the expected values.

Create a golden batch

A golden batch is a small, saved set of examples whose expected intermediate values are understood.

For a text-conditioned music model, preserve:

1. Original audio
2. Caption and metadata
3. Resampled waveform
4. Crop boundaries
5. Codec or symbolic tokens
6. Padding mask
7. Tokenised condition
8. Condition embedding shape
9. Target alignment
10. Initial loss

The same batch can be run through new code versions to reveal where behaviour changed.

Inspect data before tuning the optimiser

The model learns only from the values that reach the loss function. A correct source file can become incorrect through cropping, resampling, channel conversion, caching, tokenisation, masking, or metadata pairing.

Before training, inspect examples after every important transformation. Listen to processed audio rather than only the original file. Decode model tokens back into audio or symbolic music where possible. Print captions beside the exact audio segment they condition.

Try to overfit one tiny batch

A sufficiently expressive model should usually be able to reduce loss sharply on a tiny fixed dataset when regularisation and augmentation are disabled.

This test asks whether the complete optimisation path is connected:

1. The targets contain learnable information
2. The loss depends on the predictions
3. Gradients reach trainable parameters
4. The optimiser updates those parameters
5. The model can represent the tiny task

Passing the test does not show that the model generalises. Failing it is strong evidence that the pipeline, objective, capacity, or optimisation settings need investigation.

Loss does not decrease

A flat loss can have several causes:

1. Learning rate is too small
2. Parameters are frozen or absent from the optimiser
3. Gradients are zero, missing, or detached
4. The optimiser step is skipped
5. Targets are random, constant, or misaligned
6. The valid loss mask is empty or nearly empty
7. The model output does not influence the selected loss
8. The same stale predictions are reused
9. Integer or device conversions break the intended path
10. The model lacks enough capacity for the task

Check the parameter update itself before changing the architecture.

Learning rate problems

A learning rate that is too low can make the loss appear frozen. A learning rate that is too high can cause oscillation, sudden spikes, parameter explosion, or non-finite values.

Google's Deep Learning Tuning Playbook recommends examining loss behaviour around and above the best observed learning rate and logging the full gradient norm when investigating instability. Warm-up, clipping, optimiser changes, and architectural corrections can help in some cases, but they should follow a diagnosis rather than being added blindly.

NaN and infinite values

Non-finite values can arise from invalid arithmetic or from values exceeding the representable numerical range.

Possible causes include:

1. Division by zero
2. Logarithm of zero or a negative value
3. Invalid normalisation statistics
4. Exponential overflow
5. Very large gradients or updates
6. Unstable lower-precision operations
7. Non-finite input audio or cached features
8. Empty reductions
9. Invalid masks
10. Corrupted model or optimiser state

Find the first non-finite value rather than only the final non-finite loss.

Mixed-precision failures

When FP16 mixed precision is used, gradient scaling enlarges small gradients before the backward pass. The scaler then unscales and checks the gradients before the optimiser step.

If non-finite gradients are detected, the optimiser step can be skipped and the scale adjusted. A run that repeatedly skips steps may appear to train slowly or not at all.

For diagnosis, compare the same batch in full precision and log the gradient scale, skipped steps, unscaled gradient norm, and first non-finite operation.

Repetitive generated music

Repetition can originate during training or generation.

Training-related causes include duplicated data, narrow adaptation data, excessive training, short contexts, weak structural supervision, and a model that mainly learns common local loops.

Generation-related causes include greedy decoding, low temperature, restrictive top-k or top-p settings, an overly dominant prefix, excessive guidance, or a stop condition that never activates.

Compare teacher-forced predictions, several sampling configurations, multiple checkpoints, and nearest training examples before attributing repetition to one source.

The model ignores captions or other controls

A condition can fail at several stages:

1. The metadata does not contain the intended condition.
2. The condition is paired with the wrong audio.
3. Tokenisation produces an empty or truncated input.
4. Condition dropout removes it more often than intended.
5. The condition encoder is frozen, detached, or incompatible.
6. The fusion module does not use the resulting tensor.
7. Gradients do not reach the conditioning path.
8. The training data does not contain a reliable relationship between condition and target.
9. Generation guidance or sampling settings are unsuitable.
10. The evaluation prompt asks for attributes the model was not trained to represent.

Trace the condition from source metadata to generated output.

Real-world example: tracing MusicGen's training step

Meta's official AudioCraft solver exposes several useful debugging boundaries.

The solver:

1. Confirms that the audio batch size matches the metadata count
2. Converts metadata into conditioning attributes
3. Applies classifier-free and attribute dropout
4. Tokenises and encodes the conditions
5. Encodes audio into codec tokens
6. Constructs a padding mask from valid audio lengths
7. Produces logits for several codebooks
8. Checks that logits, targets, and masks have compatible shapes
9. Calculates cross-entropy only on valid positions
10. Runs backward propagation
11. Unscales mixed-precision gradients
12. Optionally records and clips gradient norm
13. Runs the optimiser and scheduler
14. Clears gradients
15. Raises an error when the loss is non-finite
16. Logs total and per-codebook cross-entropy and perplexity

This structure allows a failure to be located before the complete generation stage.

Silent or nearly silent output

Silence can enter the system through several routes:

1. Source audio is silent or incorrectly normalised
2. Channel conversion cancels the waveform
3. The crop selects an empty or padded region
4. Codec tokens represent silence
5. Generated tokens are invalid or dominated by a silence token pattern
6. The decoder receives the wrong shape or codebook order
7. Amplitude values are scaled incorrectly before file writing
8. The output file uses the wrong sample rate or subtype
9. Playback software reads the file incorrectly

Test each stage by saving or listening to its output.

Suspiciously good validation results

A broken split can imitate successful training.

Investigate when validation loss is unexpectedly low, generated validation pieces resemble training examples, or a small model performs implausibly well.

Check:

1. Source IDs across splits
2. Duplicate and near-duplicate clusters
3. Overlapping windows
4. Alternate arrangements or performances
5. Cached features created before splitting
6. Normalisation fitted on all data
7. Validation examples accidentally included in the sampler
8. Test results repeatedly used for model selection

Correct the split and rerun affected experiments.

Compare checkpoints and code versions

When a failure appears halfway through training, compare the last healthy checkpoint with the first unhealthy checkpoint using the same saved batch.

When a software change introduces a failure, run the golden batch across earlier and later commits. This form of bisection narrows the change interval before detailed tensor inspection begins.

Preserve optimiser state when the failure may depend on accumulated moments or learning-rate schedule.

Slow training and stalled pipelines

A run can appear frozen while it is waiting for data, synchronising devices, compiling operations, saving checkpoints, or performing an unexpectedly expensive calculation.

PyTorch Profiler can record operator time and memory. Data-loader logs can reveal slow workers or repeated restarts. Distributed debugging tools can help locate workers that reached different collective operations or stopped progressing.

Measure before assuming that the GPU or model is the bottleneck.

Confirm that the fix solved the right problem

A useful fix should pass three tests:

1. The original failure no longer occurs.
2. An unaffected control case still works.
3. The improvement persists across more than one batch, seed, or checkpoint when the failure was not inherently deterministic.

Record the before-and-after evidence. A run that merely restarts successfully has not necessarily resolved the underlying cause.

From a working training run to a trustworthy model

A model can train without crashing and still be unsuitable for release. It may sound weak, ignore controls, reproduce training material, perform differently across groups, or fail under real product conditions.

Chapter 5, Evaluating, Releasing, and Governing the Model, moves beyond training correctness to musical quality, benchmarking, safety, provenance, disclosure, monitoring, and release decisions.