2 days ago
Base Salary
$180k - $450k/yr
Responsibilities
- Own MFU and step-time targets for large-scale pretraining runs across dense and mixture-of-experts models.
- Develop or integrate fused and custom kernels for attention, normalization, and mixture-of-experts routing using CUDA, Triton, and CUTLASS.
- Optimize communication overlap, pipeline execution, activation recomputation, memory use, sharding, and data, tensor, pipeline, sequence, context, and expert parallelism.
- Profile GPU kernels, collective communication, data loading, and mixed-precision training to identify throughput and numerical-stability bottlenecks.
- Develop reusable training components for new architectures and multimodal speech, vision, and text encoders.
- Investigate performance regressions and training failures involving communication, numerical changes, hardware faults, and straggling nodes.
- Collaborate with architecture researchers and compute infrastructure engineers to align model designs, training software, and cluster capabilities.
- Improve checkpointing, restart time, failure detection, and observability, and participate in on-call support during large runs.
Requirements
- Experience developing distributed training systems and working with Megatron-LM, NeMo, or a comparable framework and its parallelism strategies.
- Strong understanding of data, tensor, pipeline, sequence, context, and expert parallelism, including communication patterns, memory tradeoffs, and failure modes.
- Experience using GPU profilers to analyze kernel execution, memory bandwidth, communication overhead, and data-loading stalls.
- Working knowledge of mixed-precision training, including BF16 and FP8, and methods for diagnosing numerical instability.
- Strong Python and C++/CUDA engineering skills, with experience building reliable tools for researchers.
- Experience developing CUDA or Triton kernels, or using CUTLASS, for attention, normalization, or mixture-of-experts models.
- Demonstrated record of measurably improving MFU or throughput on large-scale training jobs.
- Bonus qualifications include large-scale mixture-of-experts training, NCCL or topology-aware scheduling, fault-tolerant or elastic training, multimodal encoder workload balancing, Slurm, Kubernetes, and fleet-scale GPU health monitoring.
Benefits
- Full-time position.
Tech Stack
Categories
