ray.experimental.deregister_nixl_memory

ray.experimental.deregister_nixl_memory(tensor: torch.Tensor) → None

Decrements the reference count for the tensor’s NIXL memory registration added by ray.experimental.register_nixl_memory().

If the reference count reaches 0, the memory is deregistered from NIXL. This should only be called after ray.experimental.register_nixl_memory() has been called for this tensor. Any existing ray.ObjectRef instances that reference this tensor’s memory will keep the NIXL memory registration alive independently until they go out of scope.

Parameters:

tensor – A PyTorch tensor whose NIXL memory registration reference count should be decremented.

Example

# Extending the example from register_nixl_memory:
@ray.remote(num_gpus=1, enable_tensor_transport=True)
class Trainer:
    def deregister_weight(self):
        # Remove the NIXL memory registration added by register_nixl_memory.
        # The memory may still be registered if there are live ObjectRefs
        # that reference it.
        deregister_nixl_memory(self.weight)

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