Description
NVIDIA is seeking a world-class Software Engineer with AI Compiler and distributed systems experience for its Deep Learning Compiler Technology team. The team builds JaxPP, NVIDIA's open-source library for pipeline parallelism in JAX.
Responsibilities:
- Design and implement program transformations on JAX programs (jaxpr) for large-scale training
- Develop technologies and products involving optimizing compilers and runtime execution frameworks for accelerating deep learning training on large scale distributed systems
- Work on compiling a graph intermediate representation (IR) to target specific code
- Contribute to advancing the state of the art in compiler technology for delivering high performance on large-scale training workloads
- Collaborate with members of the software framework teams and the hardware architecture teams to accelerate the next generation of deep learning software
- Be a core contributor to the open-source JaxPP project
Requirements:
- Masters or PhD in Computer Science, Computer Engineering, or related field or equivalent experience
- 6+ years of relevant work or research experience in compilers, distributed systems, or deep learning training
- Compiler design and implementation
- Distributed system design, development, and deployment
- Experience with deep learning frameworks (JAX, PyTorch)
- GPU training and inference of deep learning models
- C/C++ and Python programming, with excellent software design skills
- Experience with parallel computing libraries (MPI)
- Strong communication skills
Nice to Have:
- Experience with JAX internals (jaxpr, tracing, sharding, shard_map), or with XLA, StableHLO, or MLIR
- Experience building or tuning large training jobs with pipeline, tensor, expert, and data parallelism
- Knowledge of GPU communication libraries such as NCCL, NVSHMEM, or RDMA, and of CPU and/or GPU architecture
- Contributions to open-source compiler or machine learning projects
Benefits:
- Competitive salaries
- Generous benefits package
- Equity
The expected base salary range for this position is 184,000 USD – 287,500 USD.