fix: unsafe non-blocking H2D transfers and MPS fallback timing

[musicalplatypus]

Summary

Two related MPS-specific bugs found while validating training end-to-end on Apple Silicon.

non_blocking H2D transfer race on non-CUDA devices

non_blocking=True on a .to(device, ...) H2D transfer is only safety-guaranteed on CUDA with pinned source memory. On MPS, create_data_loaders() only pins for CUDA (pin_memory=False on MPS/CPU), so an async H2D copy can race with reuse of the source buffer and corrupt the transferred tensor. Observed as AssertionError: min nan should be less than max nan inside a FakeQuantize observer — the eval-mode input tensor itself was already NaN before any quantization math ran, traced back to the async copy racing against its own CPU source buffer.

Confirmed the same precondition applies more broadly than just the training loop's evaluate_* functions: every test_onnx.py script (one per task type — timeseries classification/regression/forecasting/anomaly-detection, image classification, audio classification) and timeseries_anomalydetection/train.py's get_reconstruction_errors_stats() build their own separate DataLoaders (hardcoding pin_memory=True, or gating it on GPU count rather than CUDA specifically) and use non_blocking=True unconditionally. A runtime warning confirms pin_memory is a silent no-op on MPS: 'pin_memory' argument is set as true but not supported on MPS now, then device pinned memory won't be used. These scripts are reachable from the default pipeline — mmcli train invokes them automatically as the post-training testing step, not gated behind a separate compile step.

Fixed by gating non_blocking on device.type == 'cuda' everywhere it was previously unconditional: evaluate_classification/evaluate_forecasting/evaluate_regression/evaluate_anomalydetection (common/utils/utils.py) and all 8 files above.

PYTORCH_ENABLE_MPS_FALLBACK set too late to take effect

_get_device() (in the timeseries/vision/audio ai_modules training base classes) sets PYTORCH_ENABLE_MPS_FALLBACK=1 via os.environ.setdefault(...) when MPS is selected, so unsupported ops fall back to CPU instead of hard-failing. This env var is read once, by PyTorch, when its MPS backend registers during import torch — and torch is already imported (transitively, by the module _get_device() itself lives in) well before _get_device() ever runs. Confirmed via instrumentation that os.environ['PYTORCH_ENABLE_MPS_FALLBACK'] really was set to '1' at that point in the call stack, yet an unsupported quantization op still hard-failed with NotImplementedError instead of falling back — the flag being present in os.environ at that point has no effect on an already-initialized backend.

Fixed by setting PYTORCH_ENABLE_MPS_FALLBACK unconditionally at the top of run_tinyml_modelmaker.py — the single shared entry point for all three ai_modules, which imports only stdlib at module level — ahead of import tinyml_modelmaker (which pulls in torch transitively). Harmless on CUDA/CPU since it only changes MPS dispatch behavior. The original _get_device() calls are left in place as harmless (if individually ineffective) redundancy.

Verification

• non_blocking fix: reproduced the original crash via a full mmcli train run on Apple Silicon (quantized training + eval), confirmed it's resolved (clean run, both float and quantized ONNX exported, sane accuracy). For the 8 ONNX-runtime eval script call sites, ran the identical repro before and after the fix — both completed cleanly without crash or NaN either way, so this part is a preventive fix by pattern-match to the confirmed root cause, not something independently forced to fail-then-pass (the race is inherently probabilistic, and this code path lacks the CPU-fallback-op scheduling perturbation that made the original bug reproduce reliably).
• MPS fallback fix: reproduced the NotImplementedError with PYTORCH_ENABLE_MPS_FALLBACK unset and --training-device mps passed explicitly, confirmed clean (EXIT_CODE:0, correct fallback warning instead of a hard error, both ONNX artifacts produced) after the fix, in two separate runs.

🤖 Generated with Claude Code

[musicalplatypus]

CI failure investigated — pre-existing on upstream/main, unrelated to this PR.

Verified directly against upstream/main's own most recent CI run (same base commit, 331388a): https://github.com/TexasInstruments/tinyml-tensorlab/actions/runs/30376872256 — it already fails identically on macOS + Ubuntu with:

FAILED test_config_validation.py::...test_task_type_is_valid[google_speech_command/config_MSPM0.yaml] - unknown task_type='audio_classification'
FAILED test_config_validation.py::...test_model_name_exists_in_registry[...] - model 'DSCNN_NPU' not in registry

This PR's CI run shows only that same subset — no new failures introduced. It's fixed in #19; should resolve automatically once that merges. (Windows failures don't block merging — the workflow marks that runner continue-on-error: true.)

[musicalplatypus]

Pushed an additional commit (072ce4f) fixing two more robustness bugs found in these same files during a broader codebase review:

1. timeseries_anomalydetection/test_onnx.py: get_reconstruction_errors_stats() referenced args.gpu (singular), which the shared argparser never defines (only --gpus), crashing every invocation with AttributeError. Also fixed a related backwards rank-gated pin_memory in the same file's main(). Both now match the unconditional pin_memory=True used by every sibling test_onnx.py.
2. audio_classification/test_onnx.py: never called shutdown_data_loaders() on its DataLoader, unlike all 5 sibling scripts, leaking worker processes/semaphores whenever --workers > 0. Wrapped in try/finally matching the sibling pattern.

Both come with regression tests verified to fail pre-fix and pass post-fix. See the commit message for full detail.

[musicalplatypus]

An independent peer review (auditing whether this PR's changes and PR #22's changes to the same shared files -- `utils.py`, `timeseries_anomalydetection/train.py` -- were semantically safe to merge together, not just textually non-conflicting) found a real, independent bug already in this PR. Pushed a follow-up commit (f8ccbac):

This PR's own earlier fix gated `non_blocking` transfers on `device.type == 'cuda'` in `get_reconstruction_errors_stats()`. Before that, `non_blocking` was hardcoded `True` and the function never touched `.type`, so a plain device string worked fine. After the fix, `.type` is required -- but the call site in `main()` passed `args.device` (the raw, unconverted argparse string, `'cuda'` by default) instead of `device` (the `torch.device` already constructed by `setup_training_environment()` and in scope in the same function). `'cuda'.type` raises `AttributeError`, crashing every anomaly-detection training run -- CUDA included -- right after export, while calculating the detection threshold. Not an edge case; the normal path.

Fixed by passing the existing local `device` instead of `args.device`.

Added a test that drives the real `main()` (heavily mocked elsewhere) and asserts what it actually passes at the call site -- verified to fail pre-fix (`main()` passed the raw string `'cuda'`) and pass post-fix. Two additional tests characterize `get_reconstruction_errors_stats()`'s contract directly. Re-verified the full #20→#21→#22→#23 merge sequence still produces zero conflicts with this commit included.