JAX / Pathways

Recovery Mode Simulation & Comparison

Real-time simulation comparing McJAX vs Snapshot Replica Resize vs Snapshot Pause Resume recovery from TPU interruptions.

Cluster Healthy (2 Slices)
0.0s elapsed
Simulation Parameters
s
s
s
s
ms
st
st
x
McJAX
DUAL-SLICE
STEP
0
LOSS
6.850
Goodput
0.0%
Active dual-slice execution
Snapshot Replica Resize
DUAL-SLICE
STEP
0
LOSS
6.850
Goodput
0.0%
Active dual-slice execution
Snapshot Pause Resume
DUAL-SLICE
STEP
0
LOSS
6.850
Goodput
0.0%
Active dual-slice execution

Cluster Slice Topology & Simulated Interruption Injector

Cluster Healthy (2 Slices)

Inject interruption on Workload Slice 0 or Workload Slice 1 to test how Snapshot Replica Resize, Snapshot Pause Resume, and McJAX recover when replacement TPU slices are allocated.

Logical Workload Slices (2-Slice Mesh) Active execution mesh
WORKLOAD SLICE 0
Running on TPU Slice 0
Active Slice
WORKLOAD SLICE 1
Running on TPU Slice 1
Active Slice
Physical Cluster TPU Slices (Hardware Pool) Reconnecting allocates a random inactive slice to resume the workload
TPU SLICE 0
Active Slice
TPU SLICE 1
Active Slice
TPU SLICE 2
Spare Slice
TPU SLICE 3
Spare Slice

Training Loss

McJAX Snapshot Replica Resize Snapshot Pause Resume
X: Elapsed Time (s) • Y: Training Loss McJAX: 6.850 | Resize: 6.850 | Pause: 6.850

Estimated Goodput

McJAX Snapshot Replica Resize Snapshot Pause Resume
X: Elapsed Time (s) • Y: Cumulative Goodput (%) McJAX: 0.0% | Resize: 0.0% | Pause: 0.0%