[draft] torchdata integration #1929

andrewkho · 2024-10-30T22:37:30Z

Note! This requires torchdata nightly to be installed to work correctly.

Test multi-dataset training command:
tune run --nnodes=1 --nproc_per_node=8 lora_finetune_distributed_td --config llama3_2_vision/11B_lora_td dataloader.pin_memory=True use_torchdata=True dataloader.parallel_method=thread max_steps_per_epoch=1000 compile=True dataloader.num_workers=4

Test multi-dataset command with dataloader_only mode:
tune run --nnodes=1 --nproc_per_node=8 lora_finetune_distributed_td --config llama3_2_vision/11B_lora_td dataloader.pin_memory=True use_torchdata=True dataloader.parallel_method=thread max_steps_per_epoch=1000 profile_mode=dataloader_only compile=True dataloader.num_workers=4

Benchmarking on 8xA100

for 200 steps, full model training, batch_size: 4, gradient_accumulation_steps: 1
TWFB: Sum of Time Waiting For Batch (max across all ranks) divided by sum of step times
Single Datasets are run with OCRVQA
Multi-Dataset was run with:
ocrvqa, docvqa, dvqa, tabmwp with equal weighting for all datasets.

Trial	TWFB %	Sum DL times	sum step times
0 Single dataset (Baseline)	0.06279	14.596	232.454
1 Single dataset (TorchData threads)	0.00726	1.64746	226.780
2 Multi dataset (threads)	0.00514	2.4945	485.388
3 Multi dataset Streaming	0.01712,	8.5445	499.0716

Multi-dataset runs much slower, guessing because one of the datasets (dvqa?) requires more padding than ocrvqa.

Launch commands:
0: tune run --nnodes=1 --nproc_per_node=8 lora_finetune_distributed_td --config llama3_2_vision/11B_lora_td use_torchdata=false max_steps_per_epoch=200 compile=True
1: (same as 2 but with 3 of the datasets removed in the config)
2: tune run --nnodes=1 --nproc_per_node=8 lora_finetune_distributed_td --config llama3_2_vision/11B_lora_td use_torchdata=true max_steps_per_epoch=200 compile=True
3: tune run --nnodes=1 --nproc_per_node=8 lora_finetune_distributed_td --config llama3_2_vision/11B_lora_td use_torchdata=true max_steps_per_epoch=200 compile=True dataloader.streaming=true

Please have a look at the code set up, and how this composability can help with streaming datasets, and multi-dataset mixing. You can think of this as approaches to replace torch.utils.data.DataLoader, while introducing more flexible parallelism schemes, eg instead of just one multiprocess worker setting, you could do multi-threading, pipeline parallelism, etc. It also enables more powerful composability IMO.

I have done some single-device benchmarking on a machine with A100 40gb, both with and without the model, performance is on par or better than standard dataloader.

TODO: fill in below

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)

Please link to any issues this PR addresses.

Changelog

What are the changes made in this PR?
*

Test plan

Please make sure to do each of the following if applicable to your PR. If you're unsure about any one of these just ask and we will happily help. We also have a contributing page for some guidance on contributing.

