custom aggregators with client_states as states

Question

I want to create a custom aggregator where the state is the unique client state of each client. To initialize I can define client states as usual, and then use federated_collect to place @SERVER placement since thats what initialize_fn() wants. I can also do the same for creating new_state in the next_fn(). The problem is once I don't know how I can "broadcast" these states back into clients. Normally federated_broadcast takes say A@SERVER and then makes copies of it equal to number of clients. so for two clients it would be {A}@CLIENTS lets say (A A). What I want is to have AB@SERVER turning into (A B).

I am currently defining client states outside the aggregation process, and then passing these in run_one_round of iterative process. use federated_collect to collect these states from measurements of aggregator, and then unstack it outside. So from the outside of the federated computations it looks like

    server_state, train_metrics , client_states, aggregation_state = iterative_process.next(
                    server_state, sampled_train_data, client_states, aggregation_state)
    client_states = [x for x in client_states]

In TFF

    output = aggregation_process.next(aggregation_state, client_outputs.weights_delta, client_states)
    new_aggregation_state = output.state
    round_model_delta = output.result
    new_client_states = output.measurements

in aggregator

        measurements = tff.federated_collect(new_client_states)
        return tff.templates.MeasuredProcessOutput(
            state=new_state, result= round_model_delta, measurements=measurements)

But I am trying to define and handle these client states completely inside the aggregator so that I can plug this aggregator into tff.learning.build_federated_averaging_process like

    iterative_process = tff.learning.build_federated_averaging_process(
        model_fn,
        client_optimizer_fn=lambda: tf.keras.optimizers.SGD(learning_rate=0.02),
        server_optimizer_fn=lambda: tf.keras.optimizers.SGD(learning_rate=1.0),
        model_update_aggregation_factory=my_aggregation_factory)

Is that possible? if so how?

Zachary_Garrett_14692 · Answer

tff.federated_collect is likely not the desired tool in this situation, and it will be removed in future versions of TFF (see commit #030a406).

Alternatively, A tff.federated_computation can both take @CLIENTS parameters as input, and return @CLIENT placed values as output. Instead of collecting all the values on the server first (implying that the system is communicating the states); it maybe best to leave the values on the clients.

When executing TFF in a simulation environment (e.g. invoking a tff.Computation in a Colab notebook) a T@CLIENT placed value will be returned as a list of T objects; one for each client. This can later be used as a parameter to future tff.Computation invocation.

Example:

@tff.tf_computation(tf.int32)
def sqrt(value):
  return tf.math.sqrt(tf.cast(value, tf.float32))

@tff.federated_computation(tff.types.at_clients(tf.int32))
def federated_sqrt(values):
  return tff.federated_map(sqrt, values)

client_values = [1,2,3,4]
federated_sqrt(client_values)
>>> [,
 ,
 ,
 ]

Important caveat: the order of inputs and outputs is not necessarily guaranteed to be the same. An example of how to index and track states across invocations can be found in the tensorflow_federated/python/examples/stateful_clients/ directory inside the repository.

custom aggregators with client_states as states

1 Answers