Memory pollution while concurrent data transfer with multi cuda stream

abdula2523: <p>Hello community.</p>
<p>I’m trying to load and offload the model’s weight from the host(pinned) and GPU and from the GPU and host(pinned), with some multiple Python threads and CUDA streams.</p>
<p>But, I’ve got different inference results when transferring the model’s weights(actually layers you know) when migrate concurrently and not concurrently.</p>
<p>Also, I presume the problem is on difference between <a href="https://github.com/pytorch/pytorch/blob/main/c10/cuda/CUDACachingAllocator.cpp#L1605" rel="noopener nofollow ugc">DeviceCachingAllocator</a> 's free()<br>
and the <a href="https://github.com/pytorch/pytorch/blob/b60273b88a25e282728a036dd08d39035a369d1f/aten/src/ATen/cuda/CachingHostAllocator.cpp#L193" rel="noopener nofollow ugc">CUDAHostAllocator</a>’s free()</p>
<p>CUDAHostAllocator inserts the event to <code>cuda_events_</code> when it’s not ready.<br>
But, DeviceCachingAllocator doesn’t insert the event, <strong>even if</strong> it’s not ready.</p>
<p>Are these behaviors expected?</p>
<p>Here are some codes for simulating the problem. I know it’s…a little bit long…but you can reproduce the problem easily.<br>
and also, I’m using PyTorch 2.0.1</p>
<pre><code class="lang-auto">from queue import Queue
import threading
import torch

N_MEMORY = 10
MEMORY_LIMIT = 3 * 1024**3  # 3GB
LAYER_SIZE = 512 * 1024**2  # 512MB


class Object(object):
    pass


class FakeLoader(threading.Thread):
    def __init__(self, streams):
        threading.Thread.__init__(self)
        self.streams = streams
        self.memory_loaded = 0
        self.loaded_layer = Queue()

    def can_load(self):
        global layers

        if self.memory_loaded + LAYER_SIZE &lt; MEMORY_LIMIT:
            return True
        else:
            return False

    def load(self, id):
        global layers

        print("BEGIN LOAD", id)
        with torch.cuda.stream(self.streams["upstream"]):
            layers[id].tensor = layers[id].tensor.to("cuda:0", non_blocking=True)

        layers[id].end_load_event = torch.cuda.Event(enable_timing=True)
        layers[id].end_load_event.record(self.streams["upstream"])

        self.memory_loaded += LAYER_SIZE
        self.loaded_layer.put(id)

    def unload(self, id):
        global layers

        print("BEGIN UNLOAD", id)
        with torch.cuda.stream(self.streams["downstream"]):
            layers[id].tensor = layers[id].tensor.to("cpu", non_blocking=True)
        # !!! UNCOMMENT if you want precise result
        # self.streams["downstream"].synchronize()

        self.memory_loaded -= LAYER_SIZE

    def run(self):
        global layers

        for id in range(N_MEMORY):
            layers[id].begin_load_event = torch.cuda.Event(enable_timing=True)
            layers[id].begin_load_event.record(self.streams["upstream"])

            while True:
                if self.can_load():
                    self.load(id)

                    break
                else:
                    self.unload(self.loaded_layer.get())
            layers[id].load_lock.set()


if __name__ == "__main__":
    global layers
    layers = {}
    orginal_memorys_cpu = {}
    orginal_memorys_gpu = {}
    for id in range(N_MEMORY):
        layers[id] = Object()
        layers[id].tensor = torch.randint(
            32000, [1, 64 * 1024**2], dtype=torch.long
        )  # 512MB
        layers[id].tensor = layers[id].tensor.pin_memory()

        orginal_memorys_cpu[id] = layers[id].tensor.detach().clone().pin_memory()
        orginal_memorys_gpu[id] = layers[id].tensor.detach().clone().to("cuda:0")

    streams = {}
    streams["upstream"] = torch.cuda.Stream()
    streams["downstream"] = torch.cuda.Stream()

    # init()
    for layer in layers.values():
        layer.load_lock = threading.Event()

    loader = FakeLoader(streams)
    loader.start()
    loader.join()

    # check_result()
    for id in range(N_MEMORY):
        if layers[id].tensor.device == torch.device("cpu"):
            print("Layer", id, "is in cpu")

            if torch.allclose(
                layers[id].tensor, orginal_memorys_cpu[id], equal_nan=True
            ):
                print(f"Layer {id} is same")
            else:
                print(f"Layer {id} is different")
        elif layers[id].tensor.device == torch.device("cuda:0"):
            print("Layer", id, "is in cuda")

            if torch.allclose(
                layers[id].tensor, orginal_memorys_gpu[id], equal_nan=True
            ):
                print(f"Layer {id} is same")
            else:
                print(f"Layer {id} is different")
        else:
            print("!!! SOMETHING WRONG Device is", layers[id].tensor.device)

