Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Fix CPU and memory affinity under external resource management #3012

Open
wants to merge 1 commit into
base: main
Choose a base branch
from
Open

Fix CPU and memory affinity under external resource management #3012

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
33 changes: 33 additions & 0 deletions server/text_generation_server/models/flash_causal_lm.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -102,6 +102,39 @@ def get_sliding_windows() -> int:


def init_cpu_threads_env(rank_id: int, world_size: int):
import psutil
allowed_cpus = psutil.Process().cpu_affinity()
if len(allowed_cpus) < psutil.cpu_count(logical=True):
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

is allowed cpus logical cpu numbers? why is physical cpu numbers not used?

Copy link
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

First, we want to detect if there is CPU affinity already enforced to this process. In other words, we want to know if somebody (an external entity) has already chosen the CPUs that this process can use.

In order to check this, on line 106 we fetch currently allowed CPUs from psutil.Process().cpu_affinity(). It returns every logical CPU that this process is allowed to use.

On line 107 we compare this number to all logical CPUs in the system. (Comparing to physical CPU cores would not make any sense.) If there are less CPUs allowed to this process than there are logical CPUs in the system, we conclude that there is an external CPU manager that takes care of CPU affinity, and use exactly those CPUs that it has chosen. On the other hand, if there is no CPU manager, we assume that we should choose CPUs by ourselves, and continue with the logic that was implemented in your patch, @sywangyi.

The new logic is here: if somebody has already chosen CPUs for this process, then this process uses all of them. We trust that the external entity knows what is the best way to run this process on each hardware. If a user wants to use every hyperthread of certain CPU cores (for instance by using taskset -c ...) then we adapt to that. And if the user wants to run this process only on one hyperthread of every CPU core, then the user can choose logical CPUs correspondingly.

The reason for choosing this logic is: if we do not use all allowed CPUs, or if we continue to set memory affinity only on the node with our allowed CPUs, this process prevents running itself optimally on various platforms. There are cases were we it's best to use hyperthreads, and there are cases where it's best to use memory from a CPU-less NUMA node. This process must not pretend it knows better how to run on every platform unless it really does.

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This only works with --cpuset and doesn't seem to work with other forms of CPU control (like kube requests/limits) right ?

For reference cuda/rocm have a unified env variable used across the ecosystem to ask for visibility on the devices. Is there potentially such a variable that we could reuse here ?

In general both functions should probably be merged together. It seems we're only changing how we resolve the node_id here, right ? (So without the subset it could resolve to some range(n)).

Sidenote: Are there any reasons to use anything else as a config than the max number of threads per CPU (eventually reduced to some subset of the actual machine's physical CPUs) ? Are there benefits to using less threads than the max allowed, or more ?

Copy link
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Kube requests/limits

This PR works with Kubernetes:

In these cases, Kubernetes requests/limits affect the number of allowed CPUs. Details on how many CPUs exactly and which one depend on the CPU manager or NRI policy configuration, hardware topology and other workloads running on the same server. All these policies eventually make affect through cpuset.cpus controller in cgroups.

For reference cuda/rocm have a unified env variable used across the ecosystem to ask for visibility on the devices. Is there potentially such a variable that we could reuse here ?

I’d separate two levels here:

  1. Get the allowed CPUs from the platform. On Linux, the most reliable method for that is sched_getaffinity(), wins reading /sys/fs/cgroup or even /proc/pid/status. This is used behind the scenes by psutil. This works similarly in all cases: taskset, numactl, docker --cpuset-cpus, and the above-mentioned Kubernetes resource managers. There is no need for environment variables.
  2. Introduce workload-specific variables that tell how allowed resources must be used. These are non-standard: there is no common practice for communicating which CPU and memory must/should/could be used. In vLLM, for instance, they use LLM_CPU_OMP_THREADS_BIND environment variable, with which user can specify, for instance, “run shard-0 on CPUs 2-7 and shard-1 on CPUs 8-13” with syntax 2-7|8-13.

I agree that going to Level 2 makes sense at some point. In fact, I’m interested in prototyping text-generation-server performance so that we would tell the kernel to spread memory pages for the model across several memory controllers. This may improve performance due to more available memory bandwidth. For that, I’ll anyway have to introduce an environment variable in any case, because there is no common way (syscall and cgroups) to enforce memory policy like there is for CPU affinity.

If you like it, @Narsil, I could see if I could fit support for LLM_CPU_OMP_THREADS_BIND nicely, too. What do you think?

To keep changes comprehensible, focused and well justified, I’d prefer having this PR as a minimal bug fix to prevent resource underutilization and allow assigning multiple inferences on the same server with good performance. I can write another PR with environment variables.

In general both functions should probably be merged together. It seems we're only changing how we resolve the node_id here, right ? (So without the subset it could resolve to some range(n)).

There is a fundamental difference between the two functions, and in the two cases in general:

  1. [New] Someone who knows the platform and other workloads better than us has chosen a set of CPUs and memories for us. Use everything without questioning the choice. For instance, use both hyperthreads from every physical CPU core, or use 24 CPU cores from one socket and 12 from another socket, if that is want was given.
  2. [Old] Think that we have the whole system for running this single inference. Nobody knows better how to use resources here. Feel free to do our own tricks, like using only CPUs from one single socket.

I'm afraid that handling these quite different cases in the same function would be a bit messy. At least it got messy when I first tried implementing this fix so. There were unnecessarily many control paths through without that many common elements after all.

Furthermore, think of future fixes. For instance, supporting running on Xeons with SNC enabled. (SNC exposes every socket as multiple NUMA nodes. However, memory access across sub-NUMA nodes is way faster if they are on the same socket vs different sockets.) If we want to support optimal automatic performance in this case, the [Old] function needs to be modified so that it detects CPUs and memories are attached to the same socket, yet the OS tells they are in separate NUMA nodes. Currently [Old] function does not do it, so it uses only a fraction of socket’s CPUs and memory bandwidth when SNC is enabled. My point is: it'll be cleaner to implement this kind of fixes to Case 2 [Old] function, and keep Case 1 [New] untouched. The Case 2 function will absorb most if not all internal understanding of various platforms and how to run fast on them.

Sidenote: Are there any reasons to use anything else as a config than the max number of threads per CPU (eventually reduced to some subset of the actual machine's physical CPUs) ? Are there benefits to using less threads than the max allowed, or more ?

The easiest case first: more threads than (logical) CPUs? I cannot imagine where this would make sense. I’ve seen horrible performance (futex storms with little progress) when there have been more than one thread per logical CPU. That’s one reason for this PR to exist.

On the other hand, using less threads than there are allowed CPUs, is quite possible. Inferences are not the only things that are running on servers with hundreds of CPUs and terabytes of memory. CPUs share caches. For instance, consider a case where every block of 4 CPUs would have a common L2 cache. It might be optimal to use only one CPU per block to run something memory/cache intensive, like text-generation-server, and put only compute-intensive tasks on the remaining CPUs of each block. In this split these tasks would not interfere much each other and every CPU running inference would have maximal L2 cache, nevertheless CPUs would not be wasted. And there are even CPU designs where a block of 8 CPUs share single FPU.

Therefore, let’s try to end up with a design where users have enough power to specify the alignment they wish. Your point on environment variables (LLM_CPU_OMP_THREADS_BIND, maybe something on memory policy, too) is definitely valid for that.

_init_cpu_threads_env_use_allowed(rank_id, world_size, allowed_cpus)
else:
_init_cpu_threads_env_use_all(rank_id, world_size)

def _init_cpu_threads_env_use_allowed(rank_id: int, world_size: int, allowed_cpus: list):
import importlib.util

if os.getenv("OMP_NUM_THREADS") is None:
num_cpus_per_rank = max(int(len(allowed_cpus) / world_size), 1)
else:
num_cpus_per_rank = min(int(os.getenv("OMP_NUM_THREADS")), len(allowed_cpus))

if importlib.util.find_spec("numa") is not None:
import numa

slice_info = f"slice {rank_id+1}/{world_size} of externally allowed {len(allowed_cpus)} CPUs"
allowed_mems = numa.memory.get_membind_nodes()
cpu_start = num_cpus_per_rank * rank_id
allowed_cpus_for_rank = allowed_cpus[cpu_start : cpu_start + num_cpus_per_rank]
numa.schedule.run_on_cpus(0, *allowed_cpus_for_rank)
effective_allowed_cpus = numa.schedule.get_affinitive_cpus(0)
else:
slice_info = "externally allowed, cannot import numa for slicing"
allowed_mems = "n/a"
effective_allowed_cpus = allowed_cpus
num_threads = num_cpus_per_rank
torch.set_num_threads(num_threads)
logger.info(f"affinity={effective_allowed_cpus} ({slice_info}), membind={allowed_mems}, threads={num_threads}")

def _init_cpu_threads_env_use_all(rank_id: int, world_size: int):
import importlib.util

if importlib.util.find_spec("numa") is not None:
Expand Down