ML Platform Engineer, Distributed Training

Remote · USD 159,300-219,245 / year · 5+ years experience · Remote · Posted Oct 19, 2026

At a glance

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

As a Senior AI Engineer, you will be at the forefront of our Generative AI initiatives, taking document intelligence and RAG pipelines from initial prototype to robust, scalable production services.
You will solve for real-world constraints like latency, error handling, and cost optimization while collaborating with a diverse range of clients.

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 rather than restarting them.
Tune multi-node GPU networking (RDMA/GPUDirect, NCCL) for training throughput.
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.
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.
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.

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.
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.

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.

Show all 15 skills