fix(profiling): use SGD for pi0/pi05/pi0_fast and free CUDA cache after deterministic forward

Adam optimizer states (exp_avg + exp_avg_sq) require ~16GB extra on top of
model params and gradients for 4B parameter models, exceeding the 22GB GPU.
SGD has zero optimizer state overhead and profiling only measures
forward/backward timing anyway.

Also adds torch.cuda.empty_cache() after deterministic forward to release
transient memory before the training loop starts.

Made-with: Cursor
This commit is contained in:
Pepijn
2026-04-16 16:09:56 +02:00
parent b1e16783de
commit da7da741f1
2 changed files with 32 additions and 0 deletions
+2
View File
@@ -343,6 +343,8 @@ class TrainingProfiler:
output_dir=self._output_dir,
device_type=self._device.type,
)
if self._device.type == "cuda":
torch.cuda.empty_cache()
def __enter__(self) -> TrainingProfiler:
if self._device.type == "cuda":