Merge pull request #1176 from AI-Hypercomputer:add_flag

PiperOrigin-RevId: 716807970
AI-Hypercomputer · Jan 17, 2025 · 21706fc · 21706fc
2 parents 0b79a29 + 94d1e7b
commit 21706fc
Showing 1 changed file with 5 additions and 0 deletions.
diff --git a/benchmarks/xla_flags_library.py b/benchmarks/xla_flags_library.py
@@ -131,6 +131,11 @@
     " --xla_latency_hiding_scheduler_rerun=2"
 )
 
+# Flags to optimize pipeline parallelism over DCN with large host offloading.
+PIPELINING_FLAGS = (
+    " --xla_tpu_iova_dma_chunk_size_bytes=16777216" # breaks DMA to/from host into 16M chunks
+)
+
 # Disable bundle-aware CostModel which was causing worse perf b/357103386.
 # Some fusions in the backward pass of the model were 3x slower without this.
 DISABLE_BUNDLE_AWARE_COST_MODEL = (