Wrong number of trainable parameters printed with `strategy="deepspeed_stage_3"`

[pme0]

🐛 Bug

In the beginning of training a table with the number of trainable parameters is printed.
This is not the correct answers (66) when strategy="deepspeed_stage_3" (even when using only 1 gpu):

deepspeed_stage_3:

ddp/deepspeed_stage_1/deepspeed_stage_2:

I'm guessing this is something to do with the parameter shard-ing used in stage3...?

I'm curious to know which 2 parameters are counted, since there are no other parameters in the model.

To Reproduce

import os
import torch
from torch.utils.data import DataLoader, Dataset
from pytorch_lightning import LightningModule, Trainer

class RandomDataset(Dataset):
    def __init__(self, size, num_samples):
        self.len = num_samples
        self.data = torch.randn(num_samples, size)
    def __getitem__(self, index):
        return self.data[index]
    def __len__(self):
        return self.len
    
class BoringModel(LightningModule):
    def __init__(self):
        super().__init__()
        self.layer = torch.nn.Linear(32, 2)
    def forward(self, x):
        return self.layer(x)
    def training_step(self, batch, batch_idx):
        loss = self(batch).sum()
        return {"loss": loss}
    def validation_step(self, batch, batch_idx):
        loss = self(batch).sum()
    def configure_optimizers(self):
        return torch.optim.SGD(self.layer.parameters(), lr=0.1)
    
def run():
    train_data = DataLoader(RandomDataset(32, 64), batch_size=2)
    val_data = DataLoader(RandomDataset(32, 64), batch_size=2)
    model = BoringModel()
    params = sum(p.numel() for p in model.parameters() if p.requires_grad)
    print(f'TRAINABLE PARAMS: {params}')
    trainer = Trainer(
        default_root_dir=os.getcwd(),
        limit_train_batches=1,
        limit_val_batches=1,
        limit_test_batches=1,
        num_sanity_val_steps=0,
        max_epochs=1,
        enable_model_summary=True,
        strategy="deepspeed_stage_3",
        gpus=[0]
    )
    trainer.fit(model, train_dataloaders=train_data, val_dataloaders=val_data)

run()

Environment

• pytorch_lightning==1.5.10
• torch==1.10.2
• python==3.8.12
• deepspeed==0.5.10
• OS Linux
• CUDA/cuDNN 10.2

[rohitgr7]

tracking issue: #10291

would you like to contribute to support this?

[pme0]

Would it be possible to call model summary before the model is distributed? That is, on the pl.LightningModule instance MyModel that goes into Trainer().fit(model=MyModel).
This should work as the number of trainable params should be the same before and after distributing.

Issue [#10291] could benefit from this solution as well, as long as the model size displayed in model summary is on disk, not on memory (there will be memory overheads after distributing I suspect).

[pme0]

Would it be possible to call model summary before the model is distributed? That is, on the pl.LightningModule instance MyModel that goes into Trainer().fit(model=MyModel). This should work as the number of trainable params should be the same before and after distributing.

Issue [#10291] could benefit from this solution as well, as long as the model size displayed in model summary is on disk, not on memory (there will be memory overheads after distributing I suspect).

This is probably not enough when using configure_sharded_model() though

[stale]

This issue has been automatically marked as stale because it hasn't had any recent activity. This issue will be closed in 7 days if no further activity occurs. Thank you for your contributions, Pytorch Lightning Team!

[ar90n]

I met the same issue. And I found a workaround.
In DeepSpeed stage 3, All tensors that are handled by deepspeed are copied into its ds_tensor attribute and released.
So we can use the ds_tensor attribute instead of itself. My workaround modification codes are following.

diff --git a/pytorch_lightning/utilities/model_summary.py b/pytorch_lightning/utilities/model_summary.py
index f0419e4b9..6683f9267 100644
--- a/pytorch_lightning/utilities/model_summary.py
+++ b/pytorch_lightning/utilities/model_summary.py
@@ -32,6 +32,11 @@ warning_cache = WarningCache()
 PARAMETER_NUM_UNITS = [" ", "K", "M", "B", "T"]
 UNKNOWN_SIZE = "?"

+def _get_tensor(p) -> torch.Tensor:
+    if hasattr(p, "ds_tensor"):
+        return p.ds_tensor
+    return p
+

 class LayerSummary:
     """Summary class for a single layer in a :class:`~pytorch_lightning.core.lightning.LightningModule`. It
@@ -120,7 +125,9 @@ class LayerSummary:
     @property
     def num_parameters(self) -> int:
         """Returns the number of parameters in this module."""
-        return sum(np.prod(p.shape) if not _is_lazy_weight_tensor(p) else 0 for p in self._module.parameters())
+        return sum(
+            np.prod(_get_tensor(p).shape) if not _is_lazy_weight_tensor(p) else 0 for p in self._module.parameters()
+        )


 class ModelSummary:
@@ -225,12 +232,12 @@ class ModelSummary:

     @property
     def total_parameters(self) -> int:
-        return sum(p.numel() if not _is_lazy_weight_tensor(p) else 0 for p in self._model.parameters())
+        return sum(_get_tensor(p).numel() if not _is_lazy_weight_tensor(p) else 0 for p in self._model.parameters())

     @property
     def trainable_parameters(self) -> int:
         return sum(
-            p.numel() if not _is_lazy_weight_tensor(p) else 0 for p in self._model.parameters() if p.requires_grad
+            _get_tensor(p).numel() if not _is_lazy_weight_tensor(p) else 0 for p in self._model.parameters() if p.requires_grad
         )

     @property

And the result of the example codes that are above is the following.

  | Name  | Type   | Params
---------------------------------
0 | layer | Linear | 66    
---------------------------------
66        Trainable params
0         Non-trainable params
66        Total params
0.000     Total estimated model params size (MB)

It's no problem that the size of the Total estimated model params size is zero. Because it's smaller than 1 Kilobyte.
I think these codes are messy. Is there any good idea? I tried to look into the usage of configure_sharded_model(). But I couldn't get understood.

[paruul]

I am facing similar issue while using deepspeed_stage_1 . The number of params and trainable params are not matching for me. Here are some details. I am wondering if anyone can help. And, please let me know if any more information is required. Thanks!