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"):
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 zeroThe 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.
BatchStratifiedSampler.__init__(roll/datasets/sampler.py) already handles a domain indomain_ratiosthat has zero matching rows: it deletes the domain from bothself.domain_indicesandself.domain_ratiosand 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 prunedself.domain_ratios, so the just-deleted domain comes back.domain_listanddomain_batch_numare built from that dict, so the empty domain gets a batch count.self.domain_indicesis adefaultdict(list), so the next access to it for that domain (buildingdomain_batch_capacities) silently creates an empty list instead of raising. By__iter__, that domain has a positive batch count and zero indices, andrepeat_times = (total_required + len(indices) - 1) // len(indices)divides by zero.Repro (dataset with only domains "a" and "b"):
The print says "c is empty, delete in sampling", then
sampler.domain_ratiosstill hasc: 0.2right after__init__returns. Training can't start at all whenever adomain_ratiosconfig 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_ratioson line 39 instead of the constructor'sdomain_ratiosargument. Happy to send a PR if useful.