optimum-neuron/v0.4.0/en/training_api/transformations.md

# Model Weight Transformation Specs

The transformation specs API defines how model weights are transformed between the original Transformers implementation and the custom implementation optimized for Neuron devices. This enables automatic weight conversion during model loading and checkpoint consolidation.

## Base Classes

### ModelWeightTransformationSpec[[optimum.neuron.models.training.ModelWeightTransformationSpec]]

<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>class optimum.neuron.models.training.ModelWeightTransformationSpec</name><anchor>optimum.neuron.models.training.ModelWeightTransformationSpec</anchor><source>https://github.com/huggingface/optimum-neuron/blob/v0.4.0/optimum/neuron/models/training/transformations_utils.py#L91</source><parameters>[]</parameters></docstring>

This class defines the interface for transforming model weights between the original Transformers implementation
and the custom implementation for Neuron.



<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>adapt_peft_config</name><anchor>optimum.neuron.models.training.ModelWeightTransformationSpec.adapt_peft_config</anchor><source>https://github.com/huggingface/optimum-neuron/blob/v0.4.0/optimum/neuron/models/training/transformations_utils.py#L121</source><parameters>[{"name": "peft_config", "val": ": PeftConfig"}, {"name": "inplace", "val": ": bool = False"}]</parameters></docstring>

Adapts the PEFT config to match the custom modeling implementation.


</div>
<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>adapt_state_dict</name><anchor>optimum.neuron.models.training.ModelWeightTransformationSpec.adapt_state_dict</anchor><source>https://github.com/huggingface/optimum-neuron/blob/v0.4.0/optimum/neuron/models/training/transformations_utils.py#L157</source><parameters>[{"name": "module_fully_qualified_name", "val": ": str"}, {"name": "named_parameters", "val": ": dict[str, torch.nn.parameter.Parameter]"}, {"name": "orig_state_dict", "val": ": dict[str, torch.Tensor]"}, {"name": "upstanding_sharded_params", "val": ": dict[str, torch.Tensor]"}, {"name": "inplace", "val": ": bool = False"}]</parameters></docstring>

Transforms the state dict from the original Transformers model to match the custom modeling implementation.


</div>
<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>get_relevant_parameter_names</name><anchor>optimum.neuron.models.training.ModelWeightTransformationSpec.get_relevant_parameter_names</anchor><source>https://github.com/huggingface/optimum-neuron/blob/v0.4.0/optimum/neuron/models/training/transformations_utils.py#L107</source><parameters>[{"name": "module_fully_qualified_name", "val": ": str"}]</parameters></docstring>

Returns the set of parameter names that this spec would affect.


</div>
<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>guess_peft_type</name><anchor>optimum.neuron.models.training.ModelWeightTransformationSpec.guess_peft_type</anchor><source>https://github.com/huggingface/optimum-neuron/blob/v0.4.0/optimum/neuron/models/training/transformations_utils.py#L114</source><parameters>[{"name": "model", "val": ": Module"}, {"name": "module_fully_qualified_name", "val": ": str"}]</parameters></docstring>

Guesses the PEFT type of the module associated to the spec.


</div>
<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>to_original_peft_config</name><anchor>optimum.neuron.models.training.ModelWeightTransformationSpec.to_original_peft_config</anchor><source>https://github.com/huggingface/optimum-neuron/blob/v0.4.0/optimum/neuron/models/training/transformations_utils.py#L128</source><parameters>[{"name": "peft_config", "val": ": PeftConfig"}, {"name": "inplace", "val": ": bool = False"}]</parameters></docstring>

Restores the PEFT config to the original one that matches the original Transformers implementation.


</div>
<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>to_original_weights</name><anchor>optimum.neuron.models.training.ModelWeightTransformationSpec.to_original_weights</anchor><source>https://github.com/huggingface/optimum-neuron/blob/v0.4.0/optimum/neuron/models/training/transformations_utils.py#L207</source><parameters>[{"name": "module_fully_qualified_name", "val": ": str"}, {"name": "sharded_state_dicts", "val": ": dict[str, list[torch.Tensor]]"}, {"name": "parameters_metadata", "val": ": dict[str, dict[str, typing.Any]]"}]</parameters><paramsdesc>- **sharded_state_dicts** (dict[str, list[torch.Tensor]]) -- The sharded state dicts from the custom modeling
  implementation.
