AWS Trainium & Inferentia 文件
模型權重轉換規範
並獲得增強的文件體驗
開始使用
模型權重轉換規範
轉換規範 API 定義了模型權重如何在原始 Transformers 實現與為 Neuron 裝置最佳化的自定義實現之間進行轉換。這使得在模型載入和檢查點合併期間能夠自動進行權重轉換。
基類
ModelWeightTransformationSpec
此類定義了在原始 Transformers 實現和 Neuron 自定義實現之間轉換模型權重的介面。
調整 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 模型轉換為與自定義建模實現相匹配的形式。
返回此規範將影響的引數名稱集合。
猜測與規範關聯的模組的 PEFT 型別。
將 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]]
從自定義模型生成與此轉換規範相關的權重，以匹配原始 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
此類用於將一個模組標記為自定義模組。它用於識別那些在載入和儲存模型時需要轉換權重的模組。
轉換規範
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.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。
元資料函式
建立要與模型權重一起儲存的元資料，以便在合併分片 state dict 時能夠重構原始權重。
optimum.neuron.models.training.transformations_utils.get_tensor_model_parallel_attributes
< 原始碼 >( tensor: Tensor )
返回張量的模型並行屬性。
輔助函式
optimum.neuron.models.training.transformations_utils.get_adapter_name
< 原始碼 >( parameter_fully_qualified_name: str )