From 08c055cd72db49eac39b40285e84dc289eaf3728 Mon Sep 17 00:00:00 2001
From: Aleksa Gordic <gordicaleksa@gmail.com>
Date: Sun, 23 Jun 2024 10:46:54 +0000
Subject: [PATCH] Free up unified buffer mem

---
 llmc/zero.cuh | 1 +
 1 file changed, 1 insertion(+)

diff --git a/llmc/zero.cuh b/llmc/zero.cuh
index e64c7e6..7bcf08c 100644
--- a/llmc/zero.cuh
+++ b/llmc/zero.cuh
@@ -150,6 +150,7 @@ void multi_gpu_config_free(MultiGpuConfig* multi_gpu_config) {
     ncclCheck(ncclCommDestroy(multi_gpu_config->nccl_comm));
     cudaCheck(cudaStreamDestroy(multi_gpu_config->nccl_stream));
     cudaCheck(cudaEventDestroy(multi_gpu_config->compute_nccl_sync));
+    cudaCheck(cudaFree(multi_gpu_config->unified_buffer));
     mpiCheck(MPI_Finalize());
 #endif
 }