- **parameters_metadata** (dict[str, dict[str, Any]]) -- Metadata about the parameters in the original model.</paramsdesc><paramgroups>0</paramgroups><rettype>tuple[dict[str, torch.Tensor], list[str]]</rettype><retdesc>A tuple containing the transformed weights and a list of the
names of the parameters to remove from the final state dict.</retdesc></docstring>

Produces the weights associated to this transformation spec from the custom model to match the original
Transformers weights.








</div></div>

### ModelWeightTransformationSpecs[[optimum.neuron.models.training.ModelWeightTransformationSpecs]]

<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>class optimum.neuron.models.training.ModelWeightTransformationSpecs</name><anchor>optimum.neuron.models.training.ModelWeightTransformationSpecs</anchor><source>https://github.com/huggingface/optimum-neuron/blob/v0.4.0/optimum/neuron/models/training/transformations_utils.py#L239</source><parameters>[{"name": "module_fully_qualified_name", "val": ": str | None = None"}, {"name": "specs", "val": ": optimum.neuron.models.training.transformations_utils.ModelWeightTransformationSpec | list[optimum.neuron.models.training.transformations_utils.ModelWeightTransformationSpec] = <factory>"}]</parameters></docstring>

Defines a list of transformation specs for a given module of the model.


</div>

### CustomModule[[optimum.neuron.models.training.CustomModule]]

<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>class optimum.neuron.models.training.CustomModule</name><anchor>optimum.neuron.models.training.CustomModule</anchor><source>https://github.com/huggingface/optimum-neuron/blob/v0.4.0/optimum/neuron/models/training/transformations_utils.py#L342</source><parameters>[]</parameters></docstring>

This class is used to mark a module as a custom module. It is used to identify the modules that contain weights
that need to transformed when loading and saving the model.


</div>

## Transformation Specifications

### FusedLinearsSpec[[optimum.neuron.models.training.FusedLinearsSpec]]

<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>class optimum.neuron.models.training.FusedLinearsSpec</name><anchor>optimum.neuron.models.training.FusedLinearsSpec</anchor><source>https://github.com/huggingface/optimum-neuron/blob/v0.4.0/optimum/neuron/models/training/transformations_utils.py#L365</source><parameters>[{"name": "fused_linear_name", "val": ": str"}, {"name": "linear_names", "val": ": list[str]"}, {"name": "bias", "val": ": bool"}, {"name": "fuse_axis", "val": ": typing.Union[typing.Literal[0], typing.Literal[1], typing.Literal['column'], typing.Literal['row']]"}, {"name": "original_dims", "val": ": list[int]"}, {"name": "tp_size", "val": ": int = <factory>"}]</parameters></docstring>

Represents a transformation where multiple linear layers are fused into a single linear layer.
It can handle the case where the fused linear layer is sharded across multiple tensor parallel ranks.


</div>

### GQAQKVColumnParallelLinearSpec[[optimum.neuron.models.training.GQAQKVColumnParallelLinearSpec]]

<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>class optimum.neuron.models.training.GQAQKVColumnParallelLinearSpec</name><anchor>optimum.neuron.models.training.GQAQKVColumnParallelLinearSpec</anchor><source>https://github.com/huggingface/optimum-neuron/blob/v0.4.0/optimum/neuron/models/training/transformations_utils.py#L736</source><parameters>[{"name": "gqa_qkv_projection_name", "val": ": str"}, {"name": "query_projection_name", "val": ": str"}, {"name": "key_projection_name", "val": ": str"}, {"name": "value_projection_name", "val": ": str"}, {"name": "output_projection_name", "val": ": str"}, {"name": "num_attention_heads", "val": ": int"}, {"name": "num_key_value_heads", "val": ": int"}, {"name": "kv_size_multiplier", "val": ": int"}, {"name": "q_output_size_per_partition", "val": ": int"}, {"name": "kv_output_size_per_partition", "val": ": int"}, {"name": "fuse_qkv", "val": ": bool"}, {"name": "bias", "val": ": bool"}, {"name": "tp_size", "val": ": int = <factory>"}]</parameters></docstring>

Represents the transformation of separate query, key, and value projections into a single GQAQKVColumnParalleLinear
projection.



<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>compute_query_indices_for_rank</name><anchor>optimum.neuron.models.training.GQAQKVColumnParallelLinearSpec.compute_query_indices_for_rank</anchor><source>https://github.com/huggingface/optimum-neuron/blob/v0.4.0/optimum/neuron/models/training/transformations_utils.py#L808</source><parameters>[{"name": "tp_size", "val": ": int"}, {"name": "tp_rank", "val": ": int"}, {"name": "num_attention_heads", "val": ": int"}, {"name": "num_key_value_heads", "val": ": int"}, {"name": "kv_size_multiplier", "val": ": int"}]</parameters></docstring>

