accelerate

e07f2555 - Forward reduce_batch_size_fn through the find_executable_batch_size decorator form (#4160)

Forward reduce_batch_size_fn through the find_executable_batch_size decorator form (#4160) * Forward reduce_batch_size_fn through the decorator form find_executable_batch_size supports being used as a decorator with arguments, which 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) and reduce_batch_size_fn is not forwarded, so a custom back-off is bound on the first call, dropped by the partial, and replaced by the default 0.9 multiplier on the second. Starting at 256 with a halving strategy, the decorator tries 26 batch sizes instead of 5, which is 21 extra out-of-memory-and-retry cycles per search. Both forms now agree. Signed-off-by: Vineeth Sai <vineethsai4444@gmail.com> * Update tests/test_memory_utils.py --------- Signed-off-by: Vineeth Sai <vineethsai4444@gmail.com> Co-authored-by: Marc Sun <57196510+SunMarc@users.noreply.github.com>