Skip to content

Lhotse RIR augmentation is seeded from seed instead of shard_seed, so all ranks draw identical reverberation #16100

Description

@udsy19

Describe the bug

In nemo/collections/common/data/lhotse/dataloader.py, the RIR augmentation is the
only augmentation transform attached to the sampler that is seeded from
config.seed instead of config.shard_seed:

# nemo/collections/common/data/lhotse/dataloader.py:1011-1018
if config.rir_enabled:
    sampler = sampler.map(
        ReverbWithImpulseResponse(
            rir_recordings=RecordingSet.from_file(config.rir_path) if config.rir_path is not None else None,
            p=config.rir_prob,
            randgen=random.Random(config.seed),      # <-- here
        )
    )

Every sibling transform in the same function uses seed=config.shard_seed:

line transform seed argument
861 on-the-fly noise mix (cuts.mix) seed=config.shard_seed
939 DynamicBucketingSampler seed=config.shard_seed
962 DynamicCutSampler seed=config.shard_seed
995 LowpassUsingResampling seed=config.shard_seed
1007 ClippingTransform seed=config.shard_seed
1016 ReverbWithImpulseResponse randgen=random.Random(config.seed)
1032 Compress seed=config.shard_seed

shard_seed defaults to "trng" (:109) and is resolved inside each process
(resolve_seed("trng") is secrets.randbelow(2**31)), which is what makes ranks
differ. config.seed defaults to 0 (:162) and is resolved to a concrete
integer up front, before the sampler is built:

# :584-585  (and the same at :678 for the multi-config entry point)
config.seed = resolve_seed(config.seed)
fix_random_seed(config.seed)

so it is bit-identical on every data-parallel rank.

Net effect: every DP rank applies the same reverberation coin-flip sequence and
selects the same impulse responses for the n-th cut it sees. (Each rank still gets
a different subset of cuts — DDP sharding itself is fine — but the augmentation
decision sequence is duplicated across ranks.) Effective RIR augmentation
diversity is far lower than the config implies. Nothing errors; the only symptom is
a model that generalises worse than expected.

Steps/Code to reproduce bug

Build the sampler for rank 0 and rank 1 and compare the RNG the RIR transform holds.
This runs on CPU with no model and no real data.

import tempfile
from pathlib import Path

from lhotse import CutSet
from lhotse.testing.dummies import DummyManifest
from lhotse.dataset import LowpassUsingResampling, ReverbWithImpulseResponse
from omegaconf import OmegaConf

from nemo.collections.common.data.lhotse.dataloader import get_lhotse_dataloader_from_config


class Identity:
    def __getitem__(self, cuts):
        return cuts


root = Path(tempfile.mkdtemp())
cuts = DummyManifest(CutSet, begin_id=0, end_id=10, with_data=True)
for c in cuts:
    c.features = None
    c.custom = None
    c.supervisions[0].custom = None
p = root / "cuts.jsonl.gz"
cuts.save_audios(root / "audio").to_file(p)

config = OmegaConf.create(
    {
        "cuts_path": str(p),
        "batch_size": 2,
        "seed": 0,
        "shard_seed": "trng",   # the default
        "rir_enabled": True,
        "rir_prob": 0.5,
        "lowpass_enabled": True,   # a sibling augmentation, used as a control
        "lowpass_prob": 0.5,
    }
)


def build(rank):
    return get_lhotse_dataloader_from_config(
        config=config, global_rank=rank, world_size=2, dataset=Identity()
    ).sampler


def find(sampler, cls):
    return [t for t in sampler._transforms if isinstance(t, cls)][0]


s0, s1 = build(0), build(1)

rir0, rir1 = find(s0, ReverbWithImpulseResponse), find(s1, ReverbWithImpulseResponse)
print("RIR rank0   :", [round(rir0.random.random(), 6) for _ in range(6)])
print("RIR rank1   :", [round(rir1.random.random(), 6) for _ in range(6)])

lp0, lp1 = find(s0, LowpassUsingResampling), find(s1, LowpassUsingResampling)
print("lowpass rank0:", [round(lp0.rng.random(), 6) for _ in range(6)])
print("lowpass rank1:", [round(lp1.rng.random(), 6) for _ in range(6)])

Output (the lowpass values come from the OS TRNG and differ on every run; only the
rank0 != rank1 relationship is stable. The RIR values are random.Random(0) and are
reproducible):

RIR rank0   : [0.844422, 0.757954, 0.420572, 0.258917, 0.511275, 0.404934]
RIR rank1   : [0.844422, 0.757954, 0.420572, 0.258917, 0.511275, 0.404934]
lowpass rank0: [0.680268, 0.335713, 7.5e-05, 0.819685, 0.195258, 0.820385]
lowpass rank1: [0.620671, 0.550698, 0.337479, 0.266303, 0.173753, 0.753269]

The RIR streams are bit-identical (they are random.Random(0), i.e. the default
seed: 0). The sibling LowpassUsingResampling, seeded from shard_seed, differs
— that is the control showing the mechanism works and the harness is sound.

I also ran this with rank 0 and rank 1 in two genuinely separate OS processes and
iterated the samplers end to end. Because reverb_rir(affix_id=True) appends
_rvb to the cut id, the per-position decisions are directly readable:

rank0 cut order        : [0000, 0001, 0004, 0005, 0008, 0009, 0012, 0013, 0016, 0017]
rank1 cut order        : [0002, 0003, 0006, 0007, 0010, 0011, 0014, 0015, 0018, 0019]
rank0 reverb applied?  : [0, 0, 1, 1, 0, 1, 0, 1, 1, 0]
rank1 reverb applied?  : [0, 0, 1, 1, 0, 1, 0, 1, 1, 0]     <-- identical
rank0 lowpass applied? : [0, 0, 0, 0, 0, 1, 0, 1, 1, 0]
rank1 lowpass applied? : [1, 0, 0, 0, 0, 1, 0, 1, 0, 0]     <-- differs (control)

Expected behavior

RIR augmentation randomness is differentiated per rank/worker via shard_seed,
like every other online augmentation in the same function.

Environment overview

  • Environment location: bare-metal virtualenv
  • Method of install: source checkout of main @ 0361552a6

Environment details

  • OS: macOS 26.5.1 (darwin, arm64)
  • Python 3.10.18 / PyTorch 2.12.0 / lhotse 2.0.0a3

Additional context

CPU only; no GPU, model download, or real dataset needed.

Note that lhotse.dataset.ReverbWithImpulseResponse (the CutSet-level transform
imported here — not the same-named lhotse.augmentation AudioTransform) accepts
only a pre-built randgen: random.Random and has no seed parameter, so this
is not a kwarg rename. The siblings do self.rng = random.Random(resolve_seed(self.seed))
in their own __post_init__; the same has to be done at the call site here:

-                randgen=random.Random(config.seed),
+                randgen=random.Random(resolve_seed(config.shard_seed)),

resolve_seed is already imported in this module (:40). With an explicit integer
shard_seed this is a no-op (resolve_seed(n) == n), so reproducible runs and the
existing test_dataloader_with_synth_rir (which sets shard_seed: 0) are unaffected.

Related: open PR #14322 ("Flexible augmentation config for Lhotse dataloaders")
moves this block into a new augment.py and carries the config.seed line over
unchanged, but the new ReverbRIR dataclass it introduces there already does
rng = random.Random(resolve_seed(self.seed)) with seed: int | str = "trng" —
i.e. exactly the behaviour proposed above. This report is about the current
default path on main.

Happy to open a PR; I have the fix and a regression test ready.

Activity

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

No one assigned

    Type

    No type

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions