alibaba / alibaba/ROLL

BatchStratifiedSampler crashes with ZeroDivisionError when a domain_ratios entry has no matching rows

Open Beginner friendly
#495 0 comments 0 reactions 0 assignees View on GitHub
Dominant language
Python
Stars
3.4k
Forks
312
Avg merge
1h 2m
Merged PRs (30d)
2

Description

`BatchStratifiedSampler.__init__` (roll/datasets/sampler.py) already handles a domain in `domain_ratios` that has zero matching rows: it deletes the domain from both `self.domain_indices` and `self.domain_ratios` and prints "{key} is empty, delete in sampling." (lines 31-34).

Five lines later it undoes that: `self.domain_ratios = {key: value / sum_values for key, value in domain_ratios.items()}` (line 39) rebuilds from the original constructor argument, not from the pruned `self.domain_ratios`, so the just-deleted domain comes back. `domain_list` and `domain_batch_num` are built from that dict, so the empty domain gets a batch count. `self.domain_indices` is a `defaultdict(list)`, so the next access to it for that domain (building `domain_batch_capacities`) silently creates an empty list instead of raising. By `__iter__`, that domain has a positive batch count and zero indices, and `repeat_times = (total_required + len(indices) - 1) // len(indices)` divides by zero.

Repro (dataset with only domains "a" and "b"):

```python
ds = FakeDataset(["a"] * 10 + ["b"] * 10)
sampler = BatchStratifiedSampler(ds, domain_ratios={"a": 0.5, "b": 0.3, "c": 0.2}, batch_size=10, drop_last=True)
list(iter(sampler)) # ZeroDivisionError: integer division or modulo by zero
```

The print says "c is empty, delete in sampling", then `sampler.domain_ratios` still has `c: 0.2` right after `__init__` returns. Training can't start at all whenever a `domain_ratios` config names a domain that's absent from the current shard, which is exactly the case the deletion code was meant to handle.

Fix looks like reading from the already-pruned `self.domain_ratios` on line 39 instead of the constructor's `domain_ratios` argument. Happy to send a PR if useful.

Contributor guide

No contributing guide indexed for this repository

Research direction

Start in roll/datasets/sampler.py at BatchStratifiedSampler.__init__, especially the normalization around lines 31-39, then run the provided FakeDataset reproduction. Done means an absent domain stays removed from the sampler configuration and list(iter(sampler)) completes without ZeroDivisionError.

Written by the indexing model from the issue text.

Assessment

Tech stack
python
Domain
machine-learning
Issue type
Bug
Difficulty
1/5
Estimated time
Under an hour
Activity status
Active
Clarity
Clearly specified
Newbie friendliness
88/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.