[training] fix: Update DeepSeek-V4 FLOPs calculation

[cuichenx]

Summary

• Port DeepSeek-V4 hybrid attention FLOPs accounting from [dev] [DeepSeek-v4] Part 3: MTP support with mHC and new mHC contract NVIDIA/Megatron-LM#4518.
• Account for DSv4 projection, sparse attention, compressor, and DSA indexer terms when experimental_attention_variant == "dsv4_hybrid".
• Add focused unit coverage for the exact DSv4 hybrid formula and CSA ratio length validation including MTP layers.

Testing

• uvx pre-commit run --all-files
• uv run --no-sync python -m py_compile src/megatron/bridge/training/utils/flop_utils.py tests/unit_tests/training/utils/test_flop_utils.py
• git diff --check
• Dependency-light DSv4 hybrid FLOPs smoke test by importing flop_utils.py with local stubs

Notes

• uv run pre-commit run --all-files was attempted, but this host cannot resolve nvidia-resiliency-ext==0.6.0 because the package only publishes manylinux_2_39 wheels while the host reports manylinux_2_31.
• Focused pytest was not run in this scratch environment because importing the unit-test conftest requires additional full project dependencies such as modelopt.

[claude]

Code Review - DSv4 Hybrid FLOPs

Overall the formula porting and test structure look solid. Two inline comments posted:

1. Unsupported compress ratio values silently ignored - csa_compress_ratios entries outside {0, 4, 128} pass the length check but contribute zero FLOPs, a silent under-count. A validation guard would prevent this.
2. Misleading error message - When mtp_num_layers == 0, the error text expected num_layers + mtp_num_layers is confusing. Suggest showing both values explicitly.

Missing test coverage

The two new tests cover the happy path (exact FLOPs) and one error path (ratio length mismatch). Consider adding:

• q_lora_rank=None error path: The code raises ValueError at line 468 when q_lora_rank is None, but no test covers this.
• csa_compress_ratios=None error path: The code raises ValueError at line 490 when csa_compress_ratios is None, but no test covers this.
• No ratio==4 layers: When all layers use ratio 0 or 128, the indexer/compressor code is skipped (n_layers_r4 == 0 branch at line 535-537). This edge case is untested.
• Existing MLA regression: The else branch (standard MLA, lines 541-587) was re-indented. A quick sanity test that existing MLA FLOPs are unchanged after this refactor would add confidence (may already be covered by pre-existing tests, worth confirming).

Suggested test cases

No perf tests impacted.

[claude]

Light Code Review

Findings

Bug: unsupported compress ratios silently ignored

If csa_compress_ratios contains a value outside {0, 4, 128}, that layer is counted in neither n_layers_r0, n_layers_r4, nor n_layers_r128. It contributes zero FLOPs for sparse attention, compressor, and indexer terms -- silently under-counting. A validation after the three sum(...) lines would catch this early. See inline comment.

Test coverage gaps

The two new tests cover the happy path (exact FLOPs) and one validation (compress-ratio length). Consider adding:

• q_lora_rank is None raises ValueError
• csa_compress_ratios is None raises ValueError
• A config where n_layers_r4 == 0 (no indexer needed) -- exercises the else branch
• Missing dsa_indexer_* attrs when ratio-4 layers exist
• An unsupported compress-ratio value (if the validation above is added)

Suggested test cases

No perf tests impacted.

[cuichenx]

Addressed Claude feedback in de9a3c9f3:

• Added unsupported csa_compress_ratios validation to avoid silent FLOPs under-counting.
• Clarified the CSA ratio length error to show expected count plus num_layers and mtp_num_layers.
• Added DSv4 tests for unsupported ratios, q_lora_rank=None, csa_compress_ratios=None, and no ratio-4 layers.
• Existing standard MLA regression coverage is already present in TestMLAFlops in the same test file.

Validation rerun:

• uvx pre-commit run --all-files
• uv run --no-sync python -m py_compile src/megatron/bridge/training/utils/flop_utils.py tests/unit_tests/training/utils/test_flop_utils.py
• git diff --check
• Dependency-light DSv4 FLOPs smoke covering exact, no-ratio-4, and error paths

Known local environment blockers:

• uv run pre-commit run --all-files still fails resolving nvidia-resiliency-ext==0.6.0 for this host platform.
• Focused pytest import gets through ModelOpt after adding requests, then fails on missing transformer_engine from repo-wide conftest/import initialization.

[cuichenx]

Follow-up for the later coverage note: added missing dsa_indexer_* validation coverage in 9a252df25. The parameterized test covers missing dsa_indexer_n_heads, dsa_indexer_head_dim, and dsa_indexer_topk when ratio-4 CSA layers are present.

Validation rerun after this commit:

• uvx pre-commit run --all-files
• uv run --no-sync python -m py_compile src/megatron/bridge/training/utils/flop_utils.py tests/unit_tests/training/utils/test_flop_utils.py
• git diff --check
• Dependency-light DSv4 FLOPs smoke covering exact, no-ratio-4, unsupported ratio, missing q_lora_rank, missing csa_compress_ratios, and missing dsa_indexer_* paths