fix: random sampling in ForgetRetainDataset

[ZeguanXiao]

What does this PR do?

Fixes #139

Before submitting

• This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
• Have you gone through the contributions guide?
• Are your changes documented? Read documentation guidelines here.

[molereddy]

it would be better to use the seed from the experiment config here, rather than
int(torch.empty((), dtype=torch.int64).random_().item() to avoid introducing randomness uncontrolled by the seed.

can you try to see if you can make the experiment's cfg.seed available to this dataset class and then use seed = exp_seed + rank here?

[molereddy]

Thank you for the PR! Please see comment

[ZeguanXiao]

Thanks for the feedback! I've updated the PR accordingly. Please let me know if there are any further adjustments required.

[molereddy]

Please fix the lint errors!

[molereddy]

It is not ideal to set the seed at the exact example level. This would mean we select the same retain example index sequences even if we are using a different dataset.

Since the point is that each rank must get a different seed, imo it is better to get the rank in the global seed function: https://github.com/locuslab/open-unlearning/blob/main/src/trainer/utils.py#L8

Let me know if you see any issues.

cc @Dornavineeth

[ZeguanXiao]

@molereddy Simply modifying seed_everything() doesn’t work, because the process group is only initialized after the Accelerator is created — that is, after the trainer is initialized (in src/train.py). Calling seed_everything() again after the trainer has been initialized can achieve the desired effect. Do you think this implementation is acceptable, or do you have a better suggestion?

[ZeguanXiao]

@molereddy Currently, my implementation adds a torch.Generator instance to DataCollatorForSupervisedDataset.
After the process group is initialized (i.e., after the trainer is initialized in src/train.py),
I set the seed of the torch.Generator instance for each rank’s DataCollatorForSupervisedDataset object.

Could you please check if this approach is feasible/correct?

[filyp]

If the goal is simply to have different idx for different ranks, there's also a simpler solution:

        rank = torch.distributed.get_rank() if torch.distributed.is_initialized() else 0
        retain_idx = torch.randint(0, len(self.retain), (1,)).item()  # like before
        retain_idx = (retain_idx + rank) % len(self.retain)  # fans out the indices

Same for forget_idx. And that's the only code change that would be needed. Or for even more scrambling maybe instead of adding rank, add some hash of rank + idx. That should be as good as random, right?

[chengjiali]

Hi, is there any estimate on this?