[BUG]

Describe the bug

When deepspeed was used to simulate elasticity training in two containers, it appeared that when the worker process exited, the training step of the master process did not change.

Test script
import torch
import torchvision
import torchvision.transforms as transforms
import deepspeed
import time
import os
import torch.nn as nn
from torch.distributed.elastic.utils.data import ElasticDistributedSampler
from deepspeed.accelerator import get_accelerator
class MNIST(torch.nn.Module):
def init(self):
super(MNIST, self).init()
self.conv = torch.nn.Sequential(torch.nn.Conv2d(1, 32, 3, 1, 1),
torch.nn.ReLU(),
torch.nn.Conv2d(32, 64, 3, 1, 1),
torch.nn.ReLU(),
torch.nn.MaxPool2d(2, 2))
self.dense = torch.nn.Sequential(torch.nn.Linear(14 * 14 * 64, 1024),
torch.nn.ReLU(),
torch.nn.Dropout(p=0.2),
torch.nn.Linear(1024, 10))
def forward(self, x):
x = self.conv(x)
x = x.view(-1, 14 * 14 * 64)
x = self.dense(x)
return x
def save_checkpoint(epoch, model, optimizer, path):
torch.save({
"epoch": epoch,
"model_state_dict": model.state_dict(),
"optimize_state_dict": optimizer.state_dict(),
}, path)
def load_checkpoint(path):
checkpoint = torch.load(path)
return checkpoint
if name == "main":
deepspeed.init_distributed()
transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
])
train_dataset = torchvision.datasets.MNIST(root='./data',
train=True,
transform=transforms.ToTensor(),
download=True)
train_sampler = ElasticDistributedSampler(train_dataset)
train_loader = torch.utils.data.DataLoader(train_dataset,
batch_size=100,
num_workers=2,
pin_memory=True,
sampler=train_sampler,)
net =MNIST()
ds_config = {
"train_batch_size": 100,
"optimizer": {
"type": "Adam",
"params": {
"lr": 0.0001,
"betas": [
0.8,
0.999
],
"eps": 1e-8,
"weight_decay": 3e-7
}
},
"elasticity": {
"enabled": True,
"max_train_batch_size": 200,
"micro_batch_sizes": [100,200],
"min_gpus": 1,
"max_gpus": 2,
"min_time": 0,
"version": 0.1,
"ignore_non_elastic_batch_info": True,
}
}
model_engine, optimizer, trainloader, __ = deepspeed.initialize(model=net, model_parameters=net.parameters(), training_data=train_dataset, config=ds_config)
local_device = get_accelerator().device_name(model_engine.local_rank)
local_rank = model_engine.local_rank
criterion = nn.CrossEntropyLoss()
ckp_path = "/workspace/checkpoint.pt"
first_epoch = -1
max_epoch = 10
if os.path.exists(ckp_path):
print(f"load checkpoint from {ckp_path}")
checkpoint = load_checkpoint(ckp_path)
net.load_state_dict(checkpoint["model_state_dict"])
optimizer.load_state_dict(checkpoint["optimize_state_dict"])
first_epoch = checkpoint["epoch"]
epochs = 10
for epoch in range(first_epoch + 1, max_epoch):
# train
start = time.time()
number = 0
sum_loss = 0.0
for inputs, lables in train_loader:
optimizer.zero_grad()
outputs = net(inputs.to(local_rank))
lables = lables.to(local_rank)
loss = criterion(outputs, lables)
loss.backward()
optimizer.step()
sum_loss += loss.data
end = time.time()
number += 1
if number % 100 == 0:
print('epoch: [%d,%d], step: [%d,%d] loss:%.06f, step time:%.06f' %
(epoch + 1, 10, number+1, len(train_loader), sum_loss / len(train_loader), (end-start)/number))
save_checkpoint(epoch, net, optimizer, ckp_path)
print('Finished Training')

Test result

master logs

