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 |
|