Machine Learning Engineer — Pre-training (LLM)

Institute of Foundation Models

Sunnyvale, CA

On-site

Who this role is best for

A natural match if you have experience with distributed training frameworks and large-scale pre-training.

Best fit for

Candidates with 5+ years in distributed ML systems and a track record in large-scale pre-training
— “5+ years of experience in ML systems, infra, or distributed training”
Individuals skilled in modifying training frameworks and debugging multi-node environments
— “Experience modifying distributed ML frameworks (e.g., DeepSpeed, FSDP)”
Engineers who can translate mathematical specs into distributed implementations and build scalable infrastructure
— “Ability to implement algorithms across GPUs/nodes based on mathematical specs”

Things to consider

The role demands a high level of software engineering rigor and reliability in production code
— “Write production-quality code and tests for ML infra in PyTorch or JAX”
Candidates must be prepared to work on complex, multi-node environments with various parallelism strategies
— “Select and validate data, tensor, pipeline, expert, and context parallelism strategies”

How to stand out

Highlight experience in optimizing large-scale pre-training workflows and distributed training frameworks
— “Experience with large-scale pre-training”
Showcase your expertise in building robust data-loading and checkpointing systems for ML
— “Build and maintain data-loading and checkpoint/restart workflows”

Pace · Fast PacedCollaboration · HighAutonomy · MediumDecision Impact · Team

Derived from job-description analysis by Serendipath's career intelligence engine.

What success looks like

extend training frameworks
implement optimizers
create launch configs
build monitoring systems
write production-quality code

Typical background

machine learning engineerdistributed systems engineer

Skills & requirements

Required

Distributed Training FrameworksOptimizer ImplementationMetrics & MonitoringInfrastructure EngineeringData Loading & Checkpointing

Preferred

Mixed-precision TrainingPerformance ProfilingOpen-source Contributions

Stack & domain

Machine LearningDistributed TrainingOptimizer ImplementationMulti-node Launch ScriptsMetrics MonitoringInfrastructure EngineeringData LoadingCheckpointingLarge-scale Machine LearningPre-trainingMixed-precision TrainingPerformance ProfilingKernel FusionMemory OptimizationCUDATriton KernelCustom Training PipelinesTraining InfrastructurePerformance TuningProblem-solvingTeamworkCommunicationNegotiationInfluenceStrategic ThinkingDecision-makingTime ManagementAdaptabilityEmpathyTrust-buildingCollaborationCreativityCustomer ServicePresentationPersuasion

About the role

Original posting from Institute of Foundation Models via Lever

About the Institute of Foundation Models
We are a dedicated research lab for building, understanding, using, and risk-managing foundation models. Our mandate is to advance research, nurture the next generation of AI builders, and drive transformative contributions to a knowledge-driven economy.
As part of our team, you’ll have the opportunity to work on the core of cutting-edge foundation model training, alongside world-class researchers, data scientists, and engineers, tackling the most fundamental and impactful challenges in AI development. You will participate in the development of groundbreaking AI solutions that have the potential to reshape entire industries. Strategic and innovative problem-solving skills will be instrumental in establishing MBZUAI as a global hub for high-performance computing in deep learning, driving impactful discoveries that inspire the next generation of AI pioneers.
Key Responsibilities:
Distributed Framework Ownership – Extend or modify training frameworks (e.g., DeepSpeed, FSDP) to support new use cases and architectures.
Optimizer Implementation – Translate mathematical optimizer specs into distributed implementations.
Launch Config & Debugging – Create and debug multi-node launch scripts with flexible batch sizes, parallelism strategies, and hardware targets. Select and validate data, tensor, pipeline, expert, and context parallelism strategies as appropriate for the model and cluster.
Metrics & Monitoring – Build systems for experiment tracking, job monitoring, and logging usable by collaborators and researchers.
Infra Engineering – Write production-quality code and tests for ML infra in PyTorch or JAX; ensure reliability and maintainability at scale.
Data Loading & Checkpointing – Build and maintain data-loading and checkpoint/restart workflows, restoring model, optimizer, RNG, and data progress after interruptions.
Qualifications:
Must-Haves:
5+ years of experience in ML systems, infra, or distributed training
Experience modifying distributed ML frameworks (e.g., DeepSpeed, FSDP, FairScale, Horovod)
Strong software engineering fundamentals (Python, systems design, testing)
Proven multi-node experience (e.g., Slurm, Kubernetes, Ray) and debugging skills (e.g., NCCL/GLOO)
Ability to implement algorithms across GPUs/nodes based on mathematical specs
Experience working on an ML platform/ infrastructure, and/or distributed inference optimization team
Experience with large-scale machine learning workloads (strong ML fundamentals)
Experience with large-scale pre-training
Nice-to-Haves:
Exposure to mixed-precision training (e.g., bf16, fp8) with accuracy validation
Familiarity with performance profiling, kernel fusion, or memory optimization
Open-source contributions or published research (MLSys, ICML, NeurIPS)
CUDA or Triton kernel experience
Experience building custom training pipelines at scale and modifying them for custom needs
Deep familiarity with training infrastructure and performance tuning
Compensation:
Salary Range
The posted salary range represents the company’s good faith estimate of the compensation for this position upon hire. The actual compensation offered may vary within this range depending on individual qualifications, including but not limited to relevant skills, experience, education, certifications, geographic location, and specific business needs.
The posted salary range represents the company’s good faith estimate of the compensation for this position upon hire. The actual compensation offered may vary within this range depending on individual qualifications, including but not limited to relevant skills, experience, education, certifications, geographic location, and specific business needs.
Benefits Include
*Comprehensive medical, dental, and vision benefits
*Bonus
*401K Plan
*Generous paid time off, sick leave and holidays
*Paid Parental Leave
*Employee Assistance Program
*Life insurance and disability

Source: Institute of Foundation Models careers (Lever)