Scalene memory profiling with DDP(Lightning AI): total memory does not increase with more processes

[belfaunder]

Hello, I have a question about using Scalene to profile memory in a distributed training setup with DDP. I observed behavior that I do not fully understand:

When I run my workload with 1 process, Scalene reports about the same total memory usage as when I run it with 10 processes (all on CPU). However, when I measure memory usage with psutil, I see a clear increase, roughly from 0.8 GB to 3.7 GB. Because of this, I am unsure how to interpret Scalene’s memory numbers in a multi-process / DDP setup: does it report memory utilization only for the main process? And if so, are there settings that allow to get an aggregated memory consumption over all processes?

How to reproduce:

Run: scalene run --memory main.py,

Environment

Python: 3.12.3

lightning 2.6.1
scalene: 2.1.4
torch: 2.10.0
torchvision: 0.25.0

OS: Ubuntu 24.04 running under WSL2

Minimal reproduction example

The full code of main.py module for reproducing is based on example from Lightning AI (https://lightning.ai/docs/pytorch/stable/notebooks/lightning_examples/mnist-hello-world.html):

import os
import pytorch_lightning as pl
import torch
from torch.nn import functional as F
from torch.utils.data import DataLoader, Subset
from torchvision import transforms
from torchvision.datasets import MNIST

class MNISTModel(pl.LightningModule):
    """A PyTorch Lightning module for classifying images in the MNIST dataset."""
    def __init__(self):
        """Initializes a new instance of the MNISTModel class."""
        super().__init__()
        self.l1 = torch.nn.Linear(28 * 28, 10)
        self.leaked_losses = []

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        flattened = x.view(x.size(0), -1)
        hidden = self.l1(flattened)
        activated = torch.relu(hidden)
        return activated

    def training_step(self, batch, batch_idx: int):
        x, y = batch
        loss = F.cross_entropy(self(x), y)
        big = loss.repeat(1,1000)
        self.leaked_losses.append(big)

        return loss

    def configure_optimizers(self):
        return torch.optim.Adam(self.parameters(), lr=0.02)
    
# Init the model
mnist_model = MNISTModel()

# Init DataLoader from MNIST Dataset
train_ds = MNIST(os.environ.get("PATH_DATASETS", "."), train=True, download=True, transform=transforms.ToTensor())
train_ds_small = Subset(train_ds, range(1000))

# Create a dataloader
train_loader = DataLoader(train_ds_small, batch_size=3)

# Initialize a trainer
trainer = pl.Trainer(
    accelerator="cpu",
    strategy="ddp",
    devices=10, # or 1.
    max_epochs=3,
)

# Train the model 
trainer.fit(mnist_model, train_loader)

Screenshots of the reports for 10 and 1 process runs

Additional information

Maybe this issue from Memray is relevant, as it has an issue with following Lightning AI processes because of fork + execv (bloomberg/memray#33).