Computes the permutation for the query weight for a given TP rank.


</div>
<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>create_kv_proj_local_weight_from_regular_weight</name><anchor>optimum.neuron.models.training.GQAQKVColumnParallelLinearSpec.create_kv_proj_local_weight_from_regular_weight</anchor><source>https://github.com/huggingface/optimum-neuron/blob/v0.4.0/optimum/neuron/models/training/transformations_utils.py#L851</source><parameters>[{"name": "weight_data", "val": ": Tensor"}, {"name": "kv_size_multiplier", "val": ": int"}, {"name": "output_size_per_partition", "val": ": int"}]</parameters></docstring>

Creates the local version of the key or value projections weight for the given TP rank.


</div>
<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>create_query_or_output_projection_local_weight_from_regular_weight</name><anchor>optimum.neuron.models.training.GQAQKVColumnParallelLinearSpec.create_query_or_output_projection_local_weight_from_regular_weight</anchor><source>https://github.com/huggingface/optimum-neuron/blob/v0.4.0/optimum/neuron/models/training/transformations_utils.py#L866</source><parameters>[{"name": "weight_data", "val": ": Tensor"}, {"name": "num_attention_heads", "val": ": int"}, {"name": "num_key_value_heads", "val": ": int"}, {"name": "kv_size_multiplier", "val": ": int"}, {"name": "query_or_output_proj", "val": ": typing.Union[typing.Literal['query'], typing.Literal['output']]"}]</parameters></docstring>

Creates the local version of the query or output projections weight for the given TP rank.


</div></div>

## Utility Functions

### Weight Creation Functions[[optimum.neuron.models.training.transformations_utils.create_local_weight_with_padding]]

<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>optimum.neuron.models.training.transformations_utils.create_local_weight_with_padding</name><anchor>optimum.neuron.models.training.transformations_utils.create_local_weight_with_padding</anchor><source>https://github.com/huggingface/optimum-neuron/blob/v0.4.0/optimum/neuron/models/training/transformations_utils.py#L54</source><parameters>[{"name": "full_weight", "val": ": Tensor"}, {"name": "partition_dim", "val": ": int"}, {"name": "stride", "val": ": int"}, {"name": "out_weight", "val": ": torch.Tensor | None = None"}]</parameters></docstring>

Shards a tensor along a given axis and return a slice corresponding to the rank.
This will round up the layer to the next multiple if there is need to pad the tensor.


</div>

<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>optimum.neuron.models.training.transformations_utils.create_local_fused_weight</name><anchor>optimum.neuron.models.training.transformations_utils.create_local_fused_weight</anchor><source>https://github.com/huggingface/optimum-neuron/blob/v0.4.0/optimum/neuron/models/training/transformations_utils.py#L73</source><parameters>[{"name": "tp_rank", "val": ""}, {"name": "tp_size", "val": ""}, {"name": "individual_weights", "val": ""}, {"name": "partition_dim", "val": ""}, {"name": "fuse_axis", "val": ""}, {"name": "out_weight", "val": " = None"}]</parameters></docstring>

Shards individual weights across the tensor parallel ranks and fuses them into a single weight.


</div>

### Model-level Functions[[optimum.neuron.models.training.specialize_transformation_specs_for_model]]

<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>optimum.neuron.models.training.specialize_transformation_specs_for_model</name><anchor>optimum.neuron.models.training.specialize_transformation_specs_for_model</anchor><source>https://github.com/huggingface/optimum-neuron/blob/v0.4.0/optimum/neuron/models/training/transformations_utils.py#L1458</source><parameters>[{"name": "model", "val": ": Module"}]</parameters></docstring>


</div>

<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>optimum.neuron.models.training.adapt_peft_config_for_model</name><anchor>optimum.neuron.models.training.adapt_peft_config_for_model</anchor><source>https://github.com/huggingface/optimum-neuron/blob/v0.4.0/optimum/neuron/models/training/transformations_utils.py#L1467</source><parameters>[{"name": "model", "val": ": Module"}, {"name": "peft_config", "val": ": peft.config.PeftConfig | dict[str, peft.config.PeftConfig]"}, {"name": "inplace", "val": ": bool = False"}]</parameters></docstring>


</div>

