Research Engineer

Job description

Job responsibilities

• Design and implement scalable systems for training large ML models.
• Enable efficient workflows for data ingestion, training, and iteration.
• Develop and optimize distributed training systems across hundreds of GPUs.
• Implement strategies for parallelization, sharding, and efficient compute utilization.
• Improve training efficiency through techniques such as attention optimizations, kernel fusion, and memory management.
• Partner closely with modeling teams to accelerate iteration speed and reduce training costs.
• Build internal tools for experiment tracking, monitoring, and debugging.
• Implement systems for tracking training performance, failures, and resource utilization.

Job requirements

• Strong experience building infrastructure for large-scale ML training.
• Deep understanding of how modern LLM/VLM systems are trained and scaled.
• Proven experience setting up and scaling distributed training across hundreds of GPUs.
• Strong understanding of parallelization strategies (data, model, pipeline parallelism).
• Strong proficiency in Python programming.
• Expert-level proficiency in PyTorch and/or JAX.
• Strong understanding of techniques like attention optimization, kernel fusion, and efficient memory usage.