lausannel
commited on
Commit
•
e50cd48
1
Parent(s):
d9ab9eb
fix: change position of remove offload path and fix params in __init__()
Browse files
src/backend/moe_infinity.py
CHANGED
@@ -31,15 +31,17 @@ class MoEHFLM(HFLMWithMeasurement):
|
|
31 |
self.use_chat_template = use_chat_template
|
32 |
if "device" in kwargs:
|
33 |
kwargs.pop("device")
|
|
|
|
|
34 |
super().__init__(
|
35 |
-
*args, **kwargs, pretrained=pretrained,
|
36 |
) # Assuming HFLM accepts a 'pretrained' arg and handles it
|
37 |
# self._create_model()
|
38 |
-
shutil.rmtree(os.path.join(self.offload_path, "moe-infinity-offloads"))
|
39 |
|
40 |
def __del__(self):
|
41 |
# Clean up offloaded models from self.offload_path
|
42 |
-
|
|
|
43 |
|
44 |
def _create_model(self, *args, **kwargs):
|
45 |
"""
|
|
|
31 |
self.use_chat_template = use_chat_template
|
32 |
if "device" in kwargs:
|
33 |
kwargs.pop("device")
|
34 |
+
if os.path.exists(os.path.join(self.offload_path, "moe-infinity-offloads")):
|
35 |
+
shutil.rmtree(os.path.join(self.offload_path, "moe-infinity-offloads"))
|
36 |
super().__init__(
|
37 |
+
*args, **kwargs, pretrained=pretrained, device="cuda:0"
|
38 |
) # Assuming HFLM accepts a 'pretrained' arg and handles it
|
39 |
# self._create_model()
|
|
|
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 |
|
46 |
def _create_model(self, *args, **kwargs):
|
47 |
"""
|