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]]

class optimum.neuron.models.training.ModelWeightTransformationSpecoptimum.neuron.models.training.ModelWeightTransformationSpechttps://github.com/huggingface/optimum-neuron/blob/v0.4.0/optimum/neuron/models/training/transformations_utils.py#L91[]

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

adapt_peft_configoptimum.neuron.models.training.ModelWeightTransformationSpec.adapt_peft_confighttps://github.com/huggingface/optimum-neuron/blob/v0.4.0/optimum/neuron/models/training/transformations_utils.py#L121[{"name": "peft_config", "val": ": PeftConfig"}, {"name": "inplace", "val": ": bool = False"}]

Adapts the PEFT config to match the custom modeling implementation.

adapt_state_dictoptimum.neuron.models.training.ModelWeightTransformationSpec.adapt_state_dicthttps://github.com/huggingface/optimum-neuron/blob/v0.4.0/optimum/neuron/models/training/transformations_utils.py#L157[{"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"}]

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

get_relevant_parameter_namesoptimum.neuron.models.training.ModelWeightTransformationSpec.get_relevant_parameter_nameshttps://github.com/huggingface/optimum-neuron/blob/v0.4.0/optimum/neuron/models/training/transformations_utils.py#L107[{"name": "module_fully_qualified_name", "val": ": str"}]

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

guess_peft_typeoptimum.neuron.models.training.ModelWeightTransformationSpec.guess_peft_typehttps://github.com/huggingface/optimum-neuron/blob/v0.4.0/optimum/neuron/models/training/transformations_utils.py#L114[{"name": "model", "val": ": Module"}, {"name": "module_fully_qualified_name", "val": ": str"}]

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

to_original_peft_configoptimum.neuron.models.training.ModelWeightTransformationSpec.to_original_peft_confighttps://github.com/huggingface/optimum-neuron/blob/v0.4.0/optimum/neuron/models/training/transformations_utils.py#L128[{"name": "peft_config", "val": ": PeftConfig"}, {"name": "inplace", "val": ": bool = False"}]

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

to_original_weightsoptimum.neuron.models.training.ModelWeightTransformationSpec.to_original_weightshttps://github.com/huggingface/optimum-neuron/blob/v0.4.0/optimum/neuron/models/training/transformations_utils.py#L207[{"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]]"}]- 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.0tuple[dict[str, torch.Tensor], list[str]]A tuple containing the transformed weights and a list of the names of the parameters to remove from the final state dict.

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

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

class optimum.neuron.models.training.ModelWeightTransformationSpecsoptimum.neuron.models.training.ModelWeightTransformationSpecshttps://github.com/huggingface/optimum-neuron/blob/v0.4.0/optimum/neuron/models/training/transformations_utils.py#L239[{"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] = "}]

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

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

class optimum.neuron.models.training.CustomModuleoptimum.neuron.models.training.CustomModulehttps://github.com/huggingface/optimum-neuron/blob/v0.4.0/optimum/neuron/models/training/transformations_utils.py#L342[]

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.

Transformation Specifications

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

class optimum.neuron.models.training.FusedLinearsSpecoptimum.neuron.models.training.FusedLinearsSpechttps://github.com/huggingface/optimum-neuron/blob/v0.4.0/optimum/neuron/models/training/transformations_utils.py#L365[{"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 = "}]

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.

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

class optimum.neuron.models.training.GQAQKVColumnParallelLinearSpecoptimum.neuron.models.training.GQAQKVColumnParallelLinearSpechttps://github.com/huggingface/optimum-neuron/blob/v0.4.0/optimum/neuron/models/training/transformations_utils.py#L736[{"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 = "}]

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

compute_query_indices_for_rankoptimum.neuron.models.training.GQAQKVColumnParallelLinearSpec.compute_query_indices_for_rankhttps://github.com/huggingface/optimum-neuron/blob/v0.4.0/optimum/neuron/models/training/transformations_utils.py#L808[{"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"}]

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

create_kv_proj_local_weight_from_regular_weightoptimum.neuron.models.training.GQAQKVColumnParallelLinearSpec.create_kv_proj_local_weight_from_regular_weighthttps://github.com/huggingface/optimum-neuron/blob/v0.4.0/optimum/neuron/models/training/transformations_utils.py#L851[{"name": "weight_data", "val": ": Tensor"}, {"name": "kv_size_multiplier", "val": ": int"}, {"name": "output_size_per_partition", "val": ": int"}]

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

