Restore and broadcast benchmark.

PiperOrigin-RevId: 873114079

Author: cpgaffney1

checkpoint/orbax/checkpoint/_src/testing/benchmarks/configs/evaluations/model_sizes/llama-405b.yaml

@@ -13,11 +13,11 @@ checkpoint_config:
   path: "gs://orbax-benchmarks/checkpoints/llama-3.1-405B-checkpoints/0/items"
 
 benchmarks:
-  - generator: "orbax.checkpoint._src.testing.benchmarks.v1_benchmark.V1Benchmark"
+  - generator: "orbax.checkpoint._src.testing.benchmarks.v1.benchmark.Benchmark"
     options:
       # --- Generator Options ---
-      # These keys must match the attributes of the `V1BenchmarkOptions` class
-      # associated with the `V1Benchmark` generator.
+      # These keys must match the attributes of the `BenchmarkOptions` class
+      # associated with the `Benchmark` generator.
       async_enabled: true
       use_ocdbt: true
       use_zarr3: true

checkpoint/orbax/checkpoint/_src/testing/benchmarks/configs/evaluations/model_sizes/llama-70b.yaml

@@ -13,11 +13,11 @@ checkpoint_config:
   path: "gs://orbax-benchmarks/checkpoints/llama-3.1-70B-checkpoints/0/items"
 
 benchmarks:
-  - generator: "orbax.checkpoint._src.testing.benchmarks.v1_benchmark.V1Benchmark"
+  - generator: "orbax.checkpoint._src.testing.benchmarks.v1.benchmark.Benchmark"
     options:
       # --- Generator Options ---
-      # These keys must match the attributes of the `V1BenchmarkOptions` class
-      # associated with the `V1Benchmark` generator.
+      # These keys must match the attributes of the `BenchmarkOptions` class
+      # associated with the `Benchmark` generator.
       async_enabled: true
       use_ocdbt: true
       use_zarr3: true

checkpoint/orbax/checkpoint/_src/testing/benchmarks/configs/evaluations/model_sizes/llama-8b.yaml

@@ -12,11 +12,11 @@ checkpoint_config:
   path: "gs://orbax-benchmarks/checkpoints/llama-3.1-8B-checkpoints/0/items"
 
 benchmarks:
-  - generator: "orbax.checkpoint._src.testing.benchmarks.v1_benchmark.V1Benchmark"
+  - generator: "orbax.checkpoint._src.testing.benchmarks.v1.benchmark.Benchmark"
     options:
       # --- Generator Options ---
-      # These keys must match the attributes of the `V1BenchmarkOptions` class
-      # associated with the `V1Benchmark` generator.
+      # These keys must match the attributes of the `BenchmarkOptions` class
+      # associated with the `Benchmark` generator.
       async_enabled: true
       use_ocdbt: true
       use_zarr3: true

checkpoint/orbax/checkpoint/_src/testing/benchmarks/configs/evaluations/resharding/llama-70b_generate_4-8-4.yaml

@@ -0,0 +1,27 @@
+# The name for the entire test suite run.
+# Assumes v5p-256 (128 chips)
+suite_name: "Llama 3.1 70B"
+num_repeats: 1
+
+
+mesh_config:
+  mesh_axes: ["data", "fsdp", "tensor"]
+  # SHould match reference_sharding_path.
+  ici_parallelism: {"data": 4, "fsdp": 8, "tensor": 4}
+
+# Note: checkpoint_config field not specified.
+checkpoint_config:
+  path: "gs://orbax-benchmarks/checkpoints/llama-3.1-70B-checkpoints/0/items"
+  sharding_config_path: "gs://orbax-benchmarks/sharding-configs/llama3.1-70b-v5p-256-data-4-fsdp-8-tensor-4/abstract_state.json"
+
+benchmarks:
+  - generator: "orbax.checkpoint._src.testing.benchmarks.v1.benchmark.Benchmark"
+    options:
+      # --- Generator Options ---
+      # These keys must match the attributes of the `BenchmarkOptions` class
+      # associated with the `Benchmark` generator.
+      async_enabled: true
+      use_ocdbt: true
+      use_zarr3: true
+      use_replica_parallel: false
+      use_compression: true

checkpoint/orbax/checkpoint/_src/testing/benchmarks/configs/evaluations/resharding/llama-70b_generate_4-8-4_subchunked.yaml

