[bugfix] fix batch_sampler set_epoch

[Jintao-Huang]

#9658

[gemini-code-assist]

Code Review

This pull request updates the set_epoch method in swift/dataloader/shard.py to propagate the epoch setting to nested batch samplers. The reviewer pointed out that the proposed implementation still fails to propagate set_epoch to the underlying sampler in common scenarios, such as a standard BatchSampler wrapping a DistributedSampler. They suggested a robust, recursive traversal approach to ensure set_epoch is correctly called on all nested samplers.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

[gemini-code-assist]

The current implementation fails to propagate set_epoch to the underlying sampler in several common scenarios:

1. Standard BatchSampler wrapping a DistributedSampler:

   • self.batch_sampler is not None but does not have set_epoch or batch_sampler attributes.
   • self.sampler is None (as PyTorch's DataLoader sets self.sampler = None when batch_sampler is specified).
   • As a result, set_epoch is never called on the underlying DistributedSampler.

2. BatchSamplerShard (from HF Accelerate) wrapping a standard BatchSampler wrapping a DistributedSampler:

   • self.batch_sampler has a batch_sampler attribute, but that wrapped BatchSampler does not have set_epoch (only its sampler does).
   • Again, set_epoch is never called on the underlying DistributedSampler.

To comprehensively and robustly solve this, we should recursively traverse both batch_sampler and sampler attributes to find and call set_epoch on any objects that support it.

Suggested change

	if self.batch_sampler is not None:
	if hasattr(self.batch_sampler, 'set_epoch'):
	self.batch_sampler.set_epoch(epoch)
	if hasattr(self.batch_sampler, 'batch_sampler') and hasattr(self.batch_sampler.batch_sampler, 'set_epoch'):
	self.batch_sampler.batch_sampler.set_epoch(epoch)
	elif self.sampler is not None and hasattr(self.sampler, 'set_epoch'):
	self.sampler.set_epoch(epoch)
	visited = set()
	def _set_epoch(obj):
	if obj is None or id(obj) in visited:
	return
	visited.add(id(obj))
	if hasattr(obj, 'set_epoch'):
	obj.set_epoch(epoch)
	if hasattr(obj, 'batch_sampler'):
	_set_epoch(obj.batch_sampler)
	if hasattr(obj, 'sampler'):
	_set_epoch(obj.sampler)
	_set_epoch(self.batch_sampler)
	_set_epoch(self.sampler)