root@elastic-job-master-0:/app# deepspeed --hostfile hostfile --num_nodes 1 --min_elastic_nodes 1 --max_elastic_nodes 2 --num_gpus 1 --force_multi --elastic_training --master_port 49091 --master_addr 100.73.234.153 train.py
...
elastic-job-master-0: epoch: [1,10], step: [101,600] loss:0.152321, step time:0.040384
elastic-job-master-0: epoch: [1,10], step: [201,600] loss:0.207724, step time:0.023482
elastic-job-master-0: epoch: [1,10], step: [301,600] loss:0.251543, step time:0.017879
elastic-job-master-0: epoch: [1,10], step: [401,600] loss:0.286103, step time:0.014942
elastic-job-master-0: epoch: [1,10], step: [501,600] loss:0.314613, step time:0.013334
elastic-job-master-0: epoch: [1,10], step: [601,600] loss:0.338988, step time:0.012194
elastic-job-master-0: epoch: [2,10], step: [101,600] loss:0.020816, step time:0.008608
elastic-job-master-0: epoch: [2,10], step: [201,600] loss:0.037383, step time:0.008008
elastic-job-master-0: epoch: [2,10], step: [301,600] loss:0.054437, step time:0.007318
elastic-job-master-0: epoch: [2,10], step: [401,600] loss:0.068281, step time:0.007157
elastic-job-master-0: epoch: [2,10], step: [501,600] loss:0.081604, step time:0.006894
elastic-job-master-0: epoch: [2,10], step: [601,600] loss:0.094132, step time:0.006663
elastic-job-master-0: epoch: [3,10], step: [101,600] loss:0.011768, step time:0.008324
elastic-job-master-0: epoch: [3,10], step: [201,600] loss:0.021150, step time:0.007156
elastic-job-master-0: epoch: [3,10], step: [301,600] loss:0.032299, step time:0.006993
elastic-job-master-0: epoch: [3,10], step: [401,600] loss:0.040990, step time:0.006901
elastic-job-master-0: epoch: [3,10], step: [501,600] loss:0.049896, step time:0.006788
elastic-job-master-0: epoch: [3,10], step: [601,600] loss:0.058711, step time:0.006700
elastic-job-master-0: epoch: [4,10], step: [101,600] loss:0.008358, step time:0.007746
elastic-job-master-0: epoch: [4,10], step: [201,600] loss:0.015032, step time:0.006609
elastic-job-master-0: epoch: [4,10], step: [301,600] loss:0.023299, step time:0.006482
elastic-job-master-0: epoch: [4,10], step: [401,600] loss:0.029358, step time:0.006370
elastic-job-master-0: epoch: [4,10], step: [501,600] loss:0.036371, step time:0.006321
elastic-job-master-0: epoch: [4,10], step: [601,600] loss:0.043281, step time:0.006334
elastic-job-master-0: epoch: [5,10], step: [101,600] loss:0.006229, step time:0.008643
elastic-job-master-0: epoch: [5,10], step: [201,600] loss:0.011332, step time:0.008003
elastic-job-master-0: epoch: [5,10], step: [301,600] loss:0.018273, step time:0.007307
elastic-job-master-0: Sending process 3784 closing signal SIGTERM
elastic-job-master-0: [2023-12-07 17:36:28,348] [INFO] [real_accelerator.py:158:get_accelerator] Setting ds_accelerator to cuda (auto detect)
...
elastic-job-master-0: load checkpoint from /app/checkpoint.pt
elastic-job-master-0: epoch: [5,10], step: [101,300] loss:0.012144, step time:0.040376
elastic-job-master-0: epoch: [5,10], step: [201,300] loss:0.024505, step time:0.023769
elastic-job-master-0: epoch: [5,10], step: [301,300] loss:0.035717, step time:0.018047
elastic-job-master-0: epoch: [6,10], step: [101,300] loss:0.009512, step time:0.008259
elastic-job-master-0: epoch: [6,10], step: [201,300] loss:0.019266, step time:0.007584
elastic-job-master-0: epoch: [6,10], step: [301,300] loss:0.028346, step time:0.006997
elastic-job-master-0: epoch: [7,10], step: [101,300] loss:0.007693, step time:0.007530
elastic-job-master-0: epoch: [7,10], step: [201,300] loss:0.015422, step time:0.006844
elastic-job-master-0: epoch: [7,10], step: [301,300] loss:0.022689, step time:0.006562
elastic-job-master-0: epoch: [8,10], step: [101,300] loss:0.005978, step time:0.008205
elastic-job-master-0: epoch: [8,10], step: [201,300] loss:0.012152, step time:0.007881
elastic-job-master-0: epoch: [8,10], step: [301,300] loss:0.018234, step time:0.007441
elastic-job-master-0: epoch: [9,10], step: [101,300] loss:0.005155, step time:0.008015
elastic-job-master-0: epoch: [9,10], step: [201,300] loss:0.010217, step time:0.007272
elastic-job-master-0: epoch: [9,10], step: [301,300] loss:0.015026, step time:0.006726
elastic-job-master-0: epoch: [10,10], step: [101,300] loss:0.004170, step time:0.007917
elastic-job-master-0: epoch: [10,10], step: [201,300] loss:0.008121, step time:0.007109
elastic-job-master-0: epoch: [10,10], step: [301,300] loss:0.011786, step time:0.006650
elastic-job-master-0: Finished Training

worekr logs

root@elastic-job-worker-0:/app# deepspeed --hostfile hostfile --num_nodes 1 --min_elastic_nodes 1 --max_elastic_nodes 2 --num_gpus 1 --force_multi --elastic_training --master_port 49091 --master_addr 100.73.234.153 train.py
..
elastic-job-worker-0: Time to load fused_adam op: 0.35888099670410156 seconds
elastic-job-worker-0: epoch: [1,10], step: [101,300] loss:0.320535, step time:0.042084
elastic-job-worker-0: epoch: [1,10], step: [201,300] loss:0.431218, step time:0.024617
elastic-job-worker-0: epoch: [1,10], step: [301,300] loss:0.523071, step time:0.018793
elastic-job-worker-0: epoch: [2,10], step: [101,300] loss:0.073453, step time:0.008842
elastic-job-worker-0: epoch: [2,10], step: [201,300] loss:0.132337, step time:0.008160
elastic-job-worker-0: epoch: [2,10], step: [301,300] loss:0.184430, step time:0.007373
elastic-job-worker-0: epoch: [3,10], step: [101,300] loss:0.043537, step time:0.007847
elastic-job-worker-0: epoch: [3,10], step: [201,300] loss:0.080996, step time:0.007020
elastic-job-worker-0: epoch: [3,10], step: [301,300] loss:0.116353, step time:0.006659
elastic-job-worker-0: epoch: [4,10], step: [101,300] loss:0.029314, step time:0.008321
elastic-job-worker-0: epoch: [4,10], step: [201,300] loss:0.055802, step time:0.007983
^Z
[1]+ Stopped deepspeed --hostfile hostfile --num_nodes 1 --min_elastic_nodes 1 --max_elastic_nodes 2 --

Is it a training script test problem?

Start with the train.py reproducer and the DeepSpeed elastic-training command shown in the issue. Run the two-container setup, observe the master and worker step counts before and after the worker exits, and check whether the checkpoint reload changes the expected loader length. Done means determining whether the behavior is caused by the training script or by DeepSpeed and documenting a reproducible result.