environment variable WORLD_SIZE is incorrectly set to 1 after trainer.fit is done

### Bug description

Hello, so far I was relying on environment variables to figure out rank and world size in my DDP CUDA environment.

However if I run the attached script it seems the WORLD_SIZE variable is somehow unset by the trainer.

Relevant parts of the log (`#` are my comments)

```
# create trainer
LOCAL_RANK= 0 WORLD_SIZE=1 trainer.global_rank=0 trainer.world_size=8 --- Before fit
# dist is initialized
LOCAL_RANK= 7 WORLD_SIZE=8 trainer.global_rank=7 trainer.world_size=8 --- Before fit
LOCAL_RANK= 3 WORLD_SIZE=8 trainer.global_rank=3 trainer.world_size=8 --- Before fit
LOCAL_RANK= 5 WORLD_SIZE=8 trainer.global_rank=5 trainer.world_size=8 --- Before fit
LOCAL_RANK= 6 WORLD_SIZE=8 trainer.global_rank=6 trainer.world_size=8 --- Before fit
LOCAL_RANK= 2 WORLD_SIZE=8 trainer.global_rank=2 trainer.world_size=8 --- Before fit
LOCAL_RANK= 1 WORLD_SIZE=8 trainer.global_rank=1 trainer.world_size=8 --- Before fit
LOCAL_RANK= 4 WORLD_SIZE=8 trainer.global_rank=4 trainer.world_size=8 --- Before fit
Initializing distributed: GLOBAL_RANK: 7, MEMBER: 8/8
Initializing distributed: GLOBAL_RANK: 5, MEMBER: 6/8
Initializing distributed: GLOBAL_RANK: 3, MEMBER: 4/8
Initializing distributed: GLOBAL_RANK: 6, MEMBER: 7/8
Initializing distributed: GLOBAL_RANK: 4, MEMBER: 5/8
Initializing distributed: GLOBAL_RANK: 1, MEMBER: 2/8
Initializing distributed: GLOBAL_RANK: 2, MEMBER: 3/8
----------------------------------------------------------------------------------------------------
distributed_backend=nccl
All distributed processes registered. Starting with 8 processes
----------------------------------------------------------------------------------------------------
# trains for 2 epochs
# afterwards WORLD_SIZE is 1
LOCAL_RANK= 3 WORLD_SIZE=1 trainer.global_rank=3 trainer.world_size=8 --- After fit
LOCAL_RANK= 5 WORLD_SIZE=1 trainer.global_rank=5 trainer.world_size=8 --- After fit
LOCAL_RANK= 4 WORLD_SIZE=1 trainer.global_rank=4 trainer.world_size=8 --- After fit
LOCAL_RANK= 2 WORLD_SIZE=1 trainer.global_rank=2 trainer.world_size=8 --- After fit
LOCAL_RANK= 7 WORLD_SIZE=1 trainer.global_rank=7 trainer.world_size=8 --- After fit
LOCAL_RANK= 0 WORLD_SIZE=1 trainer.global_rank=0 trainer.world_size=8 --- After fit
LOCAL_RANK= 6 WORLD_SIZE=1 trainer.global_rank=6 trainer.world_size=8 --- After fit
LOCAL_RANK= 1 WORLD_SIZE=1 trainer.global_rank=1 trainer.world_size=8 --- After fit
LOCAL_RANK: 3 - CUDA_VISIBLE_DEVICES: [0,1,2,3,4,5,6,7]
LOCAL_RANK: 5 - CUDA_VISIBLE_DEVICES: [0,1,2,3,4,5,6,7]
LOCAL_RANK: 7 - CUDA_VISIBLE_DEVICES: [0,1,2,3,4,5,6,7]
LOCAL_RANK: 6 - CUDA_VISIBLE_DEVICES: [0,1,2,3,4,5,6,7]
LOCAL_RANK: 4 - CUDA_VISIBLE_DEVICES: [0,1,2,3,4,5,6,7]
LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0,1,2,3,4,5,6,7]
LOCAL_RANK: 1 - CUDA_VISIBLE_DEVICES: [0,1,2,3,4,5,6,7]
LOCAL_RANK: 2 - CUDA_VISIBLE_DEVICES: [0,1,2,3,4,5,6,7]
# runs the test
LOCAL_RANK= 7 WORLD_SIZE=1 trainer.global_rank=7 trainer.world_size=8 --- After test
LOCAL_RANK= 4 WORLD_SIZE=1 trainer.global_rank=4 trainer.world_size=8 --- After test
LOCAL_RANK= 2 WORLD_SIZE=1 trainer.global_rank=2 trainer.world_size=8 --- After test
LOCAL_RANK= 0 WORLD_SIZE=1 trainer.global_rank=0 trainer.world_size=8 --- After test
LOCAL_RANK= 5 WORLD_SIZE=1 trainer.global_rank=5 trainer.world_size=8 --- After test
LOCAL_RANK= 3 WORLD_SIZE=1 trainer.global_rank=3 trainer.world_size=8 --- After test
LOCAL_RANK= 1 WORLD_SIZE=1 trainer.global_rank=1 trainer.world_size=8 --- After test
LOCAL_RANK= 6 WORLD_SIZE=1 trainer.global_rank=6 trainer.world_size=8 --- After test
```

