[ONNX] dynamo_export() fails to automatically switch to external weights for large models

[borisfom]

🐛 Describe the bug

onnx.export() switches automatically to external weights when protobuf size exceeds 2G limit.
dynamo_export() does not do that.
It does generate external weights if model_with_state_dict arg is supplied (which ends up not usable, too, I will file another issue on that) - but otherwise it creates a huge .onnx file that subsequently fails to parse.
Here is a repro. You need to install NVIDIA NeMo toolkit first:
pip install --extra-index-url https://pypi.ngc.nvidia.com/ nemo_toolkit[all]

import torch
from nemo.core.classes import typecheck
from nemo.utils.export_utils import wrap_forward_method
from nemo.collections.asr.models import EncDecCTCModelBPE

# This model works (small)                                                                                                                
# model = EncDecCTCModelBPE.from_pretrained(model_name="stt_en_conformer_ctc_small")                                                      
# This model fails (large)                                                                                                                
model = EncDecCTCModelBPE.from_pretrained(model_name="nvidia/parakeet-ctc-1.1b")

model = model.eval()
wrap_forward_method(model)
model._prepare_for_export()
typecheck.set_typecheck_enabled(enabled=False)

with torch.no_grad():
    input_example = model.input_module.input_example(max_batch=4, max_dim=256)
    batch = torch.export.Dim("batch", max=64)
    dynamic_shapes={'audio_signal': { 0: batch}, 'length': {0: batch}}

    ex_model = torch.export.export(
        model, tuple(input_example),
        dynamic_shapes=dynamic_shapes,
        strict=False
    )

    ex_model = ex_model.run_decompositions()
    ex_module=ex_model.module()

    print("Running torch.onnx.dynamo_export ...")
    options = torch.onnx.ExportOptions(dynamic_shapes=True)
    ex = torch.onnx.dynamo_export(ex_module, *input_example, export_options=options)
    ex.save("pk.onnx")

 

### Versions

PyTorch nightly 05/14


cc @avikchaudhuri @gmagogsfm @zhxchen17 @tugsbayasgalan @angelayi @suo @ydwu4

[titaiwangms]

run_decomposition() is called inside dynamo_export:

pytorch/torch/onnx/_internal/fx/torch_export_graph_extractor.py

Line 93 in 796dff7

model = model.run_decompositions(options.decomposition_table)

Is feeding ExportedProgram to dynamo_export triggers the same issue? As far as I know the input model of dynamo_export should be either nn.Module, or ExportedProgram.

[borisfom]

Well, run_decompositions wasn't run in dynamo_export with neither Module nor ExportedProgram at least until yesterday - or perhaps it was run with wrong/null table:
#125894

[borisfom]

@titaiwangms : to answer your question "Is feeding ExportedProgram to dynamo_export triggers the same issue?" : no, but that version is always using external weights, even for small models. I have filed separate issue for that : #126479

[borisfom]

Note: for repro to work, onnxscript also has to be patched (I proposed a fix but it's not merged yet) :
microsoft/onnxscript#1541