RL Infrastructure Engineer

Collision; San Francisco (CA); On-site; USD 180,000 - 240,000; Full time

Job summary

Collision AI is building models and agents for cyber defense. You’ll develop infrastructure that trains them through reinforcement learning on multi-GPU clusters.

The role blends systems engineering with practical RL, ensuring training correctness while improving performance at scale. You will collaborate with researchers to diagnose issues across trainer and inference workloads, implement robust debugging workflows, and contribute to reproducible experimentation during long-running training

Qualifications

• Strong Python and PyTorch skills.
• Experience building distributed training or inference infrastructure.
• Practical understanding of RL in LLM contexts.
• Experience profiling GPU workloads and resolving ML system issues.
• Ability to investigate problems across research and infrastructure code.

Responsibilities

• Coordinate trainers, rollout workers, environment execution, and reward computation across GPU clusters.
• Improve distributed training with FSDP, tensor/pipeline parallelism, and NCCL.
• Develop monitoring and diagnostics for stale weights and drift during long runs.

Skills

Python

PyTorch

Distributed training

Reinforcement learning infrastructure

GPU profiling

Asynchronous/off-policy training

SGLang

vLLM

NCCL

CUDA/Triton

Tools

SGLang

vLLM

NCCL

CUDA

Triton

Job description

Collision AI is building models and agents for cyber defense. You’ll develop the infrastructure that trains them through reinforcement learning, helping researchers move from initial experiments to efficient, reliable runs across many GPUs.

The role combines systems engineering with a practical understanding of RL. Infrastructure problems can change learning behavior: delayed weight updates, differences between trainer and inference log probabilities, or off‑policy drift can resemble problems with the training method itself. You’ll work closely with researchers to understand these interactions and improve performance without compromising training correctness.

What you’ll work on

• Training and rollout systems. Coordinate trainers, rollout workers, environment execution, and reward computation across GPU clusters. Integrate SGLang for high‑throughput sampling, improve trainer‑to‑inference weight synchronization, and support asynchronous rollouts and off‑policy training.
• Performance at scale. Improve distributed training through FSDP, tensor and pipeline parallelism, and efficient NCCL communication. Profile bottlenecks across compute, memory, and networking, and make targeted runtime or CUDA/Triton kernel improvements where they meaningfully improve training efficiency.
• Training correctness and reliability. Develop monitoring and diagnostics that catch stale weights, trainer–inference log‑probability mismatches, and off‑policy drift during long runs. Build reproducible debugging workflows and recovery mechanisms, and use controlled comparisons to understand how infrastructure changes affect learning.

Required qualifications

• Strong Python and PyTorch skills, with experience building distributed training or inference infrastructure.
• Practical understanding of LLM reinforcement learning, including rollout policies, reward computation, and policy updates.
• Experience profiling GPU workloads and resolving performance, correctness, or reliability issues in complex ML systems.
• Ability to investigate problems across research and infrastructure code and validate changes on real training runs.
• Experience developing RL training pipelines, particularly asynchronous or off‑policy systems.
• Work on agent environments, sandboxed execution, reward models, or verification pipelines.
• Experience extending SGLang, vLLM, or comparable inference runtimes.
• Experience with FSDP, tensor or pipeline parallelism, NCCL, CUDA, or Triton.