fix(dataloader): prevent data drop and padding during validation for accurate metrics#1100
Merged
garrett4wade merged 1 commit intoinclusionAI:mainfrom Mar 30, 2026
Merged
Conversation
…accurate metrics(pre-commit checked)
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Code Review
This pull request introduces a new EvalDistributedSampler class in areal/utils/dataloader.py designed for evaluation tasks. Unlike the standard DistributedSampler, this implementation avoids padding the dataset, ensuring that each sample is evaluated exactly once across the cluster even if the dataset size is not evenly divisible by the number of replicas. The create_dataloader function was updated to use this new sampler when a ValidDatasetConfig is provided. I have no feedback to provide.
garrett4wade
approved these changes
Mar 30, 2026
Collaborator
garrett4wade
left a comment
garrett4wade
left a comment
There was a problem hiding this comment.
@Anguo-star Thanks for the fix!
FYI the new EvalDistributedSampler class may not be mandatory with the single-controller usage (i.e., "python3 train.py" instead of the legacy "python3 -m areal.infra.launcher.local train.py"), because the data is only loaded in the controller process. The fix is good for backward compatiblity. Merging anyways.
rchardx
added a commit
that referenced
this pull request
Mar 30, 2026
…rash In single-controller mode with drop_last=False (as set by PR #1100 for accurate validation metrics), the last validation batch may contain fewer sequences than dp_size. This causes balanced_greedy_partition to raise "The length of nums must be divisible by K" during evaluation dispatch. Truncate trailing items in _dispatch_tensors when the batch size is not evenly divisible by dp_size, preserving compatibility with both drop_last=False validation and strict equal-partition dispatch.
15 tasks
rchardx
added a commit
that referenced
this pull request
Mar 30, 2026
…rash In single-controller mode with drop_last=False (as set by PR #1100 for accurate validation metrics), the last validation batch may contain fewer sequences than dp_size. This causes balanced_greedy_partition to raise "The length of nums must be divisible by K" during evaluation dispatch. Truncate trailing items in _dispatch_tensors when the batch size is not evenly divisible by dp_size, preserving compatibility with both drop_last=False validation and strict equal-partition dispatch.
rchardx
added a commit
that referenced
this pull request
Mar 31, 2026
…ivisible batches When drop_last=False, the last eval batch may not be divisible by dp_size, causing balanced_greedy_partition to crash (#1100/#1109). Instead of truncating remainder items (which biases eval metrics), pad with zero-contribution dummy items before dispatch so that all DP replicas participate in collective ops without deadlock. Key changes: - Add make_dummy_eval_item in areal/utils/data.py - Add _pad_eval_batch in train_controller.py for evaluate_* methods - Move total_loss_weight assertion after all_reduce (allow local zero) - Add zero-weight fast path in FSDP/Megatron/Archon engines - Add tests for padding, dummy schema, and zero-weight loss path Refs: #1100, #1109, #898
rchardx
added a commit
that referenced
this pull request
Apr 1, 2026
…ivisible batches When drop_last=False, the last eval batch may not be divisible by dp_size, causing balanced_greedy_partition to crash. Instead of truncating remainder items (which biases metrics), pad the batch with zero-contribution dummy items before dispatch. Extend dispatch to support group_size-aware partitioning so RW chosen/rejected pairs stay on the same DP rank. Key changes: - Add make_dummy_eval_item() for zero-weight schema-compatible items - Add _pad_eval_batch() for eval-only dummy padding before dispatch - Add zero-weight fast path in FSDP/Megatron/Archon engines - Move total_loss_weight assertion after all_reduce (allow local zero) - Add group_size-aware dispatch for RW chosen/rejected pairs - Exclude dummy sequences from RW loss weight via count_nonzero - Add 20+ tests covering padding, dispatch, and zero-weight paths Refs: #1095, #1100, #898
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
fix(dataloader): use EvalDistributedSampler for validation to ensure accurate metrics
Description
This PR fixes a critical bug where the validation dataloader was producing inaccurate evaluation metrics due to hardcoded data dropping and standard distributed padding.
Previously,
create_dataloaderhardcodeddrop_last=Trueand used the standardDistributedSamplerfor all dataset configs. The standardDistributedSamplerpads the dataset to make it evenly divisible by the number of replicas, which causes some validation samples to be evaluated twice. Combined withdrop_last=True, this means the validation set was both artificially inflated (via padding) and truncated (via dropping), leading to biased evaluation results.Changes made:
EvalDistributedSampler: A custom sampler designed specifically for evaluation that prevents dataset padding. It ensures every sample in the dataset is evaluated exactly once across the cluster.create_dataloader: Added logic to check if the config is aValidDatasetConfig. If true, it dynamically applies the newEvalDistributedSamplerand forcesdrop_last=False.Related Issue
Fixes #1095
Type of Change
Checklist
pre-commit run --all-files)./docs/build_all.sh)main/review-prcommand/create-prBreaking Change Details (if applicable):
Additional Context
Need help? Check the Contributing Guide or ask in
GitHub Discussions!