diff --git a/docs/blog/tensor-parallelism-optimization.md b/docs/blog/tensor-parallelism-optimization.md new file mode 100644 index 000000000..a1a88b408 --- /dev/null +++ b/docs/blog/tensor-parallelism-optimization.md @@ -0,0 +1,27 @@ +# Optimizing Tensor Parallelism in JAX and XLA + +> **Draft** — this page is a placeholder. Content is in progress. + +## Background + +{/* Why TP matters: model sizes past single-GPU memory, the compute/communication + tradeoff versus FSDP and pipeline parallelism. */} + +## Where tensor parallelism costs performance + +{/* All-gather / reduce-scatter on the critical path, exposed communication, + small GEMM tiles at high TP degree, NVLink versus inter-node bandwidth. */} + +## Optimizations + +{/* One subsection per optimization, with the mechanism and how to enable it. */} + +## Results + +{/* Model, hardware, TP degrees swept, throughput and step-time deltas. */} + +## Reproducing + +{/* Container tag, framework config flags, launch command. */} + +## Takeaways diff --git a/docs/index.yml b/docs/index.yml index ccfac3e37..ce319a1d5 100644 --- a/docs/index.yml +++ b/docs/index.yml @@ -55,6 +55,12 @@ navigation: - page: nsys-jax path: nsys-jax.md + # ==================== Tech Blogs ==================== + - section: Tech Blogs + contents: + - page: Optimizing Tensor Parallelism in JAX and XLA + path: blog/tensor-parallelism-optimization.md + # ==================== Resiliency ==================== - section: Resiliency contents: