lllyasviel / lllyasviel/ControlNet
Running tutorial_train.py on Tesla V100 16Gb results in out-of-memory during image logging
Nobody has claimed this yet.
- Dominant language
- Python
- Stars
- 34.1k
- Forks
- 3k
- PR merge metrics
- No merged PRs in 30d
Description
Running tutorial_train.py (with fill50k dataset) on Tesla V100 16Gb results in out-of-memory during image logging (DDIM Sampler).
I am running with `batch_size=1` and `accumulate_grad_batches=4` (OOM if `batch_size>1`).
When training, 15654MiB of the total GPU memory (16160MiB) are used.
Any suggestions?
Thanks!
---
The error message:
```
File "/home/ubuntu/ControlNet/cldm/cldm.py", line 397, in log_images
samples_cfg, _ = self.sample_log(cond={"c_concat": [c_cat], "c_crossattn": [c]},
File "/home/ubuntu/anaconda3/envs/control/lib/python3.8/site-packages/torch/autograd/grad_mode.py", line 27, in decorate_context
return func(*args, **kwargs)
File "/home/ubuntu/ControlNet/cldm/cldm.py", line 413, in sample_log
samples, intermediates = ddim_sampler.sample(ddim_steps, batch_size, shape, cond, verbose=False, **kwargs)
File "/home/ubuntu/anaconda3/envs/control/lib/python3.8/site-packages/torch/autograd/grad_mode.py", line 27, in decorate_context
return func(*args, **kwargs)
File "/home/ubuntu/ControlNet/ldm/models/diffusion/ddim.py", line 103, in sample
samples, intermediates = self.ddim_sampling(conditioning, size,
File "/home/ubuntu/anaconda3/envs/control/lib/python3.8/site-packages/torch/autograd/grad_mode.py", line 27, in decorate_context
return func(*args, **kwargs)
File "/home/ubuntu/ControlNet/ldm/models/diffusion/ddim.py", line 163, in ddim_sampling
outs = self.p_sample_ddim(img, cond, ts, index=index, use_original_steps=ddim_use_original_steps,
File "/home/ubuntu/anaconda3/envs/control/lib/python3.8/site-packages/torch/autograd/grad_mode.py", line 27, in decorate_context
return func(*args, **kwargs)
File "/home/ubuntu/ControlNet/ldm/models/diffusion/ddim.py", line 211, in p_sample_ddim
model_uncond, model_t = self.model.apply_model(x_in, t_in, c_in).chunk(2)
File "/home/ubuntu/ControlNet/cldm/cldm.py", line 339, in apply_model
eps = diffusion_model(x=x_noisy, timesteps=t, context=cond_txt, control=control, only_mid_control=self.only_mid_control)
File "/home/ubuntu/anaconda3/envs/control/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1130, in _call_impl
return forward_call(*input, **kwargs)
File "/home/ubuntu/ControlNet/cldm/cldm.py", line 30, in forward
h = module(h, emb, context)
File "/home/ubuntu/anaconda3/envs/control/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1130, in _call_impl
return forward_call(*input, **kwargs)
File "/home/ubuntu/ControlNet/ldm/modules/diffusionmodules/openaimodel.py", line 84, in forward
x = layer(x, context)
File "/home/ubuntu/anaconda3/envs/control/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1130, in _call_impl
return forward_call(*input, **kwargs)
File "/home/ubuntu/ControlNet/ldm/modules/attention.py", line 334, in forward
x = block(x, context=context[i])
File "/home/ubuntu/anaconda3/envs/control/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1130, in _call_impl
return forward_call(*input, **kwargs)
File "/home/ubuntu/ControlNet/ldm/modules/attention.py", line 269, in forward
return checkpoint(self._forward, (x, context), self.parameters(), self.checkpoint)
File "/home/ubuntu/ControlNet/ldm/modules/diffusionmodules/util.py", line 114, in checkpoint
return CheckpointFunction.apply(func, len(inputs), *args)
File "/home/ubuntu/ControlNet/ldm/modules/diffusionmodules/util.py", line 129, in forward
output_tensors = ctx.run_function(*ctx.input_tensors)
File "/home/ubuntu/ControlNet/ldm/modules/attention.py", line 272, in _forward
x = self.attn1(self.norm1(x), context=context if self.disable_self_attn else None) + x
File "/home/ubuntu/anaconda3/envs/control/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1130, in _call_impl
return forward_call(*input, **kwargs)
File "/home/ubuntu/ControlNet/ldm/modules/attention.py", line 177, in forward
sim = einsum('b i d, b j d -> b i j', q, k) * self.scale
RuntimeError: CUDA out of memory. Tried to allocate 1024.00 MiB (GPU 0; 15.78 GiB total capacity; 12.46 GiB already allocated; 1024.00 MiB free; 13.61 GiB reserved in total by PyTorch) If reserved memory is >> allocated memory try setting max_split_size_mb to avoid fragmentation. See documentation for Memory Management and PYTORCH_CUDA_ALLOC_CONF
```
Contributor guide
No contributing guide indexed for this repository
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
Research direction
Reproduce tutorial_train.py with the fill50k dataset, batch_size=1, and accumulate_grad_batches=4, then start from cldm/cldm.py:log_images and sample_log in the traceback. Inspect ldm/modules/attention.py and ldm/models/diffusion/ddim.py around the failing allocation; done means image logging no longer triggers CUDA out-of-memory on the reported V100 setup.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python, pytorch
- Domain
- machine-learning, performance
- Issue type
- Bug
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Stale
- Clarity
- Needs clarification
- Newbie friendliness
- 25/100