Multi-node training with FSDP results in weird behaviour

[srikanthsrnvs]

Bug description

Here is my train script for llama 65b

import os
import sys
import time
from functools import partial
from pathlib import Path

import lightning as L
import torch
from lightning.fabric.strategies import FSDPStrategy
from prepare import (calculate_training_steps, get_batch, load_datasets,
                     prepare, validate)
from torch.distributed.fsdp.wrap import transformer_auto_wrap_policy, size_based_auto_wrap_policy
from torch.nn.utils.init import skip_init
from torch.distributed.fsdp import (
   FullyShardedDataParallel,
   CPUOffload,
)

# support running without installing as a package
wd = Path(__file__).parent.parent.resolve()
sys.path.append(str(wd))

from lit_llama.model import Block, LLaMA, LLaMAConfig
from lit_llama.utils import save_model_checkpoint

# compilation fails as it does not support torch.complex64 for RoPE
# compile = False

beta1 = 0.9
beta2 = 0.95
grad_clip = 1.0



def train(
    output_dir: str,
    data_dir: str,
    eval_interval: int = 2000,
    eval_iterations: int = 200,
    log_interval: int = 10,
    learning_rate: float = 6e-4,
    batch_size: int = 1,
    epochs: int = 3,
    context_length: int = 4096,
    model_size: str = "65B",
    weight_decay: float = 1e-1,
):
    prepped_data_dir = prepare(
        input_file=data_dir,
        tokenizer_path=f"/scratch/data/models/llama-{model_size.lower()}/tokenizer.model"
    )
    train_data, val_data = load_datasets(prepped_data_dir)
    train_steps = calculate_training_steps(train_data, batch_size, epochs, context_length)

    print("Total number of training iterations: ", train_steps)

    ############################# TRAINING PREP #########################################
    auto_wrap_policy = partial(transformer_auto_wrap_policy, transformer_layer_cls={Block})
    strategy = FSDPStrategy(auto_wrap_policy=auto_wrap_policy, activation_checkpointing=Block, limit_all_gathers=True, cpu_offload=True, precision="16")

    fabric = L.Fabric(precision="bf16-mixed", strategy=strategy)

    # fabric.launch()
    fabric.seed_everything(1337 + fabric.global_rank)

    if fabric.global_rank == 0:
        os.makedirs(output_dir, exist_ok=True)
    ######################################################################################

    config = LLaMAConfig.from_name(model_size)
    config.block_size = context_length

    with fabric.init_module(empty_init=True):
        print("Loading model...")
        model = LLaMA(config).bfloat16()
        # Main machine, no need to load the weights again
        print("Loading checkpoint...")
        checkpoint = torch.load(f'/scratch/data/models/llama-{model_size.lower()}/lit-llama.pth')
        print("Loading state dict...")
        model.load_state_dict(checkpoint, strict=False)

    fabric.barrier()
    print("Setting up modules...")
    model = fabric.setup_module(model)

    print("Setting up the optimizers...")
    optimizer = torch.optim.AdamW(model.parameters(), lr=learning_rate, weight_decay=weight_decay, betas=(beta1, beta2), foreach=False)
    optimizer = fabric.setup_optimizers(optimizer)

    loop(
        fabric,
        model,
        optimizer,
        train_data,
        val_data,
        eval_interval,
        eval_iterations,
        output_dir,
        batch_size,
        log_interval,
        train_steps
    )

############################### MAIN TRAINING LOOP ########################################
"""The training loop.

Loosely based on the nanoGPT implementation: https://github.com/karpathy/nanoGPT.
"""
def loop(
    fabric,
    model,
    optimizer,
    train_data,
    val_data,
    eval_interval,
    eval_iterations,
    output_dir,
    batch_size,
    log_interval,
    train_steps  
):
    iter_num = 0

    print("Beginning training loop ---------->")
    while True:
        # TODO: add learning rate scheduling

        # evaluate the loss on train/val sets and write checkpoints
        if iter_num > 0 and iter_num % eval_interval == 0:
            val_loss = validate(fabric, model, val_data, eval_iterations)
            fabric.print(f"step {iter_num}: val loss {val_loss:.4f}")
            fabric.print(f"Saving checkpoint to {output_dir}")
            save_model_checkpoint(fabric, model, os.path.join(output_dir, f"iter-{iter_num:06d}-ckpt.pth"))

        t0 = time.time()

        input_ids, targets = get_batch(
            fabric,
            train_data,
            block_size=model.config.block_size,  # type: ignore[union-attr,arg-type]
            batch_size=batch_size
        )
        logits = model(input_ids)
        loss = torch.nn.functional.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1), ignore_index=-1)

        fabric.backward(loss)

        # TODO: Gradient clipping
        # if grad_clip != 0.0:
        #     fabric.clip_gradients(model, optimizer, max_norm=grad_clip)

        optimizer.step()
        optimizer.zero_grad()

        dt = time.time() - t0
        if iter_num % log_interval == 0:
            fabric.print(f"iter {iter_num}: loss {loss.item():.4f}, time: {dt*1000:.2f}ms")
        iter_num += 1

        if iter_num > train_steps:
            save_model_checkpoint(fabric, model, os.path.join(output_dir, f"iter-{iter_num:06d}-ckpt.pth"))
            break

