‘tf.distribute.experimental.ParameterServerStrategy‘官方文档取之我用

0 前言

由于对ps策略具体原理不甚了解，导致编程后运行出错，分析应该是分布式数据集构建和输入时出现了错误，故对此文档进行翻译并着重理解自己需要的部分。

官方文档参考链接

1 总述

1.1 别名

也称作tf.distribute.ParameterServerStrategy，但实际上实践后，tf2.4版本并不支持，不知道更高版本是否支持没有试过（截止目前已经是tf2.9版本了）。

1.2 概述

Parameter server training是一种常用的数据并行方法，用以将模型训练扩展到多台机器上。

一个参数服务器训练集群包括workers和parameter servers。变量在parameter servers上被创建，并且在训练的每一步中都被workers读取和更新。默认情况下，workers之间不同步地独立地进行读取和更新变量。这也是为什么Parameter server training被称为异步训练的原因。

在tf2.x中，我们推荐一种基于中心协调的Parameter server training架构。coordinator使用tf.distribute.experimental.coordinator.ClusterCoordinator来协调整个集群，并且通过 tf.distribute.experimental.ParameterServerStrategy来实现在parameter servers上定义变量在workers上进行训练。

每台worker只处理来自coordinator的指令以及与ps通信，并且，workers之间并没有直接的联系。这种机制的效果是提高了集群的容错性，可允许其中的某台worker故障，而不影响整个集群的可用性。但是，coordinator和ps必须一直可用，保证集群的进展。

注意，coordinator并不是其中一个用来训练、计算的机器，相反它的工作是创建变量或数据集之类的资源，调度tf.function，以及保存检查点等等。
另外，除了worker、ps、coordinator这三种角色之外，一个可选的选项是evaluator，在集群旁边周期性地读取coordinator保存的检查点并运行进行评估。

当采用自定义循环训练方式时，ParameterServerStrategy必须要创建一个tf.distribute.experimental.coordinator.ClusterCoordinator对象。

当采用Model.fit时，目前只支持tf.keras.utils.experimental.DatasetCreator这种输入类型。

1.3 coordinator的样例代码

以下只是一个代码片段，cluster_resolver, variable_partitioner和 dataset_fn等重要参数，分别会在后面的‘集群设置’，‘变量分片’和‘数据集准备’部分详细讲到。

1.3.1 CTL

# Prepare a strategy to use with the cluster and variable partitioning info.
strategy = tf.distribute.experimental.ParameterServerStrategy(
    cluster_resolver=...,
    variable_partitioner=...)
coordinator = tf.distribute.experimental.coordinator.ClusterCoordinator(
    strategy=strategy)

# Prepare a distribute dataset that will place datasets on the workers.
distributed_dataset = coordinator.create_per_worker_dataset(dataset_fn=...)

with strategy.scope():
  model = ...
  optimizer, metrics = ...  # Keras optimizer/metrics are great choices
  checkpoint = tf.train.Checkpoint(model=model, optimizer=optimizer)
  checkpoint_manager = tf.train.CheckpointManager(
      checkpoint, checkpoint_dir, max_to_keep=2)
  # `load_checkpoint` infers initial epoch from `optimizer.iterations`.
  initial_epoch = load_checkpoint(checkpoint_manager) or 0

@tf.function
def worker_fn(iterator):

  def replica_fn(inputs):
    batch_data, labels = inputs
    # calculate gradient, applying gradient, metrics update etc.

  strategy.run(replica_fn, args=(next(iterator),))

for epoch in range(initial_epoch, num_epoch):
  distributed_iterator = iter(distributed_dataset)  # Reset iterator state.
  for step in range(steps_per_epoch):

    # Asynchronously schedule the `worker_fn` to be executed on an arbitrary
    # worker. This call returns immediately.
    coordinator.schedule(worker_fn, args=(distributed_iterator,))

  # `join` blocks until all scheduled `worker_fn`s finish execution. Once it
  # returns, we can read the metrics and save checkpoints as needed.
  coordinator.join()
  logging.info('Metric result: %r', metrics.result())
  train_accuracy.reset_states()
  checkpoint_manager.save()

1.3.2 Model.fit

# Prepare a strategy to use with the cluster and variable partitioning info.
strategy = tf.distribute.experimental.ParameterServerStrategy(
    cluster_resolver=...,
    variable_partitioner=...)

# A dataset function takes a `input_context` and returns a `Dataset`
def dataset_fn(input_context):
  dataset = tf.data.Dataset.from_tensors(...)
  return dataset.repeat().shard(...).batch(...).prefetch(...)

# With `Model.fit`, a `DatasetCreator` needs to be used.
input = tf.keras.utils.experimental.DatasetCreator(dataset_fn=...)

with strategy.scope():
  model = ...  # Make sure the `Model` is created within scope.
model.compile(optimizer="rmsprop", loss="mse", steps_per_execution=..., ...)

# Optional callbacks to checkpoint the model, back up the progress, etc.
callbacks = [tf.keras.callbacks.ModelCheckpoint(...), ...]

# `steps_per_epoch` is required with `ParameterServerStrategy`.
model.fit(input, epochs=..., steps_per_epoch=..., callbacks=callbacks)

1.3.3 两者总结对比

1. Model.fit不需要初始化创建一个coordinator实例。
2. 两者创建分布式数据集的方式不一样，CTL使用coordinator.create_per_worker_dataset(dataset_fn=...)来创建；而Model.fit使用tf.keras.utils.experimental.DatasetCreator创建。
3. ParameterServerStrategy采用Model.fit方式训练时，steps_per_epoch参数是必要的！

