phoenix-oss/llama-stack-mirror
0
0
Fork 1
mirror of https://github.com/meta-llama/llama-stack.git synced 2025-08-06 02:32:40 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

commit

Browse source
This commit is contained in:
Botao Chen 2025-01-14 18:01:49 -08:00
parent 89e3f81520
commit a77b8126d0
1 changed files with 10 additions and 0 deletions
Download patch file Download diff file Expand all files Collapse all files

10
llama_stack/providers/inline/post_training/torchtune/recipes/lora_finetuning_single_device.py
View file

@ -4,6 +4,7 @@
# This source code is licensed under the terms described in the LICENSE file in # This source code is licensed under the terms described in the LICENSE file in
# the root directory of this source tree. # the root directory of this source tree.
import gc
import logging import logging
import os import os
import time import time
@ -580,6 +581,15 @@ class LoraFinetuningSingleDevice:
checkpoint.training_metrics = training_metrics checkpoint.training_metrics = training_metrics
checkpoints.append(checkpoint) checkpoints.append(checkpoint)
# clean up the memory after training finishes
self._model.to("cpu")
del self._model
gc.collect()
torch.cuda.empty_cache()
print("Allocated:", torch.cuda.memory_allocated() / 1e6, "MB")
print("Reserved: ", torch.cuda.memory_reserved() / 1e6, "MB")
return (memory_stats, checkpoints) return (memory_stats, checkpoints)
async def validation(self) -> Tuple[float, float]: async def validation(self) -> Tuple[float, float]: