wehub/tensorflow--tensorflow
1
0
Fork 0
Code Issues Pull Requests Actions 3 Deployments Models Agent Notes Packages Projects Releases Wiki Activity

chore: import upstream snapshot with attribution
cffconvert / validate (push) Has been skipped
Details
License Check / license-check (push) Failing after 2s
Details

Browse Source
This commit is contained in:
wehub-resource-sync
2026-07-13 12:14:16 +08:00
commit 8a852e4b4e
36502 changed files with 9277225 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
third_party/xla/docs/errors/error_0100.md
Vendored
+85
View File
@@ -0,0 +1,85 @@
# Error code: E0100
**Category:** Runtime: Buffer Allocation Failure
This error indicates that XLA:TPU runtime’s memory allocator failed to find a
suitable block of memory on the accelerator’s HBM for the requested allocation.
**Sample error message:**
```
ValueError: RESOURCE_EXHAUSTED: Error allocating device buffer: Attempting to allocate 8.00M. That was not possible. There are 6.43M free.; (0x0x1_HBM0)
```
**XLA backends:** TPU
## Overview
This error is thrown on:
- Failures of user-initiated buffer allocation via
[`jax.device_put`](https://docs.jax.dev/en/latest/_autosummary/jax.device_put.html)
or
- Failures of user-scheduled program's output allocations.
These failures are typically caused due to a couple of reasons:
- **Out of Memory (OOM):** The user is trying to allocate a chunk of memory
that is larger than the total amount of free memory available on the TPU’s
HBM.
- **Memory fragmentation:** The allocation fails because **no single
contiguous free block** in the memory space is large enough to satisfy the
requested size. The total amount of free memory is sufficient for the
allocation, but it is scattered across the memory space in small,
non-contiguous blocks.
The TPU runtime has a number of mechanisms in-place to retry allocation failures
including:
- If there are queued deallocations, the runtime retries failed allocations.
- On OOMs caused by a fragmentation, the runtime can automatically trigger a
defragmentation and a retry.
- The TPU runtime prioritizes buffer allocations over keeping programs loaded.
If a buffer allocation fails due to insufficient HBM, the system will evict
loaded TPU programs until enough memory is available for the buffer.
So an error encountered after the above mitigations typically require user
action.
## Debugging
- **Reduce your model's memory footprint**
- **Decrease batch size:** Reducing the batch size directly lowers memory
usage.
- **Parameter sharding:** For very large models, use techniques like model
parallelism or sharding to distribute parameters across the HBM of
multiple TPU cores or hosts.
- **Shorten sequence/context length:** For models that operate on
sequences (like language models), reducing the input sequence length can
significantly decrease the memory footprint.
- **Buffer donation:** Utilize framework features (such as: `jax.jit(...,
donate_argnums=...)`) to signal to XLA that certain input buffers can be
overwritten and reused for outputs. Read
[Buffer donation](https://docs.jax.dev/en/latest/buffer_donation.html)
for more details.
- **Optimize checkpoint strategy:** Instead of saving the entire model
state at once, consider saving only the model weights or using a sharded
checkpointing strategy.
- **Address memory layout and padding**
- TPU memory is allocated in chunks, and padding can increase the actual
size of tensors.
- **Ensure no memory leaks**
- Ensure references to `jax.Array` objects are not being held longer than
intended. Holding on to `jax.Array` objects might prevent automatic
de-allocation even after program compilation is completed.
See also [Error code: E1000](https://openxla.org/xla/errors/error_1000) for
other strategies you can use to reduce the amount of memory each program uses.
### Tooling
Enable the `tpu_log_allocations_on_oom` flag for which the allocator will dump a
detailed report of all current allocations when an OOM occurs, which can be
invaluable for debugging.
third_party/xla/docs/errors/error_0101.md
Vendored
+75
View File
@@ -0,0 +1,75 @@
# Error code: E0101
**Category:** Runtime: Program Allocation Failure
This error indicates that the XLA runtime on a TPU device failed to load a
compiled XLA program executable into the TPU's HBM.
**Sample error message:**
```
XlaRuntimeError: RESOURCE_EXHAUSTED: Error loading program 'jit_embedding_pipeline_step_fn': Attempting to reserve 29.49G at the bottom of memory. That was not possible. There are 147.64M free, 0B reserved, and 147.64M reservable. Scope: unknown..: while running replica 0 and partition 34 of a replicated computation (other replicas may have failed as well).
```
**XLA backends:** TPU
## Overview
This error is typically caused by one of the following reasons:
- **Program size exceeds available HBM:** The compiled XLA program, including
its instructions, static data, and any embedded constants, is larger than
the total amount of free HBM currently available on the specific TPU core(s)
where the program is being loaded.
- **HBM fragmentation:** While the total free HBM on the device might be
sufficient in aggregate, it is not available in a single, contiguous block
large enough to fit the entire program.
It's important to understand how the TPU runtime prioritizes memory. Buffer
allocations are privileged over loaded programs. If a buffer allocation fails,
the runtime will evict already loaded programs from HBM to free up space. This
can lead to a situation where a program that loaded successfully before now
fails with an OOM error, because the HBM is now occupied with more data buffers.
## Debugging
- **Reduce buffer memory footprint:** Freeing up memory used by data buffers
will leave more room for the program itself.
- **Decrease batch size:** This is one of the most effective ways to
reduce the amount of memory used for activations.
- **Parameter sharding:** For very large models, use model parallelism or
sharding techniques (like FSDP or Megascale) to distribute the model's
parameters and computation across multiple TPU cores or hosts. Note that
this can increase the network communication overhead due to splitting
tensors across multiple chips.
- **Shorten sequence/context length:** For models processing sequential
data (e.g., NLP models), reducing the sequence length can significantly
decrease memory usage.
- **Buffer donation:** Use framework features (e.g., `jax.jit(...,
donate_argnums=...)`) to allow XLA to reuse the memory of input buffers
for storing output, reducing peak memory usage.
- **Reduce program’s memory requirements for temporaries**
- Reduce programs memory usage for temporaries by using the
`tpu_shared_memory_percent` flag. Note that this might negatively affect
performance.
- **Optimize execution strategy/reduce serving load**
- **Manage program loading:** If you are JIT-compiling multiple functions,
be aware that each function can result in a program being loaded. Try to
structure your workload to minimize the number of concurrently loaded
programs.
- **Ensure no memory leaks**
- Ensure references to `jax.Array` objects are not being held longer than
intended. Holding on to `jax.Array` objects might prevent automatic
de-allocation even after program compilation is completed.
See also [Error code: E1000](https://openxla.org/xla/errors/error_1000) for
other strategies you can use to reduce the amount of memory each program uses.
### Tooling
- Enable the `tpu_log_allocations_on_oom` flag for which the allocator will
dump a detailed report of all current allocations when an OOM occurs, which
can be invaluable for debugging.
- Profile Your Program: Use the JAX memory profiler or the TensorFlow profiler
to get a detailed view of your program's memory usage over time. This can
help identify unexpected peaks in memory consumption.
third_party/xla/docs/errors/error_0102.md
Vendored
+58
View File
@@ -0,0 +1,58 @@
# Error code: E0102
**Category:** Runtime: Program Input Mismatch
This error occurs when the XLA runtime detects that an input buffer provided at
execution time does not match the metadata expected by a compiled program, such
as the buffer size or PJRT memory space kind.
**Sample error message:**
```
XlaRuntimeError: INVALID_ARGUMENT: Executable(jit_embedding_pipeline_step_fn) expected parameter 2482 of size 5242880 (bf16[16,1280,40]{2,1,0:T(8,128)(2,1)}) but got buffer with incompatible size 1638400 (bf16[16,1280,40]{1,2,0:T(8,128)(2,1)}): while running replica 0 and partition 0 of a replicated computation (other replicas may have failed as well).
```
**XLA backends:** TPU
## Overview
For buffer-size mismatches, the error message indicates both the expected and
actual sizes, as well as the tensor shapes and layouts. Note that these errors
might occur even if two tensors have the same shape but their size in memory can
be different if their physical layout (how the data is tiled and arranged on the
hardware) is different. For memory-space mismatches, the error message reports
that the parameter buffer is in an unexpected memory space.
These errors are predominantly caused by:
- **Checkpoint and XLA configuration mismatch:** A model is trained and a
checkpoint is saved. The physical layout of the weights in that checkpoint
is determined by the exact XLA version and configuration (e.g. XLA flags) at
that time. Later, this checkpoint is loaded in a different environment where
the configuration has changed. A new flag, a different default value, or a
change in the model/XLA code can cause the runtime to expect a different
physical layout for the weights. When the old buffer from the checkpoint is
passed to the new compiled XLA program, the runtime throws an error.
- **Hardware/topology-specific layouts:** The XLA compiler is free to choose
different physical layouts for tensors to optimize performance on different
hardware. A layout that is optimal for v4 TPU might be different from a v5
TPU, or even for different pod slices of the same chip (e.g., 4x4x4 vs 4x8).
The error occurs when a model is compiled with an assumption about one
topology's layout, but at runtime it is scheduled on a different topology,
or there is a bug in the compiler's layout logic for a specific piece of
hardware.
## Debugging
- **Ensure configuration consistency between model export and re-runs from
checkpoints**
- Avoid using old checkpoints with new code unless you are certain that no
layout-affecting changes have been made.
- If you suspect a checkpoint/configuration mismatch, the most reliable
solution is to re-export the saved model using the exact same (and
current) codebase and configuration that you are using for inference or
fine-tuning.
- Check for configuration changes (e.g. XLA flags) between the two runs.
- **Hardware/topology-specific layouts**
- Check for hardware version and topology mismatches if switching hardware
or topologies.
third_party/xla/docs/errors/error_0200.md
Vendored
+158
View File
@@ -0,0 +1,158 @@
# Error Code: E0200
**Category:** Runtime: Core Halted Unexpectedly
This error indicates that a TPU core stopped executing
instructions prematurely. This is a fatal error state where the hardware forces
a halt due to an unrecoverable fault, a violation of hardware constraints, or a
deliberate interrupt triggered by compiler-generated runtime assertions.
**Sample error message:**
```
INTERNAL: Accelerator device halted prematurely, perhaps due to an on-device check-failure. Node 0 halted unexpectedly at tag:pc TensorCoreSequencer:1:0x1d9 ...
```
**XLA backends:** TPU
## Overview
XLA compiles JAX programs into a sequence of low-level assembly instructions.
At runtime, the TPU device executes these instructions sequentially. A
"Core Halted Unexpectedly" error occurs when the TPU hardware encounters an
unrecoverable condition that prevents further execution, forcing the core into a
fatal "HALTED" state.
Because this error can stem from physical hardware failures, compiler bugs, or
user code issues (particularly in custom kernels), you must carefully analyze
the log messages to identify the specific cause.
## Debugging
To resolve this error, you must first identify which of the three specific
scenarios caused the unexpected halt. Check your logs for the specific text
signatures described below.
- **"observed errors are: [Hardware/Network/Power]"**: This indicates a
physical infrastructure failure. → Jump to
[Scenario 1: Infrastructure failures (hardware/newtwork/power)](#scenario_1_infrastructure_failures_hardwarenetworkpower)
- **"observed errors are: [User]"**: This indicates a hardware constraint
violation. → Jump to
[Scenario 2: Hardware constraint violations](#scenario_2_hardware_constraint_violations)
- The error message contains specific details such as the following keywords:
`BoundsCheck`, `scheckne`, `scheckeq`, `schecklt`, `scheckge`,
`scheckbetween`
This indicates that a compiler-generated assertion in the compiled program
failed during execution. → Jump to
[Scenario 3: XLA compiler-generated assertion failures](#scenario_3_xla_compiler-generated_assertion_failures)
### Scenario 1: Infrastructure failures (hardware/network/power)
**Signature:** The logs explicitly state `observed errors are: [Hardware]` or
`observed errors are: [Network]` or `observed errors are: [Power]`.
This indicates a physical infrastructure failure unrelated to your software or
model logic. The TPU chip, the network fabric connecting the chips, or the
power supply has failed.
- **Retry the job:** If the issue was a transient voltage dip or network flap,
a simple retry may work.
- **Identify and remove bad nodes:** If the error persists on the same
specific task or host, the hardware is likely defective. Use your cluster
management tools to "drain"/"cordon" the affected node and restart your job
on healthy nodes.
### Scenario 2: Hardware constraint violations
**Signature:** The logs state `observed errors are: [User]`.
This indicates that the XLA compiler generated an instruction that violated a
inviolable hardware constraint (e.g., an instruction attempting to access an
out-of-bounds memory address on HBM or Scratchpad memory). While labeled "User",
this is rarely caused by high-level user code.
- **File an XLA bug:** This is likely a compiler bug, the compiler should
never emit instructions that violate hardware specs.
[Please file a bug report](https://github.com/openxla/xla).
### Scenario 3: XLA compiler-generated assertion failures
**Signature:** The error message contains specific details on the
compiler-generated assertion that is failing. Look for the for the following
keywords:
- `BoundsCheck`, `scheckne`, `scheckeq`, `schecklt`, `scheckge`,
`scheckbetween`
This indicates that a compiler-generated assertion in the compiled program
failed during execution. Analyze the specific error message to determine the
sub-type.
#### Scenario 3.A: Launch group mismatch
**Sample error message:**
```
Core halted unexpectedly: INTERNAL: Accelerator device halted prematurely, perhaps due to an on-device check-failure. Node 0 halted unexpectedly at tag:pc TensorCoreSequencer:1:0x1d9 (from TensorCoreSequencer:1:0x309): scheckne: An unexpected leader shows up in the launch group with a different launch id than the current group leader.
```
**Cause:** This error typically occurs in multi-host TPU environments. It
indicates that the TPU cores, which are expected to execute the same program in
a synchronized manner (as part of a "launch group") have become out of sync.
Specifically, a TPU core joined a synchronization group with a different program
identifier than the current group leader, suggesting inconsistent programs
across hosts.
- **Verify XLA flags:** Ensure all hosts use the exact same `XLA_FLAGS`.
- **Check for consistent JAX programs:** Check that all hosts are executing
identical JAX programs. Verify docker images, libtpu versions, etc.
#### Scenario 3.B: Bounds check failure
**Sample error message:**
```
Core halted unexpectedly: INTERNAL: Accelerator device halted prematurely, perhaps due to an on-device check-failure. Node 0 halted unexpectedly at tag:pc TensorCoreSequencer:23:0x292 (from TensorCoreSequencer:23:0xd74a): BoundsCheck 92 [deref of %s931] for %937 = dma.hbm_to_vmem [thread:$0] /*hbm=*/%s931, /*size_in_granules=*/16384, /*vmem=*/%s935, /*dst_syncflagno=*/%s860, /*src_stride=*/512, /*dst_stride=*/128, /*steps_per_stride=*/8
```
**Cause:** The program tried to access memory outside of allocated bounds. The
error message often includes details about the memory access type (e.g.,
`dma.hbm_to_vmem`) and the address calculation.
- **Debug custom kernels:** If using Pallas, check your index calculations.
Use
[`pl.debug_print`](https://docs.jax.dev/en/latest/_autosummary/jax.experimental.pallas.debug_print.html)
or
[`checkify`](https://docs.jax.dev/en/latest/_autosummary/jax.experimental.checkify.check.html)
to validate tensor indices.
- **Check sharding:** Ensure sharding annotations are consistent with tensor
shapes.
#### Scenario 3.C: Mosaic/Pallas synchronization
**Sample error message:**
```
Core halted unexpectedly: INTERNAL: Accelerator device halted prematurely, perhaps due to an on-device check-failure. Node 0 halted unexpectedly at tag:pc TensorCoreSequencer:21:0xae5 (from TensorCoreSequencer:21:0x54c5): Semaphore (scratch argument 1) has a nonzero value upon exit from a Mosaic kernel. Make sure every DMA is awaited, and every semaphore signal is paired with a wait.
```
**Cause:**
This error is specific to code generated by the Mosaic compiler (used by Pallas
JAX). It indicates a synchronization issue within a custom kernel. TPUs use
semaphores to manage dependencies (e.g., ensuring a DMA is complete before
use). This error suggests a signal on a semaphore was not properly waited upon.
- **Audit synchronization:** Ensure every `dma_start` has a corresponding
`dma_wait`.
- **Check semaphores:** Verify that semaphore signals and waits are strictly
paired.
### Uncategorized issues
If your error log does not match Scenario 1, 2, or 3 (i.e., no "observed
errors", no "scheck" tags, and no specific bounds/semaphore messages):
- **Action:** This is likely an internal XLA bug.
[Please file a bug report](https://github.com/openxla/xla).
third_party/xla/docs/errors/error_1000.md
Vendored
+279
View File
@@ -0,0 +1,279 @@
# Error code: E1000
**Category:** Compile Time: HBM OOM
This error indicates that the program requires more High Bandwidth Memory (HBM)
than is physically available on the TPU device.
**Sample error messages:**
```
RESOURCE_EXHAUSTED: TPU TensorCore Hbm usage: 34.82G, SparseCore Hbm usage 174.10G, exceeding available bytes: 95.74G
```
```
RESOURCE_EXHAUSTED: XLA:TPU compile permanent error. Ran out of memory in memory space hbm. Used 49.34G of 32.00G hbm. Exceeded hbm capacity by 17.34G.
```
**XLA backends:** TPU
## Overview
XLA performs checks to ensure that the aggregate size of all necessary static
allocations fit in the device's HBM.
The compiler manages the TPU's fixed HBM capacity for several types of
allocations:
- **Program inputs and outputs:** Training batches, optimizer states etc.
- **TPU temporaries:** Dynamic memory required for intermediate calculations
(e.g. activations, gradients, etc).
- **Compiled binary:** The machine code for both TensorCore (TC) and
SparseCore (SC).
- **System overhead:** Reserved space for the XLA Runtime (e.g. infeed buffers
on older TPU generations).
- **Constants:** Constant values embedded in the HLO IR are allocated on HBM.
- **Compiler internals:** Program level and per-HLO allocations (e.g. routing
info for nodes in the mesh).
This error occurs when the XLA compiler cannot fit all of the above allocations
into the device HBM.
## Debugging
Carefully analyze the error message and logs to determine which category of HBM
OOM below best describes your error:
- **"TC Hbm usage: X, SC Hbm usage Y":** If the error explicitly breaks down
HBM usage, the aggregate TensorCore (TC) + SparseCore (SC) usage exceeds the
HBM limit. → Jump to
[Scenario 1. Balance TC and SC HBM usage](#scenario_1_balance_tc_and_sc_hbm_usage).
- **"Ran out of memory in memory space HBM"**: Check the logs for an
enumeration of the largest allocations on HBM.
- In case one or more unexpectedly large tensors (e.g. > 50% of HBM limit)
are present → Jump to
[Scenario 2. Out of memory due to unexpectedly large allocations](#scenario_2_out_of_memory_due_to_unexpectedly_large_allocations).
- If no unexpectedly large tensors are present in the logs → Jump to
[Scenario 3. Out of memory due to aggregate allocations](#scenario_3_out_of_memory_due_to_aggregate_allocations).
---
### Scenario 1. Balance TC and SC HBM usage
If the error explicitly breaks down usage, e.g., *"TC Hbm usage: X, SC Hbm usage
Y"*, this means the aggregate TensorCore (TC) + SparseCore (SC) usage exceeds
the HBM limit. Compare the two values to identify the bottleneck:
- **High SparseCore usage**
- **Optimize HBM stack usage:** HBM stack memory consumption scales with
`feature_width`, `max_unique_nz_per_row` and `logical_replica_count`.
You can reduce peak stack usage by tuning the
`--xla_sc_num_serialized_tables_to_optimize_hbm` flag which serializes
the processing of tables. This comes at the cost of reduced parallelism.
- **Check padding overhead:** SparseCore aligns embedding tables to 32B (8
floats). Tables with small feature widths (e.g., < 8 floats) incur
significant padding overhead, wasting HBM.
- **Reduce heap usage:** High values for `maximum_parallel_iterations`
increase the amount of input data prefetched into the HBM heap. Lowering
this value can free up significant memory.
- **Verify sharding:** Ensure embedding tables are properly mod-sharded
across all chips. See
[How limits translate to tables](https://openxla.org/xla/sparsecore).
- Check out
[SC: Performance and memory bottlenecks](https://openxla.org/xla/sparsecore#10_performance_memory_bottlenecks)
for more ideas.
- **High TensorCore usage**
- Proceed to
[Scenario 2](#scenario_2_out_of_memory_due_to_unexpectedly_large_allocations).
- **Balanced**
- If neither is individually excessive but the sum is too high, you are at
the chip's capacity. You must try lowering usage of both components.
Follow recommendations in all three sections.
### Scenario 2. Out of memory due to unexpectedly large allocations
If you observe the error message *"Ran out of memory in memory space HBM"* and
one or more unexpectedly large allocations are present in the logs (> 50% of HBM
limit), it is almost never a hardware capacity issue. It is typically a
configuration error. Check the XLA label (if present) of the large allocations
for hints on their JAX source code.
- **Remove debugging artifacts**
- Using
[jax.debug.print()](https://docs.jax.dev/en/latest/_autosummary/jax.debug.print.html)
in large-scale runs can force the compiler to materialize the full
tensor in HBM to transfer it to the CPU, breaking fusion and increasing
peak memory usage. Remove any left-over `jax.debug.print()`s.
- **Fix inefficient mesh shapes or sharding**
- Incorrect mesh shapes or missing sharding annotations can cause the
compiler to default to **replication** - forcing the compiler to try to
fit really large tensors on a single chip.
- Check the shapes of the large allocations and verify sharding is
correctly specified and propagated by XLA.
### Scenario 3. Out of memory due to aggregate allocations
If you observe the error message *"Ran out of memory in memory space HBM"* and
no unexpectedly large tensors are present in the logs, the program runs out of
capacity due to the aggregate sum of allocations exceeding the HBM limit. In
this case, it is often helpful to visualize the memory profile to identify the
specific buffers contributing to the peak usage. See
[Debug OOM errors with XProf](https://openxla.org/xla/oom_debugging) for a
step-by-step guide on identifying peak memory contributors.
Once you have identified some of the top contributors, use the following steps
to optimize the memory footprint.
#### Scenario 3.A Adjust configuration
You can often resolve OOMs with these configuration adjustments:
- **Reduce batch size:** The memory needed for intermediate activations and
gradients is directly proportional to the batch size. Reducing the batch
size can often help reduce memory usage, although you may need to retune
your learning rate, momentum, or optimizer hyperparameters to maintain model
stability.
- **Donate input buffers:** When JAX executes a computation it uses buffers on
the device for all inputs and outputs. If you know that one of the inputs is
not needed after the computation, and if it matches the shape and element
type of one of the outputs, you can specify that you want the corresponding
input buffer to be donated to hold an output. This will reduce the memory
required for the execution by the size of the donated buffer. You can
achieve this by specifying
[donate_argnums](https://docs.jax.dev/en/latest/buffer_donation.html)
parameter as an argument when using `jax.jit`.
- **Enable mixed precision (bfloat16):** Use bfloat16 or quantization (int8
etc) for the largest tensors in the program if the model architecture and
quality requirements allow. Note that this change can affect model behaviour
and should be considered carefully.
##### Micro-batching (optional)
If reducing the global batch size or increasing the chip count is not viable,
and the batch size per chip is not already minimized, you can try a
micro-batching strategy:
- Split each batch into `n` micro-batches;
- For each micro-batch, process the forward and backward pass;
- Once this is done, accumulate the gradients and update the weight as a
whole.
This process reduces the activation memory as we divided each batch into `n`
micro-batches, so that the if the original batch had size `M`, the activation
memory size becomes `M/n`.
**Potential issues:** - This process increases step time as we have multiple
forward and backward passes. - If the sizes of the model and micro-batch are too
different, you may face convergence issues in your model.
#### Scenario 3.B Optimize architecture and sharding
If configuration changes are insufficient, the model topology might be too large
for the current hardware setup.
- **Use newer TPU generations:** Newer TPUs generally offer more HBM per chip;
switch to newer TPU generations if available.
- **Run on a larger chip topology:** If the model weights are too large for
the existing topology, you can try sharding them across more chips.
- **Implement advanced sharding techniques:**
- Explore more advanced data, tensor, or pipeline parallelism approaches.
- Specify
[sharding hints](https://docs.jax.dev/en/latest/notebooks/Distributed_arrays_and_automatic_parallelization.html#constraining-shardings-of-intermediates-in-jitted-code)
for intermediate values and outputs.
Note that this may cause an increase in network communication overhead due
to splitting tensors across multiple chips.
- **Use JAX host offloading:** Host offloading techniques allow the user to
offload large tensors to the host CPU memory (e.g.
[activation offloading](https://docs.jax.dev/en/latest/notebooks/host-offloading.html#activation-offloading)
and
[optimizer state offloading](https://docs.jax.dev/en/latest/notebooks/host-offloading.html#optimizer-state-offloading)).
Note that host offloading techniques can severely impact performance, since
these operations will force the system to constantly move large tensors back
and forth between TPU HBM and CPU RAM.
#### Scenario 3.C Check tensor padding and alignment
Inefficient tensor shapes are a common, silent cause of OOMs on TPUs. To get
peak performance on TPU's, XLA pads tensor dimensions—typically to multiples of
128 for the minor-most dimension and 8 for the second-minor. This padding
affects both input arrays and intermediate tensors (HLO temporaries),
potentially inflating memory usage significantly, especially with small
dimension sizes. See
[Array Layouts](https://docs.jax.dev/en/latest/pallas/tpu/details.html#array-layouts).
- **Audit shapes of large buffers:** (On TPU v5 with default layouts)
- Hovering over a buffer in
[Xprof Memory Viewer](https://openxla.org/xprof/memory_viewer#memory_viewer_components)
brings up the buffer details card which contains buffer details
including padding information.
- *Example*: A shape of `(129, 1024)` might be padded to `(256, 1024)`,
resulting in nearly 50% memory waste.
- *Correction:* A shape of `(128, 1024)` requires no padding and incurs 0%
memory waste.
- **Align dimensions:** Ensure all large tensor dimensions (batch size,
embedding dimension, hidden size) are multiples of 128. Note that this
change can affect model behaviour and should be considered carefully.
#### Scenario 3.D Tune key memory impacting XLA flags
[Key memory flags](https://openxla.org/xla/flags_guidance#memory_flags) can be
tuned to trade-off performance for lower memory usage. However, this strategy
should be used as a last resort measure since it can adversely affect
performance.
#### Scenario 3.E Tune XLA rematerialization pass/manual checkpointing
If the model is close to fitting into memory, you can use the
[jax.checkpoint](https://docs.jax.dev/en/latest/_autosummary/jax.checkpoint.html)
decorator with `jax.grad` to manually control which intermediates are saved on
the forward pass versus recomputed on the backward pass. Note that this
operation may impact performance, since you are explicitly trading compute
cycles for HBM. Check out the JAX documentation for more information: -
[Gradient checkpointing with `jax.checkpoint` (`jax.remat`)](https://docs.jax.dev/en/latest/gradient-checkpointing.html) -
[Control autodiff’s saved values with `jax.checkpoint` (aka `jax.remat`)](https://docs.jax.dev/en/latest/notebooks/autodiff_remat.html) -
[JAX Memories and Host Offloading](https://docs.jax.dev/en/latest/notebooks/host-offloading.html)
Alternatively, you can force the `XLA::Rematerialization` pass to prioritize
memory savings, potentially at the cost of slower compilations:
| Flag | Description | Impact / Trade-off |
| --- | --- | --- |
| `--xla_tpu_max_hbm_size_mib` | Manually sets the limit on HBM size used by the Rematerialization pass. | Forces the compiler to work harder to fit the program into a limit smaller than the actual physical HBM. |
| `--xla_tpu_rematerialization_algo=PEAK_PRIORITY` | Focuses efforts at the points of peak memory usage. | Can be more efficient for aggressive memory reduction than the default algorithm. |
| `--xla_tpu_rematerialization_max_block_size_limit=32` | Controls the maximum number of instructions in a block that can be rematerialized at once. | Increasing this allows for memory savings at the cost of **significantly increases compile time**. |
| `--xla_tpu_rematerialization_block_effort_factor=10.0` | Defines the amount of effort (compile time) spent searching for blocks to rematerialize. | Higher values allow a more exhaustive search for memory savings at the cost of **increased compile times**. |
| `--xla_tpu_pre_fusion_remat=true` | Enables an additional Rematerialization pass *before* the fusion pass. | Can find more memory savings, but increases compile times and may **potentially impact numerical stability**. |
Note that making changes to XLA flags should be used as a last resort measure,
since it can adversely affect performance.
#### Scenario 3.F Use advanced profiling tools
[Debug OOM errors with XProf](https://openxla.org/xla/oom_debugging.md) provides
a tutorial on using the
[XProf Memory Viewer](https://openxla.org/xprof/memory_viewer) to visualize the
compiler's view of HBM usage.
This tool allows you to see peak memory allocation and buffer lifetimes, which
is crucial for understanding exactly what consumes HBM at the point of peak
utilization. For general profiling setup, see
[Getting started with Xprof](https://openxla.org/xprof#getting_started) and
[TensorBoard Profiling](https://docs.jax.dev/en/latest/profiling.html#xprof-tensorboard-profiling).
## Summary table
The following table contains a summary of potential interventions to solve OOM
errors and information that will help you decide what to do.
Intervention | Safe to do? (Will it change the behavior of the program?) | Potential gains | Telltale signs (is this actually the bottleneck that you're experiencing?)
----------------------------------------- | ----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | --------------------------------------------------------------------------------------------------------------- | --------------------------------------------------------------------------
Using advanced sharding techniques | **Yes.** It almost never changes the numerical correctness of the experiment, although it can cause network communication overhead due to splitting tensors across multiple chips. | **Massive gains** (up to a 256x reduction) | Unexpectedly large individual allocations in the memory viewer (e.g., a single tensor replicated across all TPUs that is 256x bigger than the others). Active arrays showing as un-sharded in TensorBoard hooks.
Reducing batch size | **No.** It changes the training dynamics and usually requires retuning the learning rate. (Note: **Microbatching** is a safe alternative that reduces memory without changing behavior). | **Massive gains** (can save a factor of thousands). | "Temporaries" failing to allocate during gradient calculations. Seeing "JVP" in the operation name, and encountering many batch-size-shaped tensors in the memory profile.
Enabling mixed Precision (e.g., Bfloat16) | **Risky.** It alters numerical precision, which can change experiment results or cause the model to fail to converge entirely. | **Moderate gains** (typically a factor of 2x, as it halves memory usage). | The memory viewer confirms that the largest tensors are currently utilizing 32-bit floats (`float32`).
Manual checkpointing (`jax.checkpoint`) | **Yes.** It does not alter behavior; it merely trades computation time (flops) to save memory by recomputing tensors instead of storing them. | **Large gains** (e.g., can result in only half of the activations needing to exist in memory at the same time). | Multiple tensors of the exact same size filling up memory during a backward pass. Often accompanied by "JVP" in the operation name.
Donating input buffers (`donate_argnums`) | **Yes.** Maintains experiment integrity. If applied incorrectly, it will not corrupt data but will simply throw a clear error message. | **Marginal gains** (~1% memory savings). | There is no specific telltale sign, but it is considered a "free win" that is always worth trying.
Changing model dimensions | **No.** Directly alters the model's behavior. Modifying input or output dimensions can completely break compatibility with the dataset. | **Variable gains** depending on how drastically hidden dimensions or layers are reduced. | The Xprof memory viewer shows a large amount of memory wasted on "padding" because array dimensions are not powers of two or multiples of 128 (e.g., a dimension of 2050 instead of 2048).
Host offloading (CPU) | **Yes (numerically)**, but **No (performance)**. While mathematically safe, it is considered a "foot gun" that may cause severe speed bottlenecks due to data transfer between the CPU and TPU. | **Marginal gains** (the CPU only has about 3x the memory of the TPU). | Typically a last resort for massive optimizer states or for memory-heavy data preparation/pre-processing steps.
third_party/xla/docs/errors/error_1001.md
Vendored
+77
View File
@@ -0,0 +1,77 @@
# Error code: E1001
**Category:** Compile Time: Scoped Vmem OOM
This error indicates that the program requires more Scoped Vector Memory (Vmem)
than what was allocated.
**Sample Error Messages:**
```
RESOURCE_EXHAUSTED: Ran out of memory in memory space vmem while allocating on stack for %my-custom-kernel = bf16[2048,4096]{1,0:T(8,128)(2,1)} custom-call(...) ...
```
**XLA Backends:** TPU
## Overview
TPUs have Vector Memory (VMEM) which is a local scratchpad memory used
exclusively by the TensorCore (TC). The compiler manages Vmem for different
types of allocations:
* **Instruction-scoped allocations:** Temporary storage in Vmem while executing
a single HLO instruction. This includes operand span buffer (e.g. for double
buffering) and register spills.
* **Program-scoped allocations:** Allocations that live beyond the scope of
a single HLO instruction. These are usually HLO temporaries and intermediate
results that are inputs and/or outputs of HLO instructions.
A Compile Time Scoped Vmem OOM occurs when the instruction-scoped allocations
exceed the allocation limit for that instruction. This limit is controlled
- globally for the entire program via the flag
[--xla_tpu_scoped_vmem_limit_kib](https://openxla.org/xla/flags_guidance)
and
- per custom kernel via
[vmem_limit_bytes param](https://docs.jax.dev/en/latest/_autosummary/jax.experimental.pallas.tpu.CompilerParams.html#jax.experimental.pallas.tpu.CompilerParams.vmem_limit_bytes).
These errors are typically caused by an internal compiler bug or by a
custom kernel exceeding its allocation limit.
## Debugging
Carefully analyze the error message to identify if the error stems from a
custom kernel or a standard HLO. An error due to a custom kernel should have
the following signature:
```
Ran out of memory in memory space vmem while allocating on stack for %my-custom-call = <output-shape> custom-call(<params>), custom_call_target="tpu_custom_call" ...
```
- **Custom kernel scoped Vmem OOM**: If the error points to a custom kernel →
Jump to [Retune the Kernel](#retune_the_kernel).
- **Non-kernel Vmem issues**: If the Vmem OOM occurs due to a
non-custom-kernel op, it is likely an internal compiler bug.
[Please file a bug report on XLA](https://github.com/openxla/xla) with an
HLO dump.
### Retune the kernel
If the error originates from a custom kernel, use the following techniques to
reduce the kernel's memory requirement:
* **Adjust Block Sizes:** Reduce the block sizes (tile sizes) in your kernel
configuration, to lower Scoped Vmem usage.
* **Set Per-Kernel Scoped Vmem Limits:** Explicitly request the required
amount of memory for that specific kernel using the
[vmem_limit_bytes param](https://docs.jax.dev/en/latest/_autosummary/jax.experimental.pallas.tpu.CompilerParams.html#jax.experimental.pallas.tpu.CompilerParams.vmem_limit_bytes)
* **Modify Memory Coloring:** Explicitly color/constrain the kernel's
inputs/outputs to Vmem using
[pallas.tpu.with_memory_space_constraint](https://docs.jax.dev/en/latest/_autosummary/jax.experimental.pallas.tpu.with_memory_space_constraint.html).
Be careful not to color too many inputs outputs to Vmem, as that might cause
an overall Vmem OOM.
* **Adjust Vmem limit:** If kernel specific retuning is difficult or the issue
affects many kernels, you can adjust the global Vmem limit using the flag
[--xla_tpu_scoped_vmem_limit_kib](https://openxla.org/xla/flags_guidance).
third_party/xla/docs/errors/error_1200.md
Vendored
+60
View File
@@ -0,0 +1,60 @@
# Error code: E1200
**Category:** Compile Time: Host Offload Output Mismatch
This error occurs when a tensor explicitly offloaded to host memory is returned
as a program output, but the program's output signature is not configured to
expect host memory.
**Sample Error Messages:**
```
INVALID_ARGUMENT: Tensor which is moved to host (starting from tuple.64) is returned from the entry computation but the layout for this output is not set to host memory.
```
**XLA Backends:** TPU, GPU
## Overview
When the compiler encounters an annotation to offload a tensor to the host
(CPU), it tracks that tensor's location through the computation graph until one
of three events occurs:
1. **Move to Device:** A matching annotation moves the tensor back to the
accelerator.
2. **Host Computation:** The tensor is consumed by a host-side operation.
3. **Program End:** The tensor reaches the end of the program and becomes an
output.
This error is triggered in scenario **#3**. The tensor is physically located in
host memory at the end of the execution, but the XLA program's entry
computation signature defines that specific output as residing in **Device
Memory**. Because the compiler cannot implicitly change the entry computation's
interface, it raises an error.
## Debugging
To resolve this error, determine whether you intended for this tensor to be an
output on the Host or if it should have been moved back to the Device before
returning.
* **Intended to return on Host:** If you explicitly want this tensor to be
returned in host memory (avoiding a transfer back to device), you should
explicitly set the output memory space of the entry computation to **Host
Memory** for this specific output.
* **Intended to return on Device:** If the tensor was meant to stay on the
device or return to it before the program ends, you likely missed an
annotation. Insert a matching annotation to move the tensor back to the
device.
If the source of the offloaded tensor is unclear, or you cannot find where
the "move to device" annotation is missing, use XLA logging to trace the
instructions.
* **Enable logging:** If you are on Google Cloud TPU, rerun your program with
the following flag: `--vmodule=host_offloader=1`.
* **Analyze Logs:** Look for the "trace" output in the logs. This will show
the path of the tensor starting from the offload instruction. Use this to
pinpoint exactly where the tensor reaches the program boundary without being
moved back to the device.
third_party/xla/docs/errors/error_2001.md
Vendored
+88
View File
@@ -0,0 +1,88 @@
# Error code: E2001
**Category:** Compile Time: Unsupported RHS DataType on Hardware
This error occurs when the data type used for the **Right-Hand Side (RHS)**
operand in a matrix multiplication (e.g., `jax.lax.dot_general`, `jax.lax.conv`,
`jax.numpy.matmul`, or the `@` operator) is not natively supported by the
specific TPU generation being used.
**Sample error messages:**
```
INTERNAL: Mosaic failed to compile TPU kernel: Unsupported matmul RHS type on target: 'vector<256x256xi8>'
...
The MLIR operation involved:
%13440 = "tpu.matmul"(%13435, %13437, %13439) <dimension_numbers = #tpu.dot_dimension_numbers<...>
```
**XLA backends:** TPU
## Overview
The TPU's Matrix Multiply Unit (MXU) natively supports `Float32` operations on
all hardware generations.
However, native support for `BFloat16` and other quantized data types
(e.g., Int4, Int8, or Float8) varies by hardware generation. This error is
triggered when your kernel attempts to map a matrix multiplication to the MXU
using a data type that your specific TPU generation does not have the physical
circuitry to execute.
This error typically indicates that the compiler's **Canonicalization** pass—
which attempts to automatically convert unsupported types into supported ones
(e.g., via software emulation)—was unable to find a valid conversion rule or
was prevented from doing so because **Compatibility Mode** was disabled.
## Debugging
To resolve this error, you must align your data types with the capabilities of
your hardware. You have following options:
### 1. Cast to native types
The most reliable fix is to manually cast your operands to a hardware-supported
datatype (like `Float32` or `BFloat16` on TPU v4+) inside your kernel before the
matmul operation.
* **Why:** `Float32` is the universal data type supported natively by the
MXU on all TPU generations.
* **Trade-off:** This comes with a VPU (Vector Processing Unit) cost - the
cycles required to perform the cast, but it guarantees your kernel
will run on the current hardware.
### 2. Check Compatibility Mode
Typically the compiler can automatically handle these type mismatch issues in
**Compatibility Mode** which is enabled by default. Double check XLA configs
to make sure `--xla_mosaic_compat_mode` is not set to false.
This acts as a "polyfill," injecting software emulation sequences for operations
your hardware does not natively support.
**What Compatibility Mode enables:**
* **Mixed-precision MatMuls:** Allows mixing Integer operands with Float
accumulators by automatically inserting cast operations (e.g., extending
integers to `Float32` before the matmul).
* **Low-precision emulation:** On certain hardware generations, emulates
unsupported types like `4-bit` floating point (`4E2M1FN`) or `8-bit`
floating point (`8E4M3FN`) by extending them to supported types like
`BFloat16` or `Float32` before execution.
Note that this mode prioritizes compatibility over peak performance as emulation
requires additional instructions to convert data formats before the MXU can
operate on them.
### 3. Upgrade hardware or request support
If your algorithm strictly requires native performance for types like `Int4` or
`Float8` without the overhead of casting or emulation, you will need to run on
a newer TPU generation with native support.
**Feature request:** If you believe your hardware supports this operation, or if
the compiler is missing a valid emulation path even in Compatibility Mode,
please file a feature request. We usually guarantee that operations are forward
compatible. So if your kernel runs on a TPU generation then it should run on all
future generations, but it is not guaranteed to have emulation for older
generations (for some of which the casts would be very expensive).
third_party/xla/docs/errors/error_2002.md
Vendored
+53
View File
@@ -0,0 +1,53 @@
# Error code: E2002
**Category:** Compile Time: Mosaic Input/Output Misaligned Block and Tiling
This error occurs when the block shape of a kernel input or output does not
align with the default tiling of the datatype on the specific TPU hardware
being used.
**Sample error messages:**
```
UNIMPLEMENTED: Mosaic failed to compile TPU kernel: Failed to set window params
for input 0: Operand of shape (..., 256, 8192) has tiling (16, 128), but its
block shape (..., 8, 8192) is not divisible by tiling evenly nor matches the
full shape.
```
**XLA backends:** TPU
## Overview
Tensor cores (TC) in TPUs have two-dimensional vector registers. The two
dimensions are called **sublane** and **lane**. Since TPU compute units (e.g.,
MXU) operate at the granularity of vector registers, XLA arrays are laid out in
TPU memory in tiles.
This tiling (e.g., `8x128`) minimizes data transformations when feeding the
compute units. The exact tiling dimensions (sublanes × lanes) depend on the
hardware generation and the data type. For example, a common tiling for
most types is **8×128**.
At compile time, XLA enforces the following constraints for the minor and
2nd-minor dimensions of each kernel input/output:
1. **Divisibility:** The block dimension must be a multiple of the tile
dimension in the underlying tensor, or
2. **Full shape exception:** If the block dimension is not divisible, it must
be equal to the **full size** of that dimension in the underlying tensor.
This error is triggered when a block violates both conditions. For example,
loading a block of shape `(8, 100)` from a input of shape `(8, 1024)` on
hardware with shape `(8, 128)` tiling fails because `100` is not divisible by
`128` and `100 != 1024`. But, it would be allowed if the input shape was
`(32, 100)`.
## Debugging
To resolve this error, ensure your kernel's block shapes align with the
current hardware tiling. Modify your kernel code to align the block size such
that it is a multiple of the required tiling.
For example, if the error states the tiling is `(16, 128)` but your block shape
is `(8, 128)`, change the block spec such that the shape matches `(16, 128)`.
third_party/xla/docs/errors/error_2003.md
Vendored
+75
View File
@@ -0,0 +1,75 @@
# Error code: E2003
**Category:** Compile Time: Mosaic Unproven Memory Access Alignment
This error occurs when the compiler analyzes a memory access operation (such as
`vector.load`, `vector.store`, `tpu.load`, or `tpu.store`) and cannot
statically prove that the dynamic index used for a specific dimension is a
multiple of the required **tiling size**.
**Sample error messages:**
```
INTERNAL: Mosaic failed to compile TPU kernel: cannot statically prove that index in dimension 1 is a multiple of 128
at location: ...
The MLIR operation involved:
%14372 = "vector.load"(%14371, %93, %14363) : (memref<4x256xf32, #tpu.memory_space<vmem>>, index, index) -> vector<1x32xf32>
```
**XLA backends:** TPU
## Overview
When your kernel loads or stores a vector, the memory address
(calculated from the base pointer plus the dynamic index) must align with the
vector's **tiling size** on the hardware. For example, if a dimension is tiled
by 128 elements, the dynamic index used to access it must be `0`, `128`, `256`,
etc. Note that many operations (like vector loads and stores) have no such
requirements for static indices.
The compiler enforces this requirement using **static analysis**. It traces the
history of the index variable back through the arithmetic operations that
produced it (e.g., multiplications, additions). If the compiler cannot guarantee
(at compile time) that the resulting value will always be divisible by the
tiling size, it raises this error.
The compiler treats "proven misalignment" and "unknown alignment" identically.
So if you use an index that is mathematically guaranteed to be misaligned (e.g.,
`i * 128 + 32`), the compiler will raise the same error.
So this error can occur when
1. You use a runtime variable (dynamic index) to access memory.
2. The index calculation logic is too complex for the compiler to analyze.
3. The index is mathematically valid but lacks an explicit proof in the code.
4. Static analysis determines "proven misalignment".
## Debugging
To resolve this error you have the following options:
### 1. Assert alignment explicitly
If you know your index is valid but the compiler cannot prove it, use the
`tpu.assume_multiple` operation. This acts as a promise to the compiler that a
value is divisible by a specific factor.
### 2. Use aligned loads and rotate
In scenarios where the misalignment is intentional,
instead of loading a small, unaligned vector segment:
* Load a larger, fully aligned tile and then rotate the values by a dynamic
amount to shift the desired data into position (since vector slices with
dynamic start indices are not supported). or
* Reshape or pad the tensor so that the data starts at index 0 and the stride
between accesses matches the hardware alignment.
* Example: If you iterate over chunks of size 32 starting at offset 1,
your offsets are 1, 33, 65... (unaligned).
* Fix: Re-pack the data into a new tensor where the first chunk is at 0
and the dimension is padded to 128. Your offsets become 0, 128, 256...,
which satisfies the alignment requirement.
These methods consume more memory but often simplify kernel logic and eliminate
the need for manual alignment assertions.
third_party/xla/docs/errors/error_3000.md
Vendored
+89
View File
@@ -0,0 +1,89 @@
# Error code: E3000
**Category:** Compile Time: SparseCore Allocation Failure
This error occurs when the `XLA:SparseCore` compiler is unable to allocate a
contiguous block of memory in the specified memory space required by the current
SparseCore program.
**Sample error messages:**
```
INTERNAL:Failed to run pass pipeline. Hlo-Op: result.1:279:1: error: 'memref.alloca' op current allocation offset upper bound (140704 words) exceeds the legitimate user allocatable offset upper bound (131071 words) in memory space 201 when allocating 23440 words. result.1:279:1: note: see current operation: %232 = "memref.alloca"() <{operandSegmentSizes = array<i32: 0, 0>}> : () -> memref<23440xf32, 201>
```
**XLA backends:** TPU
## Overview
The SparseCore (SC) is a specialized processor for sparse workloads. It relies
on specific memory hierarchies to manage data movement efficiently. The XLA
compiler attempts to statically size and allocate buffers based on hardware
limits and user-defined shapes. This error indicates an **Out of Memory (OOM)**
condition during this allocation phase.
The error message typically specifies a memory space ID. Below are the common
memory spaces and their integer encodings:
| Memory Space ID | Name | Description |
| --- | --- | --- |
| **0** | **Smem** | Local Scalar Memory. Used for scalar registers and control flow. |
| **201** | **TileSpmem** | Tile-specific Scratchpad Memory. Fast, local SRAM available to a specific SC tile. |
| **202** | **Spmem** | Shared Scratchpad Memory. Used to opportunistically stage data (inputs, outputs, intermediates) to hide HBM latency. |
| **203** | **HBM** | High Bandwidth Memory. Large, shared memory used for embedding tables, heaps, and stacks. |
| **204** | **Sync Flags** | Synchronization primitives used for coordination. |
For a deep dive into SC and its memory hierarchy, refer to the
[SparseCore documentation](https://openxla.org/xla/sparsecore).
## Debugging
The resolution depends on which memory space failed the allocation.
### Scenario 1. HBM allocation failures (Memory Space ID: 203)
This error occurs if a single temporary allocation requested by the SparseCore
program is too large to fit in the available HBM. In standard embedding
workloads and SC offloaded collectives extremely large per-core partitions
or incorrect sharding specifications can force the compiler to request massive
buffers.
**Recommended actions:**
- **Check sharding:** Ensure your embedding tables and SC input/output tensors
are partitioned/sharded correctly. If a single core is responsible for too
much data, the allocation might fail.
- **Adjust limits:** Review `max_ids_per_partition` and
`max_unique_ids_per_partition`. If these are set unnecessarily high, the
compiler reserves more memory than needed. Refer to
[How limits translate to tables](https://openxla.org/xla/sparsecore#7_how_limits_translate_to_tables_on_sparsecore).
### Scenario 2. Internal memory failures (Memory Space IDs: 0, 201, 202, 204)
Allocation failures in **Smem**, **TileSpmem**, **Spmem**, or **Sync Flags**
typically occur due to compiler bugs or limitations in the allocation strategy,
where the compiler fails to account for all memory
requirements.
**Recommended actions:**
1. **Isolate the failing XLA operation:** To identify the specific SC HLO or
Mosaic kernel causing the failure, generate the intermediate compiler
representations:
- **Dump SparseCore MLIR:** Set the flag
`--xla_sc_dump_mlir_to=/path/to/dump`. This generates the MLIR of the
SparseCore program, allowing you to see which allocation size matches
the error message.
- **Dump Mosaic LLO:** For custom kernels, use
`--xla_mosaic_dump_to=/path/to/dump` to inspect all Low Level Optimizer
(LLO) programs emitted by Mosaic.
2. **Reduce scratch sizes (Pallas users):** If the failure occurs within a
Mosaic kernel, review your `scratch_shapes` configuration. Ensure that your
`pltpu.SMEM` requests fit within the hardware specifications for your
specific TPU generation.
3. **Disable collective offload:** If the error arises from a SC offloaded
collective operations, try disabling the SC offloading features:
- `--xla_tpu_enable_sparse_core_collective_offload_all_gather=false`
- `--xla_tpu_enable_sparse_core_collective_offload_all_reduce=false`
4. **File a bug:** If the above steps do not resolve the issue, it is likely a
compiler bug. [Please file a bug report](https://github.com/openxla/xla).
third_party/xla/docs/errors/error_3001.md
Vendored
+62
View File
@@ -0,0 +1,62 @@
# Error code: E3001
**Category:** CompileTime: SparseCore No Viable Logical Replica Count
This error occurs when the `XLA:SparseCore` compiler fails to determine a
valid logical replica count configuration that allows the workload to fit within
the SparseCore's local scratchpad memory (Tilespmem).
**Sample error messages:**
```
XLA:TPU compile permanent error. Compilation failure: No viable logical replica count for the embedding table with metadata: max_nz_per_row = 141352, max_unique_nz_per_row = 8, feature_width = 8, sample_count = 204800 (last tried split factor for vector splitting = 1, last tried split factor for sample dimension splitting = 1, fixed_size_allocation_bytes = 410880, row_dependent_size_allocation_bytes = 1696224, total_spmem_size_bytes = 524288) ...
```
**XLA backends:** TPU
## Overview
This error is specific to **SparseCore** use cases, particularly for Large
Embedding Models (LEMs).
The **logical replica count** is an internal compiler parameter that determines
how input batches are partitioned to manage scratchpad allocation pressure. The
compiler attempts to split the workload into smaller chunks (replicas) so that
the intermediate buffers required for each chunk fit into the SparseCore's
limited **Scratchpad Memory**. Generally, a higher logical replica count reduces
allocation pressure by processing smaller batches of data at a time.
This error indicates that even after attempting various splitting
configurations, the compiler could not find a setup where the required buffers
fit in the Tilespmem memory. The allocation size is determined by a combination
of:
- **`sample_count`**: The number of embedding lookup IDs assigned to each
SparseCore (derived from batch size).
- **`feature_width`**: The size of the embedding dimension.
- **`max_nz_per_row`**: The maximum number of non-unique embedding lookup IDs
across all SparseCores.
- **`max_unique_nz_per_row`**: The maximum number of unique embedding lookup
IDs.
## Debugging
To resolve this error, you need to reduce the memory pressure on the SparseCore
scratchpad.
### 1. Improve metadata estimations
The compiler allocates memory based on `max_nz_per_row` and
`max_unique_nz_per_row`. If these values are estimated conservatively (i.e., set
much higher than the actual data requires), the compiler will reserve
unnecessary space, causing this error. Ensure these parameters accurately
reflect the actual ID distribution of your dataset.
You can consider applying **Feedback-Directed Optimization (FDO)** to determine
optimal values for these parameters.
### 2. Reduce batch size
The `sample_count` is directly derived from your global batch size. Reducing the
batch size decreases the amount of data each SparseCore must process per step,
thereby reducing the size of the required scratchpad buffers.