NUMA Binding Utilities

Created On: Jul 25, 2025 | Last Updated On: Aug 12, 2025

class torch.numa.binding.AffinityMode(value)

See behavior description for each affinity mode in torch.distributed.run.

class torch.numa.binding.NumaOptions(affinity_mode: torch.numa.binding.AffinityMode, should_fall_back_if_binding_fails: bool = False)

affinity_mode: AffinityMode

If true, we will fall back to using the original command/entrypoint if we fail to compute NUMA bindings.

You should avoid using this option! It is only intended as a safety mechanism for facilitating mass rollouts of numa binding.

torch.numa.binding.maybe_wrap_command_args_with_numa_binding(command_args, *, gpu_index, numa_options)

Wraps command arguments with numactl to apply NUMA CPU binding.

This function prepends numactl with appropriate CPU affinity flags to the provided command arguments, binding the process to CPUs associated with the specified GPU’s NUMA node.

Parameters:

• command_args (tuple[str, ...]) – The original command arguments to wrap.

• gpu_index (int) – The index of the GPU that will be used by the subprocess.

• numa_options (NumaOptions | None) – Configuration for NUMA binding behavior. If None, returns the original command_args unchanged.

Returns:

Tuple of command arguments, potentially wrapped with numactl for NUMA binding. Returns the original command_args if numa_options is None or if binding fails and fallback is enabled.

Return type:

tuple[str, …]

torch.numa.binding.maybe_wrap_with_numa_binding(func, *, gpu_index, numa_options)

Wraps a function to apply NUMA CPU binding before execution.

This decorator applies NUMA CPU affinity to all threads in the current process before calling the wrapped function, binding them to CPUs associated with the specified GPU’s NUMA node.

Parameters:

• func (Callable[[~_TParams], _TReturn]) – The function to wrap with NUMA binding.

• gpu_index (int) – The index of the GPU that will be used.

• numa_options (NumaOptions | None) – Configuration for NUMA binding behavior. If None, returns the original function unchanged.

Returns:

A wrapped function that applies NUMA binding before execution, or the original function if numa_options is None.

Return type:

Callable[[~_TParams], _TReturn]