ML Platform Engineer, Distributed Training

Posted 2 Days Ago
Hiring Remotely in USA
Remote
159K-219K Annually
Senior level
Artificial Intelligence • Big Data • Cloud • Information Technology • Machine Learning
The Role
Build and operate distributed GPU training platforms for foundation models. Run and debug multi-node PyTorch and Ray Train workloads, tune GPU networking and scheduling, profile performance, and develop benchmarks and evaluation harnesses. Partner with infrastructure teams and client scientists to improve throughput, reliability, cost efficiency, observability, and end-to-end training pipelines on GKE.
Summary Generated by Built In
About Egen: 
 
Egen is a fast-growing and entrepreneurial company with a data-first mindset. We bring together the best engineering talent working with the most advanced technology platforms, including Google Cloud and Salesforce, to help clients drive action and impact through data and insights. We are committed to being a place where the best people choose to work so they can apply their engineering and technology expertise to envision what is next for how data and platforms can change the world for the better. We are dedicated to learning, thrive on solving tough problems, and continually innovate to achieve fast, effective results. If this describes you, we want you on our team.
 
Want to learn more about life at Egen? Check out these resources in addition to the job description.
 
Meet Egen
Life at Egen
Culture and Values at Egen
Career Development at Egen
Benefits at Egen
 
About the opportunity

We're looking for an engineer who is equally at home running a GPU training platform and the large-scale model training that runs on it. You'll be embedded with a healthcare and life-sciences AI client that trains foundation models on a large GKE-based H200 and GB200 cluster. The client is reorganizing so that its foundation-model platform and its agent platform teams work more closely together.

The role centres on the training stack and the scientists who use it. Solid infrastructure knowledge and experience are required, but the client already has strong SRE and infrastructure capacity, so running the cluster isn't the focus: you'll partner with that team and understand the training workloads well enough to tune, benchmark, and debug them end to end. You'll work alongside other Egen engineers and client scientists, and you'll own your area outright.

What You Will Do:

    Training platform
  • Work with the client's infrastructure team on the GKE training platform (GPU node pools, job queueing and scheduling, quota and capacity) from the training workload's side.
  • Diagnose and fix platform failures at the root (for example a controller that keeps running out of memory) rather than restarting them.
  • Tune multi-node GPU networking (RDMA/GPUDirect, NCCL) for training throughput.
  • Distributed training
  • Run, scale and debug multi-node distributed training with PyTorch (FSDP, DeepSpeed, Megatron or NeMo-style) and Ray Train.
  • Profile jobs and remove bottlenecks across compute, memory, input pipelines and interconnect.
  • Benchmarking and evals
  • Design and run performance benchmarks for the cluster and training stack: throughput, scaling efficiency, cost per run.
  • Build and maintain model evaluation harnesses so changes to the platform or training recipes are measured, not assumed.
  • Instrument training and the platform with telemetry and monitoring, and keep end-to-end pipeline tests running against what users actually run.
  • Ownership on the client team
  • Treat the platform as your own: proactively find issues, plug gaps and propose solutions before anyone asks, keep the team informed of status and risks, and look after the scientists who use it.
  • Explain and defend your design decisions to client engineers and scientists in detail.
  • Stay effective while the client's priorities and team structure change.

Your Technical Toolkit:

  • Core Languages: Mastery of Python and shell scripting
  • AI/LLM Ecosystem: Extensive experience with Google Gemini, GPT-4, or LLaMA; deep knowledge of Prompt Engineering and Fine-tuning

  • Data & Search: Expertise in Vector Databases (Vertex AI Vector Search, pgvector, etc.) and implementing Semantic Search

  • Infrastructure: Hands-on experience with Google Cloud (Vertex AI) and building scalable software architectures

  • Frameworks: Proficiency in LangChain, LlamaIndex, or similar orchestration layers

  • Mindset: A strong software engineering foundation—you write clean, maintainable code and understand the full SDLC

Basic Qualifications:

  • 5+ years of engineering experience, with recent hands-on work on GPU training infrastructure and large-scale training.

  • Strong Python software engineering: production-quality code, tests and code review. Coding is central to this role, not an afterthought.

  • GKE and GPU scheduling: Kubernetes on Google Cloud with GPU node pools, job queueing (Kueue or similar), and capacity and quota management.

  • Distributed PyTorch and Ray Train: has personally run and debugged multi-node training jobs.

  • RDMA and GPU networking: GPUDirect/RDMA, NCCL tuning, multi-node GB200 (or H200) topology.

  • Benchmarking and evals: has designed performance benchmarks and model evaluations, and used them to make decisions.

  • Deep enough understanding to explain how and why a system works, without leaning on AI coding tools to carry the reasoning.

  • A track record of visible ownership: present, communicative, and accountable on a demanding team.

