Loading a model changes pytorch random state

Nobody has claimed this yet.

### Bug description

Loading a model from a checkpoint `model = TestLitModel.load_from_checkpoint(ckpt_path)` changes the pytorch random state as I guess the model is initialized and then the weights are restored.
To me at least this was not clear from just calling a classmethod.

One option would be to restore the random state to what it was before the load, the other to add a hint in the documentation, if I'm not the only one for whom it is not intuitively clear.

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

master

### How to reproduce the bug

```python
import time
from typing import Any

import lightning as L
import torch
from lightning.pytorch.callbacks import ModelCheckpoint, OnExceptionCheckpoint
from lightning.pytorch.loggers import WandbLogger
from lightning.pytorch.utilities import grad_norm
from lightning.pytorch.utilities.types import STEP_OUTPUT
from torch import nn
from torch.utils.data import Dataset, DataLoader, Sampler

from src.main.ml.data.data_augmentation.helpers.random_numbers import create_rng_from_string

class TestModule(nn.Module):
def __init__(self, in_dim=512, out_dim=16):
super().__init__()
self.in_dim = in_dim
self.out_dim = out_dim
self.simple_layer = nn.Linear(self.in_dim, self.out_dim, bias=True)

def forward(self, input):
return self.simple_layer(input)

class TestBatchSampler(Sampler):
def __init__(self, step=0):
super().__init__()
self.step = step

def __len__(self) -> int:
return 1e100
# return len(self.train_allfiles)

def __iter__(self): # -> Iterator[int]:
return self

def __next__(self): # -> Iterator[int]:
return_value = self.step
self.step += 1
return [return_value]

class TestDataset(Dataset):
def __init__(self, in_dim):
super().__init__()
self.in_dim = in_dim
self.total_len = 512

def __len__(self):
return 1

def __getitem__(self, idx):
print(idx)
rng = create_rng_from_string(
str(idx) + "_"
+ "random_choice_sampler")
return torch.tensor(rng.random(self.in_dim), dtype=torch.float32)
# return torch.randn(self.in_dim)

class TestDataModule(L.LightningDataModule):
def __init__(self, start_step=0):
super().__init__()
self.in_dim = 512
self.val_batch_size = 1
self.start_step = start_step

def train_dataloader(self):
train_ds = TestDataset(self.in_dim)
train_dl = DataLoader(train_ds, batch_sampler=TestBatchSampler(step=self.start_step), num_workers=4,
shuffle=False)
return train_dl

def val_dataloader(self):
val_ds = TestDataset(self.in_dim)
val_dl = DataLoader(val_ds, batch_size=self.val_batch_size, num_workers=4, shuffle=False)
return val_dl

class TestLitModel(L.LightningModule):
def __init__(self):
super().__init__()
self.test_module_obj = TestModule(in_dim=512, out_dim=16)
self.automatic_optimization = False

def on_train_batch_end(self, outputs: STEP_OUTPUT, batch: Any, batch_idx: int) -> None:
print(f"train_batch ended:{batch_idx}")

def on_save_checkpoint(self, checkpoint):

# my own fix for reproducible checkpointing
# checkpoint['loops']['fit_loop']['epoch_loop.batch_progress']['total']['completed'] = \
# checkpoint['loops']['fit_loop']['epoch_loop.batch_progress']['total']['processed']
# checkpoint['loops']['fit_loop']['epoch_loop.batch_progress']['current']['completed'] = \
# checkpoint['loops']['fit_loop']['epoch_loop.batch_progress']['current']['processed']

checkpoint['loops']['fit_loop']['epoch_loop.batch_progress']['total']['completed'] = \
checkpoint['loops']['fit_loop']['epoch_loop.batch_progress']['total']['processed']
checkpoint['loops']['fit_loop']['epoch_loop.batch_progress']['current']['completed'] = \
checkpoint['loops']['fit_loop']['epoch_loop.batch_progress']['current']['processed']
checkpoint['loops']['fit_loop']['epoch_loop.state_dict']['_batches_that_stepped'] = \
checkpoint['loops']['fit_loop']['epoch_loop.batch_progress']['current']['processed']
print(f"creating checkpoint")

def validation_step(self, *args: Any, **kwargs: Any) -> STEP_OUTPUT:
print(f"validation step called")
return torch.tensor(1.0)

def training_step(self, batch, batch_idx):
print(f"batch_idx: {batch_idx}")

time.sleep(0.25)
optimizer = self.optimizers()

output = self.test_module_obj(batch)

loss = output.sum()

self.manual_backward(loss)

optimizer.step()

self.log_dict({"loss": loss, "batch_idx": batch_idx})
norms = grad_norm(self.test_module_obj, norm_type=2)
self.log_dict(norms)

def configure_optimizers(self):
optimizer = torch.optim.AdamW(
self.test_module_obj.parameters()
)
return optimizer

if __name__ == '__main__':
test_data_loader = TestDataModule()
test_lit_model = TestLitModel()

checkpoint_dir = 'a_test_logs_2'

checkpoint_callback = ModelCheckpoint(
dirpath=checkpoint_dir,
every_n_train_steps=10,
save_top_k=-1, )
exception_checkpoint_callback = OnExceptionCheckpoint(
dirpath=checkpoint_dir,
filename="error"
)

wandb_logger = WandbLogger(save_dir="/home/thillmann/projects/metrabs_refactoring/a_logs",
project="global_step_test",
name="global_step_test",
notes="global_step_test",
log_model="all")

trainer = L.Trainer(
log_every_n_steps=1,
callbacks=[checkpoint_callback],
max_epochs=-1,
max_steps=25,
val_check_interval=5,
logger=wandb_logger,

)

# wandb_logger.watch(model=test_lit_model.test_module_obj, log="all", log_freq=1)

###########################################################

# uncomment this to create checkpoint
# trainer.fit(test_lit_model, test_data_loader)

###########################################################

cpu_random_state = torch.get_rng_state()
print("CPU Random State:", cpu_random_state)

ckpt_path = f'{checkpoint_dir}/epoch=0-step=10.ckpt'
print("CPU Random State equal:", torch.equal(cpu_random_state, torch.get_rng_state()))

ckpt = torch.load(ckpt_path)
global_step = ckpt['global_step']

ckpt_path = f'{checkpoint_dir}/epoch=0-step=10.ckpt'
print("CPU Random State equal after pytorch load:", torch.equal(cpu_random_state, torch.get_rng_state()))

model = TestLitModel.load_from_checkpoint(ckpt_path)

ckpt_path = f'{checkpoint_dir}/epoch=0-step=10.ckpt'
cpu_random_state_4 = torch.get_rng_state()
print("CPU Random State equal after lightning load:", torch.equal(cpu_random_state, cpu_random_state_4))

test_data_loader = TestDataModule(start_step=global_step)

trainer.fit(test_lit_model,
datamodule=test_data_loader,
ckpt_path=f'{checkpoint_dir}/epoch=0-step=10.ckpt')
```

### 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):
#- 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

No source file or test is named. Start by reproducing the RNG comparison around TestLitModel.load_from_checkpoint and inspect that entry point; completion depends on deciding whether loading should preserve PyTorch RNG state or whether the behavior should be documented, with a regression check for the chosen outcome.

Written by the indexing model from the issue text.