Release accelerator model refs during cleanup

Author: xk-huang

lmms_eval/api/model.py

@@ -164,6 +164,10 @@ def set_cache_hook(self, cache_hook) -> None:
         self.cache_hook = cache_hook
 
     def clean(self):
+        accelerator = getattr(self, "accelerator", None)
+        if accelerator is not None and hasattr(accelerator, "free_memory"):
+            accelerator.free_memory()
+
         for attr_name in list(vars(self)):
             attr_value = getattr(self, attr_name)
             if isinstance(attr_value, nn.Module):

test/models/test_model_cleanup.py

@@ -0,0 +1,38 @@
+import torch
+
+from lmms_eval.api.model import lmms
+
+
+class _DummyLM(lmms):
+    def loglikelihood(self, requests):
+        return []
+
+    def generate_until(self, requests):
+        return []
+
+    def generate_until_multi_round(self, requests):
+        return []
+
+
+class _FakeAccelerator:
+    def __init__(self, model):
+        self._models = [model]
+        self.free_memory_calls = 0
+
+    def free_memory(self):
+        self.free_memory_calls += 1
+        self._models = []
+
+
+def test_clean_releases_accelerator_model_references():
+    lm = _DummyLM()
+    model = torch.nn.Linear(1, 1)
+    accelerator = _FakeAccelerator(model)
+    lm._model = model
+    lm.accelerator = accelerator
+
+    lm.clean()
+
+    assert accelerator.free_memory_calls == 1
+    assert accelerator._models == []
+    assert not hasattr(lm, "_model")