Ideas for reducing memory usage in bfs #183

Generated by ChatGPT 🤖

1. Don’t store states; store only hashes (and only what you need).
   In your code you already mostly do this (layer2 = layer2_hashes.reshape((-1,1)) when is_identity). Push this further:

• Ensure return_all_edges=False, return_all_hashes=False, and keep max_layer_size_to_store small.
• Prefer hasher.is_identity-style path whenever possible (work purely in hash space).

2. Use smaller hash / key types.
   If collisions are acceptable (often they are for “growth curve” experiments):

• Use uint32 / uint64 hashes instead of int64 (int64 is 8 bytes; uint32 is 4).
• If you currently keep int64 everywhere, switching seen/frontier to torch.int32 (or uint32 if you implement it) is an immediate ~2× win on those tensors.

3. Shard seen more aggressively (multi-GPU / multi-process / CPU RAM).
   You already shard by hash % ng. Increase num_gpus if available. If GPUs are the limiter, store shards on CPU RAM:

• Keep GPU resident only the current frontier (and maybe one recent layer).
• Keep seen_parts on CPU in sorted chunks; do GPU compute for neighbor gen + hash, then stream hashes to CPU for membership filtering.

4. Switch membership from “sorted list + searchsorted” to a bitset / bloom filter (approximate).
   This is the biggest practical memory lever if you can tolerate false positives (which only reduce exploration, not add wrong nodes).

• Bloom filter / cuckoo filter on CPU (or GPU) can represent “seen” at ~2–10 bits per element instead of 64 bits. That’s easily 10–30× smaller.
• Workflow: bloom-filter first, only for “maybe unseen” do the expensive exact check (or skip exact check entirely if you only need approximate growth).

5. Store seen as chunked on-disk / mmap arrays (if CPU RAM is the limiter).
   If GPU VRAM is the limiter but CPU RAM is large, keep seen in CPU RAM. If CPU RAM also limits:

• Store sorted hash chunks on disk via numpy.memmap and do batched merge / membership queries per chunk. This is slower but can extend feasible depth.

6. Reduce peak temporaries in neighbor generation.
   Your biggest spikes are usually:

• neighbors shape = (states_num * n_generators, state_dim) (huge)
• then hashing that
  You can cut peak memory by computing in a streaming way:
• For each generator, produce dst and immediately hash + dedup into a hash buffer, rather than materializing the full neighbors matrix.
• Even better: generate hashes directly without storing dst if your hasher can hash “gathered” data on the fly.

7. Avoid repeated torch.cat growth patterns.
   Repeated accepted = cat(old, new) is a peak-memory multiplier.

• Accumulate per-batch results into a list, and only cat once at the end of the layer (or do a k-way merge of sorted chunks).
• Same for seen_parts[g]: store as a list of sorted chunks, but periodically merge/compress them (pairwise merge) to reduce overhead and speed membership checks.

8. Use fewer bytes per state via encoding (you already have it).
   Your StringEncoder helps for state storage, but note you mostly store hashes, so it helps mainly in neighbor-gen temps. Still:

• Ensure bit_encoding_width is as small as possible.
• Ensure encoded states are stored in the compact format during BFS (avoid decoding).

9. Tune batch sizes for peak memory, not throughput.
   batch_size and hash_chunk_size directly control peak allocations.

• Smaller batches reduce peak VRAM at cost of more iterations and overhead.
• For exponential frontiers, preventing an OOM is usually worth more than peak throughput.

10. Drop exact dedup in early pipeline (approx), keep exact later.
    A cheap approximate dedup (bloom/bitset) before torch.unique / sort can reduce the size of tensors that hit expensive ops.