Fabric - unwrap_model

[pfeatherstone]

I am using Fabric to have a bit more control on the training loop. One of the things it does is export to torchscript and onnx when checkpointing. It looks like i need to convert the FabricModule wrapper to a normal nn.Module. Is there is an unwrap_model() function that takes care of spitting back the original model on cpu, with original fp32 precision while only doing so on rank zero ?
Thank you

[awaelchli]

Hey @pfeatherstone
There is no unwrapping function, but you can always get the original module back by calling

wrapped_model.module

[pfeatherstone]

Is that the only thing required? HF Accelerate has accelerator.unwrap_model(). I assumed it did a bit more than simply return wrapped_model.module, like maybe some synchronization between all nodes.

[pfeatherstone]

Basically I want to do an ONNX export when checkpointing, which requires unwrapping the model back to CPU