Job Description

Lightning AI is seeking a skilled AI Optimization Engineer to optimize training and inference workloads using the Lightning Thunder compiler and PyTorch Lightning ecosystem. This role involves deep learning research, compiler development, and large-scale system optimization to improve model performance and efficiency.

Responsibilities

• Develop performance-oriented model optimizations at graph, kernel, and system levels
• Advance the Thunder compiler with optimization passes and graph transformations
• Ensure optimizations are accessible through APIs and tooling
• Design and implement profiling and debugging tools
• Collaborate with hardware vendors and ecosystem partners
• Contribute to open-source projects
• Engage with the community on performance tuning
• Work cross-functionally with product and engineering teams

Requirements

• Strong expertise with deep learning frameworks (PyTorch, JAX, TensorFlow)
• Hands-on experience with model optimization techniques
• Deep understanding of compiler internals
• Experience with CUDA, Triton, or GPU programming models
• Knowledge of distributed systems and parallelism
• Familiarity with software engineering practices
• Proven track record contributing to open-source projects
• Excellent collaboration and communication skills
• Bachelor’s degree in Computer Science, Engineering, or related field

Benefits

• No benefits