Epic: Apple-Silicon GPU acceleration for AMICATorchNG - #83
Merged
Conversation
* Guard float32 0/0 in mu-denominator ufp/y (#75) float32 diverged to NaN on the full 30504-sample data across every seed (Newton on and off), while float64 converged. Root cause: the mu denominator sbeta*sum(ufp/y) (ufp=u*fp); at a sample sitting on a mixture mean, float32 rounds the scaled activation y to exactly 0, and fp(0)=0 for every family, so that term is 0/0=NaN and one NaN summand poisons dmu_d. float64 never rounds y to exactly 0. Diagnostics ruled out summation precision (accumulating block partials in float64, and Neumaier compensated summation, did not help) and the density / responsibilities (float64 there did not help either). Only guarding the ufp/y division does. The guard (ufp / where(y==0, 1, y), contributing 0 for the measure-zero sample) is a no-op in float64 (y is never exactly 0), so single-model #24 parity stays bit-identical, and it needs no float64, so it also stabilizes the MPS/float32 path (Apple GPUs have no FP64) -- epic #74 Phase A. Tested: float32 now converges across 5 seeds x Newton on/off on the real sample EEG, matching the float64 LL to ~5 significant digits; full non-slow torch suite green (124 passed), including the NumPy-parity and byte-identity anchors. Docs updated (perf_findings, mps_pathways, AGENTS, benchmark_gpu). * Address PR #78 review findings - Correct the guard comment: the true ufp/y limit is NOT 0 (nonzero constant at rho=2, integrable singularity diverging for rho<2), so the guard drops an unrepresentable singular term, not a removable zero. Measured: it fires <=1 sample/iteration on the sample EEG (5 of 150 iters), and float32 still matches the float64 LL to ~5 sig digits -- a bounded, negligible bias. - Scope the "fp(0)=0 for every family" claim to the supported rho>=1 (for rho<1 the GG fp is itself NaN at 0; out of scope, default minrho=1.0). - test: reference AMICATorchNG._DEGENERATE_STOP_REASONS instead of duplicating the ("nan_ll","singular_ll") tuple; shorten the float64-tracking quality check to 100 iters (it tracks float64 from iter 1; the 150-iter sweep remains the regression guard). - docs: finish the mps_pathways.md update -- intro to past tense, and Pathways B and C no longer gate on Pathway A as unmet (it is done, #75).
* Add MLX backend AMICAMLXNG (v1 MVP, #76) Epic #74 Phase C: an optional Apple-Silicon GPU backend that runs the natural-gradient EM E/M-step on the Apple GPU via MLX. v1 MVP scope: single-model, generalized Gaussian (pdftype=0), natural gradient; Newton, the other PDF families, component sharing, multi-model, outlier rejection and save/load are fast-follows (rejected with NotImplementedError). Hybrid design forced by MLX 0.32: the GPU has no float64 and all mlx.core.linalg is CPU-only. So the elementwise/matmul hot path runs on the GPU in float32 (with the Phase A ufp/y divide-by-zero guard carried over), while inv(A)/slogdet(W) run on the CPU stream, hoisted to once per iteration (measured ~42 us/iter vs a ~13 ms GPU E-pass). lgamma/digamma are absent in MLX, so the GG normalizer and the rho update are computed host-side via SciPy on the small rho array. Exactly one mx.eval per iteration bounds the lazy graph. MLX is an optional dependency (Apple Silicon only): mlx_impl is imported lazily so import pyAMICA never requires it, and the mlx extra keeps it out of the default install, so CI (ubuntu) skips the MLX tests. Tested (real sample EEG, Apple GPU): per-block sufficient stats match the NumPy float64 reference to rtol ~1e-4; the converged LL matches the PyTorch float32 backend to ~2e-6 (< 1e-2 gate); full-data fit is finite and non-degenerate. Full non-slow torch suite unaffected (124 passed). Whether MLX beats CPU/MPS is Pathway B's question. * Address PR #79 review findings Silent-failure review (parity gaps vs the torch backend's guards): - fit() now checks parameter finiteness each iteration (A/mu/alpha/beta/rho) and stops with a degenerate "nan_params" reason + nan final_ll, so a final-iteration M-step blow-up cannot complete as max_iter with silently NaN params (the torch backend backstops this in state_dict; the MLX MVP had no backstop). - Reset a NaN rho update to rho0 with a warning, matching AMICATorchNG, so it does not poison the lgamma table and every subsequent E-step unattributably. - doscaling: multiply mu by safe_scale, not raw scale, so a zero-norm (collapsed) column leaves mu unchanged like A/beta (raw scale silently zeroed mu). - Comment that a singular-A LinAlg error surfaces at the fit() mx.eval, not in _update_unmixing_matrices (lazy graph). - transform() now raises a clear NotImplementedError instead of AttributeError. Comment/doc review (accuracy): - _preprocess: MLX CPU eigh is full float64 (only the GPU stream is unsupported), so using numpy is a code-reuse choice, not a precision workaround. - Fix the A-update citation (core.py:1176-1184, not 1156-1164) and tighten the _score_gg/_log_pdf_gg/_get_block_updates line ranges. - Module docstring: state accurately which deferrals raise NotImplementedError vs are simply absent. - mps_pathways.md: reword the leftover "Cost: ... v2 option" paragraph that contradicted the "MVP LANDED" status above it. MLX tests still green (4 passed). * Omit optional MLX backend from coverage The MLX backend (pyAMICA/mlx_impl) requires MLX + an Apple GPU, so CI (ubuntu, no mlx) cannot exercise it and measured it at 0%, dropping total coverage to 73.7% and failing the --cov-fail-under=80 gate. It is covered locally by tests/mlx_tests/ on Apple hardware; omit it from the coverage metric (like the subprocess-only CLI entrypoint). Coverage returns to 82.0%.
* Add cross-platform dimension-sweep benchmark (#77) Epic #74 Phase B: measure both results and performance for every AMICA backend across CPU/MPS/CUDA/MLX on real 70-channel EEG (OpenNeuro ds002718 sub-002), sweeping channel count and n_models with component sharing on/off. benchmarks/benchmark_dimsweep.py auto-detects the host's backends (numpy/torch-cpu/torch-mps/torch-cuda/mlx), records ms/iteration (warmed, min-of-repeats) and converged LL at matched settings, and emits JSON so a Mac run and a CUDA host run merge via --report. MLX (single-model MVP) auto-excludes from multi-model/sharing configs. Data is fetched from OpenNeuro, not committed (README_dimsweep.md; benchmarks/data + result JSONs gitignored). Findings (.context/issue-77/benchmark_findings.md): - MLX is the Apple-GPU win: ~15-25 ms/it, flat across 16-70 ch, ~7x over torch-CPU and faster than an RTX 4090 (CUDA ~36 ms) at EEG scale. - PyTorch-MPS never wins (162-255 ms/it, at or worse than CPU) -- use MLX, not device="mps", on Apple hardware. - Results agree across cpu/mps/cuda/mlx and f32/f64 to ~3 digits on real data. - Multi-model has no GPU path yet (MLX MVP is single-model; MPS loses), so multi-model MLX is the top fast-follow. Docs updated: mps_pathways.md Pathway B (DONE), AGENTS.md perf note. * Match numpy block_size in the benchmark (PR #80 review) The numpy adapter didn't pass block_size, so AMICA_NumPy used its default (block_size=128, do_opt_block=True) and ran an in-fit wall-clock block-size auto-tune search inside the timed region -- unmatched vs torch/mlx (which pin block_size=512) and timing pollution that ~doubled numpy's ms/it. Pass block_size=512, do_opt_block=False. Corrected numpy: 142-622 ms/it (was 310-1350), so numpy is ~10-25x slower than MLX, not 50-70x. Results (LL) unchanged (block_size does not affect the accumulated sufficient stats). Findings table updated; the qualitative conclusions (MLX ~7x win, MPS never wins) are unaffected.
* Add multi-model support to AMICAMLXNG (#81) Epic #74 Phase D: extend the MLX Apple-GPU backend to n_models > 1, porting the multi-model machinery from AMICATorchNG -- comp_list indirection, per-model W=inv(A[:,comp_list]) + slogdet (CPU stream), cross-model responsibilities v=softmax(logV), the per-model exact-EM bias c update, and the gm-weighted A-update scattered through comp_list. Single-model (#76) stays byte-for-byte unchanged (the loop runs once, gm=1, identity comp_list). Component sharing remains a fast-follow. Validated on real sample EEG: multi-model (n_models=2) one-iteration sufficient stats match AMICATorchNG float32 to float32 precision, and the converged LL matches to ~1e-5. New tests/mlx_tests test_multimodel_matches_torch_float32; the single-model byte-identity and torch-float32 tests still pass. Benchmark (benchmark_dimsweep.py) now runs MLX in the multi-model configs: MLX wins multi-model too -- ~38-45 ms/it, ~5x over torch-CPU (MPS still loses), matching the ~7x single-model win. Findings updated (.context/issue-77). * Address PR #82 review findings Silent-failure review: - Port the dead-model warning to the multi-model c update (a zero-responsibility model kept its prior c but was surfaced-free; matches AMICATorchNG). - Add gm and c to the per-iteration mx.eval and the params_finite guard: they are new multi-model state that feeds the next E-step, so a last-iteration blow-up must be caught (nan_params) and c's cross-iteration dependency must be materialized each iteration rather than growing the lazy graph unbounded. Comment/doc review: - Fix two wrong AMICATorchNG citations (c-update 1083-1092, A-update 1231-1247) and the stale _lgamma_table shape comment ((n_mix, n_comps)). - Drop the leftover "MVP"/"v1" framing that contradicted the new multi-model scope (module/class docstrings, transform(), the module title now #76/#81). Test review: - Expand the multi-model one-iteration stat comparison to every accumulator, including dWtmp (the gm-weighted A-update input, transposed to align the (n_models,n,n) vs (n,n,n_models) layouts) and the scattered mixture stats. - Add a direct multi-model c-update check (responsibility-weighted data mean, and the two models' c must differ) and a single-model c==0 regression guard. mlx suite green (5 passed); the port is unchanged (code reviewer verified single-model byte-identity independently).
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## main #83 +/- ##
==========================================
+ Coverage 79.43% 79.44% +0.01%
==========================================
Files 14 14
Lines 1974 1975 +1
Branches 337 337
==========================================
+ Hits 1568 1569 +1
Misses 301 301
Partials 105 105
🚀 New features to boost your workflow:
|
Epic #74 final integration review (PR #83) found Phase D's multi-model MLX (#81) never propagated to docs outside its own commit, and a few untested MLX paths. Docs (the shipped state was inconsistent): update the 7 stale "single-model MVP" claims to "single- and multi-model" across AGENTS.md (architecture map + perf note), .context/mps_pathways.md (Pathway C status + fast-follow lists), benchmarks/README_dimsweep.md (backend list + example comment), pyAMICA/mlx_impl/__init__.py; add mlx_impl to README.md's file structure. The inline code docstrings were already correct; only these external docs drifted. Tests: add three MLX tests caught by the review -- the degenerate-stop path (NaN data -> nan_ll stop + nan final_ll, the machinery behind the new nan_params guard), a direct shared-seed init-parity pin vs AMICATorchNG (single- and multi-model), and the multi-model dead-model c guard (dgm[h]==0 keeps prior c). mlx suite green (8 passed). No code change; docs + tests only.
Member
Author
Integration review (3 Sonnet reviewers, worktree-isolated)The epic composes cleanly on every functional axis -- the code reviewer confirmed the Phase A Findings addressed in 1b09766 (docs + tests only, no code change): Cross-phase doc drift (code + doc reviewers): Phase D's multi-model MLX (#81) never propagated to docs outside its own commit, leaving 7 stale "single-model MVP" claims. Fixed across Test hardening (test reviewer): added the degenerate-stop path (NaN data -> Noted, not blocking (follow-ups): the non-default PDF families aren't full-data float32-regression-tested (the guard is family-independent and the |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Epic #74 delivers Apple-Silicon GPU acceleration for AMICA, governed by one hard constraint the research (#72) surfaced: Apple GPUs have no FP64 hardware (neither PyTorch MPS nor MLX offers GPU float64), so the whole path was gated on a numerically stable float32 AMICA. Four phases, each a reviewed + squash-merged PR:
0/0in the mu denominatorufp/ywhen float32 rounds an activation to exactly 0. A one-line divide-by-zero guard fixes it -- bit-identical in float64 (Verify NG init-basin vs Newton bit-parity via Fortran load_* init-matching #24 Fortran parity preserved), needs no float64, so it also unblocks MPS/MLX.AMICAMLXNG. An optional Apple-GPU backend: GPU float32 E/M-step hot path + CPU-stream linalg (hoisted per iteration) + host-side SciPy for MLX's missinglgamma/digamma. Optional dependency (CI skips it; omitted from the coverage gate).AMICAMLXNGton_models > 1; single-model stays byte-identical; the win extends to multi-model (~5x over CPU). Matches AMICATorchNG float32 (one-iter stats to float32 precision, converged LL ~1e-5).Net result
On Apple Silicon, use the MLX backend for single- and multi-model AMICA (~5-7x over CPU); avoid
device="mps"(no gain). float64-CUDA stays the bit-safe NVIDIA path. Every phase preserved #24 Fortran parity (float64 byte-identical) and the full torch suite.Merge strategy
Regular merge (NOT squash) to preserve the per-phase history above (per repo policy for epic branches). Each phase PR was already squash-merged into this branch with its own review.
Follow-ups (tracked separately)
Closes #74