Member of Technical Staff — Pretraining Infra
Own and scale the distributed training infrastructure for large-scale omni model pretraining across GPU clusters, covering job orchestration, parallelism, GPU communication, data loading, and performance optimization.
About the job
What You’ll Own
- Own the distributed training stack for omni model pretraining, from 0→1 system design to 1→10 scaling across large GPU clusters.
- Build and operate the core training runtime: job orchestration, distributed execution, checkpointing, recovery, monitoring, and debugging for long-running training jobs.
- Optimize large-scale training performance across parallelism strategy, GPU communication, memory usage, data throughput, MFU, step time, and end-to-end training efficiency.
- Build infrastructure for omni training workloads: high-throughput audio/video/text data loading, temporal alignment, variable sequence handling, multimodal synchronization, and memory-efficient training.
- Evolve the platform as model architectures, training recipes, data mixtures, sequence lengths, hardware constraints, and research directions change.
What We’re Looking For
- Hands-on experience running large-scale distributed training jobs across large GPU clusters; experience at hundreds of GPUs minimum, 1,000+ GPUs a strong plus.
- Deep understanding of distributed training mechanics: data/tensor/pipeline/sequence parallelism, gradient communication, collectives, mixed precision, activation checkpointing, optimizer state, memory pressure, and framework-level tradeoffs.
- Strong understanding of GPU communication and performance debugging: NCCL, all-reduce/all-gather/reduce-scatter, communication-computation overlap, topology, synchronization, stragglers, low MFU, OOMs, checkpoint bottlenecks, and data starvation.
- Practical experience with at least one major large-scale training stack such as Megatron, PyTorch FSDP, DeepSpeed, or equivalent internal infrastructure.
- Understanding of omni or multimodal training challenges, especially audio/video/language data, long temporal context, variable sequence lengths, modality-specific bottlenecks, and high-throughput dataloading.
- Strong software engineering fundamentals, curiosity, and adaptability to new model architectures, training frameworks, hardware constraints, and research ideas.
Bonus Points
- Prior 0→1 experience building large-scale training infrastructure or deeply modifying core training frameworks, runtimes, checkpointing, or debugging systems.
- Experience training large omni or multimodal models involving audio, video, text, or long-context temporal data.
- Experience with adjacent infrastructure areas such as RL/post-training, data infrastructure, synthetic data generation, evaluation, or serving.
- Publications or substantial open-source contributions in ML systems, distributed systems, HPC, GPU performance, or training infrastructure.
Compensation
$300,000 – $400,000 base salary, plus meaningful equity.
Skills
Pytorch Fsdp, Deepspeed, Megatron, Nccl, Data Parallelism, Tensor Parallelism, Pipeline Parallelism, Sequence Parallelism, Mixed Precision, Activation Checkpointing
Similar jobs
ML Engineering jobs
Garner Health
New York, NY
Leads and builds a team of Applied Scientists developing production algorithmic systems for healthcare optimization, LLM applications, and member engagement. Requires 6+ years of relevant industry experience, strong technical judgment, and hands-on expertise across machine learning and optimization.
OpenAI
San Francisco, CA
Build production machine learning systems for model customization, post-training, evaluation, and AWS-native API integration. The role requires 7+ years of relevant engineering experience and expertise in deep learning, transformers, LLM fine-tuning, and production ML infrastructure.
OpusClip
Mountain View, CA
Leads a hands-on AI engineering team developing, evaluating, and deploying large-scale multimodal and video models. The role combines post-training, inference optimization, product experimentation, technical roadmap ownership, and people management.
Discord
United States
Leads Discord’s Safety ML team, setting technical direction and overseeing production machine learning systems for content understanding, account integrity, and platform abuse. Requires substantial machine learning and engineering management experience, hands-on technical depth, and experience delivering ML systems at scale.
Baselayer
San Francisco, CA
Senior AI Engineer responsible for production LLM agents that enrich business identity data through web discovery, verification, classification, and risk scoring. The role requires strong asynchronous Python, agent and evaluation expertise, browser automation, and experience operating AI systems in production.