Forward reduce_batch_size_fn through the find_executable_batch_size decorator form

[vineethsaivs]

Problem

find_executable_batch_size supports being used as a decorator with arguments, and that is the form its own docstring example uses. That path goes through:

if function is None:
    return functools.partial(find_executable_batch_size, starting_batch_size=starting_batch_size)

reduce_batch_size_fn is not forwarded. It is bound on the first call, dropped by the partial, and on the second call it is None again, so the default multiply-by-0.9 back-off runs instead.

current = 256
def halve():
    global current
    current //= 2
    return current

@find_executable_batch_size(starting_batch_size=256, reduce_batch_size_fn=halve)
def train(batch_size): ...

# batch sizes tried:
#   decorator form  [256, 230, 207, 186, 167, 150, 135, 121, ...]   26 attempts
#   direct call     [256, 128, 64, 32, 16]                           5 attempts

Passing the function directly, find_executable_batch_size(train, starting_batch_size=256, reduce_batch_size_fn=halve), already works, so the two documented forms disagree.

Every one of those extra attempts is a real out-of-memory-and-retry cycle: allocate, fail, clear_device_cache(garbage_collection=True), try again. A halving strategy from 256 should cost 5 of them and costs 26. Nothing errors and nothing warns, so the only symptom is a slow start-up and a batch size that walks down in 10% steps rather than the steps the caller asked for.

reduce_batch_size_fn was added in #3071 and documented in #4051; only the forwarding through the partial was missed.

Fix

Forward the third argument.

Test

Three cases added to tests/test_memory_utils.py, next to the existing find_executable_batch_size tests, which all use the decorator form with starting_batch_size only:

• the decorator form honours a custom back-off
• the decorator form and the direct-call form agree on the same input
• the default 0.9 back-off is unchanged when no function is given

                                                 before   after
test_reduce_batch_size_fn_is_used_by_the_decorator  FAIL    pass
test_reduce_batch_size_fn_matches_the_direct_call   FAIL    pass
test_default_back_off_is_unchanged                  pass    pass
the 7 pre-existing tests in the file                pass    pass

pytest tests/test_memory_utils.py gives 2 failed / 8 passed against main and 10 passed with this change, on CPU. Both runs use src/ from a tree checked out at 16cb6eb, verified by printing accelerate.utils.memory.__file__, because an editable install elsewhere on the machine will otherwise shadow the checkout and make the before and after runs test the same code.

ruff check and ruff format --check are clean on both changed files. Note that ruff --fix also wants to remove a pre-existing pass at tests/test_memory_utils.py:143; I put it back so this diff is only the change described here.