Fix flops profiler counts for transposed convolutions

[vineethsaivs]

Symptom

The flops profiler reports zero macs for a depthwise transposed convolution, and under-counts every other grouped one. Bias flops are wrong for every strided transposed convolution.

import torch
from deepspeed.profiling.flops_profiler import FlopsProfiler

for m in [torch.nn.ConvTranspose2d(4, 6, 4, stride=2, padding=1),
          torch.nn.ConvTranspose2d(4, 4, 3, stride=2, padding=1, output_padding=1, groups=4),
          torch.nn.ConvTranspose2d(6, 6, 2, stride=2, groups=3)]:
    x = torch.randn(2, m.in_channels, 8, 8)
    p = FlopsProfiler(m); p.start_profile(); y = m(x); p.stop_profile()
    print(m.groups, p.get_total_macs(), p.get_total_flops()); p.end_profile()

layer	macs reported	macs	flops reported	flops
ConvTranspose2d(4, 6, 4, stride=2, padding=1)	49152	49152	98496	101376
ConvTranspose2d(4, 4, 3, stride=2, padding=1, output_padding=1, groups=4)	0	4608	32	11264
ConvTranspose2d(6, 6, 2, stride=2, groups=3)	0	6144	64	15360

Root cause

Two independent problems in _conv_trans_flops_compute.

1. The per-group filter count is divided by groups twice. A transposed-convolution weight is [in_channels, out_channels // groups, *kernel], not [out_channels, in_channels // groups, *kernel] as for a forward convolution. So weight.shape[1] is already out_channels // groups, and

out_channels = weight.shape[1]
...
filters_per_channel = out_channels // groups

yields out_channels // groups ** 2. For a depthwise layer that is 1 // groups, which floors to 0, so conv_per_position_macs and the whole layer collapse to zero.

2. The output shape uses the forward-convolution formula. A transposed convolution inverts it, and output_padding was accepted by the signature but never read:

output_dim = (input_dim + 2 * paddings[idx] - (dilations[idx] * (kernel_dims[idx] - 1) + 1)) // strides[idx] + 1

The bias is added once per output element, so bias_flops = out_channels * batch_size * prod(output_dims) inherits the error. For a stride-2 layer the formula returns roughly input_dim / stride where the real output is input_dim * stride, so the bias term comes out about 16x too small in 2D.

The three lines that recomputed paddings/strides/dilations as 2-tuples immediately after the loop were dead, and hardcoded 2D for a helper that also serves conv_transpose1d and conv_transpose3d, so they go with the fix.

Fix

Read filters_per_channel straight off weight.shape[1] and recover out_channels by multiplying it back up, and compute the output shape with the transposed-convolution formula, including output_padding.

The mac convention is unchanged: a transposed convolution scatters every input element across the kernel, so active_elements_count stays based on input_dims.

Test

tests/unit/profiling/flops_profiler/test_flops_profiler.py::test_conv_transpose_flops, parametrized over groups in 1/2/4 and (stride, padding, output_padding) in (1, 0, 0)/(2, 1, 1). It checks both counts against the shapes torch itself produced, so the expected values are not a copy of the implementation.

Against master: 6 failed, 2 passed (assert 0 == 9216 for the depthwise cases). With the fix: 8 passed.

$ pytest unit/profiling/flops_profiler/test_flops_profiler.py -m sequential
8 passed, 2 deselected
$ pytest unit/profiling/flops_profiler/test_flops_profiler.py
2 passed, 8 deselected

Manually checked beyond what the test parametrizes, all matching torch: ConvTranspose1d, ConvTranspose3d, dilation=2, and bias=False.

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 20cdc7b436

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

[chatgpt-codex-connector]

Add the required Signed-off-by trailer

For this non-merge commit, the commit message has no Signed-off-by: trailer; the workspace rule requires that trailer on every non-merge commit, so DCO/CI checks can reject the change even though the code is correct. Please recommit with --signoff using the configured name/email.

AGENTS.md reference: AGENTS.md:L8-L8

Useful? React with 👍 / 👎.