Fix module initialization for root module under Zero3

[Ben-Schneider-code]

What does this PR do?

Fixes #28808

Behaviour:
loaded_keys = {k.replace(f"{module_name}.", "") for k in state_dict_keys if k.startswith(f"{module_name}.")} in Trainer.py misses the case where the module_name is the empty string (the root module), as its looking for state_dict_keys that start with ".". This causes the root module (CLIPModel, AlignModel, etc.) to always have model.apply(model._initialize_weights) called on it. Which leads to initialization of an empty parameter tensor when running zero3 (that doesn't need to be initialized as it will be replaced later) in modeling_align.py. The reason this causes a crash in modeling_align.py and not other models like CLIP is because ALIGN uses xavier_uniform_ which asserts the tensor must be 2d whereas CLIP uses nn.init.normal_ which just returns another empty tensor.

Fix:
Use the whole state_dict as the root module's set of load_keys.

Please let me know if there is anything I should revise / this should be fixed in another way.😄
@ArthurZucker @muellerzr

Before submitting

• This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
• Did you read the contributor guideline,
  Pull Request section?
• Was this discussed/approved via a Github issue or the forum? Please add a link
  to it if that's the case.
• Did you make sure to update the documentation with your changes? Here are the
  documentation guidelines, and
  here are tips on formatting docstrings.
• Did you write any new necessary tests?