Accelerate 文档
流水线并行
加入 Hugging Face 社区
并获得增强的文档体验
开始使用
流水线并行
Accelerate 通过 PyTorch 的 torch.distributed.pipelining API 支持大规模训练中的流水线并行。
prepare_pippy
accelerate.prepare_pippy
< source >( model split_points: typing.Union[str, list[str], NoneType] = 'auto' no_split_module_classes: typing.Optional[list[str]] = None example_args: typing.Optional[tuple[typing.Any]] = () example_kwargs: typing.Optional[dict[str, typing.Any]] = None num_chunks: typing.Optional[int] = None gather_output: typing.Optional[bool] = False )
参数
- model (
torch.nn.Module) — 我们希望为流水线并行推理进行切分的模型 - split_points (
str或List[str]，默认为 'auto') — 如何生成切分点并在每个 GPU 上对模型进行分块。'auto' 会为任何模型找到最平衡的切分方式。否则，它应该是一个模型中用于切分的层名称列表。 - no_split_module_classes (
List[str]) — 一个包含我们不希望被切分的层类名的列表。 - example_args (模型输入的元组) — 针对单个进程、使用基于顺序的输入的模型的预期输入。如果可能，推荐使用此方法。
- example_kwargs (模型输入的字典) — 针对单个进程、使用基于字典的输入的模型的预期输入。这是一个高度受限的结构，要求在所有推理调用中都存在相同的键。除非前述条件在所有情况下都成立，否则不推荐使用。
- num_chunks (
int，默认为可用 GPU 数量) — 流水线将具有的不同阶段数。默认情况下，它会为每个 GPU 分配一个块，但这可以调整和尝试。通常情况下，应满足 num_chunks >= num_gpus。 - gather_output (
bool，默认为False) — 如果为True，最后一个 GPU（持有真实输出）的输出将被发送到所有 GPU。
为流水线并行推理包装 model。