mlx-lora-finetune · 训练项目代码检视 2026-09-19（第 1 轮）

① 项目目标与预期

② 项目简介

③ 训练流程图（模块编号 ↔ 代码文件映射）

④ 分模块核心代码

M1 alternator.sh Shell

#!/bin/bash
# 轮班调度器 v3: 单例锁 + 完成态跳过 + 日志正确重定向
# 用法: nohup bash alternator.sh > alternator.log 2>&1 &
cd "$(dirname "$0")"
PY=.venv/bin/python
SLICE=7200        # 每片段最长 2 小时
LOCK=/tmp/mlx_alternator.lock

# ---- 单例锁: 已有活着的实例则拒绝启动 ----
if [ -f "$LOCK" ] && kill -0 "$(cat $LOCK)" 2>/dev/null; then
  echo "调度器已在运行 (PID $(cat $LOCK)), 拒绝重复启动"
  exit 1
fi
echo $$ > "$LOCK"
trap 'rm -f "$LOCK"; exit 0' EXIT INT TERM

run_slice () {
  # $1=日志文件  其余=环境变量与命令; 真正重定向 stdout/stderr 到日志
  local log="$1"; shift
  env "$@" >> "$log" 2>&1 &
  local pid=$!
  local waited=0
  # 等训练退出(完成/异常提前退出立即返回), 最长 SLICE 秒强制切换
  while kill -0 "$pid" 2>/dev/null && [ "$waited" -lt "$SLICE" ]; do
    sleep 5
    waited=$((waited + 5))
  done
  kill -TERM "$pid" 2>/dev/null     # 训练脚本有 SIGTERM 优雅保存
  # 等待优雅退出最多 180s (保存 4.6GB 检查点需要时间); 卡死则 SIGKILL 兜底
  local grace=0
  while kill -0 "$pid" 2>/dev/null && [ "$grace" -lt 180 ]; do
    sleep 5
    grace=$((grace + 5))
  done
  if kill -0 "$pid" 2>/dev/null; then
    echo "[KILL] 训练未响应 SIGTERM(${grace}s), 强制 SIGKILL"
    kill -KILL "$pid" 2>/dev/null
  fi
  wait "$pid" 2>/dev/null
}

while true; do
  # A 线: TinyStories; marker 存在=已完成, 直接跳过不启动进程
  if [ -f pretrain_ckpt/.done_56570 ]; then
    echo "=== [A] TinyStories 已完成(.done_56570), 跳过 $(date '+%m-%d %H:%M') ==="
  else
    echo "=== [A] TinyStories $(date '+%m-%d %H:%M') ==="
    run_slice pretrain_60m.log \
      START_STEP=auto MAX_STEPS=56570 PEAK_LR=8e-5 \
      $PY -u pretrain.py train
  fi
  sleep 5
  # B 线: Mamba 混合架构
  if [ -f mamba_ckpt/.done_30000 ]; then
    echo "=== [B] Mamba 已完成(.done_30000), 跳过 $(date '+%m-%d %H:%M') ==="
  else
    echo "=== [B] Mamba $(date '+%m-%d %H:%M') ==="
    run_slice mamba_train.log \
      START_STEP=auto MAX_STEPS=30000 PEAK_LR=1.2e-4 \
      $PY -u mamba_hybrid.py train
  fi
  sleep 5
done

M2 watchdog.py Python

"""电子狗 v2: 锁感知 + 告警冷却 (同事故 30 分钟内不重复发)
用法: nohup .venv/bin/python -u watchdog.py > watchdog.log 2>&1 &
巡检 300s; 规则:
  1. 锁文件存在但 PID 死亡        -> 告警"调度器异常退出" + 清锁重启 (一次性)
  2. 无锁且无 alternator 进程     -> 视为调度器未运行, 静默拉起 (不算事故)
  3. 调度器活着但连续 2 轮无训练进程 -> 告警"训练异常" + 重启调度器 (带 30min 冷却)
  4. 训练进程在但日志静默 >30min  -> 告警"疑似卡死" (带 30min 冷却, 不自动杀)
"""
import os
import subprocess
import sys
import time

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
BASE = os.path.dirname(os.path.abspath(__file__))
LOCK = '/tmp/mlx_alternator.lock'
COOLDOWN = 1800  # 同类告警冷却 30 分钟

_last_alert = {}


def procs():
    return subprocess.run(['ps', 'aux'], capture_output=True, text=True).stdout


