Partitioned Optimizer Variable Reloading issue

[hmc-cs-mdrissi]

Describe the problem.

Reloading variables for optimizer with partitioning present fails to actually load variables if load is done first (deferred restoration. The code runs, but accumulator state is same as not loading at all. This is more specific issue follow up to this one.

The code linked here with 3 adjustments can be used to reproduce this. That code does,

toy_model2.build(input_shape=(None, 1))
toy_model2.optimizer.build(toy_model2.trainable_variables)  # type: ignore
print("Loading weights...")
toy_model2.load_weights(weights_path).assert_consumed()

Instead change it to do,

toy_model2.load_weights(weights_path).assert_consumed()
toy_model2.build(input_shape=(None, 1))
toy_model2.optimizer.build(toy_model2.trainable_variables)  # type: ignore
print("Loading weights...")

Also keep partitioner in second model so change,

strategy2 = tf.distribute.experimental.ParameterServerStrategy(cluster_resolver=resolver)

to

strategy2 = tf.distribute.experimental.ParameterServerStrategy(cluster_resolver=resolver, variable_partitioner=partitioner)

Lastly we need to train first model for a couple steps. Something like

toy_model.fit(tf.constant([[1], [2]]]), tf.constant([[0.0], [1.0]]), num_epochs=1, steps_per_epoch=1)

If weights are loaded before building, they fail to restore silently for partitioned optimizers. When no partitioner is used it works fine.

Lastly PS is not necessary for this issue. Just that PS is common way of using partitioner, but if you use tf.variable_creator_scope that produces sharded variables you can cause same issue without PS.

Standalone code to reproduce the issue.

from typing import Any, Mapping

import json
import os
import tempfile
from multiprocessing import Process

import tensorflow as tf

from portpicker import pick_unused_port

__spec__ = None


def create_tf_configs(worker_count: int, ps_count: int):
    """Create TF_CONFIGs for a cluster."""
    cluster_dict: dict[str, list[str]] = {}
    if worker_count:
        cluster_dict["worker"] = [f"localhost:{pick_unused_port()}" for _ in range(worker_count)]
    if ps_count:
        cluster_dict["ps"] = [f"localhost:{pick_unused_port()}" for _ in range(ps_count)]

    cluster_dict["chief"] = [f"localhost:{pick_unused_port()}"]

    tf_configs = []
    for i in range(worker_count):
        tf_configs.append({"cluster": cluster_dict, "task": {"type": "worker", "index": i}})

    for i in range(ps_count):
        tf_configs.append({"cluster": cluster_dict, "task": {"type": "ps", "index": i}})

    tf_configs.append({"cluster": cluster_dict, "task": {"type": "chief", "index": 0}})

    return tf_configs


def _create_process(tf_config: Mapping[str, Any]):
    name = tf_config["task"]["type"] + "_" + str(tf_config["task"]["index"])

    print(f"Starting {name} process...")
    os.environ["TF_CONFIG"] = json.dumps(tf_config)
    p = Process(target=run)
    p.start()


def run():
    resolver = tf.distribute.cluster_resolver.TFConfigClusterResolver()

    task_type = resolver.task_type
    if task_type in ("worker", "ps"):
        print("Starting server...")
        server = tf.distribute.Server(
            resolver.cluster_spec(),
            job_name=resolver.task_type,
            task_index=resolver.task_id,
            protocol=resolver.rpc_layer,
            start=True,
        )
        server.join()

    partitioner = tf.distribute.experimental.partitioners.MaxSizePartitioner(max_shard_bytes=100 * 16 * 4)
    strategy = tf.distribute.experimental.ParameterServerStrategy(
        cluster_resolver=resolver, variable_partitioner=partitioner
    )

    print("Building model...")
    with strategy.scope():
        toy_model = tf.keras.Sequential(
            [tf.keras.layers.Embedding(100, 32), tf.keras.layers.Dense(1, activation="sigmoid")]
        )
        toy_model.compile(loss="binary_crossentropy", optimizer=tf.optimizers.experimental.Adam())
        toy_model.build(input_shape=(None, 1))
        toy_model.optimizer.build(toy_model.trainable_variables)  # type: ignore
        toy_model.fit(tf.constant([[1], [2]]]), tf.constant([[0.0], [1.0]]), num_epochs=1, steps_per_epoch=1)
        print(toy_model.optimizer.variables) 


    print("Saving weights...")
    temp_dir = tempfile.gettempdir()
    weights_path = os.path.join(temp_dir, "model_weights")
    toy_model.save_weights(weights_path)

    strategy2 = tf.distribute.experimental.ParameterServerStrategy(cluster_resolver=resolver, variable_partitioner=partitioner)
    with strategy2.scope():
        toy_model2 = tf.keras.Sequential(
            [tf.keras.layers.Embedding(100, 32), tf.keras.layers.Dense(1, activation="sigmoid")]
        )
        toy_model2.compile(loss="binary_crossentropy", optimizer=tf.optimizers.experimental.Adam())
        toy_model2.load_weights(weights_path)
        toy_model2.build(input_shape=(None, 1))
        toy_model2.optimizer.build(toy_model2.trainable_variables)  # type: ignore
        print(toy_model2.optimizer.variables) # Inconsistent with printed optimizer variables prior.
        print("Loading weights...")
        
    print("Done!")


