yobx.torch.tiny_models#

class yobx.torch.tiny_models.ModelData(model_id: str, model: Module, export_inputs: Dict[str, Any], dynamic_shapes: Dict[str, Any], inputs_batch1: Dict[str, Any] | None = None)[source]#

Contains all the necessary information to export a model.

property dynamic_shapes_for_torch_export_export: Dict[str, Any]#: The dynamic shapes contains strings. torch.export.export() needs them to be replaced by torch.export.Dim.DYNAMIC. This is what this property is doing.

class yobx.torch.tiny_models.TinyBroadcastAddModel(*args: Any, **kwargs: Any)[source]#

A model where one output dynamic dimension becomes max(d1, d2) after a broadcast.

Inputs x (shape (batch, d1)) and y (shape (batch, d2)) are added element-wise. Broadcasting rules require that d1 == d2 or one of them equals 1 at runtime; the symbolic output shape is (batch, max(d1, d2)).

This model triggers the following when exported.

<<<

import torch
from yobx.helpers import string_type
from yobx.torch import apply_patches_for_model, use_dyn_not_str
from yobx.torch.tiny_models import TinyBroadcastAddModel

model = TinyBroadcastAddModel()
export_inputs = TinyBroadcastAddModel._export_inputs()
dynamic_shapes = use_dyn_not_str(TinyBroadcastAddModel._dynamic_shapes())

print(f"-- inputs: {string_type(export_inputs, with_shape=True)}")
print(f"-- shapes: {dynamic_shapes}")
print("--")
print("-- simple export --")
print("--")

try:
    torch.export.export(
        model,
        (),
        kwargs=export_inputs,
        dynamic_shapes=dynamic_shapes,
    )
except Exception as e:
    print(e)

print("--")
print("-- export with backed_size_oblivious=True --")
print("--")

with torch.fx.experimental._config.patch(backed_size_oblivious=True):
    try:
        torch.export.export(
            model,
            (),
            kwargs=export_inputs,
            dynamic_shapes=dynamic_shapes,
        )
    except Exception as e:
        print(e)

print("--")
print("-- patched export --")
print("--")

with (
    torch.fx.experimental._config.patch(backed_size_oblivious=True),
    apply_patches_for_model(patch_torch=True),
):
    ep = torch.export.export(
        model,
        (),
        kwargs=export_inputs,
        dynamic_shapes=dynamic_shapes,
    )
    print(ep)

>>>

    -- inputs: dict(x:T1s2x5,y:T1s2x1)
    -- shapes: {'x': {0: DimHint(DYNAMIC), 1: DimHint(DYNAMIC)}, 'y': {0: DimHint(DYNAMIC), 1: DimHint(DYNAMIC)}}
    --
    -- simple export --
    --
    --
    -- export with backed_size_oblivious=True --
    --
    The size of tensor a (s27) must match the size of tensor b (s94) at non-singleton dimension 1)
    --
    -- patched export --
    --
    ExportedProgram:
        class GraphModule(torch.nn.Module):
            def forward(self, x: "f32[s77, s27]", y: "f32[s17, s94]"):
                # File: ~/github/yet-another-onnx-builder/yobx/torch/tiny_models.py:128 in forward, code: return x + y
                add: "f32[Max(s17, s77), Max(s27, s94)]" = torch.ops.aten.add.Tensor(x, y);  x = y = None
                return (add,)
        
    Graph signature: 
        # inputs
        x: USER_INPUT
        y: USER_INPUT
        
        # outputs
        add: USER_OUTPUT
        
    Range constraints: {s77: VR[0, int_oo], s27: VR[0, int_oo], s17: VR[0, int_oo], s94: VR[1, 2]}

forward(x: Tensor, y: Tensor) → Tensor[source]#

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

yobx.torch.tiny_models.get_tiny_model(model_id, config_updates: Dict[str, Any] | None = None) → ModelData[source]#

Creates a tiny models, usually untrained to write tests. The model_id refers to what exists on HuggingFace. This functions is not expected to have fully coverage of architectures.

Supported model IDs:

"arnir0/Tiny-LLM" — a tiny LLaMA-based causal language model with a transformers.cache_utils.DynamicCache past-key-value cache.
"local/BroadcastAdd" — a minimal two-input model whose output has the symbolic shape (batch, max(d1, d2)) due to broadcasting, see yobx.torch.tiny_models.TinyBroadcastAddModel.

Parameters:

model_id – model id, see the list of supported values above
config_updates – modification to add to the configuration before creating the model

Returns:

the necessary information

Example:

import torch
from yobx.torch import get_tiny_model
from yobx.torch.torch_helper import torch_deepcopy

model_data = get_tiny_model("arnir0/Tiny-LLM")
# torch_deepcopy is needed because the cache (past_key_values) is modified in-place
# during the forward pass; reusing the same object would corrupt subsequent calls.
result = model_data.model(**torch_deepcopy(model_data.export_inputs))
ep = torch.export.export(
    model_data.model,
    (),
    kwargs=torch_deepcopy(model_data.export_inputs),
    dynamic_shapes=model_data.dynamic_shapes,
)

yobx.torch.tiny_models#

This Page