ZeRO Stage 3: Calling submodules

Hey, a less abstract description for this issue could also be "ZeRO Stage 3: Best practices for non-forward calls".
I'll start with a high-level description and high-level questions and end with issues and questions toward the more low-level stuff.

Problem Description

I have a model implementing a method generate_images which I'd like to call during training to monitor progress.

This method uses child modules of the main model and calls self multiple times afterwards (example below). This leads me to several questions.

1. How would I wrap this method to support the ZeRO 3-partitioned model?
2. Do I need to wrap the calls to child modules and if yes, how would I wrap them?
3. Can I avoid passing in the engine?

Example of the method:

class MyModel(torch.nn.Module):
    def __init__(self, child):
        self.child = child
        self.other_layers = [...]

    def forward(self, inputs):
        latent_vars = self.child(inputs)
        [...]
        return self.other_layers([...])

    @torch.no_grad()
    def generate_images(self, inputs):
        latent_vars = self.child(inputs)
        images = []
        while len(images) < 16:
            latent_vars = self(latent_vars)
            images.append(latent_vars)
        return image

How I'd like to call it:

child = [...]
model = MyModel(child)
engine = deepspeed.initialize(model, [...])

[...]  # In train loop
    image = model.generate_images(test_inputs)

A non-Solution

I found out that model already has the necessary hooks set up so I would not need to do something like this (which would only work when self.child is not partitioned by chance):

    @torch.no_grad()
    def generate_images(self, inputs):
        call_self = self
        if instance(self, deepspeed.DeepSpeedEngine):
            self = self.module
        [...]
            latent_vars = call_self(latent_vars)
        [...]

[...] # In train loop; notice I pass the engine here.
    image = MyModel.generate_images(engine, test_inputs)

Possible Issues

However, even with the hooks on model, I'm worried whether the partitioned parameter coordinator handles the self.child call at the start of the generate_images method correctly since self.child does not have the _{pre,end_of}_forward_hook which are defined in the hook setup method I linked above. These hooks are only registered for the top-level module.
There's also the mutating state in the ZeRO 3 optimizer's param_coordinator and its global FWD_MODULE_STACK.

In other words, how can I use ZeRO 3's parameter coordinator for passes through submodules of the non-wrapped model? Would I need to wrap all submodule calls in deepspeed.zero.GatheredParameters(submodule.parameters())? Would I need to manually call engine.optimizer.param_coordinator.reset_step()?

Start with the linked hook setup in deepspeed/runtime/zero/stage3.py at lines 1364-1384, then trace the generate_images example through ZeRO 3's parameter coordinator and FWD_MODULE_STACK. Check how calls to self.child and repeated self calls behave, including whether GatheredParameters or reset_step is involved. Done means a verified, documented approach for non-forward submodule calls.