mirror of
https://github.com/huggingface/lerobot.git
synced 2026-07-10 03:21:54 +00:00
Pepijn
41166b39fb
* fix(datasets): expose a generator on EpisodeAwareSampler for distributed shuffle sync In distributed training, accelerate can only synchronize the shuffle permutation across ranks when the sampler exposes a generator attribute. EpisodeAwareSampler shuffled via the global torch RNG, so disjoint batch shards relied on every rank's global CPU RNG staying in lockstep forever; any rank-asymmetric RNG consumption (e.g. eval rollouts on the main process only) silently desynced the permutations and ranks trained on overlapping/missing samples. * fix(train): seed sampler generator and gate dataset download per node - Pass a generator seeded with cfg.seed to EpisodeAwareSampler so accelerator.prepare registers it as the synchronized RNG and the shuffle order is reproducible. - Gate the initial make_dataset call on is_local_main_process instead of is_main_process: the global main process only exists on node 0, so on every other node all local ranks were downloading the dataset and building the Arrow cache concurrently.
93 lines
3.9 KiB
Python
93 lines
3.9 KiB
Python
#!/usr/bin/env python
|
|
|
|
# Copyright 2024 The HuggingFace Inc. team. All rights reserved.
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
import logging
|
|
from collections.abc import Iterator
|
|
|
|
import torch
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class EpisodeAwareSampler:
|
|
def __init__(
|
|
self,
|
|
dataset_from_indices: list[int],
|
|
dataset_to_indices: list[int],
|
|
episode_indices_to_use: list | None = None,
|
|
drop_n_first_frames: int = 0,
|
|
drop_n_last_frames: int = 0,
|
|
shuffle: bool = False,
|
|
generator: torch.Generator | None = None,
|
|
):
|
|
"""Sampler that optionally incorporates episode boundary information.
|
|
|
|
Args:
|
|
dataset_from_indices: List of indices containing the start of each episode in the dataset.
|
|
dataset_to_indices: List of indices containing the end of each episode in the dataset.
|
|
episode_indices_to_use: List of episode indices to use. If None, all episodes are used.
|
|
Assumes that episodes are indexed from 0 to N-1.
|
|
drop_n_first_frames: Number of frames to drop from the start of each episode.
|
|
drop_n_last_frames: Number of frames to drop from the end of each episode.
|
|
shuffle: Whether to shuffle the indices.
|
|
generator: Generator used for shuffling. Exposing this attribute (even when None) lets
|
|
`accelerate` register it as the synchronized RNG in distributed training, so
|
|
every rank draws the same permutation and batch shards stay disjoint. When
|
|
None, shuffling falls back to the global torch RNG.
|
|
"""
|
|
if drop_n_first_frames < 0:
|
|
raise ValueError(f"drop_n_first_frames must be >= 0, got {drop_n_first_frames}")
|
|
if drop_n_last_frames < 0:
|
|
raise ValueError(f"drop_n_last_frames must be >= 0, got {drop_n_last_frames}")
|
|
|
|
indices = []
|
|
for episode_idx, (start_index, end_index) in enumerate(
|
|
zip(dataset_from_indices, dataset_to_indices, strict=True)
|
|
):
|
|
if episode_indices_to_use is None or episode_idx in episode_indices_to_use:
|
|
ep_length = end_index - start_index
|
|
if drop_n_first_frames + drop_n_last_frames >= ep_length:
|
|
logger.warning(
|
|
"Episode %d has %d frames but drop_n_first_frames=%d and "
|
|
"drop_n_last_frames=%d removes all frames. Skipping.",
|
|
episode_idx,
|
|
ep_length,
|
|
drop_n_first_frames,
|
|
drop_n_last_frames,
|
|
)
|
|
continue
|
|
indices.extend(range(start_index + drop_n_first_frames, end_index - drop_n_last_frames))
|
|
|
|
if not indices:
|
|
raise ValueError(
|
|
"No valid frames remain after applying drop_n_first_frames and drop_n_last_frames. "
|
|
"All episodes were either filtered out or had too few frames."
|
|
)
|
|
|
|
self.indices = indices
|
|
self.shuffle = shuffle
|
|
self.generator = generator
|
|
|
|
def __iter__(self) -> Iterator[int]:
|
|
if self.shuffle:
|
|
for i in torch.randperm(len(self.indices), generator=self.generator):
|
|
yield self.indices[i]
|
|
else:
|
|
for i in self.indices:
|
|
yield i
|
|
|
|
def __len__(self) -> int:
|
|
return len(self.indices)
|