wandb sweeps accumulates GPU memory

Describe the bug When running a GPU sweep using any NERModel() or ClassificationModel() such as bert, distilbert, roberta, etc. the sweep retains some amount of GPU memory on every run. I tried to add a statement in my run functions passed to the agent to clear gpu cache but it did not work.

To Reproduce

sweep_config = {
    
    'name' : 'batch-16',
    'method': 'bayes', #grid, random, bayes
    'metric': {
      'name': 'f1_score',
      'goal': 'maximize'
      },

    'parameters': {
        
        'learning_rate': {'min': 0, 'max': 4e-4},
        'num_train_epochs': {'min': 1, 'max': 10},
        #'train_batch_size': {'values': [16, 32]},
        'weight_decay': {'min': 0, 'max': 0.1 } # {'values': [0, 0.01, 0.03]}

        },

   'early_terminate': {'type': 'hyperband', 'min_iter': 1}
}


sweep_id = wandb.sweep(sweep_config, project=wandb_project)


import gc

def run_training():

    # Initialize a new wandb run
    wandb.init()

    # Create a TransformerModel
    model = NERModel(
    model_type = model_type,
    model_name = model_name,
    args = model_args,
    sweep_config = wandb.config
    )

    # Train the model
    model.train_model(train_data = train, 
                      eval_data = val)
    
    
    # Sync wandb
    wandb.join()

    # clear gpu and cpu ram
    gc.collect()
    torch.cuda.empty_cache()

Expected behavior Cleared GPU and CPU ram on every run of sweep. As we can see below there is an incremental increase in GPU memory allocated for every run in a sweep.

Screenshots

Desktop (please complete the following information):

• OS

Additional context Add any other context about the problem here.

[akrahdan]

Hi @ThilinaRajapakse I am getting the same memory allocation errors for in colab, kaggle, and on-premise servers. Is there anything you can do to help?

[stale[bot]]

This issue has been automatically marked as stale because it has not had recent activity. It will be closed if no further activity occurs. Thank you for your contributions.