Developer chat

fastai fastai dev
666

Last few days I’ve been tracking the cause of unrecoverable Out of Memory and a mem leakage on manual interrupt of the nb run. I first found the solution to the problem, which I have been polishing for quite a while only to discard it after digging deeper and finding the cause, and then fixing the cause.

So when you get CUDA OOM and you can’t recover from it w/o restart, or when you get memory leaked when you hit stop during training, the cause is ipython. It stores the traceback of the exception. The traceback ties up the locals() and they don’t get released until… another exception occurs, at which point it frees up the old tb, which allows gc.collect() to do its work. Ouch. It was quite a journey to figure it out and I have learned a lot about python on the way.

I submitted a fix here https://github.com/ipython/ipython/pull/11572 - it seems some tests that compare the exact tb no longer match, but I trust they will figure it out. Imagine that! a one line fix and now you can OOM as much you’d like and continue running your notebook! Amazing!

If you want to understand more about the problem, I explained the nuances of the problem of saving a traceback or an exception object here.

Until ipython sorts it out, if you need a solution today, you can either do a hotfix for your installed version of ipython so you can enjoy the change now,:

curl https://github.com/ipython/ipython/commit/657cde76ad07ec5b69470758d9bb6adbae88a1da.patch > /tmp/tb-leak-fix.patch
cd $CONDA_PREFIX/lib/python3.7/site-packages/
patch -p1 < /tmp/tb-leak-fix.patch

adjust the path of course, this for python 3.7

Alternatively, here is some magic code for you:

import functools
import traceback
def get_ref_free_exc_info():
    "Free traceback from references to locals/globals to avoid circular reference leading to gc.collect() unable to reclaim memory"
    type, val, tb = sys.exc_info()
    traceback.clear_frames(tb)
    return (type, val, tb)

def gpu_mem_restore(func):
    "Reclaim GPU RAM if CUDA out of memory happened, or execution was interrupted"
    @functools.wraps(func)
    def wrapper(*args, **kwargs):
        try:
            return func(*args, **kwargs)
        except:
            type, val, tb = get_ref_free_exc_info() # must!
            raise type(val).with_traceback(tb) from None
    return wrapper

Now add before any of your functions:

@gpu_mem_restore
def fit(...

and OOM is now recoverable! And interrupts leak no memory!

Regardless of ipython’s fix this is now part of fastai, so you should be able to see the impact by just using the latest git. At the moment only functions that call fit() are positively affected.

Here is a notebook that demonstrates the OOM w/o the leak and that recovers almost 100% of memory w/o restart, using the current fastai git: https://github.com/fastai/fastai_docs/blob/master/dev_nb/mem_leaks/OOM_on_fit_recover.ipynb

And if you want to protect just a few lines of code, here is a context manager that does the same:

class gpu_mem_restore_ctx():
    " context manager to reclaim GPU RAM if CUDA out of memory happened, or execution was interrupted"
    def __enter__(self): return self
    def __exit__(self, exc_type, exc_val, exc_tb):
        if not exc_val: return True
        traceback.clear_frames(exc_tb)
        raise exc_type(exc_val).with_traceback(exc_tb) from None

So now you can do:

with gpu_mem_restore_ctx():
    learn.fit_one_cycle(1,1e-2)

with the same results. Except this one (fit functions) is already protected, this would be more useful for your custom code.

Both functions are now in https://github.com/fastai/fastai/blob/master/fastai/utils/mem.py so you will just need to from fastai.utils.mem import * before you can use them.

BTW, another workaround is to throw another exception following the OOM exception:

# cell1 - if this leads to OOM leak
learn.fit_one_cycle(1,1e-2)
# cell 2 - this will release the memory, since it will reset %tb and free its locals()
assert False, "please liberate my GPU!"

If you want a more exact case where it only recovers from OOM, but the problem remains with any other exception it’d be:

def gpu_mem_restore(func):
    "Reclaim GPU RAM if CUDA out of memory happened, or execution was interrupted"
    @functools.wraps(func)
    def wrapper(*args, **kwargs):
        oom_exc = False
        try:
            return func(*args, **kwargs)
        except RuntimeError as e:
            if "CUDA out of memory" in str(e):
                oom_exc = True
                type, val, tb = get_ref_free_exc_info() # must!
                raise type(val).with_traceback(tb) from None
            else: raise # re-raises the exact last exception
        except: raise # any other types of errors
        finally:
            if oom_exc:
                # reclaim memory
                gc.collect()
                if torch.cuda.is_available(): torch.cuda.empty_cache()
    return wrapper

(need to include the KeyboardInterrupt type in there too)

If you encounter any related issues you can discuss those here: A guide to recovering from CUDA Out of Memory and other exceptions

40 Likes