chore: import upstream snapshot with attribution

third_party/xla/docs/oom_debugging.md

@@ -0,0 +1,88 @@
+# Debug OOM errors with XProf
+
+Out Of Memory (OOM) errors occur when the accelerator's (GPU or TPU) High
+Bandwidth Memory (HBM) capacity is exhausted. Some common causes for OOM issues
+and debugging techniques are detailed in
+[E1000 - Compile Time HBM OOM documentation](./errors/error_1000.md) and
+[JAX documentation on GPU memory allocation](https://docs.jax.dev/en/latest/gpu_memory_allocation.html#common-causes-of-oom-failures).
+
+This page describes how to use **XProf's Memory Viewer tool** to visualize your
+JAX program's memory usage, identify peak usage instances, and debug OOM errors.
+This involves the following steps:
+
+1.  Run your program with
+    [`jax.profiler.trace`](https://docs.jax.dev/en/latest/_autosummary/jax.profiler.trace.html#jax.profiler.trace)
+    to capture the profile.
+2.  Start XProf in the background, and use the
+    [Memory Viewer tool](https://openxla.org/xprof/memory_viewer) to view memory
+    utilization details.
+
+## Example program
+
+The following JAX program leads to an OOM error:
+
+```python
+import jax
+from jax import random
+import jax.numpy as jnp
+
+
+@jax.profiler.trace("/tmp/xprof")
+@jax.jit
+def oom():
+    a = random.normal(random.PRNGKey(1), (327680, 327680), dtype=jnp.bfloat16)
+    return a @ a
+
+
+if __name__ == "__main__":
+    oom()
+```
+
+**Note:** Prefer `jax.profiler.trace` instead of
+`jax.profiler.start_trace`/`jax.profiler.stop_trace` because the
+`jax.profiler.trace` context manager handles profiling in an exception safe
+manner.
+
+On a TPU-machine, this program fails with:
+
+```shell
+XlaRuntimeError: RESOURCE_EXHAUSTED: Allocation (size=107374182400) would exceed memory (size=17179869184) :: #allocation7 [shape = 'u8[327680,327680]{1,0:T(8,128)(4,1)}', space=hbm, size = 0xffffffffffffffff, tag = 'output of xor_convert_fusion@{}'] :: <no-hlo-instruction>
+```
+
+(On a GPU-machine, the error looks like: `XlaRuntimeError: RESOURCE_EXHAUSTED:
+Out of memory while trying to allocate 214748364800 bytes.`)
+
+## Run XProf
+
+Install `xprof` (`pip install xprof`), and start an XProf instance specifying
+the directory where the profile is stored:
+
+```shell
+xprof --logdir=/tmp/xprof/ --port=6006
+```
+
+Go to the instance (on a local machine, at `http://localhost:6006`). In the
+*Tools* dropdown, select *Memory Viewer*, and in the Memory Viewer tool window,
+select *HBM* in the *Memory Types* dropdown (usually selected by default).
+
+![XProf Memory Viewer page for the above example program](images/oom_debugging_example_memory_viewer.png)
+
+The
+[XProf: Memory Viewer tool documentation](https://openxla.org/xprof/memory_viewer#memory_viewer_components)
+describes the components of the tool and the information presented.
+
+Focus on the *HLO Ops at Peak Memory Allocation* section that shows three buffer
+charts at the peak memory usage point. The buffer includes:
+
+-   **Program Inputs and Outputs:** Training batches, optimizer states, etc.
+-   **TensorCore and SparseCore Temporaries:** Dynamic memory required for
+    intermediate calculations (like activations, gradients, etc.)
+
+You can hover on the buffer charts to get more details about the Op like it's
+size, shape, allocation type, and more. This can help you identify and evaluate
+Ops that may have high or long-lasting temporaries, any large
+input/intermediate/output tensors that have inefficient padding, etc., that are
+contributing to the peak memory and need to be adjusted or optimized.
+
+Learn specific debugging techniques in
+[E1000: Debugging](https://openxla.org/xla/errors/error_1000#debugging).