Skip to content
CohereCohereParis, France

Member of Technical Staff, Training Infra Engineer

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.

Salary not listed
RemoteML Engineering

About the role

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.

Skills

PythonJAXPyTorchxlamlirKubernetesslurmRayDistributed Trainingsupercomputing

Similar roles

ML Engineering jobs
Airbnb

Staff Machine Learning Engineer, Traffic Intelligence

AirbnbUnited States

Architects and operates production machine-learning systems that classify web and API traffic, detect bots and scrapers, and support real-time mitigation at internet edge latency. The role requires 9+ years of applied ML experience in adversarial domains and strong expertise in evaluation, data pipelines, and large-scale systems.

212k – 265k/yrRemote9+ YOEML Engineering
Cohere

Member of Technical Staff, Agentic Environments

CohereNew York, NY

Build scalable software and tools for frontier model training, research experimentation, and production machine learning systems. The role requires strong Python and distributed-training expertise, experience with ML frameworks and infrastructure, and the ability to optimize and debug large language model systems.

Salary not listedRemoteML Engineering
Reddit

Staff Machine Learning Infrastructure Engineer, Embedding Platform

RedditUnited States

Leads the technical direction of large-scale ML infrastructure for embedding, recommendation, and personalization systems. The role requires 8+ years of ML engineering experience, expertise in deep learning and distributed training, and strong leadership across research, infrastructure, and production deployment.

253k – 355k/yrRemote8+ YOEML Engineering
Reddit

Senior Staff Machine Learning Engineer, Feed Relevance

RedditUnited States

Leads the technical direction and development of large-scale, GenAI-powered recommendation and feed-ranking systems. Requires 10+ years of industry experience in relevance-driven products, deep expertise in machine learning and recommendations, and strong organizational influence and mentoring skills.

266k – 372k/yrRemote10+ YOEML Engineering
Garner Health

Staff Applied Scientist

Garner HealthNew York, NY

Leads end-to-end development of production algorithmic systems for healthcare, spanning machine learning, optimization, and LLM applications. The player-coach role requires 6+ years of industry experience, strong problem-solving and metrics judgment, and technical leadership of a small team.

300k – 390k/yrHybrid7+ YOEML Engineering