ctde_v0.mem_probe¶
Compile-only GPU memory probe — report the XLA compiler's PLANNED peak memory for one
train_step WITHOUT running/allocating it. Safe to run alongside a busy GPU (compilation
grabs almost nothing). Use it to compare --remat off vs on for the same config.
PYTHONPATH=.:../../../FiedlerValueEstimation ~/ZymeraLab/.venv/bin/python -m ctde_v0.mem_probe --grid 32 --n-agents 10 --comm-r 5 --rollouts 16 --horizon 100 --explorer-tool frontier_attn --collision-mask on --mechanism lagrangian --conn-signal global_lambda2 --constraint-threshold 0.5 --role-picker expl_relay --w-coverage 3 --sense-walls --sense-free --boundary --cover-r 0 [--remat]
Prints temp_size (the peak scratch — this is what OOMs), plus argument/output sizes. Peak device memory ≈ temp + argument + output; the number to watch across --remat off/on is temp_size.