AWS Trainium 和 Inferentia 文档
Optimum Neuron 分布式
并获得增强的文档体验
开始使用
Optimum Neuron 分布式
optimum.neuron.distributed 模块提供了一组用于执行分布式训练和推理的工具。
并行化
选择模型特定的并行器类
每个在 optimum-neuron 中支持并行化的模型都有其自己的派生 Parallelizer 类。工厂类 ParallelizersManager 允许您轻松检索此类模型特定的 Parallelizer。
提供支持并行化的模型类型列表。
is_model_supported
< source >( model_type_or_model: typing.Union[str, transformers.modeling_utils.PreTrainedModel, optimum.neuron.utils.peft_utils.NeuronPeftModel] )
返回一个包含 3 个布尔值的元组，其中
- 第一个元素指示张量并行是否可用于此模型，
- 第二个元素指示是否可以在此模型上在张量并行之上使用序列并行，
- 第三个元素指示管道并行是否可用于此模型。
parallelizer_for_model
< source >( model_type_or_model: typing.Union[str, transformers.modeling_utils.PreTrainedModel, optimum.neuron.utils.peft_utils.NeuronPeftModel] )
返回与模型关联的并行器类。
实用工具
懒加载
当模型太大而无法装入单个设备时，通常需要分布式训练/推理。因此，需要允许延迟加载模型权重和优化器状态的工具，以避免在并行化之前内存不足。
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，则不执行延迟加载。
optimum.neuron.distributed.make_optimizer_constructor_lazy
< source >( optimizer_cls: typing.Type[ForwardRef('torch.optim.Optimizer')] )
转换优化器构造函数（优化器类），使其通过不初始化参数而变为延迟加载。这使得优化器轻量级且可用于在模型并行化后创建“真正的”优化器。