flop_counts missing for deepspeed in 2.5.5

[bnestor]

Bug description

There is no flop_counts in version 2.5.5

pytorch-lightning/src/lightning/pytorch/utilities/model_summary/model_summary.py

Line 148 in dd0a8f3

class ModelSummary:

It does exist in the main branch:

pytorch-lightning/src/lightning/pytorch/utilities/model_summary/model_summary.py

Line 314 in e088694

def flop_counts(self) -> dict[str, dict[Any, int]]:

commenting out this line:

pytorch-lightning/src/lightning/pytorch/utilities/model_summary/model_summary_deepspeed.py

Line 102 in dd0a8f3

("FLOPs", list(map(get_human_readable_count, (sum(x.values()) for x in self.flop_counts.values())))),

allows it to work

What version are you seeing the problem on?

v2.5

Reproduced in studio

No response

How to reproduce the bug

dataloader = DataLoader(my_dataset)
ptl_model = AnyModel()

trainer = Trainer(accelerator="auto", strategy="deepspeed_stage_3")

trainer.fit(ptl_model, dataloader)

Error messages and logs

Traceback (most recent call last):
  File "myfile.py", line 373, in main
    trainer.fit(ptl_model, training_dataloader, validation_dataloader)
  File "envs/torch/lib/python3.10/site-packages/pytorch_lightning/trainer/trainer.py", line 560, in fit
    call._call_and_handle_interrupt(
  File "envs/torch/lib/python3.10/site-packages/pytorch_lightning/trainer/call.py", line 48, in _call_and_handle_interrupt
    return trainer.strategy.launcher.launch(trainer_fn, *args, trainer=trainer, **kwargs)
  File "envs/torch/lib/python3.10/site-packages/pytorch_lightning/strategies/launchers/subprocess_script.py", line 105, in launch
    return function(*args, **kwargs)
  File "envs/torch/lib/python3.10/site-packages/pytorch_lightning/trainer/trainer.py", line 598, in _fit_impl
    self._run(model, ckpt_path=ckpt_path)
  File "envs/torch/lib/python3.10/site-packages/pytorch_lightning/trainer/trainer.py", line 991, in _run
    call._call_callback_hooks(self, "on_fit_start")
  File "envs/torch/lib/python3.10/site-packages/pytorch_lightning/trainer/call.py", line 228, in _call_callback_hooks
    fn(trainer, trainer.lightning_module, *args, **kwargs)
  File "envs/torch/lib/python3.10/site-packages/pytorch_lightning/callbacks/model_summary.py", line 65, in on_fit_start
    summary_data = model_summary._get_summary_data()
  File "envs/torch/lib/python3.10/site-packages/pytorch_lightning/utilities/model_summary/model_summary_deepspeed.py", line 102, in _get_summary_data
    ("FLOPs", list(map(get_human_readable_count, (sum(x.values()) for x in self.flop_counts.values())))),
AttributeError: 'DeepSpeedSummary' object has no attribute 'flop_counts'

Environment

Current environment

#- PyTorch Lightning Version (e.g., 2.5.0): 2.5.5
#- PyTorch Version (e.g., 2.5): 2.8.0
#- Python version (e.g., 3.12): 3.10
#- OS (e.g., Linux): linux
#- CUDA/cuDNN version: 12.6
#- GPU models and configuration: Nvidia rtx series
#- How you installed Lightning(`conda`, `pip`, source): pip install pytorch-lightning

More info

No response

cc @lantiga

[deependujha]

Hi @bnestor, thanks for reporting this. I was able to reproduce the error on Lightning v2.5.5.

Workaround

✅ As a temporary fix, installing Lightning from source (master branch) resolves the issue since the missing code is present there.

Root cause

• FLOP support in the model summary was introduced in PR #20868.
• FLOP information for the DeepSpeed summary was later added in PR #21100.
• However, the 2.5.2 release (June 20) and subsequent 2.5.5 do not include feat: add flops in ModelSummary columns. #20868, even though it was already merged into main by then. This leaves DeepSpeedSummary referencing flop_counts without the method defined, leading to the AttributeError.

Next step

It looks like the release branch did not pick up #20868 before the 2.5.x releases were cut. A patch release that includes #20868 should fix this for PyPI users.

compare changes: 2.5.1...2.5.2

cc: @Borda @SkafteNicki

[SkafteNicki]

Either we go for a patch release or maybe it is time for a minor release to make sure that newest release is in full sync with master?

[deependujha]

Yeah, a minor release makes more sense to get master fully in sync.

[GdoongMathew]

Because of the BC change introduced in #20868, a minor version update would make more sense.