TPU v6e ICI 带宽实测——单条 Link 93 GB/s，标称 800 GB/s 的真相

import jax
import jax.numpy as jnp
from jax.sharding import Mesh, PartitionSpec as P, NamedSharding
import numpy as np
import os
import time

# Dump HLO/LLO
os.environ['XLA_FLAGS'] = '--xla_dump_to=/tmp/xla_dump --xla_dump_hlo_as_text=true'

devices = jax.devices()
print(f"Total devices: {len(devices)}")
for d in devices:
    print(f"  {d}")

mesh_a = Mesh(np.array(devices[:2]), ('x',))
mesh_b = Mesh(np.array(devices[2:]), ('x',))
print(f"\nMesh A devices: {mesh_a.devices}")
print(f"Mesh B devices: {mesh_b.devices}")

# === Test 1: device_put cross-mesh (outside jit) ===
print("\n" + "="*60)
print("TEST 1: device_put cross-mesh transfer (outside jit)")
print("="*60)

sharding_a = NamedSharding(mesh_a, P('x'))
data_a = jax.device_put(jnp.ones((128,)), sharding_a)
print(f"Data on mesh_a: devices={data_a.devices()}, shape={data_a.shape}")

sharding_b = NamedSharding(mesh_b, P('x'))
t0 = time.time()
data_b = jax.device_put(data_a, sharding_b)
data_b.block_until_ready()
t1 = time.time()
print(f"Data on mesh_b: devices={data_b.devices()}, shape={data_b.shape}")
print(f"Transfer time: {(t1-t0)*1000:.3f} ms")

# === Test 2: Compute on mesh_a then transfer result to mesh_b ===
print("\n" + "="*60)
print("TEST 2: Compute on mesh_a, then device_put to mesh_b")
print("="*60)

@jax.jit
def compute_on_a(x):
    return x * 2.0 + 1.0

with mesh_a:
    result_a = compute_on_a(data_a)
    result_a.block_until_ready()
    print(f"Result on mesh_a: devices={result_a.devices()}")

result_b = jax.device_put(result_a, sharding_b)
result_b.block_until_ready()
print(f"Result on mesh_b: devices={result_b.devices()}")

# === Dump HLO analysis ===
print("\n" + "="*60)
print("XLA DUMP ANALYSIS")
print("="*60)

print("\n--- After optimizations HLO (all modules) ---")
os.system("for f in /tmp/xla_dump/*.after_optimizations.txt; do "
          "echo '\\n===== '$f' ====='; cat $f; done 2>/dev/null")

print("\n--- Search for send/recv/copy-to-host/infeed/outfeed ---")
os.system("grep -rni 'send\\|recv\\|infeed\\|outfeed\\|host\\|transfer-to-host' "
          "/tmp/xla_dump/*.after_optimizations.txt 2>/dev/null | head -30 "
          "|| echo 'No send/recv/host operations found in HLO'")

# === Check JAX transfer guard ===
print("\n" + "="*60)
print("TEST 3: JAX transfer guard test")
print("="*60)
try:
    jax.config.update("jax_transfer_guard", "disallow")
    data_b3 = jax.device_put(data_a, sharding_b)
    data_b3.block_until_ready()
    print("Transfer succeeded even with disallow — device_put is explicit, "
          "transfer_guard only blocks implicit transfers")
except Exception as e:
    print(f"Transfer BLOCKED: {e}")
finally:
    jax.config.update("jax_transfer_guard", "allow")

import jax
import jax.numpy as jnp
from jax.sharding import Mesh, PartitionSpec as P, NamedSharding
import numpy as np
import os

devices = jax.devices()
mesh_a = Mesh(np.array(devices[:2]), ('x',))
mesh_b = Mesh(np.array(devices[2:]), ('x',))
sharding_a = NamedSharding(mesh_a, P('x'))
sharding_b = NamedSharding(mesh_b, P('x'))

# Warmup
data_a = jax.device_put(jnp.ones((1024,)), sharding_a)
data_b = jax.device_put(data_a, sharding_b)
data_b.block_until_ready()

