ray.experimental.register_nixl_memory#

ray.experimental.register_nixl_memory(tensor: torch.Tensor) → None[source]#

Registers the tensor’s memory with NIXL and bumps the reference count so the memory region is never deregistered.

By default, the lifetime of the NIXL memory registration is tied to the ObjectRef. This means that only when the ObjectRef is created do we register the memory with NIXL and deregister it when the ObjectRef goes out of scope. However, this function can be used to pre-register a tensor’s memory with NIXL and keep it registered for the lifetime of the process which can improve performance if the same tensor is re-used in multiple RDT objects.

If called on a tensor that is already registered with NIXL, the reference count is still bumped to prevent the memory from being deregistered.

Parameters:: tensor – A PyTorch tensor whose memory should be registered with NIXL.

Example

import torch
import ray
from ray.experimental import register_nixl_memory

@ray.remote(num_gpus=1, enable_tensor_transport=True)
class Trainer:
    def __init__(self):
        self.weight = torch.randn(1000, 1000, device="cuda")
        # Pre-register the memory with NIXL for the lifetime of the process
        register_nixl_memory(self.weight)

    # Both of the below methods will use the cached NIXL memory registration on multiple calls. You can also mix them,
    # i.e. call get_weight_ref_by_rows then get_weight_ref and get_weight_ref will not trigger a new NIXL memory registration.

    # You can ray.put views to each row of the weight matrix if you want to use them separately in your code
    def get_weight_ref_by_rows(self):
        views = [self.weight[i] for i in range(1000)]
        # Each put call does not trigger a new NIXL memory registration
        return ray.put(views, _tensor_transport="nixl")

    # You can also ray.put the entire weight matrix at once
    def get_weight_ref(self):
        return ray.put(self.weight, _tensor_transport="nixl")

PublicAPI (alpha): This API is in alpha and may change before becoming stable.