Iterable Dataset #2852

felipemello1 · 2025-06-26T14:33:36Z

Context

What is the purpose of this PR? Is it to

add a new feature
fix a bug
update tests and/or documentation
other (please add here)

Enable Iterable datasets in torchtune.

CONTEXT: built on top of ongoing PR step-based-ckpt: #2384

TIps when reviewing this pr

Follow this order:

recipes/configs/llama3_2/3B_full.yaml: see the configs
torchtune/datasets/_iterable_base.py: base class for iterable dataset
torchtune/datasets/_hf_iterable.py: ds based on HF -- Can be replaced easily. Downstream does not expect HF.
torchtune/datasets/_interleaved.py: interleave the datasets
torchtune/data/_metrics.py: metrics transform to create the metrics
torchtune/data/_aggregator.py: aggregate the metrics at the recipe level
recipes/full_finetune_distributed.py: everything put together
unit tests

torchtune/datasets/_hf_iterable.py

Changelog

Datasets are infinite
User doesn't define epochs anymore, but training steps (how many times we update the optimizer)
Support for dataset mixing -- follow up PRs is to enable curriculum learning
Support for dataset metric logging -- User can understand epoch per dataset, distribution of token lens, etc. Easy to add new metrics.
HF agnostic. Even though the current dataset is HF, the dataloader, packed, datamixing, metric logging is agnostic to it
Well tested in distributed setting -- WARNING: need better testing for multiprocess dataloader. It doesnt guarantee determinism, so I postponed testing this setting

Config and builder design based on the discussions after this RFC: #2785

Next steps:
7. Gather feedback on metric logging. E.g. we can add more aggregation types.
8. Polish the code a little bit
9. Add packing from this RFC: #2819
10. Add curriculum learning
11. Docs?

Test plan

UNTESTED: resume from ckpt in the recipe. However, we have plenty of tests showing that resuming works for these iterable datasets.

…iterable_dataset_final

pytorch-bot · 2025-06-26T14:33:41Z

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/torchtune/2852

📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

❌ 3 New Failures, 4 Unrelated Failures

As of commit 72211c9 with merge base 3d73591 ():

NEW FAILURES - The following jobs have failed:

GPU tests / gpu_test (3.9, stable) (gh)
Process completed with exit code 4.
Lint / lint (3.10) (gh)
Unit Test / unit_tests (3.9) (gh)
Process completed with exit code 4.

FLAKY - The following jobs failed but were likely due to flakiness present on trunk:

Unit Test / unit_tests (3.10) (gh) (similar failure)
##[error]The operation was canceled.
Unit Test / unit_tests (3.11) (gh) (similar failure)
##[error]The operation was canceled.

BROKEN TRUNK - The following jobs failed but were present on the merge base:

👉 Rebase onto the `viable/strict` branch to avoid these failures

GPU tests / gpu_test (3.10, stable) (gh) (trunk failure)
##[error]The operation was canceled.
GPU tests / gpu_test (3.11, stable) (gh) (trunk failure)
##[error]The operation was canceled.

This comment was automatically generated by Dr. CI and updates every 15 minutes.

felipemello1 · 2025-06-26T14:42:24Z

torchtune/datasets/_slimorca.py

@@ -94,3 +95,72 @@ def slimorca_dataset(
            )
        return PackedDataset(ds, max_seq_len=tokenizer.max_seq_len)
    return ds
+
+
+def slimorca_iterable_dataset(


added here to demonstrate datamix iterable dataset with this example. Personally, i dislike exposing all of the args and defaults. I would prefer to expose only whats specific to this builder.

torchtune/datasets/_sft.py

felipemello1 · 2025-06-26T14:43:51Z

torchtune/datasets/_interleaved.py

+                logger.warning(
+                    f"Child dataset {self._datasets[ds_name].dataset_name} was exhausted. "
+                    "This is unexpected for an infinite dataset. Re-initializing its iterator."
+                )


not 100% sure i like this

Let's do this: simply have a subclass for InfiniteIterable so this is super explicit

Where did this one land? I don't see InfiniteIterable anywhere (personally I don't know enough yet to have a strong preference here, just wanna understand where things currently stand)

