Handle MCore custom fsdp checkpoint load
State: closed | Pull request: yes

# What does this PR do ?

Add a one line overview of what this PR aims to accomplish.

**Collection**: [Note which collection this PR will affect]

# Changelog 
- Add specific line by line info of high level changes in this PR.

# Usage
* You can potentially add a usage example below

```python
# Add a code snippet demonstrating how to use this 
```

# GitHub Actions CI

The Jenkins CI system has been replaced by GitHub Actions self-hosted runners.

The GitHub Actions CI will run automatically when the "Run CICD" label is added to the PR.
To re-run CI remove and add the label again.
To run CI on an untrusted fork, a NeMo user with write access must first click "Approve and run".

# Before your PR is "Ready for review"
**Pre checks**:
- [ ] Make sure you read and followed [Contributor guidelines](https://github.com/NVIDIA/NeMo/blob/main/CONTRIBUTING.md)
- [ ] Did you write any new necessary tests?
- [ ] Did you add or update any necessary documentation?
- [ ] Does the PR affect components that are optional to install? (Ex: Numba, Pynini, Apex etc)
  - [ ] Reviewer: Does the PR have correct import guards for all optional libraries?
  
**PR Type**:
- [ ] New Feature
- [ ] Bugfix
- [ ] Documentation

If you haven't finished some of the above items you can still open "Draft" PR.


## Who can review?

Anyone in the community is free to review the PR once the checks have passed. 
[Contributor guidelines](https://github.com/NVIDIA/NeMo/blob/main/CONTRIBUTING.md) contains specific people who can review PRs to various areas.

# Additional Information
* Related to # (issue)


<!-- pylint-output-strict-mode-true -->

beep boop 🤖: 🚨 The following files must be fixed before merge!

---

Your code was analyzed with PyLint. The following annotations have been identified:

```
************* Module nemo.collections.diffusion.flux_controlnet_infer
nemo/collections/diffusion/flux_controlnet_infer.py:28:0: C0301: Line too long (188/119) (line-too-long)
nemo/collections/diffusion/flux_controlnet_infer.py:26:0: C0116: Missing function or method docstring (missing-function-docstring)

-----------------------------------
Your code has been rated at 9.59/10
```

---

Thank you for improving NeMo's documentation!

<!-- pylint-output-strict-mode-false -->

beep boop 🤖: 🙏 The following files have warnings. In case you are familiar with these, please try helping us to improve the code base.

---

Your code was analyzed with PyLint. The following annotations have been identified:

```
************* Module nemo.collections.diffusion.flux_controlnet_training
nemo/collections/diffusion/flux_controlnet_training.py:36:0: C0116: Missing function or method docstring (missing-function-docstring)
nemo/collections/diffusion/flux_controlnet_training.py:18:0: W0611: Unused torch.nn imported as nn (unused-import)
nemo/collections/diffusion/flux_controlnet_training.py:22:0: W0611: Unused AutoProcessor imported from transformers (unused-import)
nemo/collections/diffusion/flux_controlnet_training.py:30:0: W0611: Unused Utils imported from nemo.collections.diffusion.utils.mcore_parallel_utils (unused-import)
************* Module nemo.collections.diffusion.flux_training
nemo/collections/diffusion/flux_training.py:36:0: C0116: Missing function or method docstring (missing-function-docstring)
nemo/collections/diffusion/flux_training.py:18:0: W0611: Unused torch.nn imported as nn (unused-import)
nemo/collections/diffusion/flux_training.py:22:0: W0611: Unused AutoProcessor imported from transformers (unused-import)
nemo/collections/diffusion/flux_training.py:30:0: W0611: Unused Utils imported from nemo.collections.diffusion.utils.mcore_parallel_utils (unused-import)
************* Module nemo.collections.diffusion.models.flux.model
nemo/collections/diffusion/models/flux/model.py:50:0: C0116: Missing function or method docstring (missing-function-docstring)
nemo/collections/diffusion/models/flux/model.py:71:0: C0115: Missing class docstring (missing-class-docstring)
nemo/collections/diffusion/models/flux/model.py:95:4: C0116: Missing function or method docstring (missing-function-docstring)
nemo/collections/diffusion/models/flux/model.py:101:0: C0115: Missing class docstring (missing-class-docstring)
nemo/collections/diffusion/models/flux/model.py:110:0: C0115: Missing class docstring (missing-class-docstring)
nemo/collections/diffusion/models/flux/model.py:159:4: C0116: Missing function or method docstring (missing-function-docstring)
nemo/collections/diffusion/models/flux/model.py:224:0: C0115: Missing class docstring (missing-class-docstring)
nemo/collections/diffusion/models/flux/model.py:246:4: C0116: Missing function or method docstring (missing-function-docstring)
nemo/collections/diffusion/models/flux/model.py:261:4: C0116: Missing function or method docstring (missing-function-docstring)
nemo/collections/diffusion/models/flux/model.py:264:4: C0116: Missing function or method docstring (missing-function-docstring)
nemo/collections/diffusion/models/flux/model.py:280:4: C0116: Missing function or method docstring (missing-function-docstring)
nemo/collections/diffusion/models/flux/model.py:297:4: C0116: Missing function or method docstring (missing-function-docstring)
nemo/collections/diffusion/models/flux/model.py:300:4: C0116: Missing function or method docstring (missing-function-docstring)
nemo/collections/diffusion/models/flux/model.py:303:4: C0116: Missing function or method docstring (missing-function-docstring)
nemo/collections/diffusion/models/flux/model.py:307:4: C0116: Missing function or method docstring (missing-function-docstring)
nemo/collections/diffusion/models/flux/model.py:312:4: C0116: Missing function or method docstring (missing-function-docstring)
nemo/collections/diffusion/models/flux/model.py:362:4: C0116: Missing function or method docstring (missing-function-docstring)
nemo/collections/diffusion/models/flux/model.py:395:4: C0116: Missing function or method docstring (missing-function-docstring)
nemo/collections/diffusion/models/flux/model.py:486:4: C0116: Missing function or method docstring (missing-function-docstring)
nemo/collections/diffusion/models/flux/model.py:490:4: C0116: Missing function or method docstring (missing-function-docstring)
nemo/collections/diffusion/models/flux/model.py:497:4: C0116: Missing function or method docstring (missing-function-docstring)
nemo/collections/diffusion/models/flux/model.py:42:0: W0611: Unused megatron_parallel imported from nemo.lightning as mp (unused-import)
************* Module nemo.collections.diffusion.models.flux.pipeline
nemo/collections/diffusion/models/flux/pipeline.py:65:0: C0301: Line too long (160/119) (line-too-long)
nemo/collections/diffusion/models/flux/pipeline.py:373:0: C0301: Line too long (165/119) (line-too-long)
nemo/collections/diffusion/models/flux/pipeline.py:382:0: C0301: Line too long (171/119) (line-too-long)
nemo/collections/diffusion/models/flux/pipeline.py:36:0: C0115: Missing class docstring (missing-class-docstring)
nemo/collections/diffusion/models/flux/pipeline.py:51:4: C0116: Missing function or method docstring (missing-function-docstring)
nemo/collections/diffusion/models/flux/pipeline.py:69:4: C0116: Missing function or method docstring (missing-function-docstring)
nemo/collections/diffusion/models/flux/pipeline.py:157:4: C0116: Missing function or method docstring (missing-function-docstring)
nemo/collections/diffusion/models/flux/pipeline.py:222:4: C0116: Missing function or method docstring (missing-function-docstring)
nemo/collections/diffusion/models/flux/pipeline.py:227:4: C0116: Missing function or method docstring (missing-function-docstring)
nemo/collections/diffusion/models/flux/pipeline.py:346:0: C0115: Missing class docstring (missing-class-docstring)
nemo/collections/diffusion/models/flux/pipeline.py:386:4: C0116: Missing function or method docstring (missing-function-docstring)
nemo/collections/diffusion/models/flux/pipeline.py:394:4: C0116: Missing function or method docstring (missing-function-docstring)
nemo/collections/diffusion/models/flux/pipeline.py:401:4: C0116: Missing function or method docstring (missing-function-docstring)
************* Module layers
nemo/collections/diffusion/models/flux_controlnet/layers.py:53:4: C0116: Missing function or method docstring (missing-function-docstring)
nemo/collections/diffusion/models/flux_controlnet/layers.py:17:0: W0611: Unused import torch (unused-import)
************* Module model
nemo/collections/diffusion/models/flux_controlnet/model.py:244:0: C0301: Line too long (120/119) (line-too-long)
nemo/collections/diffusion/models/flux_controlnet/model.py:24:0: C0116: Missing function or method docstring (missing-function-docstring)
nemo/collections/diffusion/models/flux_controlnet/model.py:30:0: C0116: Missing function or method docstring (missing-function-docstring)
nemo/collections/diffusion/models/flux_controlnet/model.py:42:0: C0115: Missing class docstring (missing-class-docstring)
nemo/collections/diffusion/models/flux_controlnet/model.py:69:0: C0115: Missing class docstring (missing-class-docstring)
nemo/collections/diffusion/models/flux_controlnet/model.py:128:4: C0116: Missing function or method docstring (missing-function-docstring)
nemo/collections/diffusion/models/flux_controlnet/model.py:137:4: C0116: Missing function or method docstring (missing-function-docstring)
nemo/collections/diffusion/models/flux_controlnet/model.py:217:0: C0115: Missing class docstring (missing-class-docstring)
nemo/collections/diffusion/models/flux_controlnet/model.py:230:0: C0115: Missing class docstring (missing-class-docstring)
nemo/collections/diffusion/models/flux_controlnet/model.py:2:0: W0611: Unused List imported from typing (unused-import)
nemo/collections/diffusion/models/flux_controlnet/model.py:4:0: W0611: Unused numpy imported as np (unused-import)
nemo/collections/diffusion/models/flux_controlnet/model.py:11:0: W0611: Unused AdaLNContinuous imported from nemo.collections.diffusion.models.dit.dit_layer_spec (unused-import)
************* Module pipeline
nemo/collections/diffusion/models/flux_controlnet/pipeline.py:52:0: C0301: Line too long (159/119) (line-too-long)
nemo/collections/diffusion/models/flux_controlnet/pipeline.py:17:0: C0115: Missing class docstring (missing-class-docstring)
nemo/collections/diffusion/models/flux_controlnet/pipeline.py:37:4: C0116: Missing function or method docstring (missing-function-docstring)
nemo/collections/diffusion/models/flux_controlnet/pipeline.py:56:4: C0116: Missing function or method docstring (missing-function-docstring)
nemo/collections/diffusion/models/flux_controlnet/pipeline.py:83:4: C0116: Missing function or method docstring (missing-function-docstring)
nemo/collections/diffusion/models/flux_controlnet/pipeline.py:158:4: C0116: Missing function or method docstring (missing-function-docstring)
nemo/collections/diffusion/models/flux_controlnet/pipeline.py:204:4: C0116: Missing function or method docstring (missing-function-docstring)
nemo/collections/diffusion/models/flux_controlnet/pipeline.py:209:4: C0116: Missing function or method docstring (missing-function-docstring)
nemo/collections/diffusion/models/flux_controlnet/pipeline.py:212:4: C0116: Missing function or method docstring (missing-function-docstring)
nemo/collections/diffusion/models/flux_controlnet/pipeline.py:1:0: W0611: Unused Any imported from typing (unused-import)
nemo/collections/diffusion/models/flux_controlnet/pipeline.py:1:0: W0611: Unused Callable imported from typing (unused-import)
nemo/collections/diffusion/models/flux_controlnet/pipeline.py:1:0: W0611: Unused Tuple imported from typing (unused-import)
************* Module nemo.collections.diffusion.utils.flux_pipeline_utils
nemo/collections/diffusion/utils/flux_pipeline_utils.py:15:0: W0611: Unused dataclass imported from dataclasses (unused-import)
nemo/collections/diffusion/utils/flux_pipeline_utils.py:16:0: W0611: Unused Callable imported from typing (unused-import)
nemo/collections/diffusion/utils/flux_pipeline_utils.py:18:0: W0611: Unused import torch (unused-import)
nemo/collections/diffusion/utils/flux_pipeline_utils.py:19:0: W0611: Unused TransformerConfig imported from megatron.core.transformer.transformer_config (unused-import)
nemo/collections/diffusion/utils/flux_pipeline_utils.py:24:0: W0611: Unused io imported from nemo.lightning (unused-import)
************* Module nemo.lightning._strategy_lib
nemo/lightning/_strategy_lib.py:574:0: C0301: Line too long (130/119) (line-too-long)
nemo/lightning/_strategy_lib.py:35:0: C0115: Missing class docstring (missing-class-docstring)
nemo/lightning/_strategy_lib.py:36:4: C0116: Missing function or method docstring (missing-function-docstring)
nemo/lightning/_strategy_lib.py:139:0: C0116: Missing function or method docstring (missing-function-docstring)
nemo/lightning/_strategy_lib.py:166:0: C0116: Missing function or method docstring (missing-function-docstring)
nemo/lightning/_strategy_lib.py:202:0: C0116: Missing function or method docstring (missing-function-docstring)
nemo/lightning/_strategy_lib.py:515:0: C0116: Missing function or method docstring (missing-function-docstring)
nemo/lightning/_strategy_lib.py:599:0: C0116: Missing function or method docstring (missing-function-docstring)
nemo/lightning/_strategy_lib.py:612:4: C0115: Missing class docstring (missing-class-docstring)
************* Module nemo.lightning.megatron_parallel
nemo/lightning/megatron_parallel.py:240:0: C0301: Line too long (127/119) (line-too-long)
nemo/lightning/megatron_parallel.py:241:0: C0301: Line too long (140/119) (line-too-long)
nemo/lightning/megatron_parallel.py:242:0: C0301: Line too long (130/119) (line-too-long)
nemo/lightning/megatron_parallel.py:551:0: C0301: Line too long (129/119) (line-too-long)
nemo/lightning/megatron_parallel.py:558:0: C0301: Line too long (135/119) (line-too-long)
nemo/lightning/megatron_parallel.py:827:0: C0301: Line too long (137/119) (line-too-long)
nemo/lightning/megatron_parallel.py:1057:0: C0301: Line too long (136/119) (line-too-long)
nemo/lightning/megatron_parallel.py:1624:0: C0301: Line too long (128/119) (line-too-long)
nemo/lightning/megatron_parallel.py:1663:0: C0301: Line too long (146/119) (line-too-long)
nemo/lightning/megatron_parallel.py:66:0: C0115: Missing class docstring (missing-class-docstring)
nemo/lightning/megatron_parallel.py:67:4: C0116: Missing function or method docstring (missing-function-docstring)
nemo/lightning/megatron_parallel.py:69:4: C0116: Missing function or method docstring (missing-function-docstring)
nemo/lightning/megatron_parallel.py:104:0: C0116: Missing function or method docstring (missing-function-docstring)
nemo/lightning/megatron_parallel.py:108:0: C0116: Missing function or method docstring (missing-function-docstring)
nemo/lightning/megatron_parallel.py:308:4: C0116: Missing function or method docstring (missing-function-docstring)
nemo/lightning/megatron_parallel.py:332:4: C0116: Missing function or method docstring (missing-function-docstring)
nemo/lightning/megatron_parallel.py:358:4: C0116: Missing function or method docstring (missing-function-docstring)
nemo/lightning/megatron_parallel.py:384:4: C0116: Missing function or method docstring (missing-function-docstring)
nemo/lightning/megatron_parallel.py:520:4: C0116: Missing function or method docstring (missing-function-docstring)
nemo/lightning/megatron_parallel.py:566:4: C0116: Missing function or method docstring (missing-function-docstring)
nemo/lightning/megatron_parallel.py:570:4: C0116: Missing function or method docstring (missing-function-docstring)
nemo/lightning/megatron_parallel.py:635:4: C0116: Missing function or method docstring (missing-function-docstring)
nemo/lightning/megatron_parallel.py:671:4: C0116: Missing function or method docstring (missing-function-docstring)
nemo/lightning/megatron_parallel.py:678:4: C0116: Missing function or method docstring (missing-function-docstring)
nemo/lightning/megatron_parallel.py:712:4: C0116: Missing function or method docstring (missing-function-docstring)
nemo/lightning/megatron_parallel.py:720:4: C0116: Missing function or method docstring (missing-function-docstring)
nemo/lightning/megatron_parallel.py:736:4: C0116: Missing function or method docstring (missing-function-docstring)
nemo/lightning/megatron_parallel.py:763:0: C0116: Missing function or method docstring (missing-function-docstring)
nemo/lightning/megatron_parallel.py:775:0: C0115: Missing class docstring (missing-class-docstring)
nemo/lightning/megatron_parallel.py:797:4: C0116: Missing function or method docstring (missing-function-docstring)
nemo/lightning/megatron_parallel.py:1317:4: C0116: Missing function or method docstring (missing-function-docstring)
nemo/lightning/megatron_parallel.py:1492:0: C0115: Missing class docstring (missing-class-docstring)
nemo/lightning/megatron_parallel.py:1498:4: C0116: Missing function or method docstring (missing-function-docstring)
nemo/lightning/megatron_parallel.py:1504:4: C0116: Missing function or method docstring (missing-function-docstring)
nemo/lightning/megatron_parallel.py:1508:4: C0116: Missing function or method docstring (missing-function-docstring)
nemo/lightning/megatron_parallel.py:1513:0: C0115: Missing class docstring (missing-class-docstring)
nemo/lightning/megatron_parallel.py:1518:0: C0115: Missing class docstring (missing-class-docstring)
nemo/lightning/megatron_parallel.py:1546:0: C0116: Missing function or method docstring (missing-function-docstring)
nemo/lightning/megatron_parallel.py:1592:8: C0116: Missing function or method docstring (missing-function-docstring)
nemo/lightning/megatron_parallel.py:1614:0: C0115: Missing class docstring (missing-class-docstring)
nemo/lightning/megatron_parallel.py:1687:0: C0115: Missing class docstring (missing-class-docstring)
nemo/lightning/megatron_parallel.py:1726:0: C0116: Missing function or method docstring (missing-function-docstring)
nemo/lightning/megatron_parallel.py:1740:0: C0116: Missing function or method docstring (missing-function-docstring)
nemo/lightning/megatron_parallel.py:51:0: W0611: Unused TrainerFn imported from pytorch_lightning.trainer.states (unused-import)

-----------------------------------
Your code has been rated at 9.36/10
```

---

Thank you for improving NeMo's documentation!