流水線並行

Accelerate 透過 PyTorch 的 torch.distributed.pipelining API 支援大規模訓練中的流水線並行。

prepare_pippy

accelerate.prepare_pippy

( 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。

< > 在 GitHub 上更新