mirror of
https://github.com/huggingface/lerobot.git
synced 2026-07-08 02:22:02 +00:00
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:
@@ -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":
|
||||
|
||||
Reference in New Issue
Block a user