fix: size-adaptive memory pool to prevent OOM on 224x224 training

MemoryPool was designed for small CIFAR-10 buffers but used the same
aggressive caching (refill=32, max=64 per class) for 224x224 tensors,
causing multi-GB allocations. Three fixes:

- Size-adaptive limits: buffers >=1MB get refill=2, max=4 (was 32/64)
- trim() now flushes thread-local caches (previously only global pool)
- Training loop scopes batch tensors so trim() can actually free them

Also adds Tensor::live_count()/live_bytes() diagnostics and
MemoryPool::cached_bytes() for tracking memory across batches.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

Author: henry

core/memory_pool.cpp

@@ -28,8 +28,20 @@ void MemoryPool::set_allocator(AllocFn alloc, FreeFn free) {
 
 namespace {
 
-const size_t kMaxPerClassPerThread = 64;
-const size_t kRefillBatch = 32;
+// Size-adaptive limits: large buffers (e.g. 128MB for [32,64,112,112])
+// should NOT fill thread-local caches with dozens of copies.
+// Threshold: 256K floats = 1MB.  Above that, cache fewer buffers.
+const size_t kLargeBufferThreshold = 256 * 1024;  // floats
+
+static size_t max_per_class(size_t cls) {
+    if (cls >= kLargeBufferThreshold) return 4;
+    return 64;
+}
+
+static size_t refill_batch(size_t cls) {
+    if (cls >= kLargeBufferThreshold) return 2;
+    return 32;
+}
 
 size_t next_power_of_two(size_t n) {
     if (n == 0) return 1;
@@ -138,7 +150,7 @@ struct ThreadCache {
             return ptr;
         }
         std::vector<float*> batch;
-        global->acquire_batch(cls, kRefillBatch, batch);
+        global->acquire_batch(cls, refill_batch(cls), batch);
         if (batch.empty()) return nullptr;
         for (size_t i = 1; i < batch.size(); ++i)
             local.push_back(batch[i]);
@@ -150,11 +162,12 @@ struct ThreadCache {
         global_ = global;
         size_t cls = MemoryPool::size_class(original_n);
         auto& local = buckets[cls];
-        if (local.size() < kMaxPerClassPerThread) {
+        size_t max_local = max_per_class(cls);
+        if (local.size() < max_local) {
             local.push_back(ptr);
             return;
         }
-        size_t drain_count = kMaxPerClassPerThread / 2;
+        size_t drain_count = max_local / 2;
         std::vector<float*> to_global;
         for (size_t i = 0; i < drain_count && !local.empty(); ++i) {
             to_global.push_back(local.back());
@@ -163,6 +176,17 @@ struct ThreadCache {
         global->release_batch(to_global, cls);
         local.push_back(ptr);
     }
+
+    // Free all cached buffers in this thread's cache directly (bypass global).
+    void flush() {
+        for (auto& kv : buckets) {
+            for (float* ptr : kv.second) {
+                if (g_free_fn) g_free_fn(ptr);
+                else std::free(ptr);
+            }
+            kv.second.clear();
+        }
+    }
 };
 
 thread_local ThreadCache t_thread_cache;
@@ -187,15 +211,34 @@ void MemoryPool::release(float* ptr, size_t original_n) {
 }
 
 void MemoryPool::trim() {
+    // Flush calling thread's local cache first (frees directly, no lock needed).
+    memory_pool_detail::t_thread_cache.flush();
+
+    // Then free global pool buckets.
     auto* p = impl();
     std::lock_guard<std::mutex> lock(p->mutex);
-    for (auto& [cls, free_list] : p->buckets) {
-        for (float* ptr : free_list) {
+    for (auto& kv : p->buckets) {
+        for (float* ptr : kv.second) {
             if (g_free_fn) g_free_fn(ptr);
             else std::free(ptr);
         }
-        free_list.clear();
+        kv.second.clear();
+    }
+}
+
+size_t MemoryPool::cached_bytes() const {
+    size_t total = 0;
+    // Thread-local cache
+    for (auto& kv : memory_pool_detail::t_thread_cache.buckets) {
+        total += kv.first * sizeof(float) * kv.second.size();
+    }
+    // Global pool
+    auto* p = const_cast<MemoryPool*>(this)->impl();
+    std::lock_guard<std::mutex> lock(p->mutex);
+    for (auto& kv : p->buckets) {
+        total += kv.first * sizeof(float) * kv.second.size();
     }
+    return total;
 }
 
 std::shared_ptr<float> MemoryPool::acquire_shared(size_t n) {

core/memory_pool.h

@@ -28,10 +28,14 @@ class MemoryPool {
     // Round requested size up to the pool's size class (for tests/tooling).
     static size_t size_class(size_t n);
 
-    // Free all cached buffers in the global pool (does not affect in-use buffers).
+    // Free all cached buffers in the global pool AND calling thread's local cache.
     // Call periodically to prevent unbounded memory growth.
     void trim();
 
+    // Return total bytes currently cached (global + calling thread's local).
+    // For diagnostics only — approximate under concurrency.
+    size_t cached_bytes() const;
+
     // Custom allocator hooks (for CUDA managed memory).
     // When set, these replace std::malloc/std::free for all pool allocations.
     using AllocFn = float*(*)(size_t n_bytes);

core/tensor.cpp

@@ -20,12 +20,26 @@
 
 static thread_local std::mt19937 rng(42);
 
-Tensor::Tensor() : requires_grad(false), data_size_(0), grad_size_(0) {}
+// Global live tensor counter (atomic for thread safety).
+static std::atomic<int64_t> g_live_tensor_count{0};
+static std::atomic<int64_t> g_live_tensor_bytes{0};
 
-Tensor::~Tensor() {}
+int64_t Tensor::live_count() { return g_live_tensor_count.load(std::memory_order_relaxed); }
+int64_t Tensor::live_bytes() { return g_live_tensor_bytes.load(std::memory_order_relaxed); }
+
+Tensor::Tensor() : requires_grad(false), data_size_(0), grad_size_(0) {
+    g_live_tensor_count.fetch_add(1, std::memory_order_relaxed);
+}
+
+Tensor::~Tensor() {
+    g_live_tensor_count.fetch_sub(1, std::memory_order_relaxed);
+    g_live_tensor_bytes.fetch_sub(static_cast<int64_t>(data_size_ + grad_size_) * sizeof(float),
+                                   std::memory_order_relaxed);
+}
 
 Tensor::Tensor(const std::vector<size_t>& shape, bool requires_grad)
     : shape(shape), requires_grad(requires_grad) {
+    g_live_tensor_count.fetch_add(1, std::memory_order_relaxed);
     size_t total = 1;
     for (auto s : shape) total *= s;
     data_size_ = total;
@@ -40,10 +54,13 @@ Tensor::Tensor(const std::vector<size_t>& shape, bool requires_grad)
     } else {
         grad_size_ = 0;
     }
+    g_live_tensor_bytes.fetch_add(static_cast<int64_t>(data_size_ + grad_size_) * sizeof(float),
+                                   std::memory_order_relaxed);
 }
 
 Tensor::Tensor(const std::vector<float>& data, const std::vector<size_t>& shape, bool requires_grad)
     : shape(shape), requires_grad(requires_grad) {
+    g_live_tensor_count.fetch_add(1, std::memory_order_relaxed);
     data_size_ = data.size();
     data_storage_ = MemoryPool::instance().acquire_shared(data_size_);
     if (!data_storage_) throw std::bad_alloc();
@@ -56,12 +73,15 @@ Tensor::Tensor(const std::vector<float>& data, const std::vector<size_t>& shape,
     } else {
         grad_size_ = 0;
     }
+    g_live_tensor_bytes.fetch_add(static_cast<int64_t>(data_size_ + grad_size_) * sizeof(float),
+                                   std::memory_order_relaxed);
 }
 
 Tensor::Tensor(std::shared_ptr<float> data_storage, size_t data_size,
                const std::vector<size_t>& shape, bool requires_grad)
     : shape(shape), requires_grad(requires_grad),
       data_storage_(std::move(data_storage)), data_size_(data_size) {
+    g_live_tensor_count.fetch_add(1, std::memory_order_relaxed);
     if (requires_grad) {
         grad_size_ = data_size_;
         grad_storage_ = MemoryPool::instance().acquire_shared(grad_size_);
@@ -70,6 +90,8 @@ Tensor::Tensor(std::shared_ptr<float> data_storage, size_t data_size,
     } else {
         grad_size_ = 0;
     }
+    g_live_tensor_bytes.fetch_add(static_cast<int64_t>(data_size_ + grad_size_) * sizeof(float),
+                                   std::memory_order_relaxed);
 }
 
 TensorPtr Tensor::create(const std::vector<size_t>& shape, bool requires_grad) {

core/tensor.h

@@ -11,6 +11,7 @@
 #include <cstdint>
 #include <cstdio>
 #include <random>
+#include <atomic>
 
 enum class DType { Float32, Float16 };
 
@@ -172,6 +173,10 @@ class Tensor : public std::enable_shared_from_this<Tensor> {
 
     void print(const char* name = nullptr) const;
 
+    // Diagnostics: live tensor count/memory (atomic, thread-safe).
+    static int64_t live_count();
+    static int64_t live_bytes();
+
 private:
     void build_topo(std::vector<Tensor*>& topo, std::vector<Tensor*>& visited);
     bool should_track_grad() const;

examples/resnet18_imagenette.cpp

@@ -647,63 +647,74 @@ int main(int argc, char* argv[]) {
             size_t batch_end = std::min(batch_start + batch_size, n_train);
             size_t bs = batch_end - batch_start;
 
-            // Assemble batch from mmap using shuffled indices
-            auto images = Tensor::create({bs, 3, 224, 224}, false);
-            auto labels = Tensor::create({bs}, false);
-
-            for (size_t i = 0; i < bs; i++) {
-                size_t idx = train_indices[batch_start + i];
-                std::memcpy(images->data() + i * img_elems,
-                             train_images_mmap.data + idx * img_elems,
-                             img_elems * sizeof(float));
-                labels->data()[i] = train_labels_mmap.data[idx];
-            }
+            float batch_loss = 0.0f;
+
+            // Scoped block: all batch tensors freed before trim()
+            {
+                // Assemble batch from mmap using shuffled indices
+                auto images = Tensor::create({bs, 3, 224, 224}, false);
+                auto labels = Tensor::create({bs}, false);
+
+                for (size_t i = 0; i < bs; i++) {
+                    size_t idx = train_indices[batch_start + i];
+                    std::memcpy(images->data() + i * img_elems,
+                                 train_images_mmap.data + idx * img_elems,
+                                 img_elems * sizeof(float));
+                    labels->data()[i] = train_labels_mmap.data[idx];
+                }
 
-            // Data augmentation: pad 16 -> random crop 224x224 -> random horizontal flip
-            auto augmented = images->pad2d(16)->random_crop(224, 224)->random_flip_horizontal(0.5f);
+                // Data augmentation: pad 16 -> random crop 224x224 -> random horizontal flip
+                auto augmented = images->pad2d(16)->random_crop(224, 224)->random_flip_horizontal(0.5f);
 
-            optimizer.zero_grad();
+                optimizer.zero_grad();
 
-            auto output = model.forward(augmented);
-            auto loss   = criterion(output, labels);
+                auto output = model.forward(augmented);
+                auto loss   = criterion(output, labels);
 
-            loss->backward();
+                loss->backward();
 
-            // Manual L2 weight decay on conv/linear weights
-            apply_weight_decay(all_params, weight_decay);
+                // Manual L2 weight decay on conv/linear weights
+                apply_weight_decay(all_params, weight_decay);
 
-            optimizer.step();
+                optimizer.step();
 
-            // Free cached memory to prevent unbounded growth with large 224x224 tensors
-            MemoryPool::instance().trim();
-
-            // Read loss scalar (handles GPU -> CPU transfer)
-            total_loss += read_loss_scalar(loss);
-            num_batches++;
+                // Extract loss and accuracy BEFORE releasing tensors
+                batch_loss = read_loss_scalar(loss);
 
-            // Compute training accuracy for this batch
-            auto cpu_output = ensure_cpu(output);
-            auto cpu_labels = ensure_cpu(labels);
-            size_t nc = cpu_output->shape[1];
-            for (size_t i = 0; i < bs; i++) {
-                size_t pred = 0;
-                float mv = cpu_output->data()[i * nc];
-                for (size_t j = 1; j < nc; j++) {
-                    if (cpu_output->data()[i * nc + j] > mv) {
-                        mv = cpu_output->data()[i * nc + j];
-                        pred = j;
+                auto cpu_output = ensure_cpu(output);
+                auto cpu_labels = ensure_cpu(labels);
+                size_t nc = cpu_output->shape[1];
+                for (size_t i = 0; i < bs; i++) {
+                    size_t pred = 0;
+                    float mv = cpu_output->data()[i * nc];
+                    for (size_t j = 1; j < nc; j++) {
+                        if (cpu_output->data()[i * nc + j] > mv) {
+                            mv = cpu_output->data()[i * nc + j];
+                            pred = j;
+                        }
                     }
+                    if (pred == static_cast<size_t>(cpu_labels->data()[i])) correct++;
+                    total++;
                 }
-                if (pred == static_cast<size_t>(cpu_labels->data()[i])) correct++;
-                total++;
             }
+            // All batch tensors (images, labels, augmented, output, loss) now freed.
+
+            // Free cached memory — thread-local + global pools.
+            MemoryPool::instance().trim();
+
+            total_loss += batch_loss;
+            num_batches++;
 
-            // Print progress
+            // Print progress + memory diagnostics
             if (epoch == start_epoch || num_batches % 20 == 0) {
                 auto batch_time = std::chrono::high_resolution_clock::now();
                 double elapsed = std::chrono::duration<double>(batch_time - epoch_start).count();
-                printf("\r  Epoch %3d: batch %3zu/%zu | %.1fs elapsed",
-                       epoch + 1, num_batches, total_batches, elapsed);
+                int64_t live = Tensor::live_count();
+                int64_t live_mb = Tensor::live_bytes() / (1024 * 1024);
+                size_t pool_mb = MemoryPool::instance().cached_bytes() / (1024 * 1024);
+                printf("\r  Epoch %3d: batch %3zu/%zu | %.1fs | tensors: %lld (%lldMB) pool: %zuMB",
+                       epoch + 1, num_batches, total_batches, elapsed,
+                       (long long)live, (long long)live_mb, pool_mb);
                 fflush(stdout);
             }
         }