Investigate Resident Memory Increase during Inference #18640
Comments
ZekunZh
added
bug
awaelchli
added
performance
help wanted
|
Thanks for investigating this and providing an example code @ZekunZh! I'm not sure we can remove |
|
I ran a couple of tests with your script, removing Lightning and only running with the raw PyTorch model: ...
for i in range(N_ITERATIONS):
torch.cuda.empty_cache()
gc.collect()
model = model.to("cuda:0")
with torch.inference_mode():
for batch in dataloader:
model(batch.to("cuda:0"))
# model.cpu()
current_memory = process.memory_info().rss
memory_usage = convert_bytes_to_megabytes(current_memory - initial_memory)
print(f"Iteration {i + 1}: Resident Memory used: {memory_usage:.3f} MB")
memory_usages.append(memory_usage)
...
(results produced with torch nightly 2.2.0.dev20230920+cu121) While the memory increase is definitely smaller, it is still a steady slope. I suppose on a production system with thousands of requests these few MB could add up. I'm definitely not familiar with memory management in Python and PyTorch, but there seems to be some hidden state somewhere that's not just in Lightning. Perhaps the impact is just amplified with Lightning and the root cause something else. |
|
Hey, I have the same problem but during training. I am currently using lightning in an active learning loop, in which I recreate the trainer in each loop. I am using lightning version import gc
import os
from copy import deepcopy
from typing import Any
import psutil
import torch
import torch.nn.functional as F
from lightning import LightningModule, Trainer
from matplotlib import pyplot as plt
from torch import nn, optim
from torch.utils.data import TensorDataset, DataLoader
from tqdm import tqdm
class MLP(LightningModule):
def __init__(
self,
n_hidden=10240 # Quite a large value to amplify the effect, my actual model has roughly the same size
):
super().__init__()
torch.manual_seed(0)
self.model = nn.Sequential(
nn.Linear(1024, n_hidden),
nn.BatchNorm1d(n_hidden),
nn.ReLU(),
nn.Linear(n_hidden, n_hidden),
nn.BatchNorm1d(n_hidden),
nn.ReLU(),
nn.Linear(n_hidden, n_hidden),
nn.BatchNorm1d(n_hidden),
nn.ReLU(),
nn.Linear(n_hidden, 2)
)
self._initial_state = deepcopy(self.state_dict())
def forward(self, x) -> Any:
return F.log_softmax(self.model(x), dim=1)
def training_step(self, batch):
x, y = batch
pred = self.forward(x)
loss = F.nll_loss(pred, y.view(-1))
return loss
def configure_optimizers(self):
return optim.Adam(self.parameters())
def reset(self):
self.load_state_dict(self._initial_state)
def convert_bytes_to_megabytes(memory_bytes):
return memory_bytes / 1024 ** 2
def run_pytorch(model, dataloader, iterations=10):
model.train()
process = psutil.Process(os.getpid())
initial_memory = process.memory_info().rss
memory_usages = []
for _ in range(iterations):
model.reset()
optimizer = optim.Adam(model.parameters())
for x, y in tqdm(dataloader):
x = x.to("cuda")
y = y.to("cuda")
pred = model(x)
loss = F.nll_loss(pred, y)
# loss = model.training_step((x, y))
optimizer.zero_grad()
loss.backward()
optimizer.step()
current_memory = process.memory_info().rss
diff = convert_bytes_to_megabytes(current_memory - initial_memory)
memory_usages.append(diff)
return memory_usages
def run_lightning(model, dataloader, iterations=10, fix=False):
process = psutil.Process(os.getpid())
initial_memory = process.memory_info().rss
memory_usages = []
for _ in range(iterations):
model.reset()
trainer = Trainer(max_epochs=1, enable_checkpointing=False, logger=False)
trainer.fit(model, train_dataloaders=dataloader)
if fix:
gc.collect()
current_memory = process.memory_info().rss
diff = convert_bytes_to_megabytes(current_memory - initial_memory)
memory_usages.append(diff)
return memory_usages
def main():
# Create some example data
torch.manual_seed(0)
X = torch.randn(size=(64, 1024))
Y = (torch.rand(size=(64,)) < 0.2).long()
dataset = TensorDataset(X, Y)
dataloader = DataLoader(
dataset,
batch_size=64,
shuffle=True,
num_workers=1
)
model = MLP()
# Remove model from cpu memory
model.to("cuda")
torch_usage = run_pytorch(model, dataloader)
lightning_usage = run_lightning(model, dataloader)
lightning_usage_fix = run_lightning(model, dataloader, fix=True)
plt.plot(torch_usage, label="PyTorch")
plt.plot(lightning_usage, label="Lightning")
plt.plot(lightning_usage_fix, label="Lightning - Fix")
plt.legend()
plt.xlabel("Iterations")
plt.ylabel("Memory usage")
plt.show()
main()
|
|
Thanks for collecting more data here. So then if Hypothetically, if we were to insert a |
|
Yes, it seems that the result is the same. |
|
@carmocca What are your thoughts on adding a |
|
I'm learning towards not adding it. Instantiating trainers like this in a loop is very unconventional and there is a cost to triggering If somebody can explain the cause of this, we would be better informed to create a fix: either by improving the reference counts or by adding this |
@carmocca Just to clarify. Above we've determined that the Trainer releases these objects. So their refcount is actually 0. It's just that the GC does not collect them from memory quick enough. By adding
I agree. In light of this I am also ok closing this issue. But for the same argumentation, I am also ok adding the |
Bug description
The memory consumption (RSS memory) continues to grow when
Traineris instantiated multiple times during the inference.In our production environment, currently we need to instantiate a
Trainerfor each request which contains 1 image. That's why we observed the OOM issue.We understand that it's might not be the best practice to use Lighting in production, any suggestions / comments are welcome ! 😃
The following curve can be reproduced with the provided python script, running 1000 iterations.
What version are you seeing the problem on?
v2.0
How to reproduce the bug
Error messages and logs
Environment
Current environment
More info
The temporary solution to fix this issue is to add
gc.collect()at the end of teardown method, while commentingself.lightning_module.cpu().Things that I've tried:
Only comment
self.lightning_module.cpu()-> not work 🛑Only comment
_optimizers_to_device(self.optimizers, torch.device("cpu"))-> not work 🛑Comment both
module to cpuandoptimiser to cpu-> not work 🛑Only add
gc.collect()-> partially work 🟡Comment
_optimizers_to_device(self.optimizers, torch.device("cpu"))+ addgc.collect()-> partially work 🟡Comment
self.lightning_module.cpu()+ addgc.collect()-> work better 🟢Comment
self.lightning_module.cpu()and_optimizers_to_device(self.optimizers, torch.device("cpu"))+ addgc.collect()-> Similar to the previous one 🟢cc @Borda
The text was updated successfully, but these errors were encountered: