本文属于机器翻译版本。若本译文内容与英语原文存在差异,则一律以英文原文为准。
SageMaker 模型并行库 v2 参考
以下是 SageMaker 模型并行库 v2 (SMPv2) 的参考资料。
SMPv2 核心功能配置参数
以下是激活和配置参数的完整列表 SageMaker 模型并行度库 v2 的核心功能 。它们必须以JSON格式编写,并在 Pyth SageMaker on 中传递给 PyTorch 估算器SDK或另存为的JSON文件。 SageMaker HyperPod
"hybrid_shard_degree":
Integer
,
"sm_activation_offloading":
Boolean
,
"activation_loading_horizon":
Integer
,
"fsdp_cache_flush_warnings":
Boolean
,
"allow_empty_shards":
Boolean
,
"tensor_parallel_degree":
Integer
,
"context_parallel_degree":
Integer
,
"expert_parallel_degree":
Integer
,
"random_seed":
Integer
如果设置为
0
,则回退到本机 PyTorch实现,而脚本API中则
tensor_parallel_degree
为 1。否则,它将
hybrid_shard_degree
根据
tensor_parallel_degree
和
world_size
计算可能的最大值。当回退到原生 PyTorch FSDP用例时,如果你使用的策略
FULL_SHARD
是的话,它会在整个集群中进行分片。GPUs如果
_HYBRID_SHARD_ZERO2
是
HYBRID_SHARD
或曾经是策略,则相当于
hybrid_shard_degree
8。当启用张量并行度时,它会根据修订后的分片。
hybrid_shard_degree
如果设置为
1
,则回退到本机 PyTorch实现,而在脚本
NO_SHARD
中,当设置API
tensor_parallel_degree
为 1 时,它将回退到本机实现。否则,它等同于
NO_SHARD
在任何给定的张量 parallel 组中。
如果设置为介于 2 和之间的整数
world_size
,则在指定的数量上进行分片。GPUs如果您未在FSDP脚本
sharding_strategy
中进行设置,则它会被重写为。
HYBRID_SHARD
如果设置了
_HYBRID_SHARD_ZERO2
,则使用
sharding_strategy
您指定的。
sm_activation_offloading
(布尔值)-指定是否启用SMP激活卸载实现。如果
False
,则卸载使用本机 PyTorch 实现。如果
True
,则使用SMP激活卸载实现。您还需要在脚本中使用 PyTorch 激活卸载包装器 (
torch.distributed.algorithms._checkpoint.checkpoint_wrapper.offload_wrapper
)。要了解更多信息,请参阅
激活分载
。默认值为
True
。
activation_loading_horizon
(整数)-一个整数,指定其激活卸载水平类型。FSDP这是输入可以同时存储在GPU内存中的检查点或已卸载层的最大数量。要了解更多信息,请参阅
激活分载
。输入值必须为正整数。默认值为
2
。
fsdp_cache_flush_warnings
(布尔值)-检测 PyTorch 内存管理器中是否发生缓存刷新并发出警告,因为它们会降低计算性能。默认值为
True
。
allow_empty_shards
(布尔值)-如果张量不可分割,则在对张量进行分片时是否允许空分片。这是针对某些场景中检查点操作期间的崩溃的实验性修复程序。禁用此功能会回到最初的 PyTorch 行为。默认值为
False
。
tensor_parallel_degree
(整数)-指定张量并行度。该值必须介于
1
和之间
world_size
。默认值为
1
。请注意,传递大于 1 的值不会自动启用上下文并行性;您还需要使用在训练
torch.sagemaker.transform
API脚本中封装模型。要了解更多信息,请参阅
张量并行性
。
context_parallel_degree
(整数)-指定上下文并行度。该值必须介于
1
和之间
world_size
,并且必须为
<=
hybrid_shard_degree
。默认值为
1
。请注意,传递大于 1 的值不会自动启用上下文并行性;您还需要使用在训练
torch.sagemaker.transform
API脚本中封装模型。要了解更多信息,请参阅
上下文并行性
。
expert_parallel_degree
(整数)-指定专家并行度。该值必须介于 1 和之间
world_size
。默认值为
1
。请注意,传递大于 1 的值不会自动启用上下文并行性;您还需要使用在训练
torch.sagemaker.transform
API脚本中封装模型。要了解更多信息,请参阅
专家平行性
。
random_seed
(整数)— 分布式模块中按SMP张量并行度或专家并行度进行随机运算的种子数。该种子被添加到张量平行或专家并行等级中,以设置每个等级的实际种子。对于每个张量平行和专家并行等级,它都是独一无二的。SMPv2 确保在张量并行和专家并行等级中生成的随机数分别与和大小写相匹配。 non-tensor-parallelism non-expert-parallelism
SMPv2
torch.sagemaker
软件包的参考资料
本节是对 SMP v2 提供的
torch.sagemaker
软件包的参考。
主题
- torch.sagemaker.delayed_param.DelayedParamIniter
- torch.sagemaker.distributed.checkpoint.state_dict_saver.async_save
- torch.sagemaker.distributed.checkpoint.state_dict_saver.maybe_finalize_async_calls
- torch.sagemaker.distributed.checkpoint.state_dict_saver.save
- torch.sagemaker.distributed.checkpoint.state_dict_loader.load
- torch.sagemaker.moe.moe_config.MoEConfig
- torch.sagemaker.nn.attn.FlashSelfAttention
- torch.sagemaker.nn.attn.FlashGroupedQueryAttention
- torch.sagemaker.nn.huggingface.llama_flashattn.LlamaFlashAttention
- torch.sagemaker.transform
- torch.sagemakerutil 函数和属性
torch.sagemaker.delayed_param.DelayedParamIniter
API用于 参数初始化延迟 申请 PyTorch模型。
class torch.sagemaker.delayed_param.DelayedParamIniter( model: nn.Module, init_method_using_config : Callable = None, verbose: bool = False,
model
(nn.Module
) — 用于封装和应用 SMP v2 延迟参数初始化功能的 PyTorch 模型。
init_method_using_config
(可调用)— 如果您使用 SMP v2 的张量并行实现或支持Hugging Face Transformer 模型与张量并SMP行度兼容,请将此参数保持为默认值,即。None
默认情况下,会DelayedParamIniter
API找出如何正确初始化给定模型。对于任何其他模型,您需要创建一个自定义参数初始化函数并将其添加到脚本中。以下代码片段是 SMP v2 为实现的默认init_method_using_config
函数。Hugging Face Transformer 模型与张量并SMP行度兼容使用以下代码片段作为参考,您可以创建自己的初始化配置函数,将其添加到脚本中,然后将其传递给的init_method_using_config
SMPDelayedParamIniter
API参数。from torch.sagemaker.utils.module_utils import empty_module_params, move_buffers_to_device # Define a custom init config function. def
custom_init_method_using_config
(module): d = torch.cuda.current_device() empty_module_params(module, device=d) if isinstance(module, (nn.Linear, Conv1D)): module.weight.data.normal_(mean=0.0, std=config.initializer_range) if module.bias is not None: module.bias.data.zero_() elif isinstance(module, nn.Embedding): module.weight.data.normal_(mean=0.0, std=config.initializer_range) if module.padding_idx is not None: module.weight.data[module.padding_idx].zero_() elif isinstance(module, nn.LayerNorm): module.weight.data.fill_(1.0) module.bias.data.zero_() elif isinstance(module, LlamaRMSNorm): module.weight.data.fill_(1.0) move_buffers_to_device(module, device=d) delayed_initer = DelayedParamIniter(model, init_method_using_config=custom_init_method_using_config
)有关上述代码片段中
torch.sagemaker.module_util
函数的更多信息,请参阅torch.sagemakerutil 函数和属性。
verbose
(布尔值)-是否在初始化和验证期间启用更详细的日志记录。默认值为False
。Methods
get_param_init_fn()
— 返回可以传递给 PyTorch FSDP包装类参数的param_init_fn
参数初始化函数。
get_post_param_init_fn()
— 返回可以传递给 PyTorch FSDP包装类参数的post_param_init_fn
参数初始化函数。当您在模型中绑定权重时,需要这样做。模型必须实现该方法tie_weights
。欲了解更多信息,请参阅中有关绑定重量的注意事项参数初始化延迟。
count_num_params
(module: nn.Module, *args: Tuple[nn.Parameter]
)-跟踪参数初始化函数正在初始化多少参数。这有助于实现以下validate_params_and_buffers_inited
方法。您通常不需要显式调用此函数,因为该validate_params_and_buffers_inited
方法在后端隐式调用此方法。
validate_params_and_buffers_inited
(enabled: bool=True
) — 这是一个上下文管理器,可帮助验证初始化的参数数量是否与模型中的参数总数相匹配。它还会验证所有参数和缓冲区现在都在GPU设备上,而不是元设备上。AssertionErrors
如果不满足这些条件,则会加注。此上下文管理器只是可选的,您无需使用此上下文管理器来初始化参数。
torch.sagemaker.distributed.checkpoint.state_dict_saver.async_save
异步保存API的条目。使用此方法将
state_dict
异步保存到指定checkpoint_id
的。def async_save( state_dict: STATE_DICT_TYPE, checkpoint_id: Union[str, os.PathLike, None] = None, storage_writer: Optional[StorageWriter] = None, planner: Optional[SavePlanner] = None, process_group: Optional[dist.ProcessGroup] = None, coordinator_rank: int = 0, queue : AsyncCallsQueue = None, sharded_strategy: Union[SaveShardedStrategy, Tuple[str, int], None] = None, wait_error_handling: bool = True, force_check_all_plans: bool = True, s3_region: Optional[str] = None, s3client_config: Optional[S3ClientConfig] = None ) -> None:
checkpoint_id
(str)-必填项。保存检查点的存储路径。
storage_writer
(StorageWriter)-可选。in 的实例StorageWriter
, PyTorch 用于执行写入操作。如果未指定,则使用的 StorageWriter
默认配置。
planner
(SavePlanner)-可选。中的一个实SavePlanner
例 PyTorch。如果未指定,则使用的 SavePlanner
默认配置。
process_group
(ProcessGroup)-可选。要处理的流程组。如果是None
,则使用默认(全局)进程组。
coordinator_rank
(int)-可选。协调员在执行集体通信运算符时的等级,例如AllReduce
。
queue
(AsyncRequestQueue)-可选。要使用的异步调度器。默认情况下,它采用全局参数DEFAULT_ASYNC_REQUEST_QUEUE
。
sharded_strategy
(PyTorchDistSaveShardedStrategy)-可选。用于保存检查点的分片策略。如果未指定,torch.sagemaker.distributed.checkpoint.state_dict_saver.PyTorchDistSaveShardedStrategy
则默认使用。
wait_error_handling
(bool)-可选。一个标志,用于指定是否等待所有等级完成错误处理。默认值为True
。
force_check_all_plans
(bool)-可选。一个标志,用于确定是否强制跨等级同步计划,即使在缓存命中时也是如此。默认值为True
。
s3_region
(str)-可选。S3 存储桶所在的区域。如果未指定,则根据推断出该checkpoint_id
区域。
s3client_config
(S3ClientConfig)-可选。为 S3 客户端公开可配置参数的数据类。如果未提供,则使用 S3 ClientConfig的默认配置。默认情况下,该 part_size
参数设置为 64MB。
torch.sagemaker.distributed.checkpoint.state_dict_saver.maybe_finalize_async_calls
此功能允许训练过程监控要完成的多个异步请求。
def maybe_finalize_async_calls( blocking=True, process_group=None ) -> List[int]:
blocking
(bool)-可选。如果是True
,它将等到所有活动请求完成。否则,它只会最终确定已经完成的异步请求。默认值为True
。
process_group
(ProcessGroup)-可选。要操作的流程组。如果设置为None
,则使用默认(全局)进程组。包含异步调用索引的列表已成功定稿。
torch.sagemaker.distributed.checkpoint.state_dict_saver.save
使用此方法将
state_dict
同步保存到指定checkpoint_id
的。def save( state_dict: STATE_DICT_TYPE, checkpoint_id: Union[str, os.PathLike, None] = None, storage_writer: Optional[StorageWriter] = None, planner: Optional[SavePlanner] = None, process_group: Optional[dist.ProcessGroup] = None, coordinator_rank: int = 0, wait_error_handling: bool = True, force_check_all_plans: bool = True, s3_region: Optional[str] = None, s3client_config: Optional[S3ClientConfig] = None ) -> None:
checkpoint_id
(str)-必填项。保存检查点的存储路径。
storage_writer
(StorageWriter)-可选。in 的实例StorageWriter
, PyTorch 用于执行写入操作。如果未指定,则使用的 StorageWriter
默认配置。
planner
(SavePlanner)-可选。中的一个实SavePlanner
例 PyTorch。如果未指定,则使用的 SavePlanner
默认配置。
process_group
(ProcessGroup)-可选。要处理的流程组。如果是None
,则使用默认(全局)进程组。
coordinator_rank
(int)-可选。协调员在执行集体通信运算符时的等级,例如AllReduce
。
wait_error_handling
(bool)-可选。一个标志,用于指定是否等待所有等级完成错误处理。默认值为True
。
force_check_all_plans
(bool)-可选。一个标志,用于确定是否强制跨等级同步计划,即使在缓存命中时也是如此。默认值为True
。
s3_region
(str)-可选。S3 存储桶所在的区域。如果未指定,则根据推断出该checkpoint_id
区域。
s3client_config
(S3ClientConfig)-可选。为 S3 客户端公开可配置参数的数据类。如果未提供,则使用 S3 ClientConfig的默认配置。默认情况下,该 part_size
参数设置为 64MB。
torch.sagemaker.distributed.checkpoint.state_dict_loader.load
加载分布式模型的状态字典 (
state_dict
)。def load( state_dict: Dict[str, Any], checkpoint_id: Union[str, os.PathLike, None] = None, storage_reader: Optional[StorageReader] = None, planner: Optional[LoadPlanner] = None, process_group: Optional[dist.ProcessGroup] = None, check_keys_matched: bool = True, coordinator_rank: int = 0, s3_region: Optional[str] = None, s3client_config: Optional[S3ClientConfig] = None ) -> None:
state_dict
(字典)-必填项。state_dict
待加载的。
checkpoint_id
(str)-必填项。检查点的 ID。的含义checkpoint_id
取决于存储空间。它可以是文件夹或文件的路径。如果存储是键值存储,则它也可以是密钥。
storage_reader
(StorageReader)-可选。in 的实例StorageReader
, PyTorch 用于执行读取操作。如果未指定,则分布式检查点将根据自动推断读取器。 checkpoint_id
如果也checkpoint_id
是None
,则会引发异常错误。
planner
(StorageReader)-可选。中的一个实LoadPlanner
例 PyTorch。如果未指定,则使用的 LoadPlanner
默认配置。
check_keys_matched
(bool)-可选。如果启用,则使用检查所有等级的state_dict
密钥是否匹配AllGather
。
s3_region
(str)-可选。S3 存储桶所在的区域。如果未指定,则根据推断出该checkpoint_id
区域。
s3client_config
(S3ClientConfig)-可选。为 S3 客户端公开可配置参数的数据类。如果未提供,则使用 S3 ClientConfig的默认配置。默认情况下,该 part_size
参数设置为 64MB。
torch.sagemaker.moe.moe_config.MoEConfig
用于设置 Mixture-of SMP-Experts (MoE) 实现的配置类。您可以通过该类指定 MoE 配置值并将其传递给
torch.sagemaker.transform
API调用。要详细了解如何使用该课程训练 MoE 模型,请参阅专家平行性。class torch.sagemaker.moe.moe_config.MoEConfig( smp_moe=True, random_seed=12345, moe_load_balancing="sinkhorn", global_token_shuffle=False, moe_all_to_all_dispatcher=True, moe_aux_loss_coeff=0.001, moe_z_loss_coeff=0.001
smp_moe
(布尔值)-是否使用 MoE 的 SMP-实现。默认值为True
。
random_seed
(整数)-专家并行分布式模块中随机运算的种子数。此种子被添加到专家平行等级中,以设置每个等级的实际种子。对于每个专家并行等级,它都是独一无二的。默认值为12345
。
moe_load_balancing
(字符串)-指定 MoE 路由器的负载平衡类型。有效的选项是aux_loss
sinkhorn
、balanced
、和none
。默认值为sinkhorn
。
global_token_shuffle
(布尔值)-是否在同一 EP 组内的 EP 等级之间洗牌。默认值为False
。
moe_all_to_all_dispatcher
(布尔值)-是否在 MoE 中使用 all-to-all 调度器进行通信。默认值为True
。
moe_aux_loss_coeff
(浮动)-辅助负载平衡损失的系数。默认值为0.001
。
moe_z_loss_coeff
(浮动)-z 损失系数。默认值为0.001
。
torch.sagemaker.nn.attn.FlashSelfAttention
API适用于与 SMP v2 FlashAttention 一起使用。
class torch.sagemaker.nn.attn.FlashSelfAttention( attention_dropout_prob: float = 0.0, scale: Optional[float] = None, triton_flash_attention: bool = False, use_alibi: bool = False,
attention_dropout_prob
(浮动)-应用于注意力的辍学概率。默认值为0.0
。
scale
(float)— 如果通过,则此比例因子将应用于 softmax。如果设置为None
(这也是默认值),则比例因子为1 / sqrt(attention_head_size)
。默认值为None
。
triton_flash_attention
(bool) — 如果通过,则使用闪光注意的 Triton 实现。这是支持带线性偏差的注意力 (ALiBi) 所必需的(参见以下use_alibi
参数)。此版本的内核不支持掉线。默认值为False
。
use_alibi
(bool) — 如果通过,则使用提供的掩码启用带线性偏差的注意力 (ALiBi)。使用时ALiBi,需要准备好注意面罩,如下所示。默认值为False
。def generate_alibi_attn_mask(attention_mask, batch_size, seq_length, num_attention_heads, alibi_bias_max=8): device, dtype = attention_mask.device, attention_mask.dtype alibi_attention_mask = torch.zeros( 1, num_attention_heads, 1, seq_length, dtype=dtype, device=device alibi_bias = torch.arange(1 - seq_length, 1, dtype=dtype, device=device).view( 1, 1, 1, seq_length m = torch.arange(1, num_attention_heads + 1, dtype=dtype, device=device) m.mul_(alibi_bias_max / num_attention_heads) alibi_bias = alibi_bias * (1.0 / (2 ** m.view(1, num_attention_heads, 1, 1))) alibi_attention_mask.add_(alibi_bias) alibi_attention_mask = alibi_attention_mask[..., :seq_length, :seq_length] if attention_mask is not None and attention_mask.bool().any(): alibi_attention_mask.masked_fill( attention_mask.bool().view(batch_size, 1, 1, seq_length), float("-inf") return alibi_attention_mask
Methods
forward(self, qkv, attn_mask=None, causal=False, cast_dtype=None, layout="b h s d")
— 常规 PyTorch 模块功能。调用 amodule(x)
时,会自动SMP运行此函数。
qkv
— 采用以下torch.Tensor
形式:(batch_size x seqlen x (3 x num_heads) x head_size)
或者(batch_size, (3 x num_heads) x seqlen x head_size)
,torch.Tensors
每个元组的形状可能相同(batch_size x seqlen x num_heads x head_size)
,或(batch_size x num_heads x seqlen x head_size)
。必须根据形状传递适当的布局参数。
attn_mask
—torch.Tensor
采用以下形式(batch_size x 1 x 1 x seqlen)
。要启用此注意掩码参数,需要triton_flash_attention=True
和use_alibi=True
。要了解如何使用此方法生成注意力掩码,请参阅中的代码示例FlashAttention。默认值为None
。
causal
— 如果设置为False
(参数的默认值),则不应用掩码。设置为时True
,该forward
方法使用标准的下三角掩码。默认值为False
。
cast_dtype
— 当设置为特定时dtype
,它会将qkv
张量转换为dtype
之前attn
的张量。这对于诸如 Hugging Fac GPT e Transformer-NeoX 模型之类的实现非常有用,该模型q
具有旋转嵌入k
和后置嵌fp32
入。如果设置为None
,则不施放任何施法。默认值为None
。
layout
(字符串)-可用值为b h s d
或b s h d
。应将其设置为传递的qkv
张量的布局,以便可以应用适当的转换。attn
默认值为b h s d
。有形状
torch.Tensor
的单曲(batch_size x num_heads x seq_len x head_size)
。
torch.sagemaker.nn.attn.FlashGroupedQueryAttention
API适用于与 SMP v2
FlashGroupedQueryAttention
一起使用。要了解有关其用法的更多信息API,请参阅使用 FlashAttention 内核进行分组查询注意。class torch.sagemaker.nn.attn.FlashGroupedQueryAttention( attention_dropout_prob: float = 0.0, scale: Optional[float] = None,
attention_dropout_prob
(浮动)-应用于注意力的辍学概率。默认值为0.0
。
scale
(float)— 如果通过,则此比例因子将应用于 softmax。如果设置为None
,1 / sqrt(attention_head_size)
则用作比例因子。默认值为None
。Methods
forward(self, q, kv, causal=False, cast_dtype=None, layout="b s h d")
— 常规 PyTorch 模块功能。调用 amodule(x)
时,会自动SMP运行此函数。
q
— 以下torch.Tensor
形式(batch_size x seqlen x num_heads x head_size)
或(batch_size x num_heads x seqlen x head_size)
. 必须根据形状传递适当的布局参数。
kv
— 以下torch.Tensor
形式(batch_size x seqlen x (2 x num_heads) x head_size)
或(batch_size, (2 x num_heads) x seqlen x head_size)
,或者由两个torch.Tensor
s 组成的元组,每个元组的形状都可能为(batch_size x seqlen x num_heads x head_size)
或(batch_size x num_heads x seqlen x head_size)
。还必须根据形状传递适当的layout
参数。
causal
— 如果设置为False
(参数的默认值),则不应用掩码。设置为时True
,该forward
方法使用标准的下三角掩码。默认值为False
。
cast_dtype
— 当设置为特定的 dtype 时,它会将qkv
张量转换为之前的 dtype。attn
这对于诸如 Hugging Fac GPT e Transformers-Neox 之类的实现非常q,k
有用,后者fp32
带有旋转嵌入功能。如果设置为None
,则不施放任何施法。默认值为None
。布局(字符串)-可用值为
"b h s d"
或"b s h d"
。应将其设置为传递的qkv
张量的布局,以便可以应用适当的转换。attn
默认值为"b h s d"
。返回一个
torch.Tensor (batch_size x num_heads x seq_len x head_size)
表示注意力计算输出的单曲。
torch.sagemaker.nn.huggingface.llama_flashattn.LlamaFlashAttention
支持 API FlashAttention Llama 模型的。这API使用的是低torch.sagemaker.nn.attn.FlashGroupedQueryAttentionAPI级的。要了解如何使用它,请参阅使用 FlashAttention 内核进行分组查询注意。
class torch.sagemaker.nn.huggingface.llama_flashattn.LlamaFlashAttention( config: LlamaConfig
config
— 美洲驼模型的 FlashAttention 配置。Methods
forward(self, hidden_states, attention_mask, position_ids, past_key_value, output_attentions, use_cache)
hidden_states
(torch.Tensor
) — 张量的隐藏状态,形式(batch_size x seq_len x num_heads x head_size)
为。
attention_mask
(torch.LongTensor
) — 掩码以避免注意以的形式填充令牌索引。(batch_size x seqlen)
默认值为None
。
position_ids
(torch.LongTensor
) — 如果不是None
,则其形式为(batch_size x seqlen)
,表示位置嵌入中每个输入序列标记的位置索引。默认值为None
。
past_key_value
(缓存)— 预先计算的隐藏状态(自我注意力块和交叉注意力块中的键和值)。默认值为None
。
output_attentions
(bool)-指示是否返回所有注意力层的注意张量。默认值为False
。
use_cache
(bool)-指示是否返回past_key_values
键值状态。默认值为False
。返回一个
torch.Tensor (batch_size x num_heads x seq_len x head_size)
表示注意力计算输出的单曲。
torch.sagemaker.transform
SMPv2 提供了此功能,
torch.sagemaker.transform()
API用于将 Hugging Face Transformer SMP 模型转换为模型实现并启用张量并行性。SMPtorch.sagemaker.transform( model: nn.Module, device: Optional[torch.device] = None, dtype: Optional[torch.dtype] = None, config: Optional[Dict] = None, load_state_dict_from_rank0: bool = False
SMPv2 Hugging Face Transformer 模型与张量并SMP行度兼容 通过将 Hugging Face Transformer 模型的配置转换为变压器配置来维护的SMP转换策略。
model
(torch.nn.Module
) — 从模型转换Hugging Face Transformer 模型与张量并SMP行度兼容并应用库的张量并行度功能。SMP
device
(torch.device
) — 如果通过,则在此设备上创建新模型。如果原始模块在元设备上有任何参数(参见参数初始化延迟),则转换后的模块也将在元设备上创建,忽略此处传递的参数。默认值为None
。
dtype
(torch.dtype
) — 如果通过,则将其设置为用于创建模型的 dtype 上下文管理器,并使用此 dtype 创建模型。这通常是不必要的,因为我们想在使用fp32
时创建模型MixedPrecision
,并且fp32
是中的 PyTorch默认 dtype。默认值为None
。
config
(dict) — 这是一本用于配置SMP变压器的字典。默认值为None
。
load_state_dict_from_rank0
(Boolean)-默认情况下,此模块使用新的权重创建模型的新实例。当此参数设置为时True
,SMP会尝试将原始 PyTorch 模型的状态字典从第 0 个等级加载到第 0 个等级所属的张量 parallel 组的变换模型中。将其设置为时True
,等级 0 在元设备上不能有任何参数。在此变换调用之后,只有第一个张量 parallel 组填充第 0 个等级的权重。你需要在FSDP包装器True
中设置为sync_module_states
,才能将这些权重从第一个张量 parallel 组传递到所有其他进程。激活此功能后,SMP库会从原始模型加载状态字典。该SMP库在变换之前获取模型state_dict
的值,将其转换为与变换模型的结构相匹配,为每个张量平行等级对其进行分片,将这种状态从第 0 等级传达给第 0 等级所属的张量并行组中的其他等级,然后加载它。默认值为False
。返回一个你可以用它来包装的变换模型 PyTorch FSDP。设置
load_state_dict_from_rank0
为时True
,涉及等级 0 的张量 parallel 组会从等级 0 的原始状态字典加载权重。在原始模型参数初始化延迟上使用时,只有这些等级才会CPUs为变换后的模型的参数和缓冲区开启实际张量。其余队伍继续在元设备上使用参数和缓冲区以节省内存。
torch.sagemaker
util 函数和属性torch.sagemaker 实用程序函数
torch.sagemaker.init(config: Optional[Union[str, Dict[str, Any]]] = None) -> None
— 使用初始化 PyTorch 训练作业。SMP
torch.sagemaker.is_initialized() -> bool
— 检查训练作业是否使用初始化。SMP当使用初始化作业 PyTorch 时回退到原生版本时SMP,某些属性不相关并变成None
,如以下属性列表所示。
torch.sagemaker.utils.module_utils.empty_module_params(module: nn.Module, device: Optional[torch.device] = None, recurse: bool = False) -> nn.Module
— 在给定的(device
如果有)上创建空参数,如果指定,则可以对所有嵌套模块进行递归。
torch.sagemaker.utils.module_utils.move_buffers_to_device(module: nn.Module, device: torch.device, recurse: bool = False) -> nn.Module
— 将模块缓冲区移动到给定的device
,如果指定,则所有嵌套模块的缓冲区可以是递归的。
torch.sagemaker.state
在 with 的初始化后拥有多个有用的属性torch.sagemaker.init
。SMP
torch.sagemaker.state.hybrid_shard_degree
(int) — 分片数据并行度,来自传递给的SMP配置中用户输入的副本。torch.sagemaker.init()
要了解更多信息,请参阅 开始使用 SageMaker 模型并行度库 v2。
torch.sagemaker.state.rank
(int)-设备的全球排名,范围为[0, world_size)
。
torch.sagemaker.state.rep_rank_process_group
(torch.distributed.ProcessGroup
) — 进程组,包括所有具有相同复制等级的设备。请注意与的微妙但根本的区别torch.sagemaker.state.tp_process_group
。当回退到原生模式时 PyTorch,它会返回None
。
torch.sagemaker.state.tensor_parallel_degree
(int) — 张量并行度,来自传递给的SMP配置中用户输入的副本。torch.sagemaker.init()
要了解更多信息,请参阅 开始使用 SageMaker 模型并行度库 v2。
torch.sagemaker.state.tp_size
(int) — 的别名torch.sagemaker.state.tensor_parallel_degree
。
torch.sagemaker.state.tp_rank
(int) — 设备的张量并行度等级在范围内[0, tp_size)
,由张量并行度和排名机制确定。
torch.sagemaker.state.tp_process_group
(torch.distributed.ProcessGroup
) — 张量并行处理组,包括在其他维度上具有相同等级(例如,分片数据并行性和复制)但张量并行等级唯一的所有设备。当回退到原生模式时 PyTorch,它会返回None
。
torch.sagemaker.state.world_size
(int)-训练中使用的设备总数。从 SMP v1 升级到 SMP v2
要从 SMP v1 移动到 SMP v2,必须更改脚本以删除 SMP v1 APIs 并应用 v2。SMP APIs我们建议您不要从 SMP v1 脚本开始,而是从 PyTorch FSDP脚本开始,然后按照中的说明进行开始使用 SageMaker 模型并行度库 v2操作。
要将 SMP v1 模型引入 SMP v2,在 SMP v1 中,您必须收集完整的模型状态字典并在模型状态字典上应用翻译函数,将其转换为 Hugging Face Transformers 模型检查点格式。然后在 SMP v2 中,如中所述使用检查点功能 SMP,你可以加载 Hugging Face Transformers 模型检查点,然后在 v2 中继续使用该 PyTorch 检查点APIs。SMP要SMP与 PyTorch FSDP模型一起使用,请务必移至 SMP v2 并更改要使用的训练脚本 PyTorch FSDP和其他最新功能。
import smdistributed.modelparallel.torch as smp # Create model model = ... model = smp.DistributedModel(model) # Run training # Save v1 full checkpoint if smp.rdp_rank() == 0: model_dict = model.state_dict(gather_to_rank0=True) # save the full model # Get the corresponding translation function in smp v1 and translate if model_type == "gpt_neox": from smdistributed.modelparallel.torch.nn.huggingface.gptneox import translate_state_dict_to_hf_gptneox translated_state_dict = translate_state_dict_to_hf_gptneox(state_dict, max_seq_len=None) # Save the checkpoint checkpoint_path = "checkpoint.pt" if smp.rank() == 0: smp.save( {"model_state_dict": translated_state_dict},