X
ML Engineer, Grok (Remote — Worldwide)
JAXPyTorchCUDADistributed TrainingLLMsNLPHPCPythonC++Reinforcement Learning
About the role
Join the core team responsible for building Grok, the world’s most powerful AI, by developing and scaling large language models on the frontier of reasoning and knowledge. As an ML Engineer at xAI, you will be tasked with training, optimizing, and deploying models across tens of thousands of GPUs using our custom compute clusters. This role is focused on pushing the limits of distributed systems and algorithmic efficiency to provide real-time, witty, and truthful information to global users.
Responsibilities
- Design and implement massive-scale pre-training runs on tens of thousands of H100/B200 GPUs.
- Develop novel architectures and optimization techniques to improve the data-efficiency and reasoning capabilities of Grok.
- Optimize the end-to-end ML pipeline, from data ingestion and tokenization to low-latency inference.
- Collaborate with the data team to curate and filter planetary-scale datasets for high-quality training.
- Periodically participate in on-call rotations to ensure the stability of long-running training jobs and live inference APIs.
- Implement and experiment with cutting-edge research papers to maintain Grok’s position at the frontier of LLM performance.
Requirements
- Proven experience training large-scale transformer models (10B+ parameters) from scratch.
- Deep expertise in distributed training frameworks such as JAX, PyTorch, or custom CUDA kernels.
- Strong understanding of high-performance computing (HPC) primitives like NCCL, InfiniBand, and RDMA.
- Experience with RLHF, DPO, or other alignment techniques to improve model reasoning and safety.
- Demonstrated ability to write production-grade Python and C++/CUDA code in a fast-paced environment.
- Bachelor’s, Master’s, or PhD in Computer Science, Physics, Mathematics, or a related quantitative field.
- Ability to work independently in a remote, high-intensity environment with a bias toward action.
Benefits
Equity, healthcare, remote-first.
Disclaimer: MMagic.ai connects talented people with AI companies around the world. While we work hard to feature quality opportunities, we don't independently verify employers, candidates, salaries, or hiring outcomes. We encourage you to research each opportunity and company before applying or making an offer.