This repository implements a set of diagnostic assertions based on Andrei Karpathy's blog post "A Recipe for Training Neural Networks". Instead of a manual checklist, we provide programmatic checks to diagnose and debug neural networks efficiently.
Training neural networks can be tricky. This toolkit saves you time by automating common diagnostic steps, allowing you to focus on the interesting parts of your model development.
Batteries included:
from neural_net_checklist import torch_diagnostics
torch_diagnostics.assert_all_for_causal_llm_cross_entropy_loss(
lambda: CausalTransformer(vocab_size),
dataloader,
embedding_layer_name="embedding",
vocab_size=vocab_size,
device="cpu",
)Output:
🔍 Checking all conditions for a causal LLM model with cross-entropy loss...
🔍 Checking loss at initialization...
Loss at initialization (4.2921) is within 75.0% of expected loss (4.0775)
✅ Loss at initialization is within the expected range.
🔍 Checking calibration at initialization...
Model is well-calibrated at initialization. Mean deviation: 0.00786806270480156
✅ Model is well-calibrated at initialization.
🔍 Checking forward batch independence...
🔧 Replaced 10 norm layers with Identity: {'LayerNorm': {'transformer.layers.3.norm2', 'transformer.layers.2.norm1', 'transformer.layers.0.norm2', 'transformer.layers.4.norm2', 'transformer.layers.0.norm1', 'transformer.layers.2.norm2', 'transformer.layers.1.norm2', 'transformer.layers.4.norm1', 'transformer.layers.1.norm1', 'transformer.layers.3.norm1'}}
✅ Forward batch independence verified.
🔍 Checking forward causal property...
🔧 Replaced 10 norm layers with Identity: {'LayerNorm': {'transformer.layers.3.norm2', 'transformer.layers.2.norm1', 'transformer.layers.0.norm2', 'transformer.layers.4.norm2', 'transformer.layers.0.norm1', 'transformer.layers.2.norm2', 'transformer.layers.1.norm2', 'transformer.layers.4.norm1', 'transformer.layers.1.norm1', 'transformer.layers.3.norm1'}}
✅ Causal independence verified.
🔍 Checking non-zero gradients...
✅ All gradients are non-zero.
🔍 Checking backward batch independence...
🔧 Replaced 10 norm layers with Identity: {'LayerNorm': {'transformer.layers.3.norm2', 'transformer.layers.2.norm1', 'transformer.layers.0.norm2', 'transformer.layers.4.norm2', 'transformer.layers.0.norm1', 'transformer.layers.2.norm2', 'transformer.layers.1.norm2', 'transformer.layers.4.norm1', 'transformer.layers.1.norm1', 'transformer.layers.3.norm1'}}
✅ Backward pass is batch independent.
🔍 Checking backward causal property...
🔧 Replaced 10 norm layers with Identity: {'LayerNorm': {'transformer.layers.3.norm2', 'transformer.layers.2.norm1', 'transformer.layers.0.norm2', 'transformer.layers.4.norm2', 'transformer.layers.0.norm1', 'transformer.layers.2.norm2', 'transformer.layers.1.norm2', 'transformer.layers.4.norm1', 'transformer.layers.1.norm1', 'transformer.layers.3.norm1'}}
✅ Backward pass exhibits causal independence.
🔍 Checking input independence baseline worse...
Loss: 3.125023: 7%|▋ | 9/128 [00:01<00:24, 4.89it/s]
Loss: 2.736129: 7%|▋ | 9/128 [00:02<00:26, 4.48it/s]
Regular loss: 2.736128568649292
Input independent loss: 3.1189355850219727
✅ Input independent baseline is worse.
🔍 Overfitting to the batch...
Loss: 0.099780: 0%| | 87/65536 [00:07<1:39:13, 10.99it/s]Loss (0.09978045523166656) below threshold (0.1) at step 87.
✅ Model can overfit to batch.
✅ All conditions for a causal LLM model with cross-entropy loss verified.
See examples/causal_transformer.py for a complete example. (Examples for ResNet on CIFAR10 and LeNet on MNIST are also included.)
We've implemented the following checks:
-
Verify loss @ init (for balanced classification tasks)
assert_balanced_classification_cross_entropy_loss_at_init -
Init well (for balanced classification tasks)
assert_balanced_classification_cross_entropy_loss_at_init -
Non-zero gradients
assert_non_zero_gradients -
Batch independence
- Forward pass (memory-efficient, but note: batchnorm breaks this naturally)
assert_batch_independence_forward - Backward pass (uses more memory, checks gradients)
assert_batch_independence_backward
- Forward pass (memory-efficient, but note: batchnorm breaks this naturally)
-
Overfit one batch
assert_overfit_one_batch -
Visualize just before the net
patch_module_raise_inputs
To run all assertions:
-
For classification tasks (e.g., computer vision):
assert_all_for_classification_cross_entropy_loss -
For causal language models:
assert_all_for_llm_cross_entropy_loss
pip install neural_net_checklistimport neural_net_checklist.torch_diagnostics as torch_diagnostics
# Assume you have a model and a DataLoader
model = YourModel()
train_loader = YourDataLoader()
# Run all checks
torch_diagnostics.assert_all_for_classification_cross_entropy_loss(model, train_loader)See the code and docstrings for more details.
Contributions are welcome! If you have ideas for additional checks or improvements, please open an issue or submit a pull request.
This project is licensed under the MIT License - see the LICENSE file for details.
Happy debugging! Remember, neural nets can be finicky, but with the right tools, we can tame them. 🧠🔧