Distributed Training & Inference Optimization Engineer (LLM) - GPU Optimization Department (GPUOD)

GPU Optimization Department (GPUOD) is responsible for the strategic management, optimization, and governance of Rakuten's company-wide AI infrastructure, ensuring high-performance, cost-efficient utilization of compute resources for machine learning workloads.

We oversee a large-scale hybrid infrastructure spanning thousands of accelerators, including the latest Hopper and upcoming Blackwell architectures.

Optimize compute resource allocation across on-premises and multi-cloud environments, maximizing efficiency for training and inference workloads.

Manage hybrid orchestration of diverse accelerator resources, ensuring seamless scalability and cost-effective deployment.

Develop and enhance frameworks for large-scale distributed training, with special focus on LLMs and generative AI.

Optimize inference performance through model optimization techniques and system-level acceleration.

Collaborate with internal teams to deliver scalable, high-availability inference services tailored to business needs.

Continuously evaluate next-generation hardware solutions, including specialized AI chips optimized for LLM workloads.

By effectively managing both conventional and specialized compute resources across on-premises and cloud environments, our team ensures Rakuten's AI ecosystem remains at the forefront of performance, reliability, and cost-efficiency.

Work on cutting-edge LLM training & inference optimization at scale.

Directly impact Rakuten’s AI infrastructure by improving efficiency and reducing costs.

Collaborate with global AI/ML teams on high-impact challenges.

Opportunity to research and implement state-of-the-art GPU optimizations.

As a GPU Training & Inference Optimization Engineer, you will focus on maximizing the performance, efficiency, and scalability of LLM training and inference workloads on Rakuten’s GPU clusters.

You will deeply optimize training frameworks (e.g., PyTorch, DeepSpeed, FSDP) and inference engines (e.g., vLLM, TensorRT-LLM, Triton, SGLang), ensuring Rakuten’s AI models run at peak efficiency.

This role requires strong expertise in GPU-accelerated ML frameworks, distributed training, and inference optimization, with a focus on reducing training time, improving GPU utilization, and minimizing inference latency.

Optimize LLM training frameworks (e.g., PyTorch, DeepSpeed, Megatron-LM, FSDP) to maximize GPU utilization and reduce training time.

Profile and optimize distributed training bottlenecks (e.g., NCCL issues, CUDA kernel efficiency, communication overhead).

Implement and tune inference optimizations (e.g., quantization, dynamic batching, KV caching) for low-latency, high-throughput LLM serving (vLLM, TensorRT-LLM, Triton, SGLang).

Collaborate with infrastructure teams to improve GPU cluster scheduling, resource allocation, and fault tolerance for large-scale training jobs.

Develop benchmarking tools to measure and improve training throughput, memory efficiency, and inference latency.

Research and apply cutting-edge techniques (e.g., mixture-of-experts, speculative decoding) to optimize LLM performance.