mlflow.transformers.log_model()执行无响应，无法完成运行

问题 #

使用mlflow.transformers.log_model()记录LoRA微调后的HuggingFace模型时，方法持续运行无法完成，无任何报错。控制台输出显示Skipping saving pretrained model weights to disk，排除权重保存导致的阻塞问题。

环境配置：

• Python 3.11.9
• 依赖版本：transformers = "^4.41.2"、mlflow = "^2.15.1"

代码示例（含伪代码）：

import mlflow
import torch
from peft import LoraConfig
from transformers import (
    AutoModelForCausalLM,
    AutoTokenizer,
    BitsAndBytesConfig,
    TrainingArguments,
)
from trl import SFTTrainer, setup_chat_format

train_dataset = ...
eval_dataset = ...

model_id = "LeoLM/leo-hessianai-7b-chat-bilingual"

# Load model and tokenizer
model = AutoModelForCausalLM.from_pretrained(
    model_id,
    device_map="auto",
    torch_dtype=torch.bfloat16,
    quantization_config=bnb_config,
)
tokenizer = AutoTokenizer.from_pretrained(model_id)
tokenizer_no_pad = AutoTokenizer.from_pretrained(model_id, add_bos_token=True)
model, tokenizer = setup_chat_format(model, tokenizer)
peft_config = LoraConfig(...)
args = TrainingArguments(...)

# Define Trainer
trainer = SFTTrainer(
    model=model,
    args=args,
    train_dataset=train_dataset,
    eval_dataset=eval_dataset,
    peft_config=peft_config,
    tokenizer=tokenizer,
    packing=True,
)

# mlflow
mlflow.set_experiment("my_experiment")
with mlflow.start_run() as run:
    mlflow.transformers.autolog()
    trainer.train()
    
     components = {
         "model": trainer.model,
         "tokenizer": tokenizer_no_pad,
     }
     # 该方法一直无法完成
     mlflow.transformers.log_model(
         transformers_model=components,
         artifact_path="model",
    )

控制台最后输出：

INFO mlflow.transformers: Overriding save_pretrained to False for PEFT models, following the Transformers behavior. The PEFT adaptor and config will be saved, but the base model weights will not and reference to the HuggingFace Hub repository will be logged instead.
Unrecognized keys in `rope_scaling` for 'rope_type'='linear': {'type'}
/mypath/llm4pa-open-source/.venv/lib/python3.11/site-packages/peft/utils/save_and_load.py:209: UserWarning: Setting `save_embedding_layers` to `True` as the embedding layer has been resized during finetuning.
  warnings.warn(
2024/08/12 18:21:14 INFO mlflow.transformers: Skipping saving pretrained model weights to disk as the save_pretrained is set to False. The reference to HuggingFace Hub repository LeoLM/leo-hessianai-7b-chat-bilingual will be logged instead.
/mypath/llm4pa-open-source/.venv/lib/python3.11/site-packages/_distutils_hack/__init__.py:26: UserWarning: Setuptools is replacing distutils.
  warnings.warn("Setuptools is replacing distutils.")

解决方法 #

1. 避免autolog与手动log_model的冲突 #

mlflow.transformers.autolog()会自动记录训练过程中的模型、指标等内容，与手动调用log_model可能产生资源占用或逻辑冲突。可以选择：

• 移除mlflow.transformers.autolog()，改为手动记录训练指标（如果需要），只保留手动log_model调用；
• 若需要保留autolog，在调用log_model前禁用自动日志：

mlflow.transformers.autolog(disable=True)

2. 使用MLFlow专门的PEFT模型记录方法 #

对于LoRA等PEFT微调模型，MLFlow提供了mlflow.peft.log_model()，对PEFT模型的支持更稳定，避免通用transformers.log_model()的潜在兼容性问题：

# 替换原log_model调用
mlflow.peft.log_model(
    peft_model=trainer.model,
    tokenizer=tokenizer_no_pad,
    artifact_path="model",
)

3. 将模型移至CPU后再记录 #

微调后的模型可能仍驻留在GPU内存中，部分情况下会导致MLFlow在处理模型时阻塞。添加以下代码将模型移至CPU：

# 在log_model前执行
trainer.model = trainer.model.to('cpu')
torch.cuda.empty_cache()  # 清理GPU缓存

4. 简化组件配置 #

先使用训练时的tokenizer（而非自定义的tokenizer_no_pad）测试log_model功能，确认基础流程正常后再替换tokenizer：

components = {
    "model": trainer.model,
    "tokenizer": tokenizer,  # 先使用训练用的tokenizer
}
mlflow.transformers.log_model(
    transformers_model=components,
    artifact_path="model",
)

5. 开启DEBUG日志定位阻塞点 #

增加日志级别，查看MLFlow内部执行细节，找到卡住的具体环节：

import logging
logging.basicConfig(level=logging.DEBUG)

内容的提问来源于stack exchange，提问作者chamaoskurumi