Explicit naming of modules

[sanchit-gandhi]

Give two modules: SimpleModule and WrapperModule, the former which takes an nn.Module as an argument, is it possible to explicitly list the nn.Module passed into SimpleModule under both the WrapperModule and SimpleModule?

Here is a code snippet that clarifies my question:

from flax import linen as nn
from jax import random
from flax.core import freeze, frozen_dict
from flax.core.frozen_dict import FrozenDict

# define a simple linear module that takes an nn.Dense module as an input
class SimpleModule(nn.Module):
    features: int
    dense1: nn.Dense
    
    def setup(self):
        # define a second dense layer
        self.dense2 = nn.Dense(self.features)

    def __call__(self, inputs, deterministic):
        x = self.dense1(inputs)
        x = self.dense2(x)
        return x

# define a wrapper module that passes an nn.Dense module to SimpleModule
class WrapperModule(nn.Module):
    features: int

    def setup(self):
        self.dense1 = nn.Dense(self.features)
        self.module = SimpleModule(features=self.features, dense1=self.dense1)

    def __call__(self, inputs):
        x = self.module(inputs)
        return x

# generate random keys
key1, key2 = random.split(random.PRNGKey(0), 2)
x = random.uniform(key1, (4,4))

# initialise the model
model = WrapperModule(features=5)
params = model.init(key2, x)

# utility function for viewing the param tree
def print_tree(d, depth=0, print_value=False):
    for k in d.keys():
        if isinstance(d[k], FrozenDict):
            print('  ' * depth, k)
            print_tree(d[k], depth + 1, print_value)
        else:
            if print_value:
                print('  ' * depth, k, d[k])
            else:
                print('  ' * depth, k)

print_tree(params)

Output:

 params
   dense1
     kernel
     bias
   module
     dense2
       kernel
       bias

With this model definition, the nn.Module dense1 is only listed under the WrapperModule. As an argument to SimpleModule, and appearing in the call function of this module, is there a way of explicitly listing it under SimpleModule too? The parameter tree would then look something as follows:

 params
   dense1
     kernel
     bias
   module
     dense1 <- listed here in addition to the wrapper
       kernel
       bias
     dense2
       kernel
       bias

The reason I ask this is that I am endeavouring to port a PyTorch model to Flax. For this, I need the PyTorch state dict to match the Flax param dict. For the demonstration example above, the equivalent PyTorch model is constructed as follows:

import torch
from torch import nn
from typing import Optional

class SimpleModule(nn.Module):
    def __init__(self, features: int, dense1: Optional[nn.Linear] = None):
        super(SimpleModule, self).__init__()
        
        if dense1 is not None:
            self.dense1 = dense1
        else:
            self.dense1 = nn.Dense(features, features)
            
        self.dense2 = nn.Linear(features, features)

    def forward(self, inputs):
        x = self.dense1(inputs)
        x = self.dense2(x)
        return x

class WrapperModule(nn.Module):
    def __init__(self, features: int):
        super(WrapperModule, self).__init__()
        self.dense1 = nn.Linear(features, features)
        self.module = SimpleModule(features=features, dense1=self.dense1)

    def forward(self, inputs):
        x = self.module(inputs)
        return x

model = WrapperModule(features=5)

print(model)

Output:

WrapperModule(
  (dense1): Linear(in_features=5, out_features=5, bias=True)
  (module): SimpleModule(
    (dense1): Linear(in_features=5, out_features=5, bias=True)
    (dense2): Linear(in_features=5, out_features=5, bias=True)
  )

We see here that dense1 is listed under both the wrapper and simple modules. For the PyTorch state dict to match the Flax param dict, we need dense1 to be listed under both modules in Flax under this name.

[marcvanzee]

Disclaimer: I am not a Pytorch expert.

It seems a Pytorch state dict and Flax parameters are essential two different things, so you shouldn't need to make them the same if you want your Flax and Pytorch models to behave the same.

Because Flax Modules are stateless, we do not use a state dict (Modules are instantiated with parameters only during apply and init, and immediately destroyed afterwards). A Flax parameter dict is exactly what the name says: a dictionary of all parameters used in the Module. While you define the Module inside WrapperModule, the parameters are only use in the submodule, so that is where they will appear in the parameter dict. If you want to have a parameter dict where they appear in both places, you need to use it in both places, but that would be a different Module.

If you want to make the state dict the same as the param dict, you can move the definition of the dense inside the submodule in the Pytorch code...

@andsteing can maybe help here as well since he wrote the Pytorch conversion HOWTO -- Andreas, maybe we can add a note about this discussion in the HOWTO?)

