FSDP: Initialisation of FSDP gives OOM for LLama-3 8B on 4 V100s

Julian_Buchel: <p>I am trying to use FSDP on 4 V100 GPUs (32GB of DRAM). My model is LLama-3 8B, so it barely fits into one GPU.</p>
<p>The code below works. I can put the model on every GPU:</p>
<pre><code class="lang-auto">import os
import functools
import torch

import torch.distributed as dist
import torch.multiprocessing as mp
from torch.utils.data.distributed import DistributedSampler
from torch.utils.data import DataLoader
from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
from torch.distributed.fsdp.fully_sharded_data_parallel import (
    CPUOffload,
    BackwardPrefetch,
    ShardingStrategy
)
from torch.distributed.fsdp.wrap import size_based_auto_wrap_policy

import transformers
from transformers import AutoModelForCausalLM, AutoTokenizer
from transformers.utils.logging import (
    enable_default_handler,
    enable_explicit_format,
)
from categories import subcategories
import datasets
from datasets import load_dataset


def setup(rank, world_size):
    os.environ['MASTER_ADDR'] = 'localhost'
    os.environ['MASTER_PORT'] = '12355'

    # initialize the process group
    dist.init_process_group("nccl", rank=rank, world_size=world_size)

def cleanup():
    dist.destroy_process_group()

def fsdp_main(rank, world_size, model):
    setup(rank, world_size)
    
    dataset_name = "cais/mmlu"
    # mmlu_task = "all"
    mmlu_task = "anatomy"
    batch_size = 1

    # create a dataset
    dataset = load_dataset(dataset_name, mmlu_task)
    for subject in (subcategories.keys() if mmlu_task == "all" else [mmlu_task]):
        if rank == 0:
            print(f"Subject {subject}")
        dataset_dev = dataset["dev"].filter(lambda x: x["subject"] == subject)
        # DEBUG this can be removed later
        sampler_dev = DistributedSampler(dataset_dev, rank=rank, num_replicas=world_size, shuffle=False)
        data_loader = DataLoader(dataset_dev, sampler=sampler_dev, batch_size=batch_size)
        
        # load the model and wrap
        torch.cuda.set_device(rank)

        tokenizer = AutoTokenizer.from_pretrained("meta-llama/Meta-Llama-3-8B")
        tokenizer.pad_token = tokenizer.eos_token

        auto_wrap_policy = functools.partial(
            size_based_auto_wrap_policy, min_num_params=20000
        )

        print(f"Rank {rank} before FSDP")
        model = model.to(rank)
        # model = FSDP(model,
        #     auto_wrap_policy=auto_wrap_policy,
        #     sharding_strategy=ShardingStrategy.FULL_SHARD,
        #     device_id=rank
        # )
        print(f"Rank {rank} after FSDP")



if __name__ == "__main__":
    # set up logger
    log_level = datasets.logging.INFO
    datasets.utils.logging.set_verbosity(log_level)
    transformers.logging.set_verbosity_info()
    enable_default_handler()
    enable_explicit_format()

    model = AutoModelForCausalLM.from_pretrained("meta-llama/Meta-Llama-3-8B", low_cpu_mem_usage=True)

    WORLD_SIZE = torch.cuda.device_count()
    print(f"WORLD SIZE {WORLD_SIZE}")

    mp.spawn(fsdp_main,
        args=(WORLD_SIZE, model),
        nprocs=WORLD_SIZE,
        join=True)
