Note
Go to the end to download the full example code.
Export with DynamicCache and dynamic shapes¶
Every LLMs implemented in transformers use cache.
One of the most used is transformers.cache_utils.DynamicCache.
The cache size is dynamic to cope with the growing context.
The example shows a tool which determines the dynamic shapes
for torch.export.export() based on a set of valid inputs.
DynamicCache¶
torch.export.export() serializes caches and any custom class
if these serialization functions are provided with is the case for
transformers.cache_utils.DynamicCache and transformers>=4.50.
The dynamic shapes must be provided following the serialized form.
import pprint
import torch
from onnx_diagnostic import doc
from onnx_diagnostic.ext_test_case import has_transformers
from onnx_diagnostic.helpers import string_type
from onnx_diagnostic.helpers.cache_helper import (
flatten_unflatten_for_dynamic_shapes,
make_dynamic_cache,
)
from onnx_diagnostic.export import ModelInputs
from onnx_diagnostic.torch_export_patches import torch_export_patches
class Model(torch.nn.Module):
def forward(self, cache, z):
return (
z
+ cache.key_cache[0]
+ cache.key_cache[1]
+ cache.value_cache[0]
+ cache.value_cache[1]
)
model = Model()
n_layers = 2
bsize, nheads, slen, dim = 2, 4, 3, 7
cache = make_dynamic_cache(
[
(torch.randn(bsize, nheads, slen, dim), torch.randn(bsize, nheads, slen, dim))
for i in range(n_layers)
]
)
z = torch.randn((1, 1, 1, 7))
model(cache, z) # to check it works.
tensor([[[[-2.7787, -4.0491, 0.5901, 1.8470, 1.2857, 3.2397, 0.7217],
[-3.6029, 0.7740, 4.2963, -1.1397, -3.0528, 2.3129, -1.6724],
[-2.5554, 5.1892, 2.2846, 0.3291, 1.8201, 2.9657, -0.5528]],
[[-4.3943, 0.7172, 3.4131, 3.5816, 2.6205, 0.8995, 0.0105],
[-0.1534, 1.2352, 0.0374, 0.8831, 0.6090, 1.4126, -0.8399],
[-1.8033, -3.3073, -2.1491, -2.6937, -0.5817, 0.1334, -1.5425]],
[[-2.3743, 2.8962, 1.2278, 1.1284, 0.0147, -3.2003, -0.5621],
[-4.1170, -2.4746, 2.1034, -0.1996, -1.6322, -3.0480, 1.6414],
[-3.9322, -1.8560, 0.3180, 0.2102, -0.7806, -4.1531, 1.0859]],
[[ 0.1846, 3.8178, -3.9087, 2.2458, -0.3611, -0.1528, -1.1948],
[-3.2047, -1.7107, 3.4738, 0.3071, 3.2227, -0.7417, -0.9667],
[-2.5967, -0.9199, 0.9515, -1.3049, -1.5334, 0.0790, -0.9437]]],
[[[ 0.2117, 2.8061, 2.2056, 0.4865, -2.0913, -0.3447, 1.1472],
[-1.5308, 4.1224, 2.0390, 2.6326, 0.6907, -2.1033, -1.4964],
[ 2.8406, 3.8819, 1.4042, 1.1646, -2.6573, 0.2164, -1.3676]],
[[-0.1977, 2.6567, 2.6405, 3.1487, 4.0104, -2.9812, -1.0712],
[-4.0589, -1.0759, 3.9501, 2.4040, 0.5540, -0.9266, -3.6990],
[-2.4832, 0.5405, 0.2627, 2.6843, 2.8122, 0.5137, 0.6774]],
[[-1.2136, 0.6183, 0.6393, 0.0591, 0.4524, 0.1407, -0.4286],
[ 0.9989, 0.0877, -2.0250, -0.0813, 1.5826, 1.2435, 3.4254],
[-0.9514, 1.5891, -3.6032, 1.4337, -1.3183, -3.8903, -0.9285]],
[[-3.0507, 0.4791, 3.4713, -0.2705, 2.7737, 0.9631, 1.9934],
[ 0.1608, 1.6646, 1.4608, 1.0155, 1.7904, 3.6289, -2.8592],
[-2.8865, 2.6105, 0.4324, 3.0624, 0.3097, -1.7066, 0.2701]]]])
The cache looks like this:
print(string_type(cache, with_shape=True))
DynamicCache(key_cache=#2[T1s2x4x3x7,T1s2x4x3x7], value_cache=#2[T1s2x4x3x7,T1s2x4x3x7])
cache2 = make_dynamic_cache(
[
(
torch.randn(bsize + 1, nheads, slen + 1, dim + 1),
torch.randn(bsize + 1, nheads, slen + 1, dim + 1),
)
for i in range(n_layers)
]
)
inputs = [
(cache, z),
(cache2, torch.randn((1, 1, 1, 8))),
]
And the second set of inputs looks like:
print(string_type(inputs[1], with_shape=True))
(DynamicCache(key_cache=#2[T1s3x4x4x8,T1s3x4x4x8], value_cache=#2[T1s3x4x4x8,T1s3x4x4x8]),T1s1x1x1x8)
Guess the dynamic shapes¶
The following tool can be used to guess the dynamic shapes
the way torch.export.export() expects them.
mi = ModelInputs(Model(), inputs)
ds = mi.guess_dynamic_shapes()
pprint.pprint(ds)
(([[{0: _DimHint(type=<_DimHintType.DYNAMIC: 3>,
min=None,
max=None,
_factory=True),
2: _DimHint(type=<_DimHintType.DYNAMIC: 3>,
min=None,
max=None,
_factory=True),
3: _DimHint(type=<_DimHintType.DYNAMIC: 3>,
min=None,
max=None,
_factory=True)},
{0: _DimHint(type=<_DimHintType.DYNAMIC: 3>,
min=None,
max=None,
_factory=True),
2: _DimHint(type=<_DimHintType.DYNAMIC: 3>,
min=None,
max=None,
_factory=True),
3: _DimHint(type=<_DimHintType.DYNAMIC: 3>,
min=None,
max=None,
_factory=True)}],
[{0: _DimHint(type=<_DimHintType.DYNAMIC: 3>,
min=None,
max=None,
_factory=True),
2: _DimHint(type=<_DimHintType.DYNAMIC: 3>,
min=None,
max=None,
_factory=True),
3: _DimHint(type=<_DimHintType.DYNAMIC: 3>,
min=None,
max=None,
_factory=True)},
{0: _DimHint(type=<_DimHintType.DYNAMIC: 3>,
min=None,
max=None,
_factory=True),
2: _DimHint(type=<_DimHintType.DYNAMIC: 3>,
min=None,
max=None,
_factory=True),
3: _DimHint(type=<_DimHintType.DYNAMIC: 3>,
min=None,
max=None,
_factory=True)}]],
{3: _DimHint(type=<_DimHintType.DYNAMIC: 3>,
min=None,
max=None,
_factory=True)}),
{})
And finally the export.
The export is simple if transformers>=4.50, otherwise,
transformers needs to be patched.
onnx_diagnostic.torch_export_patches.torch_export_patches()
registers functions to serialize DynamicCache. This one is modified to make
the shape inference implemented in torch happy.
if has_transformers("4.50"):
ep = torch.export.export(model, inputs[0], dynamic_shapes=ds[0], strict=False)
else:
with torch_export_patches(patch_transformers=True) as modificator:
ep = torch.export.export(
model, modificator(inputs[0]), dynamic_shapes=ds[0], strict=False
)
print(ep)
ExportedProgram:
class GraphModule(torch.nn.Module):
def forward(self, cache_key_cache_0: "f32[s26, 4, s28, s1]", cache_key_cache_1: "f32[s26, 4, s28, s1]", cache_value_cache_0: "f32[s26, 4, s28, s1]", cache_value_cache_1: "f32[s26, 4, s28, s1]", z: "f32[1, 1, 1, s1]"):
# File: /home/xadupre/github/onnx-diagnostic/_doc/examples/plot_export_with_dynamic_cache.py:39 in forward, code: z
add: "f32[s26, 4, s28, s1]" = torch.ops.aten.add.Tensor(z, cache_key_cache_0); z = cache_key_cache_0 = None
add_1: "f32[s26, 4, s28, s1]" = torch.ops.aten.add.Tensor(add, cache_key_cache_1); add = cache_key_cache_1 = None
add_2: "f32[s26, 4, s28, s1]" = torch.ops.aten.add.Tensor(add_1, cache_value_cache_0); add_1 = cache_value_cache_0 = None
add_3: "f32[s26, 4, s28, s1]" = torch.ops.aten.add.Tensor(add_2, cache_value_cache_1); add_2 = cache_value_cache_1 = None
return (add_3,)
Graph signature:
# inputs
cache_key_cache_0: USER_INPUT
cache_key_cache_1: USER_INPUT
cache_value_cache_0: USER_INPUT
cache_value_cache_1: USER_INPUT
z: USER_INPUT
# outputs
add_3: USER_OUTPUT
Range constraints: {s26: VR[2, int_oo], s28: VR[2, int_oo], s1: VR[2, int_oo]}
Do we need to guess?¶
Function onnx_diagnostic.helpers.string_type() is using
the serialization functions to print out the DynamicCache the was
torch.export.export() expects them.
print(string_type(cache, with_shape=True))
DynamicCache(key_cache=#2[T1s2x4x3x7,T1s2x4x3x7], value_cache=#2[T1s2x4x3x7,T1s2x4x3x7])
You can also use function
onnx_diagnostic.helpers.cache_helper.flatten_unflatten_for_dynamic_shapes()
to show a DynamicCache restructured the way torch.export.export() expects
it to be without the custom class.
print(string_type(flatten_unflatten_for_dynamic_shapes(cache), with_shape=True))
#2[#2[T1s2x4x3x7,T1s2x4x3x7],#2[T1s2x4x3x7,T1s2x4x3x7]]
This code works for any custom class if it was registered
with torch.utils._pytree.register_pytree_node().
doc.plot_legend("dynamic shapes\nfor DynamicCache", "torch.export.export", "tomato")
Total running time of the script: (0 minutes 0.371 seconds)
Related examples
Steel method forward to guess inputs and dynamic shapes (with Tiny-LLM)