feat: auto-enable torch.compile/AMP on CUDA via apply_hardware_defaults by musicalplatypus · Pull Request #23 · TexasInstruments/tinyml-tensorlab

musicalplatypus · 2026-07-30T19:44:34Z

Summary

Adds apply_hardware_defaults, which auto-enables torch.compile and native AMP when training on CUDA, unless the caller explicitly set compile_model/native_amp themselves. Wires it into timeseries init_params().

Depends on pr/compile-hardening (#22): this branch is built on top of it (all 16 of its commits, plus the 4 described below) — the auto-enable feature only makes sense once torch.compile fails safely rather than crashing training, which is what #22 hardens. The diff here will show as just the 4 new commits once #22 merges first.

`apply_hardware_defaults`

CUDA-gated (not device-gated, so it doesn't fire on MPS where torch.compile support is weaker) — sets compile_model=1 and native_amp=True as defaults, respecting any value the caller already set explicitly via the training dict.

Wiring into `init_params`

Called at the end of init_params() so compile_model/native_amp auto-enable on CUDA unless the caller explicitly set them via the training dict passed as args[0].

Known gap, not hidden: one of the three integration tests (test_init_params_respects_native_amp_false_override) currently fails on this branch alone — it surfaces a real, pre-existing, unrelated bug in ConfigDict's constructor, where a partial nested-dict override (e.g. training=dict(native_amp=False)) shallow-replaces the whole training section instead of deep-merging, dropping compile_model and all other training defaults. Confirmed via git stash comparison that this bug predates this change and isn't introduced by it — it's fixed separately in pr/configdict-deep-merge (#18). That test will pass once #18 merges too.

Robustness fix found along the way

init_params's user_training_keys computation assumed args[0] was always dict-like, but ConfigDict (which init_params feeds into) documents a YAML path string or None as valid first positional inputs too. ModelRunner.init_params is a public classmethod — nothing stops a caller from using that documented capability, which would crash immediately on args[0].get(...). Not reachable via any current real call site (all forward pre-built dict-like objects), but a real robustness gap on a public entry point. Guarded with an isinstance check.

Verification

Integration tests added proving the auto-enable wiring end-to-end (CUDA-gated default, explicit-override respected, non-dict/None first-arg no longer crashes).

🤖 Generated with Claude Code

tinyml-tinyverse's own test suite doesn't exist yet at this point in history (added by the next commit) -- this adds only the path trigger. The step that actually runs those tests is added once they exist, later in this branch, so no commit in this sequence ever references a directory that isn't there yet. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>

…all back to eager mode torch.compile() is lazy — the previous try/except only wrapped the wrap-time call, not the deferred compilation that happens on first forward. A failing compile (e.g. unsupported Triton/ptxas for the GPU's compute capability) would crash mid-training instead of falling back to eager mode. compile_model_if_enabled now accepts an optional input_shape and, when provided, runs one warmup forward pass through the compiled model before training starts. A failure at either the wrap step or the warmup step falls back to the original, uncompiled model. Backward compatible: omitting input_shape preserves the old (unguarded) behavior. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>

torch.compile's OptimizedModule wraps the original model by reference, so calling .eval() on the compiled wrapper during warmup also flips the ORIGINAL model's .training flag. If the warmup forward pass raised, the restore step was skipped (exception jumped straight to the outer except), so the fallback model was returned stuck in eval mode. Wrap the eval/forward/restore sequence in try/finally so restoration always runs. Also strengthens test_warmup_failure_falls_back_to_original_model: the old mock returned an unrelated standalone module that didn't share state/reference with the original model, so it couldn't reproduce this bug. The new version wraps the original model as a real child submodule (mirroring OptimizedModule._orig_mod) and asserts the original model's .training flag is restored after a failed warmup — confirmed this fails against the old buggy code and passes against the fix. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

…e training-mode graph

torch.compile() wraps a model in torch._dynamo.OptimizedModule, exposing the original module at ._orig_mod. Neither torch.jit.trace (quantized export) nor torch.onnx.export (float export) can trace a dynamo-optimized module directly, so export_model() crashed with "Detected that you are using FX to torch.jit.trace a dynamo-optimized function" whenever compile_model=1 actually succeeded (the hardware-defaults feature auto-enables this on CUDA). Unwrap via getattr(model, '_orig_mod', model) once before the existing deepcopy — a no-op for uncompiled models.

…op-level Single-level getattr(model, '_orig_mod', model) missed the actual failure shape: timeseries_classification always wraps the (already-compiled) model inside NeuralNetworkWithPreprocess.model after compile_model_if_enabled runs, so the compiled OptimizedModule ends up one level below the top-level model, not at it. torch.onnx.export/torch.jit.trace still crashed on the nested compiled submodule. unwrap_compiled_submodules() walks the full submodule tree recursively, replacing every torch.compile-wrapped module (top-level or nested) with its original. No-op wherever nothing is compiled. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>

…'t handle a compiled model prepare_qat_fx (FX symbolic tracing) crashes with "Detected that you are using FX to symbolically trace a dynamo-optimized function" when the model passed to it is torch.compile-wrapped (torch._dynamo.OptimizedModule). compile_model_if_enabled now skips compiling entirely whenever args.quantization is set, since compiling first would be discarded before quantization prep runs anyway. Float training (quantization=0) is unaffected. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>

…log assertion

torch.compile() wraps a model in torch._dynamo.OptimizedModule. Since setup_distributed_model only assigns model_without_ddp = model.module under DDP, in the non-DDP case model_without_ddp IS the compiled wrapper, so state_dict() emits every key prefixed _orig_mod. The downstream float->quantization weight transfer (load_weights.py) can't match those keys, silently falls back to strict=False, and discards the entire float-trained result -- no exception, no fatal error, the pipeline just quietly retrains from random init and reports success. save_checkpoint and resume_from_checkpoint now unwrap via a local getattr(model, '_orig_mod', model) at each call site, so checkpoints always carry uncompiled key names. This intentionally does not use the existing unwrap_compiled_submodules() helper, which mutates the model's submodule tree in place via setattr and would silently un-compile the live training model. Also adds weights_only=False to resume_from_checkpoint's torch.load call. This is a separate, pre-existing latent bug this fix's own test suite surfaced: checkpoint['args'] is an argparse.Namespace (every train.py passes the full args object to save_checkpoint), which is not on torch's default weights_only safe-globals list, so on torch >=2.6 the --resume path already fails before reaching load_state_dict, independent of the compile-wrapper issue. Matches the weights_only=False convention already used at every other non-tensor-only torch.load() call site in this codebase (load_weights.py, and the load_saved_model paths in each task's train.py). Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>

…ave+load ExponentialMovingAverage (AveragedModel) deep-copies its source model into self.module, so when the source was already compiled, the OptimizedModule wrapper ends up nested at model_ema.module._orig_mod -- not at model_ema._orig_mod itself, where the prior fix's top-level getattr unwrap couldn't reach it. Strip the substring from the resulting state_dict keys instead (handles any nesting depth), and remap symmetrically on load by matching against model_ema's own current key names stripped the same way. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>

…l production shape

tinyml-tinyverse/tests/ (added in this branch's earlier commits) was never executed by CI. No extra install step needed -- tinyml-tinyverse is already pip installed in the shared Install dependencies step. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>

save_checkpoint/resume_from_checkpoint (train_base.py) already stop writing _orig_mod.-prefixed keys going forward, but load_weights() -- the actual consumer for the float->quantization --weights transfer (timeseries_base.py) -- had no handling for the prefix at all, only for 'module.' (DDP). A checkpoint saved by older, unpatched code, or by any other torch.compile-using caller not covered by the train_base.py fix, would still hit the original silent strict=False fallback that discards 100% of weights with no exception. Strip _orig_mod. from incoming checkpoint data before the existing 'module.' realignment runs (it's never meaningful to preserve or match against, unlike 'module.'), and symmetrically remap onto the live model's own current key names if that model is itself currently compiled -- mirroring the same symmetric approach already used for the EMA checkpoint path. Also adds a regression test that round-trips a checkpoint through the actual load_weights() consumer (not just direct load_state_dict), closing the gap noted in review: the fix was previously verified against this consumer only via ad-hoc manual testing, not a permanent test. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>

…ompat Two independent bugs in the same function, found by adversarial security and code review: 1. Security (unsafe deserialization): the weights_only=False added to torch.load(args.resume, ...) disabled torch >=2.6's default unpickling restriction wholesale, rather than allowlisting only the one non-tensor type the checkpoint actually needs (argparse.Namespace, for checkpoint['args']). A crafted --resume checkpoint from an untrusted source could execute arbitrary code via a pickle __reduce__ payload. Fixed via torch.serialization.safe_globals([Namespace]). 2. Backward compatibility: the prior fix's unwrap only touched the live model, not the checkpoint data, so it silently assumed checkpoint['model'] always has clean (unprefixed) keys. A checkpoint written before this session's fixes existed -- with _orig_mod. keys, from a compiled model -- would now fail to load into an unwrapped model with a strict key mismatch, where it used to work by accident (both sides had matching prefixes). Fixed by extracting the same "strip from data, then remap onto the live model's actual current keys" pattern already used for the EMA branch into a shared _load_symmetric() helper, applied to both the main model and EMA uniformly. Also strengthens existing checkpoint tests to assert on .bias and full parameter sets, not just .weight -- a bug that only corrupted bias handling would have passed undetected. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>

Tests that do a genuine torch.save + resume_from_checkpoint (torch.load) round trip were passing _FakeArgs (a lightweight test stand-in) as the args parameter to save_checkpoint, which becomes checkpoint['args'] and must therefore actually be unpickled by torch.load's weights_only safety check. Only argparse.Namespace is allowlisted (matching real production usage, where args always comes from ArgumentParser.parse_args()), so these tests started failing once that check was genuinely enforced. Switched to a real Namespace() for exactly the calls that get serialized to disk; _FakeArgs remains fine for args objects that are only ever read directly as a live function parameter (.resume/.device), never pickled. Also moves the pickle-rejection test's "not on the allowlist" stand-in class to module level -- pickle cannot serialize local/nested classes at all, which made that test fail for an unrelated reason before it ever reached the weights_only check it's meant to exercise. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

Calls apply_hardware_defaults(params, user_training_keys) at the end of init_params() so compile_model/native_amp auto-enable on CUDA unless the caller explicitly set them via the training dict passed as args[0]. Adds integration tests proving the wiring end-to-end. One of the three (test_init_params_respects_native_amp_false_override) fails: it surfaces a pre-existing, unrelated bug in ConfigDict's constructor where a partial nested dict passed as an override arg (e.g. training=dict(native_amp=False)) shallow-replaces the whole 'training' section instead of deep-merging, dropping compile_model and all other training defaults. Confirmed via git-stash comparison that this bug predates this change and is not introduced by it.

…attern test_init_params_respects_native_amp_false_override was hitting a pre-existing shallow-merge bug in ConfigDict.__init__ by passing a partial training dict directly into init_params(). Rewrote it to call init_params() with zero args (like ModelRunner.init_params() in run_tinyml_modelmaker.py) then merge the user override via params.update(), which deep-merges correctly and avoids the constructor bug entirely. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>

user_training_keys assumed args[0] was always dict-like, but ConfigDict (which init_params feeds into) documents a YAML path string or None as valid first positional inputs too. ModelRunner.init_params is a public classmethod -- nothing stops a caller from using that documented capability, which would crash immediately on args[0].get(...). Not reachable via any current real call site (all forward pre-built dict-like objects), but a real robustness gap on a public entry point. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>

musicalplatypus · 2026-07-30T20:55:29Z

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.)

…ting Two independent bugs in the "Log best epoch results" section, which runs after `for epoch in range(args.start_epoch, args.epochs)`: 1. timeseries_classification/train.py logged `best['f1']` under the "AUC ROC Score" label instead of the `best['auc']` value that is actually computed and stored in `best` on every improving epoch (`best['accuracy'], best['f1'], best['auc'], ... = avg_accuracy, avg_f1, auc, ...`). 2. When the loop runs zero iterations -- e.g. `--resume` pointed at a checkpoint that already satisfies `--epochs`, a real and expected use case for re-running only the post-training export/compile steps -- the post-loop section read dict keys that are only ever assigned inside the loop body: - timeseries_classification/train.py crashed with `KeyError: 'predictions'` (best = dict(accuracy=0.0, f1=0, conf_matrix=dict(), epoch=None) has no 'predictions'/'ground_truth'/ 'auc' keys until an improving epoch runs). Now guards the whole block on `best['epoch'] is not None`, matching the `epoch=None` sentinel already present in the initial dict. - timeseries_forecasting/train.py crashed with `TypeError: 'NoneType' object is not subscriptable` on best_epoch_values['true_values'][:, :, idx] (true_values/ predictions are pre-populated with None, not omitted). Now guards the whole block on `best_epoch_values['true_values'] is not None`, extending the guard the file already used for its "Save final predictions" section to the "Per-Variable Metrics" section that actually crashes. (timeseries_regression and timeseries_anomalydetection were already safe here -- their `best` dicts only ever hold scalars with safe initial values, so no fix needed there.) Adds tests/test_train_best_epoch_bugs_timeseries.py, which drives each script's real main() through a heavily mocked model/data pipeline with args.start_epoch == args.epochs (loop runs zero iterations) and asserts main() completes without raising, plus a targeted test that runs the loop for one improving epoch with f1 and auc set to distinct values and asserts the "AUC ROC Score" log line reports auc, not f1. All three tests verified to fail pre-fix with the exact errors above and pass post-fix. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>

musicalplatypus · 2026-08-04T23:05:19Z

Pushed an additional commit (9c2f625) fixing two more bugs found in the "Log best epoch results" section of these same files during a broader codebase review:

timeseries_classification/train.py: logged best['f1'] under the "AUC ROC Score" label instead of the best['auc'] value that's actually computed and stored on every improving epoch.
Both timeseries_classification/train.py and timeseries_forecasting/train.py crashed (KeyError/TypeError) when --resume pointed at a checkpoint that already satisfies --epochs (loop runs zero iterations, so best/best_epoch_values never gets fully populated) -- a real use case for re-running only the post-training export/compile steps. Both now guard the block on whether any epoch actually ran.

Comes with a regression test suite verified to fail pre-fix with the exact errors and pass post-fix. See the commit message for full detail.

# Conflicts: # .github/workflows/test-modelmaker.yml

t5fkg8d44d-beep and others added 20 commits July 29, 2026 11:17

fix: pass input_shape to compile_model_if_enabled at all 4 call sites

3fd2bed

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

docs: clarify compile warmup only validates one graph variant, not th…

57ff586

…e training-mode graph

test: strengthen skip-compile tests with torch.compile mock and INFO-…

106e2f9

…log assertion

test: add asymmetric compiled-save/uncompiled-load EMA case — the rea…

2268e22

…l production shape

feat: add apply_hardware_defaults — auto-enable compile/AMP on CUDA

a2c05eb

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

musicalplatypus mentioned this pull request Jul 30, 2026

test: fix stale/wrong test assertions and CI trigger gap #19

Merged

This was referenced Aug 4, 2026

fix: train.py post-loop bugs in image and audio classification #32

Merged

fix: allow weights_only=False when loading --cache-dataset torch cache #31

Merged

Merge remote-tracking branch 'upstream/main' into pr/hardware-defaults

b7814cf

# Conflicts: # .github/workflows/test-modelmaker.yml

musicalplatypus force-pushed the pr/hardware-defaults branch from 9623a4e to b7814cf Compare August 6, 2026 04:41

musicalplatypus mentioned this pull request Aug 6, 2026

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

Merged

Adithya-Thonse merged commit e83aefe into TexasInstruments:main Aug 6, 2026
0 of 3 checks passed

Navigation Menu

Uh oh!

feat: auto-enable torch.compile/AMP on CUDA via apply_hardware_defaults - #23

feat: auto-enable torch.compile/AMP on CUDA via apply_hardware_defaults#23
Adithya-Thonse merged 22 commits into
TexasInstruments:mainfrom
musicalplatypus:pr/hardware-defaults

musicalplatypus commented Jul 30, 2026

Uh oh!

musicalplatypus commented Jul 30, 2026

Uh oh!

musicalplatypus commented Aug 4, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Uh oh!

Conversation

musicalplatypus commented Jul 30, 2026

Summary

apply_hardware_defaults

Wiring into init_params

Robustness fix found along the way

Verification

Uh oh!

musicalplatypus commented Jul 30, 2026

Uh oh!

musicalplatypus commented Aug 4, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

`apply_hardware_defaults`

Wiring into `init_params`