Personal Attributes:

    Nice to have
  • Healthcare or life-sciences data (genomics, clinical, biomedical), and working under HIPAA.

  • Observability and telemetry for training platforms: metrics, tracing, alerting, end-to-end pipeline testing.

Compensation & Benefits:
 
This role is eligible for our competitive salary and comprehensive benefits package to support your well-being:
- Comprehensive Health Insurance
- Paid Leave (Vacation/PTO)
- Paid Holidays
- Sick Leave
- Parental Leave 
- Bereavement Leave
- 401 (k) Employer Match
- Employee Referral Bonuses
 
Check out our complete list of benefits here - >https://egen.ai/people/#benefits
 
Important: All roles are subject to standard hiring verification practices, which may include background checks, employment verification, and other relevant checks.
 
EEO and Accommodations:
 
Egen is an equal opportunity employer and is committed to inclusion, diversity, and equity in the workplace. All qualified applicants will receive consideration for employment without regard to race, color, religion, sex, national origin, disability status, protected veterans’ status, or any other characteristic protected by federal, state, or local laws. Egen will also consider qualified applications with criminal histories, consistent with legal requirements. Egen welcomes and encourages applications from individuals with disabilities. Reasonable accommodations are available for candidates during all aspects of the selection process. Please advise the talent acquisition team if you require accommodations during the interview process.

Skills Required

  • 5+ years of engineering experience with recent hands-on GPU training infrastructure and large-scale training
  • Strong Python software engineering, including production-quality code, testing, and code review
  • Kubernetes on Google Cloud with GKE GPU node pools, job queueing such as Kueue, and capacity and quota management
  • Personally run and debugged multi-node distributed training jobs using PyTorch and Ray Train
  • Experience with GPUDirect/RDMA, NCCL tuning, and multi-node GB200 or H200 GPU topology
  • Designed performance benchmarks and model evaluations and used results to guide decisions
  • Ability to explain system behavior and technical decisions without relying on AI coding tools for reasoning
  • Track record of visible ownership, communication, and accountability on demanding teams
  • Healthcare or life-sciences data experience, including genomics, clinical, or biomedical data, and HIPAA experience
  • Observability and telemetry experience for training platforms, including metrics, tracing, alerting, and end-to-end pipeline testing
Am I A Good Fit?
beta
Get Personalized Job Insights.
Our AI-powered fit analysis compares your resume with a job listing so you know if your skills & experience align.

The Company
HQ: Naperville, IL
240 Employees
Year Founded: 2000

What We Do

Egen is a data engineering and cloud modernization firm partnering with leading Chicagoland companies to launch, scale, and modernize industry-changing technologies. We are catalysts for change who create digital breakthroughs at warp speed. Our team of cloud and data engineering experts are trusted by top clients in pursuit of the extraordinary. Our mission is to be an enabler of amazing possibilities for companies looking to use the power of cloud and data. We want to stand shoulder to shoulder with clients, as true technology partners, and make sure they succeed at what they have set out to do. We want to be disruptors, game-changers, and innovators who have played an important part in moving the world forward.

Similar Jobs

Luxury Presence Logo Luxury Presence

Product Manager

Marketing Tech • Real Estate • Software • PropTech • SEO
Easy Apply
Remote or Hybrid
United States
500 Employees
170K-200K Annually

Hex Logo Hex

Revenue Systems

Artificial Intelligence • Big Data • Software • Analytics • Business Intelligence • Big Data Analytics
Remote or Hybrid
United States
160 Employees
126K-158K Annually

Nexthink Logo Nexthink

Corporate Controller

Artificial Intelligence • Big Data • Cloud • Information Technology • Machine Learning • Software
Remote or Hybrid
United States
1200 Employees
245K-290K Annually

Qualtrics Logo Qualtrics

Front Line Sales Leader - Life Sciences - Corporate

Artificial Intelligence • HR Tech • Information Technology • Software • Business Intelligence
Remote
United States
7500 Employees
120K-145K Annually

Similar Companies Hiring

Onshore Thumbnail
Artificial Intelligence • Fintech • Software • Financial Services
New York, New York
60 Employees
Blee Thumbnail
Artificial Intelligence • Marketing Tech • Software
New York, New York
30 Employees
Vega Thumbnail
Artificial Intelligence • Automotive • Insurance • Transportation
US
65 Employees

Sign up now Access later

Create Free Account

Please log in or sign up to report this job.

Create Free Account