def alert(key, title, body):
    now = time.time()
    if now - _last_alert.get(key, 0) < COOLDOWN:
        print(f'[告警冷却中] {title}', flush=True)
        return
    _last_alert[key] = now
    try:
        import send_progress
        send_progress.send(title, f'
{body}
')
        print(f'[告警] {title}', flush=True)
    except Exception as e:
        print(f'[告警失败] {e}', flush=True)


def alternator_alive():
    if not os.path.exists(LOCK):
        return False
    try:
        pid = int(open(LOCK).read().strip())
        os.kill(pid, 0)
        return True
    except (ValueError, ProcessLookupError, PermissionError):
        return False


def training_alive():
    ps = procs()
    return ('pretrain.py train' in ps) or ('mamba_hybrid.py train' in ps)


def restart_alternator():
    subprocess.Popen(['nohup', 'bash', 'alternator.sh'],
                     cwd=BASE,
                     stdout=open(f'{BASE}/alternator.log', 'a'),
                     stderr=subprocess.STDOUT)


def main():
    print('电子狗 v2 启动: 300s 巡检, 告警冷却 30min', flush=True)
    no_training_rounds = 0
    while True:
        time.sleep(300)
        try:                                   # 主循环保护: 任何异常不能杀死看门狗
            ps = procs()
            alt_bash = 'alternator.sh' in ps
            lock_ok = alternator_alive()
            training = training_alive()

            if not alt_bash and not lock_ok:
                # 调度器没在跑: 静默拉起 (锁文件不存在不算事故)
                restart_alternator()
                print('[自愈] 调度器未运行, 已拉起', flush=True)
                no_training_rounds = 0
                continue

            if alt_bash and not lock_ok:
                alert('lock', '【mlx-lora-finetune】调度器异常退出',
                      'alternator.sh 进程消失(锁残留), 已自动重启。请检查 alternator.log。')
                restart_alternator()
                no_training_rounds = 0
                continue

            if lock_ok and not alt_bash:
                # 有锁但进程没了: 说明启动中崩溃, 清锁重启
                os.remove(LOCK)
                restart_alternator()
                print('[自愈] 调度器带死锁运行, 已重启', flush=True)
                continue

            # 调度器健康, 检查训练进程
            if not training:
                no_training_rounds += 1
                if no_training_rounds >= 2:  # 10 分钟无训练 (排除 5s 切换间隙)
                    alert('training', '【mlx-lora-finetune】训练异常告警',
                          '调度器存活但连续 10 分钟无训练进程(两片段均异常退出), 已自动重启调度器。'
                          '请检查 alternator.log 与 pretrain/mamba 日志。')
                    restart_alternator()
                    no_training_rounds = 0
            else:
                no_training_rounds = 0

            # 日志新鲜度: 只检查未完成线的活跃日志 (A 完成 marker 存在则只看 mamba)
            import glob
            logs = []
            if not os.path.exists(f'{BASE}/pretrain_ckpt/.done_56570'):
                logs += glob.glob(f'{BASE}/pretrain_60m.log')
            if not os.path.exists(f'{BASE}/mamba_ckpt/.done_30000'):
                logs += glob.glob(f'{BASE}/mamba_train.log')
            if training and logs:
                stale = max(time.time() - os.path.getmtime(l) for l in logs)
                if stale > 1800:
                    alert('stale', '【mlx-lora-finetune】训练疑似卡死',
                          f'训练进程存在但日志静默 {stale/60:.0f} 分钟。'
                          f'请人工 sample  检查(看门狗不自动杀进程)。')
        except Exception as e:
            # 看门狗自身异常: 打印并存活 (send 走 SMTP 可能网络抖动)
            import traceback
            print(f'[看门狗巡检异常, 已存活] {e}\n{traceback.format_exc()[-500:]}', flush=True)


if __name__ == '__main__':
    main()

M3 pretrain.py Python

数据管线：退化过滤 + 预分词

# ---------- 数据 ----------

def line_is_junk(line):
    """行级退化过滤: 全行扫描, 同字符连跑>30 或 多样性过低"""
    if len(set(line)) < max(6, len(line) // 25):
        return True
    run, best = 1, 1
    for a, b in zip(line, line[1:]):
        run = run + 1 if a == b else 1
        best = max(best, run)
        if best > 30:
            return True
    return False

def pretokenize(max_tokens=None):
    if max_tokens is None:
        max_tokens = int(os.environ.get('PRETOK_N', '30000000'))
    from transformers import AutoTokenizer
    tok = AutoTokenizer.from_pretrained(TOKENIZER_DIR)
    arr = np.memmap(TOKENS_BIN, dtype=np.uint32, mode='w+', shape=(max_tokens,))
    pos = 0
    buf = []
    with open(DATA_TXT) as f:
        for line in f:
            if line_is_junk(line):
                continue
            ids = tok(line, add_special_tokens=False)['input_ids'] + [tok.eos_token_id]
            vs = c['vocab_size'] if (c := CONF) else 151936
            ids = [i for i in ids if i < vs]   # 词表越界过滤, 防 NaN
            if ids and ids[-1] == tok.eos_token_id:
                pass
            else:
                ids.append(tok.eos_token_id)
            buf.extend(ids)
            if len(buf) >= 200_000:
                take = min(len(buf), max_tokens - pos)
                arr[pos:pos + take] = buf[:take]
                pos += take
                buf = buf[take:]
                if pos >= max_tokens:
                    break
                print(f'  {pos/1e6:.1f}M tokens', flush=True)
    print(f'pretokenize 完成: {pos} tokens', flush=True)


def get_batch(arr, bs, sl, rng):
    """从 memmap 随机采样 bs 条长度 sl 的序列, y 为右移一位的标签"""
    ix = rng.integers(0, len(arr) - sl - 1, size=bs)
    x = mx.array(np.stack([arr[i:i + sl] for i in ix]).astype(np.int32))
    y = mx.array(np.stack([arr[i + 1:i + sl + 1] for i in ix]).astype(np.int32))
    return x, y

分块交叉熵：248K 词表在 16GB 内存下的核心手段

    def loss_fn(model, x, y):
        h = model.hidden(x)                          # [B,S,1024] 仅 17MB
        B, S, _ = h.shape
        w = model.embed.weight.T                     # 绑定 embedding, 不再额外存 lm_head
        total = mx.zeros(())
        chunk = 128
        for i in range(0, S, chunk):                 # 分块: 每 128 token 一片
            logits = (h[:, i:i + chunk] @ w).astype(mx.float16)
            logits = mx.clip(logits, -30.0, 30.0).astype(mx.float32)  # fp16 溢出保险丝
            total = total + nn.losses.cross_entropy(logits, y[:, i:i + chunk]).sum()
        return total / (B * S)

训练主循环：断点恢复 / 完成态 / 优雅保存

def train():
    os.makedirs(CKPT_DIR, exist_ok=True)
    c = CONF
    model = Model(c)
    mx.eval(model.parameters())
    nparams = sum(p.size for _, p in tree_flatten(model.parameters()))
    print(f'模型参数: {nparams/1e6:.1f}M', flush=True)

    arr = np.memmap(TOKENS_BIN, dtype=np.uint32, mode='r')
    rng = np.random.default_rng(0)

    max_steps = int(os.environ.get('MAX_STEPS', '52500'))
    start_step = os.environ.get('START_STEP', '30000')
    if start_step == 'auto':                    # 自动找最新检查点
        import glob
        fs = glob.glob(f'{CKPT_DIR}/weights_*.safetensors')
        start_step = max((int(x.split('_')[-1].split('.')[0]) for x in fs), default=0)
    else:
        start_step = int(start_step)
    if start_step >= max_steps:                 # 已完成目标: 干净退出防崩溃循环
        marker = f'{CKPT_DIR}/.done_{max_steps}'
        if not os.path.exists(marker):          # 完成通知只发一次, 防轮班邮件轰炸
            notify('【mlx-lora-finetune】TinyStories 训练已完成目标, 无需恢复',
                   f'auto 恢复到 step {start_step} >= MAX_STEPS {max_steps}。要继续请调大 MAX_STEPS。')
            open(marker, 'w').write('done')
        print('已达标, 干净退出', flush=True)
        return
    peak_lr = float(os.environ.get('PEAK_LR', '1.2e-4'))
    ckpt = f'{CKPT_DIR}/weights_{start_step}.safetensors'
    if os.path.exists(ckpt):
        loaded = mx.load(ckpt)
        model.load_weights(list(loaded.items()))
        mx.eval(model.parameters())
        print(f'从 {ckpt} 恢复权重', flush=True)
    remaining = max_steps - start_step
    warmup = lambda it: min((it + 1) / 200, 1.0)
    base = cosine_decay(peak_lr, remaining, 3e-5)
    lr_schedule = lambda it: warmup(it) * base(it)
    opt = AdamW(learning_rate=lr_schedule, weight_decay=0.1)

    bs, sl = 2, 512
    clip = 1.0  # max grad norm

    loss_and_grad = nn.value_and_grad(model, loss_fn)

    os.makedirs(CKPT_DIR, exist_ok=True)
    from transformers import AutoTokenizer
    tok = AutoTokenizer.from_pretrained(TOKENIZER_DIR)

    cur = {'step': start_step, 'model': model}
    def on_term(signum, frame):
        """SIGTERM 优雅保存: 调度器切片到时发信号, 先落盘再退出"""
        mx.save_safetensors(f'{CKPT_DIR}/weights_{cur["step"]}.safetensors',
                            dict(tree_flatten(cur['model'].parameters())))
        print(f'[SIGTERM] 已保存 step {cur["step"]}, 退出', flush=True)
        sys.exit(0)
    signal.signal(signal.SIGTERM, on_term)

    t0 = time.time()
    tokens_done = 0
    for step in range(start_step + 1, max_steps + 1):
        x, y = get_batch(arr, bs, sl, rng)
        (l, grads) = loss_and_grad(model, x, y)
        if not mx.isfinite(l).item():           # NaN 跳批防护
            print(f'step {step}: loss 非有限, 跳过', flush=True)
            continue
        grads, gnorm = clip_grad_norm(grads, clip)
        opt.update(model, grads)
        state = [model.state, opt.state]
        mx.eval(state, l, gnorm)                # 惰性求值: 此处才真正执行计算
        cur['step'] = step
        tokens_done += bs * sl
        if step % 20 == 0:
            tps = tokens_done / (time.time() - t0)
            print(f'step {step}: loss {l.item():.3f}, gnorm {float(gnorm.item()):.2f}, '
                  f'{tps:.0f} tok/s, 累计 {tokens_done/1e6:.2f}M tokens', flush=True)
        if step % 1500 == 0 or step == max_steps:
            mx.save_safetensors(
                f'{CKPT_DIR}/weights_{step}.safetensors',
                dict(tree_flatten(model.parameters())),
            )
            import glob
            olds = sorted(glob.glob(f'{CKPT_DIR}/weights_*.safetensors'),
                          key=lambda x: int(x.split("_")[-1].split(".")[0]))
            for old in olds[:-2]:               # 滚动保留最近 2 个
                os.remove(old)
            sample(model, tok, f'--- step {step} 采样 ---')
    print('训练完成', flush=True)
    notify('【mlx-lora-finetune】TinyStories 预训练完成',
           f'累计 {tokens_done/1e6:.1f}M tokens, step {max_steps}, 检查点见 pretrain_ckpt/。日志: pretrain_60m.log')


if __name__ == '__main__':
    cmd = sys.argv[1] if len(sys.argv) > 1 else 'train'
    if cmd == 'pretokenize':
        pretokenize()
    elif cmd == 'train':
        try:
            train()
        except SystemExit:
            raise
        except Exception as e:                  # 异常邮件钩子: 含 traceback 尾部
            import traceback
            notify('【mlx-lora-finetune】TinyStories 预训练异常告警',
                   f'异常: {e}\n\n{traceback.format_exc()[-2000:]}\n\n日志: pretrain_60m.log')
            raise

M4 mamba_hybrid.py Python

SSM 单步递推 + 深度可分离因果卷积

# ---------------- Mamba SSM ----------------
def _ssm_step(h, xt, Bt, Ct, dt, negA, Dv):
    """单步递推: A_bar=exp(dt·A), B_bar=dt·B; y=C·h+D·x (D 项仅此处一次)
    negA/Dv 作为显式参数传入 -> mx.compile 全局一次编译, 所有层复用"""
    d = mx.exp(dt[..., None] * negA)                   # [B,Di,DS] 离散化 A
    new = d * h + dt[..., None] * Bt[:, None, :] * xt[..., None]
    y = (new * Ct[:, None, :]).sum(-1) + Dv * xt       # [B,Di]
    return new, y


_SSM_STEP = mx.compile(_ssm_step)      # 模块级: 避免每层每步重复编译


def ssm_scan(x, B_s, C_s, delta, A, Dv):
    """选择性 SSM 逐token递推
    x: [B,T,Di]  B/C: [B,T,DS]  delta: [B,T,Di]
    A: [Di,DS] 已为负值(调用方传 -exp(A_log), 本函数不再取 exp —— 修复双重exp)
    返回: [B,T,Di]
    """
    Bb, T, Di = x.shape
    h = mx.zeros((Bb, Di, B_s.shape[-1]), dtype=x.dtype)
    ys = []
    for t in range(T):
        h, y = _SSM_STEP(h, x[:, t], B_s[:, t], C_s[:, t], delta[:, t], A, Dv)
        ys.append(y)
    return mx.stack(ys, axis=1)


def causal_conv1d(x, weight, bias):
    """深度可分离因果卷积 k=4: x[B,L,C], weight[C,K]"""
    B, L, C = x.shape
    K = weight.shape[-1]
    pad = mx.concatenate([mx.zeros((B, K - 1, C), x.dtype), x], axis=1)
    y = mx.zeros((B, L, C), x.dtype)
    for k in range(K):
        y = y + pad[:, k:L + k, :] * weight[None, None, :, k]
    return y + bias

MambaBlock：标准 Mamba 路径（x_proj 产生 dt/B/C）

class MambaBlock(nn.Module):
    """选择性 SSM 块: 线性复杂度 O(T), 推理时 O(1) 显存/token"""

    def __init__(self, c):
        super().__init__()
        d, di, ds, r = c['hidden'], c['d_inner'], c['d_state'], c['dt_rank']
        K = c['conv_kernel']
        # in_proj 输出: [x(di) | z(di) | B(ds) | C(ds) | dt(r)]
        self.in_proj = nn.Linear(d, 2 * di + 2 * ds + r, bias=False)
        # 恒等初始化: 只有最后一个 tap=1, 初始时 conv=恒等映射 (原来0.5*4会放大2倍)
        w = np.zeros((di, K)); w[:, -1] = 1.0
        self.conv_weight = mx.array(w)
        self.conv_bias = mx.zeros(di)
        self.x_proj = nn.Linear(di, r + 2 * ds, bias=False)
        self.dt_proj = nn.Linear(r, di, bias=True)
        # S4D-real 初始化: A = -arange(1..ds)
        self.A_log = mx.log(mx.tile(mx.arange(1, ds + 1, dtype=mx.float32), (di, 1)))
        self.D = mx.ones(di)
        self.out_proj = nn.Linear(di, d, bias=False)
        # dt_bias 初始化: 使 softplus 后分布于 [1e-3, 1e-1]
        lin = np.linspace(1e-3, 1e-1, di)
        self.dt_proj.bias = mx.array(-np.log(np.exp(lin) - 1))
        self.c = {'di': di, 'ds': ds, 'r': r}   # 维度存dict: MLX Module的int实例属性会被串写!

    def __call__(self, x):
        di, ds = self.c['di'], self.c['ds']
        # 标准 Mamba: in_proj 产生 [x, z]; x 经 conv+silu; B/C/dt 由 x_proj(x_conv) 产生
        # (自审修复: 原 x_proj 是死参数, B/C 直接来自 in_proj 尾部切片偏离选择性设计)
        zx = self.in_proj(x)
        xv, z = zx[..., :di], zx[..., di:2 * di]
        xv = nn.silu(causal_conv1d(xv, self.conv_weight, self.conv_bias))
        dtr, Bq, Cq = mx.split(self.x_proj(xv), [self.c['r'], self.c['r'] + ds], axis=-1)
        dt = nn.softplus(self.dt_proj(dtr))            # [B,T,Di]
        # A = -exp(A_log) 此处取一次即可 (ssm_scan 内不再 exp); D 项在 ssm_scan 内加, 此处不重复
        y = ssm_scan(xv, Bq, Cq, dt, -mx.exp(self.A_log), self.D)
        return self.out_proj(y * nn.silu(z))

块外壳与层堆叠

class MambaBlock_(nn.Module):
    """Mamba 块外壳: PreNorm + 残差
    自审发现: mx.checkpoint 对闭包捕获参数 (a)梯度恒为零 (b)前向值错误
    -> 此前 16 个 Mamba 层实际从未被训练! 已移除 checkpoint, 用 seq=256 控内存"""
    def __init__(self, c):
        super().__init__()
        self.mixer = MambaBlock(c)
        self.norm = nn.RMSNorm(c['hidden'])

    def __call__(self, x):
        return x + self.mixer(self.norm(x))


def build_layers(c):
    """16 Mamba + 2 Attention 交错 (第 8/16 层为 Attention)"""
    layers = []
    for i in range(c['layers']):
        if (i + 1) % 8 == 0:
            layers.append(AttnBlock(c))
        else:
            layers.append(MambaBlock_(c))
    return layers


class Model(nn.Module):
    def __init__(self, c):
        super().__init__()
        self.embed = nn.Embedding(c['vocab_size'], c['hidden'])
        self.layers = build_layers(c)
        self.norm = nn.RMSNorm(c['hidden'])
        self.c = c

    def hidden(self, tokens):
        x = self.embed(tokens)
        for l in self.layers:
            x = l(x)
        return self.norm(x)

    def __call__(self, tokens):
        return self.hidden(tokens) @ self.embed.weight.T

检查点：模型 + 优化器状态 + 步数，原子写

def save_all(path, model, opt, step):
    """模型+优化器状态+步数 一起存; 原子写: 临时文件+rename, 不留半截档"""
    d = {}
    for k, v in tree_flatten(model.parameters()):
        d['m.' + k] = v
    for k, v in tree_flatten(opt.state):
        d['o.' + k] = v if hasattr(v, 'astype') else mx.array(v)
    tmp = path.replace('.safetensors', '_tmp.safetensors')
    mx.save_safetensors(tmp, d)
    os.replace(tmp, path)                # 原子替换: 要么旧档要么新档, 无中间态
    with open(f'{CKPT_DIR}/last_step.txt', 'w') as f:
        f.write(str(step))


def load_all(path, model, opt):
    """配套恢复: 模型权重 + 优化器动量/二阶矩"""
    from mlx.utils import tree_unflatten
    w = dict(mx.load(path).items())
    model.load_weights([(k[2:], v) for k, v in w.items() if k.startswith('m.')])
    o_items = [(k[2:], v) for k, v in w.items() if k.startswith('o.')]
    if o_items:
        opt.state = tree_unflatten(o_items)
    mx.eval(model.parameters())
    return len([k for k in w if k.startswith('m.')])


def valid_ckpts():
    """严格匹配 weights_<数字>.safetensors, 排除 _tmp/.corrupt, 按 step 升序"""
    import glob, re
    out = []
    for f in glob.glob(f'{CKPT_DIR}/weights_*.safetensors'):
        m = re.fullmatch(r'weights_(\d+)\.safetensors', os.path.basename(f))
        if m:
            out.append((int(m.group(1)), f))
    return sorted(out)

训练主循环：坏档降级 / 全局 LR / 磁盘水位防线

def train():
    os.makedirs(CKPT_DIR, exist_ok=True)
    c = CONF
    model = Model(c)
    start_step = os.environ.get('START_STEP', '0')
    max_steps = int(os.environ.get('MAX_STEPS', '30000'))
    if start_step == 'auto':                    # 自动找最新检查点
        start_step = latest_ckpt_step()
    else:
        start_step = int(start_step)
    if start_step >= max_steps:                 # 已完成目标: 干净退出防崩溃循环
        marker = f'{CKPT_DIR}/.done_{max_steps}'
        if not os.path.exists(marker):          # 通知只发一次, 防轮班邮件轰炸
            notify('【mlx-lora-finetune】Mamba 混合架构训练已完成目标, 无需恢复',
                   f'auto 恢复到 step {start_step} >= MAX_STEPS {max_steps}。要继续请调大 MAX_STEPS。')
            open(marker, 'w').write('done')
        print('已达标, 干净退出', flush=True)
        return
    peak_lr = float(os.environ.get('PEAK_LR', '1.2e-4'))
    # LR 调度: 以"全局步"计算, warmup 只在真正前200步生效,
    # cosine 覆盖 0->max_steps 全程, 不因切片恢复重置
    def lr_at(g):
        warm = min((g + 1) / 200, 1.0)
        progress = min(g / max(max_steps - 1, 1), 1.0)
        cos = 0.5 * (1.0 + math.cos(math.pi * progress))
        return mx.array(warm * (3e-5 + (peak_lr - 3e-5) * cos))   # 必须mx.array
    lr_schedule = lr_at
    opt = AdamW(learning_rate=lr_schedule, weight_decay=0.1)

    # 按步数降序逐个尝试加载, 坏档自动降级到次新档, 全坏从头训
    if start_step > 0:
        cands = [f for _, f in reversed(valid_ckpts())]
        restored = False
        for cand in cands:
            try:
                n = load_all(cand, model, opt)
                start_step = int(cand.split('_')[-1].split('.')[0])
                print(f'从 {cand} 恢复 ({n} tensors, 含优化器状态)', flush=True)
                restored = True
                break
            except Exception as e:
                print(f'[坏档降级] {cand}: {e}', flush=True)
                os.rename(cand, cand + '.corrupt')
        if not restored:
            start_step = 0
            print('所有检查点损坏, 从头训练', flush=True)
    mx.eval(model.parameters())
    nparams = sum(p.size for _, p in tree_flatten(model.parameters()))
    print(f'MambaHybrid 参数: {nparams/1e6:.1f}M', flush=True)

    arr = np.memmap('data_corpus/tokens_200m.npy', dtype=np.uint32, mode='r')
    rng = np.random.default_rng(0)

    bs, sl = 2, 256    # seq 256: 无 checkpoint 后控内存 (512 时 ~13.6GB)
    loss_and_grad = nn.value_and_grad(model, loss_fn)
    from transformers import AutoTokenizer
    tok = AutoTokenizer.from_pretrained(D.TOKENIZER_DIR)

    cur = {'step': start_step}

    def on_term(signum, frame):
        """SIGTERM 优雅保存: 模型 + 优化器状态 + step 一并落盘"""
        save_all(f'{CKPT_DIR}/weights_{cur["step"]}.safetensors',
                 model, opt, cur['step'])
        print(f'[SIGTERM] 已保存 step {cur["step"]} (含优化器), 退出', flush=True)
        sys.exit(0)
    signal.signal(signal.SIGTERM, on_term)

    t0, tokens_done = time.time(), 0
    for step in range(start_step + 1, max_steps + 1):
        x, y = D.get_batch(arr, bs, sl, rng)
        (l, grads) = loss_and_grad(model, x, y)
        if not mx.isfinite(l).item():           # NaN 跳批
            print(f'step {step}: 非有限, 跳过', flush=True)
            continue
        grads, gnorm = clip_grad_norm(grads, 1.0)
        opt.update(model, grads)
        state = [model.state, opt.state]
        mx.eval(state, l, gnorm)                # 惰性求值: 此处才真正执行
        tokens_done += bs * sl
        cur['step'] = step
        if step % 20 == 0:
            print(f'step {step}: loss {l.item():.3f}, lr {lr_at(step):.2e}, '
                  f'{tokens_done / (time.time() - t0):.0f} tok/s', flush=True)
        if step % 250 == 0 or step == max_steps:
            # 磁盘水位防线: 4.6GB/档, 剩余<15GB先删最老档腾位, <8GB跳过保存(不崩溃)
            import shutil
            free_gb = shutil.disk_usage(CKPT_DIR).free / 1024**3
            if free_gb < 15:
                olds = valid_ckpts()
                for _, f in olds[:-1]:        # 保留最新一个, 删其余最老
                    os.remove(f)
                    free_gb = shutil.disk_usage(CKPT_DIR).free / 1024**3
                    if free_gb >= 15:
                        break
            if free_gb < 8:
                print(f'[磁盘不足] 剩{free_gb:.1f}GB, 跳过本档保存', flush=True)
            else:
                save_all(f'{CKPT_DIR}/weights_{step}.safetensors', model, opt, step)
                # 滚动保留最近 3 个 (严格枚举)
                olds = valid_ckpts()
                for _, f in olds[:-3]:
                    os.remove(f)
                sample(model, tok, f'--- step {step} 采样 ---')
    print('训练完成', flush=True)
    notify('【mlx-lora-finetune】Mamba 混合架构训练完成',
           f'step {max_steps}, 检查点见 mamba_ckpt/。日志: mamba_train.log')


if __name__ == '__main__':
    cmd = sys.argv[1] if len(sys.argv) > 1 else 'train'
    if cmd == 'pretokenize':
        D.pretokenize()
    else:
        try:
            train()
        except SystemExit:
            raise
        except Exception as e:                  # 异常邮件钩子: 含 traceback
            import traceback
            notify('【mlx-lora-finetune】Mamba 混合架构训练异常告警',
                   f'异常: {e}\n\n{traceback.format_exc()[-2000:]}\n\n日志: mamba_train.log')
            raise

M5 send_progress.py Python

def send(subject, html_body, attachment_path=None):
    msg = MIMEMultipart("mixed")
    msg["Subject"], msg["From"], msg["To"] = subject, FROM_ADDR, TO_ADDR
    alt = MIMEMultipart("alternative")
    alt.attach(MIMEText(html_body, "html", "utf-8"))
    msg.attach(alt)
    if attachment_path:
        with open(attachment_path, "rb") as f:
            att = MIMEApplication(f.read())
        att.add_header("Content-Disposition", "attachment",
                       filename=("utf-8", "", attachment_path.split("/")[-1]))
        msg.attach(att)
    ctx = ssl.create_default_context()
    with smtplib.SMTP("smtp.163.com", 25, timeout=30) as s:  # 25端口 STARTTLS
        s.ehlo(); s.starttls(context=ctx); s.ehlo()
        s.login(SMTP_USER, SMTP_PASS)
        s.sendmail(FROM_ADDR, [TO_ADDR], msg.as_string())

M6 convert_to_hf.py Python

MX_CONF = dict(vocab_size=248077, hidden=1024, layers=18, heads=16,
               kv_heads=4, head_dim=64, ffn=2816, rope_theta=10000.0)
TOKENIZER_SRC = 'models/Qwen3.5-0.8B-4bit'


def convert(mlx_safetensors, out_dir):
    os.makedirs(out_dir, exist_ok=True)
    w = dict(mx.load(mlx_safetensors).items())
    out = {}
    out['model.embed_tokens.weight'] = w['embed.weight']
    out['model.norm.weight'] = w['norm.weight']

    def permute_rope(wt, heads, head_dim):
        # 交错配对 (2i, 2i+1) -> 半分配对 (i, i+D/2): Llama->HF 经典置换
        import numpy as np
        wa = wt.reshape(heads, head_dim, -1)
        idx = list(range(0, head_dim, 2)) + list(range(1, head_dim, 2))
        wa = wa[:, idx, :]
        return mx.array(wa.reshape(heads * head_dim, -1))

    for i in range(MX_CONF['layers']):
        p = f'layers.{i}.'
        h = f'model.layers.{i}.'
        out[h + 'self_attn.q_proj.weight'] = permute_rope(
            w[p + 'attn.q_proj.weight'], MX_CONF['heads'], MX_CONF['head_dim'])
        out[h + 'self_attn.k_proj.weight'] = permute_rope(
            w[p + 'attn.k_proj.weight'], MX_CONF['kv_heads'], MX_CONF['head_dim'])
        out[h + 'self_attn.v_proj.weight'] = w[p + 'attn.v_proj.weight']
        out[h + 'self_attn.o_proj.weight'] = w[p + 'attn.o_proj.weight']
        out[h + 'mlp.gate_proj.weight'] = w[p + 'ffn.gate_proj.weight']
        out[h + 'mlp.up_proj.weight'] = w[p + 'ffn.up_proj.weight']
        out[h + 'mlp.down_proj.weight'] = w[p + 'ffn.down_proj.weight']
        out[h + 'input_layernorm.weight'] = w[p + 'n1.weight']
        out[h + 'post_attention_layernorm.weight'] = w[p + 'n2.weight']
    # lm_head 单独存 (Hermes/llama 支持 tie; 显式写出更稳)
    out['lm_head.weight'] = w['embed.weight']
    out = {k: v.astype(mx.bfloat16) for k, v in out.items()}
    mx.save_safetensors(f'{out_dir}/model.safetensors', out)

    conf = {
        'architectures': ['LlamaForCausalLM'],
        'model_type': 'llama',
        'vocab_size': 248077,
        'hidden_size': 1024,
        'intermediate_size': 2816,
        'num_hidden_layers': 18,
        'num_attention_heads': 16,
        'num_key_value_heads': 4,
        'head_dim': 64,
        'rms_norm_eps': 1e-5,
        'rope_theta': 10000.0,
        'rope_scaling': None,
        'max_position_embeddings': 1024,
        'tie_word_embeddings': False,
        'bos_token_id': 1, 'eos_token_id': 2,
        'torch_dtype': 'bfloat16',
    }
    with open(f'{out_dir}/config.json', 'w') as f:
        json.dump(conf, f, indent=1)

    # tokenizer 文件 (从 Qwen 目录拷贝)
    for fn in ['tokenizer.json', 'tokenizer_config.json', 'chat_template.jinja']:
        src = os.path.join(TOKENIZER_SRC, fn)
        if os.path.exists(src):
            shutil.copy(src, os.path.join(out_dir, fn))
    # generation_config: eos 对齐 tokenizer (248046)
    with open(f'{out_dir}/generation_config.json', 'w') as f:
        json.dump({'eos_token_id': 248046, 'bos_token_id': 1, 'pad_token_id': 248044}, f)

    total = sum(v.size for v in out.values())
    print(f'转换完成: {out_dir} | {len(out)} tensors | {total/1e6:.0f}M params (bf16)')


if __name__ == '__main__':
    convert(sys.argv[1], sys.argv[2])

⑤ 核心超参配置表

⑥ 项目依赖清单

mlx-lora-finetune/
├── alternator.sh                 # M1 轮班调度器
├── watchdog.py                   # M2 电子狗
├── pretrain.py                   # M3 A线预训练
├── mamba_hybrid.py               # M4 B线混合架构
├── send_progress.py              # M5 邮件通知
├── convert_to_hf.py              # M6 MLX→HF 转换
├── pretrain_60m.log              # A 线训练日志（调度器重定向）
├── mamba_train.log               # B 线训练日志（调度器重定向）
├── alternator.log                # 调度器自身日志
├── watchdog.log                  # 看门狗自身日志
├── pretrain_ckpt/                # A 线检查点
│   ├── weights_<step>.safetensors
│   └── .done_56570               # 完成态 marker
├── mamba_ckpt/                   # B 线检查点
│   ├── weights_<step>.safetensors
│   ├── last_step.txt             # 最近保存的 step
│   └── .done_30000               # 完成态 marker
├── data_corpus/
│   └── tokens_200m.npy           # uint32 token memmap
├── models/
│   ├── Qwen3.5-0.8B-4bit/        # 分词器来源
│   └── hf_46b_step30k/           # M6 输出（HF 格式）
└── .venv/                        # Python 虚拟环境

本文档由训练项目代码汇总生成 · 代码基线：完整源码评审包 v2（2026-09-19）· 高亮：highlight.js 11.9.0 · 流程图：mermaid 11.4.1（neutral 主题 / curve:basis）