run pre-commit hooks and linters (make sure you've first installed via pre-commit install)
add unit tests for any new functionality
update docstrings for any new or updated methods or classes
run unit tests via pytest tests
run recipe tests via pytest tests -m integration_test
manually run any new or modified recipes with sufficient proof of correctness
include relevant commands and any other artifacts in this summary (pastes of loss curves, eval results, etc.)

UX

If your function changed a public API, please add a dummy example of what the user experience will look like when calling it.
Here is a docstring example
and a tutorial example

I did not change any public API
I have added an example to docs or docstrings

pytorch-bot · 2024-10-30T22:37:34Z

🔗 Helpful Links

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

📄 Preview Python docs built from this PR

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

❗ 1 Active SEVs

There are 1 currently active SEVs. If your PR is affected, please view them below:

GLIBC not found in Nova ROCM workflows

❌ 2 New Failures, 4 Cancelled Jobs

As of commit b1b2ab6 with merge base 4fb2464 ():

NEW FAILURES - The following jobs have failed:

Recipe Tests / recipe_test (3.11) (gh)
ERROR: THESE PACKAGES DO NOT MATCH THE HASHES FROM THE REQUIREMENTS FILE. If you have updated the package versions, please update the hashes. Otherwise, examine the package contents carefully; someone may have tampered with them.
Unit Test / unit_tests (3.9) (gh)
ERROR: Could not install packages due to an OSError: HTTPSConnectionPool(host='files.pythonhosted.org', port=443): Max retries exceeded with url: /packages/ed/a5/33cf000137545a08b0a3a6ea76c8ccbd87917f78bb5d737f9f56f3b11ef6/datasets-3.1.0-py3-none-any.whl.metadata (Caused by ResponseError('too many 502 error responses'))

CANCELLED JOBS - The following jobs were cancelled. Please retry:

Recipe Tests / recipe_test (3.10) (gh)
##[error]The operation was canceled.
Recipe Tests / recipe_test (3.9) (gh)
Unit Test / unit_tests (3.10) (gh)
##[error]The operation was canceled.
Unit Test / unit_tests (3.11) (gh)
##[error]The operation was canceled.

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

RdoubleA

Thanks for this prototype, this is helpful to see! The way I'd imagine we'd expose the torchdata dataloader would be from a builder function with a few knobs exposed with reasonable defaults:

def build_dataloader(ds, num_workers, pin_memory, prefetch, in_memory, parallel_method, ...):
	# Put together all the nodes here

# In config
dataloader:
  num_workers:
  ...

For a power user, what might they want to tune to optimize performance for their hardware and model setup?

It's also not clear to me how some media transforms/decoding might get optimized, is that just handled by the torchdata nodes automatically?

RdoubleA · 2024-10-30T22:56:40Z

recipes/lora_finetune_single_device.py

+            node = IterableWrapper(sampler)
+            node = _Mapper(node, map_fn=ds._data.__getitem__)
+            # Cut here for Streaming/Iterable dataset instead =====
+            node = _Mapper(node, map_fn=ds._prepare_sample)


is this where the transform would get parallelized?

Yes, see

_Mapper = partial( ParallelMapper, num_workers=num_workers, method=parallel_method, in_order=True, )

RdoubleA · 2024-10-30T22:57:26Z

recipes/lora_finetune_single_device.py

+                )
+            # Map style set up =======
+            node = IterableWrapper(sampler)
+            node = _Mapper(node, map_fn=ds._data.__getitem__)


So does this mean we can keep our own Dataset abstractions? What if we went with entirely Iterable datasets?

Technically yes, but this was just a way to get it working quickly, we can also refactor this into separate function. For IterableDataset, you could wrap it in IterableWrapper and then call everything underneath here, ie

node = IterableWrapper(my_iterable_dataset) node = _Mapper(node, map_fn=ds._prepare_sample) ...

joecummings · 2024-10-30T23:54:08Z

recipes/lora_finetune_single_device.py

+                            batch = next(dl_iter)
+                        dl_dt = time.perf_counter() - dl_t0
+                        idx += 1
+                    except StopIteration:


Why not create a context manager that handles this?

No good reason, hacky code is hacky :) Just wanted to see some rough numbers

ebsmothers · 2024-11-01T01:00:06Z

recipes/lora_finetune_single_device.py

+            # Map style set up =======
+            node = IterableWrapper(sampler)
+            node = _Mapper(node, map_fn=ds._data.__getitem__)
+            # Cut here for Streaming/Iterable dataset instead =====


