Skip to content

Synchronize end of loader for repeat=False #146

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.

Sign up for GitHub

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

Open
voegtlel wants to merge 6 commits into
base: develop
Choose a base branch
from
Open

Synchronize end of loader for repeat=False #146

Show file tree
Hide file tree
Changes from 5 commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
c12a188
Small test fixes
voegtlel Jun 30, 2025
476c739
Implement synchronization at the end of dataset loader with repeat=False
voegtlel Jul 1, 2025
a0402fc
Add end of dataset sync to documentation.
voegtlel Jul 1, 2025
7ec2298
Just fix
voegtlel Jul 1, 2025
64b7f6b
Fix sync test import
voegtlel Jul 1, 2025
6a8d94d
Small doc adaption
voegtlel Jul 2, 2025
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
12 changes: 9 additions & 3 deletions docs/source/advanced/epochized_blending.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -25,20 +25,23 @@ splits:
split_part: val
```

Now, the call to `get_train_dataset` requires the additional parameter `repeat=False` to interrupt iterating after one epoch:
Now, the call to `get_train_dataset` requires the additional parameter `repeat=False` to interrupt iterating after one epoch and use the `RedistributeLoader` to synchronize the last samples when some ranks have already exhausted their data:

```py
from megatron.energon import get_train_dataset, get_loader, WorkerConfig
from megatron.energon.sync_end import RedistributeLoader


loader = get_loader(get_train_dataset(
# Set repeat=False to avoid repeating the dataset.
# Also add RedistributeLoader to synchronize the end of rank exhaustion. Only works with initialized torch distributed.
loader = RedistributeLoader(get_loader(get_train_dataset(
Copy link
Collaborator

@philipp-fischer philipp-fischer Jul 2, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think it would be better if this was transparent to the user, so make get_loader handle this internally.
And since the choice of RedistributeLoader vs. StopFirstLoader actually changes the data that's being iterated, this choice should be made in the metadataset and not in the code, I think.
As a property of blend_epochized.

I.e. blend_epochized can either be a list as before (chooses default RedistributeLoader), or it can be a dict for more customization like

    blend_epochized:
      phase_out_behavior: stop_first_loader
      datasets:
        - repetitions: 5
            path: ./coco
            # ... Other parameters
        - repetitions: 2
            path: ./coyo
        - repetitions: 1
            path: ./coyo
            split_part: val

Copy link
Collaborator Author

@voegtlel voegtlel Jul 2, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Okay, we can make this be handled in get_loader. I personally would prefer it to be separate though, as it's a feature on top?

Regarding moving the configuration to the metadataset:
I see your point that this slightly modifies how the data is iterated, but I'd also argue:

  1. So far we don't really rely on torch distributed, where this piece of code is tightly bound to it.
  2. This would also disable nesting of blend_epochized, because you cannot nest different (or unconfigured) phase_out_behavior.
  3. This depends on repeat=False and doesn't make sense if repeat=True, so it's based on what the user sets in the code.
  4. At least for RedistributeLoader it should not really change the data frequency (so far the settings in the metadataset mainly focuses on data frequency / blend).
  5. If we move the boundary of metadataset to include this, then we should also have gradient accumulation, seeds, batch size, handling of incomplete batches, etc. in the config. I wouldn't want that, tbh.

Thus voting for keeping this in code, not in the metadataset config.

'metadataset.yaml',
batch_size=2,
shuffle_buffer_size=100,
max_samples_per_sequence=100,
worker_config=WorkerConfig.default_worker_config(),
repeat=False,
))
)))

# This will now stop iterating after the datasets have been iterated (coco 5 times, coyo-train 2
# times and coyo-val 1 times). Of course, the data is still being shuffled between all those
Expand All @@ -54,3 +57,6 @@ for batch in loader:

If used as dataset for `get_val_dataset`, the `repetitions` are ignored.
The metadataset would also work without setting `repeat=False`, but then the shuffle buffer will shuffle samples across bounderies of epochs.
There are two available end of iteration synchronizers:
* `RedistributeLoader`: Redistributes samples when a rank is exhausted before other ranks. On next epoch, includes the incomplete batches' samples.
* `StopFirstLoader`: Stop as soon as the first rank is exhausted. The next epoch will iterate until the next loader stops, restarting all ranks once.
21 changes: 21 additions & 0 deletions src/megatron/energon/sync_end/__init__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,21 @@
# Copyright (c) 2025, NVIDIA CORPORATION.
# SPDX-License-Identifier: BSD-3-Clause
from megatron.energon.sync_end.redistribute import RedistributeDataLoaderState, RedistributeLoader
from megatron.energon.sync_end.stop_first_end import StopFirstDataLoaderState, StopFirstLoader

"""
Provides wrappers for the dataset loaders that allow for synchronization at the end of the dataset.
I.e. if running a training with repeat=False, the loaders will typically exhaust at different times, which may require
synchronization across ranks.

The wrappers are:
- RedistributeLoader: Redistributes the last samples to the ranks that are not exhausted.
- StopFirstLoader: Stops iterating as soon as the first rank is exhausted.
"""

__all__ = [
"RedistributeLoader",
"RedistributeDataLoaderState",
"StopFirstLoader",
"StopFirstDataLoaderState",
]
Loading