SlipstreamJobs tracks this role from the company's public career site. Apply directly on the employer's site.
Sarvam AI is building India's full-stack sovereign AI platform, backed by Lightspeed, Peak XV, and Khosla Ventures. You will own the infrastructure for training the company's next generation of foundational models.
This is a deep systems role focused on distributed training at scale. You will be responsible for building and maintaining the distributed training stack across large GPU clusters, designing and implementing parallelism strategies (data, tensor, pipeline, sequence, expert), and optimizing end-to-end training throughput including kernel performance, communication overlap, memory layout, checkpointing, and data loading.
Key responsibilities include:
- Build, maintain, and continually push the limits of the distributed training stack across large GPU clusters
- Design and implement parallelism strategies and reason about which combinations make sense for different architectures and scales
- Profile and optimize end-to-end training throughput: kernel performance, communication overlap, memory layout, checkpointing, data loading
- Write and tune custom GPU kernels (CUDA, Triton) where off-the-shelf solutions leave performance on the table
- Own the reliability of long-running training jobs — fault tolerance, checkpoint integrity, deterministic restarts, automated detection of slow nodes and silent corruption
- Partner closely with researchers and data teams to ensure architectural ideas can be trained efficiently and data pipelines don't become bottlenecks
The role demands high standards: a bad MFU number, slow data loader, or network issue can cost weeks and millions of dollars. You will be expected to identify problems before they occur and fix them quickly when they do.
Requirements:
- BS or MS in Computer Science or closely related technical field (or equivalent demonstrated experience)
- 3+ years of experience building ML training infrastructure or large-scale distributed systems (exceptional early-career candidates with strong systems background considered)
- Hands-on experience training large models with distributed training frameworks (Megatron-LM, DeepSpeed, FSDP, NeMo, or equivalent); must have been on-call for a real pretraining run
- Deep working knowledge of GPU architecture, CUDA programming model, and standard profiling tools (Nsight, PyTorch profiler, etc.)
- Strong PyTorch internals knowledge; comfortable reading and modifying low-level training code, not just using high-level APIs
- Meaningful open-source contributions in training infrastructure ecosystem (Megatron, DeepSpeed, PyTorch, vLLM, Triton, NCCL, or similar)
Bonus qualifications:
- Custom CUDA or Triton kernel development with measurable wins on real workloads
- Direct experience training models at 10B+ parameters or on 1000+ GPU clusters
- Familiarity with cluster orchestration and job scheduling (Slurm, Kubernetes) at scale
- Experience with mixed precision (BF16, FP8), quantization-aware training, or other numerical work in production training runs
- First-author papers or technical reports on training systems, scaling, or model efficiency
