onnxscript 0.1.0__py3-none-any.whl

onnxscript/tools/memory_peak.py

@@ -0,0 +1,244 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+# pylint: disable=import-outside-toplevel
+from __future__ import annotations
+
+import multiprocessing
+import os
+
+
+def get_memory_rss(pid: int) -> int:
+    """
+    Returns the physical memory used by a process.
+
+    Args:
+        pid: Process id, current one is `os.getpid()`.
+
+    Returns:
+         Physical memory.
+
+    It relies on the module *psutil*.
+    """
+    import psutil
+
+    process = psutil.Process(pid)
+    mem = process.memory_info().rss
+    return mem
+
+
+class Monitor:
+    def __init__(self):
+        self.max_peak: float = 0
+        self.average: float = 0
+        self.n_measures: int = 0
+        self.begin: float = 0
+        self.end: float = 0
+
+    def to_dict(self, unit: int = 1) -> dict[str, float]:
+        funit = float(unit)
+        return dict(
+            peak=self.max_peak / funit,
+            mean=self.average * 1.0 / self.n_measures / funit,
+            n=self.n_measures / funit,
+            begin=self.begin / funit,
+            end=self.end / funit,
+        )
+
+    def __repr__(self) -> str:
+        return (
+            f"{self.__class__.__name__}(peak={self.max_peak}, "
+            f"average={self.average}, n={self.n_measures})"
+        )
+
+    def update(self, mem: float):
+        if self.n_measures == 0:
+            self.begin = mem
+        self.max_peak = max(mem, self.max_peak)
+        self.average += mem
+        self.end = mem
+        self.n_measures += 1
+
+    def send(self, conn):
+        conn.send(self.max_peak)
+        conn.send(self.average)
+        conn.send(self.n_measures)
+        conn.send(self.begin)
+        conn.send(self.end)
+
+    @classmethod
+    def recv(cls, conn) -> Monitor:
+        m = cls()
+        m.max_peak = conn.recv()
+        m.average = conn.recv()
+        m.n_measures = conn.recv()
+        m.begin = conn.recv()
+        m.end = conn.recv()
+        return m
+
+
+def _process_memory_spy(conn):
+    # Sends the value it started.
+    conn.send(-2)
+
+    # process id to spy on
+    pid = conn.recv()
+
+    # delay between two measures
+    timeout = conn.recv()
+
+    # do CUDA
+    cuda = conn.recv()
+
+    import psutil
+
+    process = psutil.Process(pid)
+
+    if cuda:
+        from pynvml import (  # type: ignore[import-not-found]
+            nvmlDeviceGetCount,
+            nvmlDeviceGetHandleByIndex,
+            nvmlDeviceGetMemoryInfo,
+            nvmlInit,
+            nvmlShutdown,
+        )
+
+        nvmlInit()
+        n_gpus = nvmlDeviceGetCount()
+        handles = [nvmlDeviceGetHandleByIndex(i) for i in range(n_gpus)]
+
+        def gpu_used():
+            return [nvmlDeviceGetMemoryInfo(h).used for h in handles]
+
+        gpus = [Monitor() for i in range(n_gpus)]
+    else:
+        gpus = []
+
+    cpu = Monitor()
+
+    conn.send(-2)
+
+    # loop
+    while True:
+        mem = process.memory_info().rss
+        cpu.update(mem)
+        if cuda:
+            for r, g in zip(gpu_used(), gpus):
+                g.update(r)
+        if conn.poll(timeout=timeout):
+            code = conn.recv()
+            if code == -3:
+                break
+
+    # final iteration
+    end = process.memory_info().rss
+    cpu.update(end)
+    if cuda:
+        for r, g in zip(gpu_used(), gpus):
+            g.update(r)
+
+    # send
+    cpu.send(conn)
+    conn.send(len(gpus))
+    for g in gpus:
+        g.send(conn)
+    if cuda:
+        nvmlShutdown()
+    conn.close()
+
+
+class MemorySpy:
+    """
+    Information about the spy. It class method `start`.
+    Method `stop` can be called to end the measure.
+
+    Args:
+        pid: process id  of the process to spy on
+        delay: spy on every delay seconds
+        cuda: enable cuda monitoring
+    """
+
+    def __init__(self, pid: int, delay: float = 0.01, cuda: bool = False):
+        self.pid = pid
+        self.delay = delay
+        self.cuda = cuda
+        self.start()
+
+    def start(self) -> MemorySpy:
+        """Starts another process and tells it to spy."""
+        self.parent_conn, self.child_conn = multiprocessing.Pipe()
+        self.child_process = multiprocessing.Process(
+            target=_process_memory_spy, args=(self.child_conn,)
+        )
+        self.child_process.start()
+        data = self.parent_conn.recv()
+        if data != -2:
+            raise RuntimeError(f"The child processing is supposed to send -2 not {data}.")
+        self.parent_conn.send(self.pid)
+        self.parent_conn.send(self.delay)
+        self.parent_conn.send(1 if self.cuda else 0)
+        data = self.parent_conn.recv()
+        if data != -2:
+            raise RuntimeError(
+                f"The child processing is supposed to send -2 again not {data}."
+            )
+        return self
+
+    def stop(self) -> dict[str, list[Monitor]]:
+        """Stops spying on."""
+        self.parent_conn.send(-3)
+
+        cpu = [Monitor.recv(self.parent_conn)]
+
+        n_gpus = self.parent_conn.recv()
+        gpus = []
+        for _ in range(n_gpus):
+            gpus.append(Monitor.recv(self.parent_conn))
+
+        self.parent_conn.close()
+        self.child_process.join()
+        res = dict(cpu=cpu)
+        if self.cuda:
+            res["gpus"] = gpus
+        return res
+
+
+def start_spying_on(
+    pid: int | None = None, delay: float = 0.01, cuda: bool = False
+) -> MemorySpy:
+    """Starts the memory spy. The function starts another
+    process spying on the one sent as an argument.
+
+    Example::
+
+    .. code-block:: python
+
+        from onnxscript.tools.memory_peak import start_spying_on, flatten
+
+        p = start_spying_on()
+        # ...
+        # code to measure
+        # ...
+        stat = p.stop()
+        print(stat)
+        print(flatten(stat))
+
+    Args:
+        pid: process id to spy or the the current one.
+        delay: delay between two measures.
+        cuda: True or False to get memory for cuda devices
+    """
+    if pid is None:
+        pid = os.getpid()
+    return MemorySpy(pid, delay, cuda)
+
+
+def flatten(ps: dict[str, list[Monitor]], prefix: str = "") -> dict[str, float]:
+    """Flattens a dictionary produced by :meth:`MemorySpy.stop`."""
+    obs = ps["cpu"][0].to_dict(unit=2**20)
+    if "gpus" in ps:
+        for i, g in enumerate(ps["gpus"]):
+            for k, v in g.to_dict(unit=2**20).items():
+                obs[f"gpu{i}_{k}"] = v
+    if prefix:
+        obs = {f"{prefix}{k}": v for k, v in obs.items()}
+    return obs

