chop.passes.graph.analysis.autosharding
autosharding_analysis_pass
- chop.passes.graph.analysis.autosharding.autosharding_analysis_pass(mg, pass_args: dict = {})[source]
Annotate the metadata of each operator in the graph with a parallelization strategy.
- Parameters:
mg (MaseGraph) – input mase graph.
pass_args (dict, optional) – pass arguments. Defaults to {}.
- Returns:
annotated mase graph.
- Return type:
The pass_args dictionary expects the following elements.
mesh_shape -> tuple : Shape of the device cluster. Should be a 2-dimensional tuple.
inter_node_bandwidth -> int : Inter-node bandwidth, i.e. between GPU nodes.
intra_node_bandwidth -> int : Intra-node bandwidth, i.e. between GPU devices in each node.
Additionally, the following elements can be passed.
algo (optional) -> str : Sharding algorithm to use. Default is “alpa”.
communications_backend (optional) -> str : Communications backend to use, e.g. “nccl” or “gloo”. Default is “nccl”.
skip_fully_replicated (optional) -> bool : If set to true, do not consider fully replicated sharding as an option for any operator.
time_limit (optional) -> int : Time limit for the ILP solver, in seconds. Default is 10000.
mip_rel_gap (optional) -> int : MIP relative gap for the ILP solver. Default is 0 (i.e. obtain full solution).
run_checks (optional) -> bool : If set to true, run checks on the autosharding solution. Default is False.
preload_solution (optional) -> bool : If set to true, preload autosharding solution from file.
ilp_solution_file (optional) -> str : File to export the autosharding solution to. Defaults to: “ilp_solution.pkl”.
alpa_autosharding_pass
- chop.passes.graph.analysis.autosharding.alpa.alpa_autosharding_pass(mg, mesh, pass_args={})[source]
A lightweight implementation of the core algorithm from the Alpa paper: https://arxiv.org/abs/2201.12023
alpa_intra_op_sharding_pass
- chop.passes.graph.analysis.autosharding.alpa_intra_operator.alpa_intra_op_sharding_pass(mg, mesh, pass_args={}, debug=False)[source]
Intra-operator auto parallelization pass from the Alpa paper: https://arxiv.org/abs/2201.12023