No public description

PiperOrigin-RevId: 886881258

Author: cpgaffney1

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

@@ -11,6 +11,7 @@ mesh_config:
 # The checkpoint configuration, shared across all generated benchmarks.
 checkpoint_config:
   path: "gs://orbax-benchmarks/checkpoints/llama-3.1-405B-checkpoints/0/items"
+  sharding_config_path: "gs://orbax-benchmarks/sharding-configs/llama3.1-405b-v5p-128-data-1-fsdp-64-tensor-1/abstract_state.json"
 
 benchmarks:
   - generator: "orbax.checkpoint._src.testing.benchmarks.v1.benchmark.Benchmark"

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

@@ -11,6 +11,7 @@ mesh_config:
 # The checkpoint configuration, shared across all generated benchmarks.
 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-32-data-1-fsdp-16-tensor-1/abstract_state.json"
 
 benchmarks:
   - generator: "orbax.checkpoint._src.testing.benchmarks.v1.benchmark.Benchmark"
@@ -23,4 +24,3 @@ benchmarks:
       use_zarr3: true
       use_replica_parallel: false
       use_compression: true
-      enable_trace: true

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

@@ -10,6 +10,7 @@ mesh_config:
 # The checkpoint configuration, shared across all generated benchmarks.
 checkpoint_config:
   path: "gs://orbax-benchmarks/checkpoints/llama-3.1-8B-checkpoints/0/items"
+  sharding_config_path: "gs://orbax-benchmarks/sharding-configs/llama3.1-8b-v5p-8-data-1-fsdp-4-tensor-1/abstract_state.json"
 
 benchmarks:
   - generator: "orbax.checkpoint._src.testing.benchmarks.v1.benchmark.Benchmark"

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

@@ -0,0 +1,26 @@
+# The name for the entire test suite run.
+# Assumes v5p-128 (64 chips)
+suite_name: "Llama 3.1 405B"
+num_repeats: 20
+
+
+mesh_config:
+  mesh_axes: ["data", "fsdp", "tensor"]
+  ici_parallelism: {"data": 1, "fsdp": 64, "tensor": 1}
+
+# The checkpoint configuration, shared across all generated benchmarks.
+checkpoint_config:
+  path: "gs://orbax-benchmarks/checkpoints/llama-405b_generate_1-64-1"
+  sharding_config_path: "gs://orbax-benchmarks/sharding-configs/llama3.1-405b-v5p-128-data-1-fsdp-64-tensor-1/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/pathways/llama-70b.yaml

@@ -0,0 +1,26 @@
+# The name for the entire test suite run.
+# Assumes v5p-32 (16 chips)
+suite_name: "Llama 3.1 70B"
+num_repeats: 20
+
+
+mesh_config:
+  mesh_axes: ["data", "fsdp", "tensor"]
+  ici_parallelism: {"data": 1, "fsdp": 16, "tensor": 1}
+
+# The checkpoint configuration, shared across all generated benchmarks.
+checkpoint_config:
+  path: "gs://orbax-benchmarks/checkpoints/llama-70b_generate_1-16-1"
+  sharding_config_path: "gs://orbax-benchmarks/sharding-configs/llama3.1-70b-v5p-32-data-1-fsdp-16-tensor-1/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/pathways/llama-8b.yaml

@@ -0,0 +1,25 @@
+# The name for the entire test suite run.
+# Assumes v5p-8 (4 chips)
+suite_name: "Llama 3.1 8B"
+num_repeats: 20
+
+mesh_config:
+  mesh_axes: ["data", "fsdp", "tensor"]
+  ici_parallelism: {"data": 1, "fsdp": 4, "tensor": 1}
+
+# The checkpoint configuration, shared across all generated benchmarks.
+checkpoint_config:
+  path: "gs://orbax-benchmarks/checkpoints/llama-8b_generate_1-4-1"
+  sharding_config_path: "gs://orbax-benchmarks/sharding-configs/llama3.1-8b-v5p-8-data-1-fsdp-4-tensor-1/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-405b_generate_1-64-1.yaml

@@ -0,0 +1,27 @@
+# The name for the entire test suite run.
+# Assumes v5p-128 (64 chips)
+suite_name: "Llama 3.1 405B"
+num_repeats: 1
+
+
+mesh_config:
+  mesh_axes: ["data", "fsdp", "tensor"]
+  # Should match reference_sharding_path.
+  ici_parallelism: {"data": 1, "fsdp": 64, "tensor": 1}
+
+# Note: checkpoint_config field not specified.
+checkpoint_config:
+  path: "gs://orbax-benchmarks/checkpoints/llama-3.1-405B-checkpoints/0/items"
+  sharding_config_path: "gs://orbax-benchmarks/sharding-configs/llama3.1-405b-v5p-128-data-1-fsdp-64-tensor-1/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_1-16-1.yaml

@@ -0,0 +1,27 @@
+# The name for the entire test suite run.
+# Assumes v5p-32 (16 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": 1, "fsdp": 16, "tensor": 1}
+
+# 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-32-data-1-fsdp-16-tensor-1/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-8b_generate_1-4-1.yaml

@@ -0,0 +1,27 @@
+# The name for the entire test suite run.
+# Assumes v5p-8 (4 chips)
+suite_name: "Llama 3.1 8B"
+num_repeats: 1
+
+
+mesh_config:
+  mesh_axes: ["data", "fsdp", "tensor"]
+  # Should match reference_sharding_path.
+  ici_parallelism: {"data": 1, "fsdp": 4, "tensor": 1}
+
+# Note: checkpoint_config field not specified.
+checkpoint_config:
+  path: "gs://orbax-benchmarks/checkpoints/llama-3.1-8B-checkpoints/0/items"
+  sharding_config_path: "gs://orbax-benchmarks/sharding-configs/llama3.1-8b-v5p-8-data-1-fsdp-4-tensor-1/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