[BUG] Inaccurate FLOPs for shared weights #7256

import torch
import torch.nn as nn

from deepspeed.profiling.flops_profiler import get_model_profile
from torch.profiler import profile, ProfilerActivity


class LinearBatchNorm(torch.nn.Module):
    def __init__(self, in_c, out_c, **kwargs):
        super().__init__(**kwargs)
        self.linear = nn.Linear(in_c, out_c)
        self.bn = nn.BatchNorm1d(out_c)

    def forward(self, x):
        x = self.linear(x)
        x = self.bn(x)
        return x


class SharedLinearLayers(torch.nn.Module):
    def __init__(self, shared: bool, **kwargs):
        super().__init__(**kwargs)

        self.linears = nn.ModuleList()

        for _ in range(3):
            self.linears.append(LinearBatchNorm(100, 100))

        if shared:
            for i in range(3):
                self.linears[i].linear = self.linears[0].linear

    def forward(self, x):
        x1, x2, x3 = torch.split(x, [100, 100, 100], dim=1)
        x = (x1, x2, x3)

        results = []
        for idx, input_feat in enumerate(x):
            output_feat = self.linears[idx](input_feat)
            results.append(output_feat)

        return tuple(results)


if __name__ == "__main__":
    input_shape = (1, 300)

    model = SharedLinearLayers(shared=True)
    if torch.cuda.is_available():
        model = model.to("cuda")
    model.eval()

    # for logging
    div = 1e3
    unit = "KFLOPs"
    print_str = "\n"

    # deepspeed
    flops, _, _ = get_model_profile(
        model=model,  # model
        input_shape=input_shape,  # input shape to the model. If specified, the model takes a tensor with this shape as the only positional argument.
        args=None,  # list of positional arguments to the model.
        kwargs=None,  # dictionary of keyword arguments to the model.
        print_profile=False,  # prints the model graph with the measured profile attached to each module
        detailed=True,  # print the detailed profile
        module_depth=-1,  # depth into the nested modules, with -1 being the inner most modules
        top_modules=1,  # the number of top modules to print aggregated profile
        warm_up=10,  # the number of warm-ups before measuring the time of each module
        as_string=False,  # print raw numbers (e.g. 1000) or as human-readable strings (e.g. 1k)
        output_file=None,  # path to the output file. If None, the profiler prints to stdout.
        ignore_modules=None,
    )  # the list of modules to ignore in the profiling
    print_str += f"deepspeed: {int(flops / div)} {unit}\n"

    # pytorch profiler
    with profile(
        activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA], with_flops=True
    ) as prof:
        model(torch.zeros(input_shape).to("cuda"))
    flops = sum(event.flops for event in prof.key_averages())
    print_str += f"torch profiler: {int(flops / div)} {unit}\n"

    print(print_str)

[2025-04-28 07:35:32,463] [INFO] [real_accelerator.py:239:get_accelerator] Setting ds_accelerator to cuda (auto detect)
--------------------------------------------------
DeepSpeed C++/CUDA extension op report
--------------------------------------------------
NOTE: Ops not installed will be just-in-time (JIT) compiled at
      runtime if needed. Op compatibility means that your system
      meet the required dependencies to JIT install the op.
--------------------------------------------------
JIT compiled ops requires ninja
ninja .................. [OKAY]
--------------------------------------------------
op name ................ installed .. compatible
--------------------------------------------------
 [WARNING]  async_io requires the dev libaio .so object and headers but these were not found.
 [WARNING]  async_io: please install the libaio-dev package with apt
 [WARNING]  If libaio is already installed (perhaps from source), try setting the CFLAGS and LDFLAGS environment variables to where it can be found.
async_io ............... [NO] ....... [NO]
fused_adam ............. [NO] ....... [OKAY]
cpu_adam ............... [NO] ....... [OKAY]
cpu_adagrad ............ [NO] ....... [OKAY]
cpu_lion ............... [NO] ....... [OKAY]
dc ..................... [NO] ....... [OKAY]
 [WARNING]  Please specify the CUTLASS repo directory as environment variable $CUTLASS_PATH
evoformer_attn ......... [NO] ....... [NO]
 [WARNING]  FP Quantizer is using an untested triton version (3.1.0), only 2.3.(0, 1) and 3.0.0 are known to be compatible with these kernels
fp_quantizer ........... [NO] ....... [NO]
fused_lamb ............. [NO] ....... [OKAY]
fused_lion ............. [NO] ....... [OKAY]
 [WARNING]  gds requires the dev libaio .so object and headers but these were not found.
 [WARNING]  gds: please install the libaio-dev package with apt
 [WARNING]  If libaio is already installed (perhaps from source), try setting the CFLAGS and LDFLAGS environment variables to where it can be found.
gds .................... [NO] ....... [NO]
transformer_inference .. [NO] ....... [OKAY]
inference_core_ops ..... [NO] ....... [OKAY]
cutlass_ops ............ [NO] ....... [OKAY]
quantizer .............. [NO] ....... [OKAY]
ragged_device_ops ...... [NO] ....... [OKAY]
ragged_ops ............. [NO] ....... [OKAY]
random_ltd ............. [NO] ....... [OKAY]
 [WARNING]  sparse_attn requires a torch version >= 1.5 and < 2.0 but detected 2.5
 [WARNING]  using untested triton version (3.1.0), only 1.0.0 is known to be compatible
sparse_attn ............ [NO] ....... [NO]
spatial_inference ...... [NO] ....... [OKAY]
transformer ............ [NO] ....... [OKAY]
stochastic_transformer . [NO] ....... [OKAY]
--------------------------------------------------
DeepSpeed general environment info:
torch install path ............... ['/opt/conda/lib/python3.11/site-packages/torch']
torch version .................... 2.5.1+cu121
deepspeed install path ........... ['/opt/conda/lib/python3.11/site-packages/deepspeed']
deepspeed info ................... 0.16.7, unknown, unknown
torch cuda version ............... 12.1
torch hip version ................ None
nvcc version ..................... 12.1
deepspeed wheel compiled w. ...... torch 2.5, cuda 12.1
shared memory (/dev/shm) size .... 188.28 GB