<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>optimum.neuron.models.training.to_original_peft_config_for_model</name><anchor>optimum.neuron.models.training.to_original_peft_config_for_model</anchor><source>https://github.com/huggingface/optimum-neuron/blob/v0.4.0/optimum/neuron/models/training/transformations_utils.py#L1484</source><parameters>[{"name": "model", "val": ": Module"}, {"name": "peft_config", "val": ": PeftConfig"}, {"name": "inplace", "val": ": bool = False"}]</parameters></docstring>


</div>

### State Dict Functions[[optimum.neuron.models.training.adapt_state_dict]]

<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>optimum.neuron.models.training.adapt_state_dict</name><anchor>optimum.neuron.models.training.adapt_state_dict</anchor><source>https://github.com/huggingface/optimum-neuron/blob/v0.4.0/optimum/neuron/models/training/transformations_utils.py#L1516</source><parameters>[{"name": "model", "val": ": Module"}, {"name": "state_dict", "val": ": dict[str, torch.Tensor]"}, {"name": "upstanding_sharded_params", "val": ": dict[str, torch.Tensor]"}, {"name": "inplace", "val": ": bool = False"}, {"name": "**peft_kwargs", "val": ": Any"}]</parameters></docstring>

Transforms the state dict from the original Transformers model to match the custom modeling implementation.


</div>

<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>optimum.neuron.models.training.to_original_weights</name><anchor>optimum.neuron.models.training.to_original_weights</anchor><source>https://github.com/huggingface/optimum-neuron/blob/v0.4.0/optimum/neuron/models/training/transformations_utils.py#L1590</source><parameters>[{"name": "transformations_specs", "val": ": list[optimum.neuron.models.training.transformations_utils.ModelWeightTransformationSpecs]"}, {"name": "sharded_state_dicts", "val": ": dict[str, list[torch.Tensor]]"}, {"name": "parameters_metadata", "val": ": dict[str, dict[str, typing.Any]]"}, {"name": "**peft_kwargs", "val": ": Any"}]</parameters></docstring>

Consolidates the sharded state dicts produced by saving the custom model into a single state dict that matches the
original Transformers model weights.


</div>

### Metadata Functions[[optimum.neuron.models.training.create_parameter_metadata]]

<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>optimum.neuron.models.training.create_parameter_metadata</name><anchor>optimum.neuron.models.training.create_parameter_metadata</anchor><source>https://github.com/huggingface/optimum-neuron/blob/v0.4.0/optimum/neuron/models/training/transformations_utils.py#L1659</source><parameters>[{"name": "model", "val": ""}]</parameters></docstring>

Creates the metadata to be saved with the model weights to be able to reconstruct the original weights when
consolidating the sharded state dicts.


</div>

<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>optimum.neuron.models.training.transformations_utils.get_tensor_model_parallel_attributes</name><anchor>optimum.neuron.models.training.transformations_utils.get_tensor_model_parallel_attributes</anchor><source>https://github.com/huggingface/optimum-neuron/blob/v0.4.0/optimum/neuron/models/training/transformations_utils.py#L1645</source><parameters>[{"name": "tensor", "val": ": Tensor"}]</parameters></docstring>

Returns the tensor model parallel attributes of a tensor.


</div>

### Helper Functions[[optimum.neuron.models.training.transformations_utils.remove_adapter_name]]

<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>optimum.neuron.models.training.transformations_utils.remove_adapter_name</name><anchor>optimum.neuron.models.training.transformations_utils.remove_adapter_name</anchor><source>https://github.com/huggingface/optimum-neuron/blob/v0.4.0/optimum/neuron/models/training/transformations_utils.py#L1501</source><parameters>[{"name": "name", "val": ": str"}]</parameters></docstring>


</div>

<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>optimum.neuron.models.training.transformations_utils.is_base_layer</name><anchor>optimum.neuron.models.training.transformations_utils.is_base_layer</anchor><source>https://github.com/huggingface/optimum-neuron/blob/v0.4.0/optimum/neuron/models/training/transformations_utils.py#L1505</source><parameters>[{"name": "name", "val": ": str"}]</parameters></docstring>


</div>

<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>optimum.neuron.models.training.transformations_utils.get_adapter_name</name><anchor>optimum.neuron.models.training.transformations_utils.get_adapter_name</anchor><source>https://github.com/huggingface/optimum-neuron/blob/v0.4.0/optimum/neuron/models/training/transformations_utils.py#L1509</source><parameters>[{"name": "parameter_fully_qualified_name", "val": ": str"}]</parameters></docstring>


</div>