TorchScript Graph: Support model > 2GB

[BowenBao]

Dependent on torch.Graph._export_onnx, large model export requires specifying a target path to store the model and weights proto. Right now this path is fixed as "".

https://github.com/microsoft/onnx-script/blob/4986a4be9f8ff1b0e2b7bdde738836574f90e900/onnxscript/function_libs/torch_aten/graph_building.py#L475

[BowenBao]

2 models (at least, some larger ones are failing for other reasons) from bench are affected. New IR and updates for exporter serialization should take care of this.

2 like: The serialized model is larger than the 2GiB limit imposed by the protobuf library. Therefore the output file must be a file path, so that the ONNX external data can be written to the same directory. Please specify the output file name.

        AlbertForMaskedLM
        AlbertForQuestionAnswering

[justinchuby]

I think we do support external data with torchscript right? Or was I hallucinating

[BowenBao]

we do but we need a str path provided. A tedious restriction that we want to eliminate.

[BowenBao]

Removing low priority now that it is the only blocker left for torch.onnx.dynamo_export w/ llama.

[justinchuby]

Will we be able to provide an option in to_model_proto to leverage what we already have in torch for this?

[BowenBao]

I believe we have many options, need to pick one that delivers most natural export api and UX.

[justinchuby]

Would you like to share some of the options?

[BowenBao]

I'm thinking of model / weight separation by default.

[wschin]

Can someone please assign this to @wschin? I have a version for demo.