[CUDA] 0.32.3: repeated forward pass of a 4-bit quantized model is not bitwise reproducible on GB10 (sm_121); not reproduced with 0.32.4

オープン
  • 何か問題がありますか？
#4,843 コメント 1 件 リアクション 0 件 担当者 0 名 GitHub で見る

メンテナーはふだん 3 日以内に返信

まだ誰も着手していません。

評価

難易度
4/5
見積もり時間
3〜5日
初心者へのやさしさ
43/100
issue の種類
バグ
明瞭さ
おおむね明確
活発さ
活発
技術スタック
cpp, python
領域
backend, machine-learning

調査の方向性

Start by comparing the CUDA behavior in mlx 0.32.3 and 0.32.4, using the provided reproduction script with the stated GB10 environment and model. Check whether repeated forward passes produce different logits or projection outputs, then trace the relevant CUDA changes between versions. Done when the reproducibility behavior and whether the 0.32.4 change accounts for it are established.

索引モデルが issue の本文から書いたものです。

説明

Describe the bug

With mlx / mlx-cuda-13 0.32.3 on an NVIDIA GB10, running the same forward pass of a 4-bit quantized
Qwen2.5-7B-Instruct several times in one process, on the same input, does not always give bitwise-equal outputs.
With 0.32.4 on the same machine, every repeat in every process we ran is bitwise equal. We have not identified the
cause; we report the observations in case they are useful (for example as a regression test), and to ask whether
run-to-run bitwise reproducibility of the CUDA backend is an intended property going forward.

To reproduce

Script below (needs only mlx, mlx-lm, numpy; synthetic token ids). It loads the Hugging Face checkpoint, quantizes it
with quantize_model(model, config, 64, 4) and, for sequence lengths 512, 2048, 4096, 4097 and 6144, runs the full
forward pass 5 times (one mx.eval + mx.synchronize() per repeat), recording the input of
layers[1].self_attn.k_proj, the k_proj/v_proj outputs and the last-position logits. It then runs only the two
projections 5 times on the input captured in the first repeat.

MLX_ENABLE_CACHE_THRASHING_CHECK=0 MLX_CUDA_GRAPH_CACHE_SIZE=100 \
  python mlx_cuda_repeat_check.py --model-dir Qwen2.5-7B-Instruct --cache-limit-gb 0.5   # HF revision a09a35458c70
mlx_cuda_repeat_check.py 
"""Run the same forward pass of a 4-bit quantized LLM several times and check whether the outputs are bitwise equal.

Needs only mlx, mlx-lm and numpy. The input is synthetic token ids (numpy default_rng). For each sequence length:
- "full": the whole forward pass, `--repeats` times (one mx.eval + mx.synchronize per repeat); records the input of
  layers[1].self_attn.k_proj, the k_proj/v_proj outputs and the last-position logits;
- "op": only k_proj and v_proj of layer 1, `--repeats` times, on the input captured in the first full repeat.
Prints, per length and tensor, how many repeats differ from the first one and the fraction of differing elements.

    python mlx_cuda_repeat_check.py --model-dir Qwen2.5-7B-Instruct [--clear-cache] [--cache-limit-gb 0]
"""
import argparse
import json
from pathlib import Path

import mlx.core as mx
import numpy as np
from mlx_lm.utils import load_model, quantize_model

LAYER = 1


def host(a: mx.array) -> np.ndarray:
    return np.array(a.view(mx.uint16) if a.dtype == mx.bfloat16 else a)


def summary(repeats: list[dict[str, np.ndarray]]) -> dict[str, dict[str, object]]:
    first = repeats[0]
    return {key: {"repeats_differing": sum(not np.array_equal(r[key], first[key]) for r in repeats[1:]),
                  "max_fraction_of_elements_differing": max(float((r[key] != first[key]).mean()) for r in repeats)}
            for key in first}


def run(model: object, vocab: int, length: int, args: argparse.Namespace) -> dict[str, object]:
    ids = mx.array(np.random.default_rng(args.seed + length).integers(0, vocab, size=(1, length)))
    attn = model.model.layers[LAYER].self_attn
    k_proj, v_proj = attn.k_proj, attn.v_proj
    seen: dict[str, mx.array] = {}

    def wrap(name: str, layer: object) -> object:
        def call(x: mx.array) -> mx.array:
            seen["x"], seen[name] = x, layer(x)
            return seen[name]
        return call

    attn.k_proj, attn.v_proj = wrap("k_proj", k_proj), wrap("v_proj", v_proj)
    full = []
    try:
        for _ in range(args.repeats):
            if args.clear_cache:
                mx.clear_cache()
            logits = model(ids)[0, -1]
            mx.eval(logits, *seen.values())
            mx.synchronize()
            full.append({"logits": host(logits.astype(mx.float32))} | {k: host(v) for k, v in seen.items()})
    finally:
        attn.k_proj, attn.v_proj = k_proj, v_proj
    x = mx.array(full[0]["x"]).view(mx.bfloat16) if full[0]["x"].dtype == np.uint16 else mx.array(full[0]["x"])
    op = []
    for _ in range(args.repeats):
        if args.clear_cache:
            mx.clear_cache()
        ys = {"k_proj": k_proj(x), "v_proj": v_proj(x)}
        mx.eval(*ys.values())
        mx.synchronize()
        op.append({k: host(v) for k, v in ys.items()})
    return {"tokens": length, "full": summary(full), "op": summary(op)}


