Customize a Strategy Builder

A Closer Look into Strategy

In AutoDist, a Strategy is a representation instructing AutoDist to transform a single-node computational graph into a distributed one.

For more technical details, one can refer to the definition page of the Strategy Protocol Buffer message.

Below is an intuitive example of a strategy for a computational graph with two variables W:0 and b:0. Each variable is assigned with a corresponding configured synchronizer within the node_config section; while the graph_config in the below example configures the data-parallel replicas of the whole graph.

node_config {
  var_name: "W:0"
  PSSynchronizer {
    reduction_destinations: "localhost:CPU:0"
    sync: true
  }
}
node_config {
  var_name: "b:0"
  AllReduceSynchronizer {
    chunk_size: 128
  }
}
graph_config {
  replicas:
    [
        "10.21.1.24:GPU:0",
        "10.21.1.24:GPU:1",
        "10.21.1.25:GPU:0",
        "10.21.1.25:GPU:1"
    ]
}

Build a Strategy

If you understand the strategy representation above, you can create your own customized StrategyBuilder just like the built-in ones. The customized strategy builder needs to follow the StrategyBuilder abstraction, with a required interface build, which takes a GraphItem and a ResourceSpec and returns a Strategy.

def build(self, graph_item: GraphItem, resource_spec: ResourceSpec) -> Strategy:

• Create a strategy representation wrapper object.

from autodist.strategy.base import Strategy
strategy = Strategy()

• Set configurations for the whole graph. For example, you can utilize the resource_spec properties to list all GPU devices for your data parallelism.

strategy.graph_config.replicas.extend([k for k, v in resource_spec.gpu_devices])

• Before configuring nodes, for example, you can utilize the graph_item methods to list all variables that you want to configure or utilize the resource_spec properties to prepare for where to put the variable states.

variables = graph_item.trainable_var_op_to_var.values()
reduction_devices = [k for k, _ in resource_spec.cpu_devices][0:1]

• Set configurations for variable nodes. Besides the below example, there are various choices of configurations listed here strategy_pb2.Strategy.Node section.

from autodist.proto import strategy_pb2
node_config = []
for var in variables:
    node = strategy_pb2.Strategy.Node()
    node.var_name = var.name
    node.PSSynchronizer.reduction_destinations.extend(reduction_devices)
    node.PSSynchronizer.local_replication = False
    node.PSSynchronizer.sync = True
    node.PSSynchronizer.staleness = 2
    node_config.append(node)

strategy.node_config.extend(node_config)

Congratulations! You have successfully created your first StrategyBuilder. AutoDist is flexible, allowing developers to create different types of strategies based on the configuration spaces, and also for auto-learning a strategy. For more developments on strategies, you could refer to other built-in strategy builders or our development reference to invent your own kernels.