Variable Partitioner.

class PartitionerConfig(partition_list=None, partition_str=None)[source]

Bases: object

Helper class to conveniently convert between partition list and partition string.


Serialize a partition list to a partition str if it is valid.


partition_list (List) – A list of integers indicating how many shards to split along each dimension.


A serialized string format of the partition list.

Return type

partition_str (str)


Deserialize a partition string to a partition list and check if it is valid.


partition_str (str) – A serialized string format of the partition list.


A valid partition list.

Return type

partition_list (List)

property partition_str[source]

The partition string indicating the partition config.



property partition_list[source]

The partition list indicating the partition config.



property num_shards[source]

The number of total partitions.



property axis[source]

The axis to partition along with.



class VariablePartitioner(key, node_config: google.protobuf.pyext._message.RepeatedScalarContainer, graph_item: autodist.graph_item.GraphItem)[source]

Bases: autodist.kernel.kernel.Kernel

Partitions a GraphItem’s variables according to the given strategy.

Essentially, this does a few things:

  1. Reads the Strategy and finds which variables should be partitioned (and how).

  2. 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.

  3. 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).

  4. 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.

classmethod apply(*args, **kwargs)[source]

Apply the Kernel transformation.