视频生成模型 · 训练 Infra 工程师
Embedding VC · Remote · San Francisco Bay Area ·
COMPENSATION
€•• – ••k
Listed on the employer's job board (Ashby)
- Fully remote
- PyTorch + CUDA
- Hundreds of GPUs
- Video generation models
- Senior level
Join the team building self-developed video generation foundation models using DiT and Flow Matching. You will construct the training platform, transforming research code into stable results across hundreds of GPUs. This role involves building reusable infrastructure rather than just using existing tools.
What you'll do
- Build training infrastructure including job scheduling and monitoring
- Implement checkpoint resumption and data or weight pipelines
- Manage distributed training on clusters with hundreds of cards
- Optimize MFU using FSDP, tensor parallelism, and Context Parallel
- Handle PB-level video data pipelines with NVDEC and VAE caching
- Improve memory and performance via FlashAttention and FP8 precision
- Ensure training stability through loss spike analysis and recovery
What you'll need
- Master PyTorch distributed systems and CUDA architecture
- Understand source code of Megatron, DeepSpeed, FSDP, or TorchTitan
- Have practical experience training on at least 256 GPUs
- Experience building training platforms or cluster scheduling tools
Nice to have
- Experience with DiT, Diffusion, or video data processing
- Written Triton or CUTLASS kernels
- Source-level knowledge of HunyuanVideo, Wan, or CogVideoX
Stack
- PyTorch
- CUDA
- FSDP
- Megatron
- DeepSpeed
- TorchTitan
- FlashAttention
- FP8
- Triton
- NVDEC
What you get
- Access to real computing power with hundreds of GPUs
- Team treats training as an engineering problem
- Space for open source contributions and publications