# Profile
TRACE_DIR = "/tmp/tpu_trace"
jax.profiler.start_trace(TRACE_DIR)
for i in range(10):
    data_a = jax.device_put(jnp.ones((4096,)), sharding_a)
    data_a.block_until_ready()
    data_b = jax.device_put(data_a, sharding_b)
    data_b.block_until_ready()
jax.profiler.stop_trace()

print(f"Trace saved to {TRACE_DIR}")
os.system(f"find {TRACE_DIR} -type f")

# 查看方式:
# 1. xprof: xprof /path/to/trace -p 8791, 浏览器打开 localhost:8791
# 2. Perfetto: 打开 https://ui.perfetto.dev, 拖入 .trace.json.gz

import jax
import jax.numpy as jnp
from jax.sharding import Mesh, PartitionSpec as P, NamedSharding
import numpy as np
import time

devices = jax.devices()
mesh_a = Mesh(np.array(devices[:2]), ('x',))
mesh_b = Mesh(np.array(devices[2:]), ('x',))
sharding_a = NamedSharding(mesh_a, P('x'))
sharding_b = NamedSharding(mesh_b, P('x'))

# Warmup
data = jax.device_put(jnp.ones((1024,)), sharding_a)
_ = jax.device_put(data, sharding_b); _.block_until_ready()

print(f"\n{'Size':>12}  {'Per Shard':>10}  {'Time (ms)':>10}  {'BW (GB/s)':>10}")
print("-" * 55)

sizes = [1024, 4096, 16384, 65536, 262144, 1048576, 4194304,
         16777216, 67108864, 268435456]

