mindspore.set_auto_parallel_context

mindspore.set_auto_parallel_context(**kwargs)[源代码]

配置自动并行,当前CPU仅支持数据并行。

应在mindspore.communication.init之前配置自动并行。

Note

配置时,必须输入配置的名称。如果某个程序具有不同并行模式下的任务,需要提前调用reset_auto_parallel_context()为下一个任务设置新的并行模式。若要设置或更改并行模式,必须在创建任何Initializer之前调用接口,否则,在编译网络时,可能会出现RuntimeError。

某些配置适用于特定的并行模式,有关详细信息,请参见下表:

Common

AUTO_PARALLEL

device_num

gradient_fp32_sync

global_rank

loss_repeated_mean

gradients_mean

auto_parallel_search_mode

parallel_mode

strategy_ckpt_load_file

all_reduce_fusion_config

strategy_ckpt_save_file

enable_parallel_optimizer

dataset_strategy

enable_alltoall

pipeline_stages

grad_accumulation_step

参数:

  • device_num (int) - 表示可用设备的编号,必须在[1,4096]范围中。默认值:1。

  • global_rank (int) - 表示全局RANK的ID,必须在[0,4095]范围中。默认值:0。

  • gradients_mean (bool) - 表示是否在梯度的 AllReduce后执行平均算子。stand_alone不支持gradients_mean。默认值:False。

  • gradient_fp32_sync (bool):在FP32中运行梯度的 AllReduce。stand_alone、data_parallel和hybrid_parallel不支持gradient_fp32_sync。默认值:True。

  • parallel_mode (str) - 有五种并行模式,分别是stand_alone、data_parallel、hybrid_parallel、semi_auto_parallel和auto_parallel。默认值:stand_alone。

    • stand_alone:单卡模式。

    • data_parallel:数据并行模式。

    • hybrid_parallel:手动实现数据并行和模型并行。

    • semi_auto_parallel:半自动并行模式。

    • auto_parallel:自动并行模式。

  • search_mode (str) - 表示有三种策略搜索模式,分别是recursive_programming,dynamic_programming和sharding_propagation。默认值:dynamic_programming。

    • recursive_programming:表示双递归搜索模式。

    • dynamic_programming:表示动态规划搜索模式。

    • sharding_propagation:表示从已配置算子的切分策略传播到所有算子。

  • auto_parallel_search_mode (str) - search_mode参数的兼容接口。将在后续的版本中删除。

  • parameter_broadcast (bool) - 表示在训练前是否广播参数。在训练之前,为了使所有设备的网络初始化参数值相同,请将设备0上的参数广播到其他设备。不同并行模式下的参数广播不同。在data_parallel模式下,除layerwise_parallel属性为True的参数外,所有参数都会被广播。在hybrid_parallel、semi_auto_parallel和auto_parallel模式下,分段参数不参与广播。默认值:False。

  • strategy_ckpt_load_file (str) - 表示用于加载并行策略checkpoint的路径。默认值:’’。

  • strategy_ckpt_save_file (str) - 表示用于保存并行策略checkpoint的路径。默认值:’’。

  • full_batch (bool) - 如果在auto_parallel模式下加载整个batch数据集,则此参数应设置为True。默认值:False。目前不建议使用该接口,建议使用dataset_strategy来替换它。

  • dataset_strategy (Union[str, tuple]) - 表示数据集分片策略。默认值:data_parallel。dataset_strategy=”data_parallel”等于full_batch=False,dataset_strategy=”full_batch”等于full_batch=True。对于通过模型并列策略加载到网络的数据集,如ds_stra ((1, 8)、(1, 8)),需要使用set_auto_parallel_context(dataset_strategy=ds_stra)。

  • enable_parallel_optimizer (bool) - 这是一个开发中的特性,它可以为数据并行训练对权重更新计算进行分片,以节省时间和内存。目前,自动和半自动并行模式支持Ascend和GPU中的所有优化器。数据并行模式仅支持Ascend中的 LambAdamWeightDecay 。默认值:False。

  • enable_alltoall (bool) - 允许在通信期间生成 AllToAll 通信算子的开关。 如果其值为 False,则将由 AllGatherSplitConcat 等通信算子的组合来代替 AllToAll 。 默认值:False。

  • all_reduce_fusion_config (list) - 通过参数索引设置 AllReduce 融合策略。仅支持ReduceOp.SUM和HCCL_WORLD_GROUP/NCCL_WORLD_GROUP。没有默认值。如果不设置,则关闭算子融合。

  • pipeline_stages (int) - 设置pipeline并行的阶段信息。这表明了设备如何单独分布在pipeline上。所有的设备将被划分为pipeline_stags个阶段。目前,这只能在启动semi_auto_parallel模式的情况下使用。默认值:1。

  • grad_accumulation_step (int) - 在自动和半自动并行模式下设置梯度的累积step。其值应为正整数。默认值:1。

  • parallel_optimizer_config (dict) - 用于开启优化器并行后的行为配置。仅在enable_parallel_optimizer=True的时候生效。目前,它支持关键字如下的关键字:

    • gradient_accumulation_shard(bool):设置累积梯度变量是否在数据并行维度上进行切分。开启后,将进一步减小模型的显存占用,但是会在反向计算梯度时引入额外的通信算子(ReduceScatter)。此配置仅在流水线并行训练和梯度累积模式下生效。默认值:True。

    • parallel_optimizer_threshold(int):设置参数切分的阈值。占用内存小于该阈值的参数不做切分。占用内存大小 = shape[0] * … * shape[n] * size(dtype)。该阈值非负。单位: KB。默认值:64。

  • comm_fusion (dict) - 用于设置通信算子的融合配置。可以同一类型的通信算子按梯度张量的大小或者顺序分块传输。输入格式为{“通信类型”: {“mode”:str, “config”: None int 或者 list}},每种通信算子的融合配置有两个键:”mode”和”config”。支持以下通信类型的融合类型和配置:

    • allreduce: 进行AllReduce算子的通信融合。”mode”包含:”auto”、”size”和”index”。在”auto”模式下,融合的是梯度变量的大小,默认值阈值为”64”MB,”config”对应的值为None。在”size”模式下,需要用户在config的字典中指定梯度大小阈值,这个值必须大于”0”MB。在”mode”为”index”时,它与”all_reduce_fusion_config”相同,用户需要给”config”传入一个列表,里面每个值表示梯度的索引。

    • allgather: 进行AllGather算子的通信融合。”mode”包含:”auto”、”size”。”auto” 和 “size”模式的配置方式与AllReduce相同。

    • reducescatter: 进行ReduceScatter算子的通信融合。”mode”包含:”auto”、”size”。”auto” 和 “size”模式的配置方式与AllReduce相同。

