Slow pack_padded_sequence/pad_packed_sequence in backward pass

dda: <p>Hi,</p>
<p>in a recent project I’ve noticed a performance impact in backward pass when packing/unpacking data for LSTMs. To demonstrate, I’ve created a simple LSTM-based network for binary sequence classification:</p>
<pre><code class="lang-auto">Net(
  (embedding): Embedding(10, 16, padding_idx=0)
  (lstm): LSTM(16, 32, batch_first=True)
  (linear): Linear(in_features=32, out_features=2, bias=True)
)
</code></pre>
<p><strong>Setup A: No padding</strong></p>
<ul>
<li>input_ids (shape: batch_size x seq_len):</li>
</ul>
<pre><code class="lang-auto">tensor([[1, 1, 1, 1, 1, 1, 1, 1, 1, 1],
        [1, 1, 1, 1, 1, 1, 1, 1, 1, 1],
        [1, 1, 1, 1, 1, 1, 1, 1, 1, 1],
        [1, 1, 1, 1, 1, 1, 1, 1, 1, 1],
        [1, 1, 1, 1, 1, 1, 1, 1, 1, 1],
        [1, 1, 1, 1, 1, 1, 1, 1, 1, 1],
        [1, 1, 1, 1, 1, 1, 1, 1, 1, 1],
        [1, 1, 1, 1, 1, 1, 1, 1, 1, 1]])
</code></pre>
<ul>
<li>
<p>input_lens (shape: batch_size):<br>
<code>tensor([10, 10, 10, 10, 10, 10, 10, 10])</code></p>
</li>
<li>
<p>target_ids (dummy values, all sequences belong to class ‘0’, shape: batch_size):<br>
<code>tensor([0, 0, 0, 0, 0, 0, 0, 0])</code></p>
</li>
</ul>
<p>Using pack/unpack only adds minimal overhead, both during forward pass and forward &amp; backward pass. This is good! (see figure below)</p>
<p><strong>Setup B: Padded data</strong></p>
<ul>
<li>input_ids (pad_id is 0, shape: batch_size x seq_len):</li>
</ul>
<pre><code class="lang-auto">tensor([[1, 1, 1, 1, 1, 1, 1, 1, 1, 1],
        [1, 1, 1, 1, 1, 1, 1, 1, 1, 0],
        [1, 1, 1, 1, 1, 1, 1, 1, 0, 0],
        [1, 1, 1, 1, 1, 1, 1, 0, 0, 0],
        [1, 1, 1, 1, 1, 1, 0, 0, 0, 0],
        [1, 1, 1, 1, 1, 0, 0, 0, 0, 0],
        [1, 1, 1, 1, 0, 0, 0, 0, 0, 0],
        [1, 1, 1, 0, 0, 0, 0, 0, 0, 0]])
</code></pre>
<ul>
<li>
<p>input_lens (shape: batch_size):<br>
<code>tensor([10,  9,  8,  7,  6,  5,  4,  3])</code></p>
</li>
<li>
<p>target_ids (dummy values, all sequences belong to class ‘0’, shape: batch_size):<br>
<code>tensor([0, 0, 0, 0, 0, 0, 0, 0])</code></p>
</li>
</ul>
<p>While the forward pass only shows minor overhead using pack/unpack (as expected), there is a performance penalty in backward pass. Is this expected? (see figure below)</p>
<p><strong>Environment</strong></p>
<ul>
<li>Torch: 1.12.1+cu116</li>
<li>OS: Ubuntu 22.04, kernel 5.15</li>
</ul>
<p><strong>Code</strong></p>
<pre><code class="lang-auto">"""
Benchmark torch pack/unpack
"""
import contextlib
import logging
from typing import Optional

import torch
from torch import nn
from torch.nn.utils.rnn import pack_padded_sequence, pad_packed_sequence
from tqdm import tqdm

import matplotlib.pyplot as plt


