[simrankaurb]

Summary

Ports the pinned memory Host-to-Device (H2D) and Device-to-Host (D2H) transfer benchmark (PinnedMemoryTransfer, SimpleTransfer, and PipelinedTransfer) and associated configurations/scripts from the tpu7x-h2dd2h branch to chs, along with jax.local_devices() sharding for single-host/node-local execution and numactl support in Dockerfile.

Code References / Source History

This PR ports the changes introduced on tpu7x-h2dd2h across the following commits:

• 9b4e8de ("Try pinned memory" by @linamy85): Added TransferStrategy hierarchy (SimpleTransfer, PipelinedTransfer, PinnedMemoryTransfer using NamedSharding(..., memory_kind="pinned_host") for H2D and parallel shard jax.device_get for D2H) in Ironwood/src/benchmark_host_device.py.
• f924a7e ("Optimize H2D/D2H transfer pipelines and add comprehensive benchmark configs" by @linamy85): Updated Ironwood/src/benchmark_host_device.py and Ironwood/configs/host_device/host_device.yaml, and added Ironwood/configs/host_device/comprehensive_experiments.yaml and Ironwood/configs/host_device/comprehensive_8dev_experiments.yaml.
• 7439b2a ("Add benchmark guide and run script" by @linamy85): Updated Ironwood/scripts/run_host_device_benchmark.sh (TPU_PREMAPPED_BUFFER_SIZE=34359738368, --numactl_interleave), and added Ironwood/guides/host_device_benchmark/README.md and Ironwood/guides/host_device_benchmark/run_comprehensive.sh.
• 38ec530 ("Allow sweeping dtype in host_device benchmarks" by @linamy85): Added configurable dtype support in Ironwood/src/benchmark_host_device.py.

Additional chs Changes in This PR

• Dockerfile:

  • Installs numactl so numactl --cpunodebind=0 --membind=0 and numactl --interleave=all runs work inside the container.

• requirements.txt:

  • Pins jax[tpu]==0.11.1 and jaxlib==0.11.1 while keeping libtpu==0.0.48 (matching main pyproject.toml), avoiding the jax 0.11.2 jax.device_get D2H and multi-device H2D performance regressions while retaining libtpu 0.0.48 for single_host_tpu_init.

  Tests: https://docs.google.com/spreadsheets/d/1VQqHS1UEvbgvyZ6ddnl6-Gy-79fwYNTkSSYopMsQnC8/edit?resourcekey=0-sFolkGOQhiSUVIthaPPszg&gid=920016885#gid=920016885

[gemini-code-assist]

Code Review

This pull request introduces a flexible benchmarking framework for Host-to-Device (H2D) and Device-to-Host (D2H) transfers, supporting multiple strategies such as simple, pipelined, and pinned memory transfers. It includes updated configuration files for comprehensive experiments and new scripts to automate execution with optional NUMA bindings. The review comments highlight critical improvements for multi-device data distribution, robustness in variance calculations, and safer execution of shell scripts when dependencies like numactl are missing.

I am having trouble creating individual review comments. Click here to see my feedback.

Ironwood/src/benchmark_host_device.py (126-132)

In a multi-device benchmark configuration, falling back to transferring the entire array to self.target_devices[0] when self.chunks_per_dev == 1 ignores the other devices. This makes the benchmark results incorrect and inconsistent for small data sizes. Instead, we should split the data across all target devices (1 chunk per device) to ensure a fair and correct multi-device transfer benchmark.

        else:
            chunks = np.array_split(host_data, self.num_devices, axis=0)
            tensors_on_device = []

            t0 = time.perf_counter()
            for idx, chunk in enumerate(chunks):
                dev = self.target_devices[idx]
                tensors_on_device.append(jax.device_put(chunk, dev))

            for device_tensor in tensors_on_device:
                device_tensor.block_until_ready()
            t1 = time.perf_counter()

            self.h2d_perf.append((t1 - t0) * 1000)
            del chunks
            return tensors_on_device

Ironwood/src/benchmark_host_device.py (101-105)

Since we are updating the fallback behavior to split across all devices instead of transferring only to a single device, we should update the warning message to accurately reflect this fallback strategy.

        if self.chunks_per_dev == 1:
            print(
                "Warning: Data size is not larger than target chunk size, "
                "falling back to 1 chunk per device."
            )

Ironwood/src/benchmark_host_device.py (317-320)

Calculating sample variance with ddof=1 when len(bw_array) <= 1 will result in a division by zero, causing a RuntimeWarning and returning nan. We should check the length of bw_array before calculating the variance to make the code more robust.

        bw_array = np.array(bw_list)
        if len(bw_array) > 1:
            sample_variance = np.var(bw_array, ddof=1)
        else:
            sample_variance = 0.0

Ironwood/guides/host_device_benchmark/run_comprehensive.sh (12-13)

If numactl is not installed on the system, the script will fail and exit immediately due to set -e. It is safer to check if numactl is available before executing it, falling back to running without it if necessary.

echo "Running 8-device benchmark with numactl interleaving..."
if command -v numactl &> /dev/null; then
    numactl --interleave=all python3 src/run_benchmark.py --config configs/host_device/comprehensive_8dev_experiments.yaml
else
    echo "Warning: numactl not found. Running 8-device benchmark without interleaving."
    python3 src/run_benchmark.py --config configs/host_device/comprehensive_8dev_experiments.yaml
fi

[linamy85]

Thanks Simran! The pinned H2D and D2H timing both look correct: allocation is outside timing, all 8 shards transfer concurrently, and block_until_ready is inside the timer. Main asks: assert the device count and fix the jax/libtpu pin. Since we're moving this to main soon, I've skipped design comments.