lllyasviel / lllyasviel/ControlNet

Running tutorial_train.py on Tesla V100 16Gb results in out-of-memory during image logging

Open
#294 1 comment 1 reaction 0 assignees View on GitHub

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

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.