Memory-Constrained Trainer
Follow the data, decisions, feedback, and validation boundaries before writing the full system.
How to read this diagram
Read left to right for the forward path: raw information becomes a representation, passes through the project’s main computational ideas, and produces an output that can be measured. Then follow the feedback path back toward the trainable or decision-making components.
Checkpointing
Activation checkpointing discards selected forward intermediates and recomputes them during backward. It trades extra arithmetic for lower peak memory and requires deterministic replay of any stochastic operation inside a checkpointed region.
Boundary check: document its accepted input, output shape, mutable state, failure modes, and the metric that proves this stage is correct before connecting it downstream.
Mixed precision
Mixed precision stores or computes many tensors in lower precision while preserving sensitive reductions and master weights at higher precision. Loss scaling protects small gradients from underflow; overflow detection decides whether an update is safe.
Boundary check: document its accepted input, output shape, mutable state, failure modes, and the metric that proves this stage is correct before connecting it downstream.
ZeRO
ZeRO removes replicated training state by partitioning optimizer tensors, gradients, and eventually parameters across workers. Each stage saves more memory while introducing communication and lifetime-management constraints.
Boundary check: document its accepted input, output shape, mutable state, failure modes, and the metric that proves this stage is correct before connecting it downstream.
Architecture review checklist
- Every arrow has a documented shape, dtype, unit, or schema.
- Training and evaluation paths cannot leak information into each other.
- Randomness is seeded and captured in experiment metadata.
- Expensive stages expose timing, memory, throughput, and error metrics.
- Each feedback loop has a stop condition and a rollback strategy.
- Small reference implementations exist for numerical comparisons.