Staff ML Engineer – AWS Trainium & SageMaker
Jobgether · United States · 1 mo ago
RemoteRemoteEngineeringFull-time
We need to output HTML fragment with sections: headings .
We need to ensure we don't miss any. Let's copy them exactly.
Accountabilities list:
1. Train and operate machine learning models using Amazon SageMaker with AWS Trainium as the underlying compute infrastructure.
2. Develop and optimize PyTorch training workloads for execution on AWS Trainium.
3. Analyze how PyTorch code compiles and executes across Trainium and NeuronCore architecture.
4. Optimize memory utilization, throughput, and other performance characteristics of training workloads.
5. Diagnose training failures and performance issues caused by hardware, compiler behavior, device configuration, data, or model code.
6. Distinguish model- and data-level problems from accelerator- and compiler-level issues during troubleshooting.
7. Translate high-level requirements for Trainium workloads into complete, production-ready training pipelines.
8. Design training workflows that balance performance, scalability, reliability, and cloud infrastructure costs.
9. Tune distributed and multi-device training workloads for throughput and cost efficiency.
10. Operate production training workloads and help ensure their reliability throughout the ML lifecycle.
11. Work directly with client engineering teams to scope, design, and deliver specialized machine learning workloads.
12. Collaborate with internal engineering teams to develop production-grade solutions rather than isolated prototypes or notebook experiments.
13. Investigate complex technical issues across the ML software and hardware stack.
14. Apply a deep understanding of accelerator behavior to improve training architecture and implementation decisions.
15. Contribute to production engineering practices around deployment, monitoring, troubleshooting, and operational reliability.
16. Help translate emerging AI infrastructure capabilities into practical production solutions.
17. Communicate technical findings and tradeoffs clearly with both technical stakeholders and client teams.
Check if any missing: yes.
Now Requirements list:
1. Strong hands-on experience with PyTorch, including a deep understanding of training workflows.
2. Experience with distributed or multi-device model training is highly desirable.
3. Production experience using Amazon SageMaker for model training and/or inference.
4. Strong understanding of how machine learning workloads interact with accelerator hardware and device-specific compilation.
5. Ability to debug issues that originate at the hardware, accelerator, compiler, or runtime layer rather than solely within model or data code.
6. Experience optimizing training workloads for performance, throughput, memory utilization, or cost.
7. Strong
for sections like About the role, Responsibilities, Requirements, Qualifications, Skills, Benefits, Pay, Schedule. Only include sections that have real specific content. No job title heading. No metadata like Job ID. Remove boilerplate: "This position is listed on behalf of a partner company..." etc. Also remove "How Jobgether Works", "Why Apply Through Jobgether?", data privacy notice, etc. Also remove application instructions. Keep substantive info: description of role, accountabilities (responsibilities), requirements, benefits. Possibly also "Qualifications" maybe same as requirements; but we can combine. We need to decide sections: About the role (summary), Responsibilities (accountabilities), Requirements (skills/experience), Benefits. Possibly also Skills separate but we can embed in Requirements. Pay and Schedule not present (no specifics). So omit those sections. We must preserve all substantive information; do not invent or summarize away detail. So we need to include all details from description, accountabilities, requirements, benefits. Must not omit any detail. Let's extract. Original text: "This position is listed on behalf of a partner company, who manages all applications and next steps. Our partner is looking for a Staff ML Engineer – AWS Trainium & SageMaker based in United States. As a Staff ML Engineer, you will design, train, optimize, and operate production machine learning workloads on AWS Trainium and Amazon SageMaker. You will work deeply across the ML stack, from PyTorch training code and distributed execution to accelerator hardware and compiler behavior. The role requires understanding how training workloads behave on custom silicon rather than treating infrastructure as a black box. You will diagnose complex training issues, optimize throughput and cost, and build reliable pipelines for real production workloads. You will work directly with engineering teams to translate business and technical requirements into scalable training solutions. The environment is highly hands-on, technical, and client-facing, with a focus on solving specialized problems that require deep engineering expertise. This is an opportunity to work at the intersection of machine learning, cloud infrastructure, distributed systems, and custom AI acceleration." Then "Accountabilities" list: - Train and operate machine learning models using Amazon SageMaker with AWS Trainium as the underlying compute infrastructure. - Develop and optimize PyTorch training workloads for execution on AWS Trainium. - Analyze how PyTorch code compiles and executes across Trainium and NeuronCore architecture. - Optimize memory utilization, throughput, and other performance characteristics of training workloads. - Diagnose training failures and performance issues caused by hardware, compiler behavior, device configuration, data, or model code. - Distinguish model- and data-level problems from accelerator- and compiler-level issues during troubleshooting. - Translate high-level requirements for Trainium workloads into complete, production-ready training pipelines. - Design training workflows that balance performance, scalability, reliability, and cloud infrastructure costs. - Tune distributed and multi-device training workloads for throughput and cost efficiency. - Operate production training workloads and help ensure their reliability throughout the ML lifecycle. - Work directly with client engineering teams to scope, design, and deliver specialized machine learning workloads. - Collaborate with internal engineering teams to develop production-grade solutions rather than isolated prototypes or notebook experiments. - Investigate complex technical issues across the ML software and hardware stack. - Apply a deep understanding of accelerator behavior to improve training architecture and implementation decisions. - Contribute to production engineering practices around deployment, monitoring, troubleshooting, and operational reliability. - Help translate emerging AI infrastructure capabilities into practical production solutions. - Communicate technical findings and tradeoffs clearly with both technical stakeholders and client teams. Then "Requirements" list: - Strong hands-on experience with PyTorch, including a deep understanding of training workflows. - Experience with distributed or multi-device model training is highly desirable. - Production experience using Amazon SageMaker for model training and/or inference. - Strong understanding of how machine learning workloads interact with accelerator hardware and device-specific compilation. - Ability to debug issues that originate at the hardware, accelerator, compiler, or runtime layer rather than solely within model or data code. - Experience optimizing training workloads for performance, throughput, memory utilization, or cost. - Strong Python programming fundamentals. - Understanding of distributed training concepts and production ML infrastructure. - Ability to design and operate end-to-end managed training pipelines. - AWS experience and familiarity with cloud-based machine learning infrastructure. - AWS Trainium or AWS Inferentia experience and familiarity with the AWS Neuron SDK is strongly preferred. - Candidates without direct Trainium or Inferentia experience may also be considered if they have deep PyTorch expertise and demonstrated ability to quickly learn new hardware targets. - Strong analytical and problem-solving skills, particularly when diagnosing complex system-level issues. - Ability to reason across multiple layers of the technology stack, from model code through frameworks, compilers, accelerators, and cloud infrastructure. - Experience working in a production engineering environment with high standards for reliability and delivery. - Strong communication and collaboration skills for working directly with client and internal engineering teams. - Comfortable operating in ambiguous environments where requirements and technical challenges may evolve. - Willingness and ability to travel approximately 20% within the United States. - Must be legally authorized to work in the United States. Then "Benefits" list: - Opportunity to work on production AI systems using AWS Trainium and Amazon SageMaker. - Exposure to custom AI accelerator hardware, compiler behavior, distributed training, and advanced ML infrastructure. - Hands-on work with complex machine learning workloads for enterprise clients. - Direct collaboration with client and internal engineering teams. - Opportunity to solve specialized technical problems that extend beyond conventional ML application development. - Production-focused engineering environment emphasizing systems that ship and operate at scale. - Approximately 20% U.S.-based travel associated with the role. - Final interview and onboarding may require onsite participation. - Professional growth through work across machine learning, cloud infrastructure, hardware acceleration, and production engineering. We need to remove "Final interview and onboarding may require onsite participation." That's maybe part of benefits but it's still substantive? It's about process; but it's still info. However it's not a benefit per se but it's still part of posting. The instruction: "Preserve ALL substantive information; do NOT invent or summarize away detail." So we should keep it. However it's not a benefit but it's still info; we could place it under a separate section like "Process" but we only have allowed headings: About the role, Responsibilities, Requirements, Qualifications, Skills, Benefits, Pay, Schedule. We could put it under Benefits as it's still a detail about the role. Or we could create a separate section like "Additional Information" but not allowed. Probably best to include it under Benefits as it's still a detail. However the instruction: "include a section ONLY when the posting has real, specific content for it — never emit a heading followed by filler/placeholder text". So we can include Benefits heading and list all benefit items, including that line. We also need to include "About the role" summary from first paragraph(s). That includes the description before Accountabilities. We'll put that in
under About the role. Responsibilities: we can use the Accountabilities list as responsibilities. Requirements: list as given. Benefits: list as given. We must not repeat job title as heading. So no heading for Staff ML Engineer. We must not include Job ID etc. We must output only HTML fragment, no /, no markdown. We'll produce something like:
About the role
...
Responsibilities
- ... ...
Requirements
-
...
Benefits
-
...