Under ddp, both of load_state_dict and on_load_from_checkpoint in a torch.nn.Module will never be called

[5o1]

Bug description

I writed some code in load_state_dict to synchronize variables from other repos that are not registered in buffer.

My code looks like this:

# in a subclass of `torch.nn.Module`
    def on_load_checkpoint(self, checkpoint):
        self.sync_buffer_to_badmodule()
        print(123)

    def load_state_dict(self, state_dict, strict = True, assign = False):
        result = super().load_state_dict(state_dict, strict, assign)
        self.sync_buffer_to_badmodule()
        print(123)
        return result

More specifically, I used kmeans from the following link in my torch.nn.Module object. The repository uses a NamedTuple to hold clustering results _result which was not registered in buffer.

torch_kmeans

I found that _result in Kmeans Module was not loaded by lightning when I loaded from lightning's ckpt. So I registered my buffer in my Module and kept track of the _result of Kmeans Module. And then I restore the my buffer to Kmeans Module in load_state_dict. However, I found that load_state_dict is not called.

But no matter how hard I try, neither of these functions will be executed when loading from a checkpoint. What am I supposed to do?

If load_state_dict was not called, why are model parameters and buffers loaded correctly? Does lightning have a mechanism for loading state_dict independent of torch.nn.Module?

What version are you seeing the problem on?

v2.5

Reproduced in studio

No response

How to reproduce the bug

Error messages and logs

# Error messages and logs here please

Environment

Current environment

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

More info

No response

cc @lantiga

[5o1]

This seems to be because under ddp, torch loads state dict through another mechanism rather than load_state_dict?

I found torch distributed documentation. Here is a note from Torch that if we use checkpoint.load_state_dict, we have to call model.load_state_dict manually again:

Distributed Checkpoint - torch.distributed.checkpoint

Did lightning initialize ddp without following the notes in the torch documentation?

torch.distributed.checkpoint.state_dict_loader.load(state_dict,*,checkpoint_id=None,storage_reader=None,planner=None,process_group=None,no_dist=False)
...
Examples

my_model = MyModule()
optimizer = Adagrad(my_model.parameters())
model_state_dict = my_model.state_dict()
fs_storage_reader = torch.distributed.checkpoint.FileSystemReader(
    "/checkpoint/1"
)

torch.distributed.checkpoint.load_state_dict(
    state_dict=model_state_dict,
    storage_reader=fs_storage_reader,
)

# module.load_state_dict() function might have customized steps
# to flush the state_dict, must call it to
# ensure correct behavior.
my_model.load_state_dict(model_state_dict)

[SkafteNicki]

Hi @5o1, thanks for raising this issue. I want to be completely sure that I understand your issue correctly. Are you saying that your Kmeans module is a child module of your lightning module or how are the different modules connected to each others? Could you provide your model setup or even a fully reproducible example?

[5o1]

Hi @5o1, thanks for raising this issue. I want to be completely sure that I understand your issue correctly. Are you saying that your Kmeans module is a child module of your lightning module or how are the different modules connected to each others? Could you provide your model setup or even a fully reproducible example?

@SkafteNicki Sure. Here is a mini runable case to describe my situation

import torch
import pytorch_lightning as pl
import torch_kmeans

n_channels = 16
batch_size = 1
dtype = torch.float32

class MynnModule(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.was_load_state_dict_called = False
        self.kmeans_kernel = torch_kmeans.KMeans(n_clusters=2, verbose=False)
        self.register_buffer("mybuffer_for_sync_kmeans_kernel", torch.tensor(0, dtype=torch.bool))

        self.dummy_parameter = torch.nn.Parameter(torch.randn((n_channels), dtype=dtype))
    
    def load_state_dict(self, state_dict, strict = True, assign = False):
        result =  super().load_state_dict(state_dict, strict, assign)
        self.was_load_state_dict_called = True

        print(f"Processing load_state_dict ...")
        # Sync my buffer to kmeans kernel
        if self.mybuffer_for_sync_kmeans_kernel:
            if self.kmeans_kernel.is_fitted:
                print("The kernel state has been saved correctly. Do nothing in load_state_dict")
            else:
                print("Do something to sync state of kmeans kernel ...")
                self.kmeans_kernel.fit(torch.zeros((batch_size, 10, n_channels), device = self.mybuffer_for_sync_kmeans_kernel.device, dtype = dtype))

        return result

    def forward(self, x: torch.Tensor):
        x = x + self.dummy_parameter

        if self.mybuffer_for_sync_kmeans_kernel:
            print(f"The flag in my buffer is True, try to access the kmeans kernel.")
            if self.kmeans_kernel.is_fitted:
                print("KMeans kernel is correctly fitted.")
                print("Cluster results:", self.kmeans_kernel.predict(x.unsqueeze(1)))
            else:
                # Check load_state_dict was called
                if self.was_load_state_dict_called:
                    raise RuntimeError("load_state_dict was called, but kmeans kernel is not fitted")
                else:
                    raise RuntimeError("load_state_dict was NOT called (my issue), and kmeans kernel is not fitted")
        else:
            print("KMeans kernel is not fitted, fitting now.")
            self.kmeans_kernel.fit(torch.zeros((batch_size, 10, n_channels), device = self.mybuffer_for_sync_kmeans_kernel.device, dtype = dtype))
            if self.kmeans_kernel.is_fitted:
                print("Kmeans kernel successfully fitted. Sync state to my buffer.")
                self.mybuffer_for_sync_kmeans_kernel.fill_(1) # Sync state to my buffer
                print("Cluster results:", self.kmeans_kernel.predict(x.unsqueeze(1)))
            else:
                raise RuntimeError("KMeans kernel fitting failed.")
            
        return x
    

class TinyDataset(torch.utils.data.Dataset):
    def __len__(self):
        return 1
    
    def __getitem__(self, idx):
        x = torch.randn((n_channels), dtype=dtype)
        y = torch.randn((n_channels), dtype=dtype)
        return x, y
    

class TinyDataModule(pl.LightningDataModule):
    def train_dataloader(self):
        return torch.utils.data.DataLoader(
            TinyDataset(),
            batch_size=1
        )

class MyLightningModule(pl.LightningModule):
    def __init__(self):
        super().__init__()
        self.model = MynnModule()
    
    def configure_optimizers(self):
        return torch.optim.AdamW(self.parameters(), 0.1)
    
    def forward(self, x):
        return self.model(x)
    
    def training_step(self, batch, batch_idx):
        x, y = batch
        output = self(x)
        loss = torch.nn.functional.mse_loss(output, y)
        return loss

model = MyLightningModule()
data = TinyDataModule()

trainer = pl.Trainer(
    accelerator="gpu",
    strategy="ddp",
    devices = [0,1],
    num_nodes = 1,
    logger = False,
    max_epochs=1,
    default_root_dir="load_state_dict_test/"
)

# Step1: Fit the model to create the initial state

trainer.fit(
    model = model,
    datamodule=data
)

# Step2: Load from ckpt to restore the situation of issue

model = MyLightningModule()
data = TinyDataModule()

trainer = pl.Trainer(
    accelerator="gpu",
    strategy="ddp",
    devices = [0,1],
    num_nodes = 1,
    logger = False,
    max_epochs=2,
    default_root_dir="load_state_dict_test/"
)

trainer.fit(
    model=model,
    datamodule=data,
    ckpt_path="load_state_dict_test/checkpoints/epoch=0-step=1.ckpt"
)

Outputs:

# Step1
...
kernel is not fitted, fitting now.
kernel is not fitted, fitting now.
Kmeans kernel successfully fitted. Sync state to my buffer.
Kmeans kernel successfully fitted. Sync state to my buffer.
`Trainer.fit` stopped: `max_epochs=1` reached.

# Step2
...
The flag in my buffer is True, try to access the kmeans kernel.
The flag in my buffer is True, try to access the kmeans kernel.
...
[rank0]: RuntimeError: load_state_dict was NOT called (my issue), and kmeans kernel is not fitted

[5o1]

I forgot @.

@SkafteNicki

[bhimrazy]

Hi @5o1 — I believe the behavior you’re seeing is due to how PyTorch handles checkpoint loading internally.

When Lightning restores from a checkpoint , it calls model.load_state_dict(...) on the top-level LightningModule. PyTorch then recursively loads parameters and buffers into each submodule by directly calling the protected method _load_from_state_dict

https://github.com/pytorch/pytorch/blob/eaa5d9d3d3dc642832b269b184f0c3ab8c990274/torch/nn/modules/module.py#L2341

— it does not seem to call each submodule’s public load_state_dict.

So probably because of this, any logic you put in MynnModule.load_state_dict will not run unless you explicitly invoke it yourself from the parent.

To ensure your sync logic runs during recursive loading, you can override _load_from_state_dict in your submodule instead. For example:

def _load_from_state_dict(
    self, state_dict, prefix, local_metadata, strict,
    missing_keys, unexpected_keys, error_msgs
):
    # Let PyTorch handle loading first
    super()._load_from_state_dict(
        state_dict, prefix, local_metadata, strict,
        missing_keys, unexpected_keys, error_msgs
    )

    self.was_load_state_dict_called = True
    print("Processing _load_from_state_dict ...")
   
  ...

This way, your post-load sync logic will always run regardless of whether you’re loading under single process or DDP.

Let me know this helps.

cc: @SkafteNicki

[SkafteNicki]

@bhimrazy are completely right and just to add, this is the relevant place in Pytorch code where the calling of child modules _load_from_state_dict method is called
https://github.com/pytorch/pytorch/blob/ba56102387ef21a3b04b357e5b183d48f0afefc7/torch/nn/modules/module.py#L2575-L2592

[bhimrazy]

This issue appears to be resolved, so we are closing it. Please feel free to reopen if any related concerns arise or additional questions come up.