Moss003MoonForCausalLM

class Moss003MoonForCausalLM(config: CollieConfig, **kwargs)

支持 3D 并行的 Moss-moon 模型。

参数:

config – CollieConfig

base_model_prefix = 'transformer'

static load_parallel_state_dict(path, config, process_exclusion=False, protocol='file', **kwargs)

从 path 中加载模型权重。path 中的模型权重应当是 huggingface 格式。

参数:

• path (str) –

• config (Union[CollieConfig, str]) –

• process_exclusion (bool, default: False) – 是否每个 rank 各自独立、互斥地加载模型权重。在模 型规模较大时，该参数可以帮助节省内存。

返回:

一个字典，每个字典都包含当前 rank 上模型需要的权重。

training: bool

static save_parallel_state_dict(state_dict, path, config, process_exclusion=False, protocol='file')

将模型权重保存到 path 路径。保存的格式同 huggingface 格式。

在保存时会在 dp rank 0 上将所有张量并行的权重合并至 tp_rank 0，然后按照流水 线的各个阶段分别保存为 sharded checkpoint 的形式。

参数:

• state_dict (dict) – 模型权重

• path (str) –

• config (CollieConfig) –

• process_exclusion (bool, default: False) – 是否每个 rank 各自独立、互斥地保存模型权重。在模 型规模较大时，该参数可以帮助节省内存。