Skip to content
目录

PreTrainedModel、生成与 KV Cache

1. PreTrainedModel 基类

models/modeling_utils.jsPreTrainedModel 继承 Callable(可像函数一样调用),构造时:

  • 保存 configsessions(多 ONNX 会话字典)、configs(如 generation_config.json);
  • 通过 resolveTypeConfig 合并 运行时行为(能否 generateforward 实现、prepare_inputs)与 session 配置MODEL_SESSION_CONFIG)。
202:233:packages/transformers/src/models/modeling_utils.js
export class PreTrainedModel extends Callable {
    main_input_name = 'input_ids';
    forward_params = ['input_ids', 'attention_mask'];
    ...
    constructor(config, sessions, configs) {
        super();

        this.config = config;
        this.sessions = sessions;
        this.configs = configs;

        const modelName = MODEL_CLASS_TO_NAME_MAPPING.get(this.constructor);
        const { typeConfig } = resolveTypeConfig(modelName, config);

        this.can_generate = typeConfig.can_generate;
        this._forward = typeConfig.forward;
        this._prepare_inputs_for_generation = typeConfig.prepare_inputs;

resolveTypeConfig 的亮点:处理 交叉架构加载(例如用 ForCausalLM 类去加载原生为 ForConditionalGeneration 的权重时),切换到原生架构的 type,并打 textOnly 标记,以便 session 与 forward 路径一致(见源码中 nativeArch 分支)。

2. 资源释放:dispose()

对每个 session 调用 release?.(),注释提到未来可用 FinalizationRegistry —— 对应 WebGPU/WASM 显存与 native handle 的释放需求。

3. 生成逻辑(高层)

库在 JS 侧复现了 Hugging Face generate大量行为

  • GenerationConfig
  • LogitsProcessorList / LogitsWarper(温度、重复惩罚、bad words 等);
  • StoppingCriteriaList
  • 与 ORT 输出对接的 sessionRun

具体实现较长,阅读时可从 PreTrainedModel.prototype.generatemodeling_utils.jsdecoder_forward / seq2seq_forward 入手。

4. DynamicCache(KV Cache)

cache_utils.js动态对象 存放 past_key_values.* 张量;提供 get_seq_length()dispose()(释放 location === 'gpu-buffer' 的张量)。

1:62:packages/transformers/src/cache_utils.js
class _DynamicCache {
    ...
    get_seq_length() {
        const self = /** @type {any} */ (this);
        for (const name in self) {
            if (name.startsWith('past_key_values.')) {
                return self[name].dims.at(-2);
            }
        }
        throw new Error('Unable to determine sequence length from the cache.');
    }
    ...
    async dispose() {
        const promises = [];
        for (const t of /** @type {Tensor[]} */ (Object.values(this))) {
            if (t.location === 'gpu-buffer') {
                promises.push(t.dispose());
            }
        }
        await Promise.all(promises);
    }
}

要点:为何用松散 key 而不是固定数组?—— ONNX 导出 各模型 past 名称与层数不同,字典更稳妥;GPU 张量需显式 dispose 避免 WebGPU OOM。

5. 与 Python 的差异提示

  • 算子支持:部分 logits 处理在 Python 可能在模型外;JS 需与导出的 ONNX 输入输出契约严格一致;
  • 精度:量化模型下 KV dtype 可能与权重不同,需读 transformers.js_config.kv_cache_dtype

下一篇:多 ONNX Session、多模态与 text-only 加载

MIT License.