op.custom_extractors provides mechanism to control extraction to NNEF.
while bypassing PyTorch full expansion of torch.Module within torch_graph
which by default use torch.jit.trace .
This may be for two main reasons
- Some layer such as LSTM/GRU have complex expension which are better handled by encapsulation instead of spreading high number of variable
- Some layer might not be serializable to .jit
- There might be some edge case where you prefer to keep full control on exported NNEF subgraph.
Modules:
| Name | Description |
|---|---|
base |
|
gdn |
Gated delta net recurrence export to tract_transformers_gdn_recurrent. |
moe |
MoE FFN export to tract_moe_ffn operator (tract_transformers extension). |
rnn |
|
Classes:
| Name | Description |
|---|---|
CausalConvUpdateReified |
Drop-in core for HF's |
GatedDeltaNetRecurrentReified |
Drop-in for HF's (chunk|recurrent)_gated_delta_rule functions. |
LSTMCellExtractor |
Decompose |
LSTMExtractor |
|
MoEFFN |
Reference MoE FFN block for export testing. |
ModuleInfoExtractor |
Class to take manual control of NNEF expansion of a nn.Module. |
Functions:
| Name | Description |
|---|---|
mark_moe_experts_for_q40 |
Mark every MoE block of |
Bases: Module
Drop-in core for HF's torch_causal_conv1d_update.
Depthwise causal conv1d + SiLU with a k-sample carry state. Eager
forward matches HF exactly but avoids F.conv1d (whose fake-tensor
meta kernel breaks under the offload tracing strategy): the conv is
an unfold + mul + sum. At export the call is reified as one
tract_transformers_causal_conv1d_update op.
Layout: hidden_states [b, C, S], conv_state [b, C, k],
weight [C, k] -> (out [b, C, S], final_state [b, C, k]).
The caller keeps responsibility for writing final_state back into
the HF cache (a plain copy_, handled by the aten copy handler).
Bases: Module
Drop-in for HF's (chunk|recurrent)_gated_delta_rule functions.
Eager forward reproduces torch_recurrent_gated_delta_rule
semantics exactly (the chunked variant computes the same recurrence
with different blocking, so substituting it is numerically
equivalent up to fp reordering). At export the whole call is
reified as one tract_transformers_gdn_recurrent op.
Bases: ModuleInfoExtractor
Decompose nn.LSTMCell into primitive NNEF ops.
Unlike nn.LSTM, an LSTMCell carries a single time-step. We emit:
preact = matmul(input, w_ih, T) + matmul(h, w_hh, T) + b_ih + b_hh
i, f, g, o = chunk(preact, 4, axis=-1)
c_new = sigmoid(f) * c + sigmoid(i) * tanh(g)
h_new = sigmoid(o) * tanh(c_new)
Input order from the user-facing wrapper is (input, h, c): the
internal nn.LSTMCell call expects (input, (h, c)) which is handled by
_call_original_mod_with_args.
Methods:
| Name | Description |
|---|---|
ordered_args |
Reorder args so the first one is |
Reorder args so the first one is input (shape (B, input_size)).
t2n's IR sometimes reorders the cell's inputs after FixedTensorList /
tuple expansion, surfacing them as e.g. (h, input, c). The cell's
input_size (= weight_ih.shape[1]) lets us pick the input tensor by
shape; the relative order of (h, c) follows the JIT graph's
prim::ListConstruct that builds hx.
MoEFFN(num_experts: int, d_model: int, d_hidden: int, k: int = 2, activation: str = 'silu', normalize_gates: bool = True, bias: bool = False)
Bases: Module
Reference MoE FFN block for export testing.
For production models use the transformers extractors directly. Expert biases are supported in eager execution but rejected on export: the operator's positional bias inputs require a SwiGLU expert layout.
Methods:
| Name | Description |
|---|---|
forward |
Forward: [T, D] -> [T, D]. |
Class to take manual control of NNEF expansion of a nn.Module.
You need to subclass it, and set MODULE_CLASS according to your targeted module.
Then write .convert_to_nnef according to your need.
Methods:
| Name | Description |
|---|---|
convert_to_nnef |
Control NNEF content to be written for each MODULE_CLASS. |
generate_in_torch_graph |
Internal method called by torch_to_nnef ir_graph. |
get_by_kind |
Get ModuleInfoExtractor by kind in torch_to_nnef internal IR. |
get_by_module |
Search if the module is one of the MODULE_CLASS registered. |
ordered_args |
Odered args for the module call. |
Control NNEF content to be written for each MODULE_CLASS.
This happen at macro level when converting from internal IR to NNEF IR stage.
This is the Core method to overwrite in subclass.
It is no different than any op implemented in torch_to_nnef
in the module
Internal method called by torch_to_nnef ir_graph.
classmethod
Get ModuleInfoExtractor by kind in torch_to_nnef internal IR.
classmethod
Search if the module is one of the MODULE_CLASS registered.
return appropriate ModuleInfoExtractor subclass if found
mark_moe_experts_for_q40(model: Module, *, quantize_q40: bool = True, layout: Optional[str] = None, module_class_names: Optional[Collection[str]] = None) -> int
Mark every MoE block of model for export-time expert handling.
Sets the module markers the tract_moe_ffn extractor consumes:
_t2n_quantize_moe_experts_q40 (quantize expert w1/w2/w3 to tract
q4_0 min-max AT EXTRACTION, so a large MoE never materializes dense
float experts) and _t2n_moe_expert_layout.
By default any module whose class the MoE extractor itself dispatches
on (i.e. with a registered weight adapter) is marked; pass
module_class_names to restrict marking to explicit class names.
Returns the number of modules marked; raises if nothing matched (silently exporting dense f16 experts is never what the caller wants).