Skip to content

Optimize CUDA Breakout rollout and training#610

Open
michaelthompsonx-lab wants to merge 1 commit into
PufferAI:5cfrom
michaelthompsonx-lab:fbr/cuda-breakout-rollout-optimizations
Open

Optimize CUDA Breakout rollout and training#610
michaelthompsonx-lab wants to merge 1 commit into
PufferAI:5cfrom
michaelthompsonx-lab:fbr/cuda-breakout-rollout-optimizations

Conversation

@michaelthompsonx-lab

Copy link
Copy Markdown

Summary

  • move Breakout physics, reset, reward, terminal, and observation generation onto CUDA
  • write rollout observations directly into BF16 storage and consume sampled precision actions without an intermediate action cast
  • compact brick state into a four-word bitset and use warp-cooperative observation writes
  • capture a complete rollout horizon per buffer and tune Breakout to two buffers
  • remove redundant saved-activation copies and accumulate the MinGRU highway gradient through GEMM beta=1
  • keep MinGRU backward checkpoints in shared memory, fuse Muon slice normalization, and parallelize PPO advantage moments
  • add deterministic CPU/CUDA Breakout parity coverage

Runtime changes are marked with fbr: comments for review. CUDA 13/NCCL build compatibility comments are intentionally untagged.

Verification

  • ./build.sh breakout --cuda-env-test && ./test_cuda_env (BF16): PASS, 65 envs x 4096 steps, 1658 completed episodes
  • ./build.sh breakout --cuda-env-test --float && ./test_cuda_env (FP32): PASS
  • ./build.sh breakout: PASS
  • ./build.sh breakout --cpu-env: PASS
  • ./build.sh breakout --profile: PASS
  • ./puffer train breakout: reaches score 864; approximately 20-22M SPS on an RTX 3080 Ti depending on GPU clock state
  • git diff --check: PASS

Notes

Experimental direct/parallel MinGRU scans, specialized small-Muon kernels, and larger-minibatch variants were profiled but intentionally excluded because they did not improve verified end-to-end training.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant