Support DDP partitioning for prebaked datasets

[kylesayrs]

Summary

• When using a prebaked dataset string (e.g. dataset="perfectblend") in a distributed setting, automatically partition the split across ranks via get_rank_partition() so each rank only loads its slice of calibration data.
• This enables DDP examples to use oneshot(model, dataset="perfectblend", ...) without manually calling load_dataset + get_rank_partition().
• Stacked on Add "perfectblend" prebaked dataset and "ultrachat" alias #3060.

Test plan

• Run DDP quantization with dataset="perfectblend" using torchrun --nproc_per_node=2
• Verify each rank loads a disjoint partition of the calibration data
• Verify non-DDP usage is unchanged

🤖 Generated with Claude Code

[coderabbitai]

Important

Review skipped

Auto reviews are disabled on this repository. Please check the settings in the CodeRabbit UI or the .coderabbit.yaml file in this repository. To trigger a single review, invoke the @coderabbitai review command.

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro Plus

Run ID: d68c75ef-86c4-4695-896e-18f277b1743d

You can disable this status message by setting the reviews.review_status to false in the CodeRabbit configuration file.

Use the checkbox below for a quick retry:

• 🔍 Trigger review

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}

[gemini-code-assist]

Code Review

This pull request introduces dataset split partitioning across ranks in a distributed data parallel (DDP) setting for prebaked (string) datasets, ensuring each rank only loads its respective slice of calibration data. The reviewer pointed out that while the dataset split is partitioned, the global dataset_args.num_calibration_samples is not updated to reflect the reduced sample count per rank. This discrepancy triggers a false warning in _make_sampler regarding insufficient samples. The reviewer provided a code suggestion to update dataset_args.num_calibration_samples based on the partitioned range for the current rank.

[gemini-code-assist]

When partitioning the dataset split across ranks in a distributed setting, the number of samples loaded per rank is reduced (e.g., to num_calibration_samples / world_size). However, dataset_args.num_calibration_samples remains at its original global value. Later in _make_sampler, this discrepancy triggers a false warning: Requested X samples but the provided dataset only has Y samples.

To prevent this false warning and ensure consistency, we should update dataset_args.num_calibration_samples to reflect the partitioned sample count for the current rank.

Suggested change

	if (
	isinstance(dataset_args.dataset, str)
	and dist.is_initialized()
	and split_str is not None
	and "[" not in split_str
	and dataset_args.num_calibration_samples is not None
	):
	split_str = get_rank_partition(
	split_str, dataset_args.num_calibration_samples
	)
	logger.info(
	f"DDP: partitioned dataset split to '{split_str}' for rank "
	f"{dist.get_rank()}/{dist.get_world_size()}"
	)
	if (
	isinstance(dataset_args.dataset, str)
	and dist.is_initialized()
	and split_str is not None
	and "[" not in split_str
	and dataset_args.num_calibration_samples is not None
	):
	start, end = _get_partition_start_end(
	dataset_args.num_calibration_samples,
	dist.get_rank(),
	dist.get_world_size(),
	)
	split_str = get_rank_partition(
	split_str, dataset_args.num_calibration_samples
	)
	dataset_args.num_calibration_samples = end - start
	logger.info(
	f"DDP: partitioned dataset split to '{split_str}' for rank "
	f"{dist.get_rank()}/{dist.get_world_size()}"
	)

[mergify]

Merge Protections

🔴 1 of 1 protections blocking · waiting on 👀 reviews

	Protection	Waiting on
🔴	Require one maintainer review	👀 reviews

🔴 Require one maintainer review

Waiting for any of

• approved-reviews-by=HDCharles
• approved-reviews-by=brian-dellabetta
• approved-reviews-by=dsikka
• approved-reviews-by=kylesayrs
• approved-reviews-by=yiliu30

This rule is failing.

All PRs must have at least one approving review from a maintainer before merging.

• any of:
  • approved-reviews-by=HDCharles
  • approved-reviews-by=brian-dellabetta
  • approved-reviews-by=dsikka
  • approved-reviews-by=kylesayrs
  • approved-reviews-by=yiliu30
• #changes-requested-reviews-by = 0