</code></pre>
<p>Thanks.</p>

ptrblck: <p>Yes, this is expected since you need to synchronize the stream before being able to use the <code>CPUTensor</code> after moving it with <code>non_blocking=True</code>.<br>
This example illustrates it:</p>
<pre><code class="lang-python">import torch
import time
from statistics import mean

def get_cycles_per_ms() -&gt; float:
    """Measure and return approximate number of cycles per millisecond for torch.cuda._sleep
    """

    def measure() -&gt; float:
        start = torch.cuda.Event(enable_timing=True)
        end = torch.cuda.Event(enable_timing=True)
        start.record()
        torch.cuda._sleep(1000000)
        end.record()
        end.synchronize()
        cycles_per_ms = 1000000 / start.elapsed_time(end)
        return cycles_per_ms

    # Get 10 values and remove the 2 max and 2 min and return the avg.
    # This is to avoid system disturbance that skew the results, e.g.
    # the very first cuda call likely does a bunch of init, which takes
    # much longer than subsequent calls.
    #
    # Tested on both Tesla V100, Quadro GP100, Titan RTX, RTX 3090 GPUs
    # and seems to return stable values.
    num = 10
    vals = []
    for _ in range(num):
        vals.append(measure())
    vals = sorted(vals)
    return mean(vals[2 : num - 2])


if __name__ == '__main__':
    seed = 0
    torch.manual_seed(seed)
    torch.cuda.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)
    
    cycles = get_cycles_per_ms()
    
    stream = torch.cuda.current_stream()

    x = torch.rand(32, 256, 220, 220, device="cuda")
    torch.cuda._sleep(int(100 * cycles))
    t = x.to(torch.device("cpu"), non_blocking=True)
    print(stream.query()) # False - work not done yet
    #stream.synchronize() # wait for stream to finish the work
    print(t.abs().sum())
    
    time.sleep(2.)
    print(stream.query()) # True - work done
    print(t.abs().sum())
</code></pre>
<p>You should see two different outputs since the needed <code>stream.synchronize()</code> is commented out.<br>
Comment it in and you should see the same results again.</p>

abdula2523: <p>Thanks for your quick and kind response <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=12" title=":slight_smile:" class="emoji" alt=":slight_smile:" loading="lazy" width="20" height="20"></p>
<p>I’ve got a question Why do they have a different approach between <code>CPUTensor</code> and <code>CUDATensor</code> ?<br>
For example, in code snippets that you provided, change <code>"cuda"</code> to <code>"cpu"</code>(<code>torch.rand()</code>), and <code>"cpu"</code> to <code>"cuda"</code>.(<code>Tensor.to()</code>), then the result is <code>True</code> and <code>True</code>.</p>
<p>This behavior on <code>CPUTensor</code> is for reducing latency or something else?</p>

ptrblck: <p>CUDA operations will be executed in the surrounding <code>CUDAStream</code>, which is the single default stream in my example. If you move a <code>CPUTensor</code> to the GPU asynchronously via <code>non_blocking=True</code>, the operation will be non-blocking w.r.t. the host, however will be performed in the surrounding stream and is thus in order. If you want async execution w.r.t. the GPU you would need to use different streams.</p>

abdula2523: <p>Thanks! . And it will be my last question. Why do they implemented different ways between DeviceHostAllocator and CUDACaching allocator?</p>

ptrblck: <p>The <code>CUDAHostAllocator</code> is responsible to allocate pinned (page-locked) memory on the host via <code>cudaHostAlloc</code> while the <code>CUDACachingAllocator</code> is responsible to allcoate and free memory on the GPU.</p>

abdula2523: <p>My actual question is why in <code>at::native::copy_kernel_cuda()</code> doesn’t call <code>CUDACachingAllocator.recordStream()</code>, but does call <code>CUDAHostAllocator.record_event()</code>.<br>
I think <code>CUDACachingAllocator.recordStream()</code> is called for synchronize, then it should call <code>CUDAHostAllocator.record_event()</code> too.</p>
<p>I think I might misunderstand something…</p>