# Member of Technical Staff, Training Infra Engineer

**Company:** [Cohere](https://hotfix.jobs/companies/cohere)
**Location:** Remote
**Role:** ML Engineering
**Skills:** Python, JAX, PyTorch, xla, mlir, Kubernetes, slurm, Ray, Distributed Training, supercomputing
**Posted:** 2025-02-20

> Build and optimize scalable infrastructure for training large frontier AI models, bridging research and production. The role requires strong software engineering, Python and ML framework expertise, and hands-on experience with distributed training at scale.

## Job Description

## Responsibilities
- Design and write high-performance, scalable software for model training.
- Improve training infrastructure and codebase performance.
- Build tools to accelerate training cycles and improve the efficacy of training infrastructure.
- Research, implement, and experiment with ideas on supercompute and data infrastructure.
- Collaborate with leading researchers to ship state-of-the-art models to production.

## Requirements
- Extremely strong software engineering skills.
- Proficiency in Python and machine learning frameworks such as JAX, PyTorch, and XLA/MLIR.
- Experience with distributed training infrastructure, including Kubernetes and Slurm.
- Experience with distributed training frameworks such as Ray.
- Experience using large-scale distributed training strategies.
- Hands-on experience training large models at scale and contributing to training infrastructure tooling or setup.

## Nice to Have
- A paper published at a top-tier venue such as NeurIPS, ICML, ICLR, AIStats, MLSys, JMLR, AAAI, Nature, COLING, ACL, or EMNLP.

## Compensation and Benefits
- Weekly lunch stipend of $75/£75 or equivalent in local currency.
- Full health and dental benefits, including a separate mental health budget.
- RRSP matching, 401K, or pension scheme.
- 100% parental leave top-up for up to 6 months for either parent.
- Annual enrichment benefits for arts and culture, fitness and wellness, quality time, and workspace improvements.
- Education and learning stipend for conferences, courses, and coaching.
- Six weeks of paid vacation (30 working days).
- Travel budget for remote employees to visit other offices, plus an annual company offsite.
- Coworking benefit for employees not near an office.
- $500 home office stipend.

## Similar roles

- [Staff Machine Learning Engineer, Traffic Intelligence](https://hotfix.jobs/jobs/cee28c3a-eb64-42f7-a584-8d2b4a0980b8) - Airbnb - Remote - $212k – $265k/yr
- [Member of Technical Staff, Agentic Environments](https://hotfix.jobs/jobs/f3ceb248-2d3e-4b28-a5d1-10891969916e) - Cohere - Remote
- [Staff Machine Learning Infrastructure Engineer, Embedding Platform](https://hotfix.jobs/jobs/9b4387bf-ac16-48ba-810b-2e6a612e3617) - Reddit - Remote - $253k – $355k/yr
- [Senior Staff Machine Learning Engineer, Feed Relevance](https://hotfix.jobs/jobs/a2f9b178-c6cf-49bf-ba93-00688d1027fb) - Reddit - Remote - $266k – $372k/yr
- [Staff Applied Scientist](https://hotfix.jobs/jobs/d518c7a7-07d4-4f06-af79-9c478df769ca) - Garner Health - New York, NY - $300k – $390k/yr

**Apply:** https://hotfix.jobs/jobs/5ef2a790-94ad-4765-90e3-4e908e439dcc
**Canonical:** https://hotfix.jobs/jobs/5ef2a790-94ad-4765-90e3-4e908e439dcc