Status of torch-mlir with pytorch2.0

[accelerate321]

Hi,
Can torch-mlir support training with pytorch2.0 torchdynamo & aot_autograd now, I try your resnet18 example, but it is an inference demo. I train the resnet18 with aot_autograd but fail. Can you describe the status of supporting pytorch2.0.

[ramiro050]

Hi @accelerate321,

Can you give more detail on the error you're seeing as well as a simple example to reproduce it? Training should be achievable with torchdynamo, although some tinkering with the graphs might be required to get it to work. See this thread for an issue someone else was facing when doing training:

#1685

[accelerate321]

I just modify the resnet18 example to training like this. I am not sure whether my code is right, or you can give me a correct demo.

...
resnet18.train()
optimizer = torch.optim.SGD(resnet18.parameters(), lr=0.01)
def training_iter_fn(batch, model, optimizer):
    optimizer.zero_grad()
    out = model(batch)
    loss = out.sum()
    loss.backward()
    optimizer.step()
    return loss

from torch._dynamo.optimizations.training import aot_autograd
compiled_model = dynamo.optimize(aot_autograd(fw_compiler=refbackend_torchdynamo_backend))(training_iter_fn)
out = compiled_model(img, resnet18, optimizer)
print(out)

The error is:

torch._dynamo.exc.BackendCompilerFailed: compiler_fn raised Exception: Please convert all Tensors to FakeTensors first or instantiate FakeTensorMode with 'allow_non_fake_inputs'. Found in aten.clone.default(*(tensor([114,  66, 217,  ...,   0,   0,   0], dtype=torch.uint8),), **{}) 

[silvasean]

There's still work to be done on TorchDynamo integration. I wouldn't expect this to work in all cases out of the box just yet.

[accelerate321]

Is there a specific time point to support the work on TorchDynamo integration?

[silvasean]

There is no specific time point yet. There is work on our end to do but also a significant number of upstream dynamo/PT2-related bugs that are not resolved upstream yet.