M3 Mac上MPS训练仅用1/10GPU核心，求多核心利用方案

问题概述 #

M3 Mac在使用PyTorch MPS训练DistilBert模型时，仅利用了1/10的GPU核心，希望提升利用率至全部或至少8个核心。

一、优化PyTorch MPS训练配置 #

MPS本身支持多核心并行，只需调整参数和环境配置即可触发：

• 增大训练批次大小
  当前per_device_train_batch_size=8过小，M3 GPU内存可支撑更大批次，建议调整为32或64（根据实际内存情况微调）：

  training_args = TrainingArguments(
      output_dir='./results',
      evaluation_strategy="epoch",
      learning_rate=2e-5,
      per_device_train_batch_size=32,  # 提升批次大小
      per_device_eval_batch_size=32,
      num_train_epochs=3,
      weight_decay=0.01,
  )
  

• 强制MPS使用多线程
  在代码开头添加环境变量设置，指定使用的核心数（M3为10核，可设为8或10）：

  import os
  os.environ["OMP_NUM_THREADS"] = "8"
  os.environ["MPS_NUM_THREADS"] = "8"
  

• 优化MPS内存与同步设置
  关闭不必要的同步检查，释放更多内存空间：

  torch.mps.set_per_process_memory_fraction(0.9)  # 允许使用90%的GPU内存
  torch.mps.empty_cache()  # 训练前清空缓存
  

• 使用梯度累积模拟大批次
  若内存不足以支撑超大批次，用梯度累积提升GPU利用率：

  training_args = TrainingArguments(
      # ...其他参数保留
      gradient_accumulation_steps=4,  # 每4个小批次累积一次梯度
  )
  

二、MLX替代方案 #

MLX是苹果官方针对Apple Silicon优化的框架，对多核心GPU的支持更原生，训练效率更高，适配步骤如下：

1. 安装依赖

   pip install mlx mlx-transformers
   

2. 适配MLX的训练代码

   import mlx.core as mx
   import mlx.nn as nn
   from mlx_transformers import DistilBertForSequenceClassification, DistilBertTokenizer
   from datasets import load_dataset
   import pandas as pd
   from sklearn.model_selection import train_test_split
   
   # 数据加载与预处理
   df = pd.read_csv('jigsaw-toxic-comment-train-processed-seqlen128.csv')
   df = df[['comment_text', 'toxic']]
   train_df, val_df = train_test_split(df, test_size=0.2, random_state=42)
   
   tokenizer = DistilBertTokenizer.from_pretrained('distilbert-base-uncased')
   
   def preprocess(examples):
       tokens = tokenizer(examples['comment_text'], padding='max_length', truncation=True, max_length=128)
       return {
           'input_ids': mx.array(tokens['input_ids']),
           'attention_mask': mx.array(tokens['attention_mask']),
           'labels': mx.array(examples['toxic'], dtype=mx.float32)
       }
   
   train_dataset = load_dataset('pandas', data_files={'train': train_df})['train'].map(preprocess, batched=True)
   val_dataset = load_dataset('pandas', data_files={'val': val_df})['train'].map(preprocess, batched=True)
   
   # 模型与训练配置
   model = DistilBertForSequenceClassification.from_pretrained('distilbert-base-uncased', num_labels=1)
   loss_fn = nn.BCEWithLogitsLoss()
   optimizer = mx.optimizers.Adam(learning_rate=2e-5)
   state = [model.state, optimizer.state]
   
   # 训练循环
   def train_step(batch):
       def loss_wrapper(model):
           logits = model(input_ids=batch['input_ids'], attention_mask=batch['attention_mask']).logits.squeeze(-1)
           return loss_fn(logits, batch['labels'])
       loss, grads = nn.value_and_grad(model, loss_wrapper)()
       optimizer.update(model, grads)
       return loss
   
   num_epochs = 3
   for epoch in range(num_epochs):
       mx.random.seed(42 + epoch)
       train_dataset.shuffle(seed=42 + epoch)
       total_loss = 0.0
       for batch in train_dataset.iter(batch_size=32):
           loss = train_step(batch)
           total_loss += loss.item()
           mx.eval(state)
       print(f"Epoch {epoch+1}, Train Loss: {total_loss/len(train_dataset):.4f}")
   
       # 验证环节
       val_loss = 0.0
       for batch in val_dataset.iter(batch_size=32):
           logits = model(input_ids=batch['input_ids'], attention_mask=batch['attention_mask']).logits.squeeze(-1)
           val_loss += loss_fn(logits, batch['labels']).item()
       print(f"Epoch {epoch+1}, Val Loss: {val_loss/len(val_dataset):.4f}")
   

   MLX会自动调度全部Apple Silicon GPU核心，无需手动配置线程数。

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