Fix issues in fused_attn_bwd

[xrennvidia]

Description

• Fix dtypes in fused_attn_bwd for CP, and make them consistent with the similar non-CP fix in [PyTorch] Fix fused attention backward's FP8 dtypes #1566 . This helps solve some convergence issues with CP+P2P and CP+A2A, with FP8 attention.
• Clone fwd scale for the recomputation in fused_attn_bwd: for non-CP, clone S's scale only; for CP, clone QKV, O and S's scales
• Fix the amax update in CP+P2P

Type of change

• Documentation change (change only to the documentation, either a fix or a new content)
• Bug fix (non-breaking change which fixes an issue)
• New feature (non-breaking change which adds functionality)
• Breaking change (fix or feature that would cause existing functionality to not work as expected)
• Infra/Build change
• Code refactoring

Checklist:

• I have read and followed the contributing guidelines
• The functionality is complete
• I have commented my code, particularly in hard-to-understand areas
• I have made corresponding changes to the documentation
• My changes generate no new warnings
• I have added tests that prove my fix is effective or that my feature works
• New and existing unit tests pass locally with my changes

[xrennvidia]

/te-ci pytorch L1

[xrennvidia]

/te-ci pytorch L1

[cyanguwa]

Thanks for the offline discussion. I think the PR looks good. I tidied up the PR description a bit and added more details. Also, the redundant None issue is being addressed by #1589 as well, so I removed it from this PR's description (no biggie, I don't think there'll be conflicts or anything). Once the user confirms convergence and the CI clears, I think you can merge. Thanks!!