Pollux: Unified World Model

Pollux is a unified world-model training and generation codebase. A collaboration between KAUST and Nucleus on world models.

Install

• Setup

Please see Install.md for a full installation.

• Train

torchrun --standalone --nnodes 1 --nproc-per-node 8 -m apps.main.train \
  config=apps/main/configs/train_bucket_256_latent_code.yaml

• Generate

python -m apps.main.generate config=apps/main/configs/eval.yaml

Docs

• MULTINODE.md - multinode support
• DEVELOP.md – development notes and cluster setup
• FA3.md – CUDA 12.8 + FlashAttention v3

Contributors