JAX PathwaysAll career routesCareer learning guide
ML & training engineer
Turn a dataset and model idea into a reproducible experiment. You define the objective, diagnose learning, keep held-out evaluation isolated and prove a saved run can continue correctly.
The work
Compare two training runs whose validation losses disagree. Trace the difference through data IDs, initialization, optimizer state and precision.
Preparation
Use foundations to become comfortable with shapes, scalar losses, gradients and explicit state. Then practice train/evaluation separation in the classifier before adding recovery and sequence models.
Your portfolio
A run dossier: data split, configuration, learning curves, error examples and a restart comparison.
Skills and evidence
Flax NNX and explicit state
Required route practice
Inspect graph/state and check a known-parameter forward pass.
Explain the independent check, repeat it under a changed condition and retain the observed result.
Practice project: stages 1.
Optax and train/eval steps
Required route practice
Verify an update independently and show evaluation does not mutate training state.
Explain the independent check, repeat it under a changed condition and retain the observed result.
Practice project: stages 2, 3.
Grain, Orbax, and XProf
Optional project extension; 4 project preparation lessons lie outside the required route
Compare resumed and uninterrupted state and inspect a real CPU trace.
Explain the independent check, repeat it under a changed condition and retain the observed result.
Practice project: stages 3, 4.
MLflow experiment provenance
Required route practice
Reload tracked metrics and identities and reproduce a paired comparison.
Explain the independent check, repeat it under a changed condition and retain the observed result.
Practice project: stages 1, 2, 3, 4.
Masked and contrastive pretraining; response-only SFT, LoRA and preference optimization; numerical architecture parity
Optional project extension; 8 project preparation lessons lie outside the required route
Defend masks, objective reductions and held-out checks; these optional objective labs do not certify a complete adapted model.
Explain the independent check, repeat it under a changed condition and retain the observed result.
Practice project: stages 1, 2, 3, 4, 5, 6, 7, 8.
Milestones
Build the shared CPU foundation
A reproducible regression fit with checked gradients and explicit random state.
Train and evaluate a small model
A reproducible classifier
Prove full-state recovery
An interrupted run restored with its full state
Diagnose and improve training speed
A trace-backed before-and-after performance report
Review your readiness
Demonstrate that your model trains, evaluates, resumes correctly, and produces reproducible results.
- How did you keep validation data and evaluation state out of training?
- Which saved values are required to reproduce the next update after interruption?
- Which error examples explain the learning curve, and does the result hold under a new seed?
- How would you recover the exact code, data and metric definition behind the selected MLflow run?
- How would you detect target leakage, an incorrect preference denominator or the first PyTorch/Flax layer mismatch?
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.