2. Add Apple Silicon (MPS) training support#8
Open
musicalplatypus wants to merge 4 commits into
Open
Conversation
musicalplatypus
changed the title
Adithya-Thonse
pushed a commit
that referenced
this pull request
Jun 12, 2026
Merge in TINYML-ALGO/tinyml-agent-skills from 2026/pranav_a to main * commit '8c3260bcdd549353b2cdbf3562bb3f32753fdddf': improving readme
Adithya-Thonse
added a commit
that referenced
this pull request
Jun 12, 2026
de8af16d Pull request #45: https://jira.itg.ti.com/browse/TINYML_ALGO-698 REVERT: e48ef1a Pull request #14: TINYML_ALGO-711: fixing readme REVERT: 16fc6a6 TINYML_ALGO-711: fixing readme REVERT: e3639d2 Pull request #13: removing pycache REVERT: f8bb3b7 removing pycache REVERT: dd38428 Pull request #12: restructuring agent skill REVERT: ff02a0e restructuring agent skill REVERT: d26c6a5 Pull request #11: fixing tiny ml name REVERT: 640ffd3 fixing tiny ml name REVERT: 4ee3a19 Pull request #10: 2026/pranav a REVERT: be83fc6 minor fixes REVERT: e3a5700 removed assets, included autoMP quant REVERT: 1af575a Pull request #9: correcting npu devices list REVERT: 31e9eb1 correcting npu devices list REVERT: 59b209b Pull request #8: improving readme REVERT: 8c3260b improving readme REVERT: 668916f Pull request #7: improving readme REVERT: 68686b3 improving readme REVERT: 814316e Pull request #6: fixes to readme and marketplace json REVERT: e4bc0b4 fixes to readme and marketplace json REVERT: 6a64208 Pull request #5: fixes to readme REVERT: 0f9c868 fixes to readme REVERT: 52f95ff Pull request #4: 2026/pranav a REVERT: 443295d fixes to readme REVERT: 1881112 fixes to readme and marketplace json REVERT: 229ab57 Pull request #3: 2026/pranav a REVERT: 6519104 minor readme fix REVERT: 38e9f9f minor readme fix REVERT: db81f81 Pull request #2: minor readme fix REVERT: 1c0737a minor readme fix REVERT: 0a0c02d Pull request #1: minor readme fix REVERT: b682335 minor readme fix REVERT: 062eb39 Initial Commit git-subtree-dir: tinyml-agent-skills git-subtree-split: de8af16d9e23de3e9bda3d811a0ebdece1178260
…ntization
Training on the Apple-Silicon MPS backend currently crashes in three places,
all reducible to two root causes.
1. float64 tensors moved to MPS before being cast. MPS does not implement
float64 ("Cannot convert a MPS Tensor to float64 dtype"), so `x.to(device)`
must not run while the tensor is still float64. The train/eval loops in
common/utils/utils.py and the calibration/eval loops in the auto-quantization
pass moved first and cast afterwards. Reorder to `.float().to(device)` so the
float32 cast happens on CPU before the transfer. This is a no-op on CUDA/CPU.
2. torcheval classification metrics on MPS. multiclass_confusion_matrix,
multiclass_f1_score and multiclass_auroc build sparse-COO tensors internally,
which MPS does not support. Move the output/target to CPU for these metric
calls; the results are scalars/small matrices, so the transfer is negligible.
With these changes the full timeseries pipeline (train -> evaluate -> QAT
quantization -> float and quantized ONNX export) completes on MPS. Note that
PyTorch itself does not implement the QAT fake-quantize op on MPS; run with
PYTORCH_ENABLE_MPS_FALLBACK=1 so that single op falls back to CPU.
musicalplatypus
force-pushed
the
pr/mps-support
branch
from
July 16, 2026 02:08
3250aef to
20aba61
Compare
The NAS architecture-search training loop (nas/train_cnn_search.py) called .cuda() unconditionally for the model, criterion, inputs and targets, and used torch.cuda.* for device selection and memory stats. This made NAS fail on any machine without CUDA (including Apple Silicon). - Add nas/utils.get_device(), resolving CUDA > MPS > CPU, and route all device placement through it. Cast to float32/long before .to(device) so MPS (which lacks float64) does not crash. Guard torch.cuda.* calls behind a cuda check. - run_tinyml_modelmaker: when nas_enabled, synthesize a minimal model description so a NAS model id (e.g. NAS_m) is not rejected as an unknown model, and tolerate a None model description downstream. - Add tests/test_nas_support.py covering model-validation bypass, fallback description, str2bool and argparse integration.
The ONNX evaluation scripts (references/*/test_onnx.py) and the anomaly- detection training script accumulated results with `torch.tensor([])`, whose default dtype is float64, and moved batches to the device before casting (`.to(device).float()` / `.to(device).long()`). On Apple-Silicon MPS both patterns raise "Cannot convert a MPS Tensor to float64 dtype". - Pin the accumulator tensors to `dtype=torch.float32`. - Cast (`.float()`/`.long()`) before `.to(device)` so the transfer never carries float64. Covers timeseries (classification, regression, forecasting, anomaly- detection) plus audio and image classification. No-op on CUDA/CPU.
musicalplatypus
force-pushed
the
pr/mps-support
branch
from
July 16, 2026 02:14
20aba61 to
9bf2289
Compare
Two undefined names on this branch, caught by static analysis:
nas/utils.py — adding get_device() displaced the trailing print() from
create_exp_dir(), leaving it stranded after get_device()'s return as dead code
referencing an undefined `path`. Restore it to create_exp_dir().
references/{audio_classification,image_classification}/test_onnx.py and
timeseries_anomalydetection/test_onnx_cls.py — each calls
.astype(np.float32) without importing numpy, raising NameError whenever the
nn_for_feature_extraction path runs. Add the import.
Verified: no undefined names remain in any file this branch changes.
musicalplatypus
force-pushed
the
pr/mps-support
branch
from
July 20, 2026 10:06
90d4827 to
7dfbfb4
Compare
musicalplatypus
pushed a commit
to musicalplatypus/tinyml-tensorlab
that referenced
this pull request
Jul 20, 2026
audio_classification/test_onnx.py calls .astype(np.float32) without importing numpy, raising NameError on the nn_for_feature_extraction path. The same fix exists on pr/mps-support, but that branch has never been merged into main, so it never propagated here. The sibling fixes for image_classification and timeseries_anomalydetection arrived via pr/macos-semaphore-fixes. Pre-existing upstream: TI's main carries the same bug. PR TexasInstruments#8 carries the fix back to them.
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
2 participants
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
Adds Apple Silicon (MPS / Metal) training support across all task types, so training runs GPU-accelerated on Macs without CUDA. Scope is limited to MPS correctness — performance tuning and macOS process fixes live in the sibling PRs in this series.
MPS device compatibility
float64. Reordered.to(device).float()/.to(device).long()to.float()/.long()before.to(device)in the training, evaluation and auto-quantization loops, and pinned empty accumulator tensors todtype=torch.float32(torch.tensor([])defaults to float64).multiclass_confusion_matrix,multiclass_f1_scoreandmulticlass_aurocbuild sparse-COO tensors that MPS does not support; compute them on CPU.references/*/test_onnx.pyand reference training scripts (timeseries classification / regression / forecasting / anomaly-detection, plus audio and image classification).NAS on MPS/CPU
.cuda()andtorch.cuda.*, so it only ran on CUDA. Addednas.utils.get_device()(CUDA > MPS > CPU) and routed all device placement and memory-stats through it, casting before transfer.run_tinyml_modelmakersynthesizes a minimal model description when NAS is enabled, so a NAS model id (e.g.NAS_m) is not rejected as an unknown model.tests/test_nas_support.py(9 tests).Note on quantization
PyTorch does not implement the QAT fake-quantize op on MPS. Run with
PYTORCH_ENABLE_MPS_FALLBACK=1so that single op falls back to CPU; with that set, the full train → eval → QAT → float & quantized ONNX export pipeline completes on MPS.Files changed (15)
Across
tinyml-tinyverse,tinyml-modeloptimizationandtinyml-modelmaker. No-op on CUDA/CPU paths.Testing
test_nas_support.py: 9 passed.