logging.basicConfig(
    format="%(asctime)s %(message)s",
    level=logging.INFO,
    datefmt="%H:%M:%S",
)

logger = logging.getLogger(__name__)


def create_batch(
    max_seq_len: int,
    batch_size: int,
    padding_mode: Optional[str],
    pad_id: int,
):
    input_ids_list = []
    input_lens_list = []
    length = max_seq_len
    for _ in range(batch_size):
        if padding_mode is None:
            input_ids_list.append([1] * max_seq_len)
            input_lens_list.append(max_seq_len)
        elif padding_mode == "triangle":
            input_ids_list.append(
                [1] * length + [pad_id] * (max_seq_len - length)
            )
            input_lens_list.append(length)

            length -= 1
            if length == 0:
                length = max_seq_len
        else:
            raise NotImplementedError(padding_mode)

    # Dummy target_id values
    target_id_list = [0] * batch_size

    # Determine sort indices for pack/unpack
    sort_ids = sorted(
        range(len(input_lens_list)),
        key=input_lens_list.__getitem__,
        reverse=True,
    )

    # Apply sorting
    input_ids_list = [input_ids_list[sort_id] for sort_id in sort_ids]
    input_lens_list = [input_lens_list[sort_id] for sort_id in sort_ids]
    target_id_list = [target_id_list[sort_id] for sort_id in sort_ids]

    # Convert to tensors
    input_ids = torch.tensor(input_ids_list)
    input_lens = torch.tensor(input_lens_list)
    target_ids = torch.tensor(target_id_list)

    return input_ids, input_lens, target_ids


class Net(nn.Module):
    def __init__(
        self,
        vocab_size: int,
        embedding_dim: int,
        hidden_dim: int,
        pad_id: int,
        out_features: int,
    ):
        super().__init__()

        self.pad_id = pad_id

        self.embedding = nn.Embedding(
            num_embeddings=vocab_size,
            embedding_dim=embedding_dim,
            padding_idx=pad_id,
        )

        self.lstm = nn.LSTM(
            input_size=embedding_dim,
            hidden_size=hidden_dim,
            num_layers=1,
            bias=True,
            batch_first=True,
            bidirectional=False,
        )

        self.linear = nn.Linear(
            in_features=hidden_dim,
            out_features=out_features,
            bias=True,
        )

    def forward(
        self,
        input_ids: torch.Tensor,
        input_lens: torch.Tensor,
        use_pack_unpack: bool,
    ):
        x = self.embedding(input_ids)

        if use_pack_unpack:
            x = pack_padded_sequence(
                input=x,
                lengths=input_lens,
                batch_first=True,
                enforce_sorted=True,
            )

        x, _ = self.lstm(x)

        if use_pack_unpack:
            x, _ = pad_packed_sequence(
                sequence=x,
                batch_first=True,
                padding_value=self.pad_id,
            )
        else:
            # Multiply with binary tensor to account for padding
            mask = (
                input_ids.not_equal(self.pad_id)
                .unsqueeze(2)
                .repeat((1, 1, x.size(2)))
            )
            x = torch.mul(x, mask)

        x = self.linear(x)

        logits = torch.sum(x, dim=1)

        return logits


