软件工程师，Trainium
Software Engineer, Trainium
关于团队
在 OpenAI，我们的使命是发现并实现通往安全、有益的人工通用智能（AGI）的路径。为此，我们认为在生成建模、强化学习、大规模优化、主动学习和其他领域需要许多技术突破。
该团队构建了性能关键的系统，使 OpenAI 的模型能够在多种 AI 加速器上高效运行。我们跨推理堆栈工作，从底层内核和编译器到模型执行，以释放底层硬件的全部功能。
关于职位
作为 Software Engineer, Trainium，您将帮助将 OpenAI 的推理工作负载带到 AWS Trainium，并构建在该平台上高效运行前沿模型所需的软件堆栈。
这是一个深入技术、跨堆栈的角色，涵盖内核、编译器和模型执行。您将负责支持 OpenAI 在 Trainium 上的推理堆栈的系统，包括开发和优化高性能内核，改进编译器支持，并实现模型前向传递的高效执行。
您将与推理、编译器、内核和 ML 系统的工程师紧密合作，识别性能瓶颈，并构建充分利用 Trainium 的软件。工作内容可能从底层硬件特定优化，到编译器和运行时改进，再到将新的模型架构集成到推理堆栈中。
如果您喜欢在 ML 系统、编译器、内核和加速器硬件的交叉点上工作，这个职位适合您。我们寻找那些自我驱动、能在抽象层之间自由操作，并热衷于为前沿规模的 AI 系统解决挑战性性能问题的工程师。
在此职位中，您将
- 为 AWS Trainium 构建和优化 OpenAI 的推理堆栈。
- 为关键模型操作和工作负载开发高性能内核。
- 扩展和改进编译器支持，以高效地针对 Trainium 硬件。
- 构建在 Trainium 上执行和优化模型前向传递所需的系统。
- 分析工作负载，识别内核、编译器生成代码、运行时和模型执行中的瓶颈。
- 与推理和 ML 系统团队合作，将新模型和架构带入 Trainium。
- 跨越硬件/软件边界，从专用 AI 加速器中解锁性能和功能。
- 负责复杂的性能分析
查看英文原文
About the Team
Our mission at OpenAI is to discover and enact the path to safe, beneficial AGI. To do this, we believe that many technical breakthroughs are needed in generative modeling, reinforcement learning, large-scale optimization, active learning, and other areas.
The team builds the performance-critical systems that allow OpenAI's models to run efficiently across a diverse set of AI accelerators. We work across the inference stack, from low-level kernels and compilers through model execution, to unlock the full capabilities of the underlying hardware.
About the Role
As a Software Engineer, Trainium, you will help bring OpenAI's inference workloads to AWS Trainium and build the software stack required to run cutting-edge frontier models efficiently on the platform.
This is a deeply technical, cross-stack role spanning kernels, compilers, and model execution. You will work on the systems needed to support OpenAI's inference stack on Trainium, including developing and optimizing high-performance kernels, improving compiler support, and enabling efficient execution of the model forward pass.
You'll work closely with engineers across inference, compilers, kernels, and ML systems to identify performance bottlenecks and build the software needed to take full advantage of Trainium. The work may range from low-level hardware-specific optimization to compiler and runtime improvements to integrating new model architectures into the inference stack.
If you enjoy working at the intersection of ML systems, compilers, kernels, and accelerator hardware, this role is for you. We're looking for engineers who are self-directed, comfortable operating across abstraction layers, and excited to solve challenging performance problems for frontier-scale AI systems.
In This Role, You Will
- Build and optimize OpenAI's inference stack for AWS Trainium.
- Develop high-performance kernels for critical model operations and workloads.
- Extend and improve compiler support to efficiently target Trainium hardware.
- Build the systems necessary to execute and optimize the model forward pass on Trainium.
- Profile workloads and identify bottlenecks across kernels, compiler-generated code, runtime, and model execution.
- Partner with inference and ML systems teams to bring new models and architectures onto Trainium.
- Work across the hardware/software boundary to unlock performance and capabilities from specialized AI accelerators.
- Own complex performance and systems problems end-to-end, from investigation through production deployment.
We're Looking for a Track Record Of
- 3+ years of relevant engineering experience, ideally in ML systems, compilers, kernels, runtimes, or performance engineering.
- Strong systems programming fundamentals and experience writing performance-critical software.
- Experience working with GPU, TPU, Trainium, or other specialized accelerator architectures.
- Ability to reason about performance across multiple layers of the stack, from hardware and kernels through compilers and ML frameworks.
- Owning technically ambiguous problems end-to-end and learning new hardware and software domains as needed.
- Bonus: experience with AWS Trainium or the AWS Neuron SDK.
- Bonus: contributions to ML frameworks such as PyTorch or JAX, compiler infrastructure such as LLVM, MLIR, XLA, or Triton, or experience developing kernels for specialized accelerators.
About OpenAI
OpenAI is an AI research and deployment company dedicated to ensuring that general-purpose artificial intelligence benefits all of humanity. We push the boundaries of the capabilities of AI systems and seek to safely deploy them to the world through our products. AI is an extremely powerful tool that must be created with safety and human needs at its core, and to achieve our mission, we must encompass and value the many different perspectives, voices, and experiences that form the full spectrum of humanity.
We are an equal opportunity employer, and we do not discriminate on the basis of race, religion, color, national origin, sex, sexual orientation, age, veteran status, disability, genetic information, or other applicable legally protected characteristic.
For additional information, please see OpenAI’s Affirmative Action and Equal Employment Opportunity Policy Statement https://cdn.openai.com/policies/eeo-policy-statement.pdf.
Background checks for applicants will be administered in accordance with applicable law, and qualified applicants with arrest or conviction records will be considered for employment consistent with those laws, including the San Francisco Fair Chance Ordinance, the Los Angeles County Fair Chance Ordinance for Employers, and the California Fair Chance Act, for US-based candidates. For unincorporated Los Angeles County workers: we reasonably believe that criminal history may have a direct, adverse and negative relationship with the following job duties, potentially resulting in the withdrawal of a conditional offer of employment: protect computer hardware entrusted to you from theft, loss or damage; return all computer hardware in your possession (including the data contained therein) upon termination of employment or end of assignment; and maintain the confidentiality of proprietary, confidential, and non-public information. In addition, job duties require access to secure and protected information technology systems and related data security obligations.
To notify OpenAI that you believe this job posting is non-compliant, please submit a report through this form https://form.asana.com/?d=57018692298241&k=5MqR40fZd7jlxVUh5J-UeA. No response will be provided to inquiries unrelated to job posting compliance.
We are committed to providing reasonable accommodations to applicants with disabilities, and requests can be made via this link https://form.asana.com/?k=bQ7w9h3iexRlicUdWRiwvg&d=57018692298241.
OpenAI Global Applicant Privacy Policy https://cdn.openai.com/policies/global-employee-and-contractor-privacy-policy.pdf
At OpenAI, we believe artificial intelligence has the potential to help people solve immense global challenges, and we want the upside of AI to be widely shared. Join us in shaping the future of technology.