ML Systems Engineer

Cohere is the leading security-first enterprise AI company. We build cutting-edge foundation AI models and end-to-end products that are designed to solve real-world business problems. We are a global technology company headquartered in Toronto with key offices in London, New York City, San Francisco, Montreal, Paris, Berlin and Seoul.

• Build and own the training framework responsible for large-scale LLM training.
• Design distributed training abstractions (data/tensor/pipeline parallelism, FSDP/ZeRO strategies, memory management, checkpointing).
• Improve training throughput and stability on multi-node clusters (e.g., GB200/300, AMD, H200/100).
• Develop and maintain tooling for monitoring, logging, debugging, and developer ergonomics.
• Collaborate closely with infra teams to ensure our cluster, container environments, and hardware configurations support high-performance training.
• Investigate and resolve performance bottlenecks across the ML systems stack.
• Build robust systems that ensure reproducible, debuggable, large-scale runs.

• Strong engineering experience in large-scale distributed training or HPC systems.
• Deep familiarity with JAX internals, distributed training libraries, or custom kernels/fused ops.
• Experience with multi-node cluster orchestration (Slurm, Ray, Kubernetes, or similar).
• Comfort debugging performance issues across CUDA/NCCL, networking, IO, and data pipelines.
• Experience working with containerized environments (Docker, Singularity/Apptainer).
• A track record of building tools that increase developer velocity for ML teams.
• Excellent judgment around trade-offs: performance vs complexity, research velocity vs maintainability.
• Strong collaboration skills — you'll work closely with infra, research, and deployment teams.
• Any of the following would also be good to have for this role:
• Experience with training LLMs or other large transformer architectures.
• Contributions to ML frameworks (PyTorch, JAX, DeepSpeed, Megatron, xFormers, etc.).
• Familiarity with evaluation and serving frameworks (vLLM, TensorRT-LLM, custom KV caches).
• Experience with data pipeline optimization, sharded datasets, or caching strategies.
• Background in performance engineering, profiling, or low-level systems.
• Bonus: paper at top-tier venues (such as NeurIPS, ICML, ICLR, AIStats, MLSys, JMLR, AAAI, Nature, COLING, ACL, EMNLP).

• A weekly lunch stipend of $75/£75 or equivalent in your local currency for lunch.
• Full health and dental benefits, including a separate budget for mental health.
• RRSP matching, 401K, Pension Scheme.
• 100% Parental Leave top-up for up to 6 months, for either parent.
• Annual enrichment benefits: Arts & culture, fitness/wellness, quality time, and a workspace improvement credit.
• Education & learning stipend for conferences, courses, and coaching.
• 6 weeks of paid vacation (30 working days!)
• Budget for traveling to other offices if you are remote, plus an annual company offsite.
• Cohere is remote-friendly, but we also have offices in Toronto, London, New York City, San Francisco, Montreal, Paris, Berlin and Seoul with more opening soon.
• For those in the office: a daily lunch program, plenty of snacks, and regular community and social events.
• For those not near an office: a co-working benefit so you can work alongside others in your city.
• Everyone receives a $500 home office stipend to set up your workspace properly.