JAX PathwaysAll career routesCareer learning guide
Accelerator & performance engineer
Measure where a workload spends time and memory, then improve the limiting operation while preserving its numerical contract. Explain when an optimization helps and when it does not.
The work
A faster-looking run excluded compilation or returned before device work completed. Rebuild the benchmark and compare identical workloads.
Preparation
Start with a correct training loop and understand its data and state. Profiling requires a reproducible workload; distributed execution adds placement and communication contracts.
Your portfolio
A correctness-backed benchmark with synchronized latency, throughput, memory and scaling limits.
Skills and evidence
XProf, HLO, and roofline reasoning
Required route practice
Read an actual CPU trace and separate compilation from steady work; target measurements are separate.
Explain the independent check, repeat it under a changed condition and retain the observed result.
Practice project: stages 4.
Meshes, sharding, and distributed recovery
Required route practice
Check count-weighted global updates and fresh-process recovery.
Explain the independent check, repeat it under a changed condition and retain the observed result.
Practice project: stages 1, 2, 3.
Pallas and reference correctness checks
Required route practice
Compare independent numerical references and declare interpreted versus target execution.
Explain the independent check, repeat it under a changed condition and retain the observed result.
Practice project: stages 1, 2, 3.
MLflow experiment provenance
Required route practice
Keep comparable configurations and observed measurements, including failed runs.
Explain the independent check, repeat it under a changed condition and retain the observed result.
Practice project: stages 1, 2, 3, 4.
Milestones
Build the shared CPU foundation
A reproducible regression fit with checked gradients and explicit random state.
Build a reference workload
A reproducible classifier
Make it recoverable
An interrupted run restored with its full state
Find its measured bottleneck
A trace-backed before-and-after performance report
Partition it across devices
A resilient sharded training experiment
Optimize a checked kernel
A measured custom kernel with correctness checks
Review your readiness
Preserve correctness, separate compile time from execution, and justify the improvement using workload measurements.
- Did every timing wait for completion and separate compilation from steady execution?
- Which trace event or memory-traffic estimate supports the claimed bottleneck?
- Does the optimization preserve results across changed shapes, and where does it stop helping?
- How would you distinguish an image/runtime change from a model optimization in your benchmark?
This is a learning route for practicing the work, not a claim of job readiness. Use the milestone coverage and a reviewer’s feedback to identify what remains to demonstrate. The capability map below separates required route practice from optional extensions. Public checks and self-reported completion are not independent review or hardware qualification.