Skip to content

Reduce CoreML conversion peak and ANE steady-state memory#1202

Open
ChinChangYang wants to merge 11 commits into
lightvector:masterfrom
ChinChangYang:feature/coreml-conversion-memory-levers
Open

Reduce CoreML conversion peak and ANE steady-state memory#1202
ChinChangYang wants to merge 11 commits into
lightvector:masterfrom
ChinChangYang:feature/coreml-conversion-memory-levers

Conversation

@ChinChangYang

@ChinChangYang ChinChangYang commented May 30, 2026

Copy link
Copy Markdown
Contributor

Summary

Cuts memory during the on-device KataGo→CoreML conversion and while running the ANE/CoreML path, with byte-identical converter output. Three changes (first commit):

  • Non-owning weight views — the converter's weight tensors become views into the parsed model instead of owning extra FP32 copies (derived/transposed tensors keep an owned buffer), dropping redundant resident weight copies during conversion. CoreML serialization is made deterministic (SetSerializationDeterministic) so output is byte-stable.
  • Streaming parser — the KataGo model parser streams the gzip through a bounded ~1 MB refill buffer instead of decompressing the whole file into memory, preserving the existing NaN/Inf weight validation.
  • Release weights on ANEModelDesc::releaseWeights() frees the in-memory weight arrays (keeping scalar shape metadata). The Metal backend calls it on the ANE (CoreML) path after converting from the model file on disk, gated by a new ComputeContext::aneOnly flag so it only fires when every configured device is ANE — the GPU/MPSGraph path keeps its weights. Serialized under computeHandleMutex; only scalar dims are read afterward.

A follow-up commit (Enforce non-owning weight-view contract at compile time) hardens the first bullet: WeightEntry now stores a non-owning view, so its backing std::vector must outlive serialization. The rvalue overloads of addConstOp/registerWeight are = deleted so a temporary can never bind to the view path (it would dangle) — temporaries are forced through the owning addOwnedConstOp/registerOwnedWeight. No codegen change: every existing call site passes an lvalue model member, verified by a full METAL rebuild plus a negative-compile test confirming the deleted overload is selected for a temporary.

Merged: Metal GPU + CoreML/ANE transformer support (#1205)

This branch now merges #1205 (transformer-net support — SiLU, GQA, learnable RoPE, RMSNorm, SwiGLU — for the Metal GPU and CoreML/ANE paths) so the two efforts reconcile in one place. All four merge conflicts were in the katagocoreml converter and came from a single root: this PR's non-owning FloatView weight registration vs #1205's per-weight FP32 precision tiers (is_fp32). Both are kept (registerWeight stamps entry.is_fp32; WeightSerializer stores fp16 iff use_fp16 && !is_fp32; addConstOp keeps the emitConstOp refactor and passes is_fp32 = (m_weight_dtype == FLOAT32)).

The merge required one integration fix (commit Conform CoreML transformer derived consts to the owned-weight + FP32 contract). The transformer attention builder emits four function-local tensors — RoPE cos/sin tables, the rotation matrix R, and per-head out-projection weight slices. After this PR made WeightEntry::data a non-owning view, those locals must:

  1. be registered through the owning addOwnedConstOp (with std::move) so the view doesn't dangle once the build function returns (serialization runs afterwards); and
  2. carry the same per-weight FP32 marking as addConstOp, because emitConstOp declares each const's dtype as m_weight_dtype. Without (2) a derived const that lands in the attention/value-head FP32 sub-region of an FP16 model is declared FP32 but stored FP16, which CoreML/ANE rejects at load ("Metadata data type does not match requested type", BNNS error −14), crashing every FP16-ANE transformer.

addOwnedConstOp/registerOwnedWeight now thread is_fp32 = (m_weight_dtype == FLOAT32), mirroring addConstOp so stored dtype always matches the declared dtype. This also fixes the same latent mismatch for addLinearOp's transposed value-head weights. The streaming parser needed no change (the new parseTransformer* use the unchanged readFloats), and ModelDesc::releaseWeights() already covers the transformer descriptors (including ropeFreqs).

Re-synced to master (74bee865). This branch has since merged current master, which subsumes #1205 and adds later Metal/CoreML transformer fixes (attention-logit bound, mask constant, fp32 value-head error reduction, extra coreml/gpu-error checks). One conflict — the katagocoreml streaming parser readFloats (this PR's bounded-buffer streaming vs master's whole-file scan) — resolved by keeping the streaming lever. Re-verified on the merged tip 2c0a4abb: runtests/runnnlayertests, a GPU+ANE testgpuerror regression across conv + all three transformer variants, memory A/B (levers intact), and the b11c768h12 v17 net — all pass, GPU near-exact and ANE FP16 within tolerance vs Eigen.

Measured (19×19, ANE path)

Resident-memory A/B on the CoreML/ANE path (metalDeviceToUseThread0=100), cold conversion every run, before vs after this PR's three memory levers. before is this branch with the levers removed (the #1205 transformer tip 39f82f6d — no non-owning views, no streaming parser, no aneOnly release); after is HEAD. Peak RSS is the /usr/bin/time -l maximum over load+convert; idle steady-state is ps-sampled after the model is loaded and released. Numbers are GB, reproduced across two runs (<0.1% spread).

Model Peak RSS (load+convert), before → after Idle steady-state RSS, before → after
Classical 40b (b40c256) 1.50 → 0.80 GB (−47%) 0.38 → 0.25 GB (−33%)
18b nbt (b18c384nbt) 0.88 → 0.48 GB (−45%) 0.59 → 0.20 GB (−66%)
Transformer (b10c384h6nbttflrs) 1.58 → 1.37 GB (−13%) 0.66 → 0.46 GB (−31%)

The peak win is largest on the classical 40b, where the conversion transients the levers target (duplicate owning weight copies + whole-gzip decompression) dominate; the idle win is largest on the 18b nbt, which has the most resident FP32 weight for the ANE releaseWeights to free.

Why the transformer's peak reduces less (−13%), and why its absolute peak is highest despite the smallest file. Both are expected and not a regression — verified by instrumented byte-accounting at serialization time:

  • The peak levers (non-owning views + streaming parser) act on host RAM during parse/convert. The transformer gives them less to remove: (i) ~39% of its converter-side weight bytes are derived tensors — RoPE cos/sin tables (the bulk), the rotation matrix, and per-head out-projection slices — that are host-computed and therefore keep an owned buffer (addOwnedConstOp), so the non-owning-view lever can dedup ~0% of them, versus ~100% of a conv net's weights (the conv nets carry only ~0.07 MB of owned tensors total). (ii) Its 38 MB file makes the streaming parser's removed whole-decompress transient small (~42 MB, vs ~187 MB for the 40b). Combined, the levers free ~80 MB here vs ~365 MB on the 40b.
  • The ~1.4 GB absolute peak is reached after the C++ converter returns (which itself peaks at only ~0.24 GB): it is CoreML/ANE compilation of the MIL graph into an ANE program. The transformer emits ~3.4× more MIL ops than the 40b (op-level attention: per-block RoPE, GQA per-head slice/concat, 361×361 softmax), and the ANE compiler's memory scales with graph complexity, not weight bytes — so the smallest model on disk has the largest compile-time footprint. This floor sits downstream of, and is unaffected by, all three (host-side) levers.

Cross-backend parity vs an Eigen reference is unchanged on both the GPU and ANE paths.

Test plan

  • METAL build; ./katago runtests (All tests passed) + ./katago runnnlayertests (14 configurations)
  • ./katago testgpuerror vs Eigen reference — GPU and ANE paths both within tolerance, no NaN
  • ANE-path RSS A/B (above); GPU path unaffected (aneOnly false)
  • Full rungpuerrortest.sh regression vs current Eigen FP32 references: Metal GPU 37/37 pass; CoreML/ANE 31/31 supported configs pass (the 6 "failures" are the expected pre-v8 SIGABRTs — run4/grun50/g103, nets predating CoreML-converter support). The two transformer nets the merge had broken on the FP16 ANE path (b7c96h3tfrs, b7c96h6kv3qk32v16 GQA) now load and match the reference to <0.0005% winrate; convnet ANE output is byte-identical and the GPU path is unchanged.

Note

The CI cache fix that was previously bundled here has been split into its own PR: #1204 (orthogonal fix for a stale-cache link failure on the build-macos-metal job). This PR is now the memory work plus the merged #1205 transformer support.

🤖 Generated with Claude Code

Cuts memory during the on-device KataGo -> CoreML conversion and while
running the ANE/CoreML path, with byte-identical converter output:

- The converter's weight tensors become non-owning views into the parsed
  model instead of owning extra FP32 copies; derived/transposed tensors keep
  an owned buffer. This drops redundant resident weight copies during
  conversion. CoreML model serialization is made deterministic
  (SetSerializationDeterministic) so the output is byte-stable.

- The KataGo model parser streams the gzip through a bounded ~1 MB refill
  buffer instead of decompressing the whole file into memory, while
  preserving the existing NaN/Inf weight validation.

- ModelDesc gains releaseWeights(), which frees the in-memory weight arrays
  (keeping scalar shape metadata). The Metal backend calls it on the ANE
  (CoreML) path after converting from the model file on disk, gated by a new
  ComputeContext::aneOnly flag so it only fires when every configured device
  is ANE -- the GPU/MPSGraph path keeps its weights. The call is serialized
  under computeHandleMutex and only scalar dims are read afterward.

Measured on b18c384nbt (19x19) over the ANE path: idle steady-state RSS
0.59 GB -> 0.19 GB; peak (load+convert) 0.87 GB -> 0.48 GB. Cross-backend
parity vs an Eigen reference is unchanged on both the GPU and ANE paths.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
ChinChangYang and others added 3 commits May 31, 2026 17:28
WeightEntry stores a non-owning view (const float*, count) into the live
KataGoModelDesc, so the backing std::vector must outlive serialization.
addConstOp/registerWeight took the data by const& and silently stored a
pointer to it; a caller passing a temporary would bind to that const& and
leave the view dangling, read much later during serialization.

Delete the rvalue overloads of both so any such call fails to compile,
forcing temporaries through addOwnedConstOp/registerOwnedWeight (which take
ownership). Named lvalues (the model-member call sites) still bind to the
const& overload, so no existing caller changes.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Own the gzFile with a custom-deleter unique_ptr so it closes on every
exit path (normal return, exception, bad_alloc); removes the manual
try/catch+gzclose in parse() and the ordering caveat on buffer allocation.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Introduce a KataGo-local non-owning FloatView for WeightEntry::data instead
of a raw const float*/size_t pair; convert to MILBlob::Util::Span only inside
WeightSerializer, keeping the MILBlob dependency out of Operations.hpp.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
ChinChangYang and others added 3 commits May 31, 2026 21:57
The ComputeHandle member-order comment claimed that declaring
mpsGraphOnlyHandle before coremlOnlyHandle is what prevents a GPU handle
from reading freed weights. That overstates the ordering's role: within a
single ComputeHandle exactly one handle is built (mutually exclusive on
gpuIdx, enforced by the ctor's exactly-one check), and releaseWeights()
only fires on an aneOnly context where no MPSGraph handle is ever built.
Reframe the declaration order as belt-and-suspenders and point at
ComputeContext::aneOnly as the actual invariant. Comment-only change.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Replace the file-local releaseXXX free functions in desc.cpp (which
reached into each desc struct's internals from outside) with
releaseWeights() member methods on each weight-bearing struct, matching
the existing OO convention used by applyScale8ToReduceActivations() and
iterConvLayers(). Each container delegates to its members; type-erased
block dispatch is inlined with the same cast pattern those methods use.

Behavior-preserving: same set of freed vectors, same block recursion,
same metaEncoderVersion guard. ModelDesc::releaseWeights() keeps its
signature, so the metalbackend.cpp call site is unchanged.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Move the 11 leaf/container releaseWeights() definitions in desc.cpp out
of the bottom cluster (inherited from the old free-function layout) and
place each immediately after its struct's last existing method, matching
the file's per-struct grouping convention used by every other method.
ModelDesc::releaseWeights() stays put, already adjacent to its siblings.

Pure relocation: function bodies and desc.h are unchanged; only two
stray double-blank lines were normalized to single. Verified clean Metal
build, testgpuerror vs Eigen reference (g170-b6c96) at <0.0004% winrate
error, and runtests all pass.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
@ChinChangYang
ChinChangYang marked this pull request as ready for review June 1, 2026 00:29
@ChinChangYang

Copy link
Copy Markdown
Contributor Author

I reviewed and approved this.

Resolve conflicts where upstream transformer support landed alongside the
CoreML/ANE memory work:

- desc.h / desc.cpp: keep both upstream getNumParameters() and the branch's
  releaseWeights() on SGFMetadataEncoderDesc. Extend releaseWeights() to the
  new transformer block kinds (TRANSFORMER_ATTENTION/FFN) in TrunkDesc and
  NestedBottleneckResidualBlockDesc, add releaseWeights() to RMSNormLayerDesc /
  TransformerRMSNormDesc / TransformerAttentionDesc / TransformerFFNDesc, and
  also release trunkTipRMSNorm -- mirroring upstream's getNumParameters()
  coverage so an ANE-path model containing transformer blocks frees its weights
  instead of hitting ASSERT_UNREACHABLE.
- metalbackend.cpp: adopt upstream's 3-arg ComputeContext ctor (NHWC mode
  removed) while keeping the branch's context->aneOnly assignment.

Verified: METAL build; runtests; runnnlayertests; testgpuerror ANE + GPU vs
Eigen reference within tolerance, no NaN; ANE-path RSS still ~0.20 GB idle /
~0.49 GB peak (memory levers preserved).

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
@ChinChangYang
ChinChangYang marked this pull request as draft June 2, 2026 15:26
ChinChangYang and others added 2 commits June 4, 2026 12:04
…into feature/coreml-conversion-memory-levers (PR lightvector#1202)

Brings Metal GPU + CoreML/ANE transformer-net support (SiLU, GQA, learnable
RoPE, RMSNorm, SwiGLU) onto the CoreML-conversion memory branch (non-owning
FloatView weight views, streaming parser, ANE weight release).

All four conflicts were in the katagocoreml converter, from one root cause:
lightvector#1202 reworked weight registration (non-owning FloatView views,
registerOwnedWeight, deleted rvalue overloads) while lightvector#1205 added per-weight
FP32 precision tiers (is_fp32). Resolved to keep both:

- Operations.{hpp,cpp}: registerWeight stamps entry.is_fp32; keep the non-owning
  view plus registerOwnedWeight.
- WeightSerializer.cpp: keep the FloatView-based count plus per-weight
  store_fp16 = use_fp16 && !is_fp32.
- MILBuilder.cpp addConstOp: keep the emitConstOp refactor and pass
  is_fp32 = (m_weight_dtype == FLOAT32).

A follow-up commit retrofits lightvector#1205's transformer derived consts to lightvector#1202's
owned-weight + FP32-marking contract (required for the FP16 ANE path).

Verified after the follow-up fix: katago runtests and runnnlayertests pass, and
testgpuerror vs fresh Eigen FP32 references passes on convnet + all three
transformer nets on both the Metal GPU and CoreML/ANE paths.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
…contract

The transformer attention builder emits four function-local std::vector<float>
tensors: RoPE cos/sin tables, the rotation matrix R, and per-head out-projection
weight slices. After merging the transformer support onto the FloatView branch,
these needed two fixes:

1. Dangling view. lightvector#1202 made WeightEntry::data a non-owning FloatView, so
   addConstOp registers a view whose backing buffer must outlive serialization.
   These locals were passed to addConstOp and would dangle once the build
   function returns (serialization runs afterwards). Route them through
   addOwnedConstOp so KataGoOps owns the buffer until serialization. (Under
   lightvector#1205's owning WeightEntry they were copied, so this only surfaces post-merge.)

2. dtype mismatch. emitConstOp declares each const's dtype as m_weight_dtype, but
   addOwnedConstOp / registerOwnedWeight stored at the global mode (is_fp32
   hardcoded false). In an FP16 model these derived consts land in the attention /
   value-head FP32 sub-region (m_weight_dtype == FLOAT32), so they were declared
   FP32 but stored FP16. CoreML/ANE then rejects the model at load ("Metadata data
   type does not match requested type", BNNS error -14), which SIGABRT'd every
   FP16 ANE transformer. Thread is_fp32 through registerOwnedWeight and have
   addOwnedConstOp pass is_fp32 = (m_weight_dtype == FLOAT32), mirroring addConstOp
   so the stored dtype always matches the declared dtype. This also fixes the same
   latent mismatch for addLinearOp's transposed value-head weights.

Verified with testgpuerror against fresh Eigen FP32 references: b7c96h3tfrs and
b7c96h6gqa, which previously SIGABRT'd on the FP16 ANE path, now load and match
to <0.0005% winrate; convnet ANE output is byte-identical and the Metal GPU path
is unchanged. katago runtests and runnnlayertests also pass.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
@ChinChangYang

Copy link
Copy Markdown
Contributor Author

Completed. Plan to set this PR to ready when #1205 is merged.

ChinChangYang added a commit to ChinChangYang/KataGo that referenced this pull request Jun 4, 2026
…contract

The transformer attention builder emits four function-local std::vector<float>
tensors: RoPE cos/sin tables, the rotation matrix R, and per-head out-projection
weight slices. After merging the transformer support onto the FloatView branch,
these needed two fixes:

1. Dangling view. lightvector#1202 made WeightEntry::data a non-owning FloatView, so
   addConstOp registers a view whose backing buffer must outlive serialization.
   These locals were passed to addConstOp and would dangle once the build
   function returns (serialization runs afterwards). Route them through
   addOwnedConstOp so KataGoOps owns the buffer until serialization. (Under
   lightvector#1205's owning WeightEntry they were copied, so this only surfaces post-merge.)

2. dtype mismatch. emitConstOp declares each const's dtype as m_weight_dtype, but
   addOwnedConstOp / registerOwnedWeight stored at the global mode (is_fp32
   hardcoded false). In an FP16 model these derived consts land in the attention /
   value-head FP32 sub-region (m_weight_dtype == FLOAT32), so they were declared
   FP32 but stored FP16. CoreML/ANE then rejects the model at load ("Metadata data
   type does not match requested type", BNNS error -14), which SIGABRT'd every
   FP16 ANE transformer. Thread is_fp32 through registerOwnedWeight and have
   addOwnedConstOp pass is_fp32 = (m_weight_dtype == FLOAT32), mirroring addConstOp
   so the stored dtype always matches the declared dtype. This also fixes the same
   latent mismatch for addLinearOp's transposed value-head weights.

Verified with testgpuerror against fresh Eigen FP32 references: b7c96h3tfrs and
b7c96h6gqa, which previously SIGABRT'd on the FP16 ANE path, now load and match
to <0.0005% winrate; convnet ANE output is byte-identical and the Metal GPU path
is unchanged. katago runtests and runnnlayertests also pass.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
(cherry picked from commit 8481a94)
ChinChangYang added a commit to ChinChangYang/KataGo that referenced this pull request Jun 4, 2026
The cherry-picked per-struct releaseWeights() refactor (44342a3/98b17ebb)
predates this branch's MLX transformer port, so it only added releaseWeights()
to the non-transformer descriptors. Extend the coverage to the transformer
descriptors present on this branch (RMSNormLayerDesc, TransformerRMSNormDesc,
TransformerAttentionDesc incl. ropeFreqs, TransformerFFNDesc) and handle
TRANSFORMER_ATTENTION_BLOCK_KIND / TRANSFORMER_FFN_BLOCK_KIND plus
trunkTipRMSNorm in the trunk release walk. Without this, releasing weights on
a transformer model would hit ASSERT_UNREACHABLE. This makes desc.cpp/desc.h
byte-identical to the lightvector#1202 feature branch.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
ChinChangYang added a commit to ChinChangYang/KataGo that referenced this pull request Jun 4, 2026
Port the lightvector#1202 ANE steady-state memory lever to the MLX backend. Add
ComputeContext::aneOnly, set in createComputeContext when every configured
device index is MLX_MUX_ANE, and call ModelDesc::releaseWeights() in
convertAndCreateCoreMLOnlyHandleMLX after the model has been converted to
CoreML on disk.

Safe because: the ANE path re-reads the model from modelPath (not the
in-memory weight arrays); the ComputeHandle ctor takes the MLX_MUX_ANE
early-return before building any MLX/GPU model (the only weight-array
consumer); only scalar dims are read afterward, which releaseWeights()
preserves; and it runs under computeHandleMutex. Mirrors the Metal backend's
aneOnly release. GPU path unaffected (aneOnly is false whenever any thread
uses MLX_MUX_GPU).

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
uestccokey pushed a commit to uestccokey/KataGo that referenced this pull request Jul 11, 2026
Brings origin/mlx-backend-squash (the MLX backend + transformer CoreML
support + refined lightvector#1202 memory levers) onto feature/mlx-app-migration.
App still builds on USE_METAL_BACKEND (the new mlx*.cpp are on disk but
not yet in the Xcode project); macOS + iOS-sim build green.

Conflict resolution (12 files):
- katagocoreml converter (6 files) + desc.cpp: took THEIRS. ios-dev had
  the earlier lightvector#1202 memory-lever lineage (raw gzFile, ptr+count
  WeightEntry, free-function releaseWeights); PR lightvector#1199 has the refined
  superset (RAII GzHandle, FloatView+is_fp32, per-struct releaseWeights,
  transformer parsing). Converter public API is unchanged, so the iOS
  CoreML bridge callers are unaffected.
- metalbackend.cpp: UNION — kept ios-dev's Task-19 cache-aware CoreML
  bridge (invokeCoreMLBridge + legacy fallback) AND adopted PR's
  context->aneOnly gate around releaseWeights().
- main.cpp: union of the USE_COREML_BACKEND and USE_MLX_BACKEND branches.
- LICENSE: union (abseil + protobuf + onnx all vendored).
- python/export_model_pytorch.py: took theirs (training-only).

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Sync PR lightvector#1202 (CoreML conversion memory levers) with current master, which
carries newer Metal GPU + CoreML/ANE transformer work (attention-logit bound,
mask constant, fp32 value-head error reduction, more coreml/gpu-error checks).

One conflict, in the katagocoreml streaming parser (KataGoParser.cpp readFloats):
master added whole-file m_buffer scanning/bounds robustness, but PR lightvector#1202 replaced
whole-file decompression with a bounded ~1 MB streaming refill buffer (m_refill /
readExact) — the memory lever this PR delivers, and the streaming design the merged
header now declares. Resolved by keeping the streaming readFloats (all three hunks
HEAD-side) and pruning master's m_buffer bounds-check; kept <limits> (used by an
auto-merged numeric_limits overflow check) and dropped unused <fstream>.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01BRLDbokeoLNwMnxgRbgQvS
@ChinChangYang
ChinChangYang marked this pull request as ready for review July 24, 2026 11:13
@ChinChangYang

Copy link
Copy Markdown
Contributor Author

Re-verified on the merged tip 2c0a4abb: runtests/runnnlayertests, a GPU+ANE testgpuerror regression across conv + all three transformer variants, memory A/B (levers intact), and the b11c768h12 v17 net — all pass, GPU near-exact and ANE FP16 within tolerance vs Eigen.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant