Skip to content

9. Export selected tensors & PEFT

Goals

At the end of this tutorial you will be able to:

  1. Export specific set of weight tensors without the graph.

Prerequisite

  • PyTorch and Python basics
  • 5 min to read this page

Sometime you wish to export only some weights, this happen for example if you fine-tuned a LLM with a PEFT technique like LoRA. In this case only a very limited set of weights are modified compared to the pre-trained neural network. So this can make sense to ship to the customer the base model only once (with the graph) at the beginning (for example app download), and then to send the LoRA weights separately as you iterate through your product (for example app update).

This patching logic is also very interesting in case you wish to allow multiple PEFT to be shipped and switched/selected at model load time.

To do so torch_to_nnef provide few convenient methods.

Command line

First if you happen to exactly want to export PEFT weights, we have a CLI for you:

# filepath can be .pth, .bin, .safetensors ...
t2n_export_peft_to_nnef \
    --read-filepath /my-little-model-file.pt \
    -o /tmp/my-outputdir

By default it exports LoRA weights, if you wish to apply it on different methods look at additional options (with --help), the core functionality behind this CLI is simple pattern matching so most of PEFT weight names matching with regex should work (DoRA, ...).

Basic API

If you wish to have programmatic control you can also use for on disk tensors:

  • torch_to_nnef.export_tensors_from_disk_to_nnef 

    export_tensors_from_disk_to_nnef(store_filepath: T.Union[Path, str], output_dir: T.Union[Path, str], filter_key: T.Optional[T.Callable[[str], bool]] = None, fn_check_found_tensors: T.Optional[T.Callable[[T.Dict[str, _Tensor]], bool]] = None) -> T.Dict[str, _Tensor]

    Export any statedict or safetensors file torch.Tensors to NNEF .dat file.

    Parameters:

    Name Type Description Default
    store_filepath Union[Path, str]

    the filepath that hold the .safetensors , .pt or .bin containing the state dict

    		 required
    output_dir Union[Path, str]

    directory to dump the NNEF tensor .dat files

    		 required
    filter_key Optional[Callable[[str], bool]]

    An optional function to filter specific keys to be exported

    		 None
    fn_check_found_tensors Optional[Callable[[Dict[str, _Tensor]], bool]]

    post checking function to ensure all requested tensors have effectively been dumped

    		 None

    Returns:

    Type Description
    Dict[str, _Tensor]

    a dict of tensor name as key and torch.Tensor values, identical to torch_to_nnef.export.export_tensors_to_nnef

    Examples:

    Simple filtered example

    >>> import tempfile
>>> from torch import nn
>>> class Mod(nn.Module):
...     def __init__(self):
...         super().__init__()
...         self.a = nn.Linear(1, 5)
...         self.b = nn.Linear(5, 1)
...
...     def forward(self, x):
...         return self.b(self.a(x))
>>> mod = Mod()
>>> pt_path = tempfile.mktemp(suffix=".pt")
>>> nnef_dir = tempfile.mkdtemp(suffix="_nnef")
>>> torch.save(mod.state_dict(), pt_path)
>>> def check(ts):
...     assert all(_.startswith("a.") for _ in ts)
>>> exported_tensors = export_tensors_from_disk_to_nnef(
...     pt_path,
...     nnef_dir,
...     lambda x: x.startswith("a."),
...     check
... )
>>> list(exported_tensors.keys())
['a.weight', 'a.bias']

Or from loaded tensors:

  • torch_to_nnef.export_tensors_to_nnef 

    export_tensors_to_nnef(name_to_torch_tensors: T.Dict[str, _Tensor], output_dir: Path) -> T.Dict[str, _Tensor]

    Export any torch.Tensors list to NNEF .dat file.

    Parameters:

    Name Type Description Default
    name_to_torch_tensors Dict[str, _Tensor]

    dict A map of name (that will be used to define .dat filename) and tensor values (that can also be special torch_to_nnef tensors)

    		 required
    output_dir Path

    directory to dump the NNEF tensor .dat files

    		 required

    Returns:

    Type Description
    Dict[str, _Tensor]

    a dict of tensor name as key and torch.Tensor values, identical to torch_to_nnef.export.export_tensors_to_nnef

    Examples:

    Simple example

    >>> import tempfile
>>> from torch import nn
>>> class Mod(nn.Module):
...     def __init__(self):
...         super().__init__()
...         self.a = nn.Linear(1, 5)
...         self.b = nn.Linear(5, 1)
...
...     def forward(self, x):
...         return self.b(self.a(x))
>>> mod = Mod()
>>> nnef_dir = tempfile.mkdtemp(suffix="_nnef")
>>> exported_tensors = export_tensors_to_nnef(
...     {k: v for k, v in mod.named_parameters() if k.startswith("b.")},
...     nnef_dir,
... )
>>> list(exported_tensors.keys())
['b.weight', 'b.bias']