Revert default behavior of get_state_dict_from_offload (#3253)

* change default to None Signed-off-by: Kyle Sayers <[email protected]> * introduce move_to_device argument Signed-off-by: Kyle Sayers <[email protected]> * remove move_to_device Signed-off-by: Kyle Sayers <[email protected]> --------- Signed-off-by: Kyle Sayers <[email protected]>
huggingface · Dec 2, 2024 · f8c77f0 · f8c77f0
1 parent b626ef5
commit f8c77f0
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/src/accelerate/utils/modeling.py b/src/accelerate/utils/modeling.py
@@ -1784,6 +1784,10 @@ def get_state_dict_from_offload(
 
     root = module_name[: module_name.rfind(".")]  # module name without .weight or .bias
 
+    # do not move parameters if the module is not offloaded
+    if not has_offloaded_params(module):
+        device_to_put_offload = None
+
     # assign the device to which the offloaded parameters will be sent
     with align_module_device(module, device_to_put_offload):
         for m_key, params in module.state_dict().items():