[Request Help] “torch._dynamo.exc.UserError"

[liye0626]

Background: I tried to integrate EAGLE-2 into ExecuTorch, but encountered some errors.
Relate code:

            if i not in noleaf_index: # An error occurred at this branch
                cid = i
                depth = position_ids_list[i]
                for j in reversed(range(depth + 1)):
                    retrieve_indices[rid][j] = cid
                    cid = mask_index_list[cid - 1]
                rid += 1

Specifically, I got an error after adding a branch in model forward

• code:

if A_Tensor: # In forward, a Boolean value or bool tensor calculated using real-time data
	...

• error:

torch._dynamo.exc.UserError: Dynamic control flow is not supported at the moment. Please use torch.cond to explicitly capture the control flow. For more information about this error, see: https://pytorch.org/docs/main/generated/exportdb/index.html#cond-operands

I refer to the above page and add the module "CondBranchNestedFunction", encountered a new error

• code:

from functorch.experimental.control_flow import cond
@torch.compile(dynamic=True, fullgraph=True) # dynamic setting True and False will both result in an error
class CondBranchNestedFunction(torch.nn.Module):
    @torch.compile(dynamic=True, fullgraph=True)
    def forward(self, tmpA, i):
        def true_fn(i):
            i+=1
            return None

        def false_fn(i):
            return None

        return cond(tmpA, true_fn, false_fn, [i])


self.condition_func = CondBranchNestedFunction()
self.condition_func = torch.compile(self.condition_func, mode='max-autotune')

tmpA = True # or bool Tensor.  In forward, a Boolean value or bool tensor calculated using real-time data
tmpB = torch.tensor(i)
self.condition_func(tmpA, tmpB)

• error:
  "torch._dynamo.exc.UncapturedHigherOrderOpError: Cond doesn't work unless it is captured completely with torch.compile."

I have tried using torch.compile in some locations. But the error did not change. I would like to ask which direction I should go to solve the problem.

Specification are as follows:

My torch version:

torch                     2.6.0.dev20241224+cpu

torchaudio                2.6.0.dev20241224+cpu

torchvision               0.22.0.dev20241224+cpu

ExecuTorch commit: 86cb5d7

The changes to ExecuTorch's llama_transformer.py(link: llama_transformer.py ) are as follows:

Add the following code to line 25 (link:llama_transformer.py)

from functorch.experimental.control_flow import cond
@torch.compile(dynamic=True, fullgraph=True) # dynamic设为True和False均会出错
class CondBranchNestedFunction(torch.nn.Module):
    @torch.compile(dynamic=True, fullgraph=True)
    def forward(self, tmpA, i):
        def true_fn(i):
            i+=1
            return None

        def false_fn(i):
            return None

        return cond(tmpA, true_fn, false_fn, [i])

Add the following code to line 425 (link:llama_transformer.py)

        self.condition_func = CondBranchNestedFunction()
        self.condition_func = torch.compile(self.condition_func, mode='max-autotune')

Add the following code to line 544(link:llama_transformer.py#L544)

        for i in range(10):
            tmpA = (i == logits).int().sum() == 0
            tmpB = torch.tensor(i)
            self.condition_func(tmpA, tmpB) 

python -m examples.models.llama.export_llama \
    --model "llama3_2" \
    --checkpoint $MODEL_PATH \
    --params $PARAM_PATH \
    -kv \
    -X \
    -d bf16 \
    --output_name=$OUTPUT_NAME \
    --use_sdpa_with_kv_cache \

Then the above error will occur

torch._dynamo.exc.UncapturedHigherOrderOpError: Cond doesn't work unless it is captured completely with torch.compile. Scroll up to find out what causes the graph break.

cc @JacobSzwejbka @angelayi @mergennachin @byjlw

[iseeyuan]

@ydwu4 , could you help on this issue since it's related to torch.cond?

[ydwu4]

Hi @liye0626 , i'm not sure about the purpose of the loop:

        for i in range(10):
            tmpA = (i == logits).int().sum() == 0
            tmpB = torch.tensor(i)
            self.condition_func(tmpA, tmpB) 

Is it that you want to count the number of "(i == logits).int().sum() == 0"? In this case, we could do

class Cond(torch.nn.Module):
    def forward(self, tmpA, i):
        def true_fn(i):
            return torch.ones([], dtype=int64)

        def false_fn(i):
            return torch.zeros([], dtype=int64)

        return cond(tmpA, true_fn, false_fn, [i])

Then sum the outputs over.

Cond is functional meaning that it shouldn't mutate the input and shouldn't have side effects. that's why the += 1 doesn't work.

[liye0626]

@ydwu4

The purpose of this loop is just an example, the actual scenario is similar but more complex.

Thank you very much for your help. In "true_fn" or "false_fn", we cannot modify the input value, but we can use the input value for calculation and return the result in the form of tensor. Is my understanding correct?

[ydwu4]

return the result in the form of tensor

yeah!

[liye0626]

@ydwu4
I added a loop in true_fn but I got the issue: torch._dynamo.exc.UncapturedHigherOrderOpError: Cond doesn't work unless it is captured completely with torch.compile. Scroll up to find out what causes the graph break.

Does “true_fn” not support "torch.arange"? Or is there something wrong with my usage for "torch.arange"?

Related code:

...
max_depth = max_depth.item()
torch._check_is_size(max_depth)
torch._check(max_depth < 10)
torch._check(max_depth > 0)

retrieve_indice = torch.full((1, max_depth),-1, dtype=torch.long)

cid = i

depth = position_ids_list[cid].item() + 1
torch._check_is_size(depth)
torch._check(depth < 10)
torch._check(depth > 0)


for j in torch.arange(depth):  #After adding this loop will result in an issue: torch._dynamo.exc.UncapturedHigherOrderOpError: Cond doesn't work unless it is captured completely with torch.compile. Scroll up to find out what causes the graph break.
    pass

[ydwu4]

can you paste the full error message? Maybe we should improve the error message and print the error directly instead of "Scroll up to find out what causes the graph break."

arrange is supported in cond's body_fn like

import torch
def g(x):
    def f(x):
        for i in torch.arange(10):
            x = x + i
        return x
    return torch.cond(x.sum() > 0, f, f, (x,))

out = torch.compile(g, fullgraph=True)(torch.randn(4, 4))
print(out)

works for me.