Add Distributed Data Parallel (DDP)

Signed-off-by: Fynn Schmitt-Ulms <[email protected]>

Author: fynnsu

src/speculators/train/training_loop.py

@@ -1,12 +1,34 @@
+import os
 import torch
 from transformers import LlamaConfig
 
 from speculators.train.eagle3.core import Eagle3DraftModel
 from speculators.train.data import Eagle3SampleFileDataset, create_collate_fn
 from torch.utils.data import DataLoader
 
+from torch.nn.parallel import DistributedDataParallel as DDP
+import torch.distributed as dist
 
-DEVICE = "cuda:0"
+def maybe_setup_distributed():
+    # Based off of https://docs.pytorch.org/tutorials/intermediate/ddp_tutorial.html#initialize-ddp-with-torch-distributed-run-torchrun
+    if "LOCAL_RANK" not in os.environ:
+        # No distributed training
+        return 0, 1, 0, False
+    local_rank = int(os.environ.get("LOCAL_RANK", 0))
+    world_size = int(os.environ.get("WORLD_SIZE", 1))
+    torch.accelerator.set_device_index(local_rank)
+    acc = torch.accelerator.current_accelerator()
+    backend = torch.distributed.get_default_backend_for_device(acc)
+    dist.init_process_group(backend)
+    rank = dist.get_rank()
+
+    print(f'Started DDP with local_rank={local_rank}, world_size={world_size}, rank={rank}')
+    return local_rank, world_size, rank, True
+
+local_rank, world_size, rank, is_distributed = maybe_setup_distributed()
+
+
+DEVICE = torch.device(local_rank)
 EPOCHS = 10
 draft_vocab_size = 5000
 verifier_vocab_size = 151936
@@ -47,6 +69,9 @@
 
 # draft_model.load_verifier_lm_head(verifier_model_name_or_path) # Doesn't work for Qwen2.5 VL, need better head loading method
 
+if is_distributed:
+    draft_model = DDP(draft_model, device_ids=[local_rank])
+opt = torch.optim.Adam(draft_model.parameters(), lr=1e-4)
 
 dataset = Eagle3SampleFileDataset(datapath=datapath, max_len=total_seq_len)
 train_loader = DataLoader(
@@ -57,7 +82,6 @@
     pin_memory=True,
     collate_fn=create_collate_fn(total_seq_len),
 )
-opt = torch.optim.Adam(draft_model.parameters(), lr=1e-4)
 
 
 def train_epoch(
@@ -67,18 +91,35 @@ def train_epoch(
     opt: torch.optim.Optimizer,
     epoch: int,
     local_rank: int,
+    is_distributed: bool,
 ):
     model.train()
 
     for batch in train_loader:
         batch = {k: v.to(local_rank) if isinstance(v, torch.Tensor) else v for k, v in batch.items()}
 
         _, loss = model(**batch, use_off_policy_tokens=True)
-        print(loss.item())
         opt.zero_grad()
         loss.backward()
         opt.step()
 
+        loss = loss.detach().clone()
+        if is_distributed:
+            # Note: this is not needed for training, just for logging
+            dist.reduce(loss, dst=0, op=dist.ReduceOp.AVG)
+
+        if local_rank == 0:
+            print(loss.item())
+
+
 
 for epoch in range(EPOCHS):
-    train_epoch(draft_model, train_loader, None, opt, epoch, DEVICE)
+    train_epoch(draft_model, train_loader, None, opt, epoch, local_rank, is_distributed)
+
+if is_distributed:
+    dist.destroy_process_group()
+    print(f'Destroyed DDP with local_rank={local_rank}, world_size={world_size}, rank={rank}')
+
+
+# RUN WITH:
+# CUDA_VISIBLE_DEVICES=0,1,2,3 torchrun --nnodes=1 --nproc_per_node=4  src/speculators/train/training_loop.py