torch_to_nnef.tensor

Modules:

Name	Description
named
offload	OffLoad Tensor.
opaque
quant	Advanced QTensor (<= 8bits) with complex quant scheme non torch native.
updater
utils

Classes:

Name	Description
NamedTensor	Tensor enriched with name attribute.
OffloadedTensor	Tensor subclass that maintains data on disk.
OpaqueTensorRef	Allow to pass through 'tracing'.
QScalePerGroupF16	f16 scale only per group.
QTensor	Common interface for all Compressed storage.
QTensorTractScaleOnly	Tract data format it serializes to: Q4_0.

Functions:

Name	Description
apply_name_to_tensor_in_module	Transform torch.Tensor or Parameters into NamedTensor.
set_opaque_tensor_in_params_as_ref	Transform OpaqueTensor Parameters into OpaqueTensorRef.

NamedTensor

NamedTensor(fp_tensor: Tensor, nnef_name: str)

Bases: Tensor

Tensor enriched with name attribute.

Attributes:

Name	Type	Description
data		Very important to keep access to all special attr of NamedTensor.

data property writable

data

Very important to keep access to all special attr of NamedTensor.

OffloadedTensor

OffloadedTensor(elem, device, offload_dir: Path, name: str, offloaded_tensor_type: Type[Tensor], force_gc_collect: bool = False)

Bases: OpaqueTensor

Tensor subclass that maintains data on disk.

It hold an virtual internal memory storage (permanent) and a temporary instantiation at each operation accessing it on targeted device.

Warning

we recommend to version of PyTorch > 1.12 for best compatibility.

Methods:

Name	Description
from_original_tensor	Take a torch.Tensor or OpaqueTensor and offload it to disk.
set_	Implement tensor-style storage replacement for offloaded payloads.
to	Change the target device when reloaded in memory.
update_values	Replace offloaded tensor by new 'values' tensor.

Attributes:

Name	Type	Description
is_meta	bool	Whether the tensor is on the meta device.

is_meta property

is_meta: bool

Whether the tensor is on the meta device.

Always False as the tensor is (off|re)loaded from disk.

from_original_tensor classmethod

from_original_tensor(tensor: Tensor, name: str, offload_dir: Optional[Path] = None, suffix_log_msg: str = '')

Take a torch.Tensor or OpaqueTensor and offload it to disk.

Parameters:

Name	Type	Description	Default
tensor	Tensor	the torch.Tensor or torch_to_nnef.tensor.OpaqueTensor to dump on disk	required
name	str	the name of the tensor that will be used to create the filename store on disk	required
offload_dir	Optional[Path]	The directory where this file will be stored (temporarly)	None
suffix_log_msg	str	Added message log suffix for context	''

set_

set_(source: Tensor, *args, **kwargs)

Implement tensor-style storage replacement for offloaded payloads.

OffloadedTensor uses a meta tensor as its in-memory shell, so PyTorch's native Tensor.set_ cannot replace its storage with a CPU tensor. Route the common param.set_(new_tensor) form through the offload store instead. This is important for quantizers that update a weight in-place before replacing it with a QTensor.

to

to(*args, **kwargs)

Change the target device when reloaded in memory.

update_values

update_values(values: Tensor, strict_shape: bool = True, strict_dtype: bool = True)

Replace offloaded tensor by new 'values' tensor.

Parameters:

Name	Type	Description	Default
values	Tensor	The tensor that will replace it on disk assertion are made to ensure same shape, dtype as prior	required
strict_shape	bool	if True (default) the shape of the new tensor must be the same as the prior one	True
strict_dtype	bool	if True (default) the dtype of the new tensor must be the same as the prior one	True

OpaqueTensorRef

OpaqueTensorRef(meta_tensor: Tensor, opaque_tensor: OpaqueTensor)

Bases: Tensor

Allow to pass through 'tracing'.

QScalePerGroupF16

QScalePerGroupF16(group_size: int, scale: Tensor, n_bits: int)

Bases: QScheme

f16 scale only per group.

Tract aligned using negative scales.

QTensor

QTensor(fp_tensor: Tensor, qscheme: QScheme, dequant_to_dtype=torch.float32, u8_compressors: Optional[List[U8Compressor]] = None)

Bases: OpaqueTensor

Common interface for all Compressed storage.

Methods:

Name	Description
to_device	Specific device handling.
write_in_file	Called at NNEF write time.

to_device

to_device(new_device)

Specific device handling.

write_in_file

write_in_file(dirpath: Union[str, Path], label: str)

Called at NNEF write time.

Each specific inference engine format should implement the file dump prefered.

QTensorTractScaleOnly

QTensorTractScaleOnly(*args, specific_machine: Optional[str] = None, **kwargs)

Bases: QTensorTract, SupportsOffloadState

Tract data format it serializes to: Q4_0.

Methods:

Name	Description
decompress	Tract dequantization depends on hardware.

decompress

decompress()

Tract dequantization depends on hardware.

Typically dequantization happen with ops in f16 on ARM and f32 (scale directly casted) on others so we overwrite the function to be consistant with tract.

apply_name_to_tensor_in_module

apply_name_to_tensor_in_module(model: Module)

Transform torch.Tensor or Parameters into NamedTensor.

This is applied at export time of torch_to_nnef Just before doing any tracing and allow to keep variable naming identical to PyTorch one

This consistent naming unlock subsequent manipulations such as LORA applications @ inference or such.

set_opaque_tensor_in_params_as_ref

set_opaque_tensor_in_params_as_ref(model: Module)

Transform OpaqueTensor Parameters into OpaqueTensorRef.

This is applied at export time of torch_to_nnef Just before doing any tracing