Exploding memory consumption when training FL model with varying number of participants per round

I'm running FL algorithm following the image classification tutorial. The number of participants vary at each round according to a predefined list of participants number.

number_of_participants_each_round = 
[108, 113, 93, 92, 114, 101, 94, 93, 107, 99, 118, 101, 114, 111, 88, 
101, 86, 96, 110, 80, 118, 84, 91, 120, 110, 109, 113, 96, 112, 107, 
119, 91, 97, 99, 97, 104, 103, 120, 89, 100, 104, 104, 103, 88, 108]

The federated data is preprocessed and batched before starting the training.

NUM_EPOCHS = 5
BATCH_SIZE = 20
SHUFFLE_BUFFER = 418
PREFETCH_BUFFER = 10

def preprocess(dataset):
    def batch_format_fn(element):
        return collections.OrderedDict(
            x=tf.reshape(element['pixels'], [-1, 784]),
            y=tf.reshape(element['label'], [-1, 1]))

    return dataset.repeat(NUM_EPOCHS).shuffle(SHUFFLE_BUFFER).batch(
        BATCH_SIZE).map(batch_format_fn).prefetch(PREFETCH_BUFFER)


def make_federated_data(client_data, client_ids):
    return [preprocess(client_data.create_tf_dataset_for_client(x)) for x in client_ids]

federated_train_data = make_federated_data(data_train, data_train.client_ids)

Participants are randomly sampled from federated_train_data[0:expected_total_clients] at each round according to the number_of_participants_each_round, then the iterative_process is executed for 45 rounds.

expected_total_clients = 500
round_nums = 45

for round_num in range(0, round_nums):
   sampled_clients = 
       np.random.choice(a=federated_train_data[0:expected_total_clients],                          
                        size=number_of_participants_each_round[round_num], 
                        replace=False)
    
   state, metrics = iterative_process.next(state, list(sampled_clients))
   print('round {:2d}, metrics={}'.format(round_num + 1, metrics))

The problem is that the VRAM usage quickly explodes after few rounds, it reaches 5.5 GB at round 6~7, and keeps increasing with an approx rate of 0.8 GB/round until the training eventually crashes at round 25~26 where the VRAM reaches 17 GB with +4000 python threads created.

Error message below

F tensorflow/core/platform/default/env.cc:72] Check failed: ret == 0 (35 vs. 0)Thread creation via pthread_create() failed.

### Troubleshooting ###

Reducing the number_of_participants_each_round to 20 allows the training to complete, but the memory consumption was still huge and growing.

Running the same code with fixed number of participants per round, memory consumption was fine with total of approx 1.5 ~ 2.0 GB VRAM throughout the entire training.

expected_total_clients = 500
fixed_client_size_per_round = 100
round_nums = 45

for round_num in range(0, round_nums):
   sampled_clients =
      np.random.choice(a=federated_train_data[0:expected_total_clients],
                       size=fixed_client_size_per_round,
                       replace=False)
    
    state, metrics = iterative_process.next(state, list(sampled_clients))
    print('round {:2d}, metrics={}'.format(round_num + 1, metrics))

Extra details:

OS: MacOS Mojave, 10.14.6
python -V: Python 3.8.5 then downgraded to Python 3.7.9
TF version: 2.4.1
TFF version: 0.18.0
Keras version: 2.4.3

Is this a normal memory behaviour or a bug? Are there any refactoring/hints to optimize memory consumption ?