Fix MFU helpers for configs without KV head count

[DaoyuanLi2816]

What does this PR do?

compute_flops_per_token and adjusted_mfu currently raise an AttributeError for GPT-NeoX configs because they do not declare num_key_value_heads. GPT-NeoX uses standard multi-head attention, so this falls back to num_attention_heads when the optional KV-head field is absent.

The regression tests compare the implicit fallback with the same config after explicitly setting num_key_value_heads to num_attention_heads.

Validation

• python -m pytest -q tests/test_utils.py -k "ComputeFlopsPerToken or ComputeMfu or AdjustedMfu" — 11 passed, 222 deselected
• python -m ruff check trl/trainer/utils.py tests/test_utils.py
• python -m ruff format --check trl/trainer/utils.py tests/test_utils.py
• git diff --check

Before submitting

• This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
• Did you read the contributor guideline, Pull Request section?
• Was this discussed/approved via a GitHub issue? Please add a link to it if that's the case.
• Did you make sure to update the documentation with your changes?
• Did you write any new necessary tests?

AI writing disclosure

We welcome the use of AI tools to help with contributions. For transparency and to help us improve our review process, please indicate the level of AI involvement in this PR.

• No AI usage: the PR was written entirely by a human.
• AI-assisted: some parts were suggested or improved by AI, but the PR was written and reviewed by a human.
• AI-generated: the PR was mostly or fully generated by an AI tool.

Who can review?

Anyone in the community is free to review the PR once the tests have passed.

---

Note

Low Risk
Narrow fallback in FLOP/MFU estimation only; no training, auth, or data-path changes.

Overview
Stops compute_flops_per_token from crashing on configs that omit num_key_value_heads (e.g. GPT-NeoX MHA). It now falls back to num_attention_heads, which also unblocks adjusted_mfu.

Tests compare the implicit fallback against the same config with the field set explicitly.

^{Reviewed by Cursor Bugbot for commit 5833714. Bugbot is set up for automated code reviews on this repo. Configure here.}

[cursor]

Cursor Bugbot has reviewed your changes using default effort and found 1 potential issue.

^{❌ Bugbot Autofix is OFF. To automatically fix reported issues with cloud agents, enable autofix in the Cursor dashboard.}

Want higher recall? High effort reviews run extra passes and find more bugs. A team admin can switch effort levels in the Cursor dashboard.

^{Reviewed by Cursor Bugbot for commit 33fe81a. Configure here.}