MLflow Transformers Flavor 完整实战指南HuggingFace Pipeline 的模型记录、加载与推理【免费下载链接】mlflowThe open source AI engineering platform for agents, LLMs, and ML models. MLflow enables teams of all sizes to debug, evaluate, monitor, and optimize production-quality AI applications while controlling costs and managing access to models and data.项目地址: https://gitcode.com/GitHub_Trending/ml/mlflow导读本文围绕 MLflow 仓库中的mlflow.transformers模块对应 API 参考文档 mlflow.transformers.rst核心实现在 mlflow/transformers/init.py系统讲解如何将 HuggingFace Transformers 生态下的Pipeline、组件字典乃至本地 checkpoint 完整地保存为 MLflow Model并支持后续以pipeline、components或pyfunc三种方式加载推理。读完本文你将掌握save_model/log_model/load_model的全部核心参数与用法、LLM v1 推理任务chat/completions/embeddings的 OpenAI 兼容输出、PEFT 微调模型的存储策略以及相关环境变量对设备与加速加载行为的影响。一、flavor 概览与支持的模型输入形态mlflow.transformers是 MLflow 为 HuggingFace/Transformer 生态提供的模型 flavor。其保存入口要求用户提交以下三种形式之一的transformers_model对应save_model/log_model的入参校验逻辑Pipeline实例直接传入transformers.pipeline(...)创建的对象例如pipeline(fill-mask, distilroberta-base)组件字典以{model: ..., tokenizer: ..., feature_extractor: ..., image_processor: ..., torch_dtype: ...}形式给出。其中model键必须指向继承自PreTrainedModelPyTorch、TFPreTrainedModelTensorFlow或FlaxPreTrainedModelFlax的模型其余组件须与模型对应的任务类型兼容。字典中不允许出现上述键之外的字段否则抛出INVALID_PARAMETER_VALUE异常见init.py 的_validate_transformers_model_dict本地/DBFS checkpoint 路径字符串指向包含config.json的模型 checkpoint 目录。此形态特别适合无法整体载入内存做序列化的大模型场景此时必须同时显式传入task参数且save_pretrained必须为True否则抛异常。保存时模块会校验模型未被以异构device_map分布式加载hf_device_map值多于一种设备即拒绝保存避免 accelerate 导致的权重写坏风险见init.py。依赖管理上get_default_pip_requirements会依据模型所属执行引擎torch/tensorflow/flax自动生成最小 pip 依赖transformers 引擎库PyTorch 引擎下额外加入torchvision与若安装了acceleratePEFT 模型则追加peftinit.py。二、保存模型save_model与log_model2.1 基本用法save_model将模型写入本地文件系统路径log_model则把模型记录为当前 run 的 artifact并支持registered_model_name注册到 Model Registry、await_registration_for等待版本就绪默认等待 5 分钟。两者共享几乎相同的参数集下面是最小示例from transformers import pipeline import mlflow fill_pipe pipeline(fill-mask, distilroberta-base) # 保存到本地路径 with mlflow.start_run(): mlflow.transformers.save_model( transformers_modelfill_pipe, pathpath/to/save/model, ) # 记录为 run artifact with mlflow.start_run(): mlflow.transformers.log_model( transformers_modelfill_pipe, namemodel, # 旧参数 artifact_path 已废弃请使用 name )组件级保存自定义 tokenizer 模型from transformers import AutoModelForMaskedLM, AutoTokenizer architecture distilroberta-base tokenizer AutoTokenizer.from_pretrained(architecture) model AutoModelForMaskedLM.from_pretrained(architecture) with mlflow.start_run(): components {model: model, tokenizer: tokenizer} mlflow.transformers.save_model( transformers_modelcomponents, pathpath/to/save/model, )本地 checkpoint 保存注意必须指定taskwith mlflow.start_run(): mlflow.transformers.save_model( transformers_modelpath/to/local/checkpoint, # 需含 config.json pathpath/to/save/model, )2.2 核心参数详解参数说明关键细节tasktransformers 任务类型或 MLflow LLM v1 推理任务传 transformers 任务如text-generation时按指定任务构建 pipeline不传时由 transformers 库推断。传llm/v1/chat、llm/v1/completions、llm/v1/embeddings时会强制使用对应 OpenAI 兼容签名。无效类型抛异常processor可选Processor子类实例多模态模型用于合并文本/图像/音频编码入口若提供 processor 保存模型将无法以Pipeline形式或 pyfunc 推理加载只能以return_typecomponents加载torch_dtype模型加载回内存时应用的 PyTorch dtype不指定则沿用模型当前 dtype用于以不同于训练时的精度保存模型model_card可选ModelCard实例来自huggingface_hub提供则随模型保存其内容不提供时尝试从基础预训练模型拉取。拉取要求安装huggingface_hub0.10.0code_paths自定义代码依赖路径列表随模型打包供加载时加入系统路径signature模型的输入输出ModelSignature可用mlflow.models.signature.infer_signature推断若提供了input_example而未给签名NLP 文本类 pipeline 会自动推断签名非文本类会告警建议显式提供input_example模型输入示例既用于签名推断也会被序列化保存_save_examplepip_requirements/extra_pip_requirements覆盖或追加 pip 依赖不指定时会通过推理带超时推断依赖并与默认依赖取并集checkpoint 路径与 PEFT 模型不做依赖推理建议显式提供conda_env自定义 conda 环境默认基于get_default_conda_env生成metadata附加元数据字典与task为 LLM v1 任务时的一致性校验有关见下文model_config推理期 pipeline 覆盖参数仅作用于 pyfunc/Spark 加载不影响load_model返回的 Pipeline例如生成类模型的top_k、num_beams、max_length、temperature、top_p、repetition_penalty。键与 pipeline/model 不兼容会在运行期抛异常保存前务必先验证prompt_template提示词模板字符串仅支持feature-extraction、fill-mask、summarization、text2text-generation、text-generation任务模板必须只含一个格式占位符{prompt}如Answer the following question concisely.\n\nQ: {prompt}\nA:否则抛INVALID_PARAMETER_VALUEtext-generation 任务带模板时默认把return_full_text设为False可用model_config覆盖save_pretrained是否将预训练权重落盘为False时仅记录 HuggingFace Hub 仓库引用与 commitsource_model_name/source_model_revision加载时从 Hub 拉取模型名不是合法 HF repo id 时会强制回退为落盘base_model_pathPEFT 模型专用本地基础模型权重目录路径设置后只保存 PEFT adapter 权重base 权重以路径引用而非落盘适用于离线/气隙环境与避免大权重重复存储目录必须含config.json且仅对 PEFT 模型有效prompts/params/tags/model_type/step/model_idlog_model的模型注册与版本管理参数由Model.log统一处理model_config的推荐实践先在内存中对 pipeline 做一次推理验证确认无异常后再保存from transformers import pipeline, AutoTokenizer task text-generation architecture gpt2 sentence_pipeline pipeline( tasktask, tokenizerAutoTokenizer.from_pretrained(architecture), modelarchitecture, ) prompts [Generative models are, Id like a coconut so that I can] model_config { top_k: 2, num_beams: 5, max_length: 30, temperature: 0.62, top_p: 0.85, repetition_penalty: 1.15, } # 先验证覆盖参数可用 sentence_pipeline(prompts, **model_config) with mlflow.start_run(): mlflow.transformers.log_model( transformers_modelsentence_pipeline, namemy_sentence_generator, tasktask, model_configmodel_config, )2.3 签名自动推断与 LLM v1 任务签名当task以llm/v1前缀开头时MLflow 会忽略用户显式传入的 signature若传入且与推断值不一致会抛异常直接套用标准的 OpenAI 兼容签名CHAT_MODEL_INPUT_SCHEMA、COMPLETIONS_MODEL_INPUT_SCHEMA、EMBEDDING_MODEL_INPUT_SCHEMA等见 llm_inference_utils.py。同时会把推理任务写入 flavor 配置的inference_task键以及 metadata 的task键若 metadata 中的任务与task参数不一致会拒绝保存init.py。对于一般 NLP pipelineinput_example会触发自动签名推断infer_or_get_default_signature。历史上generate_signature_output曾用于手工生成签名输出该函数自 3.11.0 起标记为 deprecated应改用input_example自动推断。三、加载模型load_model与三种返回形态import mlflow model_uri fruns:/{run.info.run_id}/english_to_german_translator loaded mlflow.pyfunc.load_model(model_uri) print(loaded.predict(data)) # MLflow ist großartig!3.1 参数model_uri模型位置支持本地路径、runs:/run_id/path、s3://...、mlflow-artifacts:/...等 URI 形式dst_path模型 artifact 的下载目录须已存在缺省自动创建临时目录return_typepipeline默认或componentspipeline返回由 task 决定的Pipeline对象模型 tokenizer/feature extractor/image processor 等组件。components返回组件字典例如 NLP 模型{model: BertForSequenceClassification, tokenizer: BertTokenizerFast}视觉模型返回ImageProcessor多模态模型返回FeatureExtractorTokenizer。对使用 processor 保存的模型这是唯一可行的加载方式以pipeline加载会抛BAD_REQUEST。device加载目标设备None默认、0表示默认 GPUbase_model_path加载期覆盖保存时的本地基础模型路径适用于不同环境挂载点不一致的场景仅对以base_model_path保存的模型有效kwargs透传给 transformers 加载过程的附加配置。3.2 底层加载路径从源码看model_io.py加载按保存方式分三条路径本地权重默认从 artifact 目录下的model_MODEL_BINARY_FILE_NAME与components/*目录加载模型与各组件兼容 2.4.1 之前以pipeline/*路径保存的旧格式Hub 引用save_pretrainedFalse按 flavor 配置中的source_model_revisioncommit 从 HuggingFace Hub 加载若保存时未记录 commit hash会抛INVALID_STATE异常以规避版本漂移与安全风险本地 base 路径PEFT base_model_path基础模型从外部路径加载组件从 artifact 的components目录加载。模型类解析支持自定义代码trust_remote_code优先使用 transformers 内置类名否则读取config.json的architectures与auto_map字段映射到对应 AutoClass 并启用远程代码信任model_io.py。四、pyfunc 推理输入解析与输出归一化语言类模型在保存时会自动附加 pyfunc flavorloader_modulemlflow.transformers从而获得mlflow.pyfunc.load_model、模型服务serving与 Spark UDF 能力。以下 pipeline 类型不会附加 pyfunc图像/视频类如ImageToTextPipeline、ObjectDetectionPipeline、VideoClassificationPipeline等、零样本音频分类以及 Graphormer/Informer/TimeSeriesTransformer/DecisionTransformer 等特殊模型init.py。_TransformersWrapper.predict支持的输入类型pandas.DataFrame、str、bytes、np.ndarray、List[str]、List[Dict[str, str]]、List[Dict[str, Union[str, List[str]]]]、Dict[str, Union[str, List[str]]]。内部按 pipeline 类型做针对性解析例如文本分类允许text/text_pair键的字典输入serving 场景下对 JSON 编码的字典用ast.literal_eval还原避免引号转义问题问答要求questioncontext键表问答要求table键内部把 pandas DataFrame 序列化为 JSON 字符串后喂给 pipeline零样本分类对candidate_labels做 JSON 解码输出扁平化为sequence / labels / scores的 DataFrametoken 分类NER输出被折叠为逗号分隔字符串如PRON,NOUN图像/音频支持 base64 编码数据自动解码字符串输入会校验为合法 URI 或本地文件路径避免晦涩的底层报错。推理期model_config还有两个可选控制项include_prompt默认True设为False时从生成结果中裁掉输入 prompt 前导与collapse_whitespace默认False压缩多余空白。五、LLM v1 推理任务OpenAI 兼容的 chat / completions / embeddings以taskllm/v1/chat、llm/v1/completions、llm/v1/embeddings保存模型后pyfunc 推理输出将变为 OpenAI 兼容格式见 llm_inference_utils.py。任务与 transformers pipeline 的对应关系如下LLM v1 任务底层 pipeline 任务输入键llm/v1/chattext-generationmessages消息列表经 tokenizer 的apply_chat_template转成 promptllm/v1/completionstext-generationpromptllm/v1/embeddingsfeature-extractioninput支持单个字符串或字符串列表的异构输入关键转换逻辑参数映射OpenAI 风格的max_tokens自动映射为 HuggingFace 的max_new_tokensstop序列通过 slow tokenizer 转换为StoppingCriteria列表含有无前导空格的两种匹配见_get_stopping_criteria输出包装completions返回{id, object: text_completion, created, model, choices: [{index, finish_reason, text}], usage: {prompt_tokens, completion_tokens, total_tokens}}chat的 choice 中为{role: assistant, content: ...}object为chat.completionembeddings返回{object: list, data: [{object: embedding, index, embedding}], usage: ...}finish_reason依据max_length/max_new_tokens判定stop或length签名自动套用 OpenAI 兼容输入/输出 schema推理时从输入载荷中剥离参数字段如temperature、max_tokens等并允许它们覆盖保存时的model_config。六、PEFT 模型与本地基础模型路径对 LoRA 等 PEFT 微调模型PeftModelMLflow 遵循 transformers 行为强制save_pretrainedFalse仅保存 adapter 权重到peft目录基础模型以 HuggingFace Hub 仓库引用记录。加载时用PeftModel.from_pretrained重新构建 PEFT 模型保证返回对象与保存时一致保留peft_config属性peft.py。若希望不依赖网络气隙环境或避免大权重重复存储使用base_model_path指向本地基础模型目录from peft import get_peft_model, LoraConfig from transformers import AutoModelForCausalLM, AutoTokenizer base_path /shared/models/llama-7b base_model AutoModelForCausalLM.from_pretrained(base_path) tokenizer AutoTokenizer.from_pretrained(base_path) peft_model get_peft_model(base_model, LoraConfig(...)) with mlflow.start_run(): mlflow.transformers.save_model( transformers_model{model: peft_model, tokenizer: tokenizer}, pathpath/to/save, base_model_pathbase_path, )保存时仅落盘 tokenizer 等组件flavor 配置写入local_base_model_path加载时若路径与保存时不一致可在load_model(..., base_model_path...)中覆盖。校验规则base_model_path必须是含config.json的目录且只允许配合 PEFT 模型使用。七、persist_pretrained_model补全 Hub 引用模型的权重以save_pretrainedFalse保存的模型因 artifact 中缺少完整权重无法注册到 Databricks Workspace Model Registry。persist_pretrained_model(model_uri)用于事后补全它从 HuggingFace Hub 下载权重写入 artifact并把 flavor 配置从source_model_revision引用模式更新为本地model_binary模式同时清理组件的 repo/commit 字段flavor_config.py。import mlflow with mlflow.start_run() as run: model pipeline(fill-mask, distilroberta-base) mlflow.transformers.log_model( transformers_modelmodel, namepipeline, save_pretrainedFalse ) # 补全权重前无法注册 mlflow.transformers.persist_pretrained_model(fruns:/{run.info.run_id}/pipeline) # 补全后即可注册 mlflow.register_model(fruns:/{run.info.run_id}/pipeline, fill_mask_pipeline)若权重已存在于 artifact 目录则直接跳过不做重复下载。八、autolog与相关环境变量8.1 autolog 的特殊定位mlflow.transformers.autolog()并不实现完整的 transformers 自动记录其作用在于抑制训练/评估 transformers 模型时对无关子模型的噪音式 autolog它禁用 sklearn、tensorflow、pytorch 的离散 autolog并安全 patchsetfit.Trainer/transformers.Trainer/Seq2SeqTrainer的train方法在 transformers 5.x 下还会确保MLflowCallback被注册使指标与超参数正常写入 MLflowinit.py。参数包括log_input_examples、log_model_signatures、log_models、log_datasets、disable、exclusive、disable_for_unsupported_versions、silent、extra_tags。8.2 设备与加载相关环境变量以下环境变量直接改变模型的加载与推理行为定义见 environment_variables.py消费逻辑见init.py环境变量默认值作用MLFLOW_DEFAULT_PREDICTION_DEVICE无默认预测设备整数字符串非法值时回退到 CPU-1若机器有 GPU 且未设置自动选 GPU0MLFLOW_HUGGINGFACE_USE_DEVICE_MAPFalse是否使用 accelerate 的device_map自动分配设备如多卡开启后与device参数互斥同时提供会抛异常MLFLOW_HUGGINGFACE_DEVICE_MAP_STRATEGYautodevice_map的策略值如auto、balanced等MLFLOW_HUGGINGFACE_USE_LOW_CPU_MEM_USAGETrue加载时是否启用low_cpu_mem_usage降低 CPU 内存峰值MLFLOW_HUGGINGFACE_DISABLE_ACCELERATE_FEATURESFalse禁用 accelerate 相关加载特性MLFLOW_HUGGINGFACE_MODEL_MAX_SHARD_SIZE500MBsave_pretrained的权重分片大小用于大模型分片落盘MLFLOW_INPUT_EXAMPLE_INFERENCE_TIMEOUT180秒用 input_example 推断 pip 依赖时的推理超时另外dtype/torch_dtype加载参数会依据已安装 transformers 版本自动选择正确的关键字transformers 4.56.0 将torch_dtype重命名为dtype见 torch_utils.py 的_get_torch_dtype_kwarg_name保证跨版本兼容。九、测试与验证路径仓库在 tests/transformers 下提供了完整测试覆盖可作为理解行为与复现实验的参考test_transformers_model_export.py模型导出、加载与各种 pipeline 类型的行为验证test_transformers_prompt_templating.pyprompt_template的校验、格式化与return_full_text行为test_transformers_peft_model.pyPEFT 模型与base_model_path的保存/加载test_transformers_llm_inference_utils.pyLLM v1 任务的输入预处理与 OpenAI 兼容输出test_flavor_configs.py 与 test_transformers_signature.pyflavor 配置与签名推断test_transformers_autolog.pyautolog 集成行为。此外examples/transformers 提供了simple.py、conversational.py、whisper.py、sentence_transformer.py等可直接运行的示例脚本覆盖文本生成、对话、语音识别与句子嵌入等典型场景。结语mlflow.transformers以Pipeline 优先、组件兜底、checkpoint 可选三种输入形态将 HuggingFace 生态的模型统一纳入 MLflow 的模型生命周期管理保存时自动推断引擎依赖、签名与模型卡/许可证信息加载时按保存方式选择本地/Hub/外部路径三条恢复链路推理时针对数十种 pipeline 类型做输入解析与输出归一化并通过 LLM v1 任务实现 OpenAI 兼容接口。无论是传统 NLP 模型还是大模型微调产物PEFT都能借助这一 flavor 完成记录、注册、部署与监控。【免费下载链接】mlflowThe open source AI engineering platform for agents, LLMs, and ML models. MLflow enables teams of all sizes to debug, evaluate, monitor, and optimize production-quality AI applications while controlling costs and managing access to models and data.项目地址: https://gitcode.com/GitHub_Trending/ml/mlflow创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考