Overfit one sample

Start by overfitting one sample. The loss should decrease meaningfully and then stabilize. With one sample, the gradient estimate has no sampling noise from other examples. For that tiny training distribution, the batch gradient is the population gradient.

This is the simplest consistent signal the model can receive. At higher batch sizes, the optimizer averages gradients from multiple samples that may pull in different directions. If the model cannot fit one sample, adding more samples usually makes the debugging problem harder, not easier.

If the loss cannot decrease even on one sample, the bug is probably not data diversity. Check the training loop first: gradients may not be flowing, parameters may not be passed to the optimizer, the model may be in the wrong mode, the loss may be computed from the wrong tensors, or the learning rate may be so small or so large that optimization makes no useful progress.

Check parallel loss behavior

Make sure the loss behaves the same in parallel and non-parallel setups. Distributed training, gradient accumulation, mixed precision, or multi-device reduction can silently change the effective loss scale. Before debugging model quality, verify that the same batch gives the same loss and comparable gradients in the simple single-device setup and in the parallel setup.