diff --git a/docs/examples/jax/attention.out b/docs/examples/jax/attention.out index a69cdef4e67..64afccf74e1 100644 --- a/docs/examples/jax/attention.out +++ b/docs/examples/jax/attention.out @@ -1,9 +1,9 @@ # SINGLE_GPU_OUTPUT_START Native JAX bf16 GQA + SWA: -Mean time: 5.109810829162598 ms +Mean time: 51.31755525362678 ms TE DotProductAttention GQA + SWA: -Mean time: 0.09856224060058594 ms +Mean time: 1.1168592551257461 ms # SINGLE_GPU_OUTPUT_END # MLA_OUTPUT_START diff --git a/docs/examples/jax/attention_context_parallel.out b/docs/examples/jax/attention_context_parallel.out index 100044be918..e7f10f3b9f1 100644 --- a/docs/examples/jax/attention_context_parallel.out +++ b/docs/examples/jax/attention_context_parallel.out @@ -1,16 +1,16 @@ # SINGLE_GPU_OUTPUT_START Single-GPU THD GQA + SWA: -Mean time: 126.68747901916504 ms +Mean time: 124.06466079410166 ms # SINGLE_GPU_OUTPUT_END # RING_OUTPUT_START THD CP Ring stripe_size=1: -Mean time: 57.16729164123535 ms -Speedup vs single GPU: 2.22x +Mean time: 52.09881339687854 ms +Speedup vs single GPU: 2.38x # RING_OUTPUT_END # AG_OUTPUT_START THD CP AllGather stripe_size=512: -Mean time: 53.79219055175781 ms -Speedup vs single GPU: 2.36x +Mean time: 53.20774579886347 ms +Speedup vs single GPU: 2.33x # AG_OUTPUT_END diff --git a/docs/examples/jax/attention_context_parallel.py b/docs/examples/jax/attention_context_parallel.py index c41aa6c425f..82985197d4c 100644 --- a/docs/examples/jax/attention_context_parallel.py +++ b/docs/examples/jax/attention_context_parallel.py @@ -346,11 +346,11 @@ def run_single_gpu_bench(): result = grad_fn((q, k, v), sequence_descriptor, dout) jax.block_until_ready(result) - start = time.time() + start = time.perf_counter() for _ in range(timing_iters): result = grad_fn((q, k, v), sequence_descriptor, dout) jax.block_until_ready(result) - mean_ms = (time.time() - start) * 1000 / timing_iters + mean_ms = (time.perf_counter() - start) * 1000 / timing_iters print(f"Mean time: {mean_ms} ms") return mean_ms @@ -374,7 +374,7 @@ def run_context_parallel_bench( ) jax.block_until_ready(result) - start = time.time() + start = time.perf_counter() for _ in range(timing_iters): result = grad_fn( sharded["qkv"], @@ -382,7 +382,7 @@ def run_context_parallel_bench( sharded["dout"], ) jax.block_until_ready(result) - end = time.time() + end = time.perf_counter() mean_ms = (end - start) * 1000 / timing_iters print(f"Mean time: {mean_ms} ms") diff --git a/docs/examples/jax/attention_context_parallel.rst b/docs/examples/jax/attention_context_parallel.rst index f701fe22af6..e6e38aed50d 100644 --- a/docs/examples/jax/attention_context_parallel.rst +++ b/docs/examples/jax/attention_context_parallel.rst @@ -162,7 +162,7 @@ while keeping those sharding controls visible. :start-after: # AG_OUTPUT_START :end-before: # AG_OUTPUT_END -On four GB200s, Ring is roughly **2.22x faster** and AllGather roughly **2.36x +On four GB200s, Ring is roughly **2.38x faster** and AllGather roughly **2.33x faster** than the equivalent single-GPU fused-attention forward+backward pass. These results are specific to this workload and system. Applications with long segments and wide attention windows generally have more attention computation diff --git a/docs/examples/jax/attention_single_gpu.rst b/docs/examples/jax/attention_single_gpu.rst index 2c9f5dd8c80..d07b7b5443d 100644 --- a/docs/examples/jax/attention_single_gpu.rst +++ b/docs/examples/jax/attention_single_gpu.rst @@ -80,7 +80,7 @@ implementations. :start-after: # SINGLE_GPU_OUTPUT_START :end-before: # SINGLE_GPU_OUTPUT_END -On a single GB200, this run is roughly **52x faster** for the fwd+bwd of this +On a single GB200, this run is roughly **46x faster** for the fwd+bwd of this BSHD GQA + SWA example. This compares TE ``DotProductAttention`` against the native JAX baseline above, which materializes attention scores with XLA ops; it is not a comparison against ``jax.nn.dot_product_attention(..., diff --git a/docs/examples/jax/dense.out b/docs/examples/jax/dense.out index 22b93ff04e1..29d6f4edde5 100644 --- a/docs/examples/jax/dense.out +++ b/docs/examples/jax/dense.out @@ -6,16 +6,16 @@ Variable collections: ['params'] {'params': {'Dense_0': {'kernel': ((8192, 32768), dtype('float32'))}}} bf16 baseline: -Mean time: 18.056 ms +Mean time: 18.265029820613563 ms TE MXFP8BlockScaling: -Mean time: 11.260 ms +Mean time: 10.855670000892133 ms # SINGLE_GPU_OUTPUT_END # MULTI_GPU_OUTPUT_START bf16 DP=2/TP=2: -Mean time: 5.516 ms +Mean time: 5.609957461711019 ms TE MXFP8BlockScaling DP=2/TP=2: -Mean time: 3.712 ms +Mean time: 3.6635108408518136 ms # MULTI_GPU_OUTPUT_END diff --git a/docs/examples/jax/dense.rst b/docs/examples/jax/dense.rst index 2bde67cd4c1..190f290c1e3 100644 --- a/docs/examples/jax/dense.rst +++ b/docs/examples/jax/dense.rst @@ -75,7 +75,7 @@ same input for both models. :start-after: # SINGLE_GPU_OUTPUT_START :end-before: # SINGLE_GPU_OUTPUT_END -On a single GB200, that's roughly **1.6× faster** for the fwd+bwd of one large +On a single GB200, that's roughly **1.7× faster** for the fwd+bwd of one large Dense — and the only code change was passing ``dot_general=te_dot_general_cls()`` into ``nn.Dense``. diff --git a/docs/examples/jax/quickstart_jax_utils.py b/docs/examples/jax/quickstart_jax_utils.py index 6547a5ff1a0..671eac7d8f0 100644 --- a/docs/examples/jax/quickstart_jax_utils.py +++ b/docs/examples/jax/quickstart_jax_utils.py @@ -23,8 +23,10 @@ def speedometer( warmup_iters: int = 50, rngs: Dict[str, jax.random.PRNGKey] = None, ) -> None: - """Measure average runtime for a JAX module - Perform forward and backward passes . + """Measure average forward/backward runtime for a JAX module. + + Wait for warmup and timed results so asynchronous dispatch does not skew + the elapsed wall time. """ if autocast_kwargs is None: autocast_kwargs = {"enabled": False} @@ -36,16 +38,19 @@ def speedometer( train_step_fn = create_train_step_fn(model_apply_fn, autocast_kwargs, forward_kwargs) # Warm up runs + result = None for _ in range(warmup_iters): rngs, step_rngs = _split_step_rngs(rngs) - loss, (param_grads, other_grads) = train_step_fn(variables, input, output_grad, step_rngs) + result = train_step_fn(variables, input, output_grad, step_rngs) + jax.block_until_ready(result) # Timing runs - start = time.time() + start = time.perf_counter() for _ in range(timing_iters): rngs, step_rngs = _split_step_rngs(rngs) - loss, (param_grads, other_grads) = train_step_fn(variables, input, output_grad, step_rngs) - end = time.time() + result = train_step_fn(variables, input, output_grad, step_rngs) + jax.block_until_ready(result) + end = time.perf_counter() print(f"Mean time: {(end - start) * 1000 / timing_iters} ms")