[zero3] External Parameters doc clarification

In the doc Registering External Parameters

1. at the very end it gives an example:

class ModuleZ3(torch.nn.Module):
    def __init__(self, *args):
        super().__init__(self, *args)
        self.layer1 = SomeLayer()
        self.layer2 = OtherLayer()
        deepspeed.zero.register_external_parameter(self, self.layer1.weight)

    def forward(self, input):
        x = self.layer1(input)
        # self.layer1.weight is required by self.layer2.forward
        y = self.layer2(x, self.layer1.weight)
        return y

If I understand the situation correctly that last comment:

self.layer1.weight is required by self.layer2.forward

doesn't quite fit to to what actually happens.

Isn't the point here that we need to have the weights synced in the current forward and if so shouldn't it say:

self.layer1.weight is required by ModuleZ3.forward

because this is where it's really used. The fact that it's then passed to layer2 is irrelevant because by that time it is already synced and fully present. layer2.forward is oblivious to the arg being actually a weight elsewhere.

I mean that's why ModuleZ3.__init__ calls register_external_parameter and not OtherLayer.__init__

The way I interpret the current comment is that register_external_parameter should actually happen inside OtherLayer.__init__ which I don't think is right.

---

2. And then going back to the beginning. The section starts with a nice practical example and then it drops off and shows a different example. All the suspense was lost ;)

But I find the prose after it hard to understand - it goes:

The tensor embeddings.weight is used in both embeddings.forward() and compute_logits(). We call embeddings.weight an external parameter because it is used in the training loop outside of its owning module’s forward pass.

Am I still keeping the intention correct if I rephrase it as:

The tensor embeddings.weight belongs to embeddings and is used in embeddings.forward() and LanguageModel.forward(). We call embeddings.weight an external parameter in this situation because it's used by LanguageModel.forward() to which it doesn't belong,

[all the rest of the doc]

[and then before the examples start let's finish the first example]

and therefore we need to adjust our initial example to:

class LanguageModel(torch.nn.Module):
    ...
    def __init__(self, ...):
        [....]
        self.embeddings = nn.Embedding(...)
        deepspeed.zero.register_external_parameter(self, self.embeddings.weight)
    
    def forward(self, inputs):
        embeds = self.embeddings(inputs)
        ...
        logits = compute_logits(output, self.embeddings.weight)
        ...

Here again compute_logits() makes things slightly more difficult to understand because it potentially may mean that this function could live on another layer or something like that, but again it's irrelevant since it doesn't know where its argument came from. i.e. it's only relevant inside the LanguageModel.forward. compute_logits() though is definitely less confusing than self.layer2() call from the other example, because the latter does refer to yet another sub-module.

---

The mantra I think is: if layer A's forward refers directly to layer B's weights, layer A needs to register layer's B's weights in A.__init__.

I, of course, could be wrong - just trying to match the prose to the code and vice versa.

I'd be happy to send a PR if either or both of my suggestions may help improve clarity around this topic.

Thank you!

Start with the Registering External Parameters section in zero3.html, especially the ModuleZ3 and LanguageModel examples and the surrounding explanation. Compare the prose with where external weights are used and registered; done means the examples and wording consistently explain the owning module and calling forward method.