Backward time grows linearly to the number of to zero3_consolidated_16bit_state_dict called

I'm training a model with ema states which means module state dict needs to be gathered after each step. When zero stage 3 is enabled model_engine.backward become slower as step grows.

Training loop

train_steps = 0
for _ in epoch:
    for x, y in dataset:
        ts = time()
        model_engine.backward(loss / grac_acc)
        running_loss.append(loss.detach().item() * args.grad_acc_steps)
        backward_seconds += time() - ts
        # ...
        if train_steps % grac_acc == 0:
            model_engine.step()

            if model_engine.zero_optimization_stage() == 3:
                sd = model_engine._zero3_consolidated_16bit_state_dict()
            elif is_global_rank_0():
                sd = model_engine.module_state_dict()

            if is_global_rank_0():
                update_ema(ema, sd)

            tb_writer.add_scalar(
                'Train/Backward time ms',
                backward_seconds * 1000,
                train_steps,
            )
            backward_seconds = 0
            train_steps += 1

if model_engine._zero3_consolidated_16bit_state_dict is disabled, backward time is not growing and everything works fine.

ds version is 0.12.6 Tried on 8 * a800 & 8 * 3090

Start at model_engine._zero3_consolidated_16bit_state_dict and reproduce the supplied training loop with ZeRO stage 3, recording backward time as train_steps increases. Compare runs with state-dict consolidation enabled and disabled; done means backward time no longer grows linearly with the number of calls.