I understand that the WORLD_SIZE is 1 at the beginning (dist is not initialized yet).

However I think it's wrong to set the WORLD_SIZE back to 1 after fit has ended.

The advantage of env variables is that they don't depend on passing the refererence to the trainer around. I know that they can be different depending on the job environment (slurm etc.)

If using env variables here is bad practice / not supposed to work, maybe consider updating the documentation to warn users from using them.

Best,

Simon

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

v2.4

### How to reproduce the bug

```python
conda create -n litissue python=3.10 -y
conda activate litissue
pip install lightning
pip install torchvision

run below script on >= 2 CUDA GPUs

```python
import os

import lightning as lit
import torch
import torch.nn as nn
import torch.nn.functional as F
import torchvision
from lightning.pytorch.callbacks import ModelCheckpoint, LearningRateMonitor
from lightning.pytorch.loggers import CSVLogger
from torch.optim.lr_scheduler import OneCycleLR
from torch.utils.data import DataLoader, random_split
from torchmetrics.functional import accuracy
from torchvision.datasets import CIFAR10

def main():
lit.seed_everything(7)
PATH_DATASETS = os.environ.get("PATH_DATASETS", "./temp_data")
BATCH_SIZE = 64
NUM_WORKERS = 4
NUM_GPUS = torch.cuda.device_count()
assert NUM_GPUS > 1, f"Need at least 2 GPUs, got {NUM_GPUS}"

cifar10_normalization = torchvision.transforms.Normalize(
mean=[x / 255.0 for x in [125.3, 123.0, 113.9]],
std=[x / 255.0 for x in [63.0, 62.1, 66.7]],
)

train_transforms = torchvision.transforms.Compose(
[
torchvision.transforms.RandomCrop(32, padding=4),
torchvision.transforms.RandomHorizontalFlip(),
torchvision.transforms.ToTensor(),
cifar10_normalization,
]
)
test_transforms = torchvision.transforms.Compose(
[
torchvision.transforms.ToTensor(),
cifar10_normalization,
]
)

dataset_train = CIFAR10(PATH_DATASETS, train=True, download=True, transform=train_transforms)
dataset_val = CIFAR10(PATH_DATASETS, train=True, download=True, transform=test_transforms)
dataset_train = split_dataset(dataset_train)
dataset_val = split_dataset(dataset_val, train=False)
dataset_test = CIFAR10(PATH_DATASETS, train=False, download=True, transform=test_transforms)

