Dumps intermediate results of a torch model

Looking for discrepancies is quickly annoying. Discrepancies come from two results obtained with the same models implemented in two different ways, pytorch and onnx. Models are big so where do they come from? That’s the unavoidable question. Unless there is an obvious reason, the only way is to compare intermediate outputs alon the computation. The first step into that direction is to dump the intermediate results coming from pytorch. We use onnx_diagnostic.helpers.torch_helper.steal_forward() for that.

A simple LLM Model

See onnx_diagnostic.helpers.torch_helper.dummy_llm() for its definition. It is mostly used for unit test or example.

import onnx
import torch
from onnx_array_api.plotting.graphviz_helper import plot_dot
from onnx_diagnostic import doc
from onnx_diagnostic.helpers import string_type
from onnx_diagnostic.helpers.torch_helper import dummy_llm
from onnx_diagnostic.helpers.mini_onnx_builder import create_input_tensors_from_onnx_model
from onnx_diagnostic.helpers.torch_helper import steal_forward


model, inputs, ds = dummy_llm(dynamic_shapes=True)

print(f"type(model)={type(model)}")
print(f"inputs={string_type(inputs, with_shape=True)}")
print(f"ds={string_type(ds, with_shape=True)}")

type(model)=<class 'onnx_diagnostic.helpers.torch_helper.dummy_llm.<locals>.LLM'>
inputs=(T7s2x30,)
ds=dict(input_ids:{0:Dim(batch),1:Dim(length)})

It contains the following submodules.

for name, mod in model.named_modules():
    print(f"- {name}: {type(mod)}")

- : <class 'onnx_diagnostic.helpers.torch_helper.dummy_llm.<locals>.LLM'>
- embedding: <class 'onnx_diagnostic.helpers.torch_helper.dummy_llm.<locals>.Embedding'>
- embedding.embedding: <class 'torch.nn.modules.sparse.Embedding'>
- embedding.pe: <class 'torch.nn.modules.sparse.Embedding'>
- decoder: <class 'onnx_diagnostic.helpers.torch_helper.dummy_llm.<locals>.DecoderLayer'>
- decoder.attention: <class 'onnx_diagnostic.helpers.torch_helper.dummy_llm.<locals>.MultiAttentionBlock'>
- decoder.attention.attention: <class 'torch.nn.modules.container.ModuleList'>
- decoder.attention.attention.0: <class 'onnx_diagnostic.helpers.torch_helper.dummy_llm.<locals>.AttentionBlock'>
- decoder.attention.attention.0.query: <class 'torch.nn.modules.linear.Linear'>
- decoder.attention.attention.0.key: <class 'torch.nn.modules.linear.Linear'>
- decoder.attention.attention.0.value: <class 'torch.nn.modules.linear.Linear'>
- decoder.attention.attention.1: <class 'onnx_diagnostic.helpers.torch_helper.dummy_llm.<locals>.AttentionBlock'>
- decoder.attention.attention.1.query: <class 'torch.nn.modules.linear.Linear'>
- decoder.attention.attention.1.key: <class 'torch.nn.modules.linear.Linear'>
- decoder.attention.attention.1.value: <class 'torch.nn.modules.linear.Linear'>
- decoder.attention.linear: <class 'torch.nn.modules.linear.Linear'>
- decoder.feed_forward: <class 'onnx_diagnostic.helpers.torch_helper.dummy_llm.<locals>.FeedForward'>
- decoder.feed_forward.linear_1: <class 'torch.nn.modules.linear.Linear'>
- decoder.feed_forward.relu: <class 'torch.nn.modules.activation.ReLU'>
- decoder.feed_forward.linear_2: <class 'torch.nn.modules.linear.Linear'>
- decoder.norm_1: <class 'torch.nn.modules.normalization.LayerNorm'>
- decoder.norm_2: <class 'torch.nn.modules.normalization.LayerNorm'>

Steal and dump the output of submodules

The following context spies on the intermediate results for the following module and submodules. It stores in one onnx file all the input/output for those.

with steal_forward(
    [
        ("model", model),
        ("model.decoder", model.decoder),
        ("model.decoder.attention", model.decoder.attention),
        ("model.decoder.feed_forward", model.decoder.feed_forward),
        ("model.decoder.norm_1", model.decoder.norm_1),
        ("model.decoder.norm_2", model.decoder.norm_2),
    ],
    dump_file="plot_dump_intermediate_results.inputs.onnx",
    verbose=1,
    storage_limit=2**28,
):
    model(*inputs)

+model -- stolen forward for class LLM -- iteration 0
  <- args=(T7s2x30,) --- kwargs={}
+model.decoder -- stolen forward for class DecoderLayer -- iteration 0
  <- args=(T1s2x30x16,) --- kwargs={}