def run_experiment(
    net: Net,
    min_seq_len: int,
    max_seq_len: int,
    seq_step: int,
    num_runs: int,
    batch_size: int,
    padding_mode: Optional[str],
    pad_id: int,
    use_pack_unpack: bool,
    do_backward_pass: bool,
    device: str,
):
    start = torch.cuda.Event(enable_timing=True)
    end = torch.cuda.Event(enable_timing=True)

    optimizer = torch.optim.AdamW(params=net.parameters(), lr=1e-5)
    criterion = nn.CrossEntropyLoss(ignore_index=-100)

    seq_lens = []
    elapsed_times = []
    for seq_len in tqdm(range(min_seq_len, max_seq_len + seq_step, seq_step)):
        input_ids, input_lens, target_ids = create_batch(
            max_seq_len=seq_len,
            batch_size=batch_size,
            padding_mode=padding_mode,
            pad_id=pad_id,
        )
        input_ids = input_ids.to(device)
        target_ids = target_ids.to(device)

        # Warmup to prevent powersaving interfere with benchmark
        with torch.no_grad():
            for _ in range(3):
                logits = net(input_ids, input_lens, use_pack_unpack)

        # Start benchmark
        if do_backward_pass:
            context = contextlib.suppress()
        else:
            context = torch.no_grad()

        start.record()
        with context:
            for _ in range(num_runs):
                optimizer.zero_grad()
                logits = net(input_ids, input_lens, use_pack_unpack)
                loss = criterion(logits.view(-1, 2), target_ids)
                if do_backward_pass:
                    loss.backward()

        end.record()
        torch.cuda.synchronize()
        avg_elapsed_time = start.elapsed_time(end) / num_runs

        # Collect results
        seq_lens.append(seq_len)
        elapsed_times.append(avg_elapsed_time)

    return seq_lens, elapsed_times


if __name__ == "__main__":
    # Model params
    vocab_size = 10
    embedding_dim = 16
    hidden_dim = 32
    out_features = 2
    pad_id = 0

    # Batch params
    batch_size = 512
    min_seq_len = 32
    max_seq_len = 512
    seq_step = 32

    # padding_mode = None
    padding_mode = "triangle"

    # Benchmark params
    num_runs = 20
    device = "cuda"

    # Create model
    net = Net(
        vocab_size=vocab_size,
        embedding_dim=embedding_dim,
        hidden_dim=hidden_dim,
        pad_id=pad_id,
        out_features=out_features,
    )
    net.to(device)

    # Experiments
    logger.info("Experiment 1/4: Forward pass, no pack/unpack")
    seq_lens_fw, elapsed_times_fw = run_experiment(
        net=net,
        min_seq_len=min_seq_len,
        max_seq_len=max_seq_len,
        seq_step=seq_step,
        num_runs=num_runs,
        batch_size=batch_size,
        padding_mode=padding_mode,
        pad_id=pad_id,
        use_pack_unpack=False,
        do_backward_pass=False,
        device=device,
    )

    logger.info("Experiment 2/4: Forward pass, with pack/unpack")
    (
        seq_lens_fw_use_pack_unpack,
        elapsed_times_fw_use_pack_unpack,
    ) = run_experiment(
        net=net,
        min_seq_len=min_seq_len,
        max_seq_len=max_seq_len,
        seq_step=seq_step,
        num_runs=num_runs,
        batch_size=batch_size,
        padding_mode=padding_mode,
        pad_id=pad_id,
        use_pack_unpack=True,
        do_backward_pass=False,
        device=device,
    )

    logger.info("Experiment 3/4: Forward pass &amp; backward pass, no pack/unpack")
    seq_lens_fw_bw, elapsed_times_fw_bw = run_experiment(
        net=net,
        min_seq_len=min_seq_len,
        max_seq_len=max_seq_len,
        seq_step=seq_step,
        num_runs=num_runs,
        batch_size=batch_size,
        padding_mode=padding_mode,
        pad_id=pad_id,
        use_pack_unpack=False,
        do_backward_pass=True,
        device=device,
    )

    logger.info(
        "Experiment 4/4: Forward pass &amp; backward pass, with pack/unpack"
    )
    (
        seq_lens_fw_bw_use_pack_unpack,
        elapsed_times_fw_bw_use_pack_unpack,
    ) = run_experiment(
        net=net,
        min_seq_len=min_seq_len,
        max_seq_len=max_seq_len,
        seq_step=seq_step,
        num_runs=num_runs,
        batch_size=batch_size,
        padding_mode=padding_mode,
        pad_id=pad_id,
        use_pack_unpack=True,
        do_backward_pass=True,
        device=device,
    )

    # Plot
    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(16, 6))
    fig.suptitle(
        f"Benchmark (batch_size: {batch_size}, " f"padding: {padding_mode})"
    )
    # Forward pass
    ax1.plot(
        seq_lens_fw_use_pack_unpack,
        elapsed_times_fw_use_pack_unpack,
        ".-r",
        label="with pack/unpack",
    )
    ax1.plot(
        seq_lens_fw,
        elapsed_times_fw,
        ".-g",
        label="no pack/unpack",
    )
    ax1.set_xticks(list(range(min_seq_len, max_seq_len + seq_step, seq_step)))
    ax1.set_xlabel("Sequence length")
    ax1.set_ylabel("Average runtime [ms]")
    ax1.set_title("Forward pass")
    ax1.grid(linestyle="dotted", color="gray")
    ax1.set_ylim(bottom=0)
    ax1.legend()

    # Forward pass &amp; backward pass
    ax2.plot(
        seq_lens_fw_bw_use_pack_unpack,
        elapsed_times_fw_bw_use_pack_unpack,
        ".-r",
        label="with pack/unpack",
    )
    ax2.plot(
        seq_lens_fw_bw,
        elapsed_times_fw_bw,
        ".-g",
        label="no pack/unpack",
    )
    ax2.set_xticks(list(range(min_seq_len, max_seq_len + seq_step, seq_step)))
    ax2.set_xlabel("Sequence length")
    ax2.set_ylabel("Average runtime [ms]")
    ax2.set_title("Forward pass &amp; backward pass")
    ax2.grid(linestyle="dotted", color="gray")
    ax2.set_ylim(bottom=0)
    ax2.legend()
</code></pre>

dda: <p><strong>Setup A: No padding</strong><br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/c/d/cde71c006212896b0ae77d2ae0aa19cefea44d5c.png" data-download-href="https://discuss.pytorch.org/uploads/default/cde71c006212896b0ae77d2ae0aa19cefea44d5c" title="2022-10-23_benchmark_padding_none"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/c/d/cde71c006212896b0ae77d2ae0aa19cefea44d5c_2_690x303.png" alt="2022-10-23_benchmark_padding_none" data-base62-sha1="tnv1VbOygVfqMuWgGe0h0GbEIck" width="690" height="303" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/c/d/cde71c006212896b0ae77d2ae0aa19cefea44d5c_2_690x303.png, https://discuss.pytorch.org/uploads/default/original/3X/c/d/cde71c006212896b0ae77d2ae0aa19cefea44d5c.png 1.5x, https://discuss.pytorch.org/uploads/default/original/3X/c/d/cde71c006212896b0ae77d2ae0aa19cefea44d5c.png 2x" data-dominant-color="F5F5F4"><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">2022-10-23_benchmark_padding_none</span><span class="informations">940×414 40.5 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>

dda: <p><strong>Setup B: Padded data</strong><br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/8/1/810486ed2b2fd604e03a89fdffa74658d7b6f819.png" data-download-href="https://discuss.pytorch.org/uploads/default/810486ed2b2fd604e03a89fdffa74658d7b6f819" title="2022-10-23_benchmark_padding_triangle"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/8/1/810486ed2b2fd604e03a89fdffa74658d7b6f819_2_690x303.png" alt="2022-10-23_benchmark_padding_triangle" data-base62-sha1="ipli9L6DmY6sYQe920YVsEiQ649" width="690" height="303" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/8/1/810486ed2b2fd604e03a89fdffa74658d7b6f819_2_690x303.png, https://discuss.pytorch.org/uploads/default/original/3X/8/1/810486ed2b2fd604e03a89fdffa74658d7b6f819.png 1.5x, https://discuss.pytorch.org/uploads/default/original/3X/8/1/810486ed2b2fd604e03a89fdffa74658d7b6f819.png 2x" data-dominant-color="F6F5F4"><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">2022-10-23_benchmark_padding_triangle</span><span class="informations">940×414 41.6 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>