lausannel commited on
Commit
d3ab5ad
1 Parent(s): e50cd48

fix:add resource cleanup in `__del__`

Browse files
Files changed (1) hide show
  1. src/backend/moe_infinity.py +4 -1
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