google / google/brax

Multi-Process/Multi-Controller setup

Open
#638 0 comments 0 reactions 0 assignees View on GitHub
Dominant language
Jupyter Notebook
Stars
3.2k
Forks
349
PR merge metrics
No merged PRs in 30d

Description

Jax has the ability to run in multi-controller mode: https://docs.jax.dev/en/latest/multi_process.html does brax support it?
I am scheduling training using torchx where one process is assigned a single gpu and processes need to coordinate work, is there any guide for how to use brax in this setup?

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.