Add v6e PCIe overload workaround flag (apple#955)

Author: hanzhi713

axlearn/common/compiler_options.py

@@ -74,6 +74,9 @@ def default_xla_options(
             xla_tpu_scoped_vmem_limit_kib=98304,
             # For megascale performance.
             xla_jf_crs_combiner_threshold_count=10,
+            # TODO(hanzhi-zhou): temporary workaround to avoid PCIe overload when using multi-slice
+            # v6e training caused by allreduce over DCN. This flag doesn't impact performance.
+            xla_tpu_iova_dma_chunk_size_bytes=1048576,
         )
         options.update(
             # Improved performance for v6e.