Member of Technical Staff, Training Infra Engineer

Stack

Responsibilities

• ▹Design and write high-performant, scalable software for training
• ▹Improve the training setup from an infrastructure and codebase performance standpoint
• ▹Craft and implement tools to speed up training cycles and improve training infrastructure efficacy
• ▹Research, implement and experiment with ideas on the supercompute and data infrastructure

Requirements

• ▹Extremely strong software engineering skills
• ▹Proficiency in Python and ML frameworks such as JAX, PyTorch and XLA/MLIR
• ▹Experience with distributed training infrastructures (Kubernetes, Slurm) and frameworks such as Ray
• ▹Experience using large-scale distributed training strategies
• ▹Hands-on experience training large models at scale, having contributed to training infrastructure tooling/setup

Nice to have

• ▹Paper at top-tier venues such as NeurIPS, ICML, ICLR, AIStats, MLSys, JMLR, AAAI, Nature, COLING, ACL, EMNLP

Soft skills

What we offer

• ▹Weekly lunch stipend ($75/£75 or local equivalent)
• ▹Full health and dental benefits, including a mental health budget
• ▹RRSP matching, 401(k) or pension scheme
• ▹100% parental leave top-up for up to 6 months, for either parent
• ▹Annual enrichment benefits (arts & culture, fitness/wellness, workspace improvement credit)
• ▹Education & learning stipend for conferences, courses and coaching
• ▹6 weeks of paid vacation (30 working days)
• ▹Travel budget to other offices plus an annual company offsite
• ▹$500 home office stipend and co-working benefit for remote employees

About the company

Cohere is a leading security-first enterprise AI company, building frontier foundation models and end-to-end AI products for real-world business problems. It is a global team co-headquartered in Toronto and San Francisco, with offices in London, New York City, Montreal, Seoul, Germany and Paris.