Distributed ML Infrastructure Engineer

We are looking for a Distributed ML Infrastructure Engineer to build and optimise the large-scale distributed training systems behind our foundation models. You will own the performance, scalability and reliability of training infrastructure spanning multi-node GPU and TPU clusters. This is a hands-on systems role at the core of how BrahmAI gets trained, working closely with our research teams to turn raw compute into efficient, dependable training throughput.

Responsibilities:

• Build and maintain distributed training infrastructure for large-scale AI workloads.
• Optimise training performance through profiling, memory, communication, and compute optimisation.
• Implement distributed training strategies including data, tensor, pipeline and sequence parallelism.
• Improve training throughput, scalability and reliability across multi-node clusters.
• Develop automation, monitoring, profiling and debugging tools for production training systems.
• Collaborate with research and engineering teams to deliver efficient and scalable training infrastructure.

Requirements:

• 3+ years of experience in Distributed Systems, HPC or ML Infrastructure.
• Strong proficiency in Python, Linux and shell scripting.
• Experience with multi-node GPU/TPU clusters, CUDA, PyTorch Distributed, NCCL, MPI/OpenMPI and distributed communication.
• Strong understanding of GPU/TPU architecture, distributed computing, profiling and performance optimisation.
• Experience with Docker, Kubernetes, Git and containerised development.
• Strong debugging, profiling and performance analysis skills.

Good to Have / Bonus Points:

• JAX, XLA, MaxText, XPK, PJRT and TPU training infrastructure.
• Triton, Pallas or custom kernel optimisation.
• Slurm, Ray, InfiniBand/RDMA and distributed storage systems.
• Experience with TensorBoard Profiler, Nsight Systems/Compute or XProf.
• Experience supporting large-scale LLM pretraining and distributed training at scale.