Fix resuming from checkpoint when using RayFSDPStrategy

[angerrp]

Why are these changes needed?

Restoring from a checkpoint when using FSDP is currently flawed as the state_dict keys for each layer get modified and torch can not associate the weights to the layer name when loading. The current implementation always assumes that the layer keys in the state dict are prefixed with _forward_module. and then slices the key based on the length of the prefix.

The underlying reason why we remove the _forward_module. is unclear to me but we should check if it is prefixed before removing. This is implemented in this PR and fixes the loading of checkpoints when using RayFSDPStrategy

The following is an example of the wrong state_dict keys for a checkpoint:

# Keys in checkpoint["state_dict"]
der.embed_tokens.weight
der.embed_positions.weight
der.final_layer_norm.weight
...

Correct keys:

model.model.decoder.embed_tokens.weight
model.model.decoder.embed_positions.weight
model.model.decoder.final_layer_norm.weight
...

[woshiyyya]

Hi @dabauxi , The reason why we did this is because Lightning previously have a _LightningModuleWrapperBase, which added the extra "_forward_module." prefix to the state_dict keys. Users need to manually trim it to correctly load the checkpoint. Lightning-AI/pytorch-lightning#16526

In the recent versions, lightning trimmed the prefix internally so no need to do it ourselves.

Thanks for the fix!

[woshiyyya]

LGTM

[galyna-anyscale]

@matthewdeng Please review this PR so it can be merged.

[matthewdeng]

LGTM, seems like this branching logic is needed to account for different versions of Lightning which do/don't have the prefix.