I run this script using:

lightning run model \
    --main-port=29415 \
    --node-rank=0 \
    --main-address=IP_ADDRESS \
    --accelerator=cuda \
    --devices=8 \
    --num-nodes=2 \
    train.py \
        --ARGS...

on two separate machines with a different node rank. I'm seeing two bugs, although it just might be me not seeing something

1. When training the 30B model, the weights get initialized properly on the master Node, but on the child node, all the weights are dumped into GPU 1, which gets an OOM. No other GPUs are used. This happens in the Llama(config) step.
2. When training the 65B model, the weights are initialized on node 1, but there are only 8 GPUs with 80gb each (640GB) which is not enough to load llama, and I get an OOM again. I need to shard this across both the GPUs and the Nodes, and it seems to load a copy of llama on each node rather than one llama copy and shard across the nodes.

Is there something I'm missing or is this a bug? Thanks!

What version are you seeing the problem on?

v2.0

How to reproduce the bug

See above code

Error messages and logs

CUDA oom

Environment

Current environment

#- Lightning Component (e.g. Trainer, LightningModule, LightningApp, LightningWork, LightningFlow):
#- PyTorch Lightning Version (e.g., 1.5.0):
#- Lightning App Version (e.g., 0.5.2):
#- PyTorch Version (e.g., 2.0):
#- Python version (e.g., 3.9):
#- OS (e.g., Linux):
#- CUDA/cuDNN version:
#- GPU models and configuration:
#- How you installed Lightning(`conda`, `pip`, source):
#- Running environment of LightningApp (e.g. local, cloud):

More info

No response

cc @carmocca @justusschock @awaelchli

[awaelchli]

I'll take a closer look shortly.

Two things to note I see from the code:
You are loading the checkpoint directly with torch.load(). If the checkpoint has tensors with cuda device, the tensors will land on the GPU. This could explain why you see GPU 0 memory usage during the setup phase of the script.

When training the 65B model, the weights are initialized on node 1, but there are only 8 GPUs with 80gb each (640GB) which is not enough to load llama, and I get an OOM again. I need to shard this across both the GPUs and the Nodes, and it seems to load a copy of llama on each node rather than one llama copy and shard across the nodes.

If my remark above is True (I will verify), then that also explains your OOM here. Since you do torch.load(), it will run on every process.

We have implemented fabric.load() for this reason. I will make sure to explain how to properly load models of this size into FSDP in our upcoming detailed FSDP guide (expect to see a PR sometime next week).

[srikanthsrnvs]

Thanks, Excited to use fabric! I am currently running tests and will open some PR's myself soon once this is resolved

[tsheasha]

Hey @awaelchli I am using fabric.load as follows:

        torch.set_default_dtype(torch.float16)
        with fabric.init_module(empty_init=True):
            model = LLaMA(config).bfloat16()
        torch.set_default_dtype(torch.float32)
        checkpoint = fabric.load(model_args.model_name_or_path)
        model.load_state_dict(checkpoint, strict=False)
        torch.compile(model)

the problem i have is, the path I am pointing to is the file that was generated after converting llama weights from huggingface format to lightning format based on this script. Since I am using FSDP, fabric.load fails if i don't pass a state_dict as well according to this line.

However, even if I pass {"model": model} to fabric.load, it still fails since the path I am trying to load doesn't have the _METADATA_FILENAME which means the model won't be detected as sharded and I will have to load the model on CPU once for each worker. (so for 4 workers, i need 4 x model_size of CPU RAM before the model gets sharded in GPU)

Is there a way around this?

[srikanthsrnvs]

This is the same problem I am facing ^
Currently, I have worked around this by doing