@@ -0,0 +1,28 @@
+# The name for the entire test suite run.
+# Assumes v5p-256 (128 chips)
+suite_name: "Llama 3.1 70B"
+num_repeats: 1
+
+
+mesh_config:
+  mesh_axes: ["data", "fsdp", "tensor"]
+  # SHould match reference_sharding_path.
+  ici_parallelism: {"data": 4, "fsdp": 8, "tensor": 4}
+
+# Note: checkpoint_config field not specified.
+checkpoint_config:
+  path: "gs://orbax-benchmarks/checkpoints/llama-3.1-70B-checkpoints/0/items"
+  sharding_config_path: "gs://orbax-benchmarks/sharding-configs/llama3.1-70b-v5p-256-data-4-fsdp-8-tensor-4/abstract_state.json"
+
+benchmarks:
+  - generator: "orbax.checkpoint._src.testing.benchmarks.v1.benchmark.Benchmark"
+    options:
+      # --- Generator Options ---
+      # These keys must match the attributes of the `BenchmarkOptions` class
+      # associated with the `Benchmark` generator.
+      async_enabled: true
+      use_ocdbt: true
+      use_zarr3: true
+      use_replica_parallel: false
+      use_compression: true
+      chunk_byte_size: 33554432  # 32 MiB

checkpoint/orbax/checkpoint/_src/testing/benchmarks/configs/evaluations/resharding/llama-70b_reshard_4-8-4_subchunked_to_1-128-1.yaml

@@ -0,0 +1,26 @@
+# The name for the entire test suite run.
+# Assumes v5p-256 (128 chips)
+suite_name: "llama-70b_reshard_4-8-4_subchunked_to_1-128-1"
+num_repeats: 20
+
+
+mesh_config:
+  mesh_axes: ["data", "fsdp", "tensor"]
+  # SHould match reference_sharding_path.
+  ici_parallelism: {"data": 1, "fsdp": 128, "tensor": 1}
+
+# Note: checkpoint_config field not specified.
+
+benchmarks:
+  - generator: "orbax.checkpoint._src.testing.benchmarks.v1.resharding_benchmark.ReshardingBenchmark"
+    options:
+      # --- Generator Options ---
+      # These keys must match the attributes of the `ReshardingBenchmarkOptions` class
+      # associated with the `ReshardingBenchmark` generator.
+      async_enabled: true
+      use_ocdbt: true
+      use_zarr3: true
+      use_replica_parallel: false
+      use_compression: true
+      reference_checkpoint_path: "gs://orbax-benchmarks/checkpoints/llama-70b_generate_4-8-4_subchunked/ckpt"
+      reference_sharding_path: "gs://orbax-benchmarks/sharding-configs/llama3.1-70b-v5p-256-data-1-fsdp-128-tensor-1/abstract_state.json"

checkpoint/orbax/checkpoint/_src/testing/benchmarks/configs/evaluations/resharding/llama-70b_reshard_4-8-4_subchunked_to_1-32-1.yaml

@@ -0,0 +1,26 @@
+# The name for the entire test suite run.
+# Assumes v5p-64 (32 chips)
+suite_name: "llama-70b_reshard_4-8-4_to_1-32-1"
+num_repeats: 20
+
+
+mesh_config:
+  mesh_axes: ["data", "fsdp", "tensor"]
+  # SHould match reference_sharding_path.
+  ici_parallelism: {"data": 1, "fsdp": 32, "tensor": 1}
+
+# Note: checkpoint_config field not specified.
+
+benchmarks:
+  - generator: "orbax.checkpoint._src.testing.benchmarks.v1.resharding_benchmark.ReshardingBenchmark"
+    options:
+      # --- Generator Options ---
+      # These keys must match the attributes of the `ReshardingBenchmarkOptions` class
+      # associated with the `ReshardingBenchmark` generator.
+      async_enabled: true
+      use_ocdbt: true
+      use_zarr3: true
+      use_replica_parallel: false
+      use_compression: true
+      reference_checkpoint_path: "gs://orbax-benchmarks/checkpoints/llama-70b_generate_4-8-4_subchunked/ckpt"
+      reference_sharding_path: "gs://orbax-benchmarks/sharding-configs/llama3.1-70b-v5p-64-data-1-fsdp-32-tensor-1/abstract_state.json"

checkpoint/orbax/checkpoint/_src/testing/benchmarks/configs/evaluations/resharding/llama-70b_reshard_4-8-4_to_1-128-1.yaml

