Update src/FineTune/model.py

src/FineTune/model.py

@@ -11,7 +11,7 @@ def calculate_MMD_loss(human_crit, sample_crit):
     mmd_loss = human_crit.mean() - sample_crit.mean()
     return mmd_loss
 
-def from_pretrained(cls, model_name, kwargs, cache_dir):
+def from_pretrained(cls, model_name, kwargs, cache_dir, device=None):
     # use local model if it exists
     if "/" in model_name:
         local_path = os.path.join(cache_dir, model_name.split("/")[1])
@@ -20,7 +20,18 @@ def from_pretrained(cls, model_name, kwargs, cache_dir):
 
     if os.path.exists(local_path):
         return cls.from_pretrained(local_path, **kwargs)
-    return cls.from_pretrained(model_name, **kwargs, cache_dir=cache_dir, device_map='auto')
+
+    remote_kwargs = dict(kwargs, cache_dir=cache_dir)
+    if device is not None:
+        # Pin the whole model to a single device instead of device_map='auto'.
+        # 'auto' lets accelerate split the model across CPU/disk when memory
+        # is tight at load time (common on HF Spaces' free CPU tier), which
+        # then makes the later `model.to(device)` call in load_model() raise:
+        #   "You can't move a model that has some modules offloaded to
+        #    cpu or disk."
+        # Forcing everything onto one device up front avoids that split.
+        remote_kwargs['device_map'] = {'': device}
+    return cls.from_pretrained(model_name, **remote_kwargs)
 
 model_fullnames = {  
     'gemma-1b': 'google/gemma-3-1b-pt',
@@ -76,7 +87,7 @@ class ComputeStat(nn.Module):
                 model_kwargs.update(dict(torch_dtype=torch.float16))
             if torch.__version__ >= '2.0.0' and 'gemma' in model_name:
                 model_kwargs.update({'attn_implementation': 'sdpa'})
-            model = from_pretrained(AutoModelForCausalLM, model_fullname, model_kwargs, cache_dir)
+            model = from_pretrained(AutoModelForCausalLM, model_fullname, model_kwargs, cache_dir, device=device)
             print(f'Moving model to {device}...', end='', flush=True)
             start = time.time()
             model.to(device)
@@ -179,11 +190,14 @@ class ComputeStat(nn.Module):
         model = cls(*args, **kwargs)
 
         # 2. 加载 scoring_model
+        # Pin to a single device instead of device_map="auto" -- see the
+        # comment on the module-level from_pretrained() above for why "auto"
+        # is what triggers the "modules offloaded to cpu or disk" error.
         scoring_dir = os.path.join(load_directory, "scoring_model")
         model.scoring_model = AutoPeftModelForCausalLM.from_pretrained(
-            scoring_dir, 
-            device_map="auto", 
-            low_cpu_mem_usage=True, 
+            scoring_dir,
+            device_map={'': model.device},
+            low_cpu_mem_usage=True,
             use_safetensors=True
         )
 
@@ -301,4 +315,4 @@ class ComputeStat(nn.Module):
             if buffer_name.startswith("null_distr_"):
                 domain = buffer_name.replace("null_distr_", "")
                 domains.append(domain)
-        return domains
+        return domains