Fix PyTorch DDP sampler epoch forwarding in the data loader wrappers

[Tendourisu]

Summary

Fix PyTorch DDP sampler epoch forwarding in the data loader wrappers.

In DDP mode, create_torch_data_loader() creates a DistributedSampler and passes it into TorchDataLoader, but the training loop receives the outer DataLoaderImpl. Previously neither TorchDataLoader nor DataLoaderImpl exposed set_epoch() or __len__(), so the existing guarded call in scripts/train_pytorch.py never executed. As a result, in DDP training that spans multiple full dataset cycles, the sampler kept reusing the epoch 0 index order instead of reshuffling each epoch.

Background

Relevant code paths:

• The training loop contains the guarded set_epoch() call:
  

  openpi/scripts/train_pytorch.py

  Lines 510 to 512 in 15a9616

  	# Set epoch for distributed training
  	if use_ddp and hasattr(loader, "set_epoch"):
  	loader.set_epoch(global_step // len(loader))

• A DistributedSampler is created and passed to TorchDataLoader here:
  

  openpi/src/openpi/training/data_loader.py

  Lines 305 to 337 in 15a9616

  	# Use TorchDataLoader for both frameworks
  	# For PyTorch DDP, create DistributedSampler and divide batch size by world size
  	# For JAX, divide by process count
  	sampler = None
  	if framework == "pytorch":
  	if torch.distributed.is_initialized():
  	sampler = torch.utils.data.distributed.DistributedSampler(
  	dataset,
  	num_replicas=torch.distributed.get_world_size(),
  	rank=torch.distributed.get_rank(),
  	shuffle=shuffle,
  	drop_last=True,
  	)
  	local_batch_size = batch_size // torch.distributed.get_world_size()
  	else:
  	local_batch_size = batch_size
  	else:
  	local_batch_size = batch_size // jax.process_count()
  	logging.info(f"local_batch_size: {local_batch_size}")
  	data_loader = TorchDataLoader(
  	dataset,
  	local_batch_size=local_batch_size,
  	sharding=None if framework == "pytorch" else sharding,
  	shuffle=(sampler is None and shuffle), # Don't shuffle if using sampler
  	sampler=sampler,
  	num_batches=num_batches,
  	num_workers=num_workers,
  	seed=seed,
  	framework=framework,
  	)
  	return DataLoaderImpl(data_config, data_loader)

• TorchDataLoader does not implement set_epoch() or __len__():
  

  openpi/src/openpi/training/data_loader.py

  Lines 381 to 468 in 15a9616

  	class TorchDataLoader:
  	"""Torch data loader implementation."""
  	def __init__(
  	self,
  	dataset,
  	local_batch_size: int,
  	*,
  	sharding: jax.sharding.Sharding | None = None,
  	shuffle: bool = False,
  	sampler: torch.utils.data.Sampler | None = None,
  	num_batches: int | None = None,
  	num_workers: int = 0,
  	seed: int = 0,
  	framework: str = "jax",
  	):
  	"""Create a PyTorch data loader.
  	Args:
  	dataset: The dataset to load.
  	local_batch_size: The local batch size for each process.
  	sharding: The sharding to use for the data loader.
  	shuffle: Whether to shuffle the data.
  	num_batches: If provided, determines the number of returned batches. If the
  	number is larger than the number of batches in the dataset, the data loader
  	will loop over the dataset. If not provided, will iterate over the dataset
  	indefinitely.
  	num_workers: The number of worker processes to use. If zero, the data loader will
  	execute in the main process.
  	seed: The seed to use for shuffling the data.
  	"""
  	if jax.process_count() > 1:
  	raise NotImplementedError("Data loading with multiple processes is not supported.")
  	if len(dataset) < local_batch_size:
  	raise ValueError(f"Local batch size ({local_batch_size}) is larger than the dataset size ({len(dataset)}).")
  	# Store sharding - None for PyTorch, JAX sharding for JAX
  	self._sharding = sharding
  	if sharding is None and framework == "jax":
  	# Use data parallel sharding by default for JAX only.
  	self._sharding = jax.sharding.NamedSharding(
  	jax.sharding.Mesh(jax.devices(), ("B",)),
  	jax.sharding.PartitionSpec("B"),
  	)
  	self._num_batches = num_batches
  	mp_context = None
  	if num_workers > 0:
  	mp_context = multiprocessing.get_context("spawn")
  	generator = torch.Generator()
  	generator.manual_seed(seed)
  	self._data_loader = torch.utils.data.DataLoader(
  	typing.cast(torch.utils.data.Dataset, dataset),
  	batch_size=local_batch_size,
  	shuffle=(sampler is None and shuffle), # Don't shuffle if using sampler
  	sampler=sampler,
  	num_workers=num_workers,
  	multiprocessing_context=mp_context,
  	persistent_workers=num_workers > 0,
  	collate_fn=_collate_fn,
  	worker_init_fn=_worker_init_fn,
  	drop_last=True,
  	generator=generator,
  	)
  	@property
  	def torch_loader(self) -> torch.utils.data.DataLoader:
  	return self._data_loader
  	def __iter__(self):
  	num_items = 0
  	while True:
  	data_iter = iter(self._data_loader)
  	while True:
  	if self._num_batches is not None and num_items >= self._num_batches:
  	return
  	try:
  	batch = next(data_iter)
  	except StopIteration:
  	break # We've exhausted the dataset. Create a new iterator and start over.
  	num_items += 1
  	# For JAX, convert to sharded arrays; for PyTorch, return torch tensors
  	if self._sharding is not None:
  	yield jax.tree.map(lambda x: jax.make_array_from_process_local_data(self._sharding, x), batch)
  	else:
  	yield jax.tree.map(torch.as_tensor, batch)

• The outer DataLoaderImpl wrapper also does not implement either method:
  

  openpi/src/openpi/training/data_loader.py

  Lines 530 to 540 in 15a9616

  	class DataLoaderImpl(DataLoader):
  	def __init__(self, data_config: _config.DataConfig, data_loader: TorchDataLoader | RLDSDataLoader):
  	self._data_config = data_config
  	self._data_loader = data_loader
  	def data_config(self) -> _config.DataConfig:
  	return self._data_config
  	def __iter__(self):
  	for batch in self._data_loader:
  	yield _model.Observation.from_dict(batch), batch["actions"]

Changes

• Add set_epoch() and __len__() to TorchDataLoader
• Forward set_epoch() and __len__() from DataLoaderImpl