Broadcast elementwise flops from the trailing dimension

[vineethsaivs]

Symptom

The flops profiler counts an elementwise mul/add over the wrong shape whenever the two operands differ in rank, which is exactly the common case of scaling or shifting an activation by a per-feature vector.

import torch
from deepspeed.profiling.flops_profiler import FlopsProfiler

class Elementwise(torch.nn.Module):
    def forward(self, lhs, rhs):
        return torch.mul(lhs, rhs)

for lhs_shape, rhs_shape in [((2, 3, 4), (4,)), ((4,), (2, 3, 4)), ((8,), (2, 8)), ((3, 4), (4,))]:
    model = Elementwise()
    lhs, rhs = torch.randn(*lhs_shape), torch.randn(*rhs_shape)
    prof = FlopsProfiler(model)
    prof.start_profile(); out = model(lhs, rhs); prof.stop_profile()
    print(lhs_shape, rhs_shape, prof.get_total_flops(), out.numel())
    prof.end_profile()

operands	elements touched	flops reported
[2, 3, 4] * [4]	24	48
[4] * [2, 3, 4]	24	48
[8] * [2, 8]	16	64
[3, 4] * [4]	12	16
[2, 3, 4] * [2, 3, 4]	24	24
[4, 1, 7] * [1, 5, 7]	140	140

Root cause

_elementwise_flops_compute builds the result shape by walking both operands from index 0 and padding the shorter one on the right:

for i in range(max_dim):
    in_i = input.shape[i] if i < dim_input else 1
    ot_i = other.shape[i] if i < dim_other else 1

Broadcasting lines shapes up from the trailing dimension and pads on the left. For [2, 3, 4] against [4] the loop pairs 2 with 4, then 3 with the pad, then 4 with the pad, producing [4, 3, 4] rather than [2, 3, 4].

Equal ranks are handled correctly, which is why this is easy to miss: the shape only comes out wrong when the ranks differ, and it can land either high or low depending on which dimensions get paired up. It happens to be right by coincidence when the mispaired dimensions have the same product, e.g. [16] * [32, 16] gives [16, 32].

The counter backs torch.mul, torch.Tensor.mul, torch.add and torch.Tensor.add.

Fix

torch.broadcast_shapes returns the result shape directly, so the hand-rolled loop goes. The two non-tensor branches above it are unchanged, and a 0-dim tensor operand still resolves to the other operand's shape.

Test

tests/unit/profiling/flops_profiler/test_flops_profiler.py::test_elementwise_broadcast_flops, parametrized over four mismatched-rank pairs and two matched-rank pairs (one of them exercising a size-1 dimension). The expected value is result.numel() from the tensor torch actually produced, so it is not a restatement of the implementation.

Against master: 4 failed, 4 passed (assert 48 == 24, assert 64 == 16, assert 16 == 12). 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

Also checked by hand and matching torch: a python scalar operand, Tensor.mul, torch.add, a 0-dim tensor operand, and a scalar in the first position.

Independent of #8323, which fixes _conv_trans_flops_compute in the same file. The two touch different functions and do not overlap.

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: cc816116c6

ℹ️ 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

This non-merge commit's message does not include a Signed-off-by: trailer, while the repository requires every non-merge commit to be signed off. DCO/CI checks can reject the change until the commit is recreated or amended with --signoff.

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

Useful? React with 👍 / 👎.