[sirius8050]

Summary

Configure the PyTorch thread pool of each weight fetcher based on the total
number of fetcher processes running on the host.

This prevents multiple weight fetchers from oversubscribing the available CPUs
during TorchStore tensor assembly.

The change is applied to both the vLLM v0 and v1 generator implementations.

Problem

Weight prefetching can create multiple independent _WeightFetcher processes
through n_fetcher_procs, which defaults to 8.

The parameter keys are correctly partitioned between these processes, but the
fetcher count was not passed to _WeightFetcher. As a result, each process
configured or inherited its PyTorch/OpenMP thread pool without considering the
other fetchers running on the same host.

For example, if every fetcher creates a thread pool based on the full process
CPU affinity, eight fetchers can create approximately eight times more runnable
OpenMP threads than the host can execute concurrently.

This causes CPU oversubscription during TorchStore tensor assembly. Operations
such as torch.Tensor.copy_, used while assembling tensor shards, enter native
OpenMP regions and compete with the thread pools of the other fetcher processes.
Adding more fetchers can therefore make assembly slower instead of faster.

see: https://docs.pytorch.org/docs/2.13/notes/multiprocessing.html

Solution

Pass n_fetcher_procs to every _WeightFetcher when it is spawned.

Each fetcher runs a setup endpoint once after spawning and calculates its
PyTorch thread budget from the process CPU affinity:

cores = len(os.sched_getaffinity(0))
threads_per_fetcher = max(1, cores // n_fetcher_procs)

torch.set_num_threads(threads_per_fetcher)
torch.set_num_interop_threads(threads_per_fetcher)

[meta-cla]

Hi @sirius8050!

Thank you for your pull request and welcome to our community.

Action Required

In order to merge any pull request (code, docs, etc.), we require contributors to sign our Contributor License Agreement, and we don't seem to have one on file for you.

Process

In order for us to review and merge your suggested changes, please sign at https://code.facebook.com/cla. If you are contributing on behalf of someone else (eg your employer), the individual CLA may not be sufficient and your employer may need to sign the corporate CLA.

Once the CLA is signed, our tooling will perform checks and validations. Afterwards, the pull request will be tagged with CLA signed. The tagging process may take up to 1 hour after signing. Please give it that time before contacting us about it.

If you have received this in error or have any questions, please contact us at cla@meta.com. Thanks!

[meta-cla]

Thank you for signing our Contributor License Agreement. We can now accept your code for this (and any) Meta Open Source project. Thanks!