Using empty_init results in 0 gradient

[RuABraun]

My code looks like

smallmodel = ..
with farbric.init_module(empty_init=True)
    largemodel = ..
combinedmodel = CombinedModel(smallmodel, largemodel)

combinedmodel = fabric.setup_module(combinedmodel)
optimizer = create_opt(combinedmodel, ..)
optimizer = fabric.setup_optimizers(optimizer)

This results in all parameters having 0 gradient. This changes when I remove the init_module line.

Guessing I'm using it wrong, should I be wrapping everything model related in it?

[matheus-prandini]

@RuABraun We encountered the same problem when running a model with only one GPU (SingleDeviceStrategy or auto), because its weights were initialized to zeros, which hampers learning... Simply set init_module = False to initialize the weights randomly. Using the FSDP strategy, both init_module = False and init_module = True configurations work well.