From 77e3b0ac7a09655b1cf855f48b98bb8b0cb85ae5 Mon Sep 17 00:00:00 2001 From: Jesse Gross Date: Tue, 25 Aug 2026 12:04:37 -0700 Subject: [PATCH] mlxrunner: avoid Metal GPU timeouts when loading models from slow storage Model load code eagerly evaluated every weight fold (expert stacking, gather transposes, gate/up fusing) as it was built, with the folds running on the GPU against lazily loaded tensors: Metal committed command buffers that waited on file reads, and macOS kills command buffers that stall too long, so loading a large model from a slow volume aborted with "Command buffer execution failed". The eager evals also kept every layer's fold sources alive until the post-load sweep, transiently holding roughly twice the expert weights on MoE models. Build the folds lazily and let the runner's weight eval run them, and on Metal materialize the loaded tensors with CPU reads before any weight graph exists: no command buffer is ever committed waiting on file data, at any storage speed, and fold sources free as their folds execute. CUDA loads read at dispatch and skip the pre-pass. Models no longer evaluate weights at load; on Metal, tensors the model does not retain are now read before the sweep frees them. Measured on an M5 Max, warm page cache, greedy outputs bit-identical: before after nemotron-3.5-lightning:30b-mlx 1.9s 39.7GiB 1.45s 24.7GiB qwen3.6:35b-mlx 1.27s 22.5GiB 1.1-1.2s 22.4GiB nemotron, reads at ~60MB/s aborts in 6s loads in 346s Fixes #17902 --- x/mlxrunner/runner.go | 8 ++++++++ x/models/cohere2_moe/cohere2_moe.go | 5 +---- x/models/gemma3/gemma3.go | 15 --------------- x/models/gemma4/gemma4.go | 7 +------ x/models/laguna/laguna.go | 16 ++++------------ x/models/nemotron_h/nemotron_h.go | 18 ++++-------------- x/models/qwen3_5/qwen3_5.go | 14 +++----------- 7 files changed, 21 insertions(+), 62 deletions(-) diff --git a/x/mlxrunner/runner.go b/x/mlxrunner/runner.go index f9b63510c..2db3d0f8b 100644 --- a/x/mlxrunner/runner.go +++ b/x/mlxrunner/runner.go @@ -4,6 +4,7 @@ import ( "context" "errors" "log/slog" + "maps" "net" "net/http" "slices" @@ -73,6 +74,13 @@ func (r *Runner) Load(modelName string) error { return err } + // On Metal, materialize the loaded tensors with CPU reads before any + // weight graph exists, so the weight eval never commits a command buffer + // that waits on file data. CUDA loads read at dispatch and need no pre-pass. + if mlx.MetalIsAvailable() { + mlx.Eval(slices.Collect(maps.Values(tensors))...) + } + // Assign weights to model (model-specific logic). Target and draft weights // must be loaded before sweeping so tensors from a combined manifest are // not discarded before the draft model can retain them. diff --git a/x/models/cohere2_moe/cohere2_moe.go b/x/models/cohere2_moe/cohere2_moe.go index 6dff3aa51..64f5d271f 100644 --- a/x/models/cohere2_moe/cohere2_moe.go +++ b/x/models/cohere2_moe/cohere2_moe.go @@ -401,10 +401,7 @@ func transposeExpertWeightForGatherMM(w *mlx.Array) *mlx.Array { if w == nil || !w.Valid() || w.NumDims() != 3 { return w } - t := mlx.Transpose(w, 0, 2, 1) - cloned := t.Clone() - mlx.Eval(cloned) - return cloned + return mlx.Transpose(w, 0, 2, 1).Clone() } // loadStackedProjection returns expert weights already stacked as a single 3D diff --git a/x/models/gemma3/gemma3.go b/x/models/gemma3/gemma3.go index e3ef290b3..4d3917e48 100644 --- a/x/models/gemma3/gemma3.go +++ b/x/models/gemma3/gemma3.go @@ -207,11 +207,6 @@ func precomputeGemmaScaledWeights(m *Model) { m.NormScaled = mlx.AddScalar(m.Norm.Weight, 1.0) } - var scaled []*mlx.Array - if m.NormScaled != nil { - scaled = append(scaled, m.NormScaled) - } - for _, layer := range m.Layers { if layer == nil || layer.Attention == nil { continue @@ -219,34 +214,24 @@ func precomputeGemmaScaledWeights(m *Model) { if layer.InputNorm != nil { layer.InputNormScaled = mlx.AddScalar(layer.InputNorm.Weight, 1.0) - scaled = append(scaled, layer.InputNormScaled) } if layer.PostAttnNorm != nil { layer.PostAttnNormScaled = mlx.AddScalar(layer.PostAttnNorm.Weight, 1.0) - scaled = append(scaled, layer.PostAttnNormScaled) } if layer.PreFFNorm != nil { layer.PreFFNormScaled = mlx.AddScalar(layer.PreFFNorm.Weight, 1.0) - scaled = append(scaled, layer.PreFFNormScaled) } if layer.PostFFNorm != nil { layer.PostFFNormScaled = mlx.AddScalar(layer.PostFFNorm.Weight, 1.0) - scaled = append(scaled, layer.PostFFNormScaled) } if layer.Attention.QNorm != nil { layer.Attention.QNormScaled = mlx.AddScalar(layer.Attention.QNorm.Weight, 1.0) - scaled = append(scaled, layer.Attention.QNormScaled) } if layer.Attention.KNorm != nil { layer.Attention.KNormScaled = mlx.AddScalar(layer.Attention.KNorm.Weight, 1.0) - scaled = append(scaled, layer.Attention.KNormScaled) } } - - if len(scaled) > 0 { - mlx.Eval(scaled...) - } } func newModel(root *model.Root) (base.Model, error) { diff --git a/x/models/gemma4/gemma4.go b/x/models/gemma4/gemma4.go index c52c3e9bd..f390289b3 100644 --- a/x/models/gemma4/gemma4.go +++ b/x/models/gemma4/gemma4.go @@ -185,9 +185,7 @@ func transposeForGatherMM(w *mlx.Array) *mlx.Array { if w == nil || !w.Valid() || w.NumDims() != 3 { return w } - t := mlx.Transpose(w, 0, 2, 1).Clone() - mlx.Eval(t) - return t + return mlx.Transpose(w, 0, 2, 1).Clone() } // collectExpertProjection collects per-expert tensors, stacks them, and @@ -242,15 +240,12 @@ func collectExpertProjection(tensors map[string]*mlx.Array, cfg *TextConfig, pre } stacked := mlx.Stack(weights, 0).Clone() - mlx.Eval(stacked) out := &stackedExpertResult{Weight: stacked, Bits: bits, GroupSize: groupSize, Mode: mode} if len(scales) == len(weights) { out.Scales = mlx.Stack(scales, 0).Clone() - mlx.Eval(out.Scales) } if len(biases) == len(weights) { out.Biases = mlx.Stack(biases, 0).Clone() - mlx.Eval(out.Biases) } return out } diff --git a/x/models/laguna/laguna.go b/x/models/laguna/laguna.go index 8940e93b9..07751c70e 100644 --- a/x/models/laguna/laguna.go +++ b/x/models/laguna/laguna.go @@ -590,18 +590,14 @@ func stackAndClone(parts []*mlx.Array) *mlx.Array { if len(parts) == 0 { return nil } - stacked := mlx.Stack(parts, 0).Clone() - mlx.Eval(stacked) - return stacked + return mlx.Stack(parts, 0).Clone() } func transposeExpertWeightForGatherMM(w *mlx.Array) *mlx.Array { if w == nil || !w.Valid() || w.NumDims() != 3 { return w } - t := mlx.Transpose(w, 0, 2, 1).Clone() - mlx.Eval(t) - return t + return mlx.Transpose(w, 0, 2, 1).Clone() } func transposeExpertWeightViewForGatherMM(w *mlx.Array) *mlx.Array { @@ -762,9 +758,7 @@ func fuseExpertStacks(a, b *mlx.Array, axis int) *mlx.Array { if a == nil || !a.Valid() || b == nil || !b.Valid() { return nil } - out := mlx.Concatenate([]*mlx.Array{a, b}, axis).Clone() - mlx.Eval(out) - return out + return mlx.Concatenate([]*mlx.Array{a, b}, axis).Clone() } func applyExpertGlobalScale(x, globalScale, idx *mlx.Array) *mlx.Array { @@ -1069,9 +1063,7 @@ func (m *Model) LoadWeights(tensors map[string]*mlx.Array) error { layerPrefix+".mlp.switch_mlp.e_score_correction_bias", ) if moe.EScoreCorrectionBias != nil && moe.EScoreCorrectionBias.DType() != mlx.DTypeFloat32 { - bias := moe.EScoreCorrectionBias.AsType(mlx.DTypeFloat32).Clone() - mlx.Eval(bias) - moe.EScoreCorrectionBias = bias + moe.EScoreCorrectionBias = moe.EScoreCorrectionBias.AsType(mlx.DTypeFloat32).Clone() } gateW := loadStackedProjection(tensors, cfg, useQuantizedExperts, diff --git a/x/models/nemotron_h/nemotron_h.go b/x/models/nemotron_h/nemotron_h.go index c7c86afdf..2a8bdabf7 100644 --- a/x/models/nemotron_h/nemotron_h.go +++ b/x/models/nemotron_h/nemotron_h.go @@ -354,10 +354,7 @@ func stackAndClone(parts []*mlx.Array) *mlx.Array { if len(parts) == 0 { return nil } - stacked := mlx.Stack(parts, 0) - cloned := stacked.Clone() - mlx.Eval(cloned) - return cloned + return mlx.Stack(parts, 0).Clone() } func supportsGatherQMM(mode string, bits int) bool { @@ -428,10 +425,7 @@ func transposeExpertWeightForGatherMM(w *mlx.Array) *mlx.Array { if w == nil || !w.Valid() || w.NumDims() != 3 { return w } - t := mlx.Transpose(w, 0, 2, 1) - cloned := t.Clone() - mlx.Eval(cloned) - return cloned + return mlx.Transpose(w, 0, 2, 1).Clone() } func sliceAxis(a *mlx.Array, axis int, start, stop int32) *mlx.Array { @@ -447,18 +441,14 @@ func sliceAxis(a *mlx.Array, axis int, start, stop int32) *mlx.Array { } func appendAndClone(dst, src *mlx.Array) *mlx.Array { - out := mlx.Concatenate([]*mlx.Array{dst, src}, 0).Clone() - mlx.Eval(out) - return out + return mlx.Concatenate([]*mlx.Array{dst, src}, 0).Clone() } func stackSlicesAndClone(slices []*mlx.Array) *mlx.Array { if len(slices) == 0 { return nil } - out := mlx.Stack(slices, 0).Clone() - mlx.Eval(out) - return out + return mlx.Stack(slices, 0).Clone() } func foldSharedExperts(m *SparseMoE, cfg *Config) bool { diff --git a/x/models/qwen3_5/qwen3_5.go b/x/models/qwen3_5/qwen3_5.go index f67dff614..64ca4dbbe 100644 --- a/x/models/qwen3_5/qwen3_5.go +++ b/x/models/qwen3_5/qwen3_5.go @@ -479,29 +479,21 @@ func stackAndClone(parts []*mlx.Array) *mlx.Array { if len(parts) == 0 { return nil } - stacked := mlx.Stack(parts, 0) - cloned := stacked.Clone() - mlx.Eval(cloned) - return cloned + return mlx.Stack(parts, 0).Clone() } func transposeExpertWeightForGatherMM(w *mlx.Array) *mlx.Array { if w == nil || !w.Valid() || w.NumDims() != 3 { return w } - t := mlx.Transpose(w, 0, 2, 1) - cloned := t.Clone() - mlx.Eval(cloned) - return cloned + return mlx.Transpose(w, 0, 2, 1).Clone() } func fuseExpertStacks(a, b *mlx.Array, axis int) *mlx.Array { if a == nil || !a.Valid() || b == nil || !b.Valid() { return nil } - out := mlx.Concatenate([]*mlx.Array{a, b}, axis).Clone() - mlx.Eval(out) - return out + return mlx.Concatenate([]*mlx.Array{a, b}, axis).Clone() } // fuseGateUpProjections joins gate and up stacks along the output dimension,