def main() -> None:
    p = argparse.ArgumentParser()
    p.add_argument("--model-dir", type=Path, required=True)
    p.add_argument("--lengths", default="512,2048,4096,4097,6144")
    p.add_argument("--repeats", type=int, default=5)
    p.add_argument("--seed", type=int, default=0)
    p.add_argument("--clear-cache", action="store_true", help="mx.clear_cache() before every repeat")
    p.add_argument("--cache-limit-gb", type=float, help="mx.set_cache_limit before loading")
    args = p.parse_args()
    if args.cache_limit_gb is not None:
        mx.set_cache_limit(int(args.cache_limit_gb * 2**30))
    model, config = load_model(args.model_dir)
    quantize_model(model, config, 64, 4)
    mx.eval(model.parameters())
    print(json.dumps({"mlx": mx.__version__, "device": mx.default_device().__repr__()}))
    for length in map(int, args.lengths.split(",")):
        print(json.dumps(run(model, config["vocab_size"], length, args)), flush=True)


if __name__ == "__main__":
    main()

Observed with 0.32.3 (14 processes × 5 lengths × 5 repeats)

  • Last-position logits: in every process, at least one length had a repeat that differed from the first repeat (46
    of 70 process/length combinations). When they differ, 78–89% of the 152,064 logits differ, max abs difference
    0.09–0.29.
  • In 10 of 70 combinations (mostly 6144 tokens; also 2048–4097 tokens with MLX_CUDA_USE_CUDNN_SDPA=0), the input
    of layer 1 k_proj and the k_proj/v_proj outputs also differ between repeats: 0.03–0.2% of the elements.
  • The projections alone on a fixed, already evaluated input were equal in all repeats, with one exception: one
    repeat of k_proj at 4096 tokens in one process with MLX_USE_CUDA_GRAPHS=0 MLX_CUDA_USE_CUDNN_SDPA=0.
  • Still observed with: MLX_USE_CUDA_GRAPHS=0; MLX_CUDA_USE_CUDNN_SDPA=0; both; mx.clear_cache() before each
    repeat; mx.set_cache_limit(0); allocating and releasing NaN-filled buffers before each repeat (no NaN appeared in
    any output).
  • The 0.32.3 logits equal the 0.32.4 logits in 42 of 70 repeats at 512 tokens, and in 5 of 70 at 6144 tokens.

Observed with 0.32.4 (same machine and driver; 7 processes, graphs on and off)

  • Every repeat in every process is bitwise equal, and the outputs are identical across processes.
  • The same values with the NVIDIA wheels that 0.32.4 installs (cublas 13.10.0.20, cudnn 9.28.0.42, cuda-runtime
    13.5.37) and with the NVIDIA wheels pinned to the 0.32.3 set below.
  • In LoRA fine-tuning on non-public data, two processes with the same seed diverged with 0.32.3; with 0.32.4 a
    20-step trace (CUDA graphs and cuDNN SDPA on) was bitwise equal across two processes; two 100-step runs were also
    bitwise equal (losses, gradients, weights and optimizer state at every step).

Environment

  • NVIDIA GB10 (compute capability 12.1), driver 580.178.04, Linux aarch64, Python 3.13.18
  • mlx 0.32.3 / mlx-cuda-13 0.32.3 and mlx 0.32.4 / mlx-cuda-13 0.32.4; mlx-lm 0.32.0
  • NVIDIA wheels (0.32.3 set): nvidia-cublas 13.8.2.12, nvidia-cudnn-cu13 9.27.0.42, nvidia-cuda-runtime 13.4.92,
    nvidia-cuda-nvrtc 13.4.92, nvidia-nvjitlink 13.4.92
  • MLX_ENABLE_CACHE_THRASHING_CHECK=0, MLX_CUDA_GRAPH_CACHE_SIZE=100, mx.set_cache_limit(0.5 GB) unless stated
  • We also saw run-to-run differences in LoRA fine-tuning with 0.32.3 on an RTX 5090 (compute capability 12.0); we did
    not run this script there.

Questions

  1. Is bitwise run-to-run reproducibility of the CUDA backend (same input, same process or new process) an intended
    property from 0.32.4 on?
  2. Is there a change in 0.32.4 that is expected to affect this? We would like to know whether 0.32.4 removes the
    behaviour or makes it less likely.
主要言語
C++
スター
29.1k
フォーク
2.4k
平均マージ
3日 23時間
マージ済み PR（30日）
90

環境構築

はじめの一歩

  1. issue を最後まで読み、次にプロジェクトのコントリビューションガイドを読みます。
  2. 着手することを issue にコメントします — 二人が同じ作業をするのを防げます。
  3. リポジトリをフォークし、ブランチを切って変更します。
  4. issue 番号を参照したプルリクエストを送ります。

ml-explore/mlx のほかの issue

ml-explore/mlx の issue をすべて見る

似ている issue

C++ の issue をもっと見る

新しい issue をメールで受け取る

初心者向けの GitHub issue を短くまとめたダイジェスト。