Memory leak/garbage collection issue (?) when using jax.device_get()

[mar-muel]

Hi all - I'm experiencing my RAM memory filling up over time when using jax.device_get() to move data to the host.

Here's what I'm doing: I have a while True loop that 1) loads a model (in this case a UNet), 2) trains it, 3) copies the model to host & saves a checkpoint, and 4) reset the model weights to a base model. Here's a simplified version of this code:

import logging
from flax import jax_utils
import jax
import optax
from flax.training import train_state
from diffusers import FlaxUNet2DConditionModel


logging.basicConfig(level=logging.INFO, format='%(asctime)s [%(levelname)-5.5s] [%(name)-12.12s]: %(message)s')
logger = logging.getLogger(__name__)


model_dir = '/mnt/sdb/models/sd-1.5-jax/unet'


def load_unet():
    model, model_params = FlaxUNet2DConditionModel.from_pretrained(model_dir)
    return model, model_params


def get_params_to_save(params):
    return jax.device_get(jax.tree_util.tree_map(lambda x: x[0], params))


def train(unet, unet_params):
    logger.info('Creating train state...')
    adamw = optax.adamw(learning_rate=1e-6)
    unet_state = train_state.TrainState.create(apply_fn=unet.__call__, params=unet_params, tx=adamw)

    logger.info('Replicating train state...')
    unet_state = jax_utils.replicate(unet_state)

    # train unet_state here...

    logger.info('Save checkpoint...')
    unet_params_ckpt = get_params_to_save(unet_state.params)

    # continue to save unet_params_ckpt to disk...


def main():
    logger.info('Loading model...')
    unet, unet_params = load_unet()

    while True:
        train(unet, unet_params)

        logger.info('Re-loading model...')
        unet, unet_params = load_unet()


if __name__ == "__main__":
    main()

With each iteration of the While True my RAM memory fills up by around 4GB (the size of the UNet). If I comment out get_params_to_save I don't encounter the issue. However, it seems to be not 100% consistent - so I see my memory go up at times, and sometimes it stays flat 🤔 .

My suspicion is that jax.device_get() keeps a reference to unet_state.params, which means unet_state is not properly garbage collected inside train() (or at least not consistently). I have tried to explicitly garbage collect (using gc.collect()) without luck - I assume gc.collect() only operates on Python objects, but has no influence on C/Jax structures?

I can replace the line

unet_params_ckpt = jax.device_get(jax.tree_util.tree_map(lambda x: x[0], params))

with

unet_params_ckpt = jax.tree_map(lambda x: np.asarray(x[0]), unet_state.params)

which seems to solve the issue. So I was wondering what the difference between these two lines? Could I trigger jax's internal garbage collection manually to ensure all objects created inside train() are garbage collected? I'd be super happy about any insight on this! 🙏

Side note: You may ask why I'm not just simply loading the model inside train() and then discarding it. The reason is other parts of the codebase as well as performance considerations.

[felipeangelimvieira]

May be related to felipeangelimvieira/prophetverse#94