Shortcuts

Env

class Env[源代码]

CoLLiE 的环境变量,可以从中获取各种并行的 world_size 和 rank。

调用时直接导入已经实例化好的对象 env 即可。

from collie.utils import env
print(env.dp_rank)
property seed

随机数种子

property rank

Global rank。

property local_rank

Local rank。

static barrier(group=None)[源代码]

group 上进行同步。

property world_size

分布式训练的 world size。

property pp_rank

流水线并行的 rank。

property dp_rank

数据并行的 rank。

property tp_rank

张量并行的 rank。

property mp_rank

模型并行的 rank。模型并行与数据并行相对,同时包含了张量并行和模型并行。

property is_pipeline

判断是否是流水线并行。

property pipeline_parts

返回流水线并行中模型切分的分界点,长度为 pp_size + 1

如果不存在流水线并行则返回 None

e.g. pp_size 为 2时,返回形如 [0, 13, 25] 的结果。

property pipeline_layers_idx

返回流水线并行中当前 rank 切分的模型索引。

如果不存在流水线并行返回 None

property tp_group

张量并行的通信组。

property pp_group

流水线并行的通信组。

property dp_group

数据并行的通信组。

property dp_size

数据并行的 world size。

property tp_size

张量并行的 world size。

property pp_size

流水线并行的 world size。

property is_last_stage

是否是流水线的最后一个阶段。

property is_first_stage

是否是流水线的第一个阶段。