What is the `torch.cond` way to write conditionals in PyTorch?

1

I have learnt, that if you want a graph safe way of using conditionals inside of a PyTorch model you should use torch.cond. It will reduce the painful memcpy (:scream:) events between the host and device (CUDA) which while algorithmically not a problem, performance wise makes inference very slow.

The way I have used torch.cond should be correct (right?) yet also it is a bit “ugly”. I was wondering if there are any best practices that come with using torch.cond to make the model description concise yet performative? Searching on the pytorch/examples repo, I could not find any examples.

Or is this as good as it gets? This would be a good conclusion also :grin:

Minimal example :backhand_index_pointing_down:

import torch
import torch.nn as nn


class EarlyExitMLP(nn.Module):
    def __init__(self, in_features=32, hidden=64, num_classes=10):
        super().__init__()

        # backbone with three FC layers
        self.backbone = nn.ModuleList([
            nn.Linear(in_features, hidden),
            nn.Linear(hidden, hidden),
            nn.Linear(hidden, hidden),
        ])

        # 3 early exits
        self.exit1 = nn.Linear(hidden, num_classes)
        self.exit2 = nn.Linear(hidden, num_classes)
        self.exit3 = nn.Linear(hidden, num_classes)

        # same as self.exit3
        self.final = nn.Linear(hidden, num_classes)
    
    # threshold met exit check
    def should_exit(self, logits, threshold):
        probs = torch.softmax(logits, dim=-1)
        return probs.max(dim=-1).values > threshold
    
    def result(self, logits, exit_id):
        return {"exit_id": exit_id, "logits": logits}

    def forward_early_exit(self, x, t1, t2, t3):
        # we consider batch size of one, so no variable input tensor shape complex stuff
        if x.shape[0] != 1:
            raise ValueError("Early-exit inference requires a batch size of 1.")

        x = x.flatten(1)

        h1 = torch.relu(self.backbone[0](x))
        logits1 = self.exit1(h1)
        
        # this part "works", but how to rewrite it?
        def after_exit1(h1, logits1):
            h2 = torch.relu(self.backbone[1](h1))
            logits2 = self.exit2(h2)

            def after_exit2(h2, logits2):
                h3 = torch.relu(self.backbone[2](h2))
                logits3 = self.exit3(h3)

                def after_exit3(h3, logits3):
                    final_logits = self.final(h3)
                    return self.result(final_logits, 4)

                return torch.cond(
                    self.should_exit(logits3, t3),
                    lambda h3, logits3: self.result(logits3, 3),
                    after_exit3,
                    (h3, logits3),
                )

            return torch.cond(
                self.should_exit(logits2, t2),
                lambda h2, logits2: self.result(logits2, 2),
                after_exit2,
                (h2, logits2),
            )

        return torch.cond(
            self.should_exit(logits1, t1),
            lambda h1, logits1: self.result(logits1, 1),
            after_exit1,
            (h1, logits1),
        )


# Example
model = EarlyExitMLP()
x = torch.randn(1, 32)

out = model.forward_early_exit(x, t1=0.9, t2=0.8, t3=0.7) # dummy threshold values
print(out)

I know there are some complexities when we have dynamic shapes (e.g. if some samples in a batch exit while others remain) or having multiple tensor output but here I am mainly concerned with a batch size of 1. I also understand that torch.cond is a prototype feature, yet I simply cannot find a better alternative for conditional graphs.

2

If by graph you mean cuda graph, you cannot use data-conditioned flows.
You can simply use torch.compile and pay the price of breaking the graph. Your op will be supported and won’tn expecience a crazy delay as long as the data shape of subsequent tensors is fixed.

3

If by graph you mean cuda graph, you cannot use data-conditioned flows.

Apart from writing the CUDA conditionals ourselves (e.g. IConditionLayer when using torch.compile with TensorRT) and somehow replacing the torch.cond operator is there another way to preserve the data-conditioned flows without breaking the graph?

Any reading material / references would be greatly appreciated. :nerd_face:

4

I honestly do not know the intrinsecs of TensorRT but, as per the docs you suggest, it’s not a single graph but conditional graph execution.

When you are running pytorch code in eager mode, kernels are launched by the cpu. This requires gpu-cpu sync, which imposes constrains, and launching the kernel, which takes time.

Typically, if your kernel does heavy computation, this “launching” time in neglegible against the compute time. But if your kernel is very light, it’s significant.

On the other side, part the optimizations come from kernel fusion. internally, the gpu has to copy memory from global gpu memory to shared memory (some sort of fast acces memory) and registers.
So torch.compile skips this via kernel fusion, i.e., you save this copy time + kernel launch time.

But in plain summary, this is doable because the execution flow is predefined. Memory addresses and tensor sizes. The GPU simply schedules the same operations internally with triggers upon completion, without needed to call the cpu. However, the output of a kernel cannot define what kernel will be launched next. That’s why there is no such thing as “unified graphs” with conditional execution.

What you can do is call torch.compile on the different blocks (functions that use pytorch code).

def my_func():
   ....
my_compiled_func = torch.compile(my_func,...)

Your conditional block will inevitable break the graph and the cpu will schedule the corresponding “subgraph” or block according to your condition. Which is very similar to what’s shown at TensorRT

image
image876×136 4.95 KB

Alternatively, because your code is nothing but a cascade of ops which you exit at a given point, you can reformulate your code to execute always all ops, just zeroing some stuff at the point at which you would exist, so that the blocks after “existing” behave like an identity.