Why Shuffle at the beginning of each inner epoch

[KZF-kzf]

Hello, I noticed that you performed a shuffle operation in the

Flow-Factory/src/flow_factory/trainers/nft.py

file. Could you explain the purpose of this? It seems that the Flow-GRPO code doesn't have a similar operation.

            # Shuffle samples at the beginning of each inner epoch
            perm_gen = create_generator(self.training_args.seed, self.epoch, inner_epoch)
            perm = torch.randperm(len(samples), generator=perm_gen)
            shuffled_samples = [samples[i] for i in perm]
            
            # Re-group samples into batches
            sample_batches: List[Dict[str, Union[torch.Tensor, Any, List[Any]]]] = [
                BaseSample.stack(shuffled_samples[i:i + self.training_args.per_device_batch_size])
                for i in range(0, len(shuffled_samples), self.training_args.per_device_batch_size)
            ]

[Jayce-Ping]

There is. All Flow-Factory, the original flow_grpo and nft code shuffle the samples:

https://github.com/yifan123/flow_grpo/blob/9745241cf1101407e9b518dd4f10498092150522/scripts/train_sd3.py#L831-L846

Flow-Factory. GRPO:

Flow-Factory/src/flow_factory/trainers/grpo.py

Lines 185 to 197 in 4f24c3f

	for inner_epoch in range(self.training_args.num_inner_epochs):
	# Shuffle samples at the beginning of each inner epoch
	perm_gen = create_generator(self.training_args.seed, self.epoch, inner_epoch)
	perm = torch.randperm(len(samples), generator=perm_gen)
	shuffled_samples = [samples[i] for i in perm]
	# Create batches for optimization
	# `BaseSample.stack` will try to stack all tensor fields,
	# stack non-tensor fields as a list, keep shared fields as single value
	sample_batches : List[Dict[str, Union[torch.Tensor, Any, List[Any]]]] = [
	BaseSample.stack(shuffled_samples[i:i + self.training_args.per_device_batch_size])
	for i in range(0, len(shuffled_samples), self.training_args.per_device_batch_size)
	]

NFT:

Flow-Factory/src/flow_factory/trainers/nft.py

Lines 239 to 249 in 4f24c3f

	for inner_epoch in range(self.training_args.num_inner_epochs):
	# Shuffle samples at the beginning of each inner epoch
	perm_gen = create_generator(self.training_args.seed, self.epoch, inner_epoch)
	perm = torch.randperm(len(samples), generator=perm_gen)
	shuffled_samples = [samples[i] for i in perm]
	# Re-group samples into batches
	sample_batches: List[Dict[str, Union[torch.Tensor, Any, List[Any]]]] = [
	BaseSample.stack(shuffled_samples[i:i + self.training_args.per_device_batch_size])
	for i in range(0, len(shuffled_samples), self.training_args.per_device_batch_size)
	]

I believe the purpose is just to avoid the potential bias that if all samples within a group is on the same rank.

[KZF-kzf]

There is. All Flow-Factory, the original flow_grpo and nft code shuffle the samples:

https://github.com/yifan123/flow_grpo/blob/9745241cf1101407e9b518dd4f10498092150522/scripts/train_sd3.py#L831-L846

Flow-Factory. GRPO:

Flow-Factory/src/flow_factory/trainers/grpo.py

Lines 185 to 197 in 4f24c3f

for inner_epoch in range(self.training_args.num_inner_epochs):
# Shuffle samples at the beginning of each inner epoch
perm_gen = create_generator(self.training_args.seed, self.epoch, inner_epoch)
perm = torch.randperm(len(samples), generator=perm_gen)
shuffled_samples = [samples[i] for i in perm]

 # Create batches for optimization 
 # `BaseSample.stack` will try to stack all tensor fields, 
 # stack non-tensor fields as a list, keep shared fields as single value 
 sample_batches : List[Dict[str, Union[torch.Tensor, Any, List[Any]]]] = [ 
     BaseSample.stack(shuffled_samples[i:i + self.training_args.per_device_batch_size]) 
     for i in range(0, len(shuffled_samples), self.training_args.per_device_batch_size) 
 ] 

NFT:

Flow-Factory/src/flow_factory/trainers/nft.py

Lines 239 to 249 in 4f24c3f

for inner_epoch in range(self.training_args.num_inner_epochs):
# Shuffle samples at the beginning of each inner epoch
perm_gen = create_generator(self.training_args.seed, self.epoch, inner_epoch)
perm = torch.randperm(len(samples), generator=perm_gen)
shuffled_samples = [samples[i] for i in perm]

 # Re-group samples into batches 
 sample_batches: List[Dict[str, Union[torch.Tensor, Any, List[Any]]]] = [ 
     BaseSample.stack(shuffled_samples[i:i + self.training_args.per_device_batch_size]) 
     for i in range(0, len(shuffled_samples), self.training_args.per_device_batch_size) 
 ] 

I believe the purpose is just to avoid the potential bias that if all samples within a group is on the same rank.

thanks alot!