</code></pre>
<p>However, when I uncomment,</p>
<pre><code class="lang-auto">model = FSDP(model,
    auto_wrap_policy=auto_wrap_policy,
    sharding_strategy=ShardingStrategy.FULL_SHARD,
    device_id=rank
)
</code></pre>
<p>I get an OOM exception:</p>
<pre><code class="lang-auto">torch.cuda.OutOfMemoryError: CUDA out of memory. Tried to allocate 1.96 GiB. GPU  has a total capacity of 31.74 GiB of which 1.27 GiB is free. Including non-PyTorch memory, this process has 30.47 GiB memory in use. Of the allocated memory 30.17 GiB is allocated by PyTorch, and 4.97 MiB is reserved by PyTorch but unallocated. If reserved but unallocated memory is large try setting PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True to avoid fragmentation.  See documentation for Memory Management  (https://pytorch.org/docs/stable/notes/cuda.html#environment-variables)
</code></pre>
<p>This may also be related to: <a href="https://github.com/pytorch/pytorch/issues/104026" class="inline-onebox" rel="noopener nofollow ugc">Is memory-efficient FSDP initialization intended to be possible with torch.device('meta')? · Issue #104026 · pytorch/pytorch · GitHub</a><br>
and also to: <a href="https://github.com/pytorch/pytorch/issues/90465" class="inline-onebox" rel="noopener nofollow ugc">[FSDP] Revisit meta device initialization · Issue #90465 · pytorch/pytorch · GitHub</a></p>
<p>But nothing works. Any thoughts?</p>
<p>[EDIT]<br>
This is the traceback:</p>
<pre><code class="lang-auto">Traceback (most recent call last):
  File "/gpfs/u/home/ANFM/ANFMbchl/scratch/miniconda3/envs/torch-nightly/lib/python3.10/site-packages/torch/multiprocessing/spawn.py", line 75, in _wrap
    fn(i, *args)
  File "/gpfs/u/home/ANFM/ANFMbchl/FSDP-Example/gpt2_example.py", line 68, in fsdp_main
    model = FSDP(model,
  File "/gpfs/u/home/ANFM/ANFMbchl/scratch/miniconda3/envs/torch-nightly/lib/python3.10/site-packages/torch/distributed/fsdp/fully_sharded_data_parallel.py", line 485, in __init__
    _auto_wrap(
  File "/gpfs/u/home/ANFM/ANFMbchl/scratch/miniconda3/envs/torch-nightly/lib/python3.10/site-packages/torch/distributed/fsdp/_wrap_utils.py", line 101, in _auto_wrap
    _recursive_wrap(**recursive_wrap_kwargs, **root_kwargs)  # type: ignore[arg-type]
  File "/gpfs/u/home/ANFM/ANFMbchl/scratch/miniconda3/envs/torch-nightly/lib/python3.10/site-packages/torch/distributed/fsdp/wrap.py", line 543, in _recursive_wrap
    wrapped_child, num_wrapped_params = _recursive_wrap(
  File "/gpfs/u/home/ANFM/ANFMbchl/scratch/miniconda3/envs/torch-nightly/lib/python3.10/site-packages/torch/distributed/fsdp/wrap.py", line 543, in _recursive_wrap
    wrapped_child, num_wrapped_params = _recursive_wrap(
  File "/gpfs/u/home/ANFM/ANFMbchl/scratch/miniconda3/envs/torch-nightly/lib/python3.10/site-packages/torch/distributed/fsdp/wrap.py", line 561, in _recursive_wrap
    return _wrap(module, wrapper_cls, **kwargs), nonwrapped_numel
  File "/gpfs/u/home/ANFM/ANFMbchl/scratch/miniconda3/envs/torch-nightly/lib/python3.10/site-packages/torch/distributed/fsdp/wrap.py", line 490, in _wrap
    return wrapper_cls(module, **kwargs)
  File "/gpfs/u/home/ANFM/ANFMbchl/scratch/miniconda3/envs/torch-nightly/lib/python3.10/site-packages/torch/distributed/fsdp/fully_sharded_data_parallel.py", line 511, in __init__
    _init_param_handle_from_module(
  File "/gpfs/u/home/ANFM/ANFMbchl/scratch/miniconda3/envs/torch-nightly/lib/python3.10/site-packages/torch/distributed/fsdp/_init_utils.py", line 598, in _init_param_handle_from_module
    _init_param_handle_from_params(state, managed_params, fully_sharded_module)
  File "/gpfs/u/home/ANFM/ANFMbchl/scratch/miniconda3/envs/torch-nightly/lib/python3.10/site-packages/torch/distributed/fsdp/_init_utils.py", line 610, in _init_param_handle_from_params
    handle = FlatParamHandle(
  File "/gpfs/u/home/ANFM/ANFMbchl/scratch/miniconda3/envs/torch-nightly/lib/python3.10/site-packages/torch/distributed/fsdp/_flat_param.py", line 582, in __init__
    self._init_flat_param_and_metadata(
  File "/gpfs/u/home/ANFM/ANFMbchl/scratch/miniconda3/envs/torch-nightly/lib/python3.10/site-packages/torch/distributed/fsdp/_flat_param.py", line 737, in _init_flat_param_and_metadata
    self.flat_param: FlatParameter = self.flatten_tensors_into_flat_param(
  File "/gpfs/u/home/ANFM/ANFMbchl/scratch/miniconda3/envs/torch-nightly/lib/python3.10/site-packages/torch/distributed/fsdp/_flat_param.py", line 850, in flatten_tensors_into_flat_param
    flat_param_data = self.flatten_tensors(tensors, aligned_numel)
  File "/gpfs/u/home/ANFM/ANFMbchl/scratch/miniconda3/envs/torch-nightly/lib/python3.10/site-packages/torch/distributed/fsdp/_flat_param.py", line 842, in flatten_tensors
    return torch.cat(flat_tensors, dim=0)
</code></pre>

Julian_Buchel: <p>CPU offloading also doesn’t make it work. I guess some additional parameters are initialized when creating FSDP causing the OOM. At that stage, CPU offloading probably doesn’t matter. Is that the case?</p>

lidenghui1110: <p>have you solved this problem? I meet the same issue and is begging for an answer.</p>

agu: <aside class="quote no-group" data-username="Julian_Buchel" data-post="1" data-topic="202377">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/julian_buchel/48/37532_2.png" class="avatar"> Julian_Buchel:</div>
<blockquote>
<p><code>model = model.to(rank)</code></p>
</blockquote>
</aside>
<p>Can you try to remove this call? If you pass <code>device_id</code> to FSDP, then it will move modules to the device for you.</p>