希望充分利用 Mac 的 GPU 性能，这是阿里云与苹果官方针对 Apple Silicon 芯片推荐的方案。MLX 是专为统一内存架构设计的机器学习框架，对 Qwen3 系列模型有原生支持，可以完美解决 GQA 的兼容性问题。

切换到 MLX 框架。MLX 是苹果官方为 Apple Silicon 设计的深度学习框架，对 Qwen3 的 GQA 机制有完美支持。

方案一：原生mlx-lm (轻量级，最推荐)

这是最轻量、开发效率最高的选择。主要逻辑是通过一个 YAML 配置文件来控制所有参数，完全不需要 Trainer 这类复杂的类定义。

第1步：安装环境

在终端运行以下命令

pip install mlx-lm swanlab modelscope

第2步：准备你的数据集

需要先将你的医疗数据集转换成 mlx-lm 支持的标准 JSONL 格式。你可以直接使用下面的脚本一次性完成格式转换和字段替换。

直接复制以下代码，保存为 convert_data.py 并运行：

import json
import os

# 配置系统提示词
system_prompt = '你是一个医学专家，为用户解答问题，并进行思考后回答'

def convert_to_mlx_format(input_file, output_file, replace_map):
    """
    将原始数据集转换为 mlx-lm 格式。
    """
    mlx_data = []
    with open(input_file, 'r', encoding='utf-8') as f_in:
        for line in f_in:
            try:
                original = json.loads(line.strip())
                # 构建 assistant 的内容
                assistant_content = original['answer']
                # 可选：如果数据有 'think' 字段，可以在这里处理，比如：
                # if 'think' in original and original['think']:
                #     assistant_content = f"<think>{original['think']}</think>{assistant_content}"
                
                # 构建 messages
                messages = [
                    {"role": "system", "content": system_prompt},
                    {"role": "user", "content": original['question']},
                    {"role": "assistant", "content": assistant_content}
                ]
                mlx_data.append({"messages": messages})
            except (json.JSONDecodeError, KeyError) as e:
                print(f"跳过格式错误或缺少字段的行: {e}")

    with open(output_file, 'w', encoding='utf-8') as f_out:
        for item in mlx_data:
            f_out.write(json.dumps(item, ensure_ascii=False) + '\n')
    print(f"转换完成！共处理 {len(mlx_data)} 条数据，已保存至 {output_file}")

# --- 执行转换 ---
# 确保以下文件路径与你的实际文件名一致
convert_to_mlx_format('train.json', 'train_mlx.jsonl', replace_map=None)
convert_to_mlx_format('eval.json', 'valid_mlx.jsonl', replace_map=None) # 作为验证集使用

第3步：编写配置文件

将你的 PyTorch 训练参数“翻译”成 MLX 的配置。把下面的内容保存为 config_medical.yaml

# config_medical.yaml
model: "Qwen/Qwen3-0.6B"
train: true
fine_tune_type: "lora"          # 使用 LoRA 微调，内存友好
data: "."                        # 数据集文件夹路径，会自动读取同目录下的 train.jsonl 和 valid.jsonl
iters: 542                       # 总迭代步数 (原数据集总行数)
batch_size: 1                    # 对应原 batch_size
learning_rate: 5e-5              # 对应原 learning_rate
val_batches: 25                  # 验证集迭代次数
steps_per_eval: 100              # 对应原 eval_steps
steps_per_report: 10             # 对应原 logging_steps
save_every: 400                  # 对应原 save_steps
max_seq_length: 4096             # 对应原 MAX_LENGTH
grad_checkpoint: true            # 对应原 gradient_checkpointing
lora_parameters:                 # LoRA 作用于哪些层
  keys: ["self_attn.q_proj", "self_attn.v_proj"]
  rank: 8
  scale: 20.0
  dropout: 0.0
adapter_path: "adapters"         # LoRA 权重保存路径

第4步：启动训练

在终端中运行以下命令，训练就会自动开始，并记录到 SwanLab

mlx_lm.lora --config config_medical.yaml --report-to swanlab --project-name "qwen3_0.6B_medical"

第5步：合并权重并测试

训练完成后，使用以下命令将 LoRA 权重合并回基础模型，得到一个完整的微调模型

mlx_lm.fuse --model Qwen/Qwen3-0.6B --adapter-path adapters --save-path ./final_model_merged

方案二：mlx-tune (代码改动最小)

这个方案直接兼容你原有的 PyTorch 代码结构，通过替换 import 来解决问题。

第1步：安装依赖

pip install mlx-tune swanlab

第2步：修改你的训练脚本

你原始脚本的改动幅度非常小，具体如下

# ... 之前的导入保持不变 ...
# from transformers import AutoTokenizer, AutoModelForCausalLM, Trainer, TrainingArguments, DataCollatorForSeq2Seq
# 注释掉原来的 import，改成：
from mlx_tune import FastLanguageModel as AutoModelForCausalLM
# 注意：mlx-tune 也提供了对应的 Trainer 和 TrainingArguments
from mlx_tune import Trainer, TrainingArguments
from transformers import AutoTokenizer, DataCollatorForSeq2Seq
import torch

# ... 数据加载和 tokenization 的代码可以完全保留，不用修改 ...

# ---------- 修改点 2：模型加载 ----------
# 原代码
# model = AutoModelForCausalLM.from_pretrained(model_dir, device_map='auto', torch_dtype='auto')
# 新代码
model, tokenizer = AutoModelForCausalLM.from_pretrained(
    model_dir,
    dtype=torch.float32,      # MLX 推荐使用 float32
    device_map="cpu",         # 或者 "mps"，虽然它本身就在 Apple Silicon 上跑得最好
)

# ... 后面的 Trainer 设置和训练循环基本可以保持不变 ...

这两种方案各有优势：

• mlx-lm (方案一)：最轻量，开发效率高，通过配置文件控制，适合快速实验调参。

• mlx-tune (方案二)：代码改动量最小，适合希望保留现有代码结构、最小化迁移成本的场景。