[Pallas] Add pallas distributed computation tutorial

2025-04-15 19:36:06 +00:00 · 2024-08-05 17:53:50 -07:00 · 2024-08-05 17:53:50 -07:00 · dcd186f552
commit dcd186f552
parent 500da57e91
13 changed files with 3282 additions and 0 deletions
--- a/docs/_static/pallas/distributed/all_gather.svg
+++ b/docs/_static/pallas/distributed/all_gather.svg
--- a/docs/_static/pallas/distributed/race_condition.svg
+++ b/docs/_static/pallas/distributed/race_condition.svg
--- a/docs/_static/pallas/distributed/rdma_recv.svg
+++ b/docs/_static/pallas/distributed/rdma_recv.svg
--- a/docs/_static/pallas/distributed/rdma_send.svg
+++ b/docs/_static/pallas/distributed/rdma_send.svg
--- a/docs/_static/pallas/distributed/rdma_start.svg
+++ b/docs/_static/pallas/distributed/rdma_start.svg
--- a/docs/_static/pallas/distributed/reduce_scatter_1.svg
+++ b/docs/_static/pallas/distributed/reduce_scatter_1.svg
--- a/docs/_static/pallas/distributed/reduce_scatter_2.svg
+++ b/docs/_static/pallas/distributed/reduce_scatter_2.svg
--- a/docs/_static/pallas/distributed/reduce_sum_1.svg
+++ b/docs/_static/pallas/distributed/reduce_sum_1.svg
--- a/docs/_static/pallas/distributed/reduce_sum_2.svg
+++ b/docs/_static/pallas/distributed/reduce_sum_2.svg
--- a/docs/conf.py
+++ b/docs/conf.py
@ -125,6 +125,7 @@ exclude_patterns = [
    'notebooks/*.md',
    'pallas/quickstart.md',
    'pallas/tpu/pipelining.md',
+    'pallas/tpu/distributed.md',
    'jep/9407-type-promotion.md',
    'autodidax.md',
    'sharded-computation.md',
@ -212,6 +213,7 @@ nb_execution_excludepatterns = [
    # Requires accelerators
    'pallas/quickstart.*',
    'pallas/tpu/pipelining.*',
+    'pallas/tpu/distributed.*',
    'sharded-computation.*'
 ]

--- a/docs/pallas/tpu/distributed.ipynb
+++ b/docs/pallas/tpu/distributed.ipynb
--- a/docs/pallas/tpu/distributed.md
+++ b/docs/pallas/tpu/distributed.md
--- a/docs/pallas/tpu/index.rst
+++ b/docs/pallas/tpu/index.rst
@ -8,3 +8,4 @@ TPU specific documentation.

   details
   pipelining
+   distributed