slapo.pipeline

Pipeline stage wrappers for supported frameworks.

Functions:

analyze_tie_weights(top_mod, ...)

Analyze if there is any tie weights (two weights in different module share the same memory) partitioned into different pipeline stages.

slapo.pipeline.analyze_tie_weights(top_mod, is_pipeline_partitioned)

Analyze if there is any tie weights (two weights in different module share the same memory) partitioned into different pipeline stages.

Parameters

• top_mod (torch.nn.Module) – The top-level module. This should be a top pipeline module, so 1) it should already be traced and partitioned, and 2) it should have a number of submodules that matches pipeline stages.

• is_pipeline_partitioned (bool) – Whether the module is partitioned for pipeline or not. If not, then all tie weights will have stage ID 0.

Returns

tie_groups – Mapping from the nn.Parameter object to the set of parameter names that are tied to it. The set of parameter names is a tuple of (parameter name, stage ID). The stage ID is 0 if the module is not partitioned for pipeline.

Return type

Dict[int, Set[Tuple[str, int]]]