Level

Reflection AIPosted 6mo ago

Member of Technical Staff - Research Software Engineer

Member of Technical Staff - Research Software Engineer at Reflection AI scores 100 out of 100 on AI centrality, which makes it AI Level 4 of 4 (Builds AI) on this board. The level measures how much of the work is AI, not seniority.

New York, NYleadFullTime

AI in this role

pytorchjax

Our Mission

Reflection is a research lab making intelligence open and accessible for everyone to use, customize, and build on. We build open models that let anyone control their intelligence and help shape the future of AI. Our mission: make intelligence open and accessible to all.

The Roles Mission

Bridge the gap between research and production by turning cutting-edge algorithms into scalable training systems. You will design and optimize the core infrastructure behind frontier AI models — from reinforcement learning training loops and distributed GPU training to massive-scale data pipelines.

Our systems train models across thousands of GPUs and process petabyte-scale datasets. We care deeply about numerical stability, throughput, and reproducibility.

What This Team Does

This team owns and evolves the core infrastructure behind our training systems.

We focus on:

  • Reinforcement learning training infrastructure

  • Distributed training and inference systems

  • Experiment infrastructure and reproducibility

  • Large-scale data pipelines

The goal is to build the engineering foundation that allows researchers to iterate quickly while training models at massive scale.

About the Role

You will architect and optimize the core training infrastructure that powers our models. This includes RL training loops, distributed GPU systems, and large-scale data pipelines.

You will work closely with researchers to transform new ideas into reliable, scalable training systems.

Responsibilities include:

  • Designing and optimizing large-scale training loops and data pipelines.

  • Implementing state-of-the-art techniques and ensuring they are numerically stable and computationally efficient.

  • Building internal tooling for launching, monitoring, and reproducing complex experiments.

  • Diagnosing deep bottlenecks across the training stack (GPU memory issues, communication overhead, dataloader stalls).

  • Translating research prototypes into reusable, production-grade infrastructure.

What You'll Work With

Distributed Training

  • GPU parallelism (data, tensor, pipeline, expert)

  • Large-scale distributed training infrastructure

  • Communication optimization (NCCL, RDMA, GPU interconnects)

  • FSDP / ZeRO and model sharding

Orchestration & Runtime Systems

  • Ray, Kubernetes, Slurm

  • Distributed runtimes and async systems

  • Containerization and sandboxing

Frameworks

  • PyTorch

  • JAX

  • Megatron-style training stacks

  • Triton / custom kernels

Data Infrastructure

  • Large-scale dataset curation pipelines

  • Deduplication and filtering systems

  • Tokenization and preprocessing

  • Distributed data processing frameworks

About You

  • You are a strong software engineer who speaks the language of machine learning.

  • You may not have a PhD, but you know how to implement a research paper.

  • You have deep experience in at least one of the following: Distributed Training & Inference or Data Infrastructure

  • You enjoy working at the boundary between:

    • Machine learning algorithms

    • Distributed systems

    • High-performance computing

  • You care deeply about performance, numerical stability, and reproducibility.

  • You thrive in high-agency environments and enjoy solving hard technical problems.

What We Offer:

We believe that to make intelligence open and accessible to all, you need to start at the foundation. Joining Reflection means building from the ground up as part of a talent-dense team. You will help define our future as a company, and help define the future of open foundational models.

We want you to do the most impactful work of your career with the confidence that you and the people you care about most are supported.

  • Top-tier compensation: Salary and equity structured to recognize and retain our talent globally.

  • Stock options: Everyone who joins and contributes to Reflection's success gets to share in the upside through stock options.

  • Health & wellness: Comprehensive medical, dental, vision, and life, with an annual wellness allowance.

  • Meals: Lunch and dinner are provided in the office daily.

  • Life & family: 22 weeks paid parental leave for all new birthing and non-birthing parents, including adoptive and surrogate journeys.

  • Vacation days: Unlimited paid time off in the U.S. and 30 days in the U.K.

  • Sponsorship support: We sponsor visas to help exceptional talent join our team and support long-term immigration pathways where applicable.

  • Team building: We have regular off-sites, happy hours, and team celebrations.

Export Control Notice: This position may require access to technology or source code subject to the U.S. Export Administration Regulations. Any offer of employment for this role may be conditioned on the Company's ability to provide the candidate with access to such technology or source code in compliance with applicable U.S. export control laws, which may require the Company to seek government authorization.

Prepare for this job

A free preview built only from this posting: what it asks for, what you could be asked in an interview, and how to adjust your resume.

Skills and AI tools this role asks for

PyTorchJax

Questions you could be asked

  1. What's a project where you used PyTorch hands-on?
  2. Walk me through how you've used Jax in your day-to-day work.
  3. How would you decide a model or AI system is ready to ship?
  4. Tell me about a time a model underperformed in production. How did you find out, and what did you change?

Adapt your resume

  • List these exact terms on your resume: PyTorch and Jax. An applicant tracking system matches the wording, not the idea.
  • Attach one line of real, concrete experience to at least one of them — a tool named with nothing behind it rarely survives a human read.
  • Lead with what you built, trained or shipped — this role is judged on the AI system itself, not the tools around it.

Want your resume actually rewritten for this job?

The free preview above is everything we have today. A full resume rewrite is not live yet and has no price set. Join the waitlist and we will email you if we open it.

Similar roles

Software Engineering roles rated AI Level 4 at other companies.

More jobs at Reflection AI

Related searches

Same AI level