train_dataloader = DataLoader(
dataset_train, batch_size=BATCH_SIZE, shuffle=True, num_workers=NUM_WORKERS
)
val_dataloader = DataLoader(
dataset_val, batch_size=BATCH_SIZE, shuffle=False, num_workers=NUM_WORKERS
)
test_dataloader = DataLoader(
dataset_test, batch_size=BATCH_SIZE, shuffle=False, num_workers=NUM_WORKERS
)

model = LitResnet(lr=0.05, batch_size=BATCH_SIZE)

trainer = lit.Trainer(
max_epochs=2,
accelerator="cuda",
devices=NUM_GPUS,
logger=CSVLogger(save_dir="logs/"),
callbacks=[
LearningRateMonitor(logging_interval="step"),
ModelCheckpoint(
dirpath="temp_checkpoints",
filename="{epoch}-{step}",
save_last=True,
every_n_epochs=1,
enable_version_counter=False,
),
],
)

print_with_rank(trainer, f"Before fit")
trainer.fit(model, train_dataloaders=train_dataloader, val_dataloaders=val_dataloader)
print_with_rank(trainer, f"After fit")
trainer.test(model, dataloaders=test_dataloader)
print_with_rank(trainer, f"After test")

def split_dataset(dataset, val_split=0.2, train=True):
"""Splits the dataset into train and validation set."""
len_dataset = len(dataset)
splits = get_splits(len_dataset, val_split)
dataset_train, dataset_val = random_split(
dataset, splits, generator=torch.Generator().manual_seed(42)
)

if train:
return dataset_train
return dataset_val

def get_splits(len_dataset, val_split):
"""Computes split lengths for train and validation set."""
if isinstance(val_split, int):
train_len = len_dataset - val_split
splits = [train_len, val_split]
elif isinstance(val_split, float):
val_len = int(val_split * len_dataset)
train_len = len_dataset - val_len
splits = [train_len, val_len]
else:
raise ValueError(f"Unsupported type {type(val_split)}")

return splits

