> Markdown version of [/jobs/ext/4625852-staff-ai-training-platform-engineer](https://www.wearedevelopers.com/jobs/ext/4625852-staff-ai-training-platform-engineer). Every page supports `.md` or `Accept: text/markdown`. Links point to the HTML versions so they work for humans too. Agent guide: [/agents.md](https://www.wearedevelopers.com/agents.md).

---

# Staff AI Training Platform Engineer

- **Company:** The Unconventional
- **Location:** Los Angeles, CA, United States
- **Experience:** Expert
- **Contract:** Permanent contract
- **Skills:** Artificial Intelligence, Artificial Neural Networks, C++ (Programming Language), Nvidia CUDA, Computer Programming, Computer Engineering, Software Debugging, Python (Programming Language), Machine Learning, Performance Tuning, Software Engineering, Data Streaming, System Programming, High Performance Computing, Pytorch, DeepSpeed, Large Language Models, Generative AI, Kubernetes, Information Technology, SGLang, Model Inference, Megatron
- **Published:** November 7, 2026
- **Apply:** https://www.juju.com/job/28_25303_3d6f0463eed1d38055f78636868eb4d2

## About the Role

* Education: BS in Computer Science, Physics, Electrical Engineering, or Applied Math.
 * Experience: 5+ years of experience in AI/ML engineering. Veteran of the modern ML software stack. Demonstrated ability to map state-of-the-art AI model architectures (e.g., transformers, Mixture of Experts, diffusion models) to system performance implication. Deep expertise in how models are partitioned across a cluster, with a mastery of communication primitives, and parallelism strategies.
 * Software Development: Strong programming skills in Python or C++. Proven track record of implementing, debugging, and maintaining production-grade training frameworks-such as Megatron-LM, DeepSpeed, Ray, PyTorch Lightning-turning raw compute into a reliable model-building factory., * MS/PhD or equivalent research/project experience in AI/ML or high-performance computing, with publications
 * Experience in training, post-training large-scale LLMs and generative models.
 * Experience in one or more of the following technologies: Kubernetes, GPU kernels and performance tuning, LLM inference, CUDA, Triton, NCCL, vLLM, SGLang, VeRL, TRL.

## Description

As a Member of Technical Staff, AI Training Platform, you will be a core contributor to the infrastructure that powers our model training ecosystem. You will build and scale the end-to-end tooling required to train, evaluate, and benchmark models. Your work will directly accelerate researchers' ability to map neural networks to novel hardware and push the boundaries of physics-based compute., * Training Infrastructure: Architect, scale, and maintain the core AI/ML/RL training platform and infrastructure. Design and scale multi-node distributed training systems, implementing elastic sharding and robust data streaming pipelines for fast, large-scale iteration. Implement and robust model checkpointing and recovery mechanisms.

 * Framework Development: Maintain and expand our proprietary training framework, ensuring it provides a robust and flexible foundation for all internal model training.
 * Optimization & Benchmarking: Develop and optimize kernels using low-level programming models like CUDA and Triton. Design rigorous benchmarking suites to track Model Flops Utilization (MFU), memory bandwidth, and convergence stability.
 * Hardware Evaluation Tooling: Build and iterate on tooling to enable rapid, high-level evaluation of novel hardware ideas against benchmarks.
 * Cross-Functional Collaboration: Act as a translator, discussing algorithmic trade-offs with theorists and converting model requirements into concrete specifications for infrastructure and hardware engineering teams.
 * Optimization & ML Hillclimbing: Design systems to track and visualize quality-vs-efficiency Pareto frontiers, enabling researchers to optimize models for both performance and energy consumption.