MacBook Pro M2微调tiny_starcoder_py-vi06遇MPS设备RuntimeError求助

背景 #

我正尝试在搭载M2芯片的MacBook Pro上微调tiny_starcoder_py-vi06模型，打造本地代码助手的概念验证（PoC）。

数据集格式 #

[
  { "filename": "filename.ext", "text": "const a = 1; // other code" }
]

使用代码 #

import torch
import pandas as pd
from datasets import Dataset
from transformers import AutoModelForCausalLM, AutoTokenizer, DataCollatorForLanguageModeling, Trainer, TrainingArguments

def main():
    model_name = "tiendung/tiny_starcoder_py-vi06"
    model = AutoModelForCausalLM.from_pretrained(model_name)
    tokenizer = AutoTokenizer.from_pretrained(model_name)

    if tokenizer.pad_token is None:
        tokenizer.add_special_tokens({'pad_token': tokenizer.eos_token})

    model.resize_token_embeddings(len(tokenizer))

    df = pd.read_json("collection.json")
    dataset = Dataset.from_pandas(df)

    def tokenize_function(examples):
        return tokenizer(examples["text"], padding="max_length", truncation=True)

    tokenized_datasets = dataset.map(tokenize_function, batched=True)
    tokenized_datasets.set_format(type='torch', columns=['input_ids', 'attention_mask'])

    data_collator = DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm=False)

    training_args = TrainingArguments(
        output_dir="./results",
        overwrite_output_dir=True,
        num_train_epochs=3,
        per_device_train_batch_size=8,
        save_steps=10_000,
        save_total_limit=2,
    )

    device = torch.device("mps" if torch.backends.mps.is_available() else "cpu")
    model.to(device)

    def move_to_device(batch):
        batch['input_ids'] = batch['input_ids'].to(device)
        batch['attention_mask'] = batch['attention_mask'].to(device)
        return batch

    tokenized_datasets = tokenized_datasets.map(move_to_device, batched=True)

    trainer = Trainer(
        model=model,
        args=training_args,
        data_collator=data_collator,
        train_dataset=tokenized_datasets,
    )

    trainer.train()
    trainer.save_model("fine_tuned_model")
    tokenizer.save_pretrained("fine_tuned_model")

if __name__ == "__main__":
    main()

运行错误 #

RuntimeError: Placeholder storage has not been allocated on MPS device!

问题 #

1. 使用MPS设备的配置是否存在遗漏？
2. 在MacBook M2芯片上训练模型是否需要特定配置调整？
3. 针对该问题有哪些排查建议？

环境信息 #

• 操作系统：macOS
• 芯片：Apple M2
• Transformers版本：4.42.4

1. MPS配置的遗漏点 #

手动将数据集移到MPS设备的方式存在问题，Trainer内部会自动处理设备分配，提前手动移动数据会和Trainer的设备管理逻辑冲突，引发占位符错误。另外TrainingArguments未明确开启MPS适配的优化配置，也是缺失项。

2. M2芯片上的特定配置调整 #

• 修改TrainingArguments：添加fp16=True（M2支持半精度训练，既能提升速度又能减少显存占用），无需手动移动模型和数据到MPS，Trainer会自动检测并调用MPS设备。
• 移除手动移数据代码：删除move_to_device函数及对应的map调用，避免设备逻辑冲突。
• 调小批次大小：M2显存有限，per_device_train_batch_size=8可能超出显存承载，建议先降到2或4，再根据实际显存使用情况调整。
• 适配PyTorch版本：确保使用PyTorch 2.0及以上版本，旧版本存在MPS兼容性bug。

3. 排查建议 #

• 验证MPS可用性：在代码开头添加以下代码确认MPS是否正常工作：

  print(torch.backends.mps.is_available())
  print(torch.backends.mps.is_built())
  

  若输出均为True则MPS可用，否则重新安装适配M2的PyTorch版本。
• 逐步简化代码定位：先注释训练逻辑，仅加载模型和数据集，验证模型能否正常部署到MPS，再逐步恢复训练代码，定位错误触发点。
• 监控显存占用：使用torch.mps.current_allocated_memory()和torch.mps.driver_allocated_memory()查看显存使用情况，确认是否因显存不足导致错误。
• 更新依赖库：将Transformers、Datasets、PyTorch更新到最新稳定版，修复已知的MPS兼容问题。

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