if _is_full_checkpoint(path):
            # This is inefficient, as multiple copies of the checkpoint are held in CPU memory at once.
            # There is currently no other way because `summon_full_params` does not support write-back from rank 0 only.
            checkpoint = torch.load(path, map_location="cpu")
            with FSDP.summon_full_params(module, writeback=True, rank0_only=False):
                module.load_state_dict(checkpoint, strict=strict) ```
                
In FSDP.py

[srikanthsrnvs]

Member

I'll take a closer look shortly.

Two things to note I see from the code: You are loading the checkpoint directly with torch.load(). If the checkpoint has tensors with cuda device, the tensors will land on the GPU. This could explain why you see GPU 0 memory usage during the setup phase of the script.

When training the 65B model, the weights are initialized on node 1, but there are only 8 GPUs with 80gb each (640GB) which is not enough to load llama, and I get an OOM again. I need to shard this across both the GPUs and the Nodes, and it seems to load a copy of llama on each node rather than one llama copy and shard across the nodes.

If my remark above is True (I will verify), then that also explains your OOM here. Since you do torch.load(), it will run on every process.

We have implemented fabric.load() for this reason. I will make sure to explain how to properly load models of this size into FSDP in our upcoming detailed FSDP guide (expect to see a PR sometime next week).

Another issue is FSDP with multi-node training. This has weird issues. For one, when I attempt to do this, on my second node the LLaMA(config) line seems to load the entire model into the first GPU? But on the first node it correctly loads it on the CPU. This is really strange behaviour since the code is identical on both machines.

[tsheasha]

This is the same problem I am facing ^ Currently, I have worked around this by doing

            checkpoint = torch.load(path, map_location="cpu")
            with FSDP.summon_full_params(module, writeback=True, rank0_only=False):
                module.load_state_dict(checkpoint, strict=strict) 

This won't really work around the issue that I have since you still do torch.load() on cpu which means the model will be loaded once per process and placed on cpu ram

[srikanthsrnvs]

This is the same problem I am facing ^ Currently, I have worked around this by doing

            checkpoint = torch.load(path, map_location="cpu")
            with FSDP.summon_full_params(module, writeback=True, rank0_only=False):
                module.load_state_dict(checkpoint, strict=strict) 

This won't really work around the issue that I have since you still do torch.load() on cpu which means the model will be loaded once per process and placed on cpu ram

I misunderstood your issue. How did you get around the fact that checkpoint.pop("model") does not work since model is not a valid key? The checkpoint does not contain the key "model" since it was converted to lit-llama.pth by the script you mentioned before?

[tsheasha]

I haven't managed to work around this issue yet, I was asking to see how i can transform lit-llama.pth to a sharded format that i can then use in fabric.load as sharded which takes care of loading directly into gpu in a sharded way

[awaelchli]

From the discussion I gather two points where improvement is needed:

1. Loading a raw checkpoint that wasn't produced with fabric.save() doesn't have a hierarchical structure like checkpoint["model"] and thus can't be easily loaded with fabric.load.
2. Full-state checkpoint loading is inefficient because we need to load the checkpoint into CPU memory in each process, which causes OOM on the CPU.

A quick solution to 1) could be to simply standardize the converted checkpoint in lit-gpt/lit-llama to have the structure expected by fabric.load() and standardize the loading logic across all scripts (cc @carmocca). A long term solution would be to allow loading a raw model state dict with fabric.load directly. I will open an issue for this with a proposal.

For 2), not much can be done because that's just the drawback of having the entire model consolidated into a single checkpoint file. The entire model needs to be loaded. That is also why I put a comment into the code here:
https://github.com/Lightning-AI/lightning/blob/9a2bb85d8203c2513788397f2e40847733a28eb6/src/lightning/fabric/strategies/fsdp.py#L519-L520
I will put a warning into the upcoming FSDP documentation guide about this as well.
@tsheasha proposed that maybe we could convert the single checkpoint to a sharded checkpoint, but this is easier said than done. For this reason, I propose to investigate solutions at runtime.
One solution I can think of, is to do a turn-based loading: Processes load their checkpoint to CPU sequentially rather than in parallel. It would be slower but consume less memory. Something to investigate.

[srikanthsrnvs]

@awaelchli What is the structure expected by Fabric.load? How can I convert the HF format/torch format/lit-llama.pth format to this Fabric.load format?

Also what do you think of the following script:

auto_wrap_policy = partial(transformer_auto_wrap_policy, transformer_layer_cls={Block})
    strategy = FSDPStrategy(auto_wrap_policy=auto_wrap_policy, activation_checkpointing=Block, limit_all_gathers=True)

    fabric = L.Fabric(precision="bf16-mixed", strategy=strategy)

    # fabric.launch()
    fabric.seed_everything(1337 + fabric.global_rank)

    if fabric.global_rank == 0:
        os.makedirs(output_dir, exist_ok=True)
    ######################################################################################

    config = LLaMAConfig.from_name(model_size)
    config.block_size = context_length

    with fabric.init_module(empty_init=True):
        # Main machine, no need to load the weights again

        print("Loading model...")
        model = LLaMA(config).bfloat16()
        

    print("Setting up modules...")
    model = fabric.setup_module(model)

    print("Setting up the optimizers...")
    optimizer = torch.optim.AdamW(model.parameters(), lr=learning_rate, weight_decay=weight_decay, betas=(beta1, beta2), foreach=False)
    optimizer = fabric.setup_optimizers(optimizer)

    print("Loading checkpoint...")
    fabric.load(f'/scratch/data/{model_size.lower()}.pt', state={'model': model}, strict=False)

    print("Reached barrier. Waiting on other processes..")
    fabric.barrier()

This doesn't work because Im not sure how to load the weights from the torch format into the one fabric can use.