Mac本地训练Flux模型：M1/M2芯片环境配置指南

1. 本地Mac环境下的Flux模型训练指南

作为一名长期在Mac平台上折腾机器学习模型的开发者，我深知在Apple Silicon上运行训练任务的各种痛点。最近在尝试实现Flux模型的本地训练时，我经历了无数次的失败和尝试，最终找到了一套相对可行的解决方案。本文将详细介绍如何在配备M1/M2芯片的Mac电脑上搭建Flux训练环境，并成功运行训练脚本。

重要提示：本方法对系统内存要求较高，建议使用16GB及以上统一内存的Mac设备。8GB内存的机器可能会因内存不足而无法完成训练。

2. 环境准备与工具链配置

2.1 基础环境搭建

首先需要确保你的Mac系统满足以下基本要求：

• macOS Monterey (12.3+) 或更高版本
• 已安装Homebrew包管理器
• Python 3.8+ 环境
• Git版本控制工具

打开终端，执行以下命令安装基础依赖：

bash复制# 安装Homebrew（如未安装）
/bin/bash -c "$(curl -fsSL https://raw.githubusercontent.com/Homebrew/install/HEAD/install.sh)"

# 安装Python和Git
brew install python git

2.2 克隆并配置ai-toolkit仓库

这里我们使用经过修改支持MacOS的ai-toolkit分支：

bash复制git clone https://github.com/hughescr/ai-toolkit
cd ai-toolkit
git submodule update --init --recursive

这个仓库是Ostris原版ai-toolkit训练脚本的一个分支，由Hughescr进行了MacOS适配。主要修改包括：

• 使用torch.amp替代torch.cuda.amp以支持MPS后端
• 强制使用spawn而非fork进行多进程处理
• 禁用T5量化器（MPS不兼容）
• 强制设置数据加载器的num_workers=0

2.3 创建Python虚拟环境

为了避免与系统Python环境冲突，建议创建专用虚拟环境：

bash复制python3 -m venv venv
source venv/bin/activate

3. 依赖安装与配置调整

3.1 PyTorch与核心依赖安装

安装适配Apple Silicon的PyTorch版本：

bash复制pip3 install torch

然后安装项目所需的其他依赖：

bash复制pip3 install -r requirements.txt

常见问题：如果在安装过程中遇到权限错误，可以尝试添加--user参数或使用pip install --upgrade pip确保pip版本最新。

3.2 配置文件准备

ai-toolkit提供了多种训练配置模板，位于config/examples/目录下。对于Flux模型训练，我们需要关注以下两个配置文件：

• train_lora_flux_24gb.yaml：标准Flux Dev模型训练配置
• train_lora_flux_schnell_24gb.yaml：Flux Schnell模型训练配置

建议复制其中一个模板到项目根目录并进行修改：

bash复制cp config/examples/train_lora_flux_24gb.yaml my_flux_config.yaml

关键配置项说明：

yaml复制train:
  batch_size: 4  # 根据显存调整，建议从2开始尝试
  num_workers: 0  # Mac必须设置为0
  learning_rate: 1e-5
  max_steps: 1000
  
data:
  dataset_path: "/path/to/your/dataset"  # 替换为实际数据集路径
  image_size: 512
  
output:
  save_path: "/path/to/save/model"  # 模型输出路径

4. 训练执行与问题排查

4.1 启动训练任务

由于MPS后端尚未实现所有PyTorch操作，需要通过环境变量启用CPU回退：

bash复制PYTORCH_ENABLE_MPS_FALLBACK=1 python run.py my_flux_config.yaml

4.2 常见问题与解决方案

1. 内存不足错误：

   • 症状：进程被杀死或报内存错误
   • 解决方案：减小batch_size，关闭其他内存占用大的应用

2. MPS操作未实现错误：

   • 症状：报错包含"not implemented for MPS"
   • 解决方案：确保已设置PYTORCH_ENABLE_MPS_FALLBACK=1

3. 数据加载卡死：

   • 症状：程序在数据加载阶段无响应
   • 解决方案：确认num_workers=0且数据集路径正确

4. 梯度计算错误：

   • 症状：NaN值或梯度爆炸
   • 解决方案：尝试减小学习率，或使用梯度裁剪

4.3 训练监控与优化

训练过程中可以通过以下方法监控资源使用情况：

1. 活动监视器：观察CPU、GPU和内存使用情况
2. 终端命令：

   bash复制top -o cpu  # 查看CPU使用率
   vm_stat  # 查看内存使用情况
   

对于性能优化建议：

• 使用SSD存储加速数据加载
• 保持系统凉爽以避免性能降频
• 训练期间关闭不必要的后台应用

5. 进阶配置与技巧

5.1 使用UI界面启动训练

ai-toolkit提供了图形界面，可以通过以下命令启动：

bash复制PYTORCH_ENABLE_MPS_FALLBACK=1 python flux_train_ui.py

UI界面提供了更直观的配置方式，适合不熟悉YAML配置文件的用户。

5.2 自定义模型训练

如果想训练自定义版本的Flux模型，可以修改以下关键参数：

• model_path: 指定基础模型路径
• lora_rank: 调整LoRA秩大小
• text_encoder_lr: 单独设置文本编码器学习率
• unet_lr: 单独设置UNet学习率

5.3 训练中断与恢复

如果训练意外中断，可以通过以下方式恢复：

1. 在配置文件中设置resume_from_checkpoint: true
2. 确保output.save_path指向之前保存的检查点目录
3. 重新启动训练脚本

6. 性能对比与实测数据

在我的M1 Max (32GB)设备上测试结果：

实测建议：对于16GB内存的设备，建议使用Flux Schnell配置并将batch_size设为1-2

7. 模型应用与后续处理

训练完成后，可以在Stable Diffusion等工具中使用生成的LoRA权重：

1. 将训练输出的.safetensors文件放入模型目录
2. 在生成时使用触发词激活LoRA效果
3. 可通过调整权重强度(通常0.6-0.8)获得最佳效果

对于想要进一步提升效果的开发者，可以尝试：

• 更精细的数据集清洗
• 渐进式学习率调整
• 多阶段训练策略

这套方案虽然不如CUDA环境高效，但确实为Mac用户提供了本地训练Flux模型的可行方案。随着PyTorch对MPS后端的持续优化，未来性能有望进一步提升。