+model.decoder.norm_1 -- stolen forward for class LayerNorm -- iteration 0
  <- args=(T1s2x30x16,) --- kwargs={}
  -> T1s2x30x16
-model.decoder.norm_1.
-- stores key=('model.decoder.norm_1', 0), size 3Kb -- T1s2x30x16
+model.decoder.attention -- stolen forward for class MultiAttentionBlock -- iteration 0
  <- args=(T1s2x30x16,) --- kwargs={}
  -> T1s2x30x16
-model.decoder.attention.
-- stores key=('model.decoder.attention', 0), size 3Kb -- T1s2x30x16
+model.decoder.norm_2 -- stolen forward for class LayerNorm -- iteration 0
  <- args=(T1s2x30x16,) --- kwargs={}
  -> T1s2x30x16
-model.decoder.norm_2.
-- stores key=('model.decoder.norm_2', 0), size 3Kb -- T1s2x30x16
+model.decoder.feed_forward -- stolen forward for class FeedForward -- iteration 0
  <- args=(T1s2x30x16,) --- kwargs={}
  -> T1s2x30x16
-model.decoder.feed_forward.
-- stores key=('model.decoder.feed_forward', 0), size 3Kb -- T1s2x30x16
  -> T1s2x30x16
-model.decoder.
-- stores key=('model.decoder', 0), size 3Kb -- T1s2x30x16
  -> T1s2x30x16
-model.
-- stores key=('model', 0), size 3Kb -- T1s2x30x16
-- gather stored 12 objects, size=0 Mb
-- dumps stored objects
-- done dump stored objects

Restores saved inputs/outputs

All the intermediate tensors were saved in one unique onnx model, every tensor is stored in a constant node. The model can be run with any runtime to restore the inputs and function create_input_tensors_from_onnx_model can restore their names.

saved_tensors = create_input_tensors_from_onnx_model(
    "plot_dump_intermediate_results.inputs.onnx"
)
for k, v in saved_tensors.items():
    print(f"{k} -- {string_type(v, with_shape=True)}")

('model', 0, 'I') -- ((T7s2x30,),{})
('model.decoder', 0, 'I') -- ((T1s2x30x16,),{})
('model.decoder.norm_1', 0, 'I') -- ((T1s2x30x16,),{})
('model.decoder.norm_1', 0, 'O') -- T1s2x30x16
('model.decoder.attention', 0, 'I') -- ((T1s2x30x16,),{})
('model.decoder.attention', 0, 'O') -- T1s2x30x16
('model.decoder.norm_2', 0, 'I') -- ((T1s2x30x16,),{})
('model.decoder.norm_2', 0, 'O') -- T1s2x30x16
('model.decoder.feed_forward', 0, 'I') -- ((T1s2x30x16,),{})
('model.decoder.feed_forward', 0, 'O') -- T1s2x30x16
('model.decoder', 0, 'O') -- T1s2x30x16
('model', 0, 'O') -- T1s2x30x16

Let’s explained the naming convention.

('model.decoder.norm_2', 0, 'I') -- ((T1s2x30x16,),{})
            |            |   |
            |            |   +--> input, the format is args, kwargs
            |            |
            |            +--> iteration, 0 means the first time the execution
            |                 went through that module
            |                 it is possible to call multiple times,
            |                 the model to store more
            |
            +--> the name given to function steal_forward

The same goes for output except 'I' is replaced by 'O'.

('model.decoder.norm_2', 0, 'O') -- T1s2x30x16

This trick can be used to compare intermediate results coming from pytorch to any other implementation of the same model as long as it is possible to map the stored inputs/outputs.

Conversion to ONNX

The difficult point is to be able to map the saved intermediate results to intermediate results in ONNX. Let’s create the ONNX model.

epo = torch.onnx.export(model, inputs, dynamic_shapes=ds, dynamo=True)
epo.optimize()
epo.save("plot_dump_intermediate_results.onnx")

[torch.onnx] Obtain model graph for `LLM([...]` with `torch.export.export(..., strict=False)`...
[torch.onnx] Obtain model graph for `LLM([...]` with `torch.export.export(..., strict=False)`... ✅
[torch.onnx] Run decomposition...
[torch.onnx] Run decomposition... ✅
[torch.onnx] Translate the graph into ONNX...
[torch.onnx] Translate the graph into ONNX... ✅
Applied 4 of general pattern rewrite rules.

It looks like the following.

onx = onnx.load("plot_dump_intermediate_results.onnx")
plot_dot(onx)

doc.plot_legend("steal and dump\nintermediate\nresults", "steal_forward", "blue")

Related examples

Export microsoft/phi-2

Export microsoft/phi-2

Steel method forward to guess inputs and dynamic shapes (with Tiny-LLM)

Steel method forward to guess inputs and dynamic shapes (with Tiny-LLM)

Find and fix an export issue due to dynamic shapes

Find and fix an export issue due to dynamic shapes