-
Notifications
You must be signed in to change notification settings - Fork 3.6k
Closed
Labels
Description
Bug description
Currently in the master branch, if we use RichModelSummary
with DeepSpeedStrategy
, the resulting summary table is abit off with missing Params per Device
information, like the following.
┏━━━┳━━━━━━━┳━━━━━━━━┳━━━━━━━━┳━━━━━━┳━━━━━━━┳━━━━━━━━━━┳━━━━━━━━━━━┓
┃ ┃ Name ┃ Type ┃ Params ┃ Mode ┃ FLOPs ┃ In sizes ┃ Out sizes ┃
┡━━━╇━━━━━━━╇━━━━━━━━╇━━━━━━━━╇━━━━━━╇━━━━━━━╇━━━━━━━━━━╇━━━━━━━━━━━┩
│ 0 │ layer │ Linear │ 66 │ 66 │ train │ [4, 32] │ [4, 2] │
└───┴───────┴────────┴────────┴──────┴───────┴──────────┴───────────┘
Trainable params: 66
Non-trainable params: 0
Total params: 66
Total estimated model params size (MB): 0
Modules in train mode: 1
Modules in eval mode: 0
Total FLOPs: 512
while the correct table should be like
┏━━━┳━━━━━━━┳━━━━━━━━┳━━━━━━━━┳━━━━━━━━━━┳━━━━━━━┳━━━━━━━┳━━━━━━━━━━┳━━━━━━━━━━┓
┃ ┃ ┃ ┃ ┃ Params ┃ ┃ ┃ ┃ ┃
┃ ┃ ┃ ┃ ┃ per ┃ ┃ ┃ ┃ Out ┃
┃ ┃ Name ┃ Type ┃ Params ┃ Device ┃ Mode ┃ FLOPs ┃ In sizes ┃ sizes ┃
┡━━━╇━━━━━━━╇━━━━━━━━╇━━━━━━━━╇━━━━━━━━━━╇━━━━━━━╇━━━━━━━╇━━━━━━━━━━╇━━━━━━━━━━┩
│ 0 │ layer │ Linear │ 66 │ 66 │ train │ 512 │ [4, 32] │ [4, 2] │
└───┴───────┴────────┴────────┴──────────┴───────┴───────┴──────────┴──────────┘
Trainable params: 66
Non-trainable params: 0
Total params: 66
Total estimated model params size (MB): 0
Modules in train mode: 1
Modules in eval mode: 0
Total FLOPs: 512
What version are you seeing the problem on?
master
Reproduced in studio
No response
How to reproduce the bug
Error messages and logs
# Error messages and logs here please
Environment
Current environment
#- PyTorch Lightning Version (e.g., 2.5.0):
#- PyTorch Version (e.g., 2.5):
#- Python version (e.g., 3.12):
#- OS (e.g., Linux):
#- CUDA/cuDNN version:
#- GPU models and configuration:
#- How you installed Lightning(`conda`, `pip`, source):
More info
No response
cc @lantiga