Package: dask-pytorch-ddp
dask-pytorch-ddpdask-pytorch-ddpis a Python package that makes it easy to train PyTorch models on Dask clusters using distributed data parallel.  The intended scope of the project isbootstrapping PyTorch workers on top of a Dask clusterUsing distributed data stores (e.g., S3) as normal PyTorch datasetsmechanisms for tracking and logging intermediate results, training statistics, and checkpoints.At this point, this library and examples provided are tailored to computer vision tasks, but this library is intended to be useful for any sort of PyTorch tasks. The only thing really specific to image processing is theS3ImageFolderdataset class. Implementing a PyTorch dataset (assuming map style random access) outside of images currently requires implementing__getitem__(self, idx: int):and__len__(self):We plan to add more varied examples for other use cases in the future, and welcome PRs extending functionality.Typical non-dask workflowA typical example of non-dask PyTorch usage is as follows:Loading DataCreate an dataset (ImageFolder), and wrap it in aDataLoadertransform=transforms.Compose([transforms.Resize(256),transforms.CenterCrop(250),transforms.ToTensor()])whole_dataset=ImageFolder(path,transform=transform)batch_size=100num_workers=64indices=list(range(len(data)))np.random.shuffle(indices)train_idx=indices[:num]test_idx=indices[num:num+num]train_sampler=SubsetRandomSampler(train_idx)train_loader=DataLoader(data,sampler=train_sampler,batch_size=batch_size,num_workers=num_workers)Training a ModelLoop over the dataset, and train the model by stepping the optimizerdevice=torch.device(0)net=models.resnet18(pretrained=False)model=net.to(device)device_ids=[0]criterion=nn.CrossEntropyLoss().cuda()lr=0.001optimizer=optim.SGD(model.parameters(),lr=lr,momentum=0.9)count=0forepochinrange(n_epochs):model.train()# Set model to training modeforinputs,labelsintrain_loader:inputs=inputs.to(device)labels=labels.to(device)outputs=model(inputs)_,preds=torch.max(outputs,1)loss=criterion(outputs,labels)# zero the parameter gradientsoptimizer.zero_grad()loss.backward()optimizer.step()count+=1Now on DaskWith dask_pytorch_ddp and PyTorch Distributed Data Parallel, we can train on multiple workers as follows:Loading DataLoad the dataset from S3, and explicitly set the multiprocessing context (Dask defaults to spawn, but pytorch is generally configured to use fork)fromdask_pytorch_ddp.dataimportS3ImageFolderwhole_dataset=S3ImageFolder(bucket,prefix,transform=transform)train_loader=torch.utils.data.DataLoader(whole_dataset,sampler=train_sampler,batch_size=batch_size,num_workers=num_workers,multiprocessing_context=mp.get_context('fork'))Training in ParallelWrap the training loop in a function (and add metrics logging.  Not necessary, but very useful).  Convert the model into a PyTorch Distributed Data Parallel (DDP) model which knows how to sync gradients together across workers.importuuidimportpickleimportloggingimportjsonkey=uuid.uuid4().hexrh=DaskResultsHandler(key)defrun_transfer_learning(bucket,prefix,samplesize,n_epochs,batch_size,num_workers,train_sampler):worker_rank=int(dist.get_rank())device=torch.device(0)net=models.resnet18(pretrained=False)model=net.to(device)model=DDP(model,device_ids=[0])criterion=nn.CrossEntropyLoss().cuda()lr=0.001optimizer=optim.SGD(model.parameters(),lr=lr,momentum=0.9)whole_dataset=S3ImageFolder(bucket,prefix,transform=transform)train_loader=torch.utils.data.DataLoader(whole_dataset,sampler=train_sampler,batch_size=batch_size,num_workers=num_workers,multiprocessing_context=mp.get_context('fork'))count=0forepochinrange(n_epochs):# Each epoch has a training and validation phasemodel.train()# Set model to training modeforinputs,labelsintrain_loader:dt=datetime.datetime.now().isoformat()inputs=inputs.to(device)labels=labels.to(device)outputs=model(inputs)_,preds=torch.max(outputs,1)loss=criterion(outputs,labels)# zero the parameter gradientsoptimizer.zero_grad()loss.backward()optimizer.step()count+=1# statisticsrh.submit_result(f"worker/{worker_rank}/data-{dt}.json",json.dumps({'loss':loss.item(),'epoch':epoch,'count':count,'worker':worker_rank}))if(count%100)==0andworker_rank==0:rh.submit_result(f"checkpoint-{dt}.pkl",pickle.dumps(model.state_dict()))How does it work?dask-pytorch-ddpis largely a wrapper around existingpytorchfunctionality.pytorch.distributedprovides infrastructure forDistributed Data Parallel(DDP).In DDP, you create N workers, and the 0th worker is the "master", and coordinates the synchronization of buffers and gradients.  In SGD, gradients are normally averaged between all data points in a batch.  By running batches on multiple workers, and averaging the gradients, DDP enables you to run SGD with a much bigger batch size(N * batch_size)dask-pytorch-ddpsets some environment variables to configure the "master" host and port, and then callsinit_process_groupbefore training, and callsdestroy_process_groupafter training.  This is the same process normally done manually by the data scientist.Multi GPU machinesdask_cuda_workerautomatically rotatesCUDA_VISIBLE_DEVICESfor each worker it creates (typically one per GPU).  As a result, your PyTorch code should always start with the 0th GPU.For example, if I have an 8 GPU machine, the 3rd worker will haveCUDA_VISIBLE_DEVICESset to2,3,4,5,6,7,0,1.  On that worker, if I calltorch.device(0), I will get GPU 2.What else?dask-pytorch-ddpalso implements an S3 basedImageFolder.  More distributed friendly datasets are planned.dask-pytorch-ddpalso implements a basic results aggregation framework so that it is easy to collect training metrics across different workers.  Currently, onlyDaskResultsHandlerwhich leveragesDask pub-sub communication protocolsis implemented, but an S3 based result handler is planned.Some NotesDask generally spawns processes.  PyTorch generally forks.  When using a multiprocessing enabled data loader, it is a good idea to pass theForkmultiprocessing context to force the use of Forking in the data loader.Some Dask deployments do not permit spawning processes.  To override this, you can change thedistributed.worker.daemonsetting.Environment variables are a convenient way to do this:DASK_DISTRIBUTED__WORKER__DAEMON=False