Skip to content

Fix iterable dataset dict num proc shards #8062

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
HaukurPall wants to merge 2 commits into
base: main
Choose a base branch
from
Open

Fix iterable dataset dict num proc shards #8062

Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
31a06e6
Limit num_proc based on num_shards for IterableDatasetDict push_to_hub
HaukurPall Mar 11, 2026
624e558
Merge branch 'huggingface:main' into fix-iterable-dataset-dict-num-pr…
HaukurPall Mar 12, 2026
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
15 changes: 14 additions & 1 deletion src/datasets/dataset_dict.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2646,6 +2646,19 @@ def push_to_hub(
additions = []
for split in self.keys():
logger.info(f"Pushing split {split} to the Hub.")
split_num_proc = num_proc
split_dataset = self[split]
if split_num_proc is not None and split_num_proc > split_dataset.num_shards:
logger.warning(
f"Too many num_proc: {split_num_proc} (max is dataset.num_shards={split_dataset.num_shards}). "
f"Stopping {split_num_proc - split_dataset.num_shards} processes."
)
logger.info(
f"To parallelize data loading, we give each process some shards (or data sources) to process. "
f"Therefore it's unnecessary to have a number of processes greater than dataset.num_shards={split_dataset.num_shards}. "
f"To enable more parallelism, please split the dataset in more files than {split_dataset.num_shards}."
)
split_num_proc = split_dataset.num_shards
# The split=key needs to be removed before merging
split_additions, uploaded_size, dataset_nbytes, num_examples = self[split]._push_parquet_shards_to_hub(
repo_id,
Expand All @@ -2657,7 +2670,7 @@ def push_to_hub(
# max_shard_size=max_shard_size, # TODO(QL): add arg
num_shards=num_shards.get(split),
embed_external_files=embed_external_files,
num_proc=num_proc,
num_proc=split_num_proc,
)
additions += split_additions
total_uploaded_size += uploaded_size
Expand Down