Optimum Neuron 分布式
optimum.neuron.distributed 模块提供了一组用于执行分布式训练和推理的工具。
并行化
分布式训练/推理中的主要任务是能够对模型权重、梯度和/或优化器状态等进行分片。我们构建了 Parallelizer 类来处理分片。
基础并行化器
Parallelizer 类是支持模型并行的每个模型的基类抽象类。它提供了一些方法来并行化模型，以及保存和加载分片检查点。
处理模型并行的基类抽象类。
_parallelize
< 来源 >( model: PreTrainedModel device: Optional = None parallelize_embeddings: bool = True sequence_parallel_enabled: bool = False should_parallelize_layer_predicate_func: Optional = None **parallel_layer_specific_kwargs ) → PreTrainedModel
参数
- model (
PreTrainedModel) — 需要并行的模型。 - device (
Optional[torch.device], 默认为None) — 新的并行层应放置的设备。 - parallelize_embeddings (
bool, 默认为True) — 是否应并行化嵌入。当 TP 大小无法整除词汇量大小时，可以禁用此功能。 - sequence_parallel_enabled (
bool, 默认为False) — 是否启用序列并行。 - should_parallelize_layer_predicate_func (Optional[Callable[[torch.nn.Module], bool]], 默认为
None) — 一个函数，接受一个层作为输入，并返回一个布尔值，指定是否应该并行化输入层。这在例如管道并行的情况下用于跳过不必要的并行化。 - **parallel_layer_specific_kwargs (
Dict[str, Any]) — 特定于某些并行层的关键字参数，它们将被其他并行层忽略。
返回
PreTrainedModel
并行的模型。
通过将普通层转换为其并行对应层来并行化模型。每个具体类必须实现它。
parallelize
< 源代码 >( model: Union device: Optional = None parallelize_embeddings: bool = True sequence_parallel_enabled: bool = False kv_size_multiplier: Optional = None pipeline_parallel_input_names: Union = None pipeline_parallel_num_microbatches: int = 1 pipeline_parallel_use_zero1_optimizer: bool = False pipeline_parallel_gradient_checkpointing_enabled: bool = False checkpoint_dir: Union = None num_local_ranks_per_step: int = 8 ) → PreTrainedModel
参数
- model (
Union[PreTrainedModel, NeuronPeftModel]) — 需要并行的模型。 - sequence_parallel_enabled (
bool, 默认值为False) — 是否启用序列并行。 - kv_size_multiplier (
Optional[int], 默认值为None`) — 当 TP 大小大于 KV 头数时，复制 KV 头的次数。如果未指定，将使用使 KV 头数可以被 TP 大小整除的最小倍数。 - pipeline_parallel_num_microbatches (
int, 默认值为 1) — 用于管道执行的微批次数量。 - pipeline_parallel_use_zero1_optimizer (
bool, 默认值为False) — 当使用 zero-1 优化器时，将其设置为 True，以便 PP 模型了解 zero-1 优化器将处理数据并行梯度平均。 - pipeline_parallel_gradient_checkpointing_enabled (
bool, 默认值为False) — 是否在进行管道并行时启用梯度检查点。 - checkpoint_dir (
Optional[Union[str, Path]]) — 分片检查点的路径。如果指定，检查点权重将加载到并行化的模型中。 - num_local_ranks_per_step (
int, 默认值为8) — 对应于可以同时初始化和加载模型权重的本地排名数量。如果该值小于 0，将使用最大排名数。
返回
PreTrainedModel
并行的模型。
使用 cls._parallelize() 将常规层转换为它们的并行对应层，从而并行化模型。
它还确保每个参数都已加载其权重，或者如果与之关联的预训练权重不存在，则已初始化。
optimizer_for_mp
< source >( optimizer: torch.optim.Optimizer orig_param_to_parallel_param_on_xla: Mapping ) → torch.optim.Optimizer
从现有优化器创建适用于并行化模型的优化器。
有两种情况
- 优化器是通过来自
optimum.neuron.distributed.utils.make_optimizer_constructor_lazy的延迟构造函数创建的，在这种情况下，将为张量并行创建完全预期的优化器。 - 优化器是使用常规构造函数创建的。在这种情况下，为张量并行创建的优化器尽可能接近预期，但不保证完全一致。
save_model_sharded_checkpoint
< source >( model: Union output_dir: Union optimizer: Optional = None use_xser: bool = True async_save: bool = False num_local_ranks_per_step: int = 8 )
选择模型特定的并行化器类
每个在 optimum-neuron 中支持并行化的模型都有其自己的派生 Parallelizer 类。工厂类 ParallelizersManager 允许您轻松地检索这些模型特定的 Parallelizer。
提供支持并行化的模型类型的列表。
is_model_supported
< source >( model_type_or_model: Union )
返回一个包含 3 个布尔值的元组，其中
- 第一个元素指示该模型是否可以使用张量并行，
- 第二个元素指示该模型是否可以在张量并行之上使用序列并行，
- 第三个元素指示该模型是否可以使用管道并行。
Utils
Lazy Loading
当模型太大而无法放入一个设备中时，通常需要分布式训练/推理。因此，需要允许模型权重和优化器状态延迟加载的工具，以避免在并行化之前出现内存不足错误。
optimum.neuron.distributed.lazy_load_for_parallelism
< source >( tensor_parallel_size: int = 1 pipeline_parallel_size: int = 1 )
上下文管理器，使模型并行化时模型的加载延迟。
- 每个
torch.nn.Linear都被放在torch.device("meta")设备上，这意味着实例化它不会占用内存。 - 每个
torch.nn.Embedding也被放在torch.device("meta")设备上。 - 没有真正加载状态字典，而是创建了一个权重映射并将其附加到模型。有关更多信息，请阅读
optimum.neuron.distributed.utils.from_pretrained_for_mp的文档字符串。
如果 tensor_parallel_size 和 pipeline_parallel_size 都设置为 1，则不会执行延迟加载。
将优化器构造函数（优化器类）转换为延迟函数，方法是不初始化参数。这使得优化器轻量级，并且可以在模型并行化后用来创建一个“真实”优化器。