Support higher-order while_loop and scan torch ops

[GregoryComer]

🚀 The feature, motivation and pitch

Pytorch provides several higher-level control flow ops, including while_loop and scan. These ops have to be used explicitly in the PyTorch model (Dynamo can't infer them from Python control flow), but they are useful in expressing more complex models. As a first step to try to handle complex patterns without needing to chop up the model and port logic to C++, we should support some of the higher-order ops in ET.

I believe while_loop, and scan should be a good starting point, as they are sufficient to express many of the data dependent patterns. This may allow for writing some generator loops in pure PyTorch, specially for non-multi-turn generation.

Currently, these ops are not supported on ET. It's likely that we will need both AOT and runtime work to do this. Exporting a simple cond model gives the following error. Simple scan and while_loop models also fail with various errors

Edit: From Tarun's comment, cond should be supported, but it seems like there are some issues to resolve, as I haven't been able to use cond when I've tried it. I will create specific tasks for cond. I've removed it from this issue for now.

Alternatives

No response

Additional context

No response

RFC (Optional)

No response

cc @JacobSzwejbka @angelayi

[GregoryComer]

CC @mergennachin I'm curious on your thoughts here. I'm not tagging as UX, as it's not pressing for immediate usabillity. But I'd think we can get to a future where we can express more complex data-dependent patterns in model definitions to avoid needing to chop models up and port outer loops to C++. This might help with multimodal, for example, as I believe it has nested loops. It may also allow for one-shot generation to be done completely in model (no complex runner bits needed).

[angelayi]

I believe @tarun292 previously implemented support for cond and map (code pointer).

[tarun292]

@GregoryComer cond and map are already supported in the stack. We haven't yet added support for while and scan because there haven't been any usecases that have needed this. These higher-order ops do have limitations such as mutations cannot happen in the branches which constrain their usage in GenAI use cases (mainly because intermediate memory usage goes up significantly).

So we do need to add support for while_loop and scan, but maybe before we invest resources in that we should re-write some models in python using these higher order ops and do some verification that it clearly solves problems for us.

[GregoryComer]

@GregoryComer cond and map are already supported in the stack. We haven't yet added support for while and scan because there haven't been any usecases that have needed this. These higher-order ops do have limitations such as mutations cannot happen in the branches which constrain their usage in GenAI use cases (mainly because intermediate memory usage goes up significantly).

So we do need to add support for while_loop and scan, but maybe before we invest resources in that we should re-write some models in python using these higher order ops and do some verification that it clearly solves problems for us.

Thanks. I wasn't able to use cond when I tried it, but if it's expected to work, I'll file for the specific issues I ran into. One the mutability and loops, I think it's worth looking into this further. Can we leverage buffers / IO to make this performant? I'll take a deeper look at a few of the use cases that I think might benefit from this and prototype it.

[vikasbalaga]

Hi all,
I am recently facing some issues with the presence of 'for loops' in the forward() implementation, due to which the exported model (*.pte) size is varying hugely with dummy input size. (I have raised a query here)

In this particular discussion, I came to know about while_loop and scan torch OPs (Not sure how they work)
Is there a way I can leverage them to fix my issue as I suspect my issues are due to loop unrolling being done during export() at which the control seems to be freezing without generating the *.pte model?

[vikasbalaga]

Hi again, I tried implementing my logic using the while_loop operator, but I am facing a run time issue saying :
torch._export.verifier.SpecViolationError: Node.meta item is missing val field

So, I tried to come up with a minimal script which uses the while_loop operator and also results in the above mentioned error which is as below :

import torch
from torch import nn
from torch.export import export

class ExportableLoop(nn.Module):
    def __init__(self, hidden_size, out_channels):
        super().__init__()
        self.hidden_size = hidden_size
        self.B = nn.Parameter(torch.randn(hidden_size, 1))  # (H, in_channels)
        self.C = nn.Parameter(torch.randn(out_channels, hidden_size))  # (C_out, H)
        A = torch.randn(2, hidden_size)
        self.A_real = nn.Parameter(A[0])
        self.A_imag = nn.Parameter(A[1])

    def update_state(self, h, x_t):
        # h: [B, 2, H], x_t: [B, H]
        hr, hi = h[:, 0, :], h[:, 1, :]  # [B, H]
        hrn = hr * self.A_real - hi * self.A_imag + x_t  # [B, H]
        hin = hi * self.A_real + hr * self.A_imag        # [B, H]
        hn = torch.stack([hrn, hin], dim=1)              # [B, 2, H]
        return hn, hrn

    def forward(self, u):
        # u: [B, 1, T]
        x = torch.matmul(self.B, u)  # (B, H, T)
        B, H, T = x.shape

        h = torch.zeros(B, 2, H, device=x.device, dtype=x.dtype)  # [B, 2, H]
        h_accum = torch.zeros(B, H, T, device=x.device, dtype=x.dtype)  # [B, H, T]
        i = torch.tensor(0, device=x.device, dtype=torch.int64)

        def cond(i, h, h_accum):
            return i < T

        def body(i, h, h_accum):
            x_t = torch.select(x, -1, i)  # [B, H]
            h, hr = self.update_state(h, x_t)  # h: [B, 2, H], hr: [B, H]
            h_accum = h_accum.index_copy(-1, i.unsqueeze(0), hr.unsqueeze(-1))  # [B, H, T]
            return i + 1, h, h_accum

        _, h, h_accum = torch._higher_order_ops.while_loop(cond, body, (i, h, h_accum))
        y = torch.matmul(self.C, h_accum).transpose(0, 1)  # (B, C_out, T)
        return y

# Instantiate and export
model = ExportableLoop(hidden_size=128, out_channels=10)
inp = torch.randn(1, 1, 32)  # (B, in_channels=1, T=32)
exported = export(model, (inp,))
print("Export succeeded.")

Can someone please explain this error and how to fix it?

[vikasbalaga]

I tried to debug the issues and I managed to fix them, so the model is now being exported!
But after that, as I try to transform it to an Edge IR, I am facing some "aliasing issues", which I couldn't fix!
Please find my updated code here :

import torch
from torch import nn
from torch.export import export
from executorch.exir import to_edge_transform_and_lower

class ExportableLoop(nn.Module):
    def __init__(self, hidden_size, out_channels):
        super().__init__()
        self.hidden_size = hidden_size
        self.B = nn.Parameter(torch.randn(hidden_size, 1))  # (H, in_channels)
        self.C = nn.Parameter(torch.randn(out_channels, hidden_size))  # (C_out, H)
        A = torch.randn(2, hidden_size)
        self.A_real = nn.Parameter(A[0])
        self.A_imag = nn.Parameter(A[1])

    def update_state(self, h, x_t):
        # h: [B, 2, H], x_t: [B, H]
        hr, hi = h[:, 0, :], h[:, 1, :]  # [B, H]
        hrn = hr * self.A_real - hi * self.A_imag + x_t  # [B, H]
        hin = hi * self.A_real + hr * self.A_imag        # [B, H]
        hn = torch.stack([hrn, hin], dim=1)              # [B, 2, H]
        return hn, hrn

    def forward(self, u):
        # u: [B, 1, T]
        x = torch.matmul(self.B, u)  # (B, H, T)
        B, H, T = x.shape

        h = torch.zeros(B, 2, H, device=x.device, dtype=x.dtype)  # [B, 2, H]
        h_accum = torch.zeros(B, H, T, device=x.device, dtype=x.dtype)  # [B, H, T]
        i = torch.tensor(0, device=x.device, dtype=torch.int64)

        def cond(i, h, h_accum):
            return i < T

        def body(i, h, h_accum):
            x_t = x.index_select(-1, i.unsqueeze(0)).squeeze(-1)  # ✅ safe for export
            h, hr = self.update_state(h, x_t)  # h: [B, 2, H], hr: [B, H]
            h_accum = h_accum.index_copy(-1, i.unsqueeze(0), hr.unsqueeze(-1))  # [B, H, T]
            return i + 1, h, h_accum

        _, h, h_accum = torch._higher_order_ops.while_loop(cond, body, (i, h, h_accum))
        y = torch.matmul(self.C, h_accum).transpose(0, 1)  # (B, C_out, T)
        return y

# Instantiate and export
model = ExportableLoop(hidden_size=128, out_channels=10)
inp = torch.randn(1, 1, 32)  # (B, in_channels=1, T=32)
exported = export(model, (inp,))
print("Exporting Done...")
executorch_program = to_edge_transform_and_lower(
    exported
)
print("Edge transform Done...")

out_path = "loop_model.pte"
with open(out_path, "wb") as file:
    file.write(executorch_program.buffer)

print(f"Succesfully saved model as {out_path}")

Error message is :

torch._higher_order_ops.utils.UnsupportedAliasMutationException: torch.while_loop's cond_fn might be aliasing the input!

While executing %while_loop : [num_users=3] = call_function[target=torch.ops.higher_order.while_loop](args = (%while_loop_cond_graph_0, %while_loop_body_graph_0, (%detach_, %zeros, %zeros_1), (%a_imag, %a_real, %matmul)), kwargs = {})
GraphModule: class GraphModule(torch.nn.Module):

@GregoryComer , @tarun292 , can you please help me in understanding these issues?
(I know this ticket is not appropriate one, I can create a new one if it is required)

Thanks

[GregoryComer]

@vikasbalaga Currently, the higher-order while and scan ops aren't actually supported in ExecuTorch's runtime. @JacobSzwejbka do you have any thoughts on effort needed? Since we have support for cond, seems like we should have most of the building blocks there.

On the specific exportability question, it might be worth asking in PyTorch core, as they maintain the core export logic. I'm not completely sure what's going on, but it might be related to the condition function capturing T? You might be able to add it as an explicit input to the cond function.