i made changes but didnt push them yet. I added a dummy class that does nothing:

class InfiniteTuneIterableDataset(TuneIterableDataset): """Abstract base class for infinite datasets, which yield samples indefinitely. It only purpose is to make it explicit that the dataset is expected to be infinite, i.e. it never exhausts. This is helpful to avoid complexity due to some rank hanging because of lack of data"" pass

and replaced this logger.warning with raise ValueError.

I think its better to have zero tolerance. Datasets that are not infinite need work to make sure no rank hangs.

felipemello1 · 2025-06-26T14:46:30Z

torchtune/datasets/_alpaca.py

@@ -101,3 +102,64 @@ def alpaca_dataset(
 original Alpaca dataset, `yahma/alpaca-cleaned <https://huggingface.co/datasets/yahma/alpaca-cleaned>`_.
 See the dataset page and :func:`~torchtune.datasets.alpaca_dataset` for more details.
 """
+
+
+def alpaca_iterable_dataset(


added here to demonstrate datamix iterable dataset with this example. Personally, i dislike exposing all of the args and defaults. I would prefer to expose only whats specific to this builder.

But you are doing this with ``load_dataset_kwargs, right? Or did you mean something else?

nit: it's a function, so... get_alpaca_iterable_dataset?

the get makes sense, but its not the pattern we have in tune :/

Darktex

Great PR! I mainly had a question on the interaction with packing and on the SFT transform

torchtune/data/_aggregator.py

torchtune/data/_metrics.py

Darktex · 2025-06-26T16:40:04Z

torchtune/datasets/_alpaca.py

@@ -101,3 +102,64 @@ def alpaca_dataset(
 original Alpaca dataset, `yahma/alpaca-cleaned <https://huggingface.co/datasets/yahma/alpaca-cleaned>`_.
 See the dataset page and :func:`~torchtune.datasets.alpaca_dataset` for more details.
 """
+
+
+def alpaca_iterable_dataset(


But you are doing this with ``load_dataset_kwargs, right? Or did you mean something else?

torchtune/datasets/_sft.py

torchtune/datasets/_interleaved.py

Darktex · 2025-06-26T17:53:33Z

torchtune/datasets/_interleaved.py

+                logger.warning(
+                    f"Child dataset {self._datasets[ds_name].dataset_name} was exhausted. "
+                    "This is unexpected for an infinite dataset. Re-initializing its iterator."
+                )


Let's do this: simply have a subclass for InfiniteIterable so this is super explicit

Darktex · 2025-06-26T17:59:13Z

torchtune/datasets/_iterable_base.py

+from torch.utils.data import IterableDataset
+
+
+class TuneIterableDataset(IterableDataset, ABC):


We need this guy to interact with packing and IIUC I don't believe this is currently happening?

The algo we should implement is this:

One batch can be made of multiple calls to next. We keep taking until we exceed the max seq len. When we do, we put the last one aside (we'll use it to start the next batch), pad the current one to max len and return.

The calls to next will go to the interleaved dataset, therefore we automatically construct mixed batches from multiple datasets without much effort

Also, every time we call next we should make space for logging transforms (which we are, you already wrote them). I think it's ok to make your metrics transforms and aggregators an optional property here so the semantics are clearer

we have packing here: #2819

…htune into iterable_dataset_final

ebsmothers

I read every line of this PR. (Kidding but I tried to at least look at most of the important stuff.) Thanks for taking on this massive set of changes, I think the dataset classes are a big improvement

ebsmothers · 2025-07-03T22:41:02Z

torchtune/datasets/_hf_iterable.py

+        # Load and shard dataset
+        ds = load_dataset(**load_dataset_kwargs)
+
+        # Use to_iterable_dataset for streaming datasets


I don't understand this comment. Based on L185 it seems you're using to_iterable_dataset for non-streaming datasets

ebsmothers · 2025-07-03T22:48:44Z

torchtune/datasets/_hf_iterable.py

+
+            # If the dataset is not streaming and has a defined length,
+            # we cannot have num_shards > dataset_size.
+            if not load_dataset_kwargs.get("streaming", False) and hasattr(


This is redundant inside the current if statement

ebsmothers · 2025-07-03T22:58:00Z

torchtune/datasets/_hf_iterable.py

+            # we will try 2*16 = 32 shards. Since 32 is not a multiple of 3, we will do 36 shards.
+            # Each rank gets 16 shards, each dataloader worker in that rankgets 6 shards.


this math doesn't seem quite right

Suggested change

# we will try 2*16 = 32 shards. Since 32 is not a multiple of 3, we will do 36 shards.

# Each rank gets 16 shards, each dataloader worker in that rankgets 6 shards.

# we will try 2*16 = 32 shards. Since 32 is not a multiple of 6, we will do 36 shards.

# Each rank gets 18 shards, each dataloader worker in that rank gets 6 shards.

ebsmothers · 2025-07-03T22:58:37Z

torchtune/datasets/_hf_iterable.py

+                if num_shards > dataset_size:
+                    raise ValueError(
+                        f"Number of shards ({num_shards}) is greater than the dataset size ({dataset_size})."
+                        f"Please decrease num_shards_per_rank."


nit but this can also happen due to total_workers > dataset_size (since we round up). Probably an edge case but if e.g. torch.utils.data.get_worker_info() returns a really large number then this is not actually the right fix

ebsmothers · 2025-07-03T23:05:44Z

torchtune/datasets/_hf_iterable.py

+            epoch_seed = self._seed + self._num_epochs
+            self._ds.set_epoch(epoch_seed)


I'm confused.. isn't this bumping epoch by self._seed each epoch (when it should be bumping by 1)?

Yeah, from their docs, it should be just: dataset.set_epoch(epoch)

I will make the change

ebsmothers · 2025-07-03T23:35:39Z

torchtune/data/metrics/_metric_transform.py

+            self.new_metric(
+                name="tokens_seen", value=token_len, agg_type=AggregationType.SUM
+            ),
+            self.new_metric(
+                name="seq_len", value=token_len, agg_type=AggregationType.DISTRIBUTION
+            ),


A minor thing, but to me metrics having the same value but different aggregation types should not actually be represented as distinct metrics. Like I should be able to just define how a metric is computed for a given sample, then separately choose different types of aggregation as needed

hmm, so agg_type being a List[AggregationType]?

self.new_metric( name="tokens_seen", value=token_len, agg_type=[AggregationType.SUM, AggregationType.DISTRIBUTION ),

I dont know if the extra complexity is worth it. Adding two metrics is cheap. But i guess there can be situations where i just want 'mean' and 'sum', and then if i create two metrics named "metric_mean" and "metric_sum", they would be logged as "metric_mean_mean" and "metric_sum_sum".

I have to think a bit about it

ebsmothers · 2025-07-03T23:51:38Z

torchtune/data/metrics/_metric_aggregator.py

+from torchtune.data.metrics._metric_transform import AggregationType, Metric
+
+
+class MetricsAggregator:


A high level comment: the relationship between this and the agg handlers is not super clear to me. It seems like we are using a registry pattern where the handlers are responsible for defining the actual aggregation logic. But then the all-gather happens in here. (Separately I stand by my claim that it would be better to hold off on more complex cases like distribution aggregators so as not to boil the ocean here.)

the handlers are responsible for defining the actual aggregation logic. But then the all-gather happens in here.

why is that a contradiction?

The MetricsAggregator calls the handler.finalize_local_agg

then does a single all_gather to get the results from all ranks for all metrics

Then calls handler._finalize_dist_agg([aggregated_results_per_rank]*n_ranks)

Do you wanna suggest a different way of doing it? Or is it hard to spot this pattern in the code?

my claim that it would be better to hold off on more complex cases like distribution aggregators

If we dont do aggregation across ranks, we wouldnt be able to count things like "tokens_seen", right? :/

Or do you mean that we should delete DistributionAggHandler? To clarify, this distribution has nothing to do with multiple gpus. Its just stats, e.g. std, percentiles, max, min, etc. Maybe i should rename if its causing confusion.

ebsmothers · 2025-07-03T23:52:59Z

recipes/full_finetune_distributed.py

+        if cfg.get("dataset_val") is not None:
+            raise NotImplementedError(
+                "Validation is not supported yet with iterable datasets."
+            )


Is there a specific technical reason here? Or we just haven't gotten to it yet

validation datasets are not infinite!!! Need to figure out how to solve this one, but it wont be on this PR

ebsmothers · 2025-07-03T23:54:42Z

recipes/full_finetune_distributed.py

        self.global_step = 0

+        # Step-based training support
+        self.num_training_steps = cfg.num_training_steps
+        self._dataset_metrics_log_freq = cfg.get("dataset_metrics_log_freq", 100)


Very minor but does this really need to be configured separately from log_every_n_steps? (Edit: I see you already added a TODO about this below, I guess you know my vote)

ebsmothers · 2025-07-03T23:57:57Z

recipes/full_finetune_distributed.py

+                    self._metric_logger.log_dict(log_dict, step=self.global_step)
+
+                # Log dataset metrics
+                # #TODO: it requires all_gather. Should we keep a separate log_freq for this?


Worth benchmarking, but if it's just a few floats I think it shouldn't matter

felipemello1 and others added 6 commits June 25, 2025 12:41

first commit

3cab533

update tests

2212b19

Merge remote-tracking branch 'joecummings/impl-step-based-ckpt' into …

4345832

…iterable_dataset_final

linter

2eb68b6

tests pass

2e51e04

it works

93fa743

facebook-github-bot added the CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. label Jun 26, 2025

remove code

aa9e6f4

felipemello1 changed the title ~~first commit~~ Iterable Dataset Jun 26, 2025

felipemello1 commented Jun 26, 2025

View reviewed changes

Darktex requested changes Jun 26, 2025

View reviewed changes

felipemello1 and others added 11 commits July 1, 2025 21:18

update metrics to use handlers

5b188ed

remove file after refactoring

2eab08d

add distributed tsts

58491f1

Merge branch 'iterable_dataset_final' of github.com:felipemello1/torc…

da7245d

…htune into iterable_dataset_final

tests pass

96424d0

optimize SFTOutputTransform

853147b

use ds.sampling_weight

96bc317

add sampling log to interlead dataset

3c9d161

fix nested interleave

4804663

changes to TuneIterableDataset

2fe4b40

add IterableDataset back

72211c9

ebsmothers reviewed Jul 3, 2025

View reviewed changes

		from torch.utils.data import IterableDataset


		class TuneIterableDataset(IterableDataset, ABC):

		# we will try 2*16 = 32 shards. Since 32 is not a multiple of 3, we will do 36 shards.
		# Each rank gets 16 shards, each dataloader worker in that rankgets 6 shards.

		epoch_seed = self._seed + self._num_epochs
		self._ds.set_epoch(epoch_seed)

		from torchtune.data.metrics._metric_transform import AggregationType, Metric


		class MetricsAggregator:

Iterable Dataset #2852

Are you sure you want to change the base?

Iterable Dataset #2852

Uh oh!

Conversation

felipemello1 commented Jun 26, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Context

TIps when reviewing this pr

Changelog

Test plan

Uh oh!

pytorch-bot bot commented Jun 26, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/torchtune/2852

❌ 3 New Failures, 4 Unrelated Failures

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Darktex left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

ebsmothers left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

felipemello1 Jul 4, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

felipemello1 commented Jun 26, 2025 •

edited

Loading

pytorch-bot bot commented Jun 26, 2025 •

edited

Loading

felipemello1 Jul 4, 2025 •

edited

Loading

felipemello1 Jul 4, 2025 •

edited

Loading