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/megascale/_toc.yaml
Vendored
+23
View File
@@ -0,0 +1,23 @@
# Copyright 2026 The OpenXLA Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
toc:
- heading: Megascale debugging playbook
- title: Overview
path: xla/megascale/overview
- title: Getting started
section:
- title: Debugging Megascale
path: xla/megascale/debugging_workflow
third_party/xla/docs/megascale/debugging_workflow.md
Vendored
+336
View File
@@ -0,0 +1,336 @@
# Debugging Workflow
This document describe a general workflow for debugging MXLA issues.
## Prerequisite
1. Use JAX 0.6 or up, and enable JAX distributed service. This version of JAX
contains additional logging that can help identify which workers are
experiencing issues.
2. (Optional) Generate an HLO dump using the --xla_dump_to flag when
initializing your workload. This is discussed in the [XLA
documentation](https://openxla.org/xla/hlo_dumps).
3. (Optional) Set --vmodule=real_program_continuator=1 to enable verbose
logging for the TPU program execution status.
## Flow chart
The flowchart below illustrates the debugging process. To access detailed
playbooks for each step, click on the corresponding item in the chart.
<iframe src="flow_chart.svg" width="1000" height="2000" style="border: none;"></iframe>
<!-- linter style off -->
## Hangs
### Locate the Megascale Hang Detected Error
If you see the following error message in your TPU worker logs, this means that
MXLA timed out after detecting no progress:
```text
Megascale hang detected: Timed out waiting for 4 graphs to complete at launch_id 13650. Already completed: 100. StepGloballyInProgress: true. Timeout: 1m
```
1. Workers will report errors to a coordinator for processing.
* For **Pathways** jobs: the digest can be found in the logs of
`resource_manager` job.
* For **McJAX** jobs: the logs can be found on MXLA Coordinator. This is
typically task 0 of slice 0.
2. Check logs around time of error detection, and look for `Megascale detects a
hang`.
3. Follow steps below to diagnose the issue based on the identified cause.
### Diagnosis
#### Interpreting TPU States
Before diagnosing MXLA hangs, it is important to understand the TPU states
report format. Below is a sample report:
```text
Full error digest:
Potential cause: <determined_cause>
Potential culprit workers: <task_name>
First error timestamp: <timestamp>
First error type: <error_type>
TPU states:
Launch ID: <launch_id>
Module: jit.step_fn Fingerprint: <fingerprint>
Sample worker: <task_name>@<host_name>:<tpu_chip>:<tpu_core>
Tag:PC breakdown:
<num_cores>@<location>(HLO): [<task_name>:<host_name>@<tpu_chip>:<tpu_core>, ...]
...
```
#### Bad TPU Chip (tensor core or sparse core)
```text
Megascale detects a hang that is likely caused by bad TPU chips on the following hosts. Please remove the hosts from the fleet and restart the workload. If problem persists please contact Megascale XLA team.
The host that have bad TPUs are: <host_name>
Full error digest:
Potential cause: Bad TPU chips
Potential culprit workers: <job_name>/<task_id>:<host_name>
```
This error means that the issue is potentially caused by a faulty TPU chip. The
error message should include the job information and host name of the faulty
chip. In the example above, the faulty chip is on host `<host_name>`, affecting
task `<task_id>` of the job `<job_name>`. You can configure your job to avoid
that host.
**Note:** There are some cases that the hang was caused by a XLA or custom
pallas kernel bug, but if you see the same host appearing multiple times (for
example more than 3 times) in multiple hang events, the TPU on that host is very
likely faulty.
#### Networking issue
```text
Megascale detects a hang that is likely caused by a networking issue. Please examine the underlying networking stack for the following hosts.
The hosts are: <host_name>
Full error digest:
Potential cause: Networking issue
Potential culprit workers: <host_name_1>, <host_name_2>
```
This error indicates that your job has encountered a failed network link. The
error message should include a single or pair of job name, task id, host name of
the faulty network link. In the example above, the faulty network link is
between host `host_name_1` and `host_name_2`. Sometimes RapidEye can further
localize the faulty host if a single host appears in multiple broken network
links.
#### Different modules
```text
Megascale detects a hang that is likely caused by running different modules on different devices. Please confirm that all workers is running the exact same program. It can also be caused by a hang in a subset of devices and the unaffected devices have moved on to the next program. Please inspect the digest below to further root cause the hang.
Example hosts that have different HLO modules: <host_name>
Full error digest:
Potential cause: Different module
Potential culprit workers: <host_name>
TPU stats:
<host_name>: <pc>
TPU states:
Module: jit_loss_and_grad
Fingerprint: <fingerprint>
Launch ID: 193
<tag>:<pc>(<hlo>): <host_name>
Module: jit_optimizer_apply
Fingerprint: <fingerprint>
Launch ID: 0
<tag>:<pc>(<hlo>): <host_name>
```
This error may indicate that a hang has occurred in a subset of workers, causing
those workers to be stuck at the current module while unaffected workers advance
to the next module. To identify the root cause, inspect the digest printed by
RapidEye in the logs.
The `TPU states` section of the logs shows which modules are running on which
workers. In the example above, workers are running different modules:
`jit_loss_and_grad` and `jit_optimizer_apply`.
#### Fingerprint mismatch for HLO module
```text
Megascale detects a hang that is likely caused by inconsistent HLO module compilation across workers. This is likely a bug in JAX tracing or XLA compiler. Please inspect the HLO dumps to confirm the root cause.
Example hosts that have different HLO fingerprints: <host_name>
Full error digest:
Potential cause: Fingerprint mismatch
Potential culprit workers: <host_name>
TPU stats:
Module: reduce.31
Fingerprint: <fingerprint_1>
Launch ID: 37
<tag>:<pc>(<hlo>): <host_name>
Module: reduce.31
Fingerprint: <fingerprint_2>
Launch ID: 40
<tag>:<pc>(<hlo>): <host_name>
```
This log message indicates the hang was likely caused by inconsistent HLO module
compilation across workers, possibly due to an issue with JAX tracing or the XLA
compiler. If you see this log, follow [these
steps](https://openxla.org/xla/hlo_dumps) to collect HLO dumps from the culprit
workers for further debugging.
#### Data input stall
**Note:** This error is not yet implemented.
```text
Megascale detects a hang that is likely caused by data input stall on the
following hosts. Please check the workers to make sure the data input pipeline
is working properly.
The host that have data input stalls are: <host_name>
```
This error means that all devices launched the same program, but that input data
was not provided to the program before the system timed out. To fix this issue,
confirm that:
1. The identified hosts can access the input datasource.
2. The identified hosts are properly loading/parsing the input datasource.
3. Confirm identified hosts are not throttled on reads to the input datasource.
#### Unrecoverable error
```text
Some workers have halted with an unrecoverable error:
<worker> : {some error}
Please inspect the error log of these workers:
<worker>
```
This error means that there was an issue that prevented the program from
properly executing and could not be recovered automatically. This error was
unable to be specifically categorized. Further information can be obtained from
checking the logs of the worker(s) mentioned in the error report.
If the error appears to be specific to the given machine (ex. failure to copy
data from TPU to host), then you can configure your job to avoid those hosts.
#### Unknown Error
```text
Megascale detects a hang but cannot determine the root cause. Please inspect the
full digest below.
```
This error means that there was an issue that prevented the program from
properly executing and could not be recovered automatically. This error was
unable to be specifically categorized and there is no further error information
available.
<!-- linter style on -->
## Performance
### Get an XProf session
Follow the instructions in the [XProf
documentation](https://openxla.org/xprof/capturing_profiles) to generate an
XProf trace for your problematic run.
### Check for shortage of mapped DMA buffers
The Megascale XLA runtime needs to register host memory before it can be used
for DMAs to and from the TPU. This happens early after the process starts. If
you see these registrations (`MapDmaBuffer` calls) at steady-state then it
indicates that something is wrong. Look for the presence of these calls in XProf
Trace Viewer. See the screenshot below for reference.
**Tip:** Search for the exact worker name, because there can be other workers
with similar or close names. You also search for the term “MapDmaBuffer” on the
page.
![Example xprof trace showing MapDmaBuffer
calls](./images/map_dma_buffer_example_trace.png)
If the issue is observed then try to increase the size of the premapped memory
region by increasing the value of `--megascale_grpc_premap_memory_bytes`,
restarting the job, then checking again.
### Check for memory copies during network transfers
Megascale XLA network transfers are zero-copy by design. However, there are
cases where memory copies will occur and cause degraded performance. Look for
memory copies in Megascale's "Communication Transport" traces as shown in the
example screenshot below.
![Example "Megascale: Memory Copy" trace during network
receive](./images/memory_copy_example_trace.png)
If the issue is observed then try to increase the size of the premapped memory
region by increasing the value of `--megascale_grpc_premap_memory_bytes`,
restarting the job, then checking again.
### Network Analysis
MegaScale also provides a Colab
[notebook](https://github.com/openxla/xla/blob/main/xla/megascale/tools/network_analysis_oss.ipynb)
to help analyze network performance using an XProf trace.
This tool can be used to do the following:
* Examine transfer latencies to identify potential network slowdowns or host
slowness.
* Examine transfer sizes to identify if your workload is optimized to use a
smaller number of larger transfers as opposed to a large number of small
transfers.
* Determine if your workload is unoptimized and produces bursty
collectives/always has a large number of pending collectives.
* Visualize the network throughput timeline to see if the workload is network
bound.
* Examine {source, destination} pairs to identify possible bad hardware on
individual hosts.
#### Collective Slack Too Small
One indicator that your workload is not optimized for compute/communication
overlap is seeing small slack times for a subset of collectives. This can
manifest as longer than expected `recv-done` traces in the trace viewer, or as
collectives with zero or near-zero [slack
time](https://cloud.google.com/tpu/docs/troubleshooting/troubleshoot-multislice#slack_time).
If this is the case, look towards identifying bottlenecks in your workload that
may be causing parts of your program to not overlap compute and network
communication.
#### High Network Bandwidth Demand
If you are observing long `recv-done` op latencies within your model XProf, this
could be an indication that the model is 'Bandwidth Bound' in those portions of
the step function (is blocked by available network bandwidth in the system).
You can generate a timeline of network usage for your model. If you see
consistently high network usage throughout the step, or specific regions with
large spikes, then your model may be bandwidth bound in those regions.
Use the [Network Analysis](#network-analysis) to generate a timeline of network
usage: ![example](./images/bandwidth_bound.png)
To mitigate bandwidth bound models, you can:
1. Check the [Collective Slack](#collective-slack-too-small) of your model.
Models with many collectives with low slack will have bandwidth bound
regions.
2. Confirm that the network settings are optimized.
3. Examine your model structure and data sharding to see if there are ways to
increase computation/communication overlap.
4. (Data parallel models) Confirm that you have sufficient batch size in each
local replica to overlap with the communication.
#### High Network Latency
If the bandwidth is not saturated, you may want to generate the RPC latency
timeline. If you see high RPC latencies is constantly or sporadically high, this
means that there are some issues with the MXLA RPCs.
Use the [Network Analysis](#network-analysis) to generate the RPC latency
timeline, the following example shows that there is some sporadic 200ms long
tail RPC latencies. ![example](./images/tail_network_latency.png)
#### Confirm Optimal Network Settings
High RPC latencies on Cloud environment are often caused by suboptimal TCP
configuration. Please confirm if all TCP parameters are configured
properly within the container.
If any of the TCP parameters are not correctly configured, consult Google Cloud
ML Compute Services (CMCS) team on how to configure them properly.
#### HLO Dump
Please follow [these steps](https://openxla.org/xla/hlo_dumps) to dump HLO to
the local filesystem on the TPU worker. You may need to upload the dump to GCS
in order to share them with the XLA or Megascale team.
## Straggler
**Note on Future Tooling:** Google is actively working on open-sourcing versions
of diagnostic dashboards to provide a more streamlined experience for Cloud TPU
customers to identify and diagnose stragglers. These will be available soon.
third_party/xla/docs/megascale/flow_chart.svg
Vendored
+373
View File
@@ -0,0 +1,373 @@
<?xml version="1.0" encoding="UTF-8" standalone="no"?>
<!DOCTYPE svg PUBLIC "-//W3C//DTD SVG 1.1//EN" "http://www.w3.org/Graphics/SVG/1.1/DTD/svg11.dtd">
<!-- Generated by graphviz version 2.44.1 (20201121.0304)
-->
<!-- Pages: 1 -->
<svg xmlns="http://www.w3.org/2000/svg" xmlns:xlink="http://www.w3.org/1999/xlink" width="616pt" height="1440pt" viewBox="0.00 0.00 616.00 1440.00">
<g id="graph0" class="graph" transform="scale(0.904523 0.904523) rotate(0) translate(4 1588)">
<polygon fill="white" stroke="none" points="-4,4 -4,-1588 677,-1588 677,4 -4,4"/>
<!-- start -->
<g id="googleblue" class="node">
<title>start</title>
<g id="a_googleblue"><a xlink:title=" " target="_blank">
<path fill="none" stroke="black" d="M204,-1584C204,-1584 12,-1584 12,-1584 6,-1584 0,-1578 0,-1572 0,-1572 0,-1524 0,-1524 0,-1518 6,-1512 12,-1512 12,-1512 204,-1512 204,-1512 210,-1512 216,-1518 216,-1524 216,-1524 216,-1572 216,-1572 216,-1578 210,-1584 204,-1584"/>
<text xml:space="preserve" text-anchor="middle" x="108" y="-1532.88" font-family="Times,serif" font-size="40.00">☹️</text>
</a>
</g>
</g>
<!-- mxla_job -->
<g id="googlegreen" class="node">
<title>mxla_job</title>
<g id="a_googlegreen"><a xlink:title=" " target="_blank">
<polygon fill="lightgrey" stroke="black" points="108,-1476 0,-1440 108,-1404 216,-1440 108,-1476"/>
<text xml:space="preserve" text-anchor="middle" x="108" y="-1435.7" font-family="Times,serif" font-size="14.00">Is it a multi slice job?</text>
</a>
</g>
</g>
<!-- start&#45;&gt;mxla_job -->
<g id="edge1" class="edge">
<title>start-&gt;mxla_job</title>
<path fill="none" stroke="black" d="M108,-1511.7C108,-1504.1 108,-1495.94 108,-1487.97"/>
<polygon fill="black" stroke="black" points="111.5,-1488.02 108,-1478.02 104.5,-1488.02 111.5,-1488.02"/>
</g>
<!-- hang -->
<g id="googlegreen" class="node">
<title>hang</title>
<g id="a_googlegreen"><a xlink:title=" " target="_blank">
<polygon fill="lightgrey" stroke="black" points="229,-1368 121,-1332 229,-1296 337,-1332 229,-1368"/>
<text xml:space="preserve" text-anchor="middle" x="229" y="-1327.7" font-family="Times,serif" font-size="14.00">Does the job hang?</text>
</a>
</g>
</g>
<!-- mxla_job&#45;&gt;hang -->
<g id="edge3" class="edge">
<title>mxla_job-&gt;hang</title>
<path fill="none" stroke="black" d="M137.29,-1413.34C153.64,-1399.02 174.18,-1381.03 191.67,-1365.7"/>
<polygon fill="black" stroke="black" points="193.62,-1368.64 198.84,-1359.42 189.01,-1363.38 193.62,-1368.64"/>
<text xml:space="preserve" text-anchor="middle" x="158.39" y="-1390.36" font-family="Times,serif" font-size="14.00">Yes</text>
</g>
<!-- consult_xla -->
<g id="googlered" class="node">
<title>consult_xla</title>
<g id="a_googlered"><a xlink:title=" " target="_blank">
<path fill="none" stroke="black" d="M465,-288C465,-288 273,-288 273,-288 267,-288 261,-282 261,-276 261,-276 261,-228 261,-228 261,-222 267,-216 273,-216 273,-216 465,-216 465,-216 471,-216 477,-222 477,-228 477,-228 477,-276 477,-276 477,-282 471,-288 465,-288"/>
<text xml:space="preserve" text-anchor="middle" x="369" y="-247.7" font-family="Times,serif" font-size="14.00">Consult XLA team</text>
</a>
</g>
</g>
<!-- mxla_job&#45;&gt;consult_xla -->
<g id="edge2" class="edge">
<title>mxla_job-&gt;consult_xla</title>
<path fill="none" stroke="black" d="M92.06,-1409.12C71.96,-1368.6 40,-1293.42 40,-1225 40,-1225 40,-1225 40,-467 40,-402.89 14.87,-370.51 59,-324 116.66,-263.24 163.32,-304.88 249.72,-288.22"/>
<polygon fill="black" stroke="black" points="250.24,-291.69 259.27,-286.15 248.75,-284.85 250.24,-291.69"/>
<text xml:space="preserve" text-anchor="middle" x="31.75" y="-1134.95" font-family="Times,serif" font-size="14.00">No</text>
</g>
<!-- rapideye -->
<g id="googleyellow" class="node">
<title>rapideye</title>
<g id="a_googleyellow"><a xlink:href="debugging_workflow.md#hangs" xlink:title=" " target="_parent">
<path fill="none" stroke="black" d="M325,-1260C325,-1260 133,-1260 133,-1260 127,-1260 121,-1254 121,-1248 121,-1248 121,-1200 121,-1200 121,-1194 127,-1188 133,-1188 133,-1188 325,-1188 325,-1188 331,-1188 337,-1194 337,-1200 337,-1200 337,-1248 337,-1248 337,-1254 331,-1260 325,-1260"/>
<text xml:space="preserve" text-anchor="middle" x="229" y="-1219.7" font-family="Times,serif" font-size="14.00">Use RapidEye</text>
</a>
</g>
</g>
<!-- hang&#45;&gt;rapideye -->
<g id="edge4" class="edge">
<title>hang-&gt;rapideye</title>
<path fill="none" stroke="black" d="M229,-1295.7C229,-1287.93 229,-1279.57 229,-1271.43"/>
<polygon fill="black" stroke="black" points="232.5,-1271.66 229,-1261.66 225.5,-1271.66 232.5,-1271.66"/>
<text xml:space="preserve" text-anchor="middle" x="218.5" y="-1282.65" font-family="Times,serif" font-size="14.00">Yes</text>
</g>
<!-- check_straggler -->
<g id="googleyellow" class="node">
<title>check_straggler</title>
<g id="a_googleyellow"><a xlink:href="debugging_workflow.md#straggler" xlink:title=" " target="_parent">
<path fill="none" stroke="black" d="M585,-1260C585,-1260 393,-1260 393,-1260 387,-1260 381,-1254 381,-1248 381,-1248 381,-1200 381,-1200 381,-1194 387,-1188 393,-1188 393,-1188 585,-1188 585,-1188 591,-1188 597,-1194 597,-1200 597,-1200 597,-1248 597,-1248 597,-1254 591,-1260 585,-1260"/>
<text xml:space="preserve" text-anchor="middle" x="489" y="-1219.7" font-family="Times,serif" font-size="14.00">Check for straggler</text>
</a>
</g>
</g>
<!-- hang&#45;&gt;check_straggler -->
<g id="edge5" class="edge">
<title>hang-&gt;check_straggler</title>
<path fill="none" stroke="black" d="M276.69,-1311.56C308.85,-1298.45 352.43,-1280.68 391.74,-1264.65"/>
<polygon fill="black" stroke="black" points="392.83,-1267.99 400.77,-1260.97 390.19,-1261.51 392.83,-1267.99"/>
<text xml:space="preserve" text-anchor="middle" x="331.23" y="-1272.66" font-family="Times,serif" font-size="14.00">No</text>
</g>
<!-- has_straggler -->
<g id="googlegreen" class="node">
<title>has_straggler</title>
<g id="a_googlegreen"><a xlink:title=" " target="_blank">
<polygon fill="lightgrey" stroke="black" points="489,-1152 381,-1116 489,-1080 597,-1116 489,-1152"/>
<text xml:space="preserve" text-anchor="middle" x="489" y="-1111.7" font-family="Times,serif" font-size="14.00">Has straggler?</text>
</a>
</g>
</g>
<!-- check_straggler&#45;&gt;has_straggler -->
<g id="edge6" class="edge">
<title>check_straggler-&gt;has_straggler</title>
<path fill="none" stroke="black" d="M489,-1187.7C489,-1180.1 489,-1171.94 489,-1163.97"/>
<polygon fill="black" stroke="black" points="492.5,-1164.02 489,-1154.02 485.5,-1164.02 492.5,-1164.02"/>
</g>
<!-- outkast -->
<g id="googlered" class="node">
<title>outkast</title>
<g id="a_googlered"><a xlink:title=" " target="_blank">
<path fill="none" stroke="black" d="M612,-936C612,-936 420,-936 420,-936 414,-936 408,-930 408,-924 408,-924 408,-876 408,-876 408,-870 414,-864 420,-864 420,-864 612,-864 612,-864 618,-864 624,-870 624,-876 624,-876 624,-924 624,-924 624,-930 618,-936 612,-936"/>
<text xml:space="preserve" text-anchor="middle" x="516" y="-895.7" font-family="Times,serif" font-size="14.00">Remove bad host</text>
</a>
</g>
</g>
<!-- has_straggler&#45;&gt;outkast -->
<g id="edge7" class="edge">
<title>has_straggler-&gt;outkast</title>
<path fill="none" stroke="black" d="M493.28,-1081.09C497.83,-1045.05 505.05,-987.81 510.13,-947.56"/>
<polygon fill="black" stroke="black" points="513.59,-948.03 511.37,-937.67 506.65,-947.16 513.59,-948.03"/>
<text xml:space="preserve" text-anchor="middle" x="512.92" y="-1013.33" font-family="Times,serif" font-size="14.00">Yes</text>
</g>
<!-- has_dma_map -->
<g id="googlegreen" class="node">
<title>has_dma_map</title>
<g id="a_googlegreen"><a xlink:href="debugging_workflow.md#check_for_shortage_of_mapped_dma_buffers" xlink:title=" " target="_parent">
<polygon fill="lightgrey" stroke="black" points="380,-1044 272,-1008 380,-972 488,-1008 380,-1044"/>
<text xml:space="preserve" text-anchor="middle" x="380" y="-1003.7" font-family="Times,serif" font-size="14.00">Has DMAMap in the trace?</text>
</a>
</g>
</g>
<!-- has_straggler&#45;&gt;has_dma_map -->
<g id="edge11" class="edge">
<title>has_straggler-&gt;has_dma_map</title>
<path fill="none" stroke="black" d="M462.06,-1088.8C447.91,-1075.04 430.4,-1058.01 415.19,-1043.22"/>
<polygon fill="black" stroke="black" points="417.71,-1040.79 408.1,-1036.32 412.83,-1045.8 417.71,-1040.79"/>
<text xml:space="preserve" text-anchor="middle" x="426.23" y="-1066.68" font-family="Times,serif" font-size="14.00">No</text>
</g>
<!-- problem_persist -->
<g id="googlegreen" class="node">
<title>problem_persist</title>
<g id="a_googlegreen"><a xlink:title=" " target="_blank">
<polygon fill="lightgrey" stroke="black" points="565,-180 457,-144 565,-108 673,-144 565,-180"/>
<text xml:space="preserve" text-anchor="middle" x="565" y="-139.7" font-family="Times,serif" font-size="14.00">Does problem persist?</text>
</a>
</g>
</g>
<!-- problem_persist&#45;&gt;check_straggler -->
<g id="edge9" class="edge">
<title>problem_persist-&gt;check_straggler</title>
<path fill="none" stroke="black" d="M584.96,-173.62C610.64,-213.31 652,-288.21 652,-359 652,-1009 652,-1009 652,-1009 652,-1075.76 647.37,-1099.6 606,-1152 597.29,-1163.04 586.35,-1172.76 574.68,-1181.2"/>
<polygon fill="black" stroke="black" points="572.73,-1178.3 566.45,-1186.83 576.68,-1184.08 572.73,-1178.3"/>
<text xml:space="preserve" text-anchor="middle" x="641.5" y="-739.49" font-family="Times,serif" font-size="14.00">Yes</text>
</g>
<!-- happy_user -->
<g id="googlered" class="node">
<title>happy_user</title>
<g id="a_googlered"><a xlink:title=" " target="_blank">
<path fill="none" stroke="black" d="M661,-72C661,-72 469,-72 469,-72 463,-72 457,-66 457,-60 457,-60 457,-12 457,-12 457,-6 463,0 469,0 469,0 661,0 661,0 667,0 673,-6 673,-12 673,-12 673,-60 673,-60 673,-66 667,-72 661,-72"/>
<text xml:space="preserve" text-anchor="middle" x="565" y="-16.75" font-family="Times,serif" font-size="40.00">😊</text>
</a>
</g>
</g>
<!-- problem_persist&#45;&gt;happy_user -->
<g id="edge10" class="edge">
<title>problem_persist-&gt;happy_user</title>
<path fill="none" stroke="black" d="M565,-107.7C565,-99.93 565,-91.57 565,-83.43"/>
<polygon fill="black" stroke="black" points="568.5,-83.66 565,-73.66 561.5,-83.66 568.5,-83.66"/>
<text xml:space="preserve" text-anchor="middle" x="556.75" y="-94.65" font-family="Times,serif" font-size="14.00">No</text>
</g>
<!-- check_network_config -->
<g id="googleyellow" class="node">
<title>check_network_config</title>
<g id="a_googleyellow"><a xlink:title=" " target="_blank">
<path fill="none" stroke="black" d="M272,-504C272,-504 80,-504 80,-504 74,-504 68,-498 68,-492 68,-492 68,-444 68,-444 68,-438 74,-432 80,-432 80,-432 272,-432 272,-432 278,-432 284,-438 284,-444 284,-444 284,-492 284,-492 284,-498 278,-504 272,-504"/>
<text xml:space="preserve" text-anchor="middle" x="176" y="-463.7" font-family="Times,serif" font-size="14.00">Check network config</text>
</a>
</g>
</g>
<!-- is_network_config_optimal -->
<g id="googlegreen" class="node">
<title>is_network_config_optimal</title>
<g id="a_googlegreen"><a xlink:title=" " target="_blank">
<polygon fill="lightgrey" stroke="black" points="176,-396 68,-360 176,-324 284,-360 176,-396"/>
<text xml:space="preserve" text-anchor="middle" x="176" y="-355.7" font-family="Times,serif" font-size="14.00">is_network_config_optimal</text>
</a>
</g>
</g>
<!-- check_network_config&#45;&gt;is_network_config_optimal -->
<g id="edge23" class="edge">
<title>check_network_config-&gt;is_network_config_optimal</title>
<path fill="none" stroke="black" d="M176,-431.7C176,-424.1 176,-415.94 176,-407.97"/>
<polygon fill="black" stroke="black" points="179.5,-408.02 176,-398.02 172.5,-408.02 179.5,-408.02"/>
</g>
<!-- outkast&#45;&gt;problem_persist -->
<g id="edge8" class="edge">
<title>outkast-&gt;problem_persist</title>
<path fill="none" stroke="black" d="M545.23,-863.65C552.82,-852.81 560.2,-840.44 565,-828 588.07,-768.18 584,-749.11 584,-685 584,-685 584,-685 584,-359 584,-300.4 576.85,-233.14 571.32,-189.91"/>
<polygon fill="black" stroke="black" points="574.81,-189.61 570.04,-180.15 567.87,-190.52 574.81,-189.61"/>
</g>
<!-- increase_premap_buffer_size -->
<g id="googleyellow" class="node">
<title>increase_premap_buffer_size</title>
<g id="a_googleyellow"><a xlink:title=" " target="_blank">
<path fill="none" stroke="black" d="M544,-828C544,-828 352,-828 352,-828 346,-828 340,-822 340,-816 340,-816 340,-768 340,-768 340,-762 346,-756 352,-756 352,-756 544,-756 544,-756 550,-756 556,-762 556,-768 556,-768 556,-816 556,-816 556,-822 550,-828 544,-828"/>
<text xml:space="preserve" text-anchor="middle" x="448" y="-787.7" font-family="Times,serif" font-size="14.00">Increase premap buffer size</text>
</a>
</g>
</g>
<!-- increase_premap_buffer_size&#45;&gt;problem_persist -->
<g id="edge15" class="edge">
<title>increase_premap_buffer_size-&gt;problem_persist</title>
<path fill="none" stroke="black" d="M475.56,-755.77C504.55,-715.1 546,-645.04 546,-577 546,-577 546,-577 546,-359 546,-300.4 553.15,-233.14 558.68,-189.91"/>
<polygon fill="black" stroke="black" points="562.13,-190.52 559.96,-180.15 555.19,-189.61 562.13,-190.52"/>
</g>
<!-- get_hlo_dump -->
<g id="googleyellow" class="node">
<title>get_hlo_dump</title>
<g id="a_googleyellow"><a xlink:href="debugging_workflow.md#hlo_dump" xlink:title=" " target="_parent">
<path fill="none" stroke="black" d="M506,-612C506,-612 314,-612 314,-612 308,-612 302,-606 302,-600 302,-600 302,-552 302,-552 302,-546 308,-540 314,-540 314,-540 506,-540 506,-540 512,-540 518,-546 518,-552 518,-552 518,-600 518,-600 518,-606 512,-612 506,-612"/>
<text xml:space="preserve" text-anchor="middle" x="410" y="-571.7" font-family="Times,serif" font-size="14.00">Get HLO Dump</text>
</a>
</g>
</g>
<!-- get_hlo_dump&#45;&gt;consult_xla -->
<g id="edge20" class="edge">
<title>get_hlo_dump-&gt;consult_xla</title>
<path fill="none" stroke="black" d="M405.5,-539.62C398.06,-481.21 383.21,-364.61 374.93,-299.57"/>
<polygon fill="black" stroke="black" points="378.44,-299.42 373.7,-289.94 371.49,-300.3 378.44,-299.42"/>
</g>
<!-- update_network_config -->
<g id="googleryellow" class="node">
<title>update_network_config</title>
<g id="a_googleryellow"><a xlink:title=" " target="_blank">
<path fill="none" stroke="black" d="M231,-288C231,-288 39,-288 39,-288 33,-288 27,-282 27,-276 27,-276 27,-228 27,-228 27,-222 33,-216 39,-216 39,-216 231,-216 231,-216 237,-216 243,-222 243,-228 243,-228 243,-276 243,-276 243,-282 237,-288 231,-288"/>
<text xml:space="preserve" text-anchor="middle" x="135" y="-247.7" font-family="Times,serif" font-size="14.00">Update network config</text>
</a>
</g>
</g>
<!-- update_network_config&#45;&gt;problem_persist -->
<g id="edge24" class="edge">
<title>update_network_config-&gt;problem_persist</title>
<path fill="none" stroke="black" d="M243.29,-218.32C246.22,-217.53 249.13,-216.75 252,-216 332.07,-194.97 424.63,-174.43 488.21,-160.9"/>
<polygon fill="black" stroke="black" points="488.81,-164.35 497.86,-158.86 487.35,-157.51 488.81,-164.35"/>
</g>
<!-- has_dma_map&#45;&gt;increase_premap_buffer_size -->
<g id="edge12" class="edge">
<title>has_dma_map-&gt;increase_premap_buffer_size</title>
<path fill="none" stroke="black" d="M380.12,-971.78C381.28,-942.12 385.48,-899.05 399,-864 402.45,-855.06 407.23,-846.16 412.47,-837.83"/>
<polygon fill="black" stroke="black" points="415.31,-839.87 417.93,-829.61 409.48,-836 415.31,-839.87"/>
<text xml:space="preserve" text-anchor="middle" x="378.56" y="-902.37" font-family="Times,serif" font-size="14.00">Yes</text>
</g>
<!-- has_memcpy -->
<g id="googlegreen" class="node">
<title>has_memcpy</title>
<g id="a_googlegreen"><a xlink:href="debugging_workflow.md#check_for_memory_copies_during_network_transfers" xlink:title=" " target="_parent">
<polygon fill="lightgrey" stroke="black" points="244,-936 136,-900 244,-864 352,-900 244,-936"/>
<text xml:space="preserve" text-anchor="middle" x="244" y="-895.7" font-family="Times,serif" font-size="14.00">Has 'Memory Copy' in the trace?</text>
</a>
</g>
</g>
<!-- has_dma_map&#45;&gt;has_memcpy -->
<g id="edge13" class="edge">
<title>has_dma_map-&gt;has_memcpy</title>
<path fill="none" stroke="black" d="M348.11,-982.15C329.19,-967.4 305,-948.54 284.7,-932.73"/>
<polygon fill="black" stroke="black" points="286.92,-930.01 276.88,-926.63 282.62,-935.54 286.92,-930.01"/>
<text xml:space="preserve" text-anchor="middle" x="303.53" y="-958.53" font-family="Times,serif" font-size="14.00">No</text>
</g>
<!-- has_memcpy&#45;&gt;increase_premap_buffer_size -->
<g id="edge14" class="edge">
<title>has_memcpy-&gt;increase_premap_buffer_size</title>
<path fill="none" stroke="black" d="M285.28,-877.55C309.55,-864.94 341.02,-848.59 369.82,-833.63"/>
<polygon fill="black" stroke="black" points="371.27,-836.81 378.53,-829.1 368.04,-830.6 371.27,-836.81"/>
<text xml:space="preserve" text-anchor="middle" x="322.19" y="-839.62" font-family="Times,serif" font-size="14.00">Yes</text>
</g>
<!-- no_slack_collectives -->
<g id="googlegreen" class="node">
<title>no_slack_collectives</title>
<g id="a_googlegreen"><a xlink:href="debugging_workflow.md#collective_slack_too_small" xlink:title=" " target="_parent">
<polygon fill="lightgrey" stroke="black" points="214,-828 106,-792 214,-756 322,-792 214,-828"/>
<text xml:space="preserve" text-anchor="middle" x="214" y="-787.7" font-family="Times,serif" font-size="14.00">Collective with small slack?</text>
</a>
</g>
</g>
<!-- has_memcpy&#45;&gt;no_slack_collectives -->
<g id="edge16" class="edge">
<title>has_memcpy-&gt;no_slack_collectives</title>
<path fill="none" stroke="black" d="M234.84,-866.62C232.1,-856.96 229.06,-846.22 226.17,-836"/>
<polygon fill="black" stroke="black" points="229.61,-835.31 223.52,-826.64 222.88,-837.21 229.61,-835.31"/>
<text xml:space="preserve" text-anchor="middle" x="220.8" y="-850.87" font-family="Times,serif" font-size="14.00">No</text>
</g>
<!-- high_bandwidth_demand -->
<g id="googlegreen" class="node">
<title>high_bandwidth_demand</title>
<g id="a_googlegreen"><a xlink:href="debugging_workflow.md#high_network_bandwidth_demand" xlink:title=" " target="_parent">
<polygon fill="lightgrey" stroke="black" points="214,-720 106,-684 214,-648 322,-684 214,-720"/>
<text xml:space="preserve" text-anchor="middle" x="214" y="-679.7" font-family="Times,serif" font-size="14.00">High network bandwidth demand?</text>
</a>
</g>
</g>
<!-- high_bandwidth_demand&#45;&gt;get_hlo_dump -->
<g id="edge19" class="edge">
<title>high_bandwidth_demand-&gt;get_hlo_dump</title>
<path fill="none" stroke="black" d="M254.61,-661.04C277.68,-648.56 307.3,-632.54 334.48,-617.84"/>
<polygon fill="black" stroke="black" points="335.9,-621.05 343.04,-613.22 332.57,-614.89 335.9,-621.05"/>
<text xml:space="preserve" text-anchor="middle" x="288.95" y="-623.49" font-family="Times,serif" font-size="14.00">Yes</text>
</g>
<!-- high_network_latency -->
<g id="googlegreen" class="node">
<title>high_network_latency</title>
<g id="a_googlegreen"><a xlink:href="debugging_workflow.md#high_network_latency" xlink:title=" " target="_parent">
<polygon fill="lightgrey" stroke="black" points="176,-612 68,-576 176,-540 284,-576 176,-612"/>
<text xml:space="preserve" text-anchor="middle" x="176" y="-571.7" font-family="Times,serif" font-size="14.00">High network latency?</text>
</a>
</g>
</g>
<!-- high_bandwidth_demand&#45;&gt;high_network_latency -->
<g id="edge21" class="edge">
<title>high_bandwidth_demand-&gt;high_network_latency</title>
<path fill="none" stroke="black" d="M202.7,-651.48C199.06,-641.32 194.96,-629.89 191.09,-619.1"/>
<polygon fill="black" stroke="black" points="194.5,-618.24 187.83,-610 187.91,-620.6 194.5,-618.24"/>
<text xml:space="preserve" text-anchor="middle" x="186.7" y="-634.56" font-family="Times,serif" font-size="14.00">No</text>
</g>
<!-- high_network_latency&#45;&gt;check_network_config -->
<g id="edge22" class="edge">
<title>high_network_latency-&gt;check_network_config</title>
<path fill="none" stroke="black" d="M176,-539.7C176,-531.93 176,-523.57 176,-515.43"/>
<polygon fill="black" stroke="black" points="179.5,-515.66 176,-505.66 172.5,-515.66 179.5,-515.66"/>
<text xml:space="preserve" text-anchor="middle" x="165.5" y="-526.65" font-family="Times,serif" font-size="14.00">Yes</text>
</g>
<!-- high_network_latency&#45;&gt;consult_xla -->
<g id="edge25" class="edge">
<title>high_network_latency-&gt;consult_xla</title>
<path fill="none" stroke="black" d="M224.3,-555.69C248.05,-543.96 275.31,-526.81 293,-504 340.05,-443.33 358.09,-353.68 364.93,-299.56"/>
<polygon fill="black" stroke="black" points="368.39,-300.08 366.08,-289.74 361.44,-299.27 368.39,-300.08"/>
<text xml:space="preserve" text-anchor="middle" x="322.68" y="-441.05" font-family="Times,serif" font-size="14.00">No</text>
</g>
<!-- no_slack_collectives&#45;&gt;get_hlo_dump -->
<g id="edge17" class="edge">
<title>no_slack_collectives-&gt;get_hlo_dump</title>
<path fill="none" stroke="black" d="M259.25,-770.64C282.73,-758.44 310.68,-741.17 331,-720 357.97,-691.91 378.74,-652.81 392.2,-622.6"/>
<polygon fill="black" stroke="black" points="395.29,-624.27 396.05,-613.7 388.86,-621.49 395.29,-624.27"/>
<text xml:space="preserve" text-anchor="middle" x="333.85" y="-709.34" font-family="Times,serif" font-size="14.00">Yes</text>
</g>
<!-- no_slack_collectives&#45;&gt;high_bandwidth_demand -->
<g id="edge18" class="edge">
<title>no_slack_collectives-&gt;high_bandwidth_demand</title>
<path fill="none" stroke="black" d="M214,-755.7C214,-748.1 214,-739.94 214,-731.97"/>
<polygon fill="black" stroke="black" points="217.5,-732.02 214,-722.02 210.5,-732.02 217.5,-732.02"/>
<text xml:space="preserve" text-anchor="middle" x="205.75" y="-743.05" font-family="Times,serif" font-size="14.00">No</text>
</g>
<!-- is_network_config_optimal&#45;&gt;consult_xla -->
<g id="edge27" class="edge">
<title>is_network_config_optimal-&gt;consult_xla</title>
<path fill="none" stroke="black" d="M215.99,-337.04C238.61,-324.62 267.62,-308.68 294.28,-294.04"/>
<polygon fill="black" stroke="black" points="295.96,-297.11 303.04,-289.22 292.59,-290.97 295.96,-297.11"/>
<text xml:space="preserve" text-anchor="middle" x="249.75" y="-299.43" font-family="Times,serif" font-size="14.00">Yes</text>
</g>
<!-- is_network_config_optimal&#45;&gt;update_network_config -->
<g id="edge26" class="edge">
<title>is_network_config_optimal-&gt;update_network_config</title>
<path fill="none" stroke="black" d="M163.92,-327.77C160.45,-318.81 156.6,-308.84 152.86,-299.18"/>
<polygon fill="black" stroke="black" points="156.14,-297.95 149.27,-289.88 149.61,-300.47 156.14,-297.95"/>
<text xml:space="preserve" text-anchor="middle" x="147.96" y="-294.53" font-family="Times,serif" font-size="14.00">No</text>
</g>
</g>
</svg>
Side by Side

After

Width:  |  Height:  |  Size: 22 KiB

third_party/xla/docs/megascale/images/bandwidth_bound.png
Vendored
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 79 KiB

third_party/xla/docs/megascale/images/map_dma_buffer_example_trace.png
Vendored
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 145 KiB

third_party/xla/docs/megascale/images/megascale_stats_zero_slack_example.png
Vendored
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 114 KiB

third_party/xla/docs/megascale/images/memory_copy_example_trace.png
Vendored
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 118 KiB

third_party/xla/docs/megascale/images/tail_network_latency.png
Vendored
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 248 KiB

third_party/xla/docs/megascale/overview.md
Vendored
+54
View File
@@ -0,0 +1,54 @@
<!-- linter style off -->
# Megascale XLA
MegascaleXLA is a compiler + runtime system that powers large-scale TPU
training. It implements collective communication primitives that allow multiple
TPU slices to communicate, which allows running training jobs that span beyond
the limits of a single ICI domain.
The [debugging guide](./debugging_workflow.md) discusses how to identify and
diagnose sources of performance issues such as slowness, hangs or errors in a
multi-slice job driven by Megascale.
## Terminology
- **Slice**
- A slice is a collection of chips all located inside the same TPU Pod
connected by high-speed inter chip interconnects (ICI). Slices are
described in terms of chips or TensorCores, depending on the TPU
version.
- Multislice is a group of slices, extending TPU connectivity beyond the
inter-chip interconnect (ICI) connections and leveraging the data-center
network (DCN) for transmitting data beyond a slice. Data within each slice
is still transmitted by ICI. Using this hybrid connectivity, Multislice
enables parallelism across slices and lets you use a greater number of TPU
cores for a single job than what a single slice can accommodate.
- TPUs can be used to run a job either on a single slice or multiple
slices.
- **RapidEye**
- RapidEye is a system that aims to provide global ML debugging
infrastructure to quickly identify and root cause issues caused by bad
hardware or software bugs. It monitors MegaScale jobs to automatically
detect, analyze, and classify hang events. The process involves
collecting data from all job workers, coordinating responses when a hang
occurs, and generating a summary digest file for each event.
- RapidEye is enabled by default for all multislice workloads. The
diagnosis can be found under the Pathways resource manager or MXLA
coordinator (slice0 task0 for multi-controller JAX workload). RapidEye
is also used for auto-removing bad TPUs and NICs based on fleetwide
rapideye data.
- **Megascale Collective**
- XLA collectives are supported via the Megascale XLA (MXLA) primitives
which are not directly usable by the end user. At the time of writing
this the MXLA primitives include collectives include AllGather,
AllReduce, AllToAll, ReduceScatter and OneToOne. The reduction
operations that are currently supported include summation, max/min and
product.
<!-- linter style on -->