I’m trying to change from plain DP to DDP. My code works, but is very slow.
I narrowed down the problem to a weird behaviour: the dataset, which I pass as an argument in mp.spawn
, is being pickled and unpickled every time I create new iterator from loader.
Why is this the case? Surely all arguments should be pickled and unpickled once when processes are created.
Rough code:
def train(rank, dataset):
setup_ddp(rank)
model = create_model()
loader = create_loader(dataset)
train_iter = iter(loader)
for i in range(NUM_ITERS):
try:
batch, targets = next(train_iter)
except StopIteration:
train_iter = iter(loader)
# The dataset is being pickled and unpickled every time next is called on new iterator
batch, targets = next(train_iter)
step(model, batch, targets)
cleanup_ddp(rank)
dataset = create_dataset()
mp.spawn(train, args=(dataset,), nprocs=WORLD_SIZE)