for nelems in sizes:
    total_bytes = nelems * 4
    shard_bytes = total_bytes // 2
    data_a = jax.device_put(jnp.ones((nelems,)), sharding_a)
    data_a.block_until_ready()

    # Warmup
    for _ in range(3):
        tmp = jax.device_put(data_a, sharding_b); tmp.block_until_ready()

    # Timed
    n_iters = max(1, min(100, 500_000_000 // total_bytes))
    times = []
    for _ in range(n_iters):
        t0 = time.perf_counter()
        data_b = jax.device_put(data_a, sharding_b)
        data_b.block_until_ready()
        t1 = time.perf_counter()
        times.append(t1 - t0)

    med_time = np.median(times)
    bw_gbs = total_bytes / med_time / 1e9
    label = f"{total_bytes/1024:.0f}KB" if total_bytes < 1048576 \
            else f"{total_bytes/1048576:.0f}MB"
    shard_label = f"{shard_bytes/1024:.0f}KB" if shard_bytes < 1048576 \
                  else f"{shard_bytes/1048576:.0f}MB"
    print(f"{label:>12}  {shard_label:>10}  {med_time*1000:>10.3f}  {bw_gbs:>10.2f}")

import jax
import jax.numpy as jnp
import numpy as np
import time
import concurrent.futures

devices = jax.devices()
SIZE_MB = 4096
nelems = SIZE_MB * 1024 * 1024 // 4
total_bytes = nelems * 4

def bench_single(src, dst, label, n_warmup=5, n_iters=20):
    data = jax.device_put(jnp.ones((nelems,)), devices[src])
    data.block_until_ready()
    for _ in range(n_warmup):
        t = jax.device_put(data, devices[dst]); t.block_until_ready(); del t
    times = []
    for _ in range(n_iters):
        t0 = time.perf_counter()
        t = jax.device_put(data, devices[dst]); t.block_until_ready()
        t1 = time.perf_counter()
        times.append(t1 - t0); del t
    med = np.median(times)
    bw = total_bytes / med / 1e9
    print(f"  {label}: {med*1000:.2f} ms, {bw:.1f} GB/s")
    del data
    return bw

# Baselines
print("Single link baselines:")
bw_01 = bench_single(0, 1, "chip0→chip1 (x-link)")
bw_02 = bench_single(0, 2, "chip0→chip2 (y-link)")

# Simultaneous
print("\nchip0 → chip1 AND chip0 → chip2 simultaneously:")
data_0 = jax.device_put(jnp.ones((nelems,)), devices[0])
data_0.block_until_ready()

def transfer_to(dst_idx):
    times = []
    for _ in range(20):
        t0 = time.perf_counter()
        t = jax.device_put(data_0, devices[dst_idx]); t.block_until_ready()
        t1 = time.perf_counter()
        times.append(t1 - t0); del t
    return np.median(times)

with concurrent.futures.ThreadPoolExecutor(max_workers=2) as executor:
    f1 = executor.submit(transfer_to, 1)
    f2 = executor.submit(transfer_to, 2)
    time_01, time_02 = f1.result(), f2.result()

bw_01_sim = total_bytes / time_01 / 1e9
bw_02_sim = total_bytes / time_02 / 1e9
print(f"  chip0→chip1: {bw_01_sim:.1f} GB/s")
print(f"  chip0→chip2: {bw_02_sim:.1f} GB/s")
print(f"  Total: {bw_01_sim + bw_02_sim:.1f} GB/s")
ratio = ((bw_01_sim + bw_02_sim) / 2) / bw_01
print(f"  Ratio: {ratio:.2f}x (1.0 = fully independent)")

# All links simultaneously
print("\nAll x-links simultaneously (0→1, 1→0, 2→3, 3→2):")
data = [jax.device_put(jnp.ones((nelems,)), devices[i]) for i in range(4)]
for d in data: d.block_until_ready()

pairs = [(0,1),(1,0),(2,3),(3,2)]
def do_transfer(src_dst):
    s, d = src_dst
    times = []
    for _ in range(20):
        t0 = time.perf_counter()
        t = jax.device_put(data[s], devices[d]); t.block_until_ready()
        t1 = time.perf_counter()
        times.append(t1 - t0); del t
    return np.median(times)

with concurrent.futures.ThreadPoolExecutor(max_workers=4) as executor:
    futures = {executor.submit(do_transfer, p): p for p in pairs}
    for f in concurrent.futures.as_completed(futures):
        p = futures[f]
        bw = total_bytes / f.result() / 1e9
        print(f"  chip{p[0]}→chip{p[1]}: {bw:.1f} GB/s")

import jax
import jax.numpy as jnp
from jax.sharding import Mesh, PartitionSpec as P, NamedSharding
from functools import partial
import numpy as np
import time
import concurrent.futures

try:
    from jax.shard_map import shard_map
except ImportError:
    from jax.experimental.shard_map import shard_map

devices = jax.devices()

def bench(fn, data, n_warmup=10, n_iters=30):
    for _ in range(n_warmup):
        out = fn(data); out.block_until_ready()
    times = []
    for _ in range(n_iters):
        t0 = time.perf_counter()
        out = fn(data); out.block_until_ready()
        t1 = time.perf_counter()
        times.append(t1 - t0); del out
    return np.median(times)

sizes_mb = [256, 512, 1024, 2048, 4096, 8192]

# --- psum on 2 devices (single link) ---
for dev_ids, name in [([0,1], "x-link"), ([0,2], "y-link")]:
    mesh2 = Mesh(np.array([devices[i] for i in dev_ids]), ('x',))
    sharded = NamedSharding(mesh2, P('x'))

    @partial(shard_map, mesh=mesh2, in_specs=P('x'), out_specs=P('x'))
    def do_psum(x):
        return jax.lax.psum(x, 'x')
    do_psum_jit = jax.jit(do_psum)

    print(f"\npsum 2-device ({name}):")
    for size_mb in sizes_mb:
        nelems = size_mb * 1024 * 1024 // 4
        total_bytes = nelems * 4
        data = jax.device_put(jnp.ones((nelems,)), sharded)
        data.block_until_ready()
        med = bench(do_psum_jit, data)
        print(f"  {size_mb:>6}MB: {med*1000:.2f}ms, "
              f"{total_bytes/med/1e9:.1f} GB/s")

# --- psum on 4 devices (all links) ---
mesh4 = Mesh(np.array(devices), ('x',))
sharded4 = NamedSharding(mesh4, P('x'))

@jax.jit
@partial(shard_map, mesh=mesh4, in_specs=P('x'), out_specs=P('x'))
def psum4(x):
    return jax.lax.psum(x, 'x')

print(f"\npsum 4-device (all links):")
for size_mb in sizes_mb:
    nelems = size_mb * 1024 * 1024 // 4
    total_bytes = nelems * 4
    data = jax.device_put(jnp.ones((nelems,)), sharded4)
    data.block_until_ready()
    med = bench(psum4, data)
    print(f"  {size_mb:>6}MB: {med*1000:.2f}ms, "
          f"{total_bytes/med/1e9:.1f} GB/s")

# --- device_put bidirectional ---
print(f"\ndevice_put bidirectional (0↔1):")
for size_mb in [1024, 4096, 8192]:
    nelems = size_mb * 1024 * 1024 // 4
    total_bytes = nelems * 4
    d0 = jax.device_put(jnp.ones((nelems,)), devices[0])
    d1 = jax.device_put(jnp.ones((nelems,)), devices[1])
    d0.block_until_ready(); d1.block_until_ready()

    def xfer(src_data, dst_dev):
        for _ in range(5):
            t = jax.device_put(src_data, dst_dev)
            t.block_until_ready()
        times = []
        for _ in range(20):
            t0 = time.perf_counter()
            t = jax.device_put(src_data, dst_dev)
            t.block_until_ready()
            t1 = time.perf_counter()
            times.append(t1 - t0)
        return np.median(times)

    with concurrent.futures.ThreadPoolExecutor(2) as ex:
        f01 = ex.submit(xfer, d0, devices[1])
        f10 = ex.submit(xfer, d1, devices[0])
        t01, t10 = f01.result(), f10.result()
    bw01 = total_bytes / t01 / 1e9
    bw10 = total_bytes / t10 / 1e9
    print(f"  {size_mb:>6}MB: 0→1={bw01:.1f}, 1→0={bw10:.1f}, "
          f"bidir={bw01+bw10:.1f} GB/s")

import jax
import jax.numpy as jnp
import numpy as np
import time

devices = jax.devices()
sizes_mb = [64, 128, 256, 512, 1024, 2048]

# Host → Device
print("Host → Device:")
for size_mb in sizes_mb:
    nelems = size_mb * 1024 * 1024 // 4
    total_bytes = nelems * 4
    host_data = np.ones((nelems,), dtype=np.float32)
    for _ in range(3):
        tmp = jax.device_put(host_data, devices[0])
        tmp.block_until_ready()
    times = []
    for _ in range(10):
        t0 = time.perf_counter()
        d = jax.device_put(host_data, devices[0])
        d.block_until_ready()
        t1 = time.perf_counter()
        times.append(t1 - t0)
    med = np.median(times)
    print(f"  {size_mb:>6}MB: {med*1000:.2f}ms, "
          f"{total_bytes/med/1e9:.1f} GB/s")

# Device → Host (forced copy)
print("\nDevice → Host:")
for size_mb in sizes_mb:
    nelems = size_mb * 1024 * 1024 // 4
    total_bytes = nelems * 4
    dev_data = jax.device_put(jnp.ones((nelems,)), devices[0])
    dev_data.block_until_ready()
    for _ in range(3):
        host_buf = np.empty((nelems,), dtype=np.float32)
        np.copyto(host_buf, dev_data)
    times = []
    for _ in range(10):
        host_buf = np.empty((nelems,), dtype=np.float32)
        t0 = time.perf_counter()
        np.copyto(host_buf, dev_data)
        t1 = time.perf_counter()
        times.append(t1 - t0)
    med = np.median(times)
    print(f"  {size_mb:>6}MB: {med*1000:.2f}ms, "
          f"{total_bytes/med/1e9:.1f} GB/s")