Skip to content

Stop ct.convert from rewriting the traced model's TorchScript graph - #2859

Open
Om-singhaI wants to merge 1 commit into
apple:mainfrom
Om-singhaI:fix/torchscript-graph-copy
Open

Om-singhaI wants to merge 1 commit into
apple:mainfrom
Om-singhaI:fix/torchscript-graph-copy

Conversation

@Om-singhaI

Copy link
Copy Markdown
Contributor

Fixes #2215

ct.convert changes the traced model you pass in, and for some models that means you can't save and load it afterwards:

class Net(torch.nn.Module):
    def forward(self, x):
        a, b, c = x.chunk(3)
        return (a * b) + c

traced = torch.jit.trace(Net().eval(), torch.rand(6, 4))
ct.convert(traced, inputs=[ct.TensorType(shape=(6, 4))], convert_to="milinternal")
torch.jit.save(traced, "m.pt")
torch.jit.load("m.pt")  # RuntimeError: required keyword attribute 'chunks' is undefined

_expand_and_optimize_ir takes torchscript.forward.graph and runs the inline, DCE and _jit_pass_canonicalize_graph_fuser_ops passes on it in place. That's the user's own graph, so their module ends up holding a prim::ConstantChunk node that TorchScript can't load back.

The repro uses convert_to="milinternal" so it runs without the compiled backend, but the default conversion goes through the same frontend code.

The fix runs those passes on torchscript.forward.graph.copy() instead. The converter only reads the graph this function returns, so conversion itself doesn't change. I checked traced and scripted models (chunk, nested modules with BatchNorm, a shared parameter, fork plus an if block), and the lowered graph, input names, params and buffers all come out identical to main. The user's graph is now left alone.

I added test_convert_does_not_modify_torchscript_model. On main it fails because the traced graph gains a prim::ConstantChunk node, and with the fix it passes:

pytest coremltools/converters/mil/frontend/torch/test/test_torch_conversion_api.py -k "does_not_modify_torchscript_model or traced_model_to_milinternal"
2 passed

I ran it on torch 2.8.0. I didn't run the full frontend suite locally.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

ct.convert call appears to corrupt torchscript model

1 participant