@@ -0,0 +1,26 @@
+# The name for the entire test suite run.
+# Assumes v5p-256 (128 chips)
+suite_name: "llama-70b_reshard_4-8-4_to_1-128-1"
+num_repeats: 20
+
+
+mesh_config:
+  mesh_axes: ["data", "fsdp", "tensor"]
+  # SHould match reference_sharding_path.
+  ici_parallelism: {"data": 1, "fsdp": 128, "tensor": 1}
+
+# Note: checkpoint_config field not specified.
+
+benchmarks:
+  - generator: "orbax.checkpoint._src.testing.benchmarks.v1.resharding_benchmark.ReshardingBenchmark"
+    options:
+      # --- Generator Options ---
+      # These keys must match the attributes of the `ReshardingBenchmarkOptions` class
+      # associated with the `ReshardingBenchmark` generator.
+      async_enabled: true
+      use_ocdbt: true
+      use_zarr3: true
+      use_replica_parallel: false
+      use_compression: true
+      reference_checkpoint_path: "gs://orbax-benchmarks/checkpoints/llama-70b_generate_4-8-4/ckpt"
+      reference_sharding_path: "gs://orbax-benchmarks/sharding-configs/llama3.1-70b-v5p-256-data-1-fsdp-128-tensor-1/abstract_state.json"

checkpoint/orbax/checkpoint/_src/testing/benchmarks/configs/evaluations/resharding/llama-70b_reshard_4-8-4_to_1-32-1.yaml

@@ -0,0 +1,26 @@
+# The name for the entire test suite run.
+# Assumes v5p-64 (32 chips)
+suite_name: "llama-70b_reshard_4-8-4_to_1-32-1"
+num_repeats: 20
+
+
+mesh_config:
+  mesh_axes: ["data", "fsdp", "tensor"]
+  # SHould match reference_sharding_path.
+  ici_parallelism: {"data": 1, "fsdp": 32, "tensor": 1}
+
+# Note: checkpoint_config field not specified.
+
+benchmarks:
+  - generator: "orbax.checkpoint._src.testing.benchmarks.v1.resharding_benchmark.ReshardingBenchmark"
+    options:
+      # --- Generator Options ---
+      # These keys must match the attributes of the `ReshardingBenchmarkOptions` class
+      # associated with the `ReshardingBenchmark` generator.
+      async_enabled: true
+      use_ocdbt: true
+      use_zarr3: true
+      use_replica_parallel: false
+      use_compression: true
+      reference_checkpoint_path: "gs://orbax-benchmarks/checkpoints/llama-70b_generate_4-8-4/ckpt"
+      reference_sharding_path: "gs://orbax-benchmarks/sharding-configs/llama3.1-70b-v5p-64-data-1-fsdp-32-tensor-1/abstract_state.json"

checkpoint/orbax/checkpoint/_src/testing/benchmarks/configs/evaluations/resharding/llama-70b_reshard_4-8-4_to_4-8-4.yaml

@@ -0,0 +1,27 @@
+# The name for the entire test suite run.
+# Assumes v5p-256 (128 chips)
+# This benchmark is the no-resharding baseline.
+suite_name: "llama-70b_reshard_4-8-4_to_4-8-4"
+num_repeats: 20
+
+
+mesh_config:
+  mesh_axes: ["data", "fsdp", "tensor"]
+  # SHould match reference_sharding_path.
+  ici_parallelism: {"data": 4, "fsdp": 8, "tensor": 4}
+
+# Note: checkpoint_config field not specified.
+
+benchmarks:
+  - generator: "orbax.checkpoint._src.testing.benchmarks.v1.resharding_benchmark.ReshardingBenchmark"
+    options:
+      # --- Generator Options ---
+      # These keys must match the attributes of the `ReshardingBenchmarkOptions` class
+      # associated with the `ReshardingBenchmark` generator.
+      async_enabled: true
+      use_ocdbt: true
+      use_zarr3: true
+      use_replica_parallel: false
+      use_compression: true
+      reference_checkpoint_path: "gs://orbax-benchmarks/checkpoints/llama-70b_generate_4-8-4/ckpt"
+      reference_sharding_path: "gs://orbax-benchmarks/sharding-configs/llama3.1-70b-v5p-256-data-4-fsdp-8-tensor-4/abstract_state.json"