异常:

ValueError:输入key不是自动并行上下文中的属性。

样例:

>>> import mindspore as ms
>>> ms.set_auto_parallel_context(device_num=8)
>>> ms.set_auto_parallel_context(global_rank=0)
>>> ms.set_auto_parallel_context(gradients_mean=True)
>>> ms.set_auto_parallel_context(gradient_fp32_sync=False)
>>> ms.set_auto_parallel_context(parallel_mode="auto_parallel")
>>> ms.set_auto_parallel_context(search_mode="dynamic_programming")
>>> ms.set_auto_parallel_context(auto_parallel_search_mode="dynamic_programming")
>>> ms.set_auto_parallel_context(parameter_broadcast=False)
>>> ms.set_auto_parallel_context(strategy_ckpt_load_file="./strategy_stage1.ckpt")
>>> ms.set_auto_parallel_context(strategy_ckpt_save_file="./strategy_stage1.ckpt")
>>> ms.set_auto_parallel_context(dataset_strategy=((1, 8), (1, 8)))
>>> ms.set_auto_parallel_context(enable_parallel_optimizer=False)
>>> ms.set_auto_parallel_context(enable_alltoall=False)
>>> ms.set_auto_parallel_context(all_reduce_fusion_config=[8, 160])
>>> ms.set_auto_parallel_context(pipeline_stages=2)
>>> parallel_config = {"gradient_accumulation_shard": True, "parallel_optimizer_threshold": 24}
>>> ms.set_auto_parallel_context(parallel_optimizer_config=parallel_config, enable_parallel_optimizer=True)
>>> config = {"allreduce": {"mode": "size", "config": 32}, "allgather": {"mode": "size", "config": 32}}
>>> ms.set_auto_parallel_context(comm_fusion=config)