Spaces:
Running
Running
lausannel
commited on
Commit
·
d3ab5ad
1
Parent(s):
e50cd48
fix:add resource cleanup in `__del__`
Browse files
src/backend/moe_infinity.py
CHANGED
|
@@ -17,7 +17,8 @@ class MoEHFLM(HFLMWithMeasurement):
|
|
| 17 |
self,
|
| 18 |
pretrained: str = "mistralai/Mixtral-8x7B-Instruct-v0.1",
|
| 19 |
moe_config: dict = None,
|
| 20 |
-
offload_path=os.path.expanduser("~"),
|
|
|
|
| 21 |
device_memory_ratio=0.75,
|
| 22 |
use_chat_template=True,
|
| 23 |
*args,
|
|
@@ -40,6 +41,8 @@ class MoEHFLM(HFLMWithMeasurement):
|
|
| 40 |
|
| 41 |
def __del__(self):
|
| 42 |
# Clean up offloaded models from self.offload_path
|
|
|
|
|
|
|
| 43 |
if os.path.exists(os.path.join(self.offload_path, "moe-infinity-offloads")):
|
| 44 |
shutil.rmtree(os.path.join(self.offload_path, "moe-infinity-offloads"))
|
| 45 |
|
|
|
|
| 17 |
self,
|
| 18 |
pretrained: str = "mistralai/Mixtral-8x7B-Instruct-v0.1",
|
| 19 |
moe_config: dict = None,
|
| 20 |
+
# offload_path=os.path.expanduser("~"),
|
| 21 |
+
offload_path='/mnt/raid0nvme1/xly/test-data/',
|
| 22 |
device_memory_ratio=0.75,
|
| 23 |
use_chat_template=True,
|
| 24 |
*args,
|
|
|
|
| 41 |
|
| 42 |
def __del__(self):
|
| 43 |
# Clean up offloaded models from self.offload_path
|
| 44 |
+
self._model.engine.clean_up()
|
| 45 |
+
self._model.engine.archer_engine.clean_up_resources()
|
| 46 |
if os.path.exists(os.path.join(self.offload_path, "moe-infinity-offloads")):
|
| 47 |
shutil.rmtree(os.path.join(self.offload_path, "moe-infinity-offloads"))
|
| 48 |
|