About OOM issue related to 'create_graph=True'

[DOHA-HWANG]

Dear Author,

Above all, thank you for sharing nice codes.

BTW, about quant training on CIFAR10,
Have you ever faced with OOM issues by loss.backward(create_graph=True) in update_grad_scales?
When I tried it by below args, I was faced with the "RuntimeError: CUDA out of memory" issue.

python train_quant.py --gpu_id '0'
--weight_levels 8
--act_levels 8
--baseline False
--use_hessian True
--load_pretrain True
--pretrain_path '../results/ResNet20_CIFAR10/fp/checkpoint/last_checkpoint.pth'
--log_dir '../results/ResNet20_CIFAR10/ours(hess)/W8A8/

Do you have some idea to avoid this issue?

Thank you in advance.

[Joejwu]

So have you solved the problem? I have the same problem

[kartikgupta-at-anu]

I run into same issue when training on ImageNet.

[DOHA-HWANG]

I haven't used this code recently. That's why I can't remember clearly how to avoid this problem.
However, when I checked the last private code, I disabled "--use_hessian False". (And I remember I had success training this project before.)
As far as I know, using hessian was not crucial in this project according to their paper experiments.

[kartikgupta-at-anu]

use_hessian is important if we want the scale factors in EWGS equation to be based on hessian.
I figured out a solution to this problem. The code in this part is not well-written and thus there are some references left dangling due to which it keeps on accumulating graphs. One possible solution/hack to avoid OOM issue is to do a loss.backward() in utils.py after line 116. This will release the graph after each batch.

EWGS/ImageNet/utils.py

Line 116 in 56c654c