onnxscript/tools/training_helper.py

@@ -0,0 +1,50 @@
+# -------------------------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+# --------------------------------------------------------------------------
+import torch
+from torch.onnx import ExportOptions
+from torch.onnx import _OrtBackend as OrtBackend
+from torch.onnx import _OrtBackendOptions as OrtBackendOptions
+
+
+def make_aot_ort(dynamic: bool = False):
+    """Implements an autograd backend for torch.compile based on onnxrt backend."""
+    export_options = ExportOptions(dynamic_shapes=dynamic)
+    options = OrtBackendOptions(export_options=export_options)
+    ort_backend = OrtBackend(options=options)
+    return ort_backend
+
+
+def train_loop(model, *args, loss_fn=None, optimizer=None):
+    """Implements a training loop to be used in tests."""
+
+    if loss_fn is None:
+        loss_fn = torch.nn.MSELoss()
+    if optimizer is None:
+        optimizer = torch.optim.SGD(model.parameters(), lr=1e-3)
+
+    # Set the model to training mode - important for batch normalization and dropout layers
+    # Unnecessary in this situation but added for best practices
+    model.train()
+
+    # Compute prediction and loss
+    pred = model(*args)
+    if isinstance(pred, tuple):
+        v = pred[0]
+    elif hasattr(pred, "last_hidden_state"):
+        v = pred.last_hidden_state
+    else:
+        v = pred
+    loss = loss_fn(v, torch.ones_like(v))
+
+    # Backpropagation
+    loss.backward()
+    optimizer.step()
+    # skip that part to retrieve the gradients
+    # optimizer.zero_grad()
+
+    # returns the gradients
+    res = tuple(p.grad for p in model.parameters() if p.grad is not None)
+    assert len(res) > 0, f"No gradient, loss is {loss}"
+    return res

