mirror of
https://github.com/MLSysBook/TinyTorch.git
synced 2026-03-12 03:33:33 -05:00
7d6e90c347
Created full integration test suite for KV caching module covering: Test Coverage: ✓ Linear projection integration (Q, K, V with cache) ✓ Multi-layer transformer caching (3 layers tested) ✓ Cache reset and reuse (multiple generations) ✓ Memory tracking accuracy (3 configs: tiny, small, medium) ✓ Batch inference support (parallel sequence generation) ✓ Boundary condition handling (empty, full, overflow) ✓ MultiHeadAttention compatibility Key Tests: 1. test_cache_with_linear_projections() - Verifies cache stores Linear layer Q/K/V outputs correctly - Tests autoregressive token-by-token processing - Validates cached values match original projections 2. test_cache_with_multi_layer_transformer() - Tests 3-layer transformer with cache - Verifies per-layer cache independence - Checks memory usage scales correctly 3. test_cache_reset_and_reuse() - Tests cache can handle multiple generation sequences - Verifies reset() clears state properly - Ensures new generations don't contain old data 4. test_cache_memory_tracking() - Validates memory calculation accuracy - Tests 3 model sizes (tiny, small, medium) - Ensures memory estimates are realistic 5. test_cache_with_batch_inference() - Tests 4 parallel sequences - Verifies batch dimension preserved - Ensures sequences remain independent 6. test_cache_boundary_conditions() - Empty cache retrieval - Fill to maximum capacity - Overflow protection - Invalid layer index handling 7. test_kv_cache_integration_with_attention() - Verifies compatibility with MultiHeadAttention - Tests standard attention still works - Documents integration pattern All tests follow TinyTorch testing patterns with clear output and assertions.
TinyTorch Test Suite
Comprehensive testing organized by purpose and scope.
Test Organization
📦 Module Tests (XX_modulename/)
Purpose: Test individual module functionality
Scope: Single module, isolated behavior
Example: 01_tensor/test_progressive_integration.py
These tests validate that each module works correctly in isolation.
🔗 Integration Tests (integration/)
Purpose: Test cross-module interactions
Scope: Multiple modules working together
Files:
test_gradient_flow.py- CRITICAL: Validates gradients flow through entire training stacktest_end_to_end_training.py- Full training loops (TODO)test_module_compatibility.py- Module interfaces (TODO)
Why this matters:
- Catches bugs that unit tests miss
- Validates the "seams" between modules
- Ensures training actually works end-to-end
🐛 Debugging Tests (debugging/)
Purpose: Catch common student pitfalls
Scope: Pedagogical - teaches debugging
Files:
test_gradient_vanishing.py- Detect/diagnose vanishing gradients (TODO)test_gradient_explosion.py- Detect/diagnose exploding gradients (TODO)test_common_mistakes.py- "Did you forget backward()?" style tests (TODO)
Philosophy: When these tests fail, the error message should teach the student what went wrong and how to fix it.
⚡ Autograd Edge Cases (05_autograd/)
Purpose: Stress-test autograd system
Scope: Autograd internals and edge cases
Files:
test_broadcasting.py- Broadcasting gradient bugs (TODO)test_computation_graph.py- Graph construction edge cases (TODO)test_backward_edge_cases.py- Numerical stability, etc. (TODO)
Running Tests
All tests
pytest tests/ -v
Integration tests only (recommended for debugging training issues)
pytest tests/integration/ -v
Specific test
pytest tests/integration/test_gradient_flow.py -v
Run without pytest
python tests/integration/test_gradient_flow.py
Test Philosophy
- Integration tests catch real bugs: The gradient flow test caught the exact bugs that prevented training
- Descriptive names: Test names should explain what they test
- Good error messages: When tests fail, students should understand why
- Pedagogical value: Tests teach correct usage patterns
Adding New Tests
When adding a test, ask:
- Is it testing one module? → Put in
XX_modulename/ - Is it testing modules working together? → Put in
integration/ - Is it teaching debugging? → Put in
debugging/ - Is it an autograd edge case? → Put in
05_autograd/
Most Important Tests
🔥 Must pass before merging:
integration/test_gradient_flow.py- If this fails, training is broken
📚 Module validation:
- Each module's inline tests (in
modules/source/) - Module-specific tests in
tests/XX_modulename/
Test Coverage Goals
- ✅ All tensor operations have gradient tests
- ✅ All layers compute gradients correctly
- ✅ All activations integrate with autograd
- ✅ All loss functions compute gradients
- ✅ All optimizers update parameters
- ⏳ End-to-end training converges (TODO)
- ⏳ Common pitfalls are detected (TODO)