ZeRO 3 not working with GAN example: IndexError: tuple index out of range

Using the GAN example and the following deepspeed config for ZeRO 3 offload, I get the following error:

{
  "train_batch_size": 64,
  "zero_optimization": {
    "stage": 3,
    "contiguous_gradients": true,
    "stage3_max_live_parameters": 1e9,
    "stage3_max_reuse_distance": 1e9,
    "stage3_prefetch_bucket_size": 1e7,
    "stage3_param_persistence_threshold": 1e5,
    "reduce_bucket_size": 1e7,
    "sub_group_size": 1e9,
    "offload_optimizer": {
      "device": "cpu"
    },
    "offload_param": {
      "device": "cpu"
    }
  }
}

Loading extension module utils...
Time to load utils op: 0.00021600723266601562 seconds
[2022-08-02 07:09:09,796] [INFO] [stage3.py:1834:_overflow_clean_up] [deepspeed] OVERFLOW! Rank 0 Skipping step. Attempted loss scale: 4294967296, reducing to 4294967296
[2022-08-02 07:09:09,797] [INFO] [logging.py:69:log_dist] [Rank 0] rank=0 time (ms) | backward_microstep: 261.77 | backward_inner_microstep: 249.30 | backward_allreduce_microstep: 12.28 | step_microstep: 1.40
[2022-08-02 07:09:09,797] [INFO] [logging.py:69:log_dist] [Rank 0] rank=0 time (ms) | backward: 261.77 | backward_inner: 249.32 | backward_allreduce: 12.28 | step: 1.40
[2022-08-02 07:09:09,841] [INFO] [stage3.py:1834:_overflow_clean_up] [deepspeed] OVERFLOW! Rank 0 Skipping step. Attempted loss scale: 4294967296, reducing to 4294967296
[2022-08-02 07:09:09,842] [INFO] [logging.py:69:log_dist] [Rank 0] rank=0 time (ms) | backward_microstep: 38.73 | backward_inner_microstep: 33.22 | backward_allreduce_microstep: 5.40 | step_microstep: 0.63
[2022-08-02 07:09:09,842] [INFO] [logging.py:69:log_dist] [Rank 0] rank=0 time (ms) | backward: 38.73 | backward_inner: 33.22 | backward_allreduce: 5.40 | step: 0.64
[0/1][0/3166] Loss_D: 1.8203 Loss_G: 0.6108 D(x): 0.5308 D(G(z)): 0.5933 / 0.5933
Traceback (most recent call last):
  File "gan_deepspeed_train.py", line 182, in <module>
    main()
  File "gan_deepspeed_train.py", line 179, in main
    train(args)
  File "gan_deepspeed_train.py", line 125, in train
    output = netD(real)
  File "/home/host/.local/bin/anaconda3/envs/pytorchtut/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1137, in _call_impl
    result = hook(self, input)
  File "/home/host/.local/bin/anaconda3/envs/pytorchtut/lib/python3.8/site-packages/deepspeed/utils/nvtx.py", line 11, in wrapped_fn
    return func(*args, **kwargs)
  File "/home/host/.local/bin/anaconda3/envs/pytorchtut/lib/python3.8/site-packages/deepspeed/runtime/zero/parameter_offload.py", line 324, in _pre_forward_module_hook
    self.pre_sub_module_forward_function(module)
  File "/home/host/.local/bin/anaconda3/envs/pytorchtut/lib/python3.8/site-packages/torch/autograd/grad_mode.py", line 27, in decorate_context
    return func(*args, **kwargs)
  File "/home/host/.local/bin/anaconda3/envs/pytorchtut/lib/python3.8/site-packages/deepspeed/runtime/zero/parameter_offload.py", line 444, in pre_sub_module_forward_function
    param_coordinator.trace_prologue(sub_module)
  File "/home/host/.local/bin/anaconda3/envs/pytorchtut/lib/python3.8/site-packages/deepspeed/runtime/zero/partitioned_param_coordinator.py", line 150, in trace_prologue
    if sub_module != self.__submodule_order[self.__step_id]:
IndexError: tuple index out of range

Running with the default configuration in the repo, it works just fine. Running on linux with nvidia 3060 12gb.

Reproduce the GAN example with the supplied ZeRO 3 offload configuration and compare it with the repository's default configuration. Start at gan_deepspeed_train.py, then inspect the traced calls in parameter_offload.py and partitioned_param_coordinator.py. Done means the example completes without the IndexError.