def main():
    tf_configs = create_tf_configs(2, 1)
    chief_config = tf_configs[-1]
    for tf_config in tf_configs[:-1]:
        _create_process(tf_config)

    os.environ["TF_CONFIG"] = json.dumps(chief_config)
    run()


if __name__ == "__main__":
    main()

[hmc-cs-mdrissi]

import inspect
import tensorflow as tf
from tensorflow.python.distribute.sharded_variable import ShardedVariable


SHARDED_TENSOR_PREFIX = "part_"

def calculate_shard_sizes(shard_count: int, total_size: int) -> list[int]:
    """Calculate shard sizes for a given shard count and total size.

    Any remainder is distributed to the earlier shards.
    """
    step_1 = math.ceil(total_size / shard_count)
    step_2 = math.floor(total_size / shard_count)
    count_1 = total_size % shard_count
    count_2 = shard_count - count_1

    return [step_1] * count_1 + [step_2] * count_2

def _create_variable_shards(var: tf.Variable, partitioner) -> list[tf.Variable]:
    num_shards = partitioner(var.shape, var.dtype, axis=0)
    if num_shards[0] == 1:
        return [var]

    shard_count = num_shards[0]
    shards: list[tf.Variable] = []
    start = 0
    total_size = var.shape[0]
    assert total_size is not None

    var_end_name = var.name.rsplit("/", 1)[-1].removesuffix(":0")
    for index in range(shard_count):
        shard_name = f"{var_end_name}/{SHARDED_TENSOR_PREFIX}{index}"
        size = total_size // shard_count + (1 if total_size % shard_count > index else 0)
        shards.append(tf.Variable(var[start : start + size], name=shard_name))
        start += size

    return shards

def partition_variable_creator(partitioner):
    def _creator(next_creator = None, **kwargs):
        assert next_creator is not None

        def _init_shard_fn(
            initializer: Callable[..., tf.Variable], base_name: str, index: int, start: int, end: int
        ) -> tf.Variable:
            assert "shape" in kwargs
            shard_shape = (end - start,) + kwargs["shape"][1:]

            kwargs["name"] = f"{base_name}/{SHARDED_TENSOR_PREFIX}{index}"
            kwargs["shape"] = shard_shape
            offset = (start,) + (0,) * (len(shard_shape) - 1)
            if "shard_info" in inspect.signature(initializer).parameters:
                kwargs["initial_value"] = initializer(shard_info=base.ShardInfo(shape=shard_shape, offset=offset))
            else:
                kwargs["initial_value"] = initializer(partition_shape=shard_shape, partition_offset=offset)
            var = next_creator(**kwargs)
            assert isinstance(var, tf.Variable)
            return var

        initial_value = kwargs.get("initial_value", None)
        assert initial_value is not None

        dtype = kwargs.get("dtype", None)
        shape = kwargs.get("shape", None)
        if callable(initial_value) and dtype is not None and shape is not None and len(shape) > 0:
            num_shards = partitioner(shape, dtype, axis=0)
            if num_shards[0] == 1:
                return next_creator(**kwargs)

            total_size = shape[0]
            assert total_size is not None

            shard_sizes = calculate_shard_sizes(num_shards[0], total_size)
            shards = []
            start = 0
            original_name = kwargs.get("name")
            assert original_name is not None
            for index, shard_size in enumerate(shard_sizes):
                shards.append(_init_shard_fn(initial_value, original_name, index, start, start + shard_size))
                start += shard_size

            return ShardedVariable(shards)

        var = next_creator(**kwargs)
        assert isinstance(var, tf.Variable)
        if var.shape.rank == 0:
            return var

        shards = _create_variable_shards(var, partitioner)
        if len(shards) == 1:
            return shards[0]

        return ShardedVariable(shards)

    return _creator

is some code I've been using to test this without ps strategy mixed in. Fairly similar make model, train, save, then make again and load before build will have optimizer state not reloaded when partitioner is used.