CollieModelForCausalLM

class CollieModelForCausalLM(config)

CoLLiE 的基础模型。如果要实现新的模型，必须继承该基类。

CollieModelForCausalLM 统一了非流水线模型和流水线模型的接口，并且可以执行 生成任务。

为了适应流水线线的生成过程，每个新模型除了实现基类中的抽象方法外，还需要满足：

1. 每一个 layer 包含一个 use_cache 属性来决定前向传播中是否使用 past_key_values。

2. 将 Attention 过程中产生的 key 和 value 保存在每一个 layer 的 past_key_value 属性中。

3. 将每层的 hidden_states 保存在每一个 layer 的 hidden_states 属性中。

4. 将 lm_head 的输入 hidden_states 保存在 hidden_states 属性中。也可以使 用 ColumnParallelLMHead 来自动地保存。

main_input_name = 'input_ids'

base_model_prefix = ''

generate(*args, **kwargs)

生成函数。用法同 huggingface。

classmethod from_config(config, **kwargs)

从 config 中加载一个模型。

参数:

config (Union[CollieConfig, str]) – 接受一个字符串或 CollieConfig。为字符串时，会先从 该 str 代表的路径或远程链接中加载 config，再进行初始化

abstract clean_cache()

清理 use_cache 和 hidden_states 状态的函数。

abstract set_cache(use_cache)

设置 use_cache 的函数。

参数:

use_cache – 是否在生成时使用缓存的 past_key_values。如果为 True 则会保存前向传播过程中 Attention 的 key 和 value 用于下一次 生成。可以参考 _set_use_cache() 的代码来设置。

classmethod from_pretrained(model_path_or_name, config, **kwargs)

从 model_path_or_name 中加载预训练好的模型。

参数:

• model_path_or_name (str) – huggingface 格式预训练模型的本地路径或名 称。

• config (Union[CollieConfig, str]) – 可以是字符串或者 CollieConfig。如果为字符串，则会 使用该字符串代表的模型设置；如果为 None，从 model_path_or_name 中加载模型设置。

• kwargs –

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

  其余 kwargs 的内容会用于设置 CollieConfig 的内容。

classmethod pipeline_layers(config)

获取流水线模型。

返回:

一个列表，包含一系列层；这些模型会用于初始化流水线并行模型并且 进行划分。

abstract prepare_inputs(input_ids, attention_mask=None, use_cache=None, past_key_values=None, **kwargs)

在生成过程中更新 input_ids、attention_mask 等输入参数的函数。

prepare_inputs_for_generation(input_ids, attention_mask=None, use_cache=None, past_key_values=None, **kwargs)

生成过程中更新输入和 cache 状态的函数，包含设置 use_cache 和 past_key_values 以及更新输入两个过程。

abstract static load_parallel_state_dict(path, config, process_exclusion=False, **kwargs)

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

参数:

• path (str) –

• config (Union[CollieConfig, str]) –

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

返回:

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

abstract static save_parallel_state_dict(state_dict, path, config, process_exclusion=False, **kwargs)

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

参数:

• state_dict (dict) – 模型权重

• path (str) –

• config (CollieConfig) –

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

enable_input_require_grads()

Enables the gradients for the input embeddings. This is useful for fine-tuning adapter weights while keeping the model weights fixed.

training: bool