PyTorch Profiler leaks memory

[EricWiener]

🐛 Bug

When using profiler="PyTorch", memory usage (as measured by vm_percent) will keep increasing until running out of memory.

To Reproduce

No code yet, but will try to make an example. Just wanted to make this public info.

Expected behavior

The profiler doesn't leak memory.

Environment

• PyTorch Lightning Version (e.g., 1.5.0): 1.5.8
• PyTorch Version (e.g., 1.10): 1.10.1
• Python version (e.g., 3.9): 3.6
• OS (e.g., Linux): Linux
• CUDA/cuDNN version: 11.2
• GPU models and configuration: GeForce RTX 2080 Ti
• How you installed PyTorch (conda, pip, source): pip
• If compiling from source, the output of torch.__config__.show():
• Any other relevant information:

Additional context

N/A

cc @carmocca @kaushikb11 @ninginthecloud @rohitgr7

[tchaton]

Hey @EricWiener, do you have a reproducible script ?

[EricWiener]

Hi @tchaton I wasn't able to immediately reproduce this in colab, but will hopefully be able to at some point. For reference, here are charts of Memory usage for the same job where one had the profiler and one didn't (blacked out some irrelevant info):

[EricWiener]

Is there any chance any behavior of the profiler is tied to the end of an epoch (ex. clearing out the cache or summing up metrics)? I'm training with an infinite epoch, so was thinking this could be a source of the issue

[tchaton]

Hey @EricWiener. Are you overriding on_epoch_X hook. The output of your steps would be stored automatically for you. This could explain the memory leakage.

[EricWiener]

I'm not overriding on_epoch_X hook, but I am overriding pl_module.log with:

def log(  # type: ignore[override]
        self, name: str, value: _METRIC_COLLECTION, **kwargs: Any
    ) -> None:
        if self.trainer and self.trainer.state.stage == RunningStage.TRAINING:
            kwargs.update(on_step=True)
            kwargs.update(on_epoch=False)
        super().log(name, value, **kwargs)

with the intent to turn off all epoch behavior.

[tchaton]

Hey @EricWiener,

Thanks for sharing. If you manage to make a reproducible script, this would be quite good.

[carmocca]

We should log a warning when training is configure to run an infinite epoch and the LightningModule overrides training_epoch_end.

Proposed in #11554

[tchaton]

@carmocca Mind creating an issue ?

[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!