Partitioner¶
Variable Partitioner.
-
class
PartitionerConfig(partition_list=None, partition_str=None)[source]¶ Bases:
objectHelper class to conveniently convert between partition list and partition string.
-
serialize(partition_list)[source]¶ Serialize a partition list to a partition str if it is valid.
- Parameters
partition_list (List) – A list of integers indicating how many shards to split along each dimension.
- Returns
A serialized string format of the partition list.
- Return type
partition_str (str)
-
deserialize(partition_str)[source]¶ Deserialize a partition string to a partition list and check if it is valid.
- Parameters
partition_str (str) – A serialized string format of the partition list.
- Returns
A valid partition list.
- Return type
partition_list (List)
-
-
class
VariablePartitioner(key, node_config: google.protobuf.pyext._message.RepeatedScalarContainer, graph_item: autodist.graph_item.GraphItem)[source]¶ Bases:
autodist.kernel.kernel.KernelPartitions a GraphItem’s variables according to the given strategy.
Essentially, this does a few things:
Reads the Strategy and finds which variables should be partitioned (and how).
Creates new PartitionedVariables (essentially a list of Variables) and replaces the original Variable with these new vars. Also splits the original gradient to map to each variable shard.
Deletes the original variables and recreates the Optimizer (since Optimizers can often have different logic for PartitionedVariables, and we don’t want to rewire the optimizer internals as well).
Returns a new graph and strategy modified to reflect the partitioning.
Note that this currently does not work with variables that are part of a control flow (including batchnorm). In those cases, we currently do not have a way to replicate the control flow’s FuncGraph for our new variables. Thus, we expect that any StrategyBuilder does not generate a strategy that tries to partition a control flow variable, otherwise this will likely error.