添加链接
link管理
链接快照平台
  • 输入网页链接,自动生成快照
  • 标签化管理网页链接
SageMaker 模型并行库 v2 参考 - Amazon SageMaker

本文属于机器翻译版本。若本译文内容与英语原文存在差异,则一律以英文原文为准。

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

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默认情况下,会DelayedParamIniterAPI找出如何正确初始化给定模型。对于任何其他模型,您需要创建一个自定义参数初始化函数并将其添加到脚本中。以下代码片段是 SMP v2 为实现的默认init_method_using_config函数。Hugging Face Transformer 模型与张量并SMP行度兼容使用以下代码片段作为参考,您可以创建自己的初始化配置函数,将其添加到脚本中,然后将其传递给的init_method_using_configSMPDelayedParamIniterAPI参数。

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_idNone,则会引发异常错误。

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.transformAPI调用。要详细了解如何使用该课程训练 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_losssinkhornbalanced、和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 模块功能。调用 a module(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_masktorch.Tensor 采用以下形式(batch_size x 1 x 1 x seqlen)。要启用此注意掩码参数,需要triton_flash_attention=Trueuse_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 db 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。如果设置为None1 / sqrt(attention_head_size)则用作比例因子。默认值为 None

Methods

forward(self, q, kv, causal=False, cast_dtype=None, layout="b s h d")— 常规 PyTorch 模块功能。调用 a module(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 模型转换为模型实现并启用张量并行性。SMP

torch.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.sagemakerutil 函数和属性

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},