nshtrainer 0.14.1__py3-none-any.whl → 0.14.4__py3-none-any.whl

nshtrainer/_experimental/__init__.py

@@ -1,2 +1 @@
-from .flops import MEASURE_FLOPS_AVAILABLE as MEASURE_FLOPS_AVAILABLE
-from .flops import measure_flops as measure_flops
+from lightning.fabric.utilities.throughput import measure_flops as measure_flops

nshtrainer/runner.py

@@ -3,10 +3,9 @@ import functools
 from collections.abc import Callable, Iterable, Mapping, Sequence
 from typing import Generic
 
-from nshrunner import RunInfo
+from nshrunner import RunInfo, Snapshot
 from nshrunner import Runner as _Runner
 from nshrunner._submit import screen
-from nshrunner.snapshot import SnapshotArgType
 from typing_extensions import TypeVar, TypeVarTuple, Unpack, override
 
 from .model.config import BaseConfig
@@ -93,7 +92,7 @@ class Runner(
         options: screen.ScreenJobKwargs = {},
         n_batches: int = 1,
         *,
-        snapshot: SnapshotArgType,
+        snapshot: Snapshot,
         setup_commands: Sequence[str] | None = None,
         env: Mapping[str, str] | None = None,
         transforms: list[

nshtrainer/trainer/signal_connector.py

@@ -1,5 +1,6 @@
 import logging
 import os
+import platform
 import re
 import signal
 import subprocess
@@ -25,6 +26,7 @@ log = logging.getLogger(__name__)
 
 _SIGNUM = int | signal.Signals
 _HANDLER: TypeAlias = Callable[[_SIGNUM, FrameType], Any] | int | signal.Handlers | None
+_IS_WINDOWS = platform.system() == "Windows"
 
 
 def _resolve_requeue_signals():
@@ -57,7 +59,7 @@ class _SignalConnector(_LightningSignalConnector):
         handlers: list[_HANDLER],
         replace_existing: bool = False,
     ):
-        if self._is_on_windows():
+        if _IS_WINDOWS:
             log.info(
                 f"Signal {signum.name} has no handlers or is not supported on Windows."
             )

{nshtrainer-0.14.1.dist-info → nshtrainer-0.14.4.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: nshtrainer
-Version: 0.14.1
+Version: 0.14.4
 Summary: 
 Author: Nima Shoghi
 Author-email: nimashoghi@gmail.com

{nshtrainer-0.14.1.dist-info → nshtrainer-0.14.4.dist-info}/RECORD

@@ -2,10 +2,7 @@ nshtrainer/__init__.py,sha256=39loiLLXbaGiozEsAn8mPHopxaPsek8JsgR9DD2gxtY,583
 nshtrainer/_checkpoint/loader.py,sha256=DSaNR8194kWon4O1svslNsCcN_8vlyLbF0LNCPfUpzI,13789
 nshtrainer/_checkpoint/metadata.py,sha256=onmetLp5eKbA86abq1PTkwAOO7bWj7Pa1EGUjl2TEjQ,5153
 nshtrainer/_checkpoint/saver.py,sha256=DkbCH0YeOJ71m32vAARiQdGBf0hvwwdoAV8LOFGy-0Y,1428
-nshtrainer/_experimental/__init__.py,sha256=2tQIcrWT8U8no_AeBTYnozaTmxN40kuAJdGQ4b-PoWM,120
-nshtrainer/_experimental/flops/__init__.py,sha256=edo9Ez3LlrnxkNRX9W6YBhPkRPKYGLpkpnl5gx7sEX8,1550
-nshtrainer/_experimental/flops/flop_counter.py,sha256=-sL0Fy6poXa__hyzUMdZScjPULp4coQELQpPU6p6dXU,25736
-nshtrainer/_experimental/flops/module_tracker.py,sha256=bUL-IRTd0aF_DwmXkZjHZAA31p4ZEhyqhc26XWKQUUY,4922
+nshtrainer/_experimental/__init__.py,sha256=pEXPyI184UuDHvfh4p9Kg9nQZQZI41e4_HvNd4BK-yg,81
 nshtrainer/callbacks/__init__.py,sha256=4qocBDzQbLLhhbIEfvbA3SQB_Dy9ZJH7keMwPay-ZS8,2359
 nshtrainer/callbacks/_throughput_monitor_callback.py,sha256=aJo_11rc4lo0IYOd-kHmPDtzdC4ctgXyRudkRJqH4m4,23184
 nshtrainer/callbacks/actsave.py,sha256=qbnaKts4_dvjPeAaPtv7Ds12_vEWzaHUfg_--49NB9I,4041
@@ -73,12 +70,12 @@ nshtrainer/nn/module_dict.py,sha256=NOY0B6WDTnktyWH4GthsprMQo0bpehC-hCq9SfD8paE,
 nshtrainer/nn/module_list.py,sha256=fb2u5Rqdjff8Pekyr9hkCPkBorQ-fldzzFAjsgWAm30,1719
 nshtrainer/nn/nonlinearity.py,sha256=owtU4kh4G98psD0axOJWVfBhm-OtJVgFM-TXSHmbNPU,3625
 nshtrainer/optimizer.py,sha256=kuJEA1pvB3y1FcsfhAoOJujVqEZqFHlmYO8GW6JeA1g,1527
-nshtrainer/runner.py,sha256=6qfE5FBONzD79kVHuWYKEvK0J_Qi5dMBbHQhRMmnIhE,3649
+nshtrainer/runner.py,sha256=USAjrExHkN5oVNVunsoPnLxfQrEHSaa54S3RipOe544,3605
 nshtrainer/scripts/find_packages.py,sha256=FbdlfmAefttFSMfaT0A46a-oHLP_ioaQKihwBfBeWeA,1467
 nshtrainer/trainer/__init__.py,sha256=P2rmr8oBVTHk-HJHYPcUwWqDEArMbPR4_rPpATbWK3E,40
 nshtrainer/trainer/_runtime_callback.py,sha256=sd2cUdRJG-UCdQr9ruZvEYpNGNF1t2W2fuxwwVlQD9E,4164
 nshtrainer/trainer/checkpoint_connector.py,sha256=F2tkHogbMAa5U7335sm77sZBkjEDa5v46XbJCH9Mg6c,2167
-nshtrainer/trainer/signal_connector.py,sha256=llwc8pdKAWxREFpjdi14Bpy8rGVMEJsmJx_s2p4gI8E,10689
+nshtrainer/trainer/signal_connector.py,sha256=2EzkVktlasl8PgWAKNLDZRUMY__gRlDy1HdinAU-tfU,10740
 nshtrainer/trainer/trainer.py,sha256=M97phnALfG18VxkMLoDr5AKFf4UaPBdc6S2BghdBtas,17103
 nshtrainer/util/_environment_info.py,sha256=yPtAbgjCY4tkvh5wp9sjNsF0Z45TYwzEAM_N2_b5BbY,23123
 nshtrainer/util/_useful_types.py,sha256=dwZokFkIe7M5i2GR3nQ9A1lhGw06DMAFfH5atyquqSA,8000
@@ -87,6 +84,6 @@ nshtrainer/util/seed.py,sha256=Or2wMPsnQxfnZ2xfBiyMcHFIUt3tGTNeMMyOEanCkqs,280
 nshtrainer/util/slurm.py,sha256=rofIU26z3SdL79SF45tNez6juou1cyDLz07oXEZb9Hg,1566
 nshtrainer/util/typed.py,sha256=NGuDkDzFlc1fAoaXjOFZVbmj0mRFjsQi1E_hPa7Bn5U,128
 nshtrainer/util/typing_utils.py,sha256=8ptjSSLZxlmy4FY6lzzkoGoF5fGNClo8-B_c0XHQaNU,385
-nshtrainer-0.14.1.dist-info/METADATA,sha256=V7kssNBEV-JqILCZbRujHMs4ORomy_hWU4vRK0Je2GU,860
-nshtrainer-0.14.1.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
-nshtrainer-0.14.1.dist-info/RECORD,,
+nshtrainer-0.14.4.dist-info/METADATA,sha256=Xg7ON7o-nR8N-4C8E59k2jsvIaleKPDas5VpdTVyr6s,860
+nshtrainer-0.14.4.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
+nshtrainer-0.14.4.dist-info/RECORD,,

nshtrainer/_experimental/flops/__init__.py

@@ -1,48 +0,0 @@
-from collections.abc import Callable
-
-import torch
-from lightning.fabric.utilities.imports import _TORCH_GREATER_EQUAL_2_1
-
-MEASURE_FLOPS_AVAILABLE = _TORCH_GREATER_EQUAL_2_1
-
-
-def measure_flops(
-    forward_fn: Callable[[], torch.Tensor],
-    loss_fn: Callable[[torch.Tensor], torch.Tensor] | None = None,
-    display: bool = True,
-) -> int:
-    """Utility to compute the total number of FLOPs used by a module during training or during inference.
-
-    It's recommended to create a meta-device model for this:
-
-    Example::
-
-        with torch.device("meta"):
-            model = MyModel()
-            x = torch.randn(2, 32)
-
-        model_fwd = lambda: model(x)
-        fwd_flops = measure_flops(model, model_fwd)
-
-        model_loss = lambda y: y.sum()
-        fwd_and_bwd_flops = measure_flops(model, model_fwd, model_loss)
-
-    Args:
-        model: The model whose FLOPs should be measured.
-        forward_fn: A function that runs ``forward`` on the model and returns the result.
-        loss_fn: A function that computes the loss given the ``forward_fn`` output. If provided, the loss and `backward`
-            FLOPs will be included in the result.
-
-    """
-    if not MEASURE_FLOPS_AVAILABLE:
-        raise ImportError("`measure_flops` requires PyTorch >= 2.1.")
-
-    from .flop_counter import FlopCounterMode
-
-    flop_counter = FlopCounterMode(display=display)
-    with flop_counter:
-        if loss_fn is None:
-            forward_fn()
-        else:
-            loss_fn(forward_fn()).backward()
-    return flop_counter.get_total_flops()

nshtrainer/_experimental/flops/flop_counter.py

@@ -1,787 +0,0 @@
-# type: ignore
-import warnings
-from collections import defaultdict
-from functools import wraps
-from math import prod
-from typing import Any, Dict, Iterator, List, Optional, Tuple, Union
-
-import torch
-from torch._decomp import register_decomposition
-from torch.utils._python_dispatch import TorchDispatchMode
-from torch.utils._pytree import tree_flatten, tree_map, tree_unflatten
-
-from .module_tracker import ModuleTracker
-
-__all__ = ["FlopCounterMode", "register_flop_formula"]
-
-aten = torch.ops.aten
-
-
-def get_shape(i):
-    if isinstance(i, torch.Tensor):
-        return i.shape
-    return i
-
-
-flop_registry: Dict[Any, Any] = {}
-
-
-def shape_wrapper(f):
-    @wraps(f)
-    def nf(*args, out_val=None, **kwargs):
-        args, kwargs, out_shape = tree_map(get_shape, (args, kwargs, out_val))
-        return f(*args, out_shape=out_shape, **kwargs)
-
-    return nf
-
-
-def register_flop_formula(targets, get_raw=False):
-    def register_fun(flop_formula):
-        if not get_raw:
-            flop_formula = shape_wrapper(flop_formula)
-        register_decomposition(targets, registry=flop_registry, unsafe=True)(
-            flop_formula
-        )
-        return flop_formula
-
-    return register_fun
-
-
-@register_flop_formula(aten.mm)
-def mm_flop(a_shape, b_shape, *args, out_shape=None, **kwargs) -> int:
-    """Count flops for matmul."""
-    # Inputs should be a list of length 2.
-    # Inputs contains the shapes of two matrices.
-    m, k = a_shape
-    k2, n = b_shape
-    assert k == k2
-    # NB(chilli): Should be 2 * k - 1 technically for FLOPs.
-    return m * n * 2 * k
-
-
-@register_flop_formula(aten.addmm)
-def addmm_flop(self_shape, a_shape, b_shape, out_shape=None, **kwargs) -> int:
-    """Count flops for addmm."""
-    return mm_flop(a_shape, b_shape)
-
-
-@register_flop_formula(aten.bmm)
-def bmm_flop(a_shape, b_shape, out_shape=None, **kwargs) -> int:
-    """Count flops for the bmm operation."""
-    # Inputs should be a list of length 2.
-    # Inputs contains the shapes of two tensor.
-    b, m, k = a_shape
-    b2, k2, n = b_shape
-    assert b == b2
-    assert k == k2
-    # NB(chilli): Should be 2 * k - 1 technically for FLOPs.
-    flop = b * m * n * 2 * k
-    return flop
-
-
-@register_flop_formula(aten.baddbmm)
-def baddbmm_flop(self_shape, a_shape, b_shape, out_shape=None, **kwargs) -> int:
-    """Count flops for the baddbmm operation."""
-    # Inputs should be a list of length 3.
-    # Inputs contains the shapes of three tensors.
-    return bmm_flop(a_shape, b_shape)
-
-
-def conv_flop_count(
-    x_shape: List[int],
-    w_shape: List[int],
-    out_shape: List[int],
-    transposed: bool = False,
-) -> int:
-    """Count flops for convolution.
-
-    Note only multiplication is
-    counted. Computation for bias are ignored.
-    Flops for a transposed convolution are calculated as
-    flops = (x_shape[2:] * prod(w_shape) * batch_size).
-    Args:
-        x_shape (list(int)): The input shape before convolution.
-        w_shape (list(int)): The filter shape.
-        out_shape (list(int)): The output shape after convolution.
-        transposed (bool): is the convolution transposed
-    Returns:
-        int: the number of flops
-    """
-
-    batch_size = x_shape[0]
-    conv_shape = (x_shape if transposed else out_shape)[2:]
-    c_out, c_in, *filter_size = w_shape
-
-    """
-    General idea here is that for a regular conv, for each point in the output
-    spatial dimension we convolve the filter with something (hence
-    `prod(conv_shape) * prod(filter_size)` ops). Then, this gets multiplied by
-    1. batch_size, 2. the cross product of input and weight channels.
-
-    For the transpose, it's not each point in the *output* spatial dimension but
-    each point in the *input* spatial dimension.
-    """
-    # NB(chilli): I don't think this properly accounts for padding :think:
-    # NB(chilli): Should be 2 * c_in - 1 technically for FLOPs.
-    flop = prod(conv_shape) * prod(filter_size) * batch_size * c_out * c_in * 2
-    return flop
-
-
-@register_flop_formula([aten.convolution, aten._convolution])
-def conv_flop(
-    x_shape,
-    w_shape,
-    _bias,
-    _stride,
-    _padding,
-    _dilation,
-    transposed,
-    *args,
-    out_shape=None,
-    **kwargs,
-) -> int:
-    """Count flops for convolution."""
-    return conv_flop_count(x_shape, w_shape, out_shape, transposed=transposed)
-
-
-@register_flop_formula(aten.convolution_backward)
-def conv_backward_flop(
-    grad_out_shape,
-    x_shape,
-    w_shape,
-    _bias,
-    _stride,
-    _padding,
-    _dilation,
-    transposed,
-    _output_padding,
-    _groups,
-    output_mask,
-    out_shape,
-) -> int:
-    def t(shape):
-        return [shape[1], shape[0]] + list(shape[2:])
-
-    flop_count = 0
-
-    """
-    Let's say we have a regular 1D conv
-    {A, B, C} [inp]
-    {i, j} [weight]
-    => (conv)
-    {Ai + Bj, Bi + Cj} [out]
-
-    And as a reminder, the transposed conv of the above is
-    => {Ai, Aj + Bi, Bj + Ci, Cj} [transposed conv out]
-
-    For the backwards of conv, we now have
-    {D, E} [grad_out]
-    {A, B, C} [inp]
-    {i, j} [weight]
-
-    # grad_inp as conv_transpose(grad_out, weight)
-    Let's first compute grad_inp. To do so, we can simply look at all the
-    multiplications that each element of inp is involved in. For example, A is
-    only involved in the first element of the output (and thus only depends upon
-    D in grad_out), and C is only involved in the last element of the output
-    (and thus only depends upon E in grad_out)
-
-    {Di, Dj + Ei, Ej} [grad_inp]
-
-    Note that this corresponds to the below conv_transpose. This gives us the
-    output_mask[0] branch, which is grad_inp.
-
-    {D, E} [inp (grad_out)]
-    {i, j} [weight]
-    => (conv_transpose)
-    {Di, Dj + Ei, Ej} [out (grad_inp)]
-
-    I leave the fact that grad_inp for a transposed conv is just conv(grad_out,
-    weight) as an exercise for the reader.
-
-    # grad_weight as conv(inp, grad_out)
-    To compute grad_weight, we again look at the terms in the output, which as
-    a reminder is:
-    => {Ai + Bj, Bi + Cj} [out]
-    => {D, E} [grad_out]
-    If we manually compute the gradient for the weights, we see it's
-    {AD + BE, BD + CE} [grad_weight]
-
-    This corresponds to the below conv
-    {A, B, C} [inp]
-    {D, E} [weight (grad_out)]
-    => (conv)
-    {AD + BE, BD + CE} [out (grad_weight)]
-
-    # grad_weight of transposed conv as conv(grad_out, inp)
-    As a reminder, the terms of the output of a transposed conv are:
-    => {Ai, Aj + Bi, Bj + Ci, Cj} [transposed conv out]
-    => {D, E, F, G} [grad_out]
-
-    Manually computing the gradient for the weights, we see it's
-    {AD + BE + CF, AE + BF + CG} [grad_weight]
-
-    This corresponds to the below conv
-    {D, E, F, G} [inp (grad_out)]
-    {A, B, C} [weight (inp)]
-    => (conv)
-    {AD + BE + CF, AE + BF + CG} [out (grad_weight)]
-
-    For the full backwards formula, there are also some details involving
-    transpose of the batch/channel dimensions and groups, but I skip those for
-    the sake of brevity (and they're pretty similar to matmul backwards)
-
-    Check [conv backwards decomposition as conv forwards]
-    """
-    # grad_inp as conv_transpose(grad_out, weight)
-    if output_mask[0]:
-        grad_input_shape = get_shape(out_shape[0])
-        flop_count += conv_flop_count(
-            grad_out_shape, w_shape, grad_input_shape, not transposed
-        )
-
-    if output_mask[1]:
-        grad_weight_shape = get_shape(out_shape[1])
-        if transposed:
-            # grad_weight of transposed conv as conv(grad_out, inp)
-            flop_count += conv_flop_count(
-                t(grad_out_shape), t(x_shape), t(grad_weight_shape), transposed=False
-            )
-        else:
-            # grad_weight as conv(inp, grad_out)
-            flop_count += conv_flop_count(
-                t(x_shape), t(grad_out_shape), t(grad_weight_shape), transposed=False
-            )
-
-    return flop_count
-
-
-def sdpa_flop_count(query_shape, key_shape, value_shape):
-    """
-    Count flops for self-attention.
-
-    NB: We can assume that value_shape == key_shape
-    """
-    b, h, s_q, d_q = query_shape
-    _b2, _h2, s_k, _d2 = key_shape
-    _b3, _h3, _s3, d_v = value_shape
-    assert (
-        b == _b2 == _b3 and h == _h2 == _h3 and d_q == _d2 and s_k == _s3 and d_q == _d2
-    )
-    total_flops = 0
-    # q: [b, h, s_q, d_q] @ k: [b, h, d_q, s_k] -> scores: [b, h, s_q, s_k]
-    total_flops += bmm_flop((b * h, s_q, d_q), (b * h, d_q, s_k))
-    # scores: [b, h, s_q, s_k] @ v: [b, h, s_k, d_v] -> out: [b, h, s_q, d_v]
-    total_flops += bmm_flop((b * h, s_q, s_k), (b * h, s_k, d_v))
-    return total_flops
-
-
-@register_flop_formula(
-    [
-        aten._scaled_dot_product_efficient_attention,
-        aten._scaled_dot_product_flash_attention,
-    ]
-)
-def sdpa_flop(
-    query_shape, key_shape, value_shape, *args, out_shape=None, **kwargs
-) -> int:
-    """Count flops for self-attention."""
-    # NB: We aren't accounting for causal attention here
-    return sdpa_flop_count(query_shape, key_shape, value_shape)
-
-
-def _unpack_flash_attention_nested_shapes(
-    *,
-    query,
-    key,
-    value,
-    grad_out=None,
-    cum_seq_q,
-    cum_seq_k,
-    max_q,
-    max_k,
-) -> Iterator[
-    Tuple[Tuple[int, ...], Tuple[int, ...], Tuple[int, ...], Optional[Tuple[int, ...]]]
-]:
-    """
-    Given inputs to a flash_attention_(forward|backward) kernel, this will handle behavior for
-    NestedTensor inputs by effectively unbinding the NestedTensor and yielding the shapes for
-    each batch element.
-
-    In the case that this isn't a NestedTensor kernel, then it just yields the original shapes.
-    """
-    if cum_seq_q is not None:
-        # This means we should be dealing with a Nested Jagged Tensor query.
-        # The inputs will have shape                  (sum(sequence len), heads, dimension)
-        # In comparison, non-Nested inputs have shape (batch, heads, sequence len, dimension)
-        # To deal with this, we convert to a shape of (batch, heads, max_seq_len, dimension)
-        # So the flops calculation in this case is an overestimate of the actual flops.
-        assert len(key.shape) == 3
-        assert len(value.shape) == 3
-        assert grad_out is None or grad_out.shape == query.shape
-        _, h_q, d_q = query.shape
-        _, h_k, d_k = key.shape
-        _, h_v, d_v = value.shape
-        assert cum_seq_q is not None
-        assert cum_seq_k is not None
-        assert cum_seq_q.shape == cum_seq_k.shape
-        seq_q_lengths = (cum_seq_q[1:] - cum_seq_q[:-1]).tolist()
-        seq_k_lengths = (cum_seq_k[1:] - cum_seq_k[:-1]).tolist()
-        for seq_q_len, seq_k_len in zip(seq_q_lengths, seq_k_lengths):
-            new_query_shape = (1, h_q, seq_q_len, d_q)
-            new_key_shape = (1, h_k, seq_k_len, d_k)
-            new_value_shape = (1, h_v, seq_k_len, d_v)
-            new_grad_out_shape = new_query_shape if grad_out is not None else None
-            yield new_query_shape, new_key_shape, new_value_shape, new_grad_out_shape
-        return
-
-    yield (
-        query.shape,
-        key.shape,
-        value.shape,
-        grad_out.shape if grad_out is not None else None,
-    )
-
-
-def _unpack_efficient_attention_nested_shapes(
-    *,
-    query,
-    key,
-    value,
-    grad_out=None,
-    cu_seqlens_q,
-    cu_seqlens_k,
-    max_seqlen_q,
-    max_seqlen_k,
-) -> Iterator[
-    Tuple[Tuple[int, ...], Tuple[int, ...], Tuple[int, ...], Optional[Tuple[int, ...]]]
-]:
-    """
-    Given inputs to a efficient_attention_(forward|backward) kernel, this will handle behavior for
-    NestedTensor inputs by effectively unbinding the NestedTensor and yielding the shapes for
-    each batch element.
-
-    In the case that this isn't a NestedTensor kernel, then it just yields the original shapes.
-    """
-    if cu_seqlens_q is not None:
-        # Unlike flash_attention_forward, we get a 4D tensor instead of a 3D tensor for efficient attention.
-        #
-        # This means we should be dealing with a Nested Jagged Tensor query.
-        # The inputs will have shape                  (sum(sequence len), heads, dimension)
-        # In comparison, non-Nested inputs have shape (batch, heads, sequence len, dimension)
-        # To deal with this, we convert to a shape of (batch, heads, max_seq_len, dimension)
-        # So the flops calculation in this case is an overestimate of the actual flops.
-        assert len(key.shape) == 4
-        assert len(value.shape) == 4
-        assert grad_out is None or grad_out.shape == query.shape
-        _, _, h_q, d_q = query.shape
-        _, _, h_k, d_k = key.shape
-        _, _, h_v, d_v = value.shape
-        assert cu_seqlens_q is not None
-        assert cu_seqlens_k is not None
-        assert cu_seqlens_q.shape == cu_seqlens_k.shape
-        seqlens_q = (cu_seqlens_q[1:] - cu_seqlens_q[:-1]).tolist()
-        seqlens_k = (cu_seqlens_k[1:] - cu_seqlens_k[:-1]).tolist()
-        for len_q, len_k in zip(seqlens_q, seqlens_k):
-            new_query_shape = (1, h_q, len_q, d_q)
-            new_key_shape = (1, h_k, len_k, d_k)
-            new_value_shape = (1, h_v, len_k, d_v)
-            new_grad_out_shape = new_query_shape if grad_out is not None else None
-            yield new_query_shape, new_key_shape, new_value_shape, new_grad_out_shape
-        return
-
-    yield (
-        query.shape,
-        key.shape,
-        value.shape,
-        grad_out.shape if grad_out is not None else None,
-    )
-
-
-@register_flop_formula(aten._flash_attention_forward, get_raw=True)
-def _flash_attention_forward_flop(
-    query,
-    key,
-    value,
-    cum_seq_q,
-    cum_seq_k,
-    max_q,
-    max_k,
-    *args,
-    out_shape=None,
-    **kwargs,
-) -> int:
-    """Count flops for self-attention."""
-    # NB: We aren't accounting for causal attention here
-    # in case this is a nested tensor, we unpack the individual batch elements
-    # and then sum the flops per batch element
-    sizes = _unpack_flash_attention_nested_shapes(
-        query=query,
-        key=key,
-        value=value,
-        cum_seq_q=cum_seq_q,
-        cum_seq_k=cum_seq_k,
-        max_q=max_q,
-        max_k=max_k,
-    )
-    return sum(
-        sdpa_flop_count(query_shape, key_shape, value_shape)
-        for query_shape, key_shape, value_shape, _ in sizes
-    )
-
-
-@register_flop_formula(aten._efficient_attention_forward, get_raw=True)
-def _efficient_attention_forward_flop(
-    query,
-    key,
-    value,
-    bias,
-    cu_seqlens_q,
-    cu_seqlens_k,
-    max_seqlen_q,
-    max_seqlen_k,
-    *args,
-    **kwargs,
-) -> int:
-    """Count flops for self-attention."""
-    # NB: We aren't accounting for causal attention here
-    # in case this is a nested tensor, we unpack the individual batch elements
-    # and then sum the flops per batch element
-    sizes = _unpack_efficient_attention_nested_shapes(
-        query=query,
-        key=key,
-        value=value,
-        cu_seqlens_q=cu_seqlens_q,
-        cu_seqlens_k=cu_seqlens_k,
-        max_seqlen_q=max_seqlen_q,
-        max_seqlen_k=max_seqlen_k,
-    )
-    return sum(
-        sdpa_flop_count(query_shape, key_shape, value_shape)
-        for query_shape, key_shape, value_shape, _ in sizes
-    )
-
-
-def sdpa_backward_flop_count(grad_out_shape, query_shape, key_shape, value_shape):
-    total_flops = 0
-    b, h, s_q, d_q = query_shape
-    _b2, _h2, s_k, _d2 = key_shape
-    _b3, _h3, _s3, d_v = value_shape
-    _b4, _h4, _s4, _d4 = grad_out_shape
-    assert b == _b2 == _b3 == _b4 and h == _h2 == _h3 == _h4 and d_q == _d2
-    assert d_v == _d4 and s_k == _s3 and s_q == _s4
-    total_flops = 0
-    # Step 1: We recompute the scores matrix.
-    # q: [b, h, s_q, d_q] @ k: [b, h, d_q, s_k] -> scores: [b, h, s_q, s_k]
-    total_flops += bmm_flop((b * h, s_q, d_q), (b * h, d_q, s_k))
-
-    # Step 2: We propagate the gradients through the score @ v operation.
-    # gradOut: [b, h, s_q, d_v] @ v: [b, h, d_v, s_k] -> gradScores: [b, h, s_q, s_k]
-    total_flops += bmm_flop((b * h, s_q, d_v), (b * h, d_v, s_k))
-    # scores: [b, h, s_k, s_q] @ gradOut: [b, h, s_q, d_v] -> gradV: [b, h, s_k, d_v]
-    total_flops += bmm_flop((b * h, s_k, s_q), (b * h, s_q, d_v))
-
-    # Step 3: We propagate th gradients through the k @ v operation
-    # gradScores: [b, h, s_q, s_k] @ k: [b, h, s_k, d_q] -> gradQ: [b, h, s_q, d_q]
-    total_flops += bmm_flop((b * h, s_q, s_k), (b * h, s_k, d_q))
-    # q: [b, h, d_q, s_q] @ gradScores: [b, h, s_q, s_k] -> gradK: [b, h, d_q, s_k]
-    total_flops += bmm_flop((b * h, d_q, s_q), (b * h, s_q, s_k))
-    return total_flops
-
-
-@register_flop_formula(
-    [
-        aten._scaled_dot_product_efficient_attention_backward,
-        aten._scaled_dot_product_flash_attention_backward,
-    ]
-)
-def sdpa_backward_flop(
-    grad_out_shape, query_shape, key_shape, value_shape, *args, out_shape=None, **kwargs
-) -> int:
-    """Count flops for self-attention backward."""
-    return sdpa_backward_flop_count(grad_out_shape, query_shape, key_shape, value_shape)
-
-
-@register_flop_formula(aten._flash_attention_backward, get_raw=True)
-def _flash_attention_backward_flop(
-    grad_out,
-    query,
-    key,
-    value,
-    out,  # named _out_shape to avoid kwarg collision with out_shape created in wrapper
-    logsumexp,
-    cum_seq_q,
-    cum_seq_k,
-    max_q,
-    max_k,
-    *args,
-    **kwargs,
-) -> int:
-    # in case this is a nested tensor, we unpack the individual batch elements
-    # and then sum the flops per batch element
-    shapes = _unpack_flash_attention_nested_shapes(
-        query=query,
-        key=key,
-        value=value,
-        grad_out=grad_out,
-        cum_seq_q=cum_seq_q,
-        cum_seq_k=cum_seq_k,
-        max_q=max_q,
-        max_k=max_k,
-    )
-    return sum(
-        sdpa_backward_flop_count(grad_out_shape, query_shape, key_shape, value_shape)
-        for query_shape, key_shape, value_shape, grad_out_shape in shapes
-    )
-
-
-@register_flop_formula(aten._efficient_attention_backward, get_raw=True)
-def _efficient_attention_backward_flop(
-    grad_out,
-    query,
-    key,
-    value,
-    bias,
-    out,  # named _out to avoid kwarg collision with out created in wrapper
-    cu_seqlens_q,
-    cu_seqlens_k,
-    max_seqlen_q,
-    max_seqlen_k,
-    *args,
-    **kwargs,
-) -> int:
-    # in case this is a nested tensor, we unpack the individual batch elements
-    # and then sum the flops per batch element
-    shapes = _unpack_efficient_attention_nested_shapes(
-        query=query,
-        key=key,
-        value=value,
-        grad_out=grad_out,
-        cu_seqlens_q=cu_seqlens_q,
-        cu_seqlens_k=cu_seqlens_k,
-        max_seqlen_q=max_seqlen_q,
-        max_seqlen_k=max_seqlen_k,
-    )
-    return sum(
-        sdpa_backward_flop_count(grad_out_shape, query_shape, key_shape, value_shape)
-        for query_shape, key_shape, value_shape, grad_out_shape in shapes
-    )
-
-
-flop_registry = {
-    aten.mm: mm_flop,
-    aten.addmm: addmm_flop,
-    aten.bmm: bmm_flop,
-    aten.baddbmm: baddbmm_flop,
-    aten.convolution: conv_flop,
-    aten._convolution: conv_flop,
-    aten.convolution_backward: conv_backward_flop,
-    aten._scaled_dot_product_efficient_attention: sdpa_flop,
-    aten._scaled_dot_product_flash_attention: sdpa_flop,
-    aten._scaled_dot_product_efficient_attention_backward: sdpa_backward_flop,
-    aten._scaled_dot_product_flash_attention_backward: sdpa_backward_flop,
-    aten._flash_attention_forward: _flash_attention_forward_flop,
-    aten._efficient_attention_forward: _efficient_attention_forward_flop,
-    aten._flash_attention_backward: _flash_attention_backward_flop,
-    aten._efficient_attention_backward: _efficient_attention_backward_flop,
-}
-
-
-def normalize_tuple(x):
-    if not isinstance(x, tuple):
-        return (x,)
-    return x
-
-
-# Define the suffixes for different orders of magnitude
-suffixes = ["", "K", "M", "B", "T"]
-
-
-# Thanks BingChat!
-def get_suffix_str(number):
-    # Find the index of the appropriate suffix based on the number of digits
-    # with some additional overflow.
-    # i.e. 1.01B should be displayed as 1001M, not 1.001B
-    index = max(0, min(len(suffixes) - 1, (len(str(number)) - 2) // 3))
-    return suffixes[index]
-
-
-def convert_num_with_suffix(number, suffix):
-    index = suffixes.index(suffix)
-    # Divide the number by 1000^index and format it to two decimal places
-    value = f"{number / 1000 ** index:.3f}"
-    # Return the value and the suffix as a string
-    return value + suffixes[index]
-
-
-def convert_to_percent_str(num, denom):
-    if denom == 0:
-        return "0%"
-    return f"{num / denom:.2%}"
-
-
-def _pytreeify_preserve_structure(f):
-    @wraps(f)
-    def nf(args):
-        flat_args, spec = tree_flatten(args)
-        out = f(*flat_args)
-        return tree_unflatten(out, spec)
-
-    return nf
-
-
-class FlopCounterMode(TorchDispatchMode):
-    """
-    ``FlopCounterMode`` is a context manager that counts the number of flops within its context.
-
-    It does this using a ``TorchDispatchMode``.
-
-    It also supports hierarchical output by passing a module (or list of
-    modules) to FlopCounterMode on construction. If you do not need hierarchical
-    output, you do not need to use it with a module.
-
-    Example usage
-
-    .. code-block:: python
-
-        mod = ...
-        with FlopCounterMode(mod) as flop_counter:
-            mod.sum().backward()
-
-    """
-
-    def __init__(
-        self,
-        mods: Optional[Union[torch.nn.Module, List[torch.nn.Module]]] = None,
-        depth: int = 2,
-        display: bool = True,
-        custom_mapping: Optional[Dict[Any, Any]] = None,
-    ):
-        self.flop_counts: Dict[str, Dict[Any, int]] = defaultdict(
-            lambda: defaultdict(int)
-        )
-        self.depth = depth
-        self.display = display
-        if custom_mapping is None:
-            custom_mapping = {}
-        if mods is not None:
-            warnings.warn(
-                "mods argument is not needed anymore, you can stop passing it",
-                stacklevel=2,
-            )
-        self.flop_registry = {
-            **flop_registry,
-            **{
-                k: v if getattr(v, "_get_raw", False) else shape_wrapper(v)
-                for k, v in custom_mapping.items()
-            },
-        }
-        self.mod_tracker = ModuleTracker()
-
-    def get_total_flops(self) -> int:
-        return sum(self.flop_counts["Global"].values())
-
-    def get_flop_counts(self) -> Dict[str, Dict[Any, int]]:
-        """Return the flop counts as a dictionary of dictionaries.
-
-        The outer
-        dictionary is keyed by module name, and the inner dictionary is keyed by
-        operation name.
-
-        Returns:
-            Dict[str, Dict[Any, int]]: The flop counts as a dictionary.
-        """
-        return {k: dict(v) for k, v in self.flop_counts.items()}
-
-    def get_table(self, depth=None):
-        if depth is None:
-            depth = self.depth
-        if depth is None:
-            depth = 999999
-
-        import tabulate
-
-        tabulate.PRESERVE_WHITESPACE = True
-        header = ["Module", "FLOP", "% Total"]
-        values = []
-        global_flops = self.get_total_flops()
-        global_suffix = get_suffix_str(global_flops)
-        is_global_subsumed = False
-
-        def process_mod(mod_name, depth):
-            nonlocal is_global_subsumed
-
-            total_flops = sum(self.flop_counts[mod_name].values())
-
-            is_global_subsumed |= total_flops >= global_flops
-
-            padding = " " * depth
-            values = []
-            values.append(
-                [
-                    padding + mod_name,
-                    convert_num_with_suffix(total_flops, global_suffix),
-                    convert_to_percent_str(total_flops, global_flops),
-                ]
-            )
-            for k, v in self.flop_counts[mod_name].items():
-                values.append(
-                    [
-                        padding + " - " + str(k),
-                        convert_num_with_suffix(v, global_suffix),
-                        convert_to_percent_str(v, global_flops),
-                    ]
-                )
-            return values
-
-        for mod in sorted(self.flop_counts.keys()):
-            if mod == "Global":
-                continue
-            mod_depth = mod.count(".") + 1
-            if mod_depth > depth:
-                continue
-
-            cur_values = process_mod(mod, mod_depth - 1)
-            values.extend(cur_values)
-
-        # We do a bit of messing around here to only output the "Global" value
-        # if there are any FLOPs in there that aren't already fully contained by
-        # a module.
-        if "Global" in self.flop_counts and not is_global_subsumed:
-            for idx, value in enumerate(values):
-                values[idx][0] = " " + values[idx][0]
-
-            values = process_mod("Global", 0) + values
-
-        if len(values) == 0:
-            values = [["Global", "0", "0%"]]
-
-        return tabulate.tabulate(
-            values, headers=header, colalign=("left", "right", "right")
-        )
-
-    def __enter__(self):
-        self.flop_counts.clear()
-        self.mod_tracker.__enter__()
-        super().__enter__()
-        return self
-
-    def __exit__(self, *args):
-        super().__exit__(*args)
-        self.mod_tracker.__exit__()
-        if self.display:
-            print(self.get_table(self.depth))
-
-    def __torch_dispatch__(self, func, types, args=(), kwargs=None):
-        kwargs = kwargs if kwargs else {}
-        out = func(*args, **kwargs)
-        return self._count_flops(func._overloadpacket, out, args, kwargs)
-
-    def _count_flops(self, func_packet, out, args, kwargs):
-        if func_packet in self.flop_registry:
-            flop_count_func = self.flop_registry[func_packet]
-            flop_count = flop_count_func(*args, **kwargs, out_val=out)  # type: ignore[operator]
-            for par in set(self.mod_tracker.parents):
-                self.flop_counts[par][func_packet] += flop_count
-
-        return out

nshtrainer/_experimental/flops/module_tracker.py

@@ -1,140 +0,0 @@
-import weakref
-from typing import Set
-
-import torch
-from torch.autograd.graph import register_multi_grad_hook
-from torch.nn.modules.module import (
-    register_module_forward_hook,
-    register_module_forward_pre_hook,
-)
-from torch.utils._pytree import tree_flatten
-
-__all__ = ["ModuleTracker"]
-
-
-class ModuleTracker:
-    """
-    ``ModuleTracker`` is a context manager that tracks the nn.Module hierarchy during execution
-    so that other system can query which Module is currently being executed (or its backward is being
-    executed).
-
-    You can access the ``parents`` attribute on this context manager to get the set of all the
-    Modules currently being executed via their fqn (fully qualified name, also used as the key within
-    the state_dict).
-    You can access the ``is_bw`` attribute to know if you are currently running in backward or not.
-
-    Note that ``parents`` is never empty and always contains the "Global" key. The ``is_bw`` flag
-    will remain ``True`` after the forward until another Module is executed. If you need it to be
-    more accurate, please submit an issue requesting this. Adding a map from fqn to the module instance
-    is possible but not done yet, please submit an issue requesting this if you need it.
-
-    Example usage
-
-    .. code-block:: python
-
-        mod = torch.nn.Linear(2, 2)
-
-        with ModuleTracker() as tracker:
-            # Access anything during the forward pass
-            def my_linear(m1, m2, bias):
-                print(f"Current modules: {tracker.parents}")
-                return torch.mm(m1, m2.t()) + bias
-            torch.nn.functional.linear = my_linear
-
-            mod(torch.rand(2, 2))
-
-    """
-
-    parents: Set[str]
-    """
-    A Set containing the fqn for each module currently running their forward
-    """
-
-    def __init__(self):
-        self.parents = {"Global"}
-        self._known_modules: weakref.WeakKeyDictionary = weakref.WeakKeyDictionary()
-        self._seen_modules: weakref.WeakSet = weakref.WeakSet()
-        self._has_callback = False
-
-    def _maybe_set_engine_callback(self):
-        # This assumes no concurrent calls to backward
-        if self._has_callback:
-            return
-
-        def callback():
-            self.parents = {"Global"}
-            self._has_callback = False
-
-        torch.autograd.Variable._execution_engine.queue_callback(callback)
-        self._has_callback = True
-
-    @property
-    def is_bw(self):
-        """
-        A boolean marking if this is currently running during the backward pass or not
-        """
-        return torch._C._current_graph_task_id() != -1
-
-    def _get_mod_name(self, mod):
-        if mod not in self._known_modules:
-            self._known_modules[mod] = type(mod).__name__
-        mod_name = self._known_modules[mod]
-        if mod not in self._seen_modules:
-            for name, submod in mod.named_children():
-                self._known_modules[submod] = f"{mod_name}.{name}"
-                self._get_mod_name(submod)
-            self._seen_modules.add(mod)
-        return mod_name
-
-    def _get_append_fn(self, name, is_bw):
-        def fn(*args):
-            if is_bw:
-                self._maybe_set_engine_callback()
-            if name in self.parents:
-                print(
-                    "The module hierarchy tracking seems to be messed up."
-                    "Please file a bug to PyTorch."
-                )
-            self.parents.add(name)
-
-        return fn
-
-    def _get_pop_fn(self, name, is_bw):
-        def fn(*args):
-            if name in self.parents:
-                self.parents.remove(name)
-            elif not is_bw:
-                # Due to some input/output not requiring gradients, we cannot enforce
-                # proper nesting in backward
-                raise RuntimeError(
-                    "The Module hierarchy tracking is wrong. Report a bug to PyTorch"
-                )
-
-        return fn
-
-    def _fw_pre_hook(self, mod, input):
-        name = self._get_mod_name(mod)
-        self._get_append_fn(name, False)()
-
-        args, _ = tree_flatten(input)
-        tensors = [a for a in args if isinstance(a, torch.Tensor) and a.requires_grad]
-        if tensors:
-            register_multi_grad_hook(tensors, self._get_pop_fn(name, True))
-
-    def _fw_post_hook(self, mod, input, output):
-        name = self._get_mod_name(mod)
-        self._get_pop_fn(name, False)()
-
-        args, _ = tree_flatten(output)
-        tensors = [a for a in args if isinstance(a, torch.Tensor) and a.requires_grad]
-        if tensors:
-            register_multi_grad_hook(tensors, self._get_append_fn(name, True))
-
-    def __enter__(self):
-        self._fw_pre_handle = register_module_forward_pre_hook(self._fw_pre_hook)
-        self._fw_post_handle = register_module_forward_hook(self._fw_post_hook)
-        return self
-
-    def __exit__(self, *args):
-        self._fw_pre_handle.remove()
-        self._fw_post_handle.remove()