def create_model():
model = torchvision.models.resnet18(pretrained=False, num_classes=10)
model.conv1 = nn.Conv2d(3, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
model.maxpool = nn.Identity()
return model

class LitResnet(lit.LightningModule):
def __init__(self, lr=0.05, batch_size=64):
super().__init__()
self.save_hyperparameters()
self.model = create_model()
self.batch_size = batch_size

def setup(self, stage):
print_with_rank(self.trainer, "Model setup {stage=}")

def forward(self, x):
out = self.model(x)
return F.log_softmax(out, dim=1)

def training_step(self, batch, batch_idx):
x, y = batch
logits = self(x)
loss = F.nll_loss(logits, y)
self.log("train_loss", loss)
return loss

def evaluate(self, batch, stage=None):
x, y = batch
logits = self(x)
loss = F.nll_loss(logits, y)
preds = torch.argmax(logits, dim=1)
acc = accuracy(preds, y, task="multiclass", num_classes=10)

if stage:
self.log(f"{stage}_loss", loss, prog_bar=True)
self.log(f"{stage}_acc", acc, prog_bar=True)

def validation_step(self, batch, batch_idx):
self.evaluate(batch, "val")

def test_step(self, batch, batch_idx):
self.evaluate(batch, "test")

def configure_optimizers(self):
optimizer = torch.optim.SGD(
self.parameters(),
lr=self.hparams.lr,
momentum=0.9,
weight_decay=5e-4,
)
steps_per_epoch = 45000 // self.batch_size
scheduler_dict = {
"scheduler": OneCycleLR(
optimizer,
0.1,
epochs=self.trainer.max_epochs,
steps_per_epoch=steps_per_epoch,
),
"interval": "step",
}
return {"optimizer": optimizer, "lr_scheduler": scheduler_dict}

def get_rank() -> int:
if "RANK" in os.environ:
rank = int(os.environ["RANK"])
else:
rank = int(os.environ.get("LOCAL_RANK", 0))
return rank

def get_world_size() -> int:
return int(os.environ.get("WORLD_SIZE", 1))

def print_with_rank(trainer, *args, **kwargs):
rank = get_rank()
world_size = get_world_size()
print(
f"LOCAL_RANK={rank:>2d} WORLD_SIZE={world_size} "
f"trainer.global_rank={trainer.global_rank} trainer.world_size={trainer.world_size} ---",
*args,
**kwargs,
)

if __name__ == "__main__":
main()
```
```

### Error messages and logs

see post

### Environment

Current environment

* CUDA:
- GPU:
- NVIDIA GeForce RTX 3090
- NVIDIA GeForce RTX 3090
- NVIDIA GeForce RTX 3090
- NVIDIA GeForce RTX 3090
- NVIDIA GeForce RTX 3090
- NVIDIA GeForce RTX 3090
- NVIDIA GeForce RTX 3090
- NVIDIA GeForce RTX 3090
- available: True
- version: 12.1
* Lightning:
- lightning: 2.4.0
- lightning-utilities: 0.11.6
- pytorch-lightning: 2.4.0
- torch: 2.4.0
- torchmetrics: 1.4.1
- torchvision: 0.19.0
* Packages:
- aiohappyeyeballs: 2.4.0
- aiohttp: 3.10.5
- aiosignal: 1.3.1
- async-timeout: 4.0.3
- attrs: 24.2.0
- autocommand: 2.2.2
- backports.tarfile: 1.2.0
- filelock: 3.15.4
- frozenlist: 1.4.1
- fsspec: 2024.6.1
- idna: 3.8
- importlib-metadata: 8.0.0
- importlib-resources: 6.4.0
- inflect: 7.3.1
- jaraco.context: 5.3.0
- jaraco.functools: 4.0.1
- jaraco.text: 3.12.1
- jinja2: 3.1.4
- lightning: 2.4.0
- lightning-utilities: 0.11.6 [0/8069]
- markupsafe: 2.1.5
- more-itertools: 10.3.0
- mpmath: 1.3.0
- multidict: 6.0.5
- networkx: 3.3
- numpy: 2.1.0
- nvidia-cublas-cu12: 12.1.3.1
- nvidia-cuda-cupti-cu12: 12.1.105
- nvidia-cuda-nvrtc-cu12: 12.1.105
- nvidia-cuda-runtime-cu12: 12.1.105
- nvidia-cudnn-cu12: 9.1.0.70
- nvidia-cufft-cu12: 11.0.2.54
- nvidia-curand-cu12: 10.3.2.106
- nvidia-cusolver-cu12: 11.4.5.107
- nvidia-cusparse-cu12: 12.1.0.106
- nvidia-nccl-cu12: 2.20.5
- nvidia-nvjitlink-cu12: 12.6.20
- nvidia-nvtx-cu12: 12.1.105
- ordered-set: 4.1.0
- packaging: 24.1
- pillow: 10.4.0
- pip: 24.2
- platformdirs: 4.2.2
- pytorch-lightning: 2.4.0
- pyyaml: 6.0.2
- setuptools: 72.1.0
- sympy: 1.13.2
- tomli: 2.0.1
- torch: 2.4.0
- torchmetrics: 1.4.1
- torchvision: 0.19.0
- tqdm: 4.66.5
- triton: 3.0.0
- typeguard: 4.3.0
- typing-extensions: 4.12.2
- wheel: 0.43.0
- yarl: 1.9.4
- zipp: 3.19.2
* System:
- OS: Linux
- architecture:
- 64bit
- ELF
- processor: x86_64
- python: 3.10.14
- release: 6.5.0-41-generic
- version: #41~22.04.2-Ubuntu SMP PREEMPT_DYNAMIC Mon Jun 3 11:32:55 UTC 2

### More info

_No response_

cc @lantiga @justusschock @borda

Start by running the attached reproduction with multiple CUDA GPUs and compare the environment values before and after trainer.fit and trainer.test. Trace the trainer lifecycle and distributed environment handling around those entry points, then verify that WORLD_SIZE remains consistent with trainer.world_size after training without breaking the shown rank behavior.

Written by the indexing model from the issue text.