From 4de81cd60f6bbeb29ac053082112406ecebae3c9 Mon Sep 17 00:00:00 2001
From: icn5381 <255778606+icn5381@users.noreply.github.com>
Date: Mon, 17 Aug 2026 03:19:45 +0800
Subject: [PATCH] Forward reduce_batch_size_fn in find_executable_batch_size
 decorator factory

The parenthesized decorator form dropped the reduce_batch_size_fn
argument when building its functools.partial, so a custom batch-size
reducer was silently ignored and the default multiply-by-0.9 behavior
was always used. The custom reducer is only reachable through this
factory form, which made the parameter unusable despite being
documented.
---
 src/accelerate/utils/memory.py |  6 +++++-
 tests/test_memory_utils.py     | 19 +++++++++++++++++++
 2 files changed, 24 insertions(+), 1 deletion(-)

diff --git a/src/accelerate/utils/memory.py b/src/accelerate/utils/memory.py
index a5fd255b950..8cd6cebb851 100644
--- a/src/accelerate/utils/memory.py
+++ b/src/accelerate/utils/memory.py
@@ -152,7 +152,11 @@ def find_executable_batch_size(
     ```
     """
     if function is None:
-        return functools.partial(find_executable_batch_size, starting_batch_size=starting_batch_size)
+        return functools.partial(
+            find_executable_batch_size,
+            starting_batch_size=starting_batch_size,
+            reduce_batch_size_fn=reduce_batch_size_fn,
+        )
 
     batch_size = starting_batch_size
     if reduce_batch_size_fn is None:
diff --git a/tests/test_memory_utils.py b/tests/test_memory_utils.py
index 69283ab40a8..3fb46698053 100644
--- a/tests/test_memory_utils.py
+++ b/tests/test_memory_utils.py
@@ -126,6 +126,25 @@ def mock_training_loop_function(batch_size, arg1):
         ]
         assert [bs, arg1] == [8, "hello"]
 
+    def test_custom_reduce_batch_size_fn(self):
+        # The parenthesized decorator form must forward `reduce_batch_size_fn`
+        # to the wrapped call: the custom reducer was silently dropped before,
+        # falling back to the default *0.9 behavior.
+        calls = []
+
+        def reduce_to_one():
+            calls.append("called")
+            return 1
+
+        @find_executable_batch_size(starting_batch_size=128, reduce_batch_size_fn=reduce_to_one)
+        def mock_training_loop_function(batch_size):
+            if batch_size > 1:
+                raise_fake_out_of_memory()
+            return batch_size
+
+        assert mock_training_loop_function() == 1
+        assert calls == ["called"]
+
     def test_start_zero(self):
         @find_executable_batch_size(starting_batch_size=0)
         def mock_training_loop_function(batch_size):