GPU memory usage spikes when loading models

[championsnet]

GPU usage goes really high when loading a model for fine-tuning and then goes back down, possibly due to the way the models are initialized.

"arguments": {
        "plm": "esm2_t48_15B_UR50D",
        "ft_method": "lora",
        "target_layers": "last",
        "data_type": "aav",
        "head_config": "aav_sampled_last_lora_linear_head_config.json",
        "ray_tuning": "False",
        "split": "sampled",
        "function": "fine_tuning",
        "reduction": "mean",
        "layer": "last",
        "output_dir": "/cluster/scratch/...",
        "experiment_name": "aav_sampled_esm2_t48_15B_UR50D_lora_last_mean_linear_regression",
        "experiment_dir": "/cluster/scratch/.../fine_tuning/lora_last/aav_sampled_esm2_t48_15B_UR50D_lora_last_mean_linear_regression/20241119_103612",
        "logger": "local",
        "cpus": 1,
        "gpus": "4",
        "nodes": 1,
        "beta": "True",
        "experimenting": "False",
        "zeroed": "False",
        "garbage": "False",
        "nulled": "False",
        "weights": null
    },
    "head_config": {
        "architecture_parameters": {
            "network_type": "linear",
            "output_dim": 1,
            "task": "regression",
            "dropout": 0
        },
        "training_parameters": {
            "learning_rate": 0.0001,
            "epochs": 6,
            "batch_size": 4,
            "loss_f": "mse",
            "optimizer": "adam",
            "val_split": 0.2,
            "weight_decay": 0.01,
            "warmup_steps": 0,
            "gradient_clipping": false,
            "scheduler": false,
            "scaler": false,
            "gradient_accumulation": 1,
            "early_stopping": 5,
            "epoch_sizing": 1.0,
            "model_output": "logits"
        }
    },
    "lora_config": {
        "r": 8,
        "lora_alpha": 16,
        "lora_dropout": 0.1,
        "bias": "none",
        "modules_to_save": [
            "classifier",
            "mlm",
            "lm_head"
        ]
    },