Hugging Face 的 LOGO

AWS Trainium & Inferentia 文档

模型权重转换规范

Hugging Face's logo
加入 Hugging Face 社区

并获得增强的文档体验

开始使用

模型权重转换规范

转换规范 API 定义了模型权重如何在原始 Transformers 实现与为 Neuron 设备优化的自定义实现之间进行转换。这使得在模型加载和检查点合并期间能够自动进行权重转换。

基类

ModelWeightTransformationSpec

class optimum.neuron.models.training.ModelWeightTransformationSpec

< >

( )

此类定义了在原始 Transformers 实现和 Neuron 自定义实现之间转换模型权重的接口。

adapt_peft_config

< >

( peft_config: PeftConfig inplace: bool = False )

调整 PEFT 配置以匹配自定义建模实现。

adapt_state_dict

< >

( module_fully_qualified_name: str named_parameters: dict[str, torch.nn.parameter.Parameter] orig_state_dict: dict[str, torch.Tensor] upstanding_sharded_params: dict[str, torch.Tensor] inplace: bool = False )

将 state dict 从原始 Transformers 模型转换为与自定义建模实现相匹配的形式。

get_relevant_parameter_names

< >

( module_fully_qualified_name: str )

返回此规范将影响的参数名称集合。

guess_peft_type

< >

( model: Module module_fully_qualified_name: str )

猜测与规范关联的模块的 PEFT 类型。

to_original_peft_config

< >

( peft_config: PeftConfig inplace: bool = False )

将 PEFT 配置恢复为与原始 Transformers 实现相匹配的原始配置。

to_original_weights

< >

( module_fully_qualified_name: str sharded_state_dicts: dict[str, list[torch.Tensor]] parameters_metadata: dict[str, dict[str, typing.Any]] ) → tuple[dict[str, torch.Tensor], list[str]]

参数

  • sharded_state_dicts (dict[str, list[torch.Tensor]]) — 来自自定义建模实现的分片 state dict。
  • parameters_metadata (dict[str, dict[str, Any]]) — 关于原始模型中参数的元数据。

返回

tuple[dict[str, torch.Tensor], list[str]]

一个元组，包含转换后的权重以及要从最终 state dict 中移除的参数名称列表。

从自定义模型生成与此转换规范相关的权重，以匹配原始 Transformers 权重。

ModelWeightTransformationSpecs

class optimum.neuron.models.training.ModelWeightTransformationSpecs

< >

( module_fully_qualified_name: str | None = None specs: optimum.neuron.models.training.transformations_utils.ModelWeightTransformationSpec | list[optimum.neuron.models.training.transformations_utils.ModelWeightTransformationSpec] = <factory> )

为模型的给定模块定义一个转换规范列表。

CustomModule

class optimum.neuron.models.training.CustomModule

< >

( )

此类用于将一个模块标记为自定义模块。它用于识别那些在加载和保存模型时需要转换权重的模块。

转换规范

FusedLinearsSpec

class optimum.neuron.models.training.FusedLinearsSpec

< >

( fused_linear_name: str linear_names: list[str] bias: bool fuse_axis: typing.Union[typing.Literal[0], typing.Literal[1], typing.Literal['column'], typing.Literal['row']] original_dims: list[int] tp_size: int = <factory> )

表示将多个线性层融合成一个单一线性层的转换。它可以处理融合后的线性层在多个张量并行 rank 上分片的情况。

GQAQKVColumnParallelLinearSpec

class optimum.neuron.models.training.GQAQKVColumnParallelLinearSpec

< >

( gqa_qkv_projection_name: str query_projection_name: str key_projection_name: str value_projection_name: str output_projection_name: str num_attention_heads: int num_key_value_heads: int kv_size_multiplier: int q_output_size_per_partition: int kv_output_size_per_partition: int fuse_qkv: bool bias: bool tp_size: int = <factory> )

表示将单独的查询、键和值投影转换为单个 GQAQKVColumnParalleLinear 投影。

compute_query_indices_for_rank

< >

( tp_size: int tp_rank: int num_attention_heads: int num_key_value_heads: int kv_size_multiplier: int )

计算给定 TP rank 的查询权重排列。

create_kv_proj_local_weight_from_regular_weight

< >

( weight_data: Tensor kv_size_multiplier: int output_size_per_partition: int )

为给定的 TP rank 创建键或值投影权重的本地版本。

create_query_or_output_projection_local_weight_from_regular_weight

< >

( weight_data: Tensor num_attention_heads: int num_key_value_heads: int kv_size_multiplier: int query_or_output_proj: typing.Union[typing.Literal['query'], typing.Literal['output']] )

为给定的 TP rank 创建查询或输出投影权重的本地版本。

实用函数

权重创建函数

optimum.neuron.models.training.transformations_utils.create_local_weight_with_padding

< >

( full_weight: Tensor partition_dim: int stride: int out_weight: torch.Tensor | None = None )

沿给定轴对张量进行分片，并返回与 rank 对应的切片。如果需要填充张量，此操作会将层向上取整到下一个倍数。

optimum.neuron.models.training.transformations_utils.create_local_fused_weight

< >

( tp_rank tp_size individual_weights partition_dim fuse_axis out_weight = None )

将单个权重在张量并行 ranks 上分片，并将它们融合成一个单一的权重。

模型级函数

optimum.neuron.models.training.specialize_transformation_specs_for_model

< >

( model: Module )

optimum.neuron.models.training.adapt_peft_config_for_model

< >

( model: Module peft_config: peft.config.PeftConfig | dict[str, peft.config.PeftConfig] inplace: bool = False )

optimum.neuron.models.training.to_original_peft_config_for_model

< >

( model: Module peft_config: PeftConfig inplace: bool = False )

State Dict 函数

optimum.neuron.models.training.adapt_state_dict

< >

( model: Module state_dict: dict[str, torch.Tensor] upstanding_sharded_params: dict[str, torch.Tensor] inplace: bool = False **peft_kwargs: Any )

将 state dict 从原始 Transformers 模型转换为与自定义建模实现相匹配的形式。

optimum.neuron.models.training.to_original_weights

< >

( transformations_specs: list[optimum.neuron.models.training.transformations_utils.ModelWeightTransformationSpecs] sharded_state_dicts: dict[str, list[torch.Tensor]] parameters_metadata: dict[str, dict[str, typing.Any]] **peft_kwargs: Any )

将保存自定义模型时产生的 sharded state dict 合并为一个与原始 Transformers 模型权重相匹配的 state dict。

元数据函数

optimum.neuron.models.training.create_parameter_metadata

< >

( model )

创建要与模型权重一起保存的元数据，以便在合并分片 state dict 时能够重构原始权重。

optimum.neuron.models.training.transformations_utils.get_tensor_model_parallel_attributes

< >

( tensor: Tensor )

返回张量的模型并行属性。

辅助函数

optimum.neuron.models.training.transformations_utils.remove_adapter_name

< >

( name: str )

optimum.neuron.models.training.transformations_utils.is_base_layer

< >

( name: str )

optimum.neuron.models.training.transformations_utils.get_adapter_name

< >

( parameter_fully_qualified_name: str )

←Neuron TRL Trainer Neuron 的 LoRA→