Have we tried this out on an HF dataset with streaming = True yet? (I assume it won't work out of the box yet?)

see update :D almost out-of-the-box

why do we need to wrap with IterableWrapper if the underlying dataset is IterableDataset? Also in our case, the underlying dataset would be a HF dataset class

Wrapper will make it conform to BaseNode's API so we can hold a pointer to iterator, as well as unify the state management. Subclasses need to define .iterator() instead of .__iter__().

ebsmothers

Thanks for creating this PR! I like the IterableWrapper/Batcher/etc APIs, they look pretty clean (and agree with @RdoubleA's suggestion about exposing this in builders/configs). How will this work when training on multiple datasets? Will we just apply the same set of APIs to each sampler individually?

andrewkho · 2024-11-01T22:37:12Z

Thanks for the comments y'all, I updated this with a streaming example. Test with:

tune run lora_finetune_single_device --config llama3_2_vision/11B_lora_single_device num_workers=2 pin_memory=False use_torchdata=True parallel_method=thread max_steps_per_epoch=50 dataset.streaming=True

andrewkho · 2024-11-01T22:41:35Z

@RdoubleA

For a power user, what might they want to tune to optimize performance for their hardware and model setup?

Good question, at the minimum it'd be some global worker-setting, but one option that may be worth supporting is allowing power users to define their entire pipelines entirely in config, not sure if you think this is a bad idea. eg by default, use a builder with a few knobs, but also allow the entire dataloader definition to be composable. Somewhat similar to how you enable users to pass a list of datasets right now. Something similar could be done for mixing. We'd need to be thoughtful for what the useful small-atomic units woudl be, and figure out syntax sugar.

It's also not clear to me how some media transforms/decoding might get optimized, is that just handled by the torchdata nodes automatically?

In terms of optimizing, one thing we're looking at is tf.data's autotune approach which will automatically update prefetch buffers and workers. This is not implemented, but hoping to do something along these lines which will help with the too-many-tunable-parameters problem.

joecummings · 2024-11-12T16:23:58Z

torchtune/datasets/multimodal/_the_cauldron.py

@@ -235,3 +243,73 @@ def __call__(self, sample: Mapping[str, Any]) -> Mapping[str, Any]:
    if packed:
        raise ValueError("Multimodal datasets don't support packing yet.")
    return ds
+
+
+def the_cauldron_dataset_torchdata(


How much of this could be pushed into the SFT class so that we can just reuse it for any new datasets?

A lot of it is probably re-usable, I think it should be it's own class or at least a builder, maybe like _sft.py: class SFTDatasetNode or something less terrible sounding

Actually 90% of this builder func could probably live in _sft.py as it doesn't have anything to do with the cauldron

andrewkho · 2024-11-12T16:56:03Z

recipes/lora_finetune_distributed_td.py

@@ -0,0 +1,1119 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.


Copied and modified from lora_finetune_distributed

andrewkho · 2024-11-12T16:56:26Z

recipes/configs/llama3_2_vision/11B_lora_td.yaml

@@ -0,0 +1,131 @@
+# Config for multi-device LoRA finetuning in lora_finetune_distributed_td.py


Copied and modified from 11B_lora

andrewkho · 2024-11-12T16:57:22Z

torchtune/datasets/_sft.py

@@ -117,7 +128,17 @@ def __getitem__(self, index: int) -> Dict[str, Any]:
        sample = self._data[index]
        return self._prepare_sample(sample)

-    def _prepare_sample(self, sample: Mapping[str, Any]) -> Dict[str, Any]:
+
+class PrepareSample:


Refactoring this into it's own Callable class that can be used by both the current torch.utils.Dataset and the torchdata.nodes

felipemello1

hey Andrew, thank you so much for this PR! This is such a nice feature to have.

I did a first pass. I understand that its still a draft, but thought of making comments so it could save you some time for when the PR is closer to being ready.

felipemello1 · 2024-11-13T19:52:54Z

torchtune/datasets/_sft.py

+                rank=int(os.environ.get("RANK", 0)),
+                world_size=int(os.environ.get("WORLD_SIZE", 1)),


should we use our utility instead?

torchtune/torchtune/training/_distributed.py

Line 150 in 1eb7785

def get_world_size_and_rank() -> Tuple[int, int]:

Needs to be removed

felipemello1 · 2024-11-13T20:12:33Z

recipes/lora_finetune_distributed_td.py

+                        global_streaming=streaming,
+                        global_num_workers=num_workers,
+                    ),
+                    prefetch_factor=8,


would it make sense to add it to setup_data as a default or exposed in the config? Or is this parameter not commonly changed?

Updated PR: this is now exposed

Wait, why isn't this showing the new version? Let me make sure I've pushed

felipemello1 · 2024-11-13T20:17:54Z

recipes/lora_finetune_distributed_td.py

+
+        return ds
+
+    def _setup_data_td(


This would require a bit of a refactoring, and probably not the main point of this pr, but it would be nice to have _setup_dataset and _setup_dataloader as two different methods. It should be easier to read and maintain

I've split this into setting up individual datasets (_setup_one_dataset) and the global dataloader/mixer set up.

felipemello1 · 2024-11-13T20:19:29Z

recipes/lora_finetune_distributed_td.py

+
+        # Instantiate collate_fn
+        if "left_pad_sequence" in collate_fn:
+            raise RuntimeError("left_pad_sequence collator is only for inference.")


I understand that this was already in the setup_data, but we usually try to fail fast and catch errors like this in the init.

felipemello1 · 2024-11-13T20:28:29Z

torchtune/datasets/_sft.py

+        if load_dataset_kwargs.get("streaming", False):
+            self._data = split_dataset_by_node(


I think that if we can decouple dataloader from dataset, it will be easier to maintain/work with it. For example, can we do something like:

MyDatasetDistributed = split_dataset_by_node(MyDataset) def split_dataset_by_node(...): assert hasattr(MyDataset, self._data)

or maybe SFTDataset can have getter

This needs to be removed actually

felipemello1 · 2024-11-13T20:34:40Z

recipes/lora_finetune_distributed_td.py

+        if len(cfg_datasets) == 1:
+            node = next(iter(datasets.values()))  # TODO: multi dataset
+        else:
+            node = MultiDatasetWeightedSampler(
+                source_nodes=datasets,
+                weights=weights,
+            )


n00b question: I understand that this is torchdata api, but i wonder if 'len(cfg_datasets) == 1' check should be inside of MultiDatasetWeightedSampler. I.e. do we need this if check?

felipemello1 · 2024-11-13T20:36:45Z

recipes/lora_finetune_distributed_td.py

+
+        log.info("TorchData nodes are initialized")
+
+        return node


'node' feels a bit weird, since we do: self._dataloader = self._setup_dataloader(...)

Should we rename it to dataloader?

felipemello1 · 2024-11-13T20:40:14Z

recipes/lora_finetune_distributed_td.py

+
+        # TODO: add multi-dataset mixer
+        if num_workers == 0:
+            _Mapper = Mapper


I am a bit torn about if getting mapper, sampler, pin memory, etc should be an utility shared across all recipes, or if it should be exposed. No strong opinion, just thinking outloud

felipemello1 · 2024-11-13T20:40:39Z

recipes/lora_finetune_distributed_td.py

+        if pin_memory:
+            node = PinMemory(node)
+        if num_workers > 0:
+            node = Prefetcher(node, 2)


is this PreFetcher=2 different than the previous prefect_factor=8?

felipemello1 · 2024-11-13T20:41:12Z

recipes/lora_finetune_distributed_td.py

+        All data related setup happens here. Currently this recipe only supports
+        Map-style Datasets which fit into memory and an option for random shuffling.
+        Samplers, iterable datasets, and streaming datasets are not supported.


i guess this needs to be updated

RdoubleA · 2024-11-19T20:30:47Z

recipes/configs/llama3_2_vision/11B_lora_td.yaml

+  prefetch_factor: 2
+seed: null
+
+multi_datasets:


personally this is too verbose to parse for me, and even in the recipe there are just too many nested dictionaries. Ideally, I would like to achieve this type of UI in the config for datasets:

datasets: - _component_: torchtune.datasets... weight: 1.0 subset: ... - _component_: torchtune.datasets... weight: 1.0 ...

or something similar so all I have to do is specify the dataset I want and the weight. As it is I have multi_datasets -> datasets -> dataset just to specify the dataset builder. Maybe this is very ideal, but other libraries such as Axolotl are able to do this.

I am aware there's a few challenges to having this:

MultiDatasetSampler requires passing in dictionaries for datasets and weights

weight is not a valid argument for instantiating dataset components

I'm wondering if there's a way we can do this all for the user in a builder. For example:

datasets: ListConfig for cfg_dataset in datasets: weights[k] = cfg_dataset.pop("weight") dataset[k] = Prefetcher(config.instantiate(cfg_dataset), prefetch_factor) dataloader = get_multi_datasets(datasets, weights, cfg_dataloader)

stop_criterion imo should be moved to the dataloader config

RdoubleA · 2024-11-19T20:34:03Z

recipes/lora_finetune_distributed_td.py

+        )
+        weights, datasets = {}, {}
+        cfg_datasets = cfg_multi_datasets.datasets
+        for k, cfg_and_weight in cfg_datasets.items():


all of this logic needs to be in the builder. we do not want to expose torchdata internals in each recipe. I am totally okay with creating a new file in torchtune/data that contains functions that set up torchdata dataloaders and nodes.

Another ideal that I'm curious if we can achieve, can we unify the UX for multi datasets and single datasets? i.e., if we had a get_dataloader method, you can pass in a single dataset or a multi dataset and the call is the same in the recipe regardless of what the user specifies in the config

RdoubleA · 2024-11-19T20:35:38Z

torchtune/datasets/_sft.py

+        ),
+    )
+
+    return node


this returns a node, not a dataloader right? Are users still able to access the underlying hugging face data?

RdoubleA · 2024-11-19T20:41:00Z

torchtune/datasets/_sft.py

+
+
+@requires_torchdata
+def SFTDatasetNode(  # noqa[N802]


seems we're moving away from the class-based dataset abstraction and more of a function that returns the node configured with user parameters.

Curious if it would be better UX to just abandon the SFTDataset class (after the full migration) and keep Transform classes for each kinda of dataset (Instruct, Chat, MM, SFT, Preference) which is passed into a generic node builder

RdoubleA · 2024-11-19T20:41:40Z

torchtune/datasets/_sft.py

+            dataset = dataset.shuffle(seed=seed)
+        node = IterableWrapper(dataset)
+    else:
+        sampler = DistributedSampler(


I'm confused, which of these nodes are specific to a single dataset vs global for the dataloader?

RdoubleA · 2024-11-19T20:42:53Z

torchtune/datasets/multimodal/_the_cauldron.py

+    streaming: bool = False,
+    shuffle: bool = False,
+    seed: int = 0,
+    num_workers: int = 0,


is specific params such as seed, shuffle, num_workers, etc for each individual dataset a valid use case? My understanding was you can specify this globally at the dataloader level

RdoubleA · 2024-11-19T20:44:20Z

recipes/lora_finetune_distributed_td.py

+
+        # Get global settings
+        shuffle = cfg_dataloader.shuffle
+        parallel_method = cfg_dataloader.get("parallel_method", "thread")


in general, I'm quite confused on which parameters belong in the "dataset" abstraction and which belong in the "dataloader" abstraction. As it is, it seems you are using these in both. I would prefer to make this distinction very clear, unless I am missing something you may need to configure per dataset

andrewkho added 3 commits October 30, 2024 09:28

add twfb logging

be1ca55

add dataloader and model only modes for throughput testing

845edae

couple sketches of torchdata.nodes

fb246d4

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 Oct 30, 2024

RdoubleA reviewed Oct 30, 2024

View reviewed changes

joecummings reviewed Oct 30, 2024

View reviewed changes

ebsmothers reviewed Nov 1, 2024

View reviewed changes

example set up for streaming

2f0c15f

andrewkho added 2 commits November 8, 2024 17:47

temp

7325865

new recipe ready

be1ac94

RdoubleA mentioned this pull request Nov 11, 2024

Request for PackedDataset support for mmap #1051

Open

add multi-dataset support

c6946e7

joecummings reviewed Nov 12, 2024

View reviewed changes

revert single-device recipes

7e8c802

andrewkho commented Nov 12, 2024

View reviewed changes

felipemello1 reviewed Nov 13, 2024

View reviewed changes

andrewkho added 4 commits November 13, 2024 19:02

refactor a bit

0763e6f

set seed for HF datasets distributed streaming shuffle

47e1e6a

update docstrings, delete dead code

3d25031

lint

b1b2ab6

RdoubleA reviewed Nov 19, 2024

View reviewed changes

		@@ -0,0 +1,1119 @@
		# Copyright (c) Meta Platforms, Inc. and affiliates.

		@@ -0,0 +1,131 @@
		# Config for multi-device LoRA finetuning in lora_finetune_distributed_td.py

		rank=int(os.environ.get("RANK", 0)),
		world_size=int(os.environ.get("WORLD_SIZE", 1)),

		if load_dataset_kwargs.get("streaming", False):
		self._data = split_dataset_by_node(

[draft] torchdata integration #1929

Are you sure you want to change the base?

[draft] torchdata integration #1929

Conversation

andrewkho commented Oct 30, 2024 • edited Loading

Benchmarking on 8xA100

Context

Changelog

Test plan

UX

pytorch-bot bot commented Oct 30, 2024 • edited Loading

🔗 Helpful Links

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

❗ 1 Active SEVs

❌ 2 New Failures, 4 Cancelled Jobs

RdoubleA left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

andrewkho Oct 30, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

ebsmothers left a comment

Choose a reason for hiding this comment

andrewkho commented Nov 1, 2024

andrewkho commented Nov 1, 2024

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

andrewkho Nov 12, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

felipemello1 left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

felipemello1 Nov 13, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

felipemello1 Nov 13, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

felipemello1 Nov 13, 2024 • edited Loading

Choose a reason for hiding this comment

felipemello1 Nov 13, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

andrewkho commented Oct 30, 2024 •

edited

Loading

pytorch-bot bot commented Oct 30, 2024 •

edited

Loading

andrewkho Oct 30, 2024 •

edited

Loading

andrewkho Nov 12, 2024 •

edited

Loading

felipemello1 Nov 13, 2024 •

edited

Loading

felipemello1 Nov 13, 2024 •

edited

Loading

felipemello1 Nov 13, 2024 •

edited

Loading

felipemello1 Nov 13, 2024 •

edited

Loading