Composite RoPE backward gives a large ToCopyBackward0 in profiling trace

username: cokespace2
cooked: <p>The most common used Rotary Positional Embedding is like:</p>
<pre data-code-wrap="python"><code class="lang-python">def rotate_half(t: torch.Tensor) -&gt; torch.Tensor:
    t_1, t_2 = torch.chunk(t, 2, dim=-1)
    return torch.cat((-t_2, t_1), dim=-1)

def apply_rotary_pos_emb_bshd(t: torch.Tensor, freqs: torch.Tensor):
    rot_dim = freqs.shape[-1]

    # ideally t_pass is empty so rotary pos embedding is applied to all tensor t
    t, t_pass = t[..., :rot_dim], t[..., rot_dim:]

    # first part is cosine component
    # second part is sine component, need to change signs with _rotate_half method
    cos_ = torch.cos(freqs).to(t.dtype)
    sin_ = torch.sin(freqs).to(t.dtype)

    t = (t * cos_) + (rotate_half(t) * sin_)

    return torch.cat((t, t_pass), dim=-1)
</code></pre>
<p>and when I test this composite-function with torch.profiler, the output json shocked me a lot:<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/1/5/1511f2b198a5665a98e2f8e76d8ff350f5795c88.png" data-download-href="https://discuss.pytorch.org/uploads/default/1511f2b198a5665a98e2f8e76d8ff350f5795c88" title="image"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/1/5/1511f2b198a5665a98e2f8e76d8ff350f5795c88_2_690x185.png" alt="image" data-base62-sha1="30otMquuKQRnl3kEeZ6eBvyRsGk" width="690" height="185" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/1/5/1511f2b198a5665a98e2f8e76d8ff350f5795c88_2_690x185.png, https://discuss.pytorch.org/uploads/default/optimized/3X/1/5/1511f2b198a5665a98e2f8e76d8ff350f5795c88_2_1035x277.png 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/1/5/1511f2b198a5665a98e2f8e76d8ff350f5795c88_2_1380x370.png 2x" data-dominant-color="D6DFDF"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">image</span><span class="informations">1672×449 58.2 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<p>As one can see, there is a large <strong>ToCopyBackward</strong> that takes most of the profling durations, and here’s the tested code:</p>
<pre data-code-wrap="python"><code class="lang-python">#!/usr/bin/env python
# encoding: utf-8
import torch
from torch.nn import functional as F
import time
import os
from functools import partial
import torch
import torch.distributed as dist
from torch.profiler import (
    profile,
    ProfilerActivity,
    schedule,
)


def trace_handler(profiler, file_path, op_name):
    file_path = os.path.join(file_path, f"profiling-{op_name}.trace.json")
    profiler.export_chrome_trace(file_path)


def get_profiler(file_path, op_name):
    warmup = 5
    profile_schedule = schedule(wait=2, warmup=warmup, active=1)
    profiler = profile(
        activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA],
        schedule=profile_schedule,
        record_shapes=True,
        on_trace_ready=partial(trace_handler, file_path=file_path, op_name=op_name),
        with_flops=True,
        profile_memory=True,
        with_modules=True,
    )
    return profiler

def rotate_half(t: torch.Tensor) -&gt; torch.Tensor:
    t_1, t_2 = torch.chunk(t, 2, dim=-1)
    return torch.cat((-t_2, t_1), dim=-1)


def apply_rotary_pos_emb_bshd(t: torch.Tensor, freqs: torch.Tensor):
    rot_dim = freqs.shape[-1]
    # ideally t_pass is empty so rotary pos embedding is applied to all tensor t
    t, t_pass = t[..., :rot_dim], t[..., rot_dim:]

    # first part is cosine component
    # second part is sine component, need to change signs with _rotate_half method
    cos_ = torch.cos(freqs).to(t.dtype)
    sin_ = torch.sin(freqs).to(t.dtype)

    t = (t * cos_) + (rotate_half(t) * sin_)

    return torch.cat((t, t_pass), dim=-1)

def test_ops(op_func, op_name, in_params: dict):
    # for warm up
    out = op_func(**in_params)
    loss = out.sum()
    loss.backward()

    profiler = get_profiler("/workspace", op_name)
    test_iters = 10
    torch.cuda.synchronize()
    start = time.time()
    with profiler as prof:
        for _ in range(test_iters):
            out = op_func(**in_params)
            loss = out.sum()
            loss.backward()
            prof.step()

    torch.cuda.synchronize()
    using_time = time.time() - start
    print(f'{op_name} \t cost: {using_time}')


def test_rope():
    max_seq_len = 4096
    batch_size = 10
    head_num = 32
    dim = 128 * 32
    dim = dim // head_num
    input_shape = (max_seq_len, batch_size, head_num, dim)

    input_ts = torch.randn(input_shape, dtype=torch.float32, requires_grad=True)

    freqs_cis_i = torch.randn(max_seq_len, dim)
    freqs_cis_4d = freqs_cis.reshape(max_seq_len, 1, 1, dim)

    input_data_out_F = {
        "t": input_ts.cuda(),
        "freqs": freqs_cis_4d.cuda()
    }
    test_ops(op_func=apply_rotary_pos_emb_bshd,
        op_name="rope",
        in_params=input_data_out_F,
    )


if __name__ == '__main__':
    test_rope()
</code></pre>
<p>and I run it with <code>CUDA_LAUNCH_BLOCKING=1 python test.py</code>, the cuda version is 12.2</p>
<p>Does anyone have a clue why is that happen? I tested with Conv2d and other modules or functionalities that none of them gives such output.</p>

username: soulitzer
cooked: <p>The ToCopyBackward itself is not taking all the time. Its triggering a device to host synchronization, and its waiting for all the queued up operations on device to finish.</p>

username: soulitzer
cooked: <pre><code class="lang-auto"> input_data_out_F = {
        "t": input_ts.cuda(),
        "freqs": freqs_cis_4d.cuda()
    }
</code></pre>
<p>If you want to avoid this synchronization, you can initialize your input_ts on cuda to begin with, and then you should not need these device movements.</p>

username: cokespace2
cooked: <p>Thanks! That solves the problem.</p>