1.4 worker和ps的样例代码

# Provide a `tf.distribute.cluster_resolver.ClusterResolver` that serves
# the cluster information. See below "Cluster setup" section.
cluster_resolver = ...

server = tf.distribute.Server(
    cluster_resolver.cluster_spec(),
    job_name=cluster_resolver.task_type,
    task_index=cluster_resolver.task_id,
    protocol="grpc")

# Blocking the process that starts a server from exiting.
server.join()

没有什么特别的。

1.5 集群设置

coordinator、worker、ps都需要使用tf.distribute.cluster_resolver.ClusterResolver来进行集群解析，负责提供集群信息，task类型以及目前task的index。
（有个疑问就是，Model.fit也需要吗？看样例片段没有啊）
（回答：嗷嗷，是有的。没有的是coordinator实例）

如果设置了TF_CONFIG环境变量，那么tf.distribute.cluster_resolver.TFConfigClusterResolver也需要使用。

在tf.distribute.cluster_resolver.ClusterResolver里边，需要使用if...else...语句来指定角色：coordinator、workers、parameter servers。

1.6 在strategy.scope()内创建变量

为了将变量以一种循环的方式被正式地分配到parameter servers上，变量的创建被期待在strategy.scope()这个上下文范围管理器的内部。
以下是一个样例代码：

# In this example, we're assuming having 3 ps.
strategy = tf.distribute.experimental.ParameterServerStrategy(
    cluster_resolver=...)
coordinator = tf.distribute.experimental.coordinator.ClusterCoordinator(
    strategy=strategy)

# Variables should be created inside scope to be placed on parameter servers.
# If created outside scope such as `v1` here, it would be placed on the
# coordinator.
v1 = tf.Variable(initial_value=0.0)

with strategy.scope():
  v2 = tf.Variable(initial_value=1.0)
  v3 = tf.Variable(initial_value=2.0)
  v4 = tf.Variable(initial_value=3.0)
  v5 = tf.Variable(initial_value=4.0)

# v2 through v5 are created in scope and are distributed on parameter servers.
# Default placement is round-robin but the order should not be relied on.
assert v2.device == "/job:ps/replica:0/task:0/device:CPU:0"
assert v3.device == "/job:ps/replica:0/task:1/device:CPU:0"
assert v4.device == "/job:ps/replica:0/task:2/device:CPU:0"
assert v5.device == "/job:ps/replica:0/task:0/device:CPU:0"

为什么需要定义一个coordinator实例呢？其他部分理解没有问题。

1.7 变量分片

在ps之间对大规模变量进行分片操作有助于提升训练吞吐量并缓解容量炼制。分片使得变量的不同shards可以并行计算和更新，在ps之间提供一个更好的负载均衡。

如果variable_partitioner设置在__init__中并且能满足特定的条件，在scope中定义的变量便会在ps之间以循环的方式被分片。从tf.Variable返回的变量引用将作为一种包含着分片变量的容器类型。

# Partition the embedding layer into 2 shards.
variable_partitioner = (
  tf.distribute.experimental.partitioners.MinSizePartitioner(
    min_shard_bytes=(256 << 10),
    max_shards = 2))
strategy = tf.distribute.experimental.ParameterServerStrategy(
  cluster_resolver=...,
  variable_partitioner = variable_partitioner)
with strategy.scope():
  embedding = tf.keras.layers.Embedding(input_dim=1024, output_dim=1024)
assert len(embedding.variables) == 2
assert isinstance(embedding.variables[0], tf.Variable)
assert isinstance(embedding.variables[1], tf.Variable)
assert embedding.variables[0].shape == (512, 1024)
assert embedding.variables[1].shape == (512, 1024)

变量分片的限制：

1. 在检查点保存和上传时分片的数量不能改变。
2. 在将分片后的变量保存到一个SavedModel中，这个SavedModel不能通过tf.saved_model.load进行加载。
3. Partition variable doesn’t directly work with tf.GradientTape, please use the variables attributes to get the actual variable components and use them in gradient APIs instead.

1.8 数据集准备

需要创建一个不传入参数的dataset_fn，并且返回一个tf.data.Dataset，然后将dataset_fn传入tf.distribute.experimental.coordinator中的ClusterCoordinator.create_per_worker_dataset。我们推荐将数据集打乱并不断重复来使得训练的样本尽可能地平均。

def dataset_fn():
  filenames = ...
  dataset = tf.data.Dataset.from_tensor_slices(filenames)

  # Dataset is recommended to be shuffled, and repeated.
  return dataset.shuffle(buffer_size=...).repeat().batch(batch_size=...)

coordinator =
    tf.distribute.experimental.coordinator.ClusterCoordinator(strategy=...)
distributed_dataset = coordinator.create_per_worker_dataset(dataset_fn)

可是coordinator实例不是CTL才需要的吗？

面临的限制：

1. 这个策略是实验性的，容易受到将来变化的影响。
2. 当使用Model.fit，必须使用tf.keras.utils.experimental.DatasetCreator创建数据集，并且参数steps_per_epoc在训练时必须指定。

2 参数

2.1 参数说明

参数	描述
cluster_resolver	是一个tf.distribute.cluster_resolver.ClusterResolver实例
variable_partitioner	1. 常用的一个分片器是MinSizePartitioner(min_shard_bytes = 256 << 10, max_shards = num_ps)，指定了每个shard最小256k并且每个ps最多只能获得1个shard。

2.1 参数性质

参数	性质
cluster_resolver	1. 单机器的策略通常没有此组件。 2. multi-worker策略常使用。3. ParameterServerStrategy没有提到