← All roles

AI Infra / Engineer - Large-scale Training

UMA· UMA
Apply on company site
Paris, FranceFull-timePosted 4d ago

About the role

AI Infra / Engineer focusing on large-scale training infrastructure for UMA’s humanoid robot foundation models, optimizing distributed training across GPUs/nodes, and turning experimental code into robust, scalable pipelines.

Responsibilities

  • Build distributed pretraining systems across multiple GPUs and nodes, selecting strategies such as DDP, FSDP, and tensor, pipeline, or expert parallelism as workloads require.
  • Develop RL and post-training execution infrastructure, including distributed rollouts, actor–learner coordination, policy synchronization, replay buffers, and efficient batching.
  • Maximize MFU and eliminate bottlenecks: profile the complete training loop—data loading, decoding, CPU-GPU memory transfers, and compute kernels—eliminating bottlenecks to keep MFU high and ensure accelerators operate at peak efficiency.
  • Integrate and benchmark accelerated kernels and execution techniques, such as FlashAttention, fused operations, mixed precision, and compilation; develop custom optimizations where they provide clear value.
  • Build efficient training-time data movement and buffering, including out-of-core storage, caching, prefetching, and asynchronous transfers, in partnership with the Data Infrastructure team.
  • Turn research prototypes into maintainable training systems, with reusable components that support evolving architectures and experimental workflows.
  • Diagnose distributed and numerical failures, including deadlocks, memory leaks, gradient anomalies, precision issues, and inconsistencies between single-device and distributed runs.
  • Deep numerical debugging and run determinism: diagnose complex issues across the stack while enforcing numerical determinism for reproducibility of training runs, evaluations, and experiments.
  • Build reliable checkpointing and fault-tolerant resumption, preserving model, optimizer, and relevant execution state so interrupted jobs can recover correctly.
  • Strengthen observability and reproducibility through configuration tracking, profiling, run diagnostics, and deterministic execution where practical.

Qualifications

  • Solid experience building ML training infrastructure, distributed systems, or high-performance ML software.
  • Strong Python and PyTorch, with hands-on knowledge of distributed training tools such as DDP, FSDP, or Megatron.
  • Experience with multiprocessing, asynchronous execution, and GPU memory management, and an understanding of how communication and synchronization affect scaling.
  • Practical knowledge of losses, optimizers, gradients, and numerical precision, sufficient to debug training behavior alongside researchers.
  • Experience profiling accelerator workloads with tools such as PyTorch Profiler or Nsight Systems, and translating findings into measurable improvements.
  • A pragmatic engineering mindset: turn experimental code into reliable, scalable systems while keeping research iteration fast.
  • Clear communication and a collaborative approach to working with researchers.
  • Bonus: RL infrastructure, distributed rollouts, replay buffers, or actor–learner architectures.
  • Bonus: C++, CUDA, Triton, custom kernels, or compiler-based optimization.
  • Bonus: Slurm or Kubernetes, GPU networking (InfiniBand/RoCE), or storage caching.
  • Bonus: contributions to open-source training tools or technical writing.

Skills

Required:Distributed pretrainingPerformance profilingData movement and bufferingDebugging and determinism
Nice to have:RL and post-training infrastructureAccelerated kernels
Last updated Oct 22, 2026. Interested? Apply on the company site.

About UMA

UMA - Building intelligent robots that enhance quality of life for everyone

At UMA, we develop AI and robotics dedicated to shaping a new economic and societal era. We design general-purpose mobile and humanoid robots with human-level dexterity and a deep understanding of the physical world.

AIBuilding the futureDeep LearningGeneral-purpose robotsLeRobotRoboticsheadline-ai-europe-100-2026intelligent-hardware
Company profileWebsite

More open roles at UMA