---
title: Thunder
description: API reference for all public classes and functions.
canonical: 'https://lightning.ai/docs/thunder/stable/api-reference/thunder'
---
# Thunder {#thunder}

## Compiling functions and modules {#compiling-functions-and-modules}

- [`jit`](/docs/thunder/0.2.6/api-reference/generated/thunder.jit.md) — Just-in-time compile a callable (function or model).

## Querying information on compiled functions and modules {#querying-information-on-compiled-functions-and-modules}

- [`DebugOptions`](/docs/thunder/0.2.6/api-reference/generated/thunder.DebugOptions.md) — options can be dynamically registered, currently registered ones are below
- [`compile_data`](/docs/thunder/0.2.6/api-reference/generated/thunder.compile_data.md) — Obtains the compilation data from a JITed function.
- [`compile_stats`](/docs/thunder/0.2.6/api-reference/generated/thunder.compile_stats.md) — Obtains the compilation statistics from a JITed function.
- [`last_traces`](/docs/thunder/0.2.6/api-reference/generated/thunder.last_traces.md) — Obtains the list of computation traces that have been produced for the last run of the function. This is a list
of traces mirroring the progression of transformations being applied to the trace (at in
- [`last_backward_traces`](/docs/thunder/0.2.6/api-reference/generated/thunder.last_backward_traces.md) — Obtains the list of backward traces that have been produced for the last run of the function and the selected prologue.
- [`last_prologue_traces`](/docs/thunder/0.2.6/api-reference/generated/thunder.last_prologue_traces.md) — Obtains the list of prologue traces that have been produced for the last run of the function and the selected prologue.
- [`cache_option`](/docs/thunder/0.2.6/api-reference/generated/thunder.cache_option.md) — Returns the cache options set when JITting the function.
- [`cache_hits`](/docs/thunder/0.2.6/api-reference/generated/thunder.cache_hits.md) — Returns the number of cache hits we found when running the function.
- [`cache_misses`](/docs/thunder/0.2.6/api-reference/generated/thunder.cache_misses.md) — Returns the number of cache misses we found when running the function.
- [`list_transforms`](/docs/thunder/0.2.6/api-reference/generated/thunder.list_transforms.md) — Returns the list of (explicit) transforms applied to the JITed function.
- [`last_interpreted_instructions`](/docs/thunder/0.2.6/api-reference/generated/thunder.last_interpreted_instructions.md) — Returns the list of instructions the interpreter encountered while tracing through the
user program (on the last cache miss).
- [`last_interpreter_log`](/docs/thunder/0.2.6/api-reference/generated/thunder.last_interpreter_log.md) — Returns the list of instructions and other information the interpreter encountered while tracing through the
user program (on the last cache miss).
- [`last_compile_options`](/docs/thunder/0.2.6/api-reference/generated/thunder.last_compile_options.md) — Prints how compiled options were used (or not)

## JITed Model wrapper {#jited-model-wrapper}

### `class thunder.ThunderModule(model, compiled_model_call)` {#thunder.ThunderModule}

[\[source\]](https://github.com/Lightning-AI/lightning-thunder/blob/0.2.6/thunder/core/module.py#L30-L387)

Bases: [`Module`](https://docs.pytorch.org/docs/stable/generated/torch.nn.Module.html#torch.nn.Module "(in PyTorch v2.15)")

A wrapper nn.Module subclass.

This wrapper is returned by `thunder.jit`, you would typically not instantiate it manually.

#### `get_buffer(name)` {#thunder.ThunderModule.get_buffer}

[\[source\]](https://github.com/Lightning-AI/lightning-thunder/blob/0.2.6/thunder/core/module.py#L58-L62)

Return the buffer given by `target` if it exists, otherwise throw an error.

See the docstring for `get_submodule` for a more detailed explanation of this method’s functionality as well as how to correctly specify `target`.

Parameters:

**target** – The fully-qualified string name of the buffer to look for. (See `get_submodule` for how to specify a fully-qualified string.)

Returns:

The buffer referenced by `target`

Return type:

[torch.Tensor](https://docs.pytorch.org/docs/stable/tensors.html#torch.Tensor "(in PyTorch v2.15)")

Raises:

[**AttributeError**](https://docs.python.org/3/builtins/exceptions.html#AttributeError "(in Python v3.15)") – If the target string references an invalid path or resolves to something that is not a buffer

#### `get_parameter(name)` {#thunder.ThunderModule.get_parameter}

[\[source\]](https://github.com/Lightning-AI/lightning-thunder/blob/0.2.6/thunder/core/module.py#L67-L71)

Return the parameter given by `target` if it exists, otherwise throw an error.

See the docstring for `get_submodule` for a more detailed explanation of this method’s functionality as well as how to correctly specify `target`.

Parameters:

**target** – The fully-qualified string name of the Parameter to look for. (See `get_submodule` for how to specify a fully-qualified string.)

Returns:

The Parameter referenced by `target`

Return type:

torch.nn.Parameter

Raises:

[**AttributeError**](https://docs.python.org/3/builtins/exceptions.html#AttributeError "(in Python v3.15)") – If the target string references an invalid path or resolves to something that is not an `nn.Parameter`

#### `get_submodule(name)` {#thunder.ThunderModule.get_submodule}

[\[source\]](https://github.com/Lightning-AI/lightning-thunder/blob/0.2.6/thunder/core/module.py#L73-L77)

Return the submodule given by `target` if it exists, otherwise throw an error.

For example, let’s say you have an `nn.Module` `A` that looks like this:

```text
A(
    (net_b): Module(
        (net_c): Module(
            (conv): Conv2d(16, 33, kernel_size=(3, 3), stride=(2, 2))
        )
        (linear): Linear(in_features=100, out_features=200, bias=True)
    )
)
```

(The diagram shows an `nn.Module` `A`. `A` which has a nested submodule `net_b`, which itself has two submodules `net_c` and `linear`. `net_c` then has a submodule `conv`.)

To check whether or not we have the `linear` submodule, we would call `get_submodule("net_b.linear")`. To check whether we have the `conv` submodule, we would call `get_submodule("net_b.net_c.conv")`.

The runtime of `get_submodule` is bounded by the degree of module nesting in `target`. A query against `named_modules` achieves the same result, but it is O(N) in the number of transitive modules. So, for a simple check to see if some submodule exists, `get_submodule` should always be used.

Parameters:

**target** – The fully-qualified string name of the submodule to look for. (See above example for how to specify a fully-qualified string.)

Returns:

The submodule referenced by `target`

Return type:

[torch.nn.Module](https://docs.pytorch.org/docs/stable/generated/torch.nn.Module.html#torch.nn.Module "(in PyTorch v2.15)")

Raises:

[**AttributeError**](https://docs.python.org/3/builtins/exceptions.html#AttributeError "(in Python v3.15)") – If at any point along the path resulting from the target string the (sub)path resolves to a non-existent attribute name or an object that is not an instance of `nn.Module`.

#### `load_state_dict(state_dict, strict=True, assign=False)` {#thunder.ThunderModule.load_state_dict}

[\[source\]](https://github.com/Lightning-AI/lightning-thunder/blob/0.2.6/thunder/core/module.py#L267-L338)

Loads the state dict to a transformed module.

Parameters:

*   **state\_dict** ([`Mapping`](https://docs.python.org/3/library/collections.abc.html#collections.abc.Mapping "(in Python v3.15)")\[[`str`](https://docs.python.org/3/builtins/stdtypes.html#str "(in Python v3.15)"), [`Any`](https://docs.python.org/3/library/typing.html#typing.Any "(in Python v3.15)")\]) – the state dict to load.
    
*   **strict** ([`bool`](https://docs.python.org/3/builtins/functions.html#bool "(in Python v3.15)")) – error on missing / unused state dict members
    
*   **assign** ([`bool`](https://docs.python.org/3/builtins/functions.html#bool "(in Python v3.15)")) – assign the state dict tensors instead of copying the data
    
*   **state\_dict** ([_Mapping_](https://docs.python.org/3/library/collections.abc.html#collections.abc.Mapping "(in Python v3.15)")_\[_[_str_](https://docs.python.org/3/builtins/stdtypes.html#str "(in Python v3.15)")_,_ [_Any_](https://docs.python.org/3/library/typing.html#typing.Any "(in Python v3.15)")_\]_) –
    
*   **strict** ([_bool_](https://docs.python.org/3/builtins/functions.html#bool "(in Python v3.15)")) –
    
*   **assign** ([_bool_](https://docs.python.org/3/builtins/functions.html#bool "(in Python v3.15)")) –
    

This is much simpler than PyTorch’s load\_state\_dict. (Regarding hooks, customization etc.)

#### `named_buffers(prefix='', recurse=True, remove_duplicate=True, *, persistent=None)` {#thunder.ThunderModule.named_buffers}

[\[source\]](https://github.com/Lightning-AI/lightning-thunder/blob/0.2.6/thunder/core/module.py#L111-L123)

Return an iterator over module buffers, yielding both the name of the buffer as well as the buffer itself.

Parameters:

*   **prefix** ([_str_](https://docs.python.org/3/builtins/stdtypes.html#str "(in Python v3.15)")) – prefix to prepend to all buffer names.
    
*   **recurse** ([_bool_](https://docs.python.org/3/builtins/functions.html#bool "(in Python v3.15)")_,_ _optional_) – if True, then yields buffers of this module and all submodules. Otherwise, yields only buffers that are direct members of this module. Defaults to True.
    
*   **remove\_duplicate** ([_bool_](https://docs.python.org/3/builtins/functions.html#bool "(in Python v3.15)")_,_ _optional_) – whether to remove the duplicated buffers in the result. Defaults to True.
    

Yields:

_(str, torch.Tensor)_ – Tuple containing the name and buffer

Example:

```python
>>> # xdoctest: +SKIP("undefined vars")
>>> for name, buf in self.named_buffers():
>>>     if name in ['running_var']:
>>>         print(buf.size())
```

#### `named_parameters(prefix='', recurse=True, remove_duplicate=True)` {#thunder.ThunderModule.named_parameters}

[\[source\]](https://github.com/Lightning-AI/lightning-thunder/blob/0.2.6/thunder/core/module.py#L102-L109)

Return an iterator over module parameters, yielding both the name of the parameter as well as the parameter itself.

Parameters:

*   **prefix** ([_str_](https://docs.python.org/3/builtins/stdtypes.html#str "(in Python v3.15)")) – prefix to prepend to all parameter names.
    
*   **recurse** ([_bool_](https://docs.python.org/3/builtins/functions.html#bool "(in Python v3.15)")) – if True, then yields parameters of this module and all submodules. Otherwise, yields only parameters that are direct members of this module.
    
*   **remove\_duplicate** ([_bool_](https://docs.python.org/3/builtins/functions.html#bool "(in Python v3.15)")_,_ _optional_) – whether to remove the duplicated parameters in the result. Defaults to True.
    

Yields:

_(str, Parameter)_ – Tuple containing the name and parameter

Example:

```python
>>> # xdoctest: +SKIP("undefined vars")
>>> for name, param in self.named_parameters():
>>>     if name in ['bias']:
>>>         print(param.size())
```

#### `no_sync()` {#thunder.ThunderModule.no_sync}

[\[source\]](https://github.com/Lightning-AI/lightning-thunder/blob/0.2.6/../home/runner/.local/share/uv/python/cpython-3.10-linux-x86_64-gnu/lib/python3.10/contextlib.py#L340-L382)

Context manager to disable gradient synchronization in data parallel mode.

This context manager is intended to be used in conjunction with [`torch.nn.parallel.DistributedDataParallel`](https://docs.pytorch.org/docs/stable/generated/torch.nn.parallel.DistributedDataParallel.html#torch.nn.parallel.DistributedDataParallel "(in PyTorch v2.15)") to disable gradient synchronization in the backward pass. It will not have any effect when used with other modules.

Note

This could lead to different accumulated gradients with `torch.nn.parallel.distributed.DistributedDataParallel.no_sync`. PyTorch’s gradient synchronization is implemented by applying all-reduce to gradient buckets of `torch.nn.Parameter.grad`. Thus the `no_sync` context leads to $`\text{AllReduce} \left( \sum_{i = 0}^{\text{ga_steps}} g_i \right)`$ where $`\text{ga_steps}`$ means the number of gradient accumulation steps. In contrast, this synchronizes accumulated gradients when exiting, leading to $`\text{AllReduce} \left( \sum_{i = 0}^{\text{ga_steps - 1}} g_i \right) + \text{AllReduce}(g_{\text{ga_steps}})`$.

Warning

You must reuse this context manager in each group of gradient accumulation iterations since gradients will get synchronized on context manager exit.

```python
with model.no_sync():
    for _ in range(len(gradient_accumulation_iters)):
        loss(model(x)).backward()  # uses no-sync-backward trace
loss(model(x)).backward()  # uses the regular backward trace
optimizer.step()
```

#### `original_state_dict(*, destination=None, prefix='', keep_vars=False)` {#thunder.ThunderModule.original_state_dict}

[\[source\]](https://github.com/Lightning-AI/lightning-thunder/blob/0.2.6/thunder/core/module.py#L230-L265)

Returns the state dict of the transformed [`ThunderModule`](#thunder.ThunderModule "thunder.ThunderModule") with the reverse transform applied.

For example, [`ThunderModule.state_dict()`](#thunder.ThunderModule.state_dict "thunder.ThunderModule.state_dict") returns a state dict of sharded tensors if a model is [`thunder.distributed.fsdp()`](/docs/thunder/0.2.6/distributed/generated/thunder.distributed.fsdp.md#thunder.distributed.fsdp "thunder.distributed.fsdp") applied while [`ThunderModule.original_state_dict()`](#thunder.ThunderModule.original_state_dict "thunder.ThunderModule.original_state_dict") returns a state dict of unsharded tensors.

Parameters:

*   **destination** ([`Optional`](https://docs.python.org/3/library/typing.html#typing.Optional "(in Python v3.15)")\[[`dict`](https://docs.python.org/3/builtins/stdtypes.html#dict "(in Python v3.15)")\[[`str`](https://docs.python.org/3/builtins/stdtypes.html#str "(in Python v3.15)"), [`Any`](https://docs.python.org/3/library/typing.html#typing.Any "(in Python v3.15)")\]\]) – if given, use this mutable mapping as the dict container.
    
*   **prefix** ([`str`](https://docs.python.org/3/builtins/stdtypes.html#str "(in Python v3.15)")) – a prefix for the keys.
    
*   **keep\_vars** ([`bool`](https://docs.python.org/3/builtins/functions.html#bool "(in Python v3.15)")) – do not detach
    
*   **destination** ([_Union_](https://docs.python.org/3/library/typing.html#typing.Union "(in Python v3.15)")_\[__None__,_ [_dict_](https://docs.python.org/3/builtins/stdtypes.html#dict "(in Python v3.15)")_\[_[_str_](https://docs.python.org/3/builtins/stdtypes.html#str "(in Python v3.15)")_,_ [_Any_](https://docs.python.org/3/library/typing.html#typing.Any "(in Python v3.15)")_\]__\]_) –
    
*   **prefix** ([_str_](https://docs.python.org/3/builtins/stdtypes.html#str "(in Python v3.15)")) –
    
*   **keep\_vars** ([_bool_](https://docs.python.org/3/builtins/functions.html#bool "(in Python v3.15)")) –
    

Return type:

[`dict`](https://docs.python.org/3/builtins/stdtypes.html#dict "(in Python v3.15)")\[[`str`](https://docs.python.org/3/builtins/stdtypes.html#str "(in Python v3.15)"), [`Any`](https://docs.python.org/3/library/typing.html#typing.Any "(in Python v3.15)")\]

#### `state_dict(*, destination=None, prefix='', keep_vars=False)` {#thunder.ThunderModule.state_dict}

[\[source\]](https://github.com/Lightning-AI/lightning-thunder/blob/0.2.6/thunder/core/module.py#L192-L228)

Returns the state dict of the (transformed) Thunder module.

Parameters:

*   **destination** – if given, use this mutable mapping as the dict container.
    
*   **prefix** – a prefix for the keys.
    
*   **keep\_vars** – do not detach
    

Note that this is similar but rather more rudimentary than PyTorch’s state\_dict (e.g. no hook support yet).

---

**Navigation**

- Parent: [API reference](/docs/thunder/0.2.6/api-reference.md)