add stdit model
State: open | Pull request: yes

# What does this PR do ?

Add STDiT(Spatial-Temporal DiT) model for nemo training. Support megatron parallel pattern in training pipeline.

**Collection**: [Note which collection this PR will affect]

# Changelog 
- Add stdit model
- Add stdit layer
- Add stdit embedding 
- Add stdit training factory function
- Support tp/sp/cp/dsp/pp/fsdp parallel pattern

# 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**:
- [x] 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**:
- [x] 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.models.stdit.spatial_attention
nemo/collections/diffusion/models/stdit/spatial_attention.py:115:4: C0116: Missing function or method docstring (missing-function-docstring)
************* Module nemo.collections.diffusion.models.stdit.stdit_embeddings
nemo/collections/diffusion/models/stdit/stdit_embeddings.py:163:4: C0116: Missing function or method docstring (missing-function-docstring)
nemo/collections/diffusion/models/stdit/stdit_embeddings.py:435:4: C0116: Missing function or method docstring (missing-function-docstring)
nemo/collections/diffusion/models/stdit/stdit_embeddings.py:450:0: C0116: Missing function or method docstring (missing-function-docstring)
nemo/collections/diffusion/models/stdit/stdit_embeddings.py:5:0: W0611: Unused numpy imported as np (unused-import)
************* Module nemo.collections.diffusion.models.stdit.stdit_layer_spec
nemo/collections/diffusion/models/stdit/stdit_layer_spec.py:55:0: C0301: Line too long (137/119) (line-too-long)
nemo/collections/diffusion/models/stdit/stdit_layer_spec.py:58:0: C0301: Line too long (129/119) (line-too-long)
nemo/collections/diffusion/models/stdit/stdit_layer_spec.py:73:0: C0301: Line too long (137/119) (line-too-long)
nemo/collections/diffusion/models/stdit/stdit_layer_spec.py:76:0: C0301: Line too long (129/119) (line-too-long)
nemo/collections/diffusion/models/stdit/stdit_layer_spec.py:90:0: C0301: Line too long (152/119) (line-too-long)
nemo/collections/diffusion/models/stdit/stdit_layer_spec.py:445:0: C0301: Line too long (129/119) (line-too-long)
nemo/collections/diffusion/models/stdit/stdit_layer_spec.py:104:4: C0116: Missing function or method docstring (missing-function-docstring)
nemo/collections/diffusion/models/stdit/stdit_layer_spec.py:132:4: C0116: Missing function or method docstring (missing-function-docstring)
nemo/collections/diffusion/models/stdit/stdit_layer_spec.py:141:4: C0116: Missing function or method docstring (missing-function-docstring)
nemo/collections/diffusion/models/stdit/stdit_layer_spec.py:145:4: C0116: Missing function or method docstring (missing-function-docstring)
nemo/collections/diffusion/models/stdit/stdit_layer_spec.py:149:4: C0116: Missing function or method docstring (missing-function-docstring)
nemo/collections/diffusion/models/stdit/stdit_layer_spec.py:153:4: C0116: Missing function or method docstring (missing-function-docstring)
nemo/collections/diffusion/models/stdit/stdit_layer_spec.py:157:4: C0116: Missing function or method docstring (missing-function-docstring)
nemo/collections/diffusion/models/stdit/stdit_layer_spec.py:165:4: C0116: Missing function or method docstring (missing-function-docstring)
nemo/collections/diffusion/models/stdit/stdit_layer_spec.py:171:0: C0115: Missing class docstring (missing-class-docstring)
nemo/collections/diffusion/models/stdit/stdit_layer_spec.py:283:4: C0116: Missing function or method docstring (missing-function-docstring)
nemo/collections/diffusion/models/stdit/stdit_layer_spec.py:3:0: W0611: Unused Literal imported from typing (unused-import)
************* Module nemo.collections.diffusion.models.stdit.stdit_model
nemo/collections/diffusion/models/stdit/stdit_model.py:41:0: C0301: Line too long (123/119) (line-too-long)
nemo/collections/diffusion/models/stdit/stdit_model.py:42:0: C0301: Line too long (125/119) (line-too-long)
nemo/collections/diffusion/models/stdit/stdit_model.py:2:0: W0611: Unused Dict imported from typing (unused-import)
nemo/collections/diffusion/models/stdit/stdit_model.py:6:0: W0611: Unused torch.nn imported as nn (unused-import)
nemo/collections/diffusion/models/stdit/stdit_model.py:8:0: W0611: Unused parallel_state imported from megatron.core (unused-import)
nemo/collections/diffusion/models/stdit/stdit_model.py:9:0: W0611: Unused ShardedStateDict imported from megatron.core.dist_checkpointing.mapping (unused-import)
nemo/collections/diffusion/models/stdit/stdit_model.py:17:0: W0611: Unused make_sharded_tensor_for_checkpoint imported from megatron.core.utils (unused-import)
************* Module nemo.collections.diffusion.data.diffusion_fake_datamodule
nemo/collections/diffusion/data/diffusion_fake_datamodule.py:147:0: C0115: Missing class docstring (missing-class-docstring)
nemo/collections/diffusion/data/diffusion_fake_datamodule.py:267:0: C0115: Missing class docstring (missing-class-docstring)

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

---

Mitigation guide:

* Add sensible and useful docstrings to functions and methods
* For trivial methods like getter/setters, consider adding `# pylint: disable=C0116` inside the function itself
* To disable multiple functions/methods at once, put a `# pylint: disable=C0116` before the first and a `# pylint: enable=C0116` after the last.

By applying these rules, we reduce the occurance of this message in future.

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.models.model
nemo/collections/diffusion/models/model.py:77:0: C0116: Missing function or method docstring (missing-function-docstring)
nemo/collections/diffusion/models/model.py:410:0: C0115: Missing class docstring (missing-class-docstring)
nemo/collections/diffusion/models/model.py:420:0: C0115: Missing class docstring (missing-class-docstring)
nemo/collections/diffusion/models/model.py:430:0: C0115: Missing class docstring (missing-class-docstring)
nemo/collections/diffusion/models/model.py:465:4: C0116: Missing function or method docstring (missing-function-docstring)
nemo/collections/diffusion/models/model.py:487:4: C0116: Missing function or method docstring (missing-function-docstring)
************* Module nemo.collections.diffusion.train
nemo/collections/diffusion/train.py:99:0: C0116: Missing function or method docstring (missing-function-docstring)
nemo/collections/diffusion/train.py:240:0: C0116: Missing function or method docstring (missing-function-docstring)
nemo/collections/diffusion/train.py:275:0: C0116: Missing function or method docstring (missing-function-docstring)
nemo/collections/diffusion/train.py:540:0: C0116: Missing function or method docstring (missing-function-docstring)
nemo/collections/diffusion/train.py:33:0: W0611: Unused STDiT3BConfig imported from nemo.collections.diffusion.models.model (unused-import)

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

---

Mitigation guide:

* Add sensible and useful docstrings to functions and methods
* For trivial methods like getter/setters, consider adding `# pylint: disable=C0116` inside the function itself
* To disable multiple functions/methods at once, put a `# pylint: disable=C0116` before the first and a `# pylint: enable=C0116` after the last.

By applying these rules, we reduce the occurance of this message in future.

Thank you for improving NeMo's documentation!