Same retain samples across ranks

[ZeguanXiao]

Information

• The official example scripts
• My own modified scripts

Tasks

• An officially supported task
• My own task or dataset (give details below)

Reproduction

Any official scripts.

Expected behavior

During distributed training with multiple GPUs, all ranks produce identical retain loss values while forget loss values differ as expected. This indicates that all ranks are sampling the same retain data samples, which breaks the randomness assumption in the unlearning process.

The issue is located in src/data/unlearn.py in the ForgetRetainDataset.__getitem__() method:

def __getitem__(self, idx):
    item = {}
    if self.anchor == "forget":
        item["forget"] = self.forget[idx]  # Sequential access - different across ranks ✓
        if self.retain:
            retain_idx = torch.randint(0, len(self.retain), (1,)).item()  # ❌ Problem here
            item["retain"] = self.retain[retain_idx]

The problem: torch.randint() uses the global PyTorch random number generator, which is synchronized across all ranks. This causes all ranks to generate the same sequence of random indices, leading to identical retain sample selection.

[Dornavineeth]

This is a good catch. We can set the seed here based on the ranks (based on GPT). Can you test this and raise a PR?

g = torch.Generator()
rank = torch.distributed.get_rank() if torch.distributed.is_initialized() else 0
seed = int(torch.empty((), dtype=torch.int64).random_().item() + rank)
g.manual_seed(seed)
retain_idx = torch.randint(0, len(self.retain), (1,), generator=g).item()

[ZeguanXiao]

I have conducted preliminary tests, and the fix does indeed work. However, it should be noted that this may cause the experimental results to differ from those before the fix, which is expected. Some experimental hyperparameters may need to be re-tuned or searched.