Lightning-AI / Lightning-AI/pytorch-lightning

Loading checkpoint before fabric.setup(model) gets abnormal loss when using fabric.init_module()

Open
#20,490 4 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

bug ver: 2.4.x
Dominant language
Python
Stars
31.4k
Forks
3.8k
Avg merge
6d 7h
Merged PRs (30d)
6

Description

### Bug description

Init model with `fabric.init_module(True)` and load checkpoint **after** `model = fabric.setup(model)`, the training loss is normal
```
with fabric.init_module(empty_init=(fabric.world_size > 1)):
model = GPT(config)
model = fabric.setup(model)
load_checkpoint(fabric, model, checkpoint_path)

step = 1 | loss train: 0.8448048233985901
step = 2 | loss train: 1.3229767084121704
step = 3 | loss train: 1.2647839784622192
step = 4 | loss train: 1.287076711654663
step = 5 | loss train: 1.0357563495635986
```
but when loading checkpoint **before** `model = fabric.setup(model)`, get loss much larger
```
with fabric.init_module(empty_init=(fabric.world_size > 1)):
model = GPT(config)
load_checkpoint(fabric, model, checkpoint_path)
model = fabric.setup(model)

step = 1 | loss train: 12.027938842773438
step = 2 | loss train: 12.051375389099121
step = 3 | loss train: 12.112957954406738
step = 4 | loss train: 12.08558177947998
step = 5 | loss train: 12.089488983154297
```

Another phenomenon is that, if not using `fabric.init_module()`, I can get normal loss when loading checkpoint before `fabric.setup(model)`,
```
# with fabric.init_module(empty_init=(fabric.world_size > 1)):
if True:
model = GPT(config)
load_checkpoint(fabric, model, checkpoint_path)
model = fabric.setup(model)

step = 1 | loss train: 0.8447667956352234
step = 2 | loss train: 1.3229438066482544
step = 3 | loss train: 1.2663335800170898
step = 4 | loss train: 1.2902932167053223
step = 5 | loss train: 1.035811185836792
```
So how to load hf models converted by `litgpt.scripts.convert_hf_checkpoint` in a correct way?

### What version are you seeing the problem on?

v2.4

### How to reproduce the bug

```python
from pathlib import Path

import torch
import lightning as L
from lightning.fabric.strategies import FSDPStrategy

from litgpt.args import TrainArgs
from litgpt.config import Config
from litgpt.model import GPT, Block
from litgpt.data import Alpaca2k
from litgpt.tokenizer import Tokenizer
from litgpt.utils import (
chunked_cross_entropy,
load_checkpoint,
num_parameters,
get_default_supported_precision,
)

def get_lr_scheduler(optimizer, warmup_steps: int, max_steps: int):
# linear warmup followed by cosine annealing
scheduler1 = torch.optim.lr_scheduler.LambdaLR(optimizer, lambda step: step / warmup_steps)
scheduler2 = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=(max_steps - warmup_steps))
return torch.optim.lr_scheduler.SequentialLR(optimizer, [scheduler1, scheduler2], milestones=[warmup_steps])

def main(
checkpoint_dir: Path,
devices: int = 8,
num_nodes: int = 1,
precision: str = "bf16-true",
seed: int = 1337,
) -> None:
torch.set_float32_matmul_precision("high")

train_args = TrainArgs(
save_interval = 1000,
log_interval = 1,
global_batch_size = 64,
micro_batch_size = 4,
lr_warmup_steps = 1000,
epochs = 10,
max_steps = 10000,
)

strategy = FSDPStrategy(
auto_wrap_policy={Block},
activation_checkpointing_policy={Block},
state_dict_type="full",
limit_all_gathers=True,
cpu_offload=False,
)

fabric = L.Fabric(
accelerator="cuda",
devices=devices,
num_nodes=num_nodes,
strategy=strategy,
precision=precision,
)
fabric.launch()
fabric.seed_everything(seed) # same seed for every process to init model (FSDP)

dataset = Alpaca2k()
tokenizer = Tokenizer(str(checkpoint_dir))
dataset.connect(tokenizer, batch_size=train_args.micro_batch_size, max_seq_length=512)
with fabric.rank_zero_first():
dataset.prepare_data()
dataset.setup()
dataloader = dataset.train_dataloader()
dataloader = fabric.setup_dataloaders(dataloader)

checkpoint_path = str(checkpoint_dir / "lit_model.pth")
config = Config.from_file(checkpoint_dir / "model_config.yaml")
with fabric.init_module(empty_init=(fabric.world_size > 1)):
model = GPT(config)
fabric.print(f"Number of trainable parameters: {num_parameters(model, requires_grad=True):,}")
# load_checkpoint(fabric, model, checkpoint_path)
model = fabric.setup(model)
load_checkpoint(fabric, model, checkpoint_path)

optimizer = torch.optim.AdamW(model.parameters(), lr=0.0002, weight_decay=0.0, betas=(0.9, 0.95))
optimizer = fabric.setup_optimizers(optimizer)
scheduler = get_lr_scheduler(optimizer, warmup_steps=train_args.lr_warmup_steps, max_steps=train_args.max_steps)

model.train()
for epoch in range(train_args.epochs):
for step, batch in enumerate(dataloader, 1):
input, target = batch["input_ids"], batch["labels"]
logits = model(input)
loss = chunked_cross_entropy(logits[..., :-1, :], target[..., 1:])
fabric.backward(loss)

optimizer.step()
optimizer.zero_grad()
scheduler.step()
fabric.print(f"{step = } | loss train: {loss.detach().item()}")

if __name__ == "__main__":
checkpoint_dir = Path("./Qwen2.5-1.5B/")

main(checkpoint_dir)
```

### Error messages and logs

```
# Error messages and logs here please
```

### Environment

Current environment

```
#- PyTorch Lightning Version (e.g., 2.4.0):
#- PyTorch Version (e.g., 2.4.1):
#- Python version (e.g., 3.10):
#- OS (e.g., Linux):
#- CUDA/cuDNN version:12.1
#- GPU models and configuration:
#- How you installed Lightning(`conda`, `pip`, source): pip
```

### More info

_No response_

Contributor guide

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Research direction

Reproduce the two loading orders in the provided Fabric/FSDP example, starting with litgpt.utils.load_checkpoint and the fabric.init_module and fabric.setup entry points. Compare this with checkpoints produced by litgpt.scripts.convert_hf_checkpoint, and verify that loading the converted checkpoint before setup produces the expected training loss without changing the non-empty-init behavior.

Written by the indexing model from the issue text.

Assessment

Tech stack
python, pytorch
Domain
distributed-systems, machine-learning
Issue type
Bug
Difficulty
4/5
Estimated time
3-5 days
Activity status
Stale
Clarity
Mostly clear
Newbie friendliness
35/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.