Use TraceML With Your Stack

TraceML can be added to a normal PyTorch training loop or used through a framework-specific entry point. Pick the path that matches how your training job already runs.

Integration support matrix

The manifest at docs/data/integration_support.json is the source of truth for this table. A status is evidence, not a promise: CI tested means the linked job installs the extra and runs the linked real-framework test. documented recipe and experimental are not end-to-end validation.

Stack	Recommended API	Tested dependency range	CPU / GPU	Single- / multi-process	Multi-node	Emitted signals	Example	Validation level	Limitations
Plain PyTorch loop	traceml.init() + traceml.trace_step(model)	CI resolves torch>=2.5.0 (not pinned; this is not a broad version matrix).	CPU: CI tested — Ubuntu integration job GPU: Not claimed — no GPU job	Single: CI tested — launcher smoke run Multi: Not claimed — no matrix evidence	Not claimed — no matrix evidence	The cited smoke test asserts a final summary; it does not declare a per-signal contract.	quickstart.py	CI tested CI job: test-integration; real PyTorch launcher smoke test	Distributed Training · #141 · #320 · #323
Hugging Face Trainer	traceml_ai.integrations.huggingface.init() + TraceMLTrainerCallback	CI resolves torch>=2.5.0, transformers (unbounded), and accelerate>=0.26.0; not a pinned compatibility matrix.	CPU: CI tested — real Trainer run GPU: Not claimed — no GPU job	Single: CI tested — CPU only Multi: Documented recipe — no real distributed job	Experimental — no reproducible job	dataloader fetch with the recommended callback; step, forward, backward, and optimizer through a real Trainer conformance run.	huggingface_trainer_minimal.py	CI tested CI job: test-integration; real Transformers Trainer test	Hugging Face limitations · #141 · #320 · #323
Hugging Face Accelerate	traceml.init() + traceml.trace_step(accelerator.unwrap_model(model))	CI resolves torch>=2.5.0 and accelerate>=0.26.0; not a pinned compatibility matrix.	CPU: CI tested — real Accelerator.prepare/backward run GPU: Not claimed — no GPU job	Single: CI tested — CPU only Multi: Documented recipe — no real distributed job	Experimental — no reproducible job	step, forward, backward, optimizer, and dataloader fetch on the cited real Accelerate run.	accelerate_minimal.py	CI tested CI job: test-integration; real Accelerate prepare/backward test	Accelerate limitations · #141 · #320 · #323
PyTorch Lightning	traceml_ai.integrations.lightning.init() + TraceMLCallback	CI resolves lightning>=2.6.0 (not pinned; this is not a broad version matrix). The cited tests run under both the lightning and the legacy pytorch_lightning namespace.	CPU: CI tested — Ubuntu integration job GPU: Documented recipe — no GPU job	Single: CI tested — real Trainer.fit() CPU run Multi: Documented recipe — no real distributed job	Experimental — no reproducible job	tests/integrations/test_telemetry_conformance.py asserts step, forward, backward, optimizer, and DataLoader fetch on a CPU run, and that a missing init() leaves the fetch stream dark; H2D is GPU-only and not asserted in CI.	lightning_minimal.py	CI tested CI job: test-integration; real Lightning Trainer.fit() tests	Lightning guide · #141 · #320 · #323
RF-DETR	traceml_ai.integrations.rfdetr.init() + model.train()	Dedicated CPU CI pins rfdetr[train]==1.10.1. CPU correctness, offline Nano and Gloo checks also passed locally against RF-DETR development commit 0ed5be8 with torch 2.11 / Lightning 2.6.1; this is not a broad compatibility range.	CPU: CI tested — dedicated RF-DETR CPU job GPU: Documented recipe — no CUDA/NCCL validation	Single: CI tested — parity and real Nano smoke Multi: CI tested — two CPU/Gloo ranks with lightweight RF-DETR fixture	Experimental — no physical multi-node validation	CPU tests assert training step, inner-model forward, backward, optimizer and DataLoader fetch, excluding evaluation. H2D and CUDA memory require GPU validation.	rfdetr_minimal.py	CI tested CI job: test-rfdetr; RF-DETR parity, Nano and Gloo smoke tests	RF-DETR limitations · #141 · #323
Ray Train	TraceMLTorchTrainer + TraceMLRayConfig	Documented install requirement: ray[train]>=2.0.0. No real-Ray CI-tested range.	CPU: Documented recipe — no real-Ray job GPU: Documented recipe — no GPU job	Single: Documented recipe — local Ray recipe Multi: Documented recipe — worker recipe only	Experimental — no reproducible job	Guide documents step, forward, backward, optimizer, and wrapped Ray Data fetch timing; no real-Ray CI signal assertion exists.	torchtrainer_minimal.py	documented recipe	Ray guide · #141 · #320 · #323
DeepSpeed	traceml.init() + traceml.trace_step(model_engine.module)	Documented install requirement: deepspeed>=0.10.0. No real-DeepSpeed CI-tested range.	CPU: Unsupported — the guide requires CUDA GPU: Documented recipe — no CUDA CI job	Single: Documented recipe — no real engine run Multi: Documented recipe — no real distributed job	Experimental — no reproducible job	Guide documents DataLoader fetch, H2D, forward, backward, optimizer, step, and memory; no real-DeepSpeed CI signal assertion exists.	deepspeed_minimal.py	documented recipe	DeepSpeed limitations · #141 · #320 · #323
W&B / MLflow	traceml.summary()	No W&B or MLflow dependency is installed or run in CI.	CPU: Documented recipe — tracker client not run GPU: Not applicable — tracker export is device-independent	Single: Documented recipe — no tracker-client CI run Multi: Not claimed — no matrix evidence	Not claimed — no matrix evidence	No tracker telemetry is emitted by TraceML; the documented integration exports TraceML's compact summary dictionary.	summary_logging_minimal.py	documented recipe	W&B / MLflow guide · #323

Evidence limits

• This matrix makes no GPU, multi-process, or multi-node claim without a corresponding reproducible job.
• Lightning, Ray, and DeepSpeed tests in the current integration suite include mocked, lazy-import, or skipped paths; they are therefore not classified as end-to-end validation here.
• Run python tools/integration_support_matrix.py --check after editing the manifest, examples, or CI install extras. The check fails if a CI tested row loses its cited extra, test, guide, example, or generated documentation row.