Perf enhancement

Is your feature request related to a problem? Please describe.
Use llama3 70b pretraining/finetuning as an example:

Detail doc
Baseline is 34% MFU.

• flash attention optimization
  • avoid duplicative attention mask generation #615
    • ~6.6% rel or +2.2% MFU here
    • will cover the models with pp activated only
  • pos_token_id processing, importing, sanitation check, conditional flash attention call
    • anti compilation including dynamic shape, torch where, conditional function here
    • SOL perform gain ~4.6% rel or +1.6% MFU here at the costs of flexibility
    • a super hacky commit which need some design for trade off and performance
    • it would help most of the models.
• evenly distributed pp
  • 25.11 https://github.com/NVIDIA-NeMo/Automodel/pull/790

  • num_virtual_stages = math.ceil(num_layers / layers_per_stage) + 1 this above calculation might prevent us from achieving evenly distributing pp/vpp stages? For example, with an 80-layer model, PP=4, and VPP=5, the ideal total number of virtual stages would be 20 (4 layers per stage). We could then add embeddings/lm_head to the first and last stages. However, this formula makes it impossible to calculate a result of 20. The closest is layer_per_stage as 4 we got num_virtual_stages=21 and later num_virtual_stages=24 with round up to multiple of PP. This could negatively impact pipeline balance. Looking at the following module_names_per_stage, first 8 stages have 4 layers with the following 16 stages have 3 layers instead.

• LoRA optimization: 22.74% and target is 32.23% MFU
  • improved from 17.79% 27.8% rel. improvement since last release, 7% rel. from norm fusion, and the other 20% rel. from transformers bump up.
    target is 32.23% MFU still 30% perf gap, mainly relevant to exposed TP communication. which I hope to cover in next release.
  • asyncTP for remaining ~30% rel. perf gap
  • 26.02NormFusion expected ~5% rel. / actual ~7% rel. https://github.com/NVIDIA-NeMo/Automodel/pull/1120
  • 25.11 Better Sharding Parallelism estimated ~20% rel. / actual 40% rel. https://github.com/NVIDIA-NeMo/Automodel/pull/733
  • 25.11 Combined Projection estimated ~30% rel. / actual 40% rel. https://github.com/NVIDIA-NeMo/Automodel/pulls?q=is%3Amerged+is%3Apr+author%3AZhiyuLi-Nvidia+
  • 25.11 Benchmark with PP Pipeline Parallelism Support ~40% rel. https://github.com/NVIDIA-NeMo/Automodel/pull/790
• [bug] regression (with the same recipe) between
  • https://github.com/NVIDIA-NeMo/Automodel/commit/8f3ad3380b2eff4f541dedeb18004295d572f84e 34% MFU with ~22s step time exp log
  • https://github.com/NVIDIA-NeMo/Automodel/commit/2c236e081e49e97ddaff4202a8556a30a664a362 18% MFU with ~40s step time exp log
    • why tps per gpu is the same between 2 runs in the log.

Describe the solution you'd like
A clear and concise description of what you want to happen.

Describe alternatives you've considered
A clear and concise description of any alternative solutions or features you've considered.

Additional context
Add any other context or screenshots about the feature request here.