Accelerate documentation

The inference API

You are viewing v0.31.0 version. A newer version v1.2.1 is available.
Hugging Face's logo
Join the Hugging Face community

and get access to the augmented documentation experience

to get started

The inference API

These docs refer to the PiPPy integration.

accelerate.prepare_pippy

< >

( model split_points: Union = 'auto' no_split_module_classes: Optional = None example_args: Optional = () example_kwargs: Optional = None num_chunks: Optional = None gather_output: Optional = False )

Parameters

  • model (torch.nn.Module) — A model we want to split for pipeline-parallel inference
  • split_points (str or List[str], defaults to ‘auto’) — How to generate the split points and chunk the model across each GPU. ‘auto’ will find the best balanced split given any model. Should be a list of layer names in the model to split by otherwise.
  • no_split_module_classes (List[str]) — A list of class names for layers we don’t want to be split.
  • example_args (tuple of model inputs) — The expected inputs for the model that uses order-based inputs. Recommended to use this method if possible.
  • example_kwargs (dict of model inputs) — The expected inputs for the model that uses dictionary-based inputs. This is a highly limiting structure that requires the same keys be present at all inference calls. Not recommended unless the prior condition is true for all cases.
  • num_chunks (int, defaults to the number of available GPUs) — The number of different stages the Pipeline will have. By default it will assign one chunk per GPU, but this can be tuned and played with. In general one should have num_chunks >= num_gpus.
  • gather_output (bool, defaults to False) — If True, the output from the last GPU (which holds the true outputs) is sent across to all GPUs.

Wraps model for pipeline parallel inference.

< > Update on GitHub