[training] fix flops computation

[jepio]

What does this PR do ?

This PR fixes two discrepancies found when analyzing GPT-OSS FLOPs numbers coming from MBridge:

• gated linear units were only being counted correctly when using SWIGLU activations. Other gated linear unit activation units were using an expansion factor of 2 rather than 3. Correct this to 3 for all GLU. For GPT-OSS 20B this is on the order of 10% of total FLOPS during a training step.
• sliding window attention was undercounting FLOPS because we were dividing by 2. Expand the formula for the two regimes so that we correctly compute those FLOPS. This has a minimal effect but worth including for correctness sake.

Changelog

• fix FLOPS accounting for gated linear units
• fix FLOPS accounting for sliding window attention with small windows

GitHub Actions CI

See the CI section in the Contributing doc for how to trigger the CI. A Nvidia developer will need to approve and trigger the CI for external contributors.

Before your PR is "Ready for review"

Pre checks:

• Make sure you read and followed Contributor guidelines
• 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?

If you haven't finished some of the above items you can still open "Draft" PR.

Additional Information

• Related to # (issue)

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[claude]

Review -- fix flops computation

The two fixes are directionally correct:

• GLU fix: Removing the activation_func == F.silu guard so all gated linear units (not just SwiGLU) get expansion factor 3.
• SWA context fix: Replacing the min(W, S) upper-bound approximation with the exact average causal SWA context, and removing the spurious /2 that was double-counting the causal mask.

1. Existing unit tests will break (blocking)
At least two tests in test_flop_utils.py compute expected values using the old formulas and will fail with this change:

• test_swa_exact_delta (line 983) -- hardcodes core_diff = Q * (S - W), which no longer matches.
• test_swa_no_effect_when_window_ge_seq (line 958) -- asserts flops_swa == flops_full, but the new W >= T branch uses (T+1)/2 while full_core uses T/2, producing a mismatch of Q per SWA layer.
  These tests need to be updated in this PR.

2. W >= T inconsistency (see inline comment)
When the window covers the full sequence, swa_core should equal full_core. The else branch currently uses (effective_seq_length + 1) / 2 (exact) while full_core uses core_attn_seq_factor / 2 (approximate). Suggest using core_attn_seq_factor / 2 in the else branch so the two paths agree -- this also correctly threads through the seqlen_squared_sum override.

3. Missing test for non-silu GLU (the motivating fix)
The GLU fix is the larger of the two corrections (about 10 percent of total FLOPs per the PR description), but there is no test exercising gated_linear_unit=True with a non-silu activation (e.g., F.gelu). The existing test_swiglu_scaling_factor only tests with the default F.silu. Please add a test for GLU with a non-silu activation confirming expansion factor 3.

Suggested test cases
No perf tests impacted.
Unit tests to add/update in tests/unit_tests/training/utils/test_flop_utils.py:

• Update test_swa_exact_delta with the new SWA context/core formulas
• Update test_swa_no_effect_when_window_ge_seq to match the aligned W >= T behavior
• Add a test for GLU with non-silu activation (gated_linear_unit=True, activation_func=F.gelu) confirming expansion factor 3

[jepio]

Suggested test cases No perf tests impacted. Unit tests to add/update in tests/unit_tests/training/utils/test_flop_utils.py:

• Update test_swa_exact_delta with the new SWA context/core formulas
• Update test_swa_no_effect_when_window_ge_seq to match the aligned W >= T behavior
• Add a test for GLU with non-silu activation (gated_linear_unit=True, activation_func=F.gelu) confirming expansion factor 3

Did all 3.

[jepio]

@claude do you want to review again

[yaoyu-33]

/ok to test 8729946