Distributed PyTorch

Launching PyTorch Distributed Training with Dragon

PyTorch distributed training needs each worker to agree on three things before the first collective call happens: how many ranks are participating, which rank the current worker owns, and how all ranks will find the same rendezvous address. Dragon helps with that setup, but two different layers are involved.

ProcessGroup is Dragon’s general native process orchestration API. It gives you explicit control over placement and lifecycle, and it is the right fit when you are launching mixed workloads or building a lower-level orchestration layer yourself.

CollectiveGroup() is the higher-level helper used in the examples below. It builds on Dragon’s runtime services to launch one training worker per policy, provides a RankInfo object inside each worker, and supplies the rank, world size, and rendezvous information that PyTorch needs for torch.distributed.init_process_group.

If your goal is “launch one distributed training worker per GPU and initialize PyTorch correctly,” start with CollectiveGroup. Reach for ProcessGroup when you need more manual control over how those workers are created or placed.

Listing 59 Setting up NCCL backend for PyTorch Distributed Training with CollectiveGroup
 1from dragon.native.machine import System
 2from dragon.ai.collective_group import CollectiveGroup, RankInfo
 3
 4import torch
 5import torch.distributed as dist
 6
 7
 8def train():
 9    rank_info = RankInfo()
10    rank = rank_info.my_rank
11    master_addr = rank_info.master_addr
12    master_port = rank_info.master_port
13    world_size = rank_info.world_size
14
15    dist.init_process_group(
16        backend="nccl",
17        init_method=f"tcp://{master_addr}:{master_port}",
18        world_size=world_size,
19        rank=rank,
20    )
21
22    device = torch.device("cuda")  # the provided Policy already sets which GPU id to use
23    tensor = torch.ones(1, device=device) * rank
24    dist.all_reduce(tensor, op=dist.ReduceOp.SUM)
25
26    print(f"Rank {rank}: Tensor after all_reduce = {tensor.item()}")
27
28    dist.destroy_process_group()
29
30if __name__ == "__main__":
31
32    gpu_policies = System().gpu_policies()
33    pg = CollectiveGroup(
34            training_fn=train,
35            training_args=None,
36            training_kwargs=None,
37            policies=gpu_policies,
38            hide_stderr=False,
39            port=29500,
40        )
41    pg.init()
42    pg.start()
43    pg.join()
44    pg.close()

Loading Training Data with PyTorch

Distributed training requires each rank to consume a different, non-overlapping shard of the training dataset. PyTorch’s DistributedSampler handles this automatically: it partitions the dataset indices across world_size ranks so that every GPU receives a unique slice of each epoch’s data.

When using Dragon’s CollectiveGroup(), the RankInfo helper provides the rank, world size, and master address that DistributedSampler needs. No manual rank-file management is required.

Listing 60 Distributed DataLoader with DistributedSampler inside a CollectiveGroup worker
 1from dragon.ai.collective_group import CollectiveGroup, RankInfo
 2from dragon.native.machine import System
 3
 4import torch
 5import torch.distributed as dist
 6from torch.utils.data import DataLoader, TensorDataset
 7from torch.utils.data.distributed import DistributedSampler
 8
 9
10def training_fn():
11    rank_info = RankInfo()
12    rank        = rank_info.my_rank
13    world_size  = rank_info.world_size
14    master_addr = rank_info.master_addr
15    master_port = rank_info.master_port
16
17    dist.init_process_group(
18        backend="nccl",
19        init_method=f"tcp://{master_addr}:{master_port}",
20        world_size=world_size,
21        rank=rank,
22    )
23
24    # Build a synthetic dataset — replace with your real dataset
25    num_samples = 10_000
26    input_size  = 128
27    X = torch.randn(num_samples, input_size)
28    y = torch.randint(0, 10, (num_samples,))
29    dataset = TensorDataset(X, y)
30
31    # DistributedSampler ensures each rank sees a unique shard of the dataset
32    sampler = DistributedSampler(
33        dataset,
34        num_replicas=world_size,
35        rank=rank,
36        shuffle=True,
37        drop_last=True,
38    )
39    loader = DataLoader(dataset, batch_size=64, sampler=sampler, num_workers=0)
40
41    # Call sampler.set_epoch(epoch) at the start of each epoch so that shuffling
42    # differs between epochs across all ranks.
43    for epoch in range(5):
44        sampler.set_epoch(epoch)
45        for batch_X, batch_y in loader:
46            # Move data to this rank's GPU (Policy already sets GPU affinity)
47            batch_X = batch_X.to("cuda")
48            batch_y = batch_y.to("cuda")
49            # ... forward pass, loss, backward, optimizer step ...
50
51        if rank == 0:
52            print(f"Epoch {epoch + 1} complete", flush=True)
53
54    dist.destroy_process_group()
55
56
57if __name__ == "__main__":
58    gpu_policies = System().gpu_policies()
59    pg = CollectiveGroup(
60        training_fn=training_fn,
61        training_args=None,
62        training_kwargs=None,
63        policies=gpu_policies,
64        hide_stderr=False,
65        port=29500,
66    )
67    pg.init()
68    pg.start()
69    pg.join()
70    pg.close()

Key points:

  • One DataLoader per rank — each rank constructs its own DataLoader backed by the same full dataset. The DistributedSampler partitions the indices so that ranks never load the same sample within an epoch.

  • ``sampler.set_epoch(epoch)`` — call this at the start of every epoch to ensure that the random shuffle seeds differ between epochs while remaining consistent across ranks.

  • GPU affinity — Dragon’s Policy already pins each worker to its own GPU, so torch.device("cuda") always refers to the correct device.