use sampler in train_ddp.py

mreso · mreso · commit c4e58bb36e75 · 2025-01-15T16:26:46.000-08:00
diff --git a/train_ddp.py b/train_ddp.py
@@ -48,12 +48,13 @@ def main() -> None:
         rank=0,
         # for DDP we can use replica groups of size 1, FSDP/PP/CP would need more.
         num_replicas=1,
+        shuffle=True,
     )
 
     # This uses the torchdata StatefulDataLoader to be able to checkpoint and
     # restore the per worker dataloader position.
     trainloader = StatefulDataLoader(
-        trainset, batch_size=64, shuffle=True, num_workers=2
+        trainset, batch_size=64, num_workers=2, sampler=sampler
     )
 
     def load_state_dict(state_dict):

Original file line number	Diff line number	Diff line change
`@@ -48,12 +48,13 @@ def main() -> None:`
`48`	`48`	`rank=0,`
`49`	`49`	`# for DDP we can use replica groups of size 1, FSDP/PP/CP would need more.`
`50`	`50`	`num_replicas=1,`
	`51`	`+ shuffle=True,`
`51`	`52`	`)`
`52`	`53`
`53`	`54`	`# This uses the torchdata StatefulDataLoader to be able to checkpoint and`
`54`	`55`	`# restore the per worker dataloader position.`
`55`	`56`	`trainloader = StatefulDataLoader(`
`56`		`- trainset, batch_size=64, shuffle=True, num_workers=2`
	`57`	`+ trainset, batch_size=64, num_workers=2, sampler=sampler`
`57`	`58`	`)`
`58`	`59`
`59`	`60`	`def load_state_dict(state_dict):`