Software Engineer, Distributed Training

River AI is building personal AI systems through personal hardware, training infrastructure, next-generation user interfaces, and deep learning research. The Software Engineer, Distributed Training will build and optimize distributed training engines for large GPU clusters, including fine-tuning and reinforcement learning, while ensuring numerical correctness, reliability, and efficient compute utilization.

Responsibilities:

• 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