Do not materialize entire randperm in RandomSampler

[aviverma01]

In our DDP training workloads, each rank was initializing a RandomSampler for a dataset with a length of 3.5 billion items. We noticed that when this sampler was in scope, gc.collect calls were taking on the order of seconds to run, which would slow down the entire training iteration. This is because when we call torch.randperm(n).tolist(), we create a python list of 3.5 billion items, which massively slows down the periodic mark & sweep garbage collection.

This PR swaps out the .tolist() call with a .numpy() call and manually calls .item() on each element as it is being requested. This has two benefits:

1. The first call to RandomSampler::__next__ should be about twice as fast, since .numpy does not copy the contents of the original tensor
2. The runtime of gc.collect() calls no longer scales linearly with the size of the dataset passed to RandomSampler

I've attached some timeit samples to illustrate the speedups with this Pr:

Main (no GC):  51.72115747816861
Main (10 GC calls) 83.61965207383037
PR (no GC) 33.06403830461204
PR (10 GC calls) 33.959467427805066

Code

from timeit import timeit


baseline_no_gc = """
import torch

n = int(1e9)
steps = n // 100

x = torch.randperm(n).tolist()
x_iter = iter(x)

for i in range(steps):
    next(x_iter)
"""


baseline_gc = """
import torch
import gc
n = int(1e9)
steps = n // 100
gc_every = steps // 10

x = torch.randperm(n).tolist()
x_iter = iter(x)

for i in range(steps):
    next(x_iter)
    if i % gc_every == 0:
        gc.collect()
"""


numpy_no_gc = """
import torch
n = int(1e9)
steps = n // 100

x = torch.randperm(n).numpy()
x_iter = (i.item() for i in x)

for i in range(steps):
    next(x_iter)
"""

numpy_gc = """
import torch
import gc
n = int(1e9)
steps = n // 100
gc_every = steps // 10

x = torch.randperm(n).numpy()
x_iter = (i.item() for i in x)

for i in range(steps):
    next(x_iter)
    if i % gc_every == 0:
        gc.collect()
"""


if __name__ == "__main__":
    print("Main (no GC): ", timeit(baseline_no_gc, number=1))
    print("Main (10 GC calls)", timeit(baseline_gc, number=1))
    print("PR (no GC)",  timeit(numpy_no_gc, number=1))
    print("PR (10 GC calls)", timeit(numpy_gc, number=1))

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/103339

• 📄 Preview Python docs built from this PR
• 📄 Preview C++ docs built from this PR
• ❓ Need help or want to give feedback on the CI? Visit the bot commands wiki or our office hours

Note: Links to docs will display an error until the docs builds have been completed.

✅ No Failures

As of commit 8f5e445 with merge base bc2caa7 ():
💚 Looks good so far! There are no failures yet. 💚

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[linux-foundation-easycla]

The committers listed above are authorized under a signed CLA.

• ✅ login: aviverma01 (7cf7c49, 6af9178, d9f151b, fe86601, 8f5e445)

[vadimkantorov]

one minor issue is that repeated .item() are slow #29973

so maybe tolist() is actually not that bad? (espeically if indices itself is materialized as tensor)

[aviverma01]

Thanks for the suggestion! I've updated the code to call .numpy on the tensor before iterating on it, which should avoid the slow .item() calls

[vadimkantorov]

Well, if you can afford an extra allocation, why not just yield from indices.tolist()? Because the indices Python list would take too much memory?

I don't know what's the current state of affaires on obligation of numpy dependency

[vadimkantorov]

Maybe worth filing a feature request to iterate Python items? Or supporting memoryview on tensors (so that it can be iterated)...

[aviverma01]

I've added some comments in the PR description, but the main issue is that calling .tolist() on a torch tensor of size 1billion+ adds a massive garbage collection overhead because we just allocated billions of individual python int objects that need to be managed separately. By using a numpy array instead, the garbage collector only needs to keep track of 1 object regardless of the dataset size, making garbage collection much faster during training.

[vadimkantorov]

i created also #103352

[vadimkantorov]

could this just be yield from torch.randperm(n, generator=generator).numpy() (+ some indexing) keeping the existing terse syntax?

[aviverma01]

Done in 7a7a1e9, note that there's an additional map call compared to the original comment to ensure that the type matches

[kit1980]

@aviverma01 Should we also change the other branch, yield from torch.randint(high=n, size=(32,), dtype=torch.int64, generator=generator).tolist() ?

[aviverma01]

Done in b450690

[facebook-github-bot]

@kit1980 has imported this pull request. If you are a Meta employee, you can view this diff on Phabricator.

[kit1980]

I've triggered more tests and also imported this internally to make sure nothing breaks.

[aviverma01]

I've triggered more tests and also imported this internally to make sure nothing breaks.

Thanks @kit1980, would you able to help fix the "Meta Internal-Only Changes Check"? Also I think some of the tests may be flakey. Would it be possible to re-kick the failing tests?

[kit1980]

@pytorchbot rebase

[pytorchmergebot]

@pytorchbot started a rebase job onto refs/remotes/origin/viable/strict. Check the current status here

[pytorchmergebot]

Successfully rebased patch-1 onto refs/remotes/origin/viable/strict, please pull locally before adding more changes (for example, via git checkout patch-1 && git pull --rebase)

[aviverma01]

@kit1980 Looks like there may be some remaining flakey tests, and I believe the "Meta Internal-Only" Changes Check is still failing. Any chance you could help/show me how to fix it?

[facebook-github-bot]

@kit1980 has imported this pull request. If you are a Meta employee, you can view this diff on Phabricator.

[kit1980]

@mergebot merge -i

[aviverma01]

@mergebot merge -i

@kit1980 do you know why the pytorchbot didn't pick up the merge request?

[kit1980]

@aviverma01 sorry, I misspelled the bot name.

[kit1980]

@pytorchmergebot merge

[pytorchmergebot]

Merge started

Your change will be merged once all checks pass (ETA 0-4 Hours).

Learn more about merging in the wiki.

Questions? Feedback? Please reach out to the PyTorch DevX Team

Advanced Debugging

Check the merge workflow status
here

[qqaatw]

Hi @aviverma01 @kit1980, the change this PR introduces seems to block the usage of manually specifying DataLoader's generator with a non-CPU device due to the numpy() operations:

gen = torch.Generator(device=torch.device("mps:0"))
data_loader = data.DataLoader(dataset, batch_size=8, shuffle=True, generator=gen)

The error:

TypeError: can't convert mps:0 device type tensor to numpy. Use Tensor.cpu() to copy the tensor to host memory first.

Was it intended to be? Or do you have any idea on this?

Thanks :)

[kit1980]

@pytorchbot revert -m "Cause issues on MPS, and also fails without numpy" -c nosignal

[kit1980]

I'm reverting this.

I've realized there is another issue with the PR, in fails without numpy, which is an optional dependency actually.

[pytorchmergebot]

@pytorchbot successfully started a revert job. Check the current status here.
Questions? Feedback? Please reach out to the PyTorch DevX Team

[pytorchmergebot]

@aviverma01 your PR has been successfully reverted.