onnxscript/tools/transformers_models/__init__.py

@@ -0,0 +1,190 @@
+# -------------------------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+# --------------------------------------------------------------------------
+# pylint: disable=import-outside-toplevel
+from __future__ import annotations
+
+import random
+from typing import Any, Sequence
+
+import onnx
+import onnx.inliner
+import torch
+
+import onnxscript.optimizer
+import onnxscript.rewriter
+
+
+def export_to_onnx(
+    model: Any,
+    *args: Sequence[Any],
+    optimize: bool = True,
+    export_api: bool = True,
+    no_grad: bool = False,
+) -> onnx.ModelProto:
+    """
+    Export a model to ONNX.
+    If optimize is True, it calls *onnxscript.optimizer.optimize*,
+    *onnxscript.rewriter.rewriter*, *onnx.inliner.inline_local_functions*.
+    If *export_api* is True, the function uses ``torch.onnx.export``
+    and not ``torch.onnx.dynamo_export``.
+    """
+    if no_grad:
+        with torch.no_grad():
+            if export_api:
+                prog = torch.onnx.export(model, args, dynamo=True)  # pylint: disable=no-value-for-parameter
+            else:
+                prog = torch.onnx.dynamo_export(model, *args)
+    else:
+        if export_api:
+            prog = torch.onnx.export(model, args, dynamo=True)  # pylint: disable=no-value-for-parameter
+        else:
+            prog = torch.onnx.dynamo_export(model, *args)
+    assert prog is not None
+    model_proto = prog.model_proto
+    if optimize:
+        model_proto = onnxscript.optimizer.optimize(
+            model_proto,
+            num_iterations=2,
+            onnx_shape_inference=True,
+        )
+        model_proto = onnxscript.rewriter.rewrite(model_proto)
+        model_proto = onnx.inliner.inline_local_functions(model_proto)
+    return model_proto
+
+
+def ids_tensor(
+    shape: Sequence[int],
+    vocab_size: int,
+    rng: random.Random | None = None,
+    name: str | None = None,
+):
+    """Creates a random int32 tensor of the shape within the vocab size."""
+    del name  # unused
+
+    if rng is None:
+        rng = random.Random()
+
+    total_dims = 1
+    for dim in shape:
+        total_dims *= dim
+
+    values = []
+    for _ in range(total_dims):
+        values.append(rng.randint(0, vocab_size - 1))
+
+    return torch.tensor(data=values, dtype=torch.long).view(shape).contiguous()
+
+
+def get_input_dims_for_llm(
+    dynamic_shapes: bool, warmup: int, repeat: int
+) -> list[tuple[int, int]]:
+    """Returns input dimensions for model such as llama, phi, ..."""
+    if not dynamic_shapes:
+        return [(2, 1024)] * (warmup + repeat)
+    w = [(2, 1024), (3, 1024), (2, 1096)] * warmup
+    w = w[:warmup]
+    r = [(2, 1024), (3, 1024), (4, 1024), (2, 1096), (2, 1112)] * repeat
+    r = r[:repeat]
+    return w + r
+
+
+def get_model_and_inputs(
+    model: str,
+    config: str,
+    dynamic_shapes: bool,
+    device: str = "cpu",
+    num_hidden_layers: int = 1,
+    with_mask: bool = True,
+    implementation: str = "eager",
+    dtype: str | None = None,
+    warmup: int = 5,
+    repeat: int = 10,
+) -> tuple[Any, list[tuple[torch.Tensor, ...]], dict | None]:
+    """
+    Returns a model and a couple of dummy inputs.
+
+    Args:
+        model: model name, 'phi', 'llama', 'phi3', ...
+        config: 'small', 'medium', 'large', ...
+        dynamic_shapes: dynamic or static shapes
+        device: 'cpu' or 'cuda'
+        num_hidden_layers: Number of hidden layers.
+        with_mask: One input or two inputs.
+        implementation: eager or sdpa
+        warmup: Number of inputs to generate.
+        repeat: Number of inputs to generate for repeat.
+        dtype: If specified, cast the model and the inputs into this type.
+
+    Returns:
+        model and list of inputs
+    """
+    if model == "llama":
+        import onnxscript.tools.transformers_models.llama as m_llama
+
+        tmodel, inputs, dynamic_shapes_def = m_llama.get_llama_model_from_config(
+            warmup=warmup,
+            repeat=repeat,
+            implementation=implementation,
+            with_mask=with_mask,
+            num_hidden_layers=num_hidden_layers,
+            dynamic_shapes=dynamic_shapes,
+            config=config,
+        )
+
+    elif model == "mistral":
+        import onnxscript.tools.transformers_models.mistral as m_mistral
+
+        tmodel, inputs, dynamic_shapes_def = m_mistral.get_mistral_model_from_config(
+            warmup=warmup,
+            repeat=repeat,
+            implementation=implementation,
+            with_mask=with_mask,
+            num_hidden_layers=num_hidden_layers,
+            dynamic_shapes=dynamic_shapes,
+            config=config,
+        )
+
+    elif model == "phi":
+        import onnxscript.tools.transformers_models.phi as m_phi
+
+        tmodel, inputs, dynamic_shapes_def = m_phi.get_phi_model_from_config(
+            warmup=warmup,
+            repeat=repeat,
+            implementation=implementation,
+            with_mask=with_mask,
+            num_hidden_layers=num_hidden_layers,
+            dynamic_shapes=dynamic_shapes,
+            config=config,
+        )
+
+    elif model == "phi3":
+        import onnxscript.tools.transformers_models.phi3 as m_phi3
+
+        tmodel, inputs, dynamic_shapes_def = m_phi3.get_phi3_model_from_config(
+            warmup=warmup,
+            repeat=repeat,
+            implementation=implementation,
+            with_mask=with_mask,
+            num_hidden_layers=num_hidden_layers,
+            dynamic_shapes=dynamic_shapes,
+            config=config,
+        )
+
+    else:
+        raise ValueError(f"Model {model!r} is unknown.")
+
+    if dtype is not None:
+        dt = getattr(torch, dtype)
+        tmodel = tmodel.to(dt)
+        inputs = [
+            tuple((i if i.dtype in {torch.int64, torch.int32} else i.to(dt)) for i in inp)
+            for inp in inputs
+        ]
+
+    if device == "cuda":
+        tmodel = tmodel.to("cuda")
+        inputs = [tuple(i.to("cuda") for i in inp) for inp in inputs]
+
+    return tmodel, inputs, dynamic_shapes_def