[sanchit-gandhi]

Thanks for the reply @marcvanzee! The problem I'm having is in trying to port a PyTorch model into Flax - because there is an inherently different structure between the state-dict of the PyTorch model and the parameter tree of the Flax one, the two models are incompatible. What I'm hoping to be able to do is define a Flax model with module names and structures that match one-to-one with the PyTorch ones. This will then facilitate the conversion of weights from PyTorch to Flax using the following weight conversion script: https://github.com/huggingface/transformers/blob/b24201fa44e1a14e83be890dcbc231e926c37ec1/src/transformers/modeling_flax_pytorch_utils.py#L41

While you define the Module inside WrapperModule, the parameters are only use in the submodule, so that is where they will appear in the parameter dict.

In my previous code snippet, I defined the module dense1 under the WrapperModule. I then passed this model into the SimpleModule - only here is it actually used. Since dense1 is only used in the call method for the SimpleModule, should it not appear in the parameter dict here? If we inspect the parameter tree, we see that dense1 is in fact listed under the WrapperModule, and not the SimpleModule as anticipated:

 params
   dense1
     kernel
     bias
   module <- dense1 not listed under the SimpleModule
     dense2
       kernel
       bias

If you want to make the state dict the same as the param dict, you can move the definition of the dense inside the submodule in the Pytorch code...

 params <- dense1 would then be omitted from here, which we (unfortunately!) need to match the PT state-dict
   module
     dense1
       kernel
       bias
     dense2
       kernel
       bias

Thank you for your help! Much appreciated!

[marcvanzee]

Sorry, I meant the params will only appear where the Module is registered, which is in this case in WrapperModule's setup method.

So in your case the Torch model is given (and fixed), and you are trying to implement a Flax model that will fit the torch state_dict?

I don't think there currently is a clean way of doing this in Flax natively for the example you give above. Our parameter tree represents a shared Module with a single entry in the variable dict, while Pytorch seems to add this entry everywhere the Module is defined and used.

This does seem like a common use case and i think we should try to find a solution for it. I see two possible solutions:

1. Write some utilty function that converts a Pytorch state dict to a Pytorch parameter dict, where a unique Module is only represented where it is defined.
2. (Somewhat Dirty) Use stubs in the Flax code that will make sure it matches with the Pytorch state. These stubs are unused.

For (1), I do not know enough about Pytorch to know how much effort this would take. Do you have an idea? Would be a very useful thing for us to implement!

For (2), this could look as below in your code. It isn't very clean, but maybe eonough if you just want to get something working quickly. Note that it should not be a big problem in performance if you jit you code since XLA will use dead code elimination to optimize the unused calls out.

class SimpleModule(nn.Module):
    features: int
    outer_dense: nn.Dense
    
    def setup(self):
        self.dense1 = nn.Dense(self.outer_dense.features)  # Stub
        self.dense2 = nn.Dense(self.features)

    def __call__(self, inputs):
        _ = self.dense1(inputs)  # Just to make it appear in the variable dict.
        x = self.outer_dense(inputs)
        x = self.dense2(x)
        return x

class WrapperModule(nn.Module):
    features: int

    def setup(self):
        self.dense1 = nn.Dense(self.features)
        self.module = SimpleModule(features=self.features, outer_dense=self.dense1)

    def __call__(self, inputs):
        x = self.module(inputs)
        return x

[sanchit-gandhi]

Thank you for the detailed reply!

So in your case the Torch model is given (and fixed), and you are trying to implement a Flax model that will fit the torch state_dict?

Exactly that - we have PyTorch weights associated with a large PyTorch model that we are trying to port into Flax. In order to do so, we need to define a Flax model that is 'compatible' with the torch state_dict so to speak.

The hacky workaround with stubs works! Thanks for highlighting this workaround! But I'm wary that adding the dummy layers confuses and convolutes the code base somewhat. Modifying the PyTorch-to-Flax conversion script to handle this case is probably the cleanest and hence preferred option here. I'll take a look into that and let you know how I go about solving the issue!

[marcvanzee]

Definitely, it is a very ugly hack 😄. You could perhaps add a flag ignore_missing that will not throw an error if a variable is missing on the Flax side, or something like missing_filter that contains specific paths to ignore on the Flax side.

Yes please let me know! Curious to hear how this pans out.