Memory Leak when instantiating Fabric multiple times

[vkakerbeck]

Bug description

As mentioned here I am trying to instantiate multiple Fabric instances in a loop (one per iteration). However, after every iteration, the memory consumption goes up.
Is there some fabric.teardown() like method that could help here? Could this memory increase be caused by something else?

This is the memory increase with 3 iterations:

For reference here you can see when a new iteration starts (step count resets):

What version are you seeing the problem on?

v2.0

How to reproduce the bug

I tried to make a self-contained example of my problem (derived from lit-gpt code but simplified to illustrate the problem). Parameters are just chosen to show the problem (using large enough network). devices set to 1 for debugging but also shows up for larger number.


import shutil
import warnings
from pathlib import Path
from types import SimpleNamespace

import lightning as L
import torch
import torch.nn as nn
import wandb
from lightning.fabric.accelerators.mps import MPSAccelerator
from lightning.fabric.strategies import DeepSpeedStrategy


def main():
    config = SimpleNamespace(
        input_size=1024,
        output_size=1024,
        hidden_size=8192,
        num_layers=5,
        devices=1,
        ds_config=None,
        precision="bf16-true",
        entity="x",
        wandb_project="project",
        run_name="memory_test",
        checkpoint_path=None,
        learning_rate=1e-4,
        weight_decay=1e-5,
        max_steps=1000,
        batch_size=64,
        log_interval=1,
        out_dir=Path("./out"),
    )

    batch_size_per_device = config.batch_size / config.devices
    micro_batch_size = 1
    config.gradient_accumulation_steps = int(batch_size_per_device // micro_batch_size)

    for epoch in range(3):
        config.checkpoint_path = run_epoch(config)


def run_epoch(config):
    # Load Fabric
    fabric = L.Fabric(
        devices=config.devices,
        strategy=(
            DeepSpeedStrategy(config=config.ds_config) if config.devices > 1 else "auto"
        ),
        precision=config.precision,
    )
    fabric.launch()
    fabric.seed_everything(1337 + fabric.global_rank)

    if fabric.global_rank == 0:
        wandb.init(
            entity=config.entity,
            project=config.wandb_project,
            name=config.run_name,
            config=config,
        )

    with fabric.device:
        torch.set_default_tensor_type(torch.HalfTensor)
        model = LargeMODEL(config).bfloat16()
        torch.set_default_tensor_type(torch.FloatTensor)
        # Load checkpoint if this is not the first epoch
        if config.checkpoint_path is not None:
            checkpoint_path = config.checkpoint_path
            checkpoint = torch.load(checkpoint_path)
            model.load_state_dict(checkpoint, strict=False)

    # Setup model and optimizer in fabric
    optimizer = torch.optim.AdamW(
        model.parameters(), lr=config.learning_rate, weight_decay=config.weight_decay
    )
    model, optimizer = fabric.setup(model, optimizer)

    train_data = load_datasets(config)

    # Train the model
    train(
        config,
        fabric,
        model,
        optimizer,
        train_data,
    )

    # Save the final checkpoint at the end of training
    save_path = config.out_dir / "model_trained.pth"
    fabric.print(f"Saving weights to {str(save_path)!r}")
    print(save_path)
    save_model_checkpoint(fabric, model, save_path)
    return save_path


def train(config, fabric, model, optimizer, train_data):
    step_count = 0
    max_iters = int(config.max_steps * config.gradient_accumulation_steps)
    for iter_num in range(max_iters):
        input_ids, targets = get_batch(fabric, train_data, config.batch_size)
        with fabric.no_backward_sync(
            model, enabled=((iter_num + 1) % config.gradient_accumulation_steps != 0)
        ):
            logits = model(input_ids)
            loss = loss_fn(logits, targets)
            fabric.backward(loss / config.gradient_accumulation_steps)
        if (iter_num + 1) % config.gradient_accumulation_steps == 0:
            optimizer.step()
            optimizer.zero_grad()
            step_count += 1
            fabric.call("on_train_batch_end", model=model)

        # Report performance to command line and wandb
        if step_count % config.log_interval == 0:
            if fabric.global_rank == 0:
                wandb.log(
                    {
                        "iter": iter_num,
                        "step": step_count,
                        "train/loss": loss,
                        "train/lr": optimizer.param_groups[0]["lr"],
                    }
                )


def save_model_checkpoint(fabric, model, file_path: Path):
    file_path = Path(file_path)
    # Ensure the directory exists
    file_path.parent.mkdir(parents=True, exist_ok=True)

    if isinstance(fabric.strategy, DeepSpeedStrategy):
        from deepspeed.utils.zero_to_fp32 import (
            get_fp32_state_dict_from_zero_checkpoint,
        )

        tmp_path = file_path.with_suffix(".tmp")
        fabric.save(tmp_path, {"model": model})
        fabric.barrier()
        if fabric.global_rank == 0:
            state_dict = get_fp32_state_dict_from_zero_checkpoint(tmp_path)
            torch.save(state_dict, file_path)
            shutil.rmtree(tmp_path)
    else:
        if fabric.global_rank == 0:
            state_dict = model.state_dict()
            torch.save(state_dict, file_path)
        fabric.barrier()

class LargeMODEL(nn.Module):
    def __init__(self, config):
        super(LargeMODEL, self).__init__()

        layers = []
        input_size = config.input_size
        for _ in range(config.num_layers):
            layers.append(nn.Linear(input_size, config.hidden_size))
            layers.append(nn.ReLU())
            input_size = config.hidden_size

        layers.append(nn.Linear(config.hidden_size, config.output_size))

        self.net = nn.Sequential(*layers)

    def forward(self, x):
        return self.net(x)


def load_datasets(config):
    # For the sake of this demonstration, let's return some random data
    data = torch.randn(1000, config.input_size), torch.randn(1000, config.output_size)
    return data


def get_batch(fabric, data, batch_size):
    input_ids, targets = data
    ix = torch.randint(len(data), (batch_size,))
    x, y = input_ids[ix], targets[ix]
    # Return the first batch_size samples from the dataset
    if isinstance(fabric.accelerator, MPSAccelerator):
        x, y = fabric.to_device((x, y))
    else:
        x, y = fabric.to_device((x.pin_memory(), y.pin_memory()))
    return x, y


def loss_fn(logits, targets):
    return ((logits - targets) ** 2).mean()


if __name__ == "__main__":
    # Uncomment this line if you see an error:
    # "Expected is_sm80 to be true, but got false"
    # torch.backends.cuda.enable_flash_sdp(False)
    torch.set_float32_matmul_precision("high")

    # from jsonargparse.cli import CLI
    warnings.filterwarnings(
        # false positive using deepspeed:
        # https://github.com/Lightning-AI/lightning/pull/17761#discussion_r1219705307
        "ignore",
        message="Remove `.no_backward_sync()` from your code",
    )
    main()

### Error messages and logs

No error message in the example code above. If the network size increases I get the following warning:

[WARNING] [stage3.py:1898:step] 14 pytorch allocator cache flushes since last step. this happens when there is high memory pressure and is detrimental to performance. if this is happening frequently consider adjusting settings to reduce memory consumption. If you are unable to make the cache flushes go away consider adding get_accelerator().empty_cache() calls in your training loop to ensure that all ranks flush their caches at the same time

If the size increases further or I use less GPUs I simply get an OOM error.


### Environment

<details>
  <summary>Current environment</summary>

`
#- Lightning Component (e.g. Trainer, LightningModule, LightningApp, LightningWork, LightningFlow):
#- PyTorch Lightning Version (e.g., 1.5.0): 2.0.6
#- Lightning App Version (e.g., 0.5.2): 2.1.0.dev0
#- PyTorch Version (e.g., 2.0): 2.1.0.dev20230801+cu118
#- Python version (e.g., 3.9): 3.10.10
#- OS (e.g., Linux): Linux
#- CUDA/cuDNN version: 11.8
#- GPU models and configuration: NVIDIA A100-SXM4-80GB
#- How you installed Lightning(`conda`, `pip`, source): pip
#- Running environment of LightningApp (e.g. local, cloud): 

More info

No response

cc @carmocca @justusschock @awaelchli @Borda

[awaelchli]

@vkakerbeck I copied your code and ran it for several minutes.

I get this for PyTorch 2.0.1

And I get this (similar as you reported but not quite the same) for PyTorch 2.1 nightly:

In absolute numbers, it seems that 2.1 nightly seems to use less memory overall, which is the good news. I'm not sure what the cause for your fluctuations and jumps are. I couldn't reproduce them despite using the same environment (A100, PyTorch 2.0.1 and 2.1 nightly, Python 3.10, Lightning 2.1dev). I'll let you know if I find anything else.

[vkakerbeck]

Thanks for looking into this! Did you run the code until completion? Does the spike in memory around 5 minutes on PyTorch 2.1 nightly correspond to when the next fabric instance is initialized?

I wrapped the call to run_epoch into its own process using pythons multiprocessing library so that the process is completely terminated after the function is executed. This seems to solve the memory leak (when calling run_epoch the second and third time, memory use is the same and doesn't increase with every call). So there definitely seems to be a memory leak issue somewhere in this function... I'm not sure where though.

[ipoletaev]

Another way to reproduce the leaks. Leak-s because it looks like both VRAM and RAM leak :(

import lightning.fabric as lightning_fabric
import peft
import torch
import transformers
from lightning.fabric import strategies

DS_CONFIG = {
    "train_micro_batch_size_per_gpu": 1,
    "gradient_accumulation_steps": 16,
    "gradient_clipping": 1.0,
    "zero_allow_untested_optimizer": True,
    "bf16": {
        "enabled": True,
        "loss_scale": 0,
        "initial_scale_power": 16,
        "loss_scale_window": 1000,
        "hysteresis": 2,
        "min_loss_scale": 1
    },
    "zero_optimization": {
        "stage": 2,
        "allgather_partitions": True,
        "reduce_scatter": True,
        "reduce_bucket_size": 1e9,
        "allgather_bucket_size": 1e9,
        "overlap_comm": True,
        "contiguous_gradients": True
    }
}


def apply_lora(model: torch.nn.Module) -> torch.nn.Module:
    for _, param in model.named_parameters():
        param.requires_grad = False

    peft_config = peft.LoraConfig(
        task_type=peft.TaskType.CAUSAL_LM,
        target_modules=["q_proj", "k_proj", "v_proj", "o_proj"],
        r=8,
        lora_alpha=16,
        lora_dropout=0.0)
    return peft.get_peft_model(model, peft_config=peft_config)


for i in range(10):
    print("--" * 30 + "Attempt #{}".format(i) + "--" * 30)
    fabric = lightning_fabric.Fabric(
        precision='bf16',
        accelerator="gpu",
        strategy=strategies.DeepSpeedStrategy(config=DS_CONFIG))
    fabric.launch()
    with fabric.init_module(empty_init=True):
        model = transformers.AutoModelForCausalLM.from_pretrained(
            "meta-llama/Llama-2-7b-hf",
            quantization_config=transformers.BitsAndBytesConfig(
                load_in_4bit=True, bnb_4bit_quant_type="nf4"))
    model = apply_lora(model)
    optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
    model, optimizer = fabric.setup(model, optimizer)

Public LLaMA and LoRA don't matter. Even a simpler model leads to the same issue - but only if DeepSpeed is used.