DOC: document missing parameters in load_accelerator_state, find_executable_batch_size, and send_to_device (#4051)

Three public API functions were missing parameters from their docstrings:

- load_accelerator_state(): dataloaders was a required parameter with no
  documentation despite being present in the function signature.

- find_executable_batch_size(): reduce_batch_size_fn was undocumented.
  It allows users to provide a custom batch size reduction strategy
  instead of the default multiply-by-0.9 behavior.

- send_to_device(): non_blocking and skip_keys were both undocumented.
  non_blocking enables async GPU transfers; skip_keys lets callers
  exclude specific dict keys from device transfer.

Author: kratos0718

src/accelerate/checkpointing.py

@@ -204,6 +204,8 @@ def load_accelerator_state(
             A list of optimizer instances
         schedulers (`List[torch.optim.lr_scheduler._LRScheduler]`):
             A list of learning rate schedulers
+        dataloaders (`List[torch.utils.data.DataLoader]`):
+            A list of dataloader instances used in your program
         process_index (`int`):
             The current process index in the Accelerator state
         scaler (`torch.amp.GradScaler`, *optional*):

src/accelerate/utils/memory.py

@@ -132,6 +132,10 @@ def find_executable_batch_size(
             A function to wrap
         starting_batch_size (`int`, *optional*):
             The batch size to try and fit into memory
+        reduce_batch_size_fn (`callable`, *optional*):
+            A function to determine the new batch size after an out-of-memory error. If not
+            provided, the batch size is multiplied by 0.9 on each failure. The function takes
+            no arguments and should return the new (reduced) batch size as an `int`.
 
     Example:
 

src/accelerate/utils/operations.py

@@ -142,6 +142,12 @@ def send_to_device(tensor, device, non_blocking=False, skip_keys=None):
             The data to send to a given device.
         device (`torch.device`):
             The device to send the data to.
+        non_blocking (`bool`, *optional*, defaults to `False`):
+            If `True`, the transfer to the device is performed asynchronously, which can overlap
+            data movement with computation. Only effective when the device supports it (e.g. CUDA).
+        skip_keys (`str` or `List[str]`, *optional*):
+            A key or list of keys in a dictionary `tensor` whose values should not be sent to
+            the given `device`. Entries with these keys are left on their original device.
 
     Returns:
         The same data structure as `tensor` with all tensors sent to the proper device.