Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add support for variable length dataloaders in DDP #3416

Merged
merged 9 commits into from
Jun 24, 2024
Merged
Changes from 4 commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
22 changes: 22 additions & 0 deletions composer/trainer/trainer.py
Original file line number Diff line number Diff line change
Expand Up @@ -3640,6 +3640,16 @@ def _iter_dataloader(self, trainer_mode: TrainerMode):
else:
dataloader_iter = itertools.islice(self.state.dataloader, int(self.state.dataloader_len))

# Track if iteration has finished (used for distributed training when we have variable length dataloaders)
# 0 = not finished, 1 = finished (using integer tensors so we can use dist.all_reduce)
iter_finished = torch.zeros(1, dtype=torch.uint8)
iter_finished = self.state.device.tensor_to_device(iter_finished)
JAEarly marked this conversation as resolved.
Show resolved Hide resolved

# Initialize batch to avoid "referenced before assignment" warnings
# Unique sentinel value to differentiate uninitialized state and dataloader yielding None
sentinel = object()
batch = sentinel

while True:
try:
# [BEFORE/AFTER]_DATALOADER only runs while training
Expand All @@ -3655,7 +3665,19 @@ def _iter_dataloader(self, trainer_mode: TrainerMode):
# Otherwise, we will encounter an error at the start of the next epoch when
# Event.BEFORE_DATALOADER tries to start an unfinished marker.
self.engine.run_marker_only_event(Event.AFTER_DATALOADER)
# Mark iteration as finished - don't break yet as we need to sync across ranks
iter_finished += 1

# Sync iter finished across ranks
dist.all_reduce(iter_finished, reduce_operation='MAX')
# If any rank has finished, stop all rank iterations
if iter_finished.item() == 1:
break

if batch is sentinel:
raise RuntimeError(
"Batch should have been assigned or loop should have been broken. This shouldn't happen!",
)
mvpatel2000 marked this conversation as resolved.
Show resolved Hide resolved
yield batch

def _use_closures(self) -> bool:
Expand Down
Loading