MLflow Transformers 功能
MLflow Transformers 功能提供与 Hugging Face Transformers 库的原生集成，支持 NLP、音频、视觉和多模态任务的模型日志记录、加载和推理。
主要特性
- Pipeline 和组件日志记录：保存完整的 Pipeline 或单个模型组件
- PyFunc 集成：使用标准化的推理接口部署模型
- PEFT 支持：对参数高效微调（LoRA、QLoRA 等）的原生支持
- Prompt 模板：与 Pipeline 一起保存和管理 Prompt 模板
- 自动元数据日志记录：自动记录模型卡和元数据
- 灵活的推理配置：通过
model_config和签名参数自定义模型行为
安装
bash
pip install mlflow transformers
基本用法
日志记录 Pipeline
python
import mlflow
from transformers import pipeline
# Create a text generation pipeline
text_gen = pipeline("text-generation", model="gpt2")
# Log the pipeline
with mlflow.start_run():
mlflow.transformers.log_model(
transformers_model=text_gen,
name="model",
)
加载和推理
python
# Load as native transformers
model = mlflow.transformers.load_model("runs:/<run_id>/model")
result = model("Hello, how are you?")
# Load as PyFunc
pyfunc_model = mlflow.pyfunc.load_model("runs:/<run_id>/model")
result = pyfunc_model.predict("Hello, how are you?")
使用 HuggingFace Trainer 进行自动日志记录
当使用 HuggingFace Trainer 类进行微调时，您可以通过在 TrainingArguments 中设置 report_to="mlflow" 来启用向 MLflow 的自动日志记录
python
from transformers import TrainingArguments, Trainer
training_args = TrainingArguments(
output_dir="./results",
report_to="mlflow", # Enable MLflow logging
# ... other training arguments
)
trainer = Trainer(
model=model,
args=training_args,
train_dataset=train_dataset,
# ... other trainer arguments
)
trainer.train()
这将自动将训练指标、超参数和模型检查点记录到您当前的 MLflow 运行中。
教程
快速入门
微调
高级用例
重要注意事项
PyFunc 限制
- 并非所有 Pipeline 类型都支持 PyFunc 推理
- 某些输出（例如，额外的分数、引用）可能无法捕获
- 支持音频和文本 LLM；视觉和多模态模型需要原生加载
- 请参阅指南了解支持的 Pipeline 类型
输入/输出类型
PyFunc 的输入和输出格式可能与原生 Pipeline 不同。确保与您的数据处理工作流程兼容。
模型配置
当同时提供 ModelSignature 和 model_config 时，ModelSignature 中的参数会覆盖 model_config 中的参数。
处理大型模型
对于具有数十亿参数的模型，MLflow 提供优化技术来减少内存使用并加快日志记录速度。请参阅大型模型指南。
任务
task 参数决定输入/输出格式。MLflow 支持原生的 Transformers 任务，以及 llm/v1/chat 和 llm/v1/completions 等高级任务，以实现兼容 OpenAI 的推理。请参阅任务指南。