torch._dynamo.exc.Unsupported: dynamic shapes: arange

🐛 Describe the bug

While giving a try with pytorch2 on OpenNMT-py

using these two lines:

    rawmodel = build_model(model_opt, opt, vocabs, checkpoint)
    model = torch.compile(rawmodel, fullgraph=True, backend='nvprims_aten')

Error logs

getting this:

from user code:
   File "/home/vincent/nlp/OpenNMT-py/onmt/encoders/transformer.py", line 126, in forward
    mask = ~sequence_mask(src_len).unsqueeze(1)
  File "/home/vincent/nlp/OpenNMT-py/onmt/utils/misc.py", line 58, in sequence_mask
    return (torch.arange(0, max_len, device=lengths.device)

Set torch._dynamo.config.verbose=True for more information


You can suppress this exception and fall back to eager by setting:
    torch._dynamo.config.suppress_errors = True

Minified repro

No response

[vince62s]

because of this: https://github.com/pytorch/pytorch/issues/90170

not so easy to run a minifier but will try.

[vince62s]

okay, I managed to tweak triton to recognize ptxas 7.8 / cuda 11.8 and inductor mode does not trigger any error. However nothing seems happening got a loop with both the TypedStorage warning and: [2022-12-09 13:33:41,820] torch._inductor.lowering: [WARNING] using triton random, expect difference from eager no log from my training loop. nvidia-smi show some activty (both ram and util) but without any other message difficult to go further.

EDIT: I’ll try module by module to see where the problem could be.