Nvidia is hiring a
Senior Software Engineer, Frameworks
We are hiring Senior engineers to passionately continue developing our AI platforms, with a heavy emphasis on large language models (LLMs) such as PAXML. Work with a dedicated group of engineers in a fast-paced production environment to make a tangible impact to Deep Learning Frameworks, particularly JAX. Come join us in these exciting times and make a sizable difference in the exploding world of AI! As an all-stack engineer with a focus on the back-end libraries and compilers, you will play a meaningful role in NVIDIA's efforts in contributing to the Deep Learning ecosystem.
What You Will Be Doing:
Collaborate on design and development of Deep Learning platform solutions with our Google partners
Design and build core Frameworks and Compilers and drive peak performance on NVIDIA GPUs
Evaluate and improve the end-to-end performance of generative AI models, especially large-language models (e.g. PAXML), on NVIDIA hardware
Work with AI applied researchers and leaders to build highly-efficient ML models
Build tools that will increase the efficiency of teams developing AI-based systems.
Work to bridge the gap between the latest in numerical computing, simulation and deep learning research and their applications in real world products.
Impact the design of next generation of accelerated computing software and hardware
What We Need To See:
BS in Computer Science or Computer Engineering or related field (or equivalent experience) with 8+ years relevant experience
Experienced in C++, Python, and MLIR
Experience with at least one machine learning framework, e.g. JAX, PyTorch, TensorFlow
Familiar with such concepts as asynchronous computing, sharding, collectives, various flavors of parallelism, GPU kernels, ML building blocks (e.g. Multi-Head Attention, MLP)
Proven technical foundation in CPU or GPU architectures, numeric libraries, modular software design.
Highly motivated with excellent verbal and written communication skills.
Ability to work collaboratively with multi-functional teams, principles and architects. Coordinates effectively across organizational boundaries.
Ways To Stand Out From The Crowd:
Prior experience in MLIR, XLA, JAX, LLVM, and CUDA
Firm grasp of distribution computation such as data, model, ZeRO, and pipeline parallelism
Being intimately familiar with Google C++ Style Guide and having contributed to OSS
Collaborative and production-minded demeanor
You will also be eligible for equity and benefits.
Please mention that you found the job on ARVR OK. Thanks.