torch.manual_seed leaks memory

[sniklaus]

🐛 Bug

Calling torch.manual_seed(...) (or torch.cuda.manual_seed_all(...)) leaks memory, quite a bit in fact. I am resetting the seed in my data loader for each sample and over the course of a few hours there are tens of gigabytes that are being leaked.

To Reproduce

for i in range(1000000000): torch.manual_seed(i)

And watch the memory consumption of the process slowly creep up.

Expected behavior

The memory consumption to remain constant, as is the case with for i in range(1000000000): random.seed(i).

Environment

PyTorch version: 1.8.1+cu102
Is debug build: False
CUDA used to build PyTorch: 10.2
ROCM used to build PyTorch: N/A

OS: Ubuntu 20.04.2 LTS (x86_64)
GCC version: (Ubuntu 9.3.0-17ubuntu1~20.04) 9.3.0
Clang version: 10.0.0-4ubuntu1 
CMake version: version 3.19.6

Python version: 3.8 (64-bit runtime)
Is CUDA available: True
CUDA runtime version: 10.1.243
GPU models and configuration: GPU 0: GeForce GTX 1650
Nvidia driver version: 460.39
cuDNN version: Could not collect
HIP runtime version: N/A
MIOpen runtime version: N/A

cc @ezyang @gchanan @zou3519 @bdhirsh @jbschlosser @anjali411 @pbelevich

[ngimel]

High pri to investigate leak.

[eqy]

FWIW I think this is not just a bug with manual_seed, but also in other generator functions. On upstream I see increasing memory usage with torch.seed() as well.

[kshitij12345]

Turns out that the problem is with the lazy initialization of cuda.

pytorch/torch/cuda/__init__.py

Lines 114 to 123 in 389380f

	def _lazy_call(callable):
	if is_initialized():
	callable()
	else:
	# TODO(torch_deploy): this accesses linecache, which attempts to read the
	# file system to get traceback info. Patch linecache or do something
	# else here if this ends up being important.
	# Don't store the actual traceback to avoid memory cycle
	_queued_calls.append((callable, traceback.format_stack()))

If cuda is not initialized, any calls wrapped in _lazy_call get appended to _queued_calls (in this case, all calls to manual_seed were queued). Memory doesn't grow if we call torch.cuda.init() before the loop.

Also many generator related functions besides manual_seed as noted by @eqy also call _lazy_call.
Ref : https://github.com/pytorch/pytorch/blob/master/torch/cuda/random.py

To this end, I think we can choose to fix this in 1 of the 3 ways.

1. Once a threshold for queued calls is reached, just init the CUDA. (this maybe bad because maybe they never intended to use cuda)
2. Once a threshold for queued calls is reached, warn the user that this is queuing is happening and they should init cuda (maybe?)
3. Once a threshold for queued calls is reached, raise an error with msg. (raising an error for implementation detail doesn't seem like a good option)

I think second is ok. Because, they maybe unknowingly queuing the calls on CUDA by using manual_seed or other similar function, so getting a warning makes sense.

---

(Also we should empty _queued_calls once cuda is initialized)

pytorch/torch/cuda/__init__.py

Lines 176 to 187 in 389380f

	_tls.is_initializing = True
	try:
	for queued_call, orig_traceback in _queued_calls:
	try:
	queued_call()
	except Exception as e:
	msg = (f"CUDA call failed lazily at initialization with error: {str(e)}\n\n"
	f"CUDA call was originally invoked at:\n\n{orig_traceback}")
	raise DeferredCudaCallError(msg) from e
	finally:
	delattr(_tls, 'is_initializing')
	_initialized = True

[sniklaus]

Thanks for looking into this!

But why does torch.manual_seed() have anything to do with CUDA? For torch.cuda.manual_seed() sure, but not the former.

None of the proposed solutions seem good to me. Consider calling torch.manual_seed() repeatedly in the data loader, 1) wouldn't work due to requiring CUDA, 2) wouldn't work since it doesn't fix the issue, and 3) just makes the data loader fail.

[kshitij12345]

Agree that none of the proposed solution actually fix the problem. But updating manual_seed to not update CUDA's seed would be BC-Breaking. However, I agree that logically it makes sense to decouple them.

Also if we decide not to do decouple, the documentation should be updated to reflect this behaviour of setting seed for CUDA generators as well.
https://pytorch.org/docs/stable/generated/torch.manual_seed.html?highlight=manual_seed#torch.manual_seed

EDIT: Look at PR #2811 for more idea about the motivation

This behavior of setting seed for CUDA has been there since 0.2

cc: @ngimel

[ezyang]

There is a simple fix: seeding is memoryless, so instead of tracking every modification to seed, track the "latest" CUDA seed only, and apply that on its own at the end. This won't fix leaks for other API calls but it seems more urgent to fix this one first.

@kshitij12345 do you think that's enough detail to prepare a PR?

[kshitij12345]

Ah. That makes sense. Thanks @ezyang