River AI Verified

AI Infrastructure, Model Serving

Software Engineer, Distributed Training

Name: RiseMe
Availability: InStock

Palo Alto, CAOnsiteFull Time$200,000–$420,000 /yrPosted 1 month ago

Is this role right for you?

Upload your resume and get a skill-by-skill breakdown — see exactly where you match, where you're close, and what to highlight. Not a mystery percentage.

Get a tailored resume highlighting what this role needs.

Required skills

distributed training
large GPU clusters
gradient computation
optimizer updates
rollout coordination
checkpoint recovery
Python
PyTorch
JAX
backpropagation

Nice to have

reinforcement-learning infrastructure
rollout generation
asynchronous training
tensor parallelism
pipeline parallelism
expert parallelism

Role summary

River AI is seeking exceptional systems engineers to build and optimize the distributed training engines for their AI platform. This role involves making fine-tuning and reinforcement learning fast, reliable, and numerically correct across large GPU clusters. You will own the execution of training workloads, including gradient computation, optimizer updates, and checkpoint recovery. The position requires strong proficiency in Python and deep learning frameworks like PyTorch or JAX, with a solid understanding of distributed systems, numerical methods, and GPU memory management. Experience with reinforcement learning infrastructure and various parallelism techniques is preferred. The role is based in Palo Alto, CA, with competitive compensation and visa sponsorship available.

At River AI, our mission is to create personal AI owned and shaped by each individual. To achieve this, we are rewriting the entire stack from scratch: personal hardware for local inference, bespoke training infrastructure, next-generation UIs, and frontier deep learning research.

Who we are

We are scientists, engineers, and builders from the industry's top tech companies and AI labs. We bring a proven track record of scaling consumer systems for hundreds of millions of users and architecting the pre-training infrastructure behind today's frontier models.

About the Role

We are looking for exceptional systems engineers to build the distributed training engines behind the River API. Your goal is to make fine-tuning and reinforcement learning fast, numerically correct, and reliable across large GPU clusters.

You will own the execution of training workloads, including gradient computation, optimizer updates, rollout coordination, and checkpoint recovery. Working closely with researchers and inference engineers, you will bring new learning methods into production and improve how efficiently models use compute.

What You’ll Do

Build and optimize distributed training for large dense and mixture-of-experts models, including low-rank adapter training.

Improve reinforcement-learning pipelines by coordinating sampling, reward computation, training updates, and weight transfer.

Optimize GPU memory use, parallelism, and communication to increase training throughput.

Implement reliable checkpointing, resumption, and worker recovery while preserving consistent training state.

Validate losses, gradients, and optimizer behavior, and diagnose numerical or distributed execution failures.

Partner with researchers to implement new algorithms and make them accessible through the River API.

Skills & Qualifications

Minimum Qualifications:

Bachelor’s degree in Computer Science, Computer Engineering, or equivalent practical industry experience.

Hands-on experience building or substantially improving distributed model-training systems.

Strong proficiency in Python and a modern deep-learning framework, such as PyTorch or JAX.

Solid understanding of backpropagation, optimizers, mixed-precision training, and GPU memory management.

Strong debugging skills across concurrent execution, collective communication, and distributed failure recovery.

A highly collaborative mindset and a bias for action to push boundaries across the stack.

Preferred Qualifications: (We encourage you to apply even if you don't meet all of these)

Experience with reinforcement-learning infrastructure, rollout generation, or asynchronous training.

Familiarity with tensor, pipeline, expert, or data parallelism and their performance tradeoffs.

Work on LoRA, mixture-of-experts training, distributed optimizers, or activation checkpointing.

Experience with NCCL, communication profiling, and overlapping computation with data transfers.

Proficiency in C++, Rust, or CUDA, with experience investigating performance below the framework layer.

Contributions to training frameworks or a track record of operating large training runs.

Logistics & Benefits

Location: Palo Alto, California.

Compensation: Depending on experience and skills the expected base pay is $200,000 - $420,000 USD per year, plus equity.

Benefits: Comprehensive health, dental, and vision insurance; unlimited PTO; and relocation assistance as needed.

Visa Sponsorship: We sponsor visas and are committed to supporting the process for the right candidate.

Ready to apply?

You'll be redirected to River AI's application page.