onnxscript/tools/transformers_models/llama.py

@@ -0,0 +1,168 @@
+# -------------------------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+# --------------------------------------------------------------------------
+# pylint: disable=import-outside-toplevel
+from __future__ import annotations
+
+from typing import Any, Sequence
+
+import torch
+
+import onnxscript.tools.transformers_models
+
+
+def get_llama_model(
+    input_dims: Sequence[tuple[int, int]] = ((2, 8), (4, 7), (9, 15)),
+    hidden_size: int = 16,
+    num_hidden_layers: int = 1,
+    vocab_size: int = 1024,
+    intermediate_size: int = 16,
+    max_position_embeddings: int = 1024,
+    num_attention_heads: int = 2,
+    _attn_implementation: str = "eager",  # needed value to remove graph breaks
+    with_mask: bool = True,
+) -> tuple[Any, list[tuple[torch.Tensor, ...]], dict]:
+    """
+    Returns a model.
+    See `LlamaConfig
+    <https://huggingface.co/docs/transformers/main/en/model_doc/llama#transformers.LlamaConfig>`_.
+    The parameters are chosen for a unit test configuration.
+    """
+    from transformers import LlamaConfig
+    from transformers.models.llama.modeling_llama import LlamaModel
+
+    dynamic_shapes = {0: {0: "batch", 1: "length"}}
+    if with_mask:
+        dynamic_shapes.update({1: {0: "batch", 1: "length"}})
+
+    config = LlamaConfig(
+        num_hidden_layers=num_hidden_layers,
+        vocab_size=vocab_size,
+        hidden_size=hidden_size,
+        intermediate_size=intermediate_size,
+        max_position_embeddings=max_position_embeddings,
+        num_attention_heads=num_attention_heads,
+    )
+    if _attn_implementation:
+        config._attn_implementation = _attn_implementation  # pylint: disable=protected-access
+
+    if with_mask:
+
+        class LlamaModelWrapperMask(torch.nn.Module):
+            def __init__(self, config):
+                super().__init__()
+                self.model = LlamaModel(config)
+
+            def forward(self, input_ids, attention_mask):
+                model_output = self.model(
+                    input_ids, attention_mask=attention_mask, use_cache=False
+                )
+                return model_output.to_tuple()
+
+        def generate_example_inputs_mask(batch: int, seq: int, vocab_size: int):
+            input_ids = onnxscript.tools.transformers_models.ids_tensor(
+                [batch, seq], vocab_size
+            )
+            input_mask = torch.tril(torch.ones(batch, seq, dtype=torch.float32))
+            assert input_mask.dtype == torch.float32
+            return input_ids, input_mask
+
+        example_args_collection = []
+        for b, s in input_dims:
+            example_args_collection.append(generate_example_inputs_mask(b, s, vocab_size))
+
+        return LlamaModelWrapperMask(config), example_args_collection, dynamic_shapes
+
+    # no mask
+
+    class LlamaModelWrapper(torch.nn.Module):
+        def __init__(self, config):
+            super().__init__()
+            self.model = LlamaModel(config)
+
+        def forward(self, input_ids):
+            model_output = self.model(input_ids, use_cache=False)
+            return model_output.to_tuple()
+
+    def generate_example_inputs(batch: int, seq: int, vocab_size: int):
+        input_ids = onnxscript.tools.transformers_models.ids_tensor([batch, seq], vocab_size)
+        return (input_ids,)
+
+    example_args_collection = []
+    for b, s in input_dims:
+        example_args_collection.append(generate_example_inputs(b, s, vocab_size))
+
+    return LlamaModelWrapper(config), example_args_collection, dynamic_shapes
+
+
+def get_llama_model_from_config(
+    warmup: int = 5,
+    repeat: int = 10,
+    config: str = "small",
+    num_hidden_layers: int = 1,
+    implementation: str = "eager",
+    dynamic_shapes: bool = False,
+    with_mask: bool = True,
+) -> tuple[Any, list[tuple[torch.Tensor, ...]], dict]:
+    """
+    Returns a model Phi to test or benchmark.
+
+    Args:
+        warmup: Number of inputs to generate.
+        repeat: Number of inputs to generate for repeat.
+        config: small, medium or large
+        num_hidden_layers: Number of hidden layers.
+        implementation: eager or sdpa
+        with_mask: One or two inputs.
+        dynamic_shapes: dynamic shapes or not
+
+    Returns:
+        Model and list of inputs.
+    """
+    if config == "small":
+        conf_dict = dict(
+            input_dims=onnxscript.tools.transformers_models.get_input_dims_for_llm(
+                dynamic_shapes, warmup, repeat
+            ),
+            hidden_size=16,
+            num_hidden_layers=num_hidden_layers,
+            vocab_size=1024,
+            intermediate_size=16,
+            max_position_embeddings=1024,
+            num_attention_heads=2,
+            _attn_implementation=implementation,
+            with_mask=with_mask,
+        )
+    elif config == "medium":
+        conf_dict = dict(
+            input_dims=onnxscript.tools.transformers_models.get_input_dims_for_llm(
+                dynamic_shapes, warmup, repeat
+            ),
+            hidden_size=1024,
+            num_hidden_layers=num_hidden_layers,
+            vocab_size=1024,
+            intermediate_size=1024,
+            max_position_embeddings=1024,
+            num_attention_heads=2,
+            _attn_implementation=implementation,
+            with_mask=with_mask,
+        )
+    elif config in ("large", "default"):
+        conf_dict = dict(
+            input_dims=onnxscript.tools.transformers_models.get_input_dims_for_llm(
+                dynamic_shapes, warmup, repeat
+            ),
+            hidden_size=4096,
+            num_hidden_layers=num_hidden_layers,
+            vocab_size=32000,
+            intermediate_size=11008,
+            max_position_embeddings=2048,
+            num_attention_heads=32,
+            _attn_implementation=implementation,
+            with_mask=with_mask,
+        )
+    else:
+        raise ValueError(f"Unexpected configuration {config!r}.")
+
+    return get_llama_model(**conf_dict)  # type: ignore[arg-type]