Fabric and multiple models

[pfeatherstone]

When calling fabric.setup() there isn't an option to pass in multiple models. My current hack is to put them into an nn.ModuleDict but then the returned _FabricModule object isn't subscriptable like nn.ModuleDict. Is there a workaround? Note, the equivalent function with Huggingface Accelerate allows this.

[francotheengineer]

Have you tried this guide?:
https://lightning.ai/docs/fabric/stable/advanced/multiple_setup.html#multiple-models-multiple-optimizers
You should be able to call fabric.setup multiple times, each with a new model, and optimiser if needed.

[pfeatherstone]

But then I need one optimiser per model...

[ringohoffman]

@pfeatherstone I think you should open an issue for better visibility. I have the exact same request.

[awaelchli]

Regarding the module dict, I opened an issue here #18427. I think it should be possible to support that on the _FabricModule. In the mean time, you can always access the original module via model.module and then use it like a dict/list.