create_query_or_output_projection_local_weight_from_regular_weightoptimum.neuron.models.training.GQAQKVColumnParallelLinearSpec.create_query_or_output_projection_local_weight_from_regular_weighthttps://github.com/huggingface/optimum-neuron/blob/v0.4.0/optimum/neuron/models/training/transformations_utils.py#L866[{"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']]"}]

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

Utility Functions

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

optimum.neuron.models.training.transformations_utils.create_local_weight_with_paddingoptimum.neuron.models.training.transformations_utils.create_local_weight_with_paddinghttps://github.com/huggingface/optimum-neuron/blob/v0.4.0/optimum/neuron/models/training/transformations_utils.py#L54[{"name": "full_weight", "val": ": Tensor"}, {"name": "partition_dim", "val": ": int"}, {"name": "stride", "val": ": int"}, {"name": "out_weight", "val": ": torch.Tensor | None = None"}]

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.

optimum.neuron.models.training.transformations_utils.create_local_fused_weightoptimum.neuron.models.training.transformations_utils.create_local_fused_weighthttps://github.com/huggingface/optimum-neuron/blob/v0.4.0/optimum/neuron/models/training/transformations_utils.py#L73[{"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"}]

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

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

optimum.neuron.models.training.specialize_transformation_specs_for_modeloptimum.neuron.models.training.specialize_transformation_specs_for_modelhttps://github.com/huggingface/optimum-neuron/blob/v0.4.0/optimum/neuron/models/training/transformations_utils.py#L1458[{"name": "model", "val": ": Module"}]

optimum.neuron.models.training.adapt_peft_config_for_modeloptimum.neuron.models.training.adapt_peft_config_for_modelhttps://github.com/huggingface/optimum-neuron/blob/v0.4.0/optimum/neuron/models/training/transformations_utils.py#L1467[{"name": "model", "val": ": Module"}, {"name": "peft_config", "val": ": peft.config.PeftConfig | dict[str, peft.config.PeftConfig]"}, {"name": "inplace", "val": ": bool = False"}]

optimum.neuron.models.training.to_original_peft_config_for_modeloptimum.neuron.models.training.to_original_peft_config_for_modelhttps://github.com/huggingface/optimum-neuron/blob/v0.4.0/optimum/neuron/models/training/transformations_utils.py#L1484[{"name": "model", "val": ": Module"}, {"name": "peft_config", "val": ": PeftConfig"}, {"name": "inplace", "val": ": bool = False"}]

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

optimum.neuron.models.training.adapt_state_dictoptimum.neuron.models.training.adapt_state_dicthttps://github.com/huggingface/optimum-neuron/blob/v0.4.0/optimum/neuron/models/training/transformations_utils.py#L1516[{"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"}]

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

optimum.neuron.models.training.to_original_weightsoptimum.neuron.models.training.to_original_weightshttps://github.com/huggingface/optimum-neuron/blob/v0.4.0/optimum/neuron/models/training/transformations_utils.py#L1590[{"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"}]

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

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

optimum.neuron.models.training.create_parameter_metadataoptimum.neuron.models.training.create_parameter_metadatahttps://github.com/huggingface/optimum-neuron/blob/v0.4.0/optimum/neuron/models/training/transformations_utils.py#L1659[{"name": "model", "val": ""}]

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

optimum.neuron.models.training.transformations_utils.get_tensor_model_parallel_attributesoptimum.neuron.models.training.transformations_utils.get_tensor_model_parallel_attributeshttps://github.com/huggingface/optimum-neuron/blob/v0.4.0/optimum/neuron/models/training/transformations_utils.py#L1645[{"name": "tensor", "val": ": Tensor"}]

Returns the tensor model parallel attributes of a tensor.

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

optimum.neuron.models.training.transformations_utils.remove_adapter_nameoptimum.neuron.models.training.transformations_utils.remove_adapter_namehttps://github.com/huggingface/optimum-neuron/blob/v0.4.0/optimum/neuron/models/training/transformations_utils.py#L1501[{"name": "name", "val": ": str"}]

optimum.neuron.models.training.transformations_utils.is_base_layeroptimum.neuron.models.training.transformations_utils.is_base_layerhttps://github.com/huggingface/optimum-neuron/blob/v0.4.0/optimum/neuron/models/training/transformations_utils.py#L1505[{"name": "name", "val": ": str"}]

optimum.neuron.models.training.transformations_utils.get_adapter_nameoptimum.neuron.models.training.transformations_utils.get_adapter_namehttps://github.com/huggingface/optimum-neuron/blob/v0.4.0/optimum/neuron/models/training/transformations_utils.py#L1509[{"name": "parameter_fully_qualified_name", "val": ": str"}]