add gc collect in packing

intel · Feb 13, 2025 · 8e39e07 · 8e39e07
1 parent edc2b03
commit 8e39e07
Show file tree

Hide file tree

Showing 3 changed files with 9 additions and 1 deletion.
diff --git a/auto_round/export/export_to_autogptq/export.py b/auto_round/export/export_to_autogptq/export.py
@@ -11,7 +11,7 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-
+import gc
 
 # MIT License
 #
@@ -116,6 +116,8 @@ def pack_layer(name, model, layer_config, backend, pbar):
             qlayer.pack(layer, scale, zero, None)
         qlayer.to(device)
         pbar.update(1)
+        if pbar.n%50==0:
+            gc.collect()
 
 
 def save_quantized_as_autogptq(output_dir, inplace=True, backend="auto_gptq:exllamav2",

diff --git a/auto_round/export/export_to_autoround/export.py b/auto_round/export/export_to_autoround/export.py
@@ -151,6 +151,9 @@ def pack_layer(name, model, layer_config, backend, pbar):
             set_module(model, name, qlayer)
             clear_memory()
         pbar.update(1)
+        if pbar.n%50==0:
+            import  gc
+            gc.collect()
 
 
 def save_quantized_as_autoround(output_dir, inplace=True, backend="auto_round:exllamav2", **kwargs):

diff --git a/auto_round/export/export_to_awq/export.py b/auto_round/export/export_to_awq/export.py
@@ -66,6 +66,9 @@ def pack_layer(name, model, layer_config, backend, pbar):
         )
         set_module(model, name, q_linear)
         pbar.update(1)
+        if pbar.n % 50 == 0:
+            import gc
+            gc.collect()
 
 
 def save_quantized_as_autoawq(output_dir, inplace=True, **kwargs):