Software Engineer, ML Training Optimization
Foster City, CA
- Pay
USD 177,000–270,000/year — pay source
Salary 177,000 – 270,000 USD per year About Zoox
Read the full posting- Work setting
On-site — work setting source
Work setup in the employer’s posting: onsite
From the employer’s posting- Employment
Full-time — employment source
Employment type Full-time
Employment type in the employer’s posting: Full-time
From the employer’s posting
What you’ll work on
Full postingOptimize distributed training efficiency using parallelism, mixed precision, and gradient checkpointing
Improve RL framework performance for driving foundation models
Profile and debug training throughput bottlenecks across the stack
Implement custom kernels and build benchmarks to validate performance gains
Define requirements and architecture for large-scale training infrastructure with researchers and engineers
From the employer’s posting
Design, implement, and operate cutting-edge ML training performance optimization techniques
Work on RL framework and RL performance optimization for driving foundation model
Profile and debug training throughput bottlenecks across the compute, compiler, framework, and cluster-networking stack.
Implement custom kernels and build benchmarks to track and validate training-performance optimizations across the stack.
Collaborate closely with ML researchers, data engineers, and hardware engineers to define requirements and align on architectural decisions for large-scale training infrastructure.
What you’ll bring
All qualificationsCore experience
- 4+ years of total experience, including 2+ years working on large-scale model training platforms.
- Experience with distributed training frameworks (Megatron-LM, TorchTitan), efficiently utilizing GPU/accelerator clusters at scale.
- Experience with RL framework such as Verl/MILES and RL algorithms(GRPO, PPO)
- Experience with collective communication libraries (e.g., NCCL) and optimizing communication/compute overlap in multi-node training.
- Experience using profiling tools (NVIDIA Nsight, PyTorch Profiler, or XLA/JAX profiling) to identify and resolve training bottlenecks.
- Proficient in Python or C++.
Qualification wording
4+ years of total experience, including 2+ years working on large-scale model training platforms.
Experience with distributed training frameworks (Megatron-LM, TorchTitan), efficiently utilizing GPU/accelerator clusters at scale.
Experience with RL framework such as Verl/MILES and RL algorithms(GRPO, PPO)
Experience with collective communication libraries (e.g., NCCL) and optimizing communication/compute overlap in multi-node training.
Experience using profiling tools (NVIDIA Nsight, PyTorch Profiler, or XLA/JAX profiling) to identify and resolve training bottlenecks.
Proficient in Python or C++.
Tools in this posting
- C++
- Python
- PyTorch
Source — Tool mentions in context
- Experience using profiling tools (NVIDIA Nsight, PyTorch Profiler, or XLA/JAX profiling) to identify and resolve training bottlenecks. - Proficient in Python or C++. - Experience with ML compilers, performance tooling, or hardware/software co-design (e.g., GPU/TPU/accelerator architectures).
- Experience with collective communication libraries (e.g., NCCL) and optimizing communication/compute overlap in multi-node training. - Experience using profiling tools (NVIDIA Nsight, PyTorch Profiler, or XLA/JAX profiling) to identify and resolve training bottlenecks. - Proficient in Python or C++.
About Amazon
Zoox is developing the first ground-up, fully autonomous vehicle fleet and the supporting ecosystem required to bring this technology to market.
In the employer’s words · Read in context
Posting record
Report an errorActively hiringFirst seen by us 13 days ago, still listed, and not a repost of an earlier posting.
- Status in our records
- Live
- First seen by us
- Oct 16, 2026 · 13 days ago
- Recorded sightings
- 29
- Last seen by us
- Oct 30, 2026
These dates show when we found the posting. Check the employer’s website to confirm it is still accepting applications.
Your next step
- Have your resume and examples of relevant work ready.
- Check the listed location, eligibility and core experience before starting.
Complete your application on jobs.lever.co. The employer’s form will show what is required.
Already applied? Track this application
Job description
In this role, you will:
-
Design, implement, and operate cutting-edge ML training performance optimization techniques — distributed training (data/tensor/pipeline parallelism), mixed-precision, gradient checkpointing, and GPU kernels — to scale our VLM, VLA, and Foundational models training efficiency.
-
Work on RL framework and RL performance optimization for driving foundation model
-
Profile and debug training throughput bottlenecks across the compute, compiler, framework, and cluster-networking stack.
-
Collaborate closely with ML researchers, data engineers, and hardware engineers to define requirements and align on architectural decisions for large-scale training infrastructure.
-
Implement custom kernels and build benchmarks to track and validate training-performance optimizations across the stack.
Qualifications
-
4+ years of total experience, including 2+ years working on large-scale model training platforms.
-
Experience with distributed training frameworks (Megatron-LM, TorchTitan), efficiently utilizing GPU/accelerator clusters at scale.
-
Experience with RL framework such as Verl/MILES and RL algorithms(GRPO, PPO)
-
Experience with collective communication libraries (e.g., NCCL) and optimizing communication/compute overlap in multi-node training.
-
Experience using profiling tools (NVIDIA Nsight, PyTorch Profiler, or XLA/JAX profiling) to identify and resolve training bottlenecks.
-
Proficient in Python or C++.
-
Experience with ML compilers, performance tooling, or hardware/software co-design (e.g., GPU/TPU/accelerator architectures).
Note: You do not have to meet all the requirements below to be considered for this position:
Base Salary Range
There are three major components to compensation for this position: salary, Amazon Restricted Stock Units (RSUs), and Zoox Stock Appreciation Rights. A sign-on bonus may be offered as part of the compensation package. The listed range applies only to the base salary. Compensation will vary based on geographic location and level. Leveling, as well as positioning within a level, is determined by a range of factors, including, but not limited to, a candidate's relevant years of experience, domain knowledge, and interview performance. The salary range listed in this posting is representative of the range of levels Zoox is considering for this position.
Zoox also offers a comprehensive package of benefits, including paid time off (e.g. sick leave, vacation, bereavement), unpaid time off, Zoox Stock Appreciation Rights, Amazon RSUs, health insurance, long-term care insurance, long-term and short-term disability insurance, and life insurance.
Salary
177,000 – 270,000 USD per year
Employment type
Full-time
Source details
Source notes
Source excerptsSelected passages from the saved posting. Check the full description for conditions and exceptions.
- Pay
Salary 177,000 – 270,000 USD per year About Zoox
- Location & working pattern
Foster City, CA
Working pattern and location restrictions need checking in the full posting.
- Work authorization
No clear work-authorization passage found. Eligibility is unconfirmed.
See how this role fits your experience
Add your resume to compare the role’s scope, tools and requirements with your experience.