From 1683fb7ecb8399fdaff64fd5e680329b220dc4e4 Mon Sep 17 00:00:00 2001 From: Mohit Khatwani Date: Mon, 9 Feb 2026 03:45:01 +0000 Subject: [PATCH] diloco trainer --- .../base_requirements/requirements.txt | 1 + .../cuda12-requirements.txt | 1 + .../tpu-requirements.txt | 1 + dependencies/requirements/requirements.txt | 1 + .../vllm/maxtext_vllm_adapter/config.json | 58 ++++ src/MaxText/sharding.py | 8 +- src/MaxText/train_compile.py | 35 ++- src/maxtext/common/data_loader.py | 6 +- src/maxtext/configs/base.yml | 11 +- src/maxtext/configs/types.py | 24 ++ src/maxtext/trainers/diloco/__init__.py | 13 + src/maxtext/trainers/diloco/diloco.py | 279 +++++++++++++++++ src/maxtext/utils/maxtext_utils.py | 9 +- src/maxtext/utils/train_utils.py | 44 ++- tests/unit/diloco_test.py | 287 ++++++++++++++++++ .../tpu7x-16/slice_1/named_shardings.json | 168 ++++++++++ .../tpu7x-16/slice_4/named_shardings.json | 168 ++++++++++ .../v5p-16/slice_1/named_shardings.json | 168 ++++++++++ .../v5p-16/slice_4/named_shardings.json | 168 ++++++++++ .../v6e-16/slice_1/named_shardings.json | 168 ++++++++++ .../v6e-16/slice_4/named_shardings.json | 168 ++++++++++ .../tpu7x-16/slice_1/named_shardings.json | 252 +++++++++++++++ .../tpu7x-16/slice_4/named_shardings.json | 252 +++++++++++++++ .../v5p-16/slice_1/named_shardings.json | 252 +++++++++++++++ .../v5p-16/slice_4/named_shardings.json | 252 +++++++++++++++ .../v6e-16/slice_1/named_shardings.json | 252 +++++++++++++++ .../v6e-16/slice_4/named_shardings.json | 252 +++++++++++++++ .../tpu7x-16/slice_1/named_shardings.json | 84 +++++ .../tpu7x-16/slice_4/named_shardings.json | 84 +++++ .../v5p-16/slice_1/named_shardings.json | 84 +++++ .../v5p-16/slice_4/named_shardings.json | 84 +++++ .../v6e-16/slice_1/named_shardings.json | 84 +++++ .../v6e-16/slice_4/named_shardings.json | 84 +++++ 33 files changed, 3782 insertions(+), 20 deletions(-) create mode 100644 src/MaxText/integration/vllm/maxtext_vllm_adapter/config.json create mode 100644 src/maxtext/trainers/diloco/__init__.py create mode 100644 src/maxtext/trainers/diloco/diloco.py create mode 100644 tests/unit/diloco_test.py diff --git a/dependencies/requirements/base_requirements/requirements.txt b/dependencies/requirements/base_requirements/requirements.txt index 582d99c3d7..c40252cfc1 100644 --- a/dependencies/requirements/base_requirements/requirements.txt +++ b/dependencies/requirements/base_requirements/requirements.txt @@ -4,6 +4,7 @@ array-record cloud-accelerator-diagnostics cloud-tpu-diagnostics datasets +drjax flax gcsfs google-api-python-client diff --git a/dependencies/requirements/generated_requirements/cuda12-requirements.txt b/dependencies/requirements/generated_requirements/cuda12-requirements.txt index 00efbc3b1c..9879536ab1 100644 --- a/dependencies/requirements/generated_requirements/cuda12-requirements.txt +++ b/dependencies/requirements/generated_requirements/cuda12-requirements.txt @@ -40,6 +40,7 @@ dill>=0.4.0 distlib>=0.4.0 dm-tree>=0.1.9 docstring-parser>=0.17.0 +drjax>=0.1.4 editdistance>=0.8.1 einops>=0.8.1 einshape>=1.0 diff --git a/dependencies/requirements/generated_requirements/tpu-requirements.txt b/dependencies/requirements/generated_requirements/tpu-requirements.txt index 1e16576363..21f5668c98 100644 --- a/dependencies/requirements/generated_requirements/tpu-requirements.txt +++ b/dependencies/requirements/generated_requirements/tpu-requirements.txt @@ -41,6 +41,7 @@ dill>=0.4.0 distlib>=0.4.0 dm-tree>=0.1.9 docstring-parser>=0.17.0 +drjax>=0.1.4 editdistance>=0.8.1 einops>=0.8.1 einshape>=1.0 diff --git a/dependencies/requirements/requirements.txt b/dependencies/requirements/requirements.txt index 439e0e3a75..7ae9f9114a 100644 --- a/dependencies/requirements/requirements.txt +++ b/dependencies/requirements/requirements.txt @@ -4,6 +4,7 @@ array-record cloud-accelerator-diagnostics cloud-tpu-diagnostics datasets +drjax>=0.1.4 flax gcsfs google-api-python-client diff --git a/src/MaxText/integration/vllm/maxtext_vllm_adapter/config.json b/src/MaxText/integration/vllm/maxtext_vllm_adapter/config.json new file mode 100644 index 0000000000..04603b729c --- /dev/null +++ b/src/MaxText/integration/vllm/maxtext_vllm_adapter/config.json @@ -0,0 +1,58 @@ +{ + "architectures": [ + "MaxTextForCausalLM" + ], + "attention_bias": false, + "attention_dropout": 0.0, + "auto_map": { + "AutoConfig": "configuration_deepseek.DeepseekV3Config", + "AutoModel": "modeling_deepseek.DeepseekV3Model", + "AutoModelForCausalLM": "modeling_deepseek.DeepseekV3ForCausalLM" + }, + "bos_token_id": 0, + "eos_token_id": 1, + "ep_size": 1, + "first_k_dense_replace": 3, + "hidden_act": "silu", + "hidden_size": 7168, + "initializer_range": 0.02, + "intermediate_size": 18432, + "kv_lora_rank": 512, + "max_position_embeddings": 163840, + "model_type": "deepseek_v3", + "moe_intermediate_size": 2048, + "moe_layer_freq": 1, + "n_group": 8, + "n_routed_experts": 256, + "n_shared_experts": 1, + "norm_topk_prob": true, + "num_attention_heads": 128, + "num_experts_per_tok": 8, + "num_hidden_layers": 61, + "num_key_value_heads": 128, + "num_nextn_predict_layers": 1, + "q_lora_rank": 1536, + "qk_nope_head_dim": 128, + "qk_rope_head_dim": 64, + "rms_norm_eps": 1e-06, + "rope_scaling": { + "beta_fast": 32, + "beta_slow": 1, + "factor": 40, + "mscale": 1.0, + "mscale_all_dim": 1.0, + "original_max_position_embeddings": 4096, + "type": "yarn" + }, + "rope_theta": 10000, + "routed_scaling_factor": 2.5, + "scoring_func": "sigmoid", + "tie_word_embeddings": false, + "topk_group": 4, + "topk_method": "noaux_tc", + "torch_dtype": "bfloat16", + "transformers_version": "4.33.1", + "use_cache": true, + "v_head_dim": 128, + "vocab_size": 129280 +} \ No newline at end of file diff --git a/src/MaxText/sharding.py b/src/MaxText/sharding.py index ed4967dbab..d5eb12ad53 100644 --- a/src/MaxText/sharding.py +++ b/src/MaxText/sharding.py @@ -36,7 +36,13 @@ def get_input_data_sharding(config, mesh): """Get the input data sharding for the model""" - return create_sharding(mesh, config.input_data_sharding_logical_axes, rules=config.logical_axis_rules) + if config.enable_diloco: + data_sharding = create_sharding( + mesh, ["diloco"] + config.input_data_sharding_logical_axes, rules=config.logical_axis_rules + ) + else: + data_sharding = create_sharding(mesh, config.input_data_sharding_logical_axes, rules=config.logical_axis_rules) + return data_sharding def maybe_shard_with_name(inputs, named_sharding, shard_mode, debug_sharding=False, extra_stack_level=0): diff --git a/src/MaxText/train_compile.py b/src/MaxText/train_compile.py index a5c88350ef..2d21e9cc54 100644 --- a/src/MaxText/train_compile.py +++ b/src/MaxText/train_compile.py @@ -24,6 +24,7 @@ from typing import Sequence import os import pickle +import functools from absl import app @@ -45,6 +46,7 @@ from maxtext.utils import gcs_utils from maxtext.utils import max_utils from maxtext.utils import maxtext_utils +from maxtext.trainers.diloco import diloco # pylint: disable=too-many-positional-arguments @@ -235,13 +237,32 @@ def main(argv: Sequence[str]) -> None: # Get data sharding data_sharding = sharding.get_input_data_sharding(config, topology_mesh) - - # Get function to compile and shardings - func_to_compile, in_shard, out_shard, static_argnums, donate_argnums = ( - maxtext_utils.get_functional_train_with_signature( - train.train_step, data_sharding, state_mesh_shardings, model, config - ) - ) + if config.enable_diloco: + # Build abstract DiLoCo state and shardings for AOT compilation + abstract_state = shaped_train_args[0] + diloco_state, state_mesh_shardings, inner_state_shardings = diloco.build_abstract_diloco_state( + config, abstract_state, state_mesh_shardings, topology_mesh + ) + shaped_train_args = (diloco_state, shaped_train_args[1], shaped_train_args[2]) + + # Wrap train_step with diloco + train_step_partial = functools.partial(train.train_step, model, config, inner_state_shardings, None) + train_step_fn = diloco.build_diloco_train_step(config, train_step_partial) + + # For DiLoCo, the train_step_fn is already fully wrapped and takes (state, batch, prng) + func_to_compile = train_step_fn + func_to_compile.__name__ = "train_step" + in_shard = (state_mesh_shardings, data_sharding, None) # State, batch, rng + out_shard = (state_mesh_shardings, None) # State, metrics + static_argnums = () + donate_argnums = 0 + else: + # Get function to compile and shardings + func_to_compile, in_shard, out_shard, static_argnums, donate_argnums = ( + maxtext_utils.get_functional_train_with_signature( + train.train_step, data_sharding, state_mesh_shardings, model, config + ) + ) # print weights sharding info under debug sharding mode if config.debug_sharding: diff --git a/src/maxtext/common/data_loader.py b/src/maxtext/common/data_loader.py index 83f73f5c48..274495d897 100644 --- a/src/maxtext/common/data_loader.py +++ b/src/maxtext/common/data_loader.py @@ -25,6 +25,7 @@ maybe_record_goodput, ) from maxtext.utils import exceptions +from maxtext.trainers.diloco import diloco class DataLoader: @@ -70,10 +71,13 @@ def load_next_batch_pre_sharding(self): def load_next_batch(self, *args, **kwargs): """Loads the next batch with sharding hint""" - return jax.device_put( + example_batch = jax.device_put( self.load_next_batch_pre_sharding(), self.input_data_shardings, ) + if self.config.enable_diloco: + example_batch = diloco.reshape_first_axis_with_diloco(self.config.num_diloco_replicas, example_batch) + return example_batch def check_example_batch(self): if self.config.max_checkify: diff --git a/src/maxtext/configs/base.yml b/src/maxtext/configs/base.yml index fc836aeb44..0e059108b2 100644 --- a/src/maxtext/configs/base.yml +++ b/src/maxtext/configs/base.yml @@ -400,7 +400,7 @@ hardware: 'tpu' # Supported hardware types are 'tpu', 'gpu', 'gpu_multiprocess' # Parallelism shard_mode: "auto" # can be either auto or explicit -mesh_axes: ['data', 'stage', 'fsdp', 'fsdp_transpose', 'sequence', 'context', 'context_autoregressive', 'tensor', 'tensor_transpose', 'tensor_sequence', 'expert', 'autoregressive'] +mesh_axes: ['diloco', 'data', 'stage', 'fsdp', 'fsdp_transpose', 'sequence', 'context', 'context_autoregressive', 'tensor', 'tensor_transpose', 'tensor_sequence', 'expert', 'autoregressive'] logical_axis_rules: [ ['activation_batch', ['data', 'fsdp', 'fsdp_transpose', 'expert']], ['activation_batch_no_exp', ['data', 'fsdp', 'fsdp_transpose']], @@ -483,6 +483,7 @@ logical_axis_rules: [ ['paged_kv_head_dim_size', []], ['dense_layers', []], ['moe_layers', []], + ['diloco', 'diloco'], ] # Axes used for DCN must be earlier in this list than ICI, see (b/339009148) for details data_sharding: [['data', 'stage', 'fsdp', 'fsdp_transpose', 'sequence', 'context', 'context_autoregressive', 'tensor', 'tensor_transpose', 'tensor_sequence', 'expert', 'autoregressive']] @@ -495,6 +496,7 @@ sharding_tolerance: 0.02 # value to auto-shard based on available slices and devices. # By default, product of the DCN axes should equal number of slices # and product of the ICI axes should equal number of devices per slice. +dcn_diloco_parallelism: 1 dcn_data_parallelism: -1 # recommended DCN axis to be auto-sharded dcn_fsdp_parallelism: 1 dcn_fsdp_transpose_parallelism: 1 @@ -507,6 +509,7 @@ dcn_tensor_sequence_parallelism: 1 # never recommended dcn_pipeline_parallelism: 1 dcn_expert_parallelism: 1 dcn_autoregressive_parallelism: 1 # never recommended +ici_diloco_parallelism: 1 ici_data_parallelism: 1 ici_fsdp_parallelism: -1 # recommended ICI axis to be auto-sharded ici_fsdp_transpose_parallelism: 1 @@ -738,6 +741,12 @@ enable_data_shuffling: True data_shuffle_seed: 0 init_weights_seed: 0 +# DiLoCo params. +enable_diloco: False +diloco_sync_period: 36 +diloco_outer_lr: 0.3 +diloco_outer_momentum: 0.9 + # You may disable clipping by setting gradient_clipping_threshold to zero. gradient_clipping_threshold: 1.0 diff --git a/src/maxtext/configs/types.py b/src/maxtext/configs/types.py index 1043dbc3ed..b2c17923bb 100644 --- a/src/maxtext/configs/types.py +++ b/src/maxtext/configs/types.py @@ -784,6 +784,7 @@ class LayoutAndSharding(BaseModel): class DcnParallelism(BaseModel): """Parallelism dimensions across the DCN (Data Center Network).""" + dcn_diloco_parallelism: int = Field(1, description="DCN axis for Diloco parallelism.") dcn_data_parallelism: int = Field(-1, description="DCN axis for data parallelism.") dcn_fsdp_parallelism: int = Field(1, description="DCN axis for FSDP.") dcn_fsdp_transpose_parallelism: int = Field(1, description="DCN axis for FSDP transpose.") @@ -803,6 +804,7 @@ class DcnParallelism(BaseModel): class IciParallelism(BaseModel): """Parallelism dimensions within the ICI (Inter-Chip Interconnect).""" + ici_diloco_parallelism: int = Field(1, description="ICI axis for Diloco parallelism.") ici_data_parallelism: int = Field(1, description="ICI axis for data parallelism.") ici_fsdp_parallelism: int = Field(-1, description="ICI axis for FSDP.") ici_fsdp_transpose_parallelism: int = Field(1, description="ICI axis for FSDP transpose.") @@ -1082,6 +1084,15 @@ class ManifoldConstrainedHyperConnections(BaseModel): sinkhorn_iterations: PositiveInt = Field(20, description="The number of iterations for the Sinkhorn-Knopp algorithm.") +class DilocoParams(BaseModel): + """Diloco Hyperparameters""" + + enable_diloco: bool = Field(False, description="Enable Diloco parallelism") + diloco_sync_period: int = Field(36, description="Diloco sync period.") + diloco_outer_lr: float = Field(0.3, description="learning rate for outer optimizer.") + diloco_outer_momentum: float = Field(0.9, description="momentum for outer optimizer.") + + class Optimizer(BaseModel): """Configuration for the optimizer and learning rate schedule.""" @@ -1632,6 +1643,11 @@ class DerivedValues(BaseModel): description="Effective number of query heads, scaled by `global_parameter_scale`.", ) + num_diloco_replicas: None | int = Field( + None, + description="The number of diloco replicas, derived from ICI and DCN values.", + ) + ici_parallelism: None | list[int] = Field( None, description="Aggregated list of all ICI parallelism values for legacy compatibility.", @@ -1779,6 +1795,7 @@ class MaxTextConfig( RematAndOffload, TrainingLoop, ManifoldConstrainedHyperConnections, + DilocoParams, Optimizer, AdamW, Muon, @@ -2375,6 +2392,7 @@ def calculate_global_batch_sizes(per_device_batch_size, expansion_factor, num_de # Create the ici_parallelism and dcn_parallelism lists for legacy compatibility. if self.using_pipeline_parallelism and self.mesh_axes and self.mesh_axes[0] == "stage": self.ici_parallelism = [ + self.ici_diloco_parallelism, self.ici_pipeline_parallelism, self.ici_data_parallelism, self.ici_fsdp_parallelism, @@ -2389,6 +2407,7 @@ def calculate_global_batch_sizes(per_device_batch_size, expansion_factor, num_de self.ici_autoregressive_parallelism, ] self.dcn_parallelism = [ + self.dcn_diloco_parallelism, self.dcn_pipeline_parallelism, self.dcn_data_parallelism, self.dcn_fsdp_parallelism, @@ -2404,6 +2423,7 @@ def calculate_global_batch_sizes(per_device_batch_size, expansion_factor, num_de ] else: ici_map = { + "diloco": self.ici_diloco_parallelism, "data": self.ici_data_parallelism, "stage": self.ici_pipeline_parallelism, "fsdp": self.ici_fsdp_parallelism, @@ -2422,6 +2442,7 @@ def calculate_global_batch_sizes(per_device_batch_size, expansion_factor, num_de self.ici_parallelism = [ici_map[axis] for axis in self.mesh_axes] dcn_map = { + "diloco": self.dcn_diloco_parallelism, "data": self.dcn_data_parallelism, "stage": self.dcn_pipeline_parallelism, "fsdp": self.dcn_fsdp_parallelism, @@ -2439,6 +2460,9 @@ def calculate_global_batch_sizes(per_device_batch_size, expansion_factor, num_de } self.dcn_parallelism = [dcn_map[axis] for axis in self.mesh_axes] + # Diloco params + self.num_diloco_replicas = int(self.ici_diloco_parallelism * self.dcn_diloco_parallelism) + # Final string-to-enum conversions if they haven't been coerced by pydantic yet. if isinstance(self.decoder_block, str): self.decoder_block = DecoderBlockType(self.decoder_block.lower()) diff --git a/src/maxtext/trainers/diloco/__init__.py b/src/maxtext/trainers/diloco/__init__.py new file mode 100644 index 0000000000..5c7e6e3878 --- /dev/null +++ b/src/maxtext/trainers/diloco/__init__.py @@ -0,0 +1,13 @@ +# Copyright 2023-2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. diff --git a/src/maxtext/trainers/diloco/diloco.py b/src/maxtext/trainers/diloco/diloco.py new file mode 100644 index 0000000000..d12ec0c65e --- /dev/null +++ b/src/maxtext/trainers/diloco/diloco.py @@ -0,0 +1,279 @@ +# Copyright 2025 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""An implementation of Distributed Low-Communication (DiLoCo) training. + +This module contains implementations of: + +- DiLoCo: Distributed Low-Communication Training of Language Models + https://arxiv.org/abs/2311.08105 +- Streaming DiLoCo with overlapping communication: Towards a Distributed Free Lunch + https://arxiv.org/abs/2501.18512 +""" + +from collections.abc import Sequence +from typing import Any, Callable + +import drjax +from flax import struct +from flax.training import train_state +import jax +import jax.numpy as jnp +from jaxtyping import Array, Int32, Key, PyTree, UInt32 +import optax + +from MaxText import pyconfig + +Batch = Any +Params = PyTree +Metrics = PyTree +OptState = optax.OptState +InnerOptStates = optax.OptState +PRNGKey = Key[Array, ""] | UInt32[Array, "2"] +Step = Int32[Array, ""] + + +class DiLoCoTrainState(struct.PyTreeNode): + """The state of the DiLoCo training process. + + Attributes: + inner_state: A `flax.training.train_state.TrainState` of the state for each + step of the inner optimization. All arrays are expected to have a leading + dimension with size of the number of diloco replicas so that training + steps can be mapped over this dimension. + params: A PyTree of the global model weights. These will mimic a + sub-PyTree in `inner_state`, which rank-1 shape. + outer_opt_state: The state for the outer Nesterov momentum optimizer. + step: The step counter of the training process. + """ + + inner_state: train_state.TrainState + params: Params + outer_opt_state: OptState + step: Step + + +def add_diloco_to_sharding(pytree): + """ + Recursively traverses a PyTree and prepends 'diloco' to the PartitionSpec + of any NamedSharding object that doesn't have an empty PartitionSpec. + """ + + def map_fn(leaf): + if isinstance(leaf, jax.sharding.NamedSharding): + new_spec = jax.sharding.PartitionSpec("diloco", *leaf.spec) + return jax.sharding.NamedSharding(mesh=leaf.mesh, spec=new_spec) + return leaf + + return jax.tree_util.tree_map(map_fn, pytree) + + +def reshape_first_axis_with_diloco(num_diloco_replicas: int, pytree: PyTree) -> PyTree: + """Reshapes the first dimension of each array in the PyTree to include a DiLoCo axis. + + This function takes a a batch of data represented as a PyTree + and reshapes the leading dimension of each array within it. The purpose is + to introduce a new 'diloco' axis, which is used for distributing data + across DiLoCo replicas. + + Args: + num_diloco_replicas: The number of DiLoCo replicas. This determines the + size of the new leading dimension. + pytree: The input PyTree, where each array is expected to have a batch + dimension as its first axis. + + Returns: + A new PyTree with the same structure as the input, but with each array's + first dimension reshaped to `(num_diloco_replicas, original_batch_dim // num_diloco_replicas, ...)`. + The sharding specification is also updated to include the 'diloco' axis. + """ + + def extend_pspec(pspec: jax.sharding.PartitionSpec | Sequence[str | Sequence[str]] = ()) -> jax.sharding.PartitionSpec: + if tuple(*pspec)[0] == "diloco": + # pull out diloco axis if already present + return jax.sharding.PartitionSpec("diloco", (*pspec[0][1:],), (*pspec[1:],)) + return jax.sharding.PartitionSpec("diloco", *pspec) + + def reshape_for_diloco(arr): + batch_dim, *example_shape = arr.shape + diloco_shape = (num_diloco_replicas, batch_dim // num_diloco_replicas, *example_shape) + s = arr.sharding + s = jax.sharding.NamedSharding(mesh=s.mesh, spec=extend_pspec(s.spec)) + return jax.lax.with_sharding_constraint(jnp.reshape(arr, shape=diloco_shape), s) + + return jax.tree.map(reshape_for_diloco, pytree) + + +def build_abstract_diloco_state( + config: "pyconfig.HyperParameters", + abstract_state: PyTree, + state_mesh_shardings: PyTree, + mesh: jax.sharding.Mesh, +) -> tuple[DiLoCoTrainState, DiLoCoTrainState]: + """Build abstract DiLoCo state and shardings for AOT compilation. + + This function creates abstract (shape-only) DiLoCo state suitable for + ahead-of-time compilation, where we don't have actual arrays. + + Args: + config: The config used to set up training. + abstract_state: Abstract train state (ShapeDtypeStruct objects). + state_mesh_shardings: Shardings for the regular train state. + mesh: The mesh for sharding. + + Returns: + A tuple of (abstract_diloco_state, diloco_state_shardings). + """ + + # Create inner state with diloco dimension prepended to all arrays + def add_diloco_dim(x): + if hasattr(x, "shape") and hasattr(x, "dtype"): + new_shape = (config.num_diloco_replicas,) + tuple(x.shape) + return jax.ShapeDtypeStruct(new_shape, x.dtype) + return x + + inner_state = jax.tree.map(add_diloco_dim, abstract_state) + + # Create outer optimizer state shape using eval_shape + outer_optimizer = optax.sgd( + config.diloco_outer_lr, + momentum=config.diloco_outer_momentum, + nesterov=True, + ) + outer_opt_state = jax.eval_shape(outer_optimizer.init, abstract_state.params) + + # Create abstract step + abstract_step = jax.ShapeDtypeStruct((), jnp.int32) + + # Build abstract DiLoCo state + diloco_state = DiLoCoTrainState( + inner_state=inner_state, + params=abstract_state.params, + outer_opt_state=outer_opt_state, + step=abstract_step, + ) + + # Build shardings + inner_state_shardings = add_diloco_to_sharding(state_mesh_shardings) + outer_opt_state_sharding = jax.tree.map( + lambda _: jax.sharding.NamedSharding(mesh, jax.sharding.PartitionSpec()), + outer_opt_state, + ) + diloco_state_shardings = DiLoCoTrainState( + inner_state=inner_state_shardings, + params=state_mesh_shardings.params, + outer_opt_state=outer_opt_state_sharding, + step=None, + ) + + return diloco_state, diloco_state_shardings, inner_state_shardings + + +def build_diloco_state( + config: "pyconfig.HyperParameters", + initialize_state: Callable[[], train_state.TrainState], +) -> tuple[DiLoCoTrainState, PyTree]: + """Given a non-DiLoCo train state, construct a DiLoCo training state.""" + outer_optimizer = optax.sgd( + config.diloco_outer_lr, + momentum=config.diloco_outer_momentum, + nesterov=True, + ) + + @drjax.program(placements={"diloco": config.num_diloco_replicas}) + def init_diloco_state() -> tuple[DiLoCoTrainState, PyTree]: + state = initialize_state() + # Inner state must be broadcast across clients. + inner_state = drjax.broadcast(state) + # Outer state retains a single copy of the model parameters and optimizer state. + outer_params = state.params + outer_opt_state = outer_optimizer.init(outer_params) + outer_opt_state_sharding = jax.tree_util.tree_map(lambda x: x.sharding, outer_opt_state) + return ( + DiLoCoTrainState(inner_state=inner_state, params=outer_params, outer_opt_state=outer_opt_state, step=state.step), + outer_opt_state_sharding, + ) + + return init_diloco_state() + + +def build_diloco_train_step( + config: pyconfig.HyperParameters, + train_step: Callable[[train_state.TrainState, Batch, PRNGKey], tuple[train_state.TrainState, Metrics]], +) -> Callable[[DiLoCoTrainState, Batch, PRNGKey], tuple[DiLoCoTrainState, Metrics]]: + """Convert a local state and train step into DiLoCo-compatible versions. + + This is an implementation of the original (non-streaming) DiLoCo algorithm + which syncs all model parameters across the replicas every + `config.diloco_sync_period` steps, treating the difference accumulated over + non-sync steps as a pseudo gradient and applying SGD with Nesterov momentum on + the "global" model. + + Args: + config: The config used to set up training. + train_step: A local train step. This will be executed independently within + each replica. + """ + outer_optimizer = optax.sgd( + config.diloco_outer_lr, + momentum=config.diloco_outer_momentum, + nesterov=True, + ) + + def synchronize(state): + # Calculate the delta between the current replica's state and the global + # state (since last synchronization). + broadcast_outer_params = drjax.broadcast(state.params) + model_delta = jax.tree.map(lambda x, y: y - x, state.inner_state.params, broadcast_outer_params) + # Treat the average delta as the outer optimizer's gradient and apply to + # the global (outer) model params. + averaged_pseudo_grad = drjax.reduce_mean(model_delta) + updates, new_opt_state = outer_optimizer.update(averaged_pseudo_grad, state.outer_opt_state, state.params) + new_outer_params = optax.apply_updates(state.params, updates) + # Replace inner model params with the new global model params. + # NOTE: inner optimizer state is retained despite the change in parameters, + # see section 6.1 in https://arxiv.org/pdf/2311.08105. + new_inner_state = drjax.map_fn(lambda state: state.replace(params=new_outer_params), state.inner_state) + return state.replace( + params=new_outer_params, + outer_opt_state=new_opt_state, + inner_state=new_inner_state, + ) + + def typed_reduce_mean(in_tree): + total = drjax.reduce_sum(in_tree) + avg = jax.tree.map(lambda x: (x / config.num_diloco_replicas).astype(x.dtype), total) + return avg + + @drjax.program(placements={"diloco": config.num_diloco_replicas}) + def diloco_train_step(state, batch, prng): + # Broadcast the RNG across replicas. + broadcast_rng = drjax.broadcast(prng) + inner_state, metrics = drjax.map_fn(train_step, (state.inner_state, batch, broadcast_rng)) + avg_metrics = typed_reduce_mean(metrics) + state = state.replace( + inner_state=inner_state, + step=inner_state.step[0], + ) + # Either synchronize the model, or no-op, depending on whether the current + # step falls on the synchronization period. + state = jax.lax.cond( + inner_state.step[0] % config.diloco_sync_period == 0, + synchronize, + lambda x: x, # no-op + state, + ) + return state, avg_metrics + + return diloco_train_step diff --git a/src/maxtext/utils/maxtext_utils.py b/src/maxtext/utils/maxtext_utils.py index 197b23fcdb..c4bb32aae0 100644 --- a/src/maxtext/utils/maxtext_utils.py +++ b/src/maxtext/utils/maxtext_utils.py @@ -128,7 +128,14 @@ def get_reorder_callable(cp_size, shard_mode): def get_shaped_batch(config): """Return the shape of the batch - this is what eval_shape would return for the output of create_data_iterator, but eval_shape doesn't work, see b/306901078.""" - batch_shape = (config.global_batch_size_to_load, config.max_target_length) + if config.enable_diloco: + batch_shape = ( + config.num_diloco_replicas, + config.global_batch_size_to_load // config.num_diloco_replicas, + config.max_target_length, + ) + else: + batch_shape = (config.global_batch_size_to_load, config.max_target_length) shaped_batch = {} shaped_batch["inputs"] = jax.ShapeDtypeStruct(batch_shape, jnp.int32) shaped_batch["inputs_position"] = jax.ShapeDtypeStruct(batch_shape, jnp.int32) diff --git a/src/maxtext/utils/train_utils.py b/src/maxtext/utils/train_utils.py index b53926aee6..ecc66aa9c2 100644 --- a/src/maxtext/utils/train_utils.py +++ b/src/maxtext/utils/train_utils.py @@ -17,6 +17,8 @@ import os import jax +import functools +from flax.linen import partitioning as nn_partitioning from MaxText import sharding from MaxText import optimizers from MaxText.rampup_batch import create_rampup_manager @@ -28,6 +30,7 @@ from maxtext.utils import max_utils from maxtext.utils import maxtext_utils from maxtext.utils import model_creation_utils +from maxtext.trainers.diloco import diloco def create_training_tools(config, model, mesh): @@ -83,15 +86,22 @@ def create_training_tools(config, model, mesh): def jit_train_step(config, model, state, state_mesh_shardings, data_sharding, train_step, params_shardings): """Returns a JIT-compiled train step function, which is loaded from a file if specified in the config.""" - ( - functional_train, - in_shardings, - out_shardings, - static_argnums, - donate_argnums, - ) = maxtext_utils.get_functional_train_with_signature( - train_step, data_sharding, state_mesh_shardings, model, config, params_shardings - ) + if config.enable_diloco: + functional_train = train_step + in_shardings = (state_mesh_shardings, data_sharding, None) # State, batch, rng + out_shardings = (state_mesh_shardings, None) # State, metrics + static_argnums = () # We partial out the static argnums of model and config + donate_argnums = 0 # This is the index of the state - we allow the compiler to make use of this memory. + else: + ( + functional_train, + in_shardings, + out_shardings, + static_argnums, + donate_argnums, + ) = maxtext_utils.get_functional_train_with_signature( + train_step, data_sharding, state_mesh_shardings, model, config, params_shardings + ) # Define the compilation of functional_train, either by loading the compiled version or wrapping a new one in a jit if config.compiled_trainstep_file != "": @@ -147,6 +157,9 @@ def jit_train_and_eval_step( params_shardings=None, ): """Returns a JIT-compiled train and eval step function.""" + if config.enable_diloco: + train_step_partial = functools.partial(train_step, model, config, state_mesh_shardings, params_shardings) + train_step = diloco.build_diloco_train_step(config, train_step_partial) data_sharding = sharding.get_input_data_sharding(config, mesh) p_train_step = jit_train_step(config, model, state, state_mesh_shardings, data_sharding, train_step, params_shardings) p_eval_step = None @@ -211,6 +224,19 @@ def setup_train_loop(config, recorder, devices=None): model, data_iterator, tx, config, init_rng, mesh, checkpoint_manager ) + if config.enable_diloco: + with jax.set_mesh(mesh), nn_partitioning.axis_rules(config.logical_axis_rules): + state, outer_opt_state_sharding = diloco.build_diloco_state(config, lambda: state) + + # create state_mesh_shardings for the DilocoState + inner_state_shardings = diloco.add_diloco_to_sharding(state_mesh_shardings) + state_mesh_shardings = diloco.DiLoCoTrainState( + inner_state_shardings, + state_mesh_shardings.params, + outer_opt_state_sharding, + jax.sharding.NamedSharding(mesh=state_mesh_shardings.step.mesh, spec=jax.sharding.PartitionSpec()), + ) + # TODO(aireenmei, hengtaoguo): support sharding in vit for multimodal if not config.using_pipeline_parallelism and not config.use_multimodal: # The vocab tensor(s) of shape [vocab, embed] (and transpose) are not sharded by stage diff --git a/tests/unit/diloco_test.py b/tests/unit/diloco_test.py new file mode 100644 index 0000000000..d80020d244 --- /dev/null +++ b/tests/unit/diloco_test.py @@ -0,0 +1,287 @@ +# Copyright 2025 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Tests for the DiLoCo implementation in diloco.py""" + + +import os +import unittest +from tempfile import gettempdir + +import chex +from flax.experimental import nnx +from flax.training import train_state +import jax +import jax.numpy as jnp +import jax.sharding +import numpy as np +import optax +import pytest + +from MaxText.pyconfig import initialize_pydantic +from MaxText.train_compile import main as train_compile_main +from maxtext.trainers.diloco import diloco +from tests.utils.test_helpers import get_test_config_path + + +class SimpleNNXModel(nnx.Module): + """A simple state for testing a minimal model.""" + + def __init__(self, *, rngs: nnx.Rngs): + self.dense = nnx.Linear( + 2, + 1, + kernel_init=nnx.initializers.constant(jnp.asarray([[2.0], [1.0]])), + bias_init=nnx.initializers.ones_init(), + rngs=rngs, + ) + + def __call__(self, x): + return self.dense(x) + + +class DiLoCoTest(unittest.TestCase): + + @pytest.mark.tpu_only + def test_diloco_training_simulation_with_mesh(self): + """Runs a simulation of DiLoCo training on a mesh and asserts correctness.""" + num_replicas = 2 + num_steps = 4 + + devices = jax.devices() + if len(devices) < num_replicas: + self.skipTest(f"Test requires {num_replicas} devices, but only {len(devices)} are available.") + + mesh_devices = np.array(devices[:num_replicas]).reshape(1, num_replicas) + mesh = jax.sharding.Mesh(mesh_devices, axis_names=("data", "diloco")) + + test_config = initialize_pydantic( + [ + "", + get_test_config_path(), + f"dcn_diloco_parallelism={num_replicas}", + "ici_diloco_parallelism=1", + "diloco_outer_momentum=0.9", + "diloco_outer_lr=1.0", + f"diloco_sync_period={num_steps-1}", + ] + ) + + with mesh: + tx = optax.sgd(learning_rate=0.1) + rngs = nnx.Rngs(params=jax.random.key(seed=42)) + model = SimpleNNXModel(rngs=rngs) + graphdef, params = nnx.split(model) + + def nnx_apply_fn(params, inputs): + model_replica = nnx.merge(graphdef, params) + return model_replica(inputs) + + # 2. Vmap this new wrapper function + vmapped_apply = jax.vmap(nnx_apply_fn, in_axes=(None, 0)) + + def _test_train_step(state: train_state.TrainState, batch, prng_key: diloco.PRNGKey): + """A simple MSE loss train step to enable numerics testing.""" + del prng_key + + def loss_fn(params, batch): + inputs, labels = batch + logits = vmapped_apply(params, inputs) + residual = logits - labels + sq_residual = jnp.square(residual) + msq_residual = jnp.mean(sq_residual) + return msq_residual + + loss, grad = jax.value_and_grad(loss_fn)(state.params, batch) + return state.apply_gradients(grads=grad), loss + + initial_test_state = train_state.TrainState.create( + apply_fn=vmapped_apply, + params=params, + tx=tx, + ) + + diloco_test_state, _ = diloco.build_diloco_state(test_config, lambda: initial_test_state) + chex.assert_equal(diloco_test_state.step, 0) + chex.assert_trees_all_equal(diloco_test_state.params, initial_test_state.params) + + diloco_train_step = diloco.build_diloco_train_step(test_config, _test_train_step) + inputs = jnp.array( + [ + [[0.0, 1.0], [1.0, 0.0]], # First replica inputs. + [[1.0, 0.0], [0.0, 1.0]], # Second replica inputs. + ] + ) + labels = jnp.array( + [ + [[1.0], [2.0]], # First replica labels. + [[2.0], [3.0]], # Second replica labels. + ] + ) + + sharding = jax.sharding.NamedSharding(mesh, jax.sharding.PartitionSpec(None, "diloco")) + inputs = jax.device_put(inputs, sharding) + labels = jax.device_put(labels, sharding) + + # Run the first step (no synchronization). + # Replica 0: + # Data: [[0, 1], [1, 0]] + # Labels: [[1], [2]] + # Weights: w = [[2], [1]] + # Bias: b = [1] + # Loss = mean((y - pred)^2) = + # = mean( ([[1], [2]] - (x . w + b)) ^ 2 ) ) + # = mean( ([[1], [2]] - ([[0, 1], [1, 0]] . [[2], [1]] + [1])) ^ 2 ) + # = mean( ([[1], [2]] - [[2], [3]]) ^ 2 ) + # = mean( ([-1, 1]) ^ 2 ) = mean( [1, 1] ) + # = 1.0 + # + # Replica 1: + # Data: [[1, 0], [0, 1]] + # Labels: [[2], [3]] + # Weights: w = [[2], [1]] + # Bias: b = [1] + # Loss = mean((y - pred)^2) = + # = mean( ([[2], [3]] - (x . w + b)) ^ 2 ) ) + # = mean( ([[2], [3]] - ([[1, 0], [0, 1]] . [[2], [1]] + [1])) ^ 2 ) + # = mean( ([[2], [3]] - [[3], [2]]) ^ 2 ) + # = mean( ([-1, 1]) ^ 2 ) = mean( [1, 1] ) + # = 1.0 + diloco_test_state, loss = diloco_train_step(diloco_test_state, (inputs, labels), jax.random.key(seed=42)) + chex.assert_equal(diloco_test_state.step, 1.0) + chex.assert_equal(loss, 1.0) + # Assert no updates to the global model yet (no synchronization) + chex.assert_trees_all_equal(diloco_test_state.params, initial_test_state.params) + + # Run the second step (no synchronization). + # Replica 0: + # Data: [[0, 1], [1, 0]] + # Labels: [[1], [2]] + # Weights: w = [[1.9], [0.9]] + # Bias: b = [0.8] + # Loss = mean((y - pred)^2) = + # = mean( ([[1], [2]] - (x . w + b)) ^ 2 ) ) + # = mean( ([[1], [2]] - ([[0, 1], [1, 0]] . [[1.9], [0.9]] + [0.8])) ^ 2 ) + # = mean( ([[1], [2]] - [[1.7], [2.7]]) ^ 2 ) + # = mean( ([-0.7, 0.7]) ^ 2 ) = mean( [0.49, 0.49] ) + # = 0.49 + # + # Replica 1: + # Data: [[1, 0], [0, 1]] + # Labels: [[2], [3]] + # Weights: w = [[1.9], [1.1]] + # Bias: b = [1] + # Loss = mean((y - pred)^2) = + # = mean( ([[2], [3]] - (x . w + b)) ^ 2 ) ) + # = mean( ([[2], [3]] - ([[1, 0], [0, 1]] . [[1.9], [1.1]] + [1])) ^ 2 ) + # = mean( ([[2], [3]] - [[2.9], [2.1]]) ^ 2 ) + # = mean( ([-0.9, 0.9]) ^ 2 ) = mean( [0.81, 0.81] ) + # = 0.81 + diloco_test_state, loss = diloco_train_step(diloco_test_state, (inputs, labels), jax.random.key(seed=42)) + chex.assert_equal(diloco_test_state.step, 2.0) + chex.assert_trees_all_close(loss, 0.65) + # Assert no updates to the global model yet (no synchronization) + chex.assert_trees_all_equal(diloco_test_state.params, initial_test_state.params) + + # Run the third step, which synchronizes afterwards. + # Replica 0: + # Data: [[0, 1], [1, 0]] + # Labels: [[1], [2]] + # Weights: w = [[1.83], [0.83]] + # Bias: b = [0.66] + # Loss = mean((y - pred)^2) = + # = mean( ([[1], [2]] - (x . w + b)) ^ 2 ) ) + # = mean( ([[1], [2]] - ([[0, 1], [1, 0]] . [[1.83], [0.83]] + [0.66])) ^ 2 ) + # = mean( ([[1], [2]] - [[1.49], [2.49]]) ^ 2 ) + # = mean( ([-0.49, 0.49]) ^ 2 ) = mean( [0.2401, 0.2401] ) + # = 0.2401 + # + # Replica 1: + # Data: [[1, 0], [0, 1]] + # Labels: [[2], [3]] + # Weights: w = [[1.81], [1.19]] + # Bias: b = [1.] + # Loss = mean((y - pred)^2) = + # = mean( ([[2], [3]] - (x . w + b)) ^ 2 ) ) + # = mean( ([[2], [3]] - ([[1, 0], [0, 1]] . [[1.81], [1.19]] + [1])) ^ 2 ) + # = mean( ([[2], [3]] - [[2.81], [2.19]]) ^ 2 ) + # = mean( ([-0.81, 0.81]) ^ 2 ) = mean( [0.6561, 0.6561] ) + # = 0.6561 + # + # After these are averaged, the model differences are computed to create a + # pseudo-gradient update to the outer_params and applied via a momentum + # based outer optimizer. + diloco_test_state, loss = diloco_train_step(diloco_test_state, (inputs, labels), jax.random.key(seed=42)) + chex.assert_equal(diloco_test_state.step, 3.0) + chex.assert_trees_all_close(loss, 0.4481) + # Assert that inner and outer parameters are all equal now that + # synchronization has happened. + chex.assert_trees_all_equal( + diloco_test_state.params, + jax.tree.map(lambda arr: arr[0, ...], diloco_test_state.inner_state.params), + ) + chex.assert_trees_all_equal( + diloco_test_state.params, + jax.tree.map(lambda arr: arr[1, ...], diloco_test_state.inner_state.params), + ) + + # Run the fourth step (no synchronization). + # Replica 0: + # Data: [[0, 1], [1, 0]] + # Labels: [[1], [2]] + # Weights: w = [[1.5345], [1.0494]] + # Bias: b = [0.5839] + # Loss = mean((y - pred)^2) = + # = mean( ([[1], [2]] - (x . w + b)) ^ 2 ) ) + # = mean( ([[1], [2]] - ([[0, 1], [1, 0]] . [[1.5345], [1.0494]]] + [0.5839])) ^ 2 ) + # = mean( ([[1], [2]] - [[1.6333], [2.1184]]) ^ 2 ) + # = mean( ([-0.6333, 0.1184]) ^ 2 ) = mean( [0.4010, 0.0140] ) + # ~ 0.2075 + # + # Replica 1: + # Data: [[1, 0], [0, 1]] + # Labels: [[2], [3]] + # Weights: w = [[1.5345], [1.0494]] + # Bias: b = [0.5839] + # Loss = mean((y - pred)^2) = + # = mean( ([[2], [3]] - (x . w + b)) ^ 2 ) ) + # = mean( ([[2], [3]] - ([[1, 0], [0, 1]] . [[1.5345], [1.0494]] + [0.5839])) ^ 2 ) + # = mean( ([[2], [3]] - [[2.1184], [1.6333]]) ^ 2 ) + # = mean( ([-0.1184, 1.3667]) ^ 2 ) = mean( [0.0140, 1.8678] ) + # ~ 0.94 + step_three_outer_params = diloco_test_state.params + diloco_test_state, loss = diloco_train_step(diloco_test_state, (inputs, labels), jax.random.key(seed=42)) + chex.assert_equal(diloco_test_state.step, 4.0) + chex.assert_trees_all_close(loss, 0.574244) + # Assert no updates to the global model since previous step (no + # synchronization). + chex.assert_trees_all_equal(diloco_test_state.params, step_three_outer_params) + + @pytest.mark.tpu_only + def test_diloco_two_slices(self): + temp_dir = gettempdir() + compiled_trainstep_file = os.path.join(temp_dir, "test_compiled_diloco.pickle") + train_compile_main( + ( + None, + get_test_config_path(), + f"compiled_trainstep_file={compiled_trainstep_file}", + "compile_topology=tpu7x-8", + "compile_topology_num_slices=2", + "ici_fsdp_parallelism=-1", + "dcn_diloco_parallelism=2", + "enable_diloco=true", + "model_name=gemma2-2b", + ) + ) diff --git a/tests/utils/sharding_info/deepseek2-16b/tpu7x-16/slice_1/named_shardings.json b/tests/utils/sharding_info/deepseek2-16b/tpu7x-16/slice_1/named_shardings.json index ed09ed2037..1fd6ceb6fd 100644 --- a/tests/utils/sharding_info/deepseek2-16b/tpu7x-16/slice_1/named_shardings.json +++ b/tests/utils/sharding_info/deepseek2-16b/tpu7x-16/slice_1/named_shardings.json @@ -2,6 +2,7 @@ ".step": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -16,6 +17,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -36,6 +38,7 @@ ".params/['params']/['decoder']/['decoder_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -50,6 +53,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -77,6 +81,7 @@ ".params/['params']/['decoder']/['dense_layers']/['mlp']/['wi_0']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -91,6 +96,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -130,6 +136,7 @@ ".params/['params']/['decoder']/['dense_layers']/['mlp']/['wi_1']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -144,6 +151,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -183,6 +191,7 @@ ".params/['params']/['decoder']/['dense_layers']/['mlp']/['wo']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -197,6 +206,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -236,6 +246,7 @@ ".params/['params']/['decoder']/['dense_layers']/['post_self_attention_layer_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -250,6 +261,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -279,6 +291,7 @@ ".params/['params']/['decoder']/['dense_layers']/['pre_self_attention_layer_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -293,6 +306,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -322,6 +336,7 @@ ".params/['params']/['decoder']/['dense_layers']/['self_attention']/['kv_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -336,6 +351,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -365,6 +381,7 @@ ".params/['params']/['decoder']/['dense_layers']/['self_attention']/['out']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -379,6 +396,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -420,6 +438,7 @@ ".params/['params']/['decoder']/['dense_layers']/['self_attention']/['query']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -434,6 +453,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -475,6 +495,7 @@ ".params/['params']/['decoder']/['dense_layers']/['self_attention']/['wkv_a']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -489,6 +510,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -524,6 +546,7 @@ ".params/['params']/['decoder']/['dense_layers']/['self_attention']/['wkv_b']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -538,6 +561,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -579,6 +603,7 @@ ".params/['params']/['decoder']/['logits_dense']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -593,6 +618,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -630,6 +656,7 @@ ".params/['params']/['decoder']/['moe_layers']/['DeepSeekMoeBlock_0']/['MoeBlock_0']/['gate']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -644,6 +671,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -679,6 +707,7 @@ ".params/['params']/['decoder']/['moe_layers']/['DeepSeekMoeBlock_0']/['MoeBlock_0']/['wi_0']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -693,6 +722,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -733,6 +763,7 @@ ".params/['params']/['decoder']/['moe_layers']/['DeepSeekMoeBlock_0']/['MoeBlock_0']/['wi_1']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -747,6 +778,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -787,6 +819,7 @@ ".params/['params']/['decoder']/['moe_layers']/['DeepSeekMoeBlock_0']/['MoeBlock_0']/['wo']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -801,6 +834,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -841,6 +875,7 @@ ".params/['params']/['decoder']/['moe_layers']/['DeepSeekMoeBlock_0']/['shared_experts']/['wi_0']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -855,6 +890,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -894,6 +930,7 @@ ".params/['params']/['decoder']/['moe_layers']/['DeepSeekMoeBlock_0']/['shared_experts']/['wi_1']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -908,6 +945,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -947,6 +985,7 @@ ".params/['params']/['decoder']/['moe_layers']/['DeepSeekMoeBlock_0']/['shared_experts']/['wo']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -961,6 +1000,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -1000,6 +1040,7 @@ ".params/['params']/['decoder']/['moe_layers']/['post_self_attention_layer_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1014,6 +1055,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -1043,6 +1085,7 @@ ".params/['params']/['decoder']/['moe_layers']/['pre_self_attention_layer_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1057,6 +1100,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -1086,6 +1130,7 @@ ".params/['params']/['decoder']/['moe_layers']/['self_attention']/['kv_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1100,6 +1145,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -1129,6 +1175,7 @@ ".params/['params']/['decoder']/['moe_layers']/['self_attention']/['out']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1143,6 +1190,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -1184,6 +1232,7 @@ ".params/['params']/['decoder']/['moe_layers']/['self_attention']/['query']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1198,6 +1247,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -1239,6 +1289,7 @@ ".params/['params']/['decoder']/['moe_layers']/['self_attention']/['wkv_a']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1253,6 +1304,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -1288,6 +1340,7 @@ ".params/['params']/['decoder']/['moe_layers']/['self_attention']/['wkv_b']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1302,6 +1355,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -1343,6 +1397,7 @@ ".params/['params']/['token_embedder']/['embedding']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1357,6 +1412,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -1394,6 +1450,7 @@ ".opt_state/[0]/.count": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1408,6 +1465,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -1428,6 +1486,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['decoder_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1442,6 +1501,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -1469,6 +1529,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['dense_layers']/['mlp']/['wi_0']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1483,6 +1544,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -1522,6 +1584,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['dense_layers']/['mlp']/['wi_1']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1536,6 +1599,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -1575,6 +1639,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['dense_layers']/['mlp']/['wo']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1589,6 +1654,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -1628,6 +1694,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['dense_layers']/['post_self_attention_layer_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1642,6 +1709,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -1671,6 +1739,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['dense_layers']/['pre_self_attention_layer_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1685,6 +1754,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -1714,6 +1784,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['dense_layers']/['self_attention']/['kv_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1728,6 +1799,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -1757,6 +1829,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['dense_layers']/['self_attention']/['out']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1771,6 +1844,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -1812,6 +1886,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['dense_layers']/['self_attention']/['query']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1826,6 +1901,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -1867,6 +1943,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['dense_layers']/['self_attention']/['wkv_a']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1881,6 +1958,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -1916,6 +1994,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['dense_layers']/['self_attention']/['wkv_b']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1930,6 +2009,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -1971,6 +2051,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['logits_dense']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1985,6 +2066,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -2022,6 +2104,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['moe_layers']/['DeepSeekMoeBlock_0']/['MoeBlock_0']/['gate']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2036,6 +2119,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -2071,6 +2155,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['moe_layers']/['DeepSeekMoeBlock_0']/['MoeBlock_0']/['wi_0']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2085,6 +2170,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -2125,6 +2211,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['moe_layers']/['DeepSeekMoeBlock_0']/['MoeBlock_0']/['wi_1']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2139,6 +2226,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -2179,6 +2267,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['moe_layers']/['DeepSeekMoeBlock_0']/['MoeBlock_0']/['wo']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2193,6 +2282,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -2233,6 +2323,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['moe_layers']/['DeepSeekMoeBlock_0']/['shared_experts']/['wi_0']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2247,6 +2338,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -2286,6 +2378,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['moe_layers']/['DeepSeekMoeBlock_0']/['shared_experts']/['wi_1']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2300,6 +2393,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -2339,6 +2433,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['moe_layers']/['DeepSeekMoeBlock_0']/['shared_experts']/['wo']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2353,6 +2448,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -2392,6 +2488,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['moe_layers']/['post_self_attention_layer_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2406,6 +2503,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -2435,6 +2533,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['moe_layers']/['pre_self_attention_layer_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2449,6 +2548,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -2478,6 +2578,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['moe_layers']/['self_attention']/['kv_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2492,6 +2593,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -2521,6 +2623,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['moe_layers']/['self_attention']/['out']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2535,6 +2638,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -2576,6 +2680,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['moe_layers']/['self_attention']/['query']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2590,6 +2695,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -2631,6 +2737,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['moe_layers']/['self_attention']/['wkv_a']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2645,6 +2752,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -2680,6 +2788,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['moe_layers']/['self_attention']/['wkv_b']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2694,6 +2803,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -2735,6 +2845,7 @@ ".opt_state/[0]/.mu/['params']/['token_embedder']/['embedding']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2749,6 +2860,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -2786,6 +2898,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['decoder_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2800,6 +2913,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -2827,6 +2941,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['dense_layers']/['mlp']/['wi_0']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2841,6 +2956,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -2880,6 +2996,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['dense_layers']/['mlp']/['wi_1']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2894,6 +3011,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -2933,6 +3051,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['dense_layers']/['mlp']/['wo']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2947,6 +3066,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -2986,6 +3106,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['dense_layers']/['post_self_attention_layer_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3000,6 +3121,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -3029,6 +3151,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['dense_layers']/['pre_self_attention_layer_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3043,6 +3166,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -3072,6 +3196,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['dense_layers']/['self_attention']/['kv_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3086,6 +3211,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -3115,6 +3241,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['dense_layers']/['self_attention']/['out']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3129,6 +3256,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -3170,6 +3298,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['dense_layers']/['self_attention']/['query']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3184,6 +3313,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -3225,6 +3355,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['dense_layers']/['self_attention']/['wkv_a']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3239,6 +3370,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -3274,6 +3406,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['dense_layers']/['self_attention']/['wkv_b']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3288,6 +3421,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -3329,6 +3463,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['logits_dense']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3343,6 +3478,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -3380,6 +3516,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['moe_layers']/['DeepSeekMoeBlock_0']/['MoeBlock_0']/['gate']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3394,6 +3531,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -3429,6 +3567,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['moe_layers']/['DeepSeekMoeBlock_0']/['MoeBlock_0']/['wi_0']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3443,6 +3582,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -3483,6 +3623,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['moe_layers']/['DeepSeekMoeBlock_0']/['MoeBlock_0']/['wi_1']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3497,6 +3638,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -3537,6 +3679,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['moe_layers']/['DeepSeekMoeBlock_0']/['MoeBlock_0']/['wo']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3551,6 +3694,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -3591,6 +3735,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['moe_layers']/['DeepSeekMoeBlock_0']/['shared_experts']/['wi_0']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3605,6 +3750,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -3644,6 +3790,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['moe_layers']/['DeepSeekMoeBlock_0']/['shared_experts']/['wi_1']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3658,6 +3805,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -3697,6 +3845,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['moe_layers']/['DeepSeekMoeBlock_0']/['shared_experts']/['wo']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3711,6 +3860,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -3750,6 +3900,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['moe_layers']/['post_self_attention_layer_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3764,6 +3915,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -3793,6 +3945,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['moe_layers']/['pre_self_attention_layer_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3807,6 +3960,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -3836,6 +3990,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['moe_layers']/['self_attention']/['kv_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3850,6 +4005,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -3879,6 +4035,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['moe_layers']/['self_attention']/['out']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3893,6 +4050,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -3934,6 +4092,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['moe_layers']/['self_attention']/['query']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3948,6 +4107,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -3989,6 +4149,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['moe_layers']/['self_attention']/['wkv_a']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -4003,6 +4164,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -4038,6 +4200,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['moe_layers']/['self_attention']/['wkv_b']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -4052,6 +4215,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -4093,6 +4257,7 @@ ".opt_state/[0]/.nu/['params']/['token_embedder']/['embedding']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -4107,6 +4272,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -4144,6 +4310,7 @@ ".opt_state/[2]/.count": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -4158,6 +4325,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, diff --git a/tests/utils/sharding_info/deepseek2-16b/tpu7x-16/slice_4/named_shardings.json b/tests/utils/sharding_info/deepseek2-16b/tpu7x-16/slice_4/named_shardings.json index a7fa362422..5b2ab94daf 100644 --- a/tests/utils/sharding_info/deepseek2-16b/tpu7x-16/slice_4/named_shardings.json +++ b/tests/utils/sharding_info/deepseek2-16b/tpu7x-16/slice_4/named_shardings.json @@ -2,6 +2,7 @@ ".step": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -16,6 +17,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -36,6 +38,7 @@ ".params/['params']/['decoder']/['decoder_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -50,6 +53,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -77,6 +81,7 @@ ".params/['params']/['decoder']/['dense_layers']/['mlp']/['wi_0']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -91,6 +96,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -130,6 +136,7 @@ ".params/['params']/['decoder']/['dense_layers']/['mlp']/['wi_1']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -144,6 +151,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -183,6 +191,7 @@ ".params/['params']/['decoder']/['dense_layers']/['mlp']/['wo']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -197,6 +206,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -236,6 +246,7 @@ ".params/['params']/['decoder']/['dense_layers']/['post_self_attention_layer_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -250,6 +261,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -279,6 +291,7 @@ ".params/['params']/['decoder']/['dense_layers']/['pre_self_attention_layer_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -293,6 +306,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -322,6 +336,7 @@ ".params/['params']/['decoder']/['dense_layers']/['self_attention']/['kv_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -336,6 +351,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -365,6 +381,7 @@ ".params/['params']/['decoder']/['dense_layers']/['self_attention']/['out']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -379,6 +396,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -420,6 +438,7 @@ ".params/['params']/['decoder']/['dense_layers']/['self_attention']/['query']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -434,6 +453,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -475,6 +495,7 @@ ".params/['params']/['decoder']/['dense_layers']/['self_attention']/['wkv_a']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -489,6 +510,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -524,6 +546,7 @@ ".params/['params']/['decoder']/['dense_layers']/['self_attention']/['wkv_b']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -538,6 +561,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -579,6 +603,7 @@ ".params/['params']/['decoder']/['logits_dense']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -593,6 +618,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -630,6 +656,7 @@ ".params/['params']/['decoder']/['moe_layers']/['DeepSeekMoeBlock_0']/['MoeBlock_0']/['gate']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -644,6 +671,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -679,6 +707,7 @@ ".params/['params']/['decoder']/['moe_layers']/['DeepSeekMoeBlock_0']/['MoeBlock_0']/['wi_0']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -693,6 +722,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -733,6 +763,7 @@ ".params/['params']/['decoder']/['moe_layers']/['DeepSeekMoeBlock_0']/['MoeBlock_0']/['wi_1']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -747,6 +778,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -787,6 +819,7 @@ ".params/['params']/['decoder']/['moe_layers']/['DeepSeekMoeBlock_0']/['MoeBlock_0']/['wo']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -801,6 +834,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -841,6 +875,7 @@ ".params/['params']/['decoder']/['moe_layers']/['DeepSeekMoeBlock_0']/['shared_experts']/['wi_0']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -855,6 +890,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -894,6 +930,7 @@ ".params/['params']/['decoder']/['moe_layers']/['DeepSeekMoeBlock_0']/['shared_experts']/['wi_1']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -908,6 +945,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -947,6 +985,7 @@ ".params/['params']/['decoder']/['moe_layers']/['DeepSeekMoeBlock_0']/['shared_experts']/['wo']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -961,6 +1000,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -1000,6 +1040,7 @@ ".params/['params']/['decoder']/['moe_layers']/['post_self_attention_layer_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1014,6 +1055,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -1043,6 +1085,7 @@ ".params/['params']/['decoder']/['moe_layers']/['pre_self_attention_layer_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1057,6 +1100,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -1086,6 +1130,7 @@ ".params/['params']/['decoder']/['moe_layers']/['self_attention']/['kv_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1100,6 +1145,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -1129,6 +1175,7 @@ ".params/['params']/['decoder']/['moe_layers']/['self_attention']/['out']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1143,6 +1190,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -1184,6 +1232,7 @@ ".params/['params']/['decoder']/['moe_layers']/['self_attention']/['query']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1198,6 +1247,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -1239,6 +1289,7 @@ ".params/['params']/['decoder']/['moe_layers']/['self_attention']/['wkv_a']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1253,6 +1304,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -1288,6 +1340,7 @@ ".params/['params']/['decoder']/['moe_layers']/['self_attention']/['wkv_b']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1302,6 +1355,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -1343,6 +1397,7 @@ ".params/['params']/['token_embedder']/['embedding']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1357,6 +1412,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -1394,6 +1450,7 @@ ".opt_state/[0]/.count": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1408,6 +1465,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -1428,6 +1486,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['decoder_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1442,6 +1501,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -1469,6 +1529,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['dense_layers']/['mlp']/['wi_0']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1483,6 +1544,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -1522,6 +1584,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['dense_layers']/['mlp']/['wi_1']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1536,6 +1599,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -1575,6 +1639,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['dense_layers']/['mlp']/['wo']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1589,6 +1654,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -1628,6 +1694,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['dense_layers']/['post_self_attention_layer_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1642,6 +1709,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -1671,6 +1739,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['dense_layers']/['pre_self_attention_layer_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1685,6 +1754,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -1714,6 +1784,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['dense_layers']/['self_attention']/['kv_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1728,6 +1799,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -1757,6 +1829,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['dense_layers']/['self_attention']/['out']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1771,6 +1844,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -1812,6 +1886,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['dense_layers']/['self_attention']/['query']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1826,6 +1901,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -1867,6 +1943,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['dense_layers']/['self_attention']/['wkv_a']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1881,6 +1958,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -1916,6 +1994,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['dense_layers']/['self_attention']/['wkv_b']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1930,6 +2009,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -1971,6 +2051,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['logits_dense']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1985,6 +2066,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -2022,6 +2104,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['moe_layers']/['DeepSeekMoeBlock_0']/['MoeBlock_0']/['gate']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2036,6 +2119,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -2071,6 +2155,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['moe_layers']/['DeepSeekMoeBlock_0']/['MoeBlock_0']/['wi_0']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2085,6 +2170,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -2125,6 +2211,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['moe_layers']/['DeepSeekMoeBlock_0']/['MoeBlock_0']/['wi_1']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2139,6 +2226,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -2179,6 +2267,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['moe_layers']/['DeepSeekMoeBlock_0']/['MoeBlock_0']/['wo']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2193,6 +2282,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -2233,6 +2323,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['moe_layers']/['DeepSeekMoeBlock_0']/['shared_experts']/['wi_0']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2247,6 +2338,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -2286,6 +2378,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['moe_layers']/['DeepSeekMoeBlock_0']/['shared_experts']/['wi_1']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2300,6 +2393,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -2339,6 +2433,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['moe_layers']/['DeepSeekMoeBlock_0']/['shared_experts']/['wo']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2353,6 +2448,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -2392,6 +2488,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['moe_layers']/['post_self_attention_layer_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2406,6 +2503,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -2435,6 +2533,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['moe_layers']/['pre_self_attention_layer_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2449,6 +2548,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -2478,6 +2578,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['moe_layers']/['self_attention']/['kv_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2492,6 +2593,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -2521,6 +2623,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['moe_layers']/['self_attention']/['out']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2535,6 +2638,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -2576,6 +2680,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['moe_layers']/['self_attention']/['query']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2590,6 +2695,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -2631,6 +2737,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['moe_layers']/['self_attention']/['wkv_a']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2645,6 +2752,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -2680,6 +2788,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['moe_layers']/['self_attention']/['wkv_b']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2694,6 +2803,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -2735,6 +2845,7 @@ ".opt_state/[0]/.mu/['params']/['token_embedder']/['embedding']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2749,6 +2860,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -2786,6 +2898,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['decoder_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2800,6 +2913,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -2827,6 +2941,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['dense_layers']/['mlp']/['wi_0']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2841,6 +2956,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -2880,6 +2996,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['dense_layers']/['mlp']/['wi_1']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2894,6 +3011,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -2933,6 +3051,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['dense_layers']/['mlp']/['wo']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2947,6 +3066,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -2986,6 +3106,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['dense_layers']/['post_self_attention_layer_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3000,6 +3121,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -3029,6 +3151,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['dense_layers']/['pre_self_attention_layer_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3043,6 +3166,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -3072,6 +3196,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['dense_layers']/['self_attention']/['kv_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3086,6 +3211,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -3115,6 +3241,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['dense_layers']/['self_attention']/['out']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3129,6 +3256,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -3170,6 +3298,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['dense_layers']/['self_attention']/['query']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3184,6 +3313,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -3225,6 +3355,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['dense_layers']/['self_attention']/['wkv_a']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3239,6 +3370,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -3274,6 +3406,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['dense_layers']/['self_attention']/['wkv_b']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3288,6 +3421,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -3329,6 +3463,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['logits_dense']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3343,6 +3478,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -3380,6 +3516,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['moe_layers']/['DeepSeekMoeBlock_0']/['MoeBlock_0']/['gate']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3394,6 +3531,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -3429,6 +3567,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['moe_layers']/['DeepSeekMoeBlock_0']/['MoeBlock_0']/['wi_0']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3443,6 +3582,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -3483,6 +3623,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['moe_layers']/['DeepSeekMoeBlock_0']/['MoeBlock_0']/['wi_1']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3497,6 +3638,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -3537,6 +3679,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['moe_layers']/['DeepSeekMoeBlock_0']/['MoeBlock_0']/['wo']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3551,6 +3694,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -3591,6 +3735,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['moe_layers']/['DeepSeekMoeBlock_0']/['shared_experts']/['wi_0']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3605,6 +3750,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -3644,6 +3790,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['moe_layers']/['DeepSeekMoeBlock_0']/['shared_experts']/['wi_1']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3658,6 +3805,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -3697,6 +3845,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['moe_layers']/['DeepSeekMoeBlock_0']/['shared_experts']/['wo']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3711,6 +3860,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -3750,6 +3900,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['moe_layers']/['post_self_attention_layer_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3764,6 +3915,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -3793,6 +3945,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['moe_layers']/['pre_self_attention_layer_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3807,6 +3960,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -3836,6 +3990,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['moe_layers']/['self_attention']/['kv_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3850,6 +4005,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -3879,6 +4035,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['moe_layers']/['self_attention']/['out']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3893,6 +4050,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -3934,6 +4092,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['moe_layers']/['self_attention']/['query']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3948,6 +4107,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -3989,6 +4149,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['moe_layers']/['self_attention']/['wkv_a']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -4003,6 +4164,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -4038,6 +4200,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['moe_layers']/['self_attention']/['wkv_b']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -4052,6 +4215,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -4093,6 +4257,7 @@ ".opt_state/[0]/.nu/['params']/['token_embedder']/['embedding']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -4107,6 +4272,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -4144,6 +4310,7 @@ ".opt_state/[2]/.count": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -4158,6 +4325,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, diff --git a/tests/utils/sharding_info/deepseek2-16b/v5p-16/slice_1/named_shardings.json b/tests/utils/sharding_info/deepseek2-16b/v5p-16/slice_1/named_shardings.json index a7e781f9c3..72cbbdea66 100644 --- a/tests/utils/sharding_info/deepseek2-16b/v5p-16/slice_1/named_shardings.json +++ b/tests/utils/sharding_info/deepseek2-16b/v5p-16/slice_1/named_shardings.json @@ -2,6 +2,7 @@ ".step": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -16,6 +17,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -36,6 +38,7 @@ ".params/['params']/['decoder']/['decoder_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -50,6 +53,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -77,6 +81,7 @@ ".params/['params']/['decoder']/['dense_layers']/['mlp']/['wi_0']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -91,6 +96,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -130,6 +136,7 @@ ".params/['params']/['decoder']/['dense_layers']/['mlp']/['wi_1']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -144,6 +151,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -183,6 +191,7 @@ ".params/['params']/['decoder']/['dense_layers']/['mlp']/['wo']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -197,6 +206,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -236,6 +246,7 @@ ".params/['params']/['decoder']/['dense_layers']/['post_self_attention_layer_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -250,6 +261,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -279,6 +291,7 @@ ".params/['params']/['decoder']/['dense_layers']/['pre_self_attention_layer_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -293,6 +306,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -322,6 +336,7 @@ ".params/['params']/['decoder']/['dense_layers']/['self_attention']/['kv_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -336,6 +351,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -365,6 +381,7 @@ ".params/['params']/['decoder']/['dense_layers']/['self_attention']/['out']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -379,6 +396,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -420,6 +438,7 @@ ".params/['params']/['decoder']/['dense_layers']/['self_attention']/['query']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -434,6 +453,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -475,6 +495,7 @@ ".params/['params']/['decoder']/['dense_layers']/['self_attention']/['wkv_a']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -489,6 +510,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -524,6 +546,7 @@ ".params/['params']/['decoder']/['dense_layers']/['self_attention']/['wkv_b']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -538,6 +561,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -579,6 +603,7 @@ ".params/['params']/['decoder']/['logits_dense']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -593,6 +618,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -630,6 +656,7 @@ ".params/['params']/['decoder']/['moe_layers']/['DeepSeekMoeBlock_0']/['MoeBlock_0']/['gate']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -644,6 +671,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -679,6 +707,7 @@ ".params/['params']/['decoder']/['moe_layers']/['DeepSeekMoeBlock_0']/['MoeBlock_0']/['wi_0']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -693,6 +722,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -733,6 +763,7 @@ ".params/['params']/['decoder']/['moe_layers']/['DeepSeekMoeBlock_0']/['MoeBlock_0']/['wi_1']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -747,6 +778,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -787,6 +819,7 @@ ".params/['params']/['decoder']/['moe_layers']/['DeepSeekMoeBlock_0']/['MoeBlock_0']/['wo']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -801,6 +834,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -841,6 +875,7 @@ ".params/['params']/['decoder']/['moe_layers']/['DeepSeekMoeBlock_0']/['shared_experts']/['wi_0']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -855,6 +890,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -894,6 +930,7 @@ ".params/['params']/['decoder']/['moe_layers']/['DeepSeekMoeBlock_0']/['shared_experts']/['wi_1']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -908,6 +945,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -947,6 +985,7 @@ ".params/['params']/['decoder']/['moe_layers']/['DeepSeekMoeBlock_0']/['shared_experts']/['wo']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -961,6 +1000,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -1000,6 +1040,7 @@ ".params/['params']/['decoder']/['moe_layers']/['post_self_attention_layer_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1014,6 +1055,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -1043,6 +1085,7 @@ ".params/['params']/['decoder']/['moe_layers']/['pre_self_attention_layer_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1057,6 +1100,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -1086,6 +1130,7 @@ ".params/['params']/['decoder']/['moe_layers']/['self_attention']/['kv_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1100,6 +1145,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -1129,6 +1175,7 @@ ".params/['params']/['decoder']/['moe_layers']/['self_attention']/['out']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1143,6 +1190,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -1184,6 +1232,7 @@ ".params/['params']/['decoder']/['moe_layers']/['self_attention']/['query']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1198,6 +1247,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -1239,6 +1289,7 @@ ".params/['params']/['decoder']/['moe_layers']/['self_attention']/['wkv_a']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1253,6 +1304,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -1288,6 +1340,7 @@ ".params/['params']/['decoder']/['moe_layers']/['self_attention']/['wkv_b']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1302,6 +1355,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -1343,6 +1397,7 @@ ".params/['params']/['token_embedder']/['embedding']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1357,6 +1412,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -1394,6 +1450,7 @@ ".opt_state/[0]/.count": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1408,6 +1465,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -1428,6 +1486,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['decoder_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1442,6 +1501,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -1469,6 +1529,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['dense_layers']/['mlp']/['wi_0']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1483,6 +1544,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -1522,6 +1584,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['dense_layers']/['mlp']/['wi_1']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1536,6 +1599,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -1575,6 +1639,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['dense_layers']/['mlp']/['wo']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1589,6 +1654,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -1628,6 +1694,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['dense_layers']/['post_self_attention_layer_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1642,6 +1709,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -1671,6 +1739,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['dense_layers']/['pre_self_attention_layer_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1685,6 +1754,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -1714,6 +1784,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['dense_layers']/['self_attention']/['kv_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1728,6 +1799,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -1757,6 +1829,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['dense_layers']/['self_attention']/['out']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1771,6 +1844,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -1812,6 +1886,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['dense_layers']/['self_attention']/['query']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1826,6 +1901,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -1867,6 +1943,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['dense_layers']/['self_attention']/['wkv_a']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1881,6 +1958,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -1916,6 +1994,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['dense_layers']/['self_attention']/['wkv_b']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1930,6 +2009,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -1971,6 +2051,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['logits_dense']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1985,6 +2066,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -2022,6 +2104,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['moe_layers']/['DeepSeekMoeBlock_0']/['MoeBlock_0']/['gate']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2036,6 +2119,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -2071,6 +2155,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['moe_layers']/['DeepSeekMoeBlock_0']/['MoeBlock_0']/['wi_0']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2085,6 +2170,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -2125,6 +2211,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['moe_layers']/['DeepSeekMoeBlock_0']/['MoeBlock_0']/['wi_1']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2139,6 +2226,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -2179,6 +2267,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['moe_layers']/['DeepSeekMoeBlock_0']/['MoeBlock_0']/['wo']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2193,6 +2282,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -2233,6 +2323,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['moe_layers']/['DeepSeekMoeBlock_0']/['shared_experts']/['wi_0']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2247,6 +2338,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -2286,6 +2378,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['moe_layers']/['DeepSeekMoeBlock_0']/['shared_experts']/['wi_1']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2300,6 +2393,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -2339,6 +2433,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['moe_layers']/['DeepSeekMoeBlock_0']/['shared_experts']/['wo']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2353,6 +2448,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -2392,6 +2488,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['moe_layers']/['post_self_attention_layer_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2406,6 +2503,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -2435,6 +2533,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['moe_layers']/['pre_self_attention_layer_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2449,6 +2548,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -2478,6 +2578,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['moe_layers']/['self_attention']/['kv_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2492,6 +2593,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -2521,6 +2623,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['moe_layers']/['self_attention']/['out']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2535,6 +2638,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -2576,6 +2680,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['moe_layers']/['self_attention']/['query']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2590,6 +2695,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -2631,6 +2737,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['moe_layers']/['self_attention']/['wkv_a']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2645,6 +2752,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -2680,6 +2788,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['moe_layers']/['self_attention']/['wkv_b']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2694,6 +2803,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -2735,6 +2845,7 @@ ".opt_state/[0]/.mu/['params']/['token_embedder']/['embedding']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2749,6 +2860,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -2786,6 +2898,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['decoder_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2800,6 +2913,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -2827,6 +2941,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['dense_layers']/['mlp']/['wi_0']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2841,6 +2956,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -2880,6 +2996,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['dense_layers']/['mlp']/['wi_1']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2894,6 +3011,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -2933,6 +3051,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['dense_layers']/['mlp']/['wo']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2947,6 +3066,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -2986,6 +3106,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['dense_layers']/['post_self_attention_layer_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3000,6 +3121,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -3029,6 +3151,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['dense_layers']/['pre_self_attention_layer_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3043,6 +3166,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -3072,6 +3196,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['dense_layers']/['self_attention']/['kv_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3086,6 +3211,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -3115,6 +3241,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['dense_layers']/['self_attention']/['out']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3129,6 +3256,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -3170,6 +3298,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['dense_layers']/['self_attention']/['query']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3184,6 +3313,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -3225,6 +3355,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['dense_layers']/['self_attention']/['wkv_a']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3239,6 +3370,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -3274,6 +3406,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['dense_layers']/['self_attention']/['wkv_b']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3288,6 +3421,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -3329,6 +3463,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['logits_dense']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3343,6 +3478,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -3380,6 +3516,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['moe_layers']/['DeepSeekMoeBlock_0']/['MoeBlock_0']/['gate']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3394,6 +3531,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -3429,6 +3567,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['moe_layers']/['DeepSeekMoeBlock_0']/['MoeBlock_0']/['wi_0']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3443,6 +3582,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -3483,6 +3623,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['moe_layers']/['DeepSeekMoeBlock_0']/['MoeBlock_0']/['wi_1']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3497,6 +3638,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -3537,6 +3679,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['moe_layers']/['DeepSeekMoeBlock_0']/['MoeBlock_0']/['wo']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3551,6 +3694,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -3591,6 +3735,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['moe_layers']/['DeepSeekMoeBlock_0']/['shared_experts']/['wi_0']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3605,6 +3750,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -3644,6 +3790,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['moe_layers']/['DeepSeekMoeBlock_0']/['shared_experts']/['wi_1']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3658,6 +3805,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -3697,6 +3845,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['moe_layers']/['DeepSeekMoeBlock_0']/['shared_experts']/['wo']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3711,6 +3860,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -3750,6 +3900,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['moe_layers']/['post_self_attention_layer_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3764,6 +3915,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -3793,6 +3945,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['moe_layers']/['pre_self_attention_layer_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3807,6 +3960,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -3836,6 +3990,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['moe_layers']/['self_attention']/['kv_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3850,6 +4005,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -3879,6 +4035,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['moe_layers']/['self_attention']/['out']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3893,6 +4050,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -3934,6 +4092,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['moe_layers']/['self_attention']/['query']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3948,6 +4107,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -3989,6 +4149,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['moe_layers']/['self_attention']/['wkv_a']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -4003,6 +4164,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -4038,6 +4200,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['moe_layers']/['self_attention']/['wkv_b']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -4052,6 +4215,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -4093,6 +4257,7 @@ ".opt_state/[0]/.nu/['params']/['token_embedder']/['embedding']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -4107,6 +4272,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -4144,6 +4310,7 @@ ".opt_state/[2]/.count": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -4158,6 +4325,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, diff --git a/tests/utils/sharding_info/deepseek2-16b/v5p-16/slice_4/named_shardings.json b/tests/utils/sharding_info/deepseek2-16b/v5p-16/slice_4/named_shardings.json index 19cd50adc3..65120bac91 100644 --- a/tests/utils/sharding_info/deepseek2-16b/v5p-16/slice_4/named_shardings.json +++ b/tests/utils/sharding_info/deepseek2-16b/v5p-16/slice_4/named_shardings.json @@ -2,6 +2,7 @@ ".step": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -16,6 +17,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -36,6 +38,7 @@ ".params/['params']/['decoder']/['decoder_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -50,6 +53,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -77,6 +81,7 @@ ".params/['params']/['decoder']/['dense_layers']/['mlp']/['wi_0']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -91,6 +96,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -130,6 +136,7 @@ ".params/['params']/['decoder']/['dense_layers']/['mlp']/['wi_1']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -144,6 +151,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -183,6 +191,7 @@ ".params/['params']/['decoder']/['dense_layers']/['mlp']/['wo']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -197,6 +206,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -236,6 +246,7 @@ ".params/['params']/['decoder']/['dense_layers']/['post_self_attention_layer_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -250,6 +261,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -279,6 +291,7 @@ ".params/['params']/['decoder']/['dense_layers']/['pre_self_attention_layer_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -293,6 +306,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -322,6 +336,7 @@ ".params/['params']/['decoder']/['dense_layers']/['self_attention']/['kv_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -336,6 +351,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -365,6 +381,7 @@ ".params/['params']/['decoder']/['dense_layers']/['self_attention']/['out']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -379,6 +396,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -420,6 +438,7 @@ ".params/['params']/['decoder']/['dense_layers']/['self_attention']/['query']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -434,6 +453,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -475,6 +495,7 @@ ".params/['params']/['decoder']/['dense_layers']/['self_attention']/['wkv_a']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -489,6 +510,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -524,6 +546,7 @@ ".params/['params']/['decoder']/['dense_layers']/['self_attention']/['wkv_b']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -538,6 +561,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -579,6 +603,7 @@ ".params/['params']/['decoder']/['logits_dense']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -593,6 +618,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -630,6 +656,7 @@ ".params/['params']/['decoder']/['moe_layers']/['DeepSeekMoeBlock_0']/['MoeBlock_0']/['gate']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -644,6 +671,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -679,6 +707,7 @@ ".params/['params']/['decoder']/['moe_layers']/['DeepSeekMoeBlock_0']/['MoeBlock_0']/['wi_0']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -693,6 +722,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -733,6 +763,7 @@ ".params/['params']/['decoder']/['moe_layers']/['DeepSeekMoeBlock_0']/['MoeBlock_0']/['wi_1']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -747,6 +778,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -787,6 +819,7 @@ ".params/['params']/['decoder']/['moe_layers']/['DeepSeekMoeBlock_0']/['MoeBlock_0']/['wo']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -801,6 +834,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -841,6 +875,7 @@ ".params/['params']/['decoder']/['moe_layers']/['DeepSeekMoeBlock_0']/['shared_experts']/['wi_0']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -855,6 +890,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -894,6 +930,7 @@ ".params/['params']/['decoder']/['moe_layers']/['DeepSeekMoeBlock_0']/['shared_experts']/['wi_1']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -908,6 +945,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -947,6 +985,7 @@ ".params/['params']/['decoder']/['moe_layers']/['DeepSeekMoeBlock_0']/['shared_experts']/['wo']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -961,6 +1000,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -1000,6 +1040,7 @@ ".params/['params']/['decoder']/['moe_layers']/['post_self_attention_layer_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1014,6 +1055,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -1043,6 +1085,7 @@ ".params/['params']/['decoder']/['moe_layers']/['pre_self_attention_layer_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1057,6 +1100,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -1086,6 +1130,7 @@ ".params/['params']/['decoder']/['moe_layers']/['self_attention']/['kv_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1100,6 +1145,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -1129,6 +1175,7 @@ ".params/['params']/['decoder']/['moe_layers']/['self_attention']/['out']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1143,6 +1190,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -1184,6 +1232,7 @@ ".params/['params']/['decoder']/['moe_layers']/['self_attention']/['query']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1198,6 +1247,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -1239,6 +1289,7 @@ ".params/['params']/['decoder']/['moe_layers']/['self_attention']/['wkv_a']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1253,6 +1304,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -1288,6 +1340,7 @@ ".params/['params']/['decoder']/['moe_layers']/['self_attention']/['wkv_b']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1302,6 +1355,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -1343,6 +1397,7 @@ ".params/['params']/['token_embedder']/['embedding']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1357,6 +1412,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -1394,6 +1450,7 @@ ".opt_state/[0]/.count": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1408,6 +1465,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -1428,6 +1486,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['decoder_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1442,6 +1501,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -1469,6 +1529,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['dense_layers']/['mlp']/['wi_0']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1483,6 +1544,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -1522,6 +1584,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['dense_layers']/['mlp']/['wi_1']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1536,6 +1599,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -1575,6 +1639,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['dense_layers']/['mlp']/['wo']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1589,6 +1654,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -1628,6 +1694,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['dense_layers']/['post_self_attention_layer_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1642,6 +1709,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -1671,6 +1739,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['dense_layers']/['pre_self_attention_layer_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1685,6 +1754,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -1714,6 +1784,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['dense_layers']/['self_attention']/['kv_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1728,6 +1799,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -1757,6 +1829,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['dense_layers']/['self_attention']/['out']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1771,6 +1844,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -1812,6 +1886,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['dense_layers']/['self_attention']/['query']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1826,6 +1901,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -1867,6 +1943,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['dense_layers']/['self_attention']/['wkv_a']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1881,6 +1958,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -1916,6 +1994,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['dense_layers']/['self_attention']/['wkv_b']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1930,6 +2009,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -1971,6 +2051,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['logits_dense']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1985,6 +2066,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -2022,6 +2104,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['moe_layers']/['DeepSeekMoeBlock_0']/['MoeBlock_0']/['gate']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2036,6 +2119,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -2071,6 +2155,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['moe_layers']/['DeepSeekMoeBlock_0']/['MoeBlock_0']/['wi_0']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2085,6 +2170,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -2125,6 +2211,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['moe_layers']/['DeepSeekMoeBlock_0']/['MoeBlock_0']/['wi_1']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2139,6 +2226,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -2179,6 +2267,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['moe_layers']/['DeepSeekMoeBlock_0']/['MoeBlock_0']/['wo']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2193,6 +2282,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -2233,6 +2323,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['moe_layers']/['DeepSeekMoeBlock_0']/['shared_experts']/['wi_0']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2247,6 +2338,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -2286,6 +2378,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['moe_layers']/['DeepSeekMoeBlock_0']/['shared_experts']/['wi_1']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2300,6 +2393,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -2339,6 +2433,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['moe_layers']/['DeepSeekMoeBlock_0']/['shared_experts']/['wo']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2353,6 +2448,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -2392,6 +2488,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['moe_layers']/['post_self_attention_layer_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2406,6 +2503,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -2435,6 +2533,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['moe_layers']/['pre_self_attention_layer_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2449,6 +2548,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -2478,6 +2578,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['moe_layers']/['self_attention']/['kv_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2492,6 +2593,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -2521,6 +2623,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['moe_layers']/['self_attention']/['out']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2535,6 +2638,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -2576,6 +2680,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['moe_layers']/['self_attention']/['query']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2590,6 +2695,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -2631,6 +2737,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['moe_layers']/['self_attention']/['wkv_a']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2645,6 +2752,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -2680,6 +2788,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['moe_layers']/['self_attention']/['wkv_b']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2694,6 +2803,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -2735,6 +2845,7 @@ ".opt_state/[0]/.mu/['params']/['token_embedder']/['embedding']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2749,6 +2860,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -2786,6 +2898,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['decoder_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2800,6 +2913,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -2827,6 +2941,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['dense_layers']/['mlp']/['wi_0']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2841,6 +2956,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -2880,6 +2996,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['dense_layers']/['mlp']/['wi_1']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2894,6 +3011,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -2933,6 +3051,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['dense_layers']/['mlp']/['wo']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2947,6 +3066,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -2986,6 +3106,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['dense_layers']/['post_self_attention_layer_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3000,6 +3121,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -3029,6 +3151,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['dense_layers']/['pre_self_attention_layer_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3043,6 +3166,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -3072,6 +3196,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['dense_layers']/['self_attention']/['kv_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3086,6 +3211,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -3115,6 +3241,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['dense_layers']/['self_attention']/['out']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3129,6 +3256,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -3170,6 +3298,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['dense_layers']/['self_attention']/['query']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3184,6 +3313,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -3225,6 +3355,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['dense_layers']/['self_attention']/['wkv_a']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3239,6 +3370,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -3274,6 +3406,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['dense_layers']/['self_attention']/['wkv_b']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3288,6 +3421,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -3329,6 +3463,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['logits_dense']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3343,6 +3478,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -3380,6 +3516,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['moe_layers']/['DeepSeekMoeBlock_0']/['MoeBlock_0']/['gate']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3394,6 +3531,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -3429,6 +3567,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['moe_layers']/['DeepSeekMoeBlock_0']/['MoeBlock_0']/['wi_0']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3443,6 +3582,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -3483,6 +3623,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['moe_layers']/['DeepSeekMoeBlock_0']/['MoeBlock_0']/['wi_1']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3497,6 +3638,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -3537,6 +3679,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['moe_layers']/['DeepSeekMoeBlock_0']/['MoeBlock_0']/['wo']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3551,6 +3694,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -3591,6 +3735,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['moe_layers']/['DeepSeekMoeBlock_0']/['shared_experts']/['wi_0']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3605,6 +3750,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -3644,6 +3790,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['moe_layers']/['DeepSeekMoeBlock_0']/['shared_experts']/['wi_1']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3658,6 +3805,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -3697,6 +3845,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['moe_layers']/['DeepSeekMoeBlock_0']/['shared_experts']/['wo']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3711,6 +3860,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -3750,6 +3900,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['moe_layers']/['post_self_attention_layer_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3764,6 +3915,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -3793,6 +3945,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['moe_layers']/['pre_self_attention_layer_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3807,6 +3960,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -3836,6 +3990,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['moe_layers']/['self_attention']/['kv_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3850,6 +4005,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -3879,6 +4035,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['moe_layers']/['self_attention']/['out']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3893,6 +4050,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -3934,6 +4092,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['moe_layers']/['self_attention']/['query']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3948,6 +4107,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -3989,6 +4149,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['moe_layers']/['self_attention']/['wkv_a']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -4003,6 +4164,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -4038,6 +4200,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['moe_layers']/['self_attention']/['wkv_b']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -4052,6 +4215,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -4093,6 +4257,7 @@ ".opt_state/[0]/.nu/['params']/['token_embedder']/['embedding']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -4107,6 +4272,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -4144,6 +4310,7 @@ ".opt_state/[2]/.count": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -4158,6 +4325,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, diff --git a/tests/utils/sharding_info/deepseek2-16b/v6e-16/slice_1/named_shardings.json b/tests/utils/sharding_info/deepseek2-16b/v6e-16/slice_1/named_shardings.json index ed09ed2037..1fd6ceb6fd 100644 --- a/tests/utils/sharding_info/deepseek2-16b/v6e-16/slice_1/named_shardings.json +++ b/tests/utils/sharding_info/deepseek2-16b/v6e-16/slice_1/named_shardings.json @@ -2,6 +2,7 @@ ".step": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -16,6 +17,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -36,6 +38,7 @@ ".params/['params']/['decoder']/['decoder_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -50,6 +53,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -77,6 +81,7 @@ ".params/['params']/['decoder']/['dense_layers']/['mlp']/['wi_0']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -91,6 +96,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -130,6 +136,7 @@ ".params/['params']/['decoder']/['dense_layers']/['mlp']/['wi_1']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -144,6 +151,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -183,6 +191,7 @@ ".params/['params']/['decoder']/['dense_layers']/['mlp']/['wo']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -197,6 +206,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -236,6 +246,7 @@ ".params/['params']/['decoder']/['dense_layers']/['post_self_attention_layer_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -250,6 +261,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -279,6 +291,7 @@ ".params/['params']/['decoder']/['dense_layers']/['pre_self_attention_layer_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -293,6 +306,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -322,6 +336,7 @@ ".params/['params']/['decoder']/['dense_layers']/['self_attention']/['kv_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -336,6 +351,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -365,6 +381,7 @@ ".params/['params']/['decoder']/['dense_layers']/['self_attention']/['out']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -379,6 +396,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -420,6 +438,7 @@ ".params/['params']/['decoder']/['dense_layers']/['self_attention']/['query']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -434,6 +453,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -475,6 +495,7 @@ ".params/['params']/['decoder']/['dense_layers']/['self_attention']/['wkv_a']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -489,6 +510,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -524,6 +546,7 @@ ".params/['params']/['decoder']/['dense_layers']/['self_attention']/['wkv_b']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -538,6 +561,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -579,6 +603,7 @@ ".params/['params']/['decoder']/['logits_dense']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -593,6 +618,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -630,6 +656,7 @@ ".params/['params']/['decoder']/['moe_layers']/['DeepSeekMoeBlock_0']/['MoeBlock_0']/['gate']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -644,6 +671,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -679,6 +707,7 @@ ".params/['params']/['decoder']/['moe_layers']/['DeepSeekMoeBlock_0']/['MoeBlock_0']/['wi_0']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -693,6 +722,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -733,6 +763,7 @@ ".params/['params']/['decoder']/['moe_layers']/['DeepSeekMoeBlock_0']/['MoeBlock_0']/['wi_1']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -747,6 +778,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -787,6 +819,7 @@ ".params/['params']/['decoder']/['moe_layers']/['DeepSeekMoeBlock_0']/['MoeBlock_0']/['wo']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -801,6 +834,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -841,6 +875,7 @@ ".params/['params']/['decoder']/['moe_layers']/['DeepSeekMoeBlock_0']/['shared_experts']/['wi_0']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -855,6 +890,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -894,6 +930,7 @@ ".params/['params']/['decoder']/['moe_layers']/['DeepSeekMoeBlock_0']/['shared_experts']/['wi_1']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -908,6 +945,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -947,6 +985,7 @@ ".params/['params']/['decoder']/['moe_layers']/['DeepSeekMoeBlock_0']/['shared_experts']/['wo']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -961,6 +1000,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -1000,6 +1040,7 @@ ".params/['params']/['decoder']/['moe_layers']/['post_self_attention_layer_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1014,6 +1055,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -1043,6 +1085,7 @@ ".params/['params']/['decoder']/['moe_layers']/['pre_self_attention_layer_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1057,6 +1100,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -1086,6 +1130,7 @@ ".params/['params']/['decoder']/['moe_layers']/['self_attention']/['kv_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1100,6 +1145,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -1129,6 +1175,7 @@ ".params/['params']/['decoder']/['moe_layers']/['self_attention']/['out']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1143,6 +1190,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -1184,6 +1232,7 @@ ".params/['params']/['decoder']/['moe_layers']/['self_attention']/['query']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1198,6 +1247,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -1239,6 +1289,7 @@ ".params/['params']/['decoder']/['moe_layers']/['self_attention']/['wkv_a']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1253,6 +1304,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -1288,6 +1340,7 @@ ".params/['params']/['decoder']/['moe_layers']/['self_attention']/['wkv_b']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1302,6 +1355,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -1343,6 +1397,7 @@ ".params/['params']/['token_embedder']/['embedding']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1357,6 +1412,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -1394,6 +1450,7 @@ ".opt_state/[0]/.count": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1408,6 +1465,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -1428,6 +1486,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['decoder_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1442,6 +1501,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -1469,6 +1529,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['dense_layers']/['mlp']/['wi_0']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1483,6 +1544,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -1522,6 +1584,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['dense_layers']/['mlp']/['wi_1']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1536,6 +1599,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -1575,6 +1639,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['dense_layers']/['mlp']/['wo']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1589,6 +1654,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -1628,6 +1694,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['dense_layers']/['post_self_attention_layer_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1642,6 +1709,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -1671,6 +1739,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['dense_layers']/['pre_self_attention_layer_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1685,6 +1754,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -1714,6 +1784,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['dense_layers']/['self_attention']/['kv_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1728,6 +1799,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -1757,6 +1829,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['dense_layers']/['self_attention']/['out']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1771,6 +1844,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -1812,6 +1886,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['dense_layers']/['self_attention']/['query']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1826,6 +1901,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -1867,6 +1943,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['dense_layers']/['self_attention']/['wkv_a']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1881,6 +1958,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -1916,6 +1994,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['dense_layers']/['self_attention']/['wkv_b']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1930,6 +2009,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -1971,6 +2051,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['logits_dense']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1985,6 +2066,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -2022,6 +2104,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['moe_layers']/['DeepSeekMoeBlock_0']/['MoeBlock_0']/['gate']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2036,6 +2119,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -2071,6 +2155,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['moe_layers']/['DeepSeekMoeBlock_0']/['MoeBlock_0']/['wi_0']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2085,6 +2170,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -2125,6 +2211,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['moe_layers']/['DeepSeekMoeBlock_0']/['MoeBlock_0']/['wi_1']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2139,6 +2226,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -2179,6 +2267,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['moe_layers']/['DeepSeekMoeBlock_0']/['MoeBlock_0']/['wo']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2193,6 +2282,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -2233,6 +2323,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['moe_layers']/['DeepSeekMoeBlock_0']/['shared_experts']/['wi_0']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2247,6 +2338,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -2286,6 +2378,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['moe_layers']/['DeepSeekMoeBlock_0']/['shared_experts']/['wi_1']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2300,6 +2393,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -2339,6 +2433,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['moe_layers']/['DeepSeekMoeBlock_0']/['shared_experts']/['wo']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2353,6 +2448,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -2392,6 +2488,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['moe_layers']/['post_self_attention_layer_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2406,6 +2503,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -2435,6 +2533,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['moe_layers']/['pre_self_attention_layer_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2449,6 +2548,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -2478,6 +2578,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['moe_layers']/['self_attention']/['kv_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2492,6 +2593,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -2521,6 +2623,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['moe_layers']/['self_attention']/['out']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2535,6 +2638,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -2576,6 +2680,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['moe_layers']/['self_attention']/['query']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2590,6 +2695,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -2631,6 +2737,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['moe_layers']/['self_attention']/['wkv_a']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2645,6 +2752,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -2680,6 +2788,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['moe_layers']/['self_attention']/['wkv_b']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2694,6 +2803,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -2735,6 +2845,7 @@ ".opt_state/[0]/.mu/['params']/['token_embedder']/['embedding']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2749,6 +2860,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -2786,6 +2898,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['decoder_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2800,6 +2913,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -2827,6 +2941,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['dense_layers']/['mlp']/['wi_0']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2841,6 +2956,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -2880,6 +2996,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['dense_layers']/['mlp']/['wi_1']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2894,6 +3011,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -2933,6 +3051,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['dense_layers']/['mlp']/['wo']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2947,6 +3066,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -2986,6 +3106,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['dense_layers']/['post_self_attention_layer_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3000,6 +3121,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -3029,6 +3151,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['dense_layers']/['pre_self_attention_layer_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3043,6 +3166,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -3072,6 +3196,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['dense_layers']/['self_attention']/['kv_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3086,6 +3211,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -3115,6 +3241,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['dense_layers']/['self_attention']/['out']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3129,6 +3256,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -3170,6 +3298,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['dense_layers']/['self_attention']/['query']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3184,6 +3313,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -3225,6 +3355,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['dense_layers']/['self_attention']/['wkv_a']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3239,6 +3370,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -3274,6 +3406,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['dense_layers']/['self_attention']/['wkv_b']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3288,6 +3421,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -3329,6 +3463,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['logits_dense']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3343,6 +3478,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -3380,6 +3516,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['moe_layers']/['DeepSeekMoeBlock_0']/['MoeBlock_0']/['gate']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3394,6 +3531,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -3429,6 +3567,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['moe_layers']/['DeepSeekMoeBlock_0']/['MoeBlock_0']/['wi_0']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3443,6 +3582,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -3483,6 +3623,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['moe_layers']/['DeepSeekMoeBlock_0']/['MoeBlock_0']/['wi_1']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3497,6 +3638,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -3537,6 +3679,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['moe_layers']/['DeepSeekMoeBlock_0']/['MoeBlock_0']/['wo']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3551,6 +3694,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -3591,6 +3735,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['moe_layers']/['DeepSeekMoeBlock_0']/['shared_experts']/['wi_0']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3605,6 +3750,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -3644,6 +3790,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['moe_layers']/['DeepSeekMoeBlock_0']/['shared_experts']/['wi_1']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3658,6 +3805,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -3697,6 +3845,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['moe_layers']/['DeepSeekMoeBlock_0']/['shared_experts']/['wo']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3711,6 +3860,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -3750,6 +3900,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['moe_layers']/['post_self_attention_layer_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3764,6 +3915,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -3793,6 +3945,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['moe_layers']/['pre_self_attention_layer_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3807,6 +3960,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -3836,6 +3990,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['moe_layers']/['self_attention']/['kv_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3850,6 +4005,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -3879,6 +4035,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['moe_layers']/['self_attention']/['out']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3893,6 +4050,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -3934,6 +4092,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['moe_layers']/['self_attention']/['query']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3948,6 +4107,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -3989,6 +4149,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['moe_layers']/['self_attention']/['wkv_a']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -4003,6 +4164,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -4038,6 +4200,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['moe_layers']/['self_attention']/['wkv_b']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -4052,6 +4215,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -4093,6 +4257,7 @@ ".opt_state/[0]/.nu/['params']/['token_embedder']/['embedding']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -4107,6 +4272,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -4144,6 +4310,7 @@ ".opt_state/[2]/.count": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -4158,6 +4325,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, diff --git a/tests/utils/sharding_info/deepseek2-16b/v6e-16/slice_4/named_shardings.json b/tests/utils/sharding_info/deepseek2-16b/v6e-16/slice_4/named_shardings.json index a7fa362422..5b2ab94daf 100644 --- a/tests/utils/sharding_info/deepseek2-16b/v6e-16/slice_4/named_shardings.json +++ b/tests/utils/sharding_info/deepseek2-16b/v6e-16/slice_4/named_shardings.json @@ -2,6 +2,7 @@ ".step": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -16,6 +17,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -36,6 +38,7 @@ ".params/['params']/['decoder']/['decoder_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -50,6 +53,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -77,6 +81,7 @@ ".params/['params']/['decoder']/['dense_layers']/['mlp']/['wi_0']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -91,6 +96,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -130,6 +136,7 @@ ".params/['params']/['decoder']/['dense_layers']/['mlp']/['wi_1']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -144,6 +151,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -183,6 +191,7 @@ ".params/['params']/['decoder']/['dense_layers']/['mlp']/['wo']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -197,6 +206,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -236,6 +246,7 @@ ".params/['params']/['decoder']/['dense_layers']/['post_self_attention_layer_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -250,6 +261,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -279,6 +291,7 @@ ".params/['params']/['decoder']/['dense_layers']/['pre_self_attention_layer_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -293,6 +306,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -322,6 +336,7 @@ ".params/['params']/['decoder']/['dense_layers']/['self_attention']/['kv_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -336,6 +351,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -365,6 +381,7 @@ ".params/['params']/['decoder']/['dense_layers']/['self_attention']/['out']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -379,6 +396,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -420,6 +438,7 @@ ".params/['params']/['decoder']/['dense_layers']/['self_attention']/['query']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -434,6 +453,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -475,6 +495,7 @@ ".params/['params']/['decoder']/['dense_layers']/['self_attention']/['wkv_a']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -489,6 +510,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -524,6 +546,7 @@ ".params/['params']/['decoder']/['dense_layers']/['self_attention']/['wkv_b']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -538,6 +561,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -579,6 +603,7 @@ ".params/['params']/['decoder']/['logits_dense']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -593,6 +618,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -630,6 +656,7 @@ ".params/['params']/['decoder']/['moe_layers']/['DeepSeekMoeBlock_0']/['MoeBlock_0']/['gate']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -644,6 +671,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -679,6 +707,7 @@ ".params/['params']/['decoder']/['moe_layers']/['DeepSeekMoeBlock_0']/['MoeBlock_0']/['wi_0']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -693,6 +722,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -733,6 +763,7 @@ ".params/['params']/['decoder']/['moe_layers']/['DeepSeekMoeBlock_0']/['MoeBlock_0']/['wi_1']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -747,6 +778,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -787,6 +819,7 @@ ".params/['params']/['decoder']/['moe_layers']/['DeepSeekMoeBlock_0']/['MoeBlock_0']/['wo']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -801,6 +834,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -841,6 +875,7 @@ ".params/['params']/['decoder']/['moe_layers']/['DeepSeekMoeBlock_0']/['shared_experts']/['wi_0']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -855,6 +890,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -894,6 +930,7 @@ ".params/['params']/['decoder']/['moe_layers']/['DeepSeekMoeBlock_0']/['shared_experts']/['wi_1']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -908,6 +945,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -947,6 +985,7 @@ ".params/['params']/['decoder']/['moe_layers']/['DeepSeekMoeBlock_0']/['shared_experts']/['wo']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -961,6 +1000,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -1000,6 +1040,7 @@ ".params/['params']/['decoder']/['moe_layers']/['post_self_attention_layer_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1014,6 +1055,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -1043,6 +1085,7 @@ ".params/['params']/['decoder']/['moe_layers']/['pre_self_attention_layer_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1057,6 +1100,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -1086,6 +1130,7 @@ ".params/['params']/['decoder']/['moe_layers']/['self_attention']/['kv_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1100,6 +1145,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -1129,6 +1175,7 @@ ".params/['params']/['decoder']/['moe_layers']/['self_attention']/['out']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1143,6 +1190,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -1184,6 +1232,7 @@ ".params/['params']/['decoder']/['moe_layers']/['self_attention']/['query']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1198,6 +1247,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -1239,6 +1289,7 @@ ".params/['params']/['decoder']/['moe_layers']/['self_attention']/['wkv_a']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1253,6 +1304,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -1288,6 +1340,7 @@ ".params/['params']/['decoder']/['moe_layers']/['self_attention']/['wkv_b']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1302,6 +1355,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -1343,6 +1397,7 @@ ".params/['params']/['token_embedder']/['embedding']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1357,6 +1412,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -1394,6 +1450,7 @@ ".opt_state/[0]/.count": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1408,6 +1465,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -1428,6 +1486,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['decoder_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1442,6 +1501,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -1469,6 +1529,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['dense_layers']/['mlp']/['wi_0']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1483,6 +1544,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -1522,6 +1584,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['dense_layers']/['mlp']/['wi_1']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1536,6 +1599,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -1575,6 +1639,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['dense_layers']/['mlp']/['wo']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1589,6 +1654,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -1628,6 +1694,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['dense_layers']/['post_self_attention_layer_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1642,6 +1709,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -1671,6 +1739,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['dense_layers']/['pre_self_attention_layer_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1685,6 +1754,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -1714,6 +1784,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['dense_layers']/['self_attention']/['kv_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1728,6 +1799,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -1757,6 +1829,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['dense_layers']/['self_attention']/['out']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1771,6 +1844,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -1812,6 +1886,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['dense_layers']/['self_attention']/['query']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1826,6 +1901,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -1867,6 +1943,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['dense_layers']/['self_attention']/['wkv_a']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1881,6 +1958,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -1916,6 +1994,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['dense_layers']/['self_attention']/['wkv_b']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1930,6 +2009,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -1971,6 +2051,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['logits_dense']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1985,6 +2066,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -2022,6 +2104,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['moe_layers']/['DeepSeekMoeBlock_0']/['MoeBlock_0']/['gate']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2036,6 +2119,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -2071,6 +2155,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['moe_layers']/['DeepSeekMoeBlock_0']/['MoeBlock_0']/['wi_0']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2085,6 +2170,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -2125,6 +2211,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['moe_layers']/['DeepSeekMoeBlock_0']/['MoeBlock_0']/['wi_1']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2139,6 +2226,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -2179,6 +2267,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['moe_layers']/['DeepSeekMoeBlock_0']/['MoeBlock_0']/['wo']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2193,6 +2282,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -2233,6 +2323,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['moe_layers']/['DeepSeekMoeBlock_0']/['shared_experts']/['wi_0']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2247,6 +2338,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -2286,6 +2378,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['moe_layers']/['DeepSeekMoeBlock_0']/['shared_experts']/['wi_1']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2300,6 +2393,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -2339,6 +2433,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['moe_layers']/['DeepSeekMoeBlock_0']/['shared_experts']/['wo']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2353,6 +2448,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -2392,6 +2488,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['moe_layers']/['post_self_attention_layer_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2406,6 +2503,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -2435,6 +2533,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['moe_layers']/['pre_self_attention_layer_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2449,6 +2548,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -2478,6 +2578,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['moe_layers']/['self_attention']/['kv_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2492,6 +2593,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -2521,6 +2623,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['moe_layers']/['self_attention']/['out']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2535,6 +2638,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -2576,6 +2680,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['moe_layers']/['self_attention']/['query']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2590,6 +2695,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -2631,6 +2737,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['moe_layers']/['self_attention']/['wkv_a']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2645,6 +2752,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -2680,6 +2788,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['moe_layers']/['self_attention']/['wkv_b']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2694,6 +2803,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -2735,6 +2845,7 @@ ".opt_state/[0]/.mu/['params']/['token_embedder']/['embedding']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2749,6 +2860,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -2786,6 +2898,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['decoder_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2800,6 +2913,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -2827,6 +2941,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['dense_layers']/['mlp']/['wi_0']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2841,6 +2956,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -2880,6 +2996,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['dense_layers']/['mlp']/['wi_1']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2894,6 +3011,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -2933,6 +3051,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['dense_layers']/['mlp']/['wo']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2947,6 +3066,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -2986,6 +3106,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['dense_layers']/['post_self_attention_layer_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3000,6 +3121,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -3029,6 +3151,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['dense_layers']/['pre_self_attention_layer_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3043,6 +3166,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -3072,6 +3196,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['dense_layers']/['self_attention']/['kv_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3086,6 +3211,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -3115,6 +3241,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['dense_layers']/['self_attention']/['out']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3129,6 +3256,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -3170,6 +3298,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['dense_layers']/['self_attention']/['query']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3184,6 +3313,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -3225,6 +3355,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['dense_layers']/['self_attention']/['wkv_a']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3239,6 +3370,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -3274,6 +3406,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['dense_layers']/['self_attention']/['wkv_b']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3288,6 +3421,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -3329,6 +3463,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['logits_dense']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3343,6 +3478,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -3380,6 +3516,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['moe_layers']/['DeepSeekMoeBlock_0']/['MoeBlock_0']/['gate']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3394,6 +3531,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -3429,6 +3567,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['moe_layers']/['DeepSeekMoeBlock_0']/['MoeBlock_0']/['wi_0']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3443,6 +3582,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -3483,6 +3623,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['moe_layers']/['DeepSeekMoeBlock_0']/['MoeBlock_0']/['wi_1']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3497,6 +3638,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -3537,6 +3679,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['moe_layers']/['DeepSeekMoeBlock_0']/['MoeBlock_0']/['wo']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3551,6 +3694,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -3591,6 +3735,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['moe_layers']/['DeepSeekMoeBlock_0']/['shared_experts']/['wi_0']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3605,6 +3750,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -3644,6 +3790,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['moe_layers']/['DeepSeekMoeBlock_0']/['shared_experts']/['wi_1']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3658,6 +3805,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -3697,6 +3845,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['moe_layers']/['DeepSeekMoeBlock_0']/['shared_experts']/['wo']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3711,6 +3860,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -3750,6 +3900,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['moe_layers']/['post_self_attention_layer_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3764,6 +3915,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -3793,6 +3945,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['moe_layers']/['pre_self_attention_layer_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3807,6 +3960,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -3836,6 +3990,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['moe_layers']/['self_attention']/['kv_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3850,6 +4005,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -3879,6 +4035,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['moe_layers']/['self_attention']/['out']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3893,6 +4050,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -3934,6 +4092,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['moe_layers']/['self_attention']/['query']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3948,6 +4107,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -3989,6 +4149,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['moe_layers']/['self_attention']/['wkv_a']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -4003,6 +4164,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -4038,6 +4200,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['moe_layers']/['self_attention']/['wkv_b']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -4052,6 +4215,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -4093,6 +4257,7 @@ ".opt_state/[0]/.nu/['params']/['token_embedder']/['embedding']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -4107,6 +4272,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -4144,6 +4310,7 @@ ".opt_state/[2]/.count": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -4158,6 +4325,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, diff --git a/tests/utils/sharding_info/gpt-oss-20b/tpu7x-16/slice_1/named_shardings.json b/tests/utils/sharding_info/gpt-oss-20b/tpu7x-16/slice_1/named_shardings.json index 6a4eb12a10..78e42a8848 100644 --- a/tests/utils/sharding_info/gpt-oss-20b/tpu7x-16/slice_1/named_shardings.json +++ b/tests/utils/sharding_info/gpt-oss-20b/tpu7x-16/slice_1/named_shardings.json @@ -2,6 +2,7 @@ ".step": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -16,6 +17,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -36,6 +38,7 @@ ".params/['params']/['decoder']/['decoder_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -50,6 +53,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -77,6 +81,7 @@ ".params/['params']/['decoder']/['layers']/['layers_0']/['GptOssAttention']/['key']/['bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -91,6 +96,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -124,6 +130,7 @@ ".params/['params']/['decoder']/['layers']/['layers_0']/['GptOssAttention']/['key']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -138,6 +145,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -179,6 +187,7 @@ ".params/['params']/['decoder']/['layers']/['layers_0']/['GptOssAttention']/['out']/['bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -193,6 +202,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -226,6 +236,7 @@ ".params/['params']/['decoder']/['layers']/['layers_0']/['GptOssAttention']/['out']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -240,6 +251,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -281,6 +293,7 @@ ".params/['params']/['decoder']/['layers']/['layers_0']/['GptOssAttention']/['query']/['bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -295,6 +308,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -328,6 +342,7 @@ ".params/['params']/['decoder']/['layers']/['layers_0']/['GptOssAttention']/['query']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -342,6 +357,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -383,6 +399,7 @@ ".params/['params']/['decoder']/['layers']/['layers_0']/['GptOssAttention']/['sinks']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -397,6 +414,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -423,6 +441,7 @@ ".params/['params']/['decoder']/['layers']/['layers_0']/['GptOssAttention']/['value']/['bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -437,6 +456,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -470,6 +490,7 @@ ".params/['params']/['decoder']/['layers']/['layers_0']/['GptOssAttention']/['value']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -484,6 +505,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -525,6 +547,7 @@ ".params/['params']/['decoder']/['layers']/['layers_0']/['GptOssMlp']/['gate']/['bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -539,6 +562,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -565,6 +589,7 @@ ".params/['params']/['decoder']/['layers']/['layers_0']/['GptOssMlp']/['gate']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -579,6 +604,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -614,6 +640,7 @@ ".params/['params']/['decoder']/['layers']/['layers_0']/['GptOssMlp']/['wi_0']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -628,6 +655,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -668,6 +696,7 @@ ".params/['params']/['decoder']/['layers']/['layers_0']/['GptOssMlp']/['wi_0_bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -682,6 +711,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -714,6 +744,7 @@ ".params/['params']/['decoder']/['layers']/['layers_0']/['GptOssMlp']/['wi_1']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -728,6 +759,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -768,6 +800,7 @@ ".params/['params']/['decoder']/['layers']/['layers_0']/['GptOssMlp']/['wi_1_bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -782,6 +815,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -814,6 +848,7 @@ ".params/['params']/['decoder']/['layers']/['layers_0']/['GptOssMlp']/['wo']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -828,6 +863,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -868,6 +904,7 @@ ".params/['params']/['decoder']/['layers']/['layers_0']/['GptOssMlp']/['wo_bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -882,6 +919,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -913,6 +951,7 @@ ".params/['params']/['decoder']/['layers']/['layers_0']/['post_self_attention_layer_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -927,6 +966,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -956,6 +996,7 @@ ".params/['params']/['decoder']/['layers']/['layers_0']/['pre_self_attention_layer_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -970,6 +1011,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -999,6 +1041,7 @@ ".params/['params']/['decoder']/['layers']/['layers_1']/['GptOssAttention']/['key']/['bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1013,6 +1056,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -1046,6 +1090,7 @@ ".params/['params']/['decoder']/['layers']/['layers_1']/['GptOssAttention']/['key']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1060,6 +1105,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -1101,6 +1147,7 @@ ".params/['params']/['decoder']/['layers']/['layers_1']/['GptOssAttention']/['out']/['bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1115,6 +1162,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -1148,6 +1196,7 @@ ".params/['params']/['decoder']/['layers']/['layers_1']/['GptOssAttention']/['out']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1162,6 +1211,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -1203,6 +1253,7 @@ ".params/['params']/['decoder']/['layers']/['layers_1']/['GptOssAttention']/['query']/['bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1217,6 +1268,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -1250,6 +1302,7 @@ ".params/['params']/['decoder']/['layers']/['layers_1']/['GptOssAttention']/['query']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1264,6 +1317,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -1305,6 +1359,7 @@ ".params/['params']/['decoder']/['layers']/['layers_1']/['GptOssAttention']/['sinks']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1319,6 +1374,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -1345,6 +1401,7 @@ ".params/['params']/['decoder']/['layers']/['layers_1']/['GptOssAttention']/['value']/['bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1359,6 +1416,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -1392,6 +1450,7 @@ ".params/['params']/['decoder']/['layers']/['layers_1']/['GptOssAttention']/['value']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1406,6 +1465,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -1447,6 +1507,7 @@ ".params/['params']/['decoder']/['layers']/['layers_1']/['GptOssMlp']/['gate']/['bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1461,6 +1522,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -1487,6 +1549,7 @@ ".params/['params']/['decoder']/['layers']/['layers_1']/['GptOssMlp']/['gate']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1501,6 +1564,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -1536,6 +1600,7 @@ ".params/['params']/['decoder']/['layers']/['layers_1']/['GptOssMlp']/['wi_0']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1550,6 +1615,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -1590,6 +1656,7 @@ ".params/['params']/['decoder']/['layers']/['layers_1']/['GptOssMlp']/['wi_0_bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1604,6 +1671,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -1636,6 +1704,7 @@ ".params/['params']/['decoder']/['layers']/['layers_1']/['GptOssMlp']/['wi_1']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1650,6 +1719,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -1690,6 +1760,7 @@ ".params/['params']/['decoder']/['layers']/['layers_1']/['GptOssMlp']/['wi_1_bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1704,6 +1775,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -1736,6 +1808,7 @@ ".params/['params']/['decoder']/['layers']/['layers_1']/['GptOssMlp']/['wo']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1750,6 +1823,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -1790,6 +1864,7 @@ ".params/['params']/['decoder']/['layers']/['layers_1']/['GptOssMlp']/['wo_bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1804,6 +1879,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -1835,6 +1911,7 @@ ".params/['params']/['decoder']/['layers']/['layers_1']/['post_self_attention_layer_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1849,6 +1926,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -1878,6 +1956,7 @@ ".params/['params']/['decoder']/['layers']/['layers_1']/['pre_self_attention_layer_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1892,6 +1971,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -1921,6 +2001,7 @@ ".params/['params']/['decoder']/['logits_dense']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1935,6 +2016,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -1972,6 +2054,7 @@ ".params/['params']/['token_embedder']/['embedding']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1986,6 +2069,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -2023,6 +2107,7 @@ ".opt_state/[0]/.count": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2037,6 +2122,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -2057,6 +2143,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['decoder_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2071,6 +2158,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -2098,6 +2186,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['layers_0']/['GptOssAttention']/['key']/['bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2112,6 +2201,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -2145,6 +2235,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['layers_0']/['GptOssAttention']/['key']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2159,6 +2250,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -2200,6 +2292,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['layers_0']/['GptOssAttention']/['out']/['bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2214,6 +2307,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -2247,6 +2341,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['layers_0']/['GptOssAttention']/['out']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2261,6 +2356,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -2302,6 +2398,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['layers_0']/['GptOssAttention']/['query']/['bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2316,6 +2413,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -2349,6 +2447,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['layers_0']/['GptOssAttention']/['query']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2363,6 +2462,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -2404,6 +2504,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['layers_0']/['GptOssAttention']/['sinks']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2418,6 +2519,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -2444,6 +2546,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['layers_0']/['GptOssAttention']/['value']/['bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2458,6 +2561,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -2491,6 +2595,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['layers_0']/['GptOssAttention']/['value']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2505,6 +2610,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -2546,6 +2652,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['layers_0']/['GptOssMlp']/['gate']/['bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2560,6 +2667,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -2586,6 +2694,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['layers_0']/['GptOssMlp']/['gate']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2600,6 +2709,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -2635,6 +2745,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['layers_0']/['GptOssMlp']/['wi_0']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2649,6 +2760,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -2689,6 +2801,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['layers_0']/['GptOssMlp']/['wi_0_bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2703,6 +2816,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -2735,6 +2849,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['layers_0']/['GptOssMlp']/['wi_1']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2749,6 +2864,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -2789,6 +2905,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['layers_0']/['GptOssMlp']/['wi_1_bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2803,6 +2920,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -2835,6 +2953,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['layers_0']/['GptOssMlp']/['wo']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2849,6 +2968,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -2889,6 +3009,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['layers_0']/['GptOssMlp']/['wo_bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2903,6 +3024,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -2934,6 +3056,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['layers_0']/['post_self_attention_layer_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2948,6 +3071,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -2977,6 +3101,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['layers_0']/['pre_self_attention_layer_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2991,6 +3116,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -3020,6 +3146,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['layers_1']/['GptOssAttention']/['key']/['bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3034,6 +3161,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -3067,6 +3195,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['layers_1']/['GptOssAttention']/['key']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3081,6 +3210,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -3122,6 +3252,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['layers_1']/['GptOssAttention']/['out']/['bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3136,6 +3267,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -3169,6 +3301,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['layers_1']/['GptOssAttention']/['out']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3183,6 +3316,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -3224,6 +3358,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['layers_1']/['GptOssAttention']/['query']/['bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3238,6 +3373,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -3271,6 +3407,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['layers_1']/['GptOssAttention']/['query']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3285,6 +3422,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -3326,6 +3464,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['layers_1']/['GptOssAttention']/['sinks']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3340,6 +3479,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -3366,6 +3506,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['layers_1']/['GptOssAttention']/['value']/['bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3380,6 +3521,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -3413,6 +3555,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['layers_1']/['GptOssAttention']/['value']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3427,6 +3570,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -3468,6 +3612,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['layers_1']/['GptOssMlp']/['gate']/['bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3482,6 +3627,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -3508,6 +3654,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['layers_1']/['GptOssMlp']/['gate']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3522,6 +3669,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -3557,6 +3705,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['layers_1']/['GptOssMlp']/['wi_0']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3571,6 +3720,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -3611,6 +3761,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['layers_1']/['GptOssMlp']/['wi_0_bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3625,6 +3776,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -3657,6 +3809,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['layers_1']/['GptOssMlp']/['wi_1']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3671,6 +3824,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -3711,6 +3865,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['layers_1']/['GptOssMlp']/['wi_1_bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3725,6 +3880,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -3757,6 +3913,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['layers_1']/['GptOssMlp']/['wo']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3771,6 +3928,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -3811,6 +3969,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['layers_1']/['GptOssMlp']/['wo_bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3825,6 +3984,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -3856,6 +4016,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['layers_1']/['post_self_attention_layer_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3870,6 +4031,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -3899,6 +4061,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['layers_1']/['pre_self_attention_layer_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3913,6 +4076,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -3942,6 +4106,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['logits_dense']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3956,6 +4121,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -3993,6 +4159,7 @@ ".opt_state/[0]/.mu/['params']/['token_embedder']/['embedding']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -4007,6 +4174,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -4044,6 +4212,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['decoder_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -4058,6 +4227,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -4085,6 +4255,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['layers_0']/['GptOssAttention']/['key']/['bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -4099,6 +4270,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -4132,6 +4304,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['layers_0']/['GptOssAttention']/['key']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -4146,6 +4319,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -4187,6 +4361,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['layers_0']/['GptOssAttention']/['out']/['bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -4201,6 +4376,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -4234,6 +4410,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['layers_0']/['GptOssAttention']/['out']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -4248,6 +4425,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -4289,6 +4467,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['layers_0']/['GptOssAttention']/['query']/['bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -4303,6 +4482,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -4336,6 +4516,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['layers_0']/['GptOssAttention']/['query']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -4350,6 +4531,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -4391,6 +4573,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['layers_0']/['GptOssAttention']/['sinks']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -4405,6 +4588,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -4431,6 +4615,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['layers_0']/['GptOssAttention']/['value']/['bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -4445,6 +4630,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -4478,6 +4664,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['layers_0']/['GptOssAttention']/['value']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -4492,6 +4679,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -4533,6 +4721,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['layers_0']/['GptOssMlp']/['gate']/['bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -4547,6 +4736,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -4573,6 +4763,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['layers_0']/['GptOssMlp']/['gate']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -4587,6 +4778,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -4622,6 +4814,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['layers_0']/['GptOssMlp']/['wi_0']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -4636,6 +4829,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -4676,6 +4870,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['layers_0']/['GptOssMlp']/['wi_0_bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -4690,6 +4885,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -4722,6 +4918,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['layers_0']/['GptOssMlp']/['wi_1']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -4736,6 +4933,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -4776,6 +4974,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['layers_0']/['GptOssMlp']/['wi_1_bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -4790,6 +4989,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -4822,6 +5022,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['layers_0']/['GptOssMlp']/['wo']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -4836,6 +5037,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -4876,6 +5078,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['layers_0']/['GptOssMlp']/['wo_bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -4890,6 +5093,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -4921,6 +5125,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['layers_0']/['post_self_attention_layer_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -4935,6 +5140,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -4964,6 +5170,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['layers_0']/['pre_self_attention_layer_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -4978,6 +5185,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -5007,6 +5215,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['layers_1']/['GptOssAttention']/['key']/['bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -5021,6 +5230,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -5054,6 +5264,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['layers_1']/['GptOssAttention']/['key']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -5068,6 +5279,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -5109,6 +5321,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['layers_1']/['GptOssAttention']/['out']/['bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -5123,6 +5336,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -5156,6 +5370,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['layers_1']/['GptOssAttention']/['out']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -5170,6 +5385,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -5211,6 +5427,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['layers_1']/['GptOssAttention']/['query']/['bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -5225,6 +5442,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -5258,6 +5476,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['layers_1']/['GptOssAttention']/['query']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -5272,6 +5491,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -5313,6 +5533,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['layers_1']/['GptOssAttention']/['sinks']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -5327,6 +5548,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -5353,6 +5575,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['layers_1']/['GptOssAttention']/['value']/['bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -5367,6 +5590,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -5400,6 +5624,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['layers_1']/['GptOssAttention']/['value']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -5414,6 +5639,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -5455,6 +5681,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['layers_1']/['GptOssMlp']/['gate']/['bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -5469,6 +5696,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -5495,6 +5723,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['layers_1']/['GptOssMlp']/['gate']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -5509,6 +5738,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -5544,6 +5774,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['layers_1']/['GptOssMlp']/['wi_0']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -5558,6 +5789,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -5598,6 +5830,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['layers_1']/['GptOssMlp']/['wi_0_bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -5612,6 +5845,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -5644,6 +5878,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['layers_1']/['GptOssMlp']/['wi_1']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -5658,6 +5893,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -5698,6 +5934,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['layers_1']/['GptOssMlp']/['wi_1_bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -5712,6 +5949,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -5744,6 +5982,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['layers_1']/['GptOssMlp']/['wo']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -5758,6 +5997,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -5798,6 +6038,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['layers_1']/['GptOssMlp']/['wo_bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -5812,6 +6053,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -5843,6 +6085,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['layers_1']/['post_self_attention_layer_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -5857,6 +6100,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -5886,6 +6130,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['layers_1']/['pre_self_attention_layer_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -5900,6 +6145,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -5929,6 +6175,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['logits_dense']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -5943,6 +6190,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -5980,6 +6228,7 @@ ".opt_state/[0]/.nu/['params']/['token_embedder']/['embedding']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -5994,6 +6243,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -6031,6 +6281,7 @@ ".opt_state/[2]/.count": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -6045,6 +6296,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, diff --git a/tests/utils/sharding_info/gpt-oss-20b/tpu7x-16/slice_4/named_shardings.json b/tests/utils/sharding_info/gpt-oss-20b/tpu7x-16/slice_4/named_shardings.json index fffa91ebe5..ed765f1d18 100644 --- a/tests/utils/sharding_info/gpt-oss-20b/tpu7x-16/slice_4/named_shardings.json +++ b/tests/utils/sharding_info/gpt-oss-20b/tpu7x-16/slice_4/named_shardings.json @@ -2,6 +2,7 @@ ".step": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -16,6 +17,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -36,6 +38,7 @@ ".params/['params']/['decoder']/['decoder_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -50,6 +53,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -77,6 +81,7 @@ ".params/['params']/['decoder']/['layers']/['layers_0']/['GptOssAttention']/['key']/['bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -91,6 +96,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -124,6 +130,7 @@ ".params/['params']/['decoder']/['layers']/['layers_0']/['GptOssAttention']/['key']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -138,6 +145,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -179,6 +187,7 @@ ".params/['params']/['decoder']/['layers']/['layers_0']/['GptOssAttention']/['out']/['bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -193,6 +202,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -226,6 +236,7 @@ ".params/['params']/['decoder']/['layers']/['layers_0']/['GptOssAttention']/['out']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -240,6 +251,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -281,6 +293,7 @@ ".params/['params']/['decoder']/['layers']/['layers_0']/['GptOssAttention']/['query']/['bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -295,6 +308,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -328,6 +342,7 @@ ".params/['params']/['decoder']/['layers']/['layers_0']/['GptOssAttention']/['query']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -342,6 +357,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -383,6 +399,7 @@ ".params/['params']/['decoder']/['layers']/['layers_0']/['GptOssAttention']/['sinks']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -397,6 +414,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -423,6 +441,7 @@ ".params/['params']/['decoder']/['layers']/['layers_0']/['GptOssAttention']/['value']/['bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -437,6 +456,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -470,6 +490,7 @@ ".params/['params']/['decoder']/['layers']/['layers_0']/['GptOssAttention']/['value']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -484,6 +505,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -525,6 +547,7 @@ ".params/['params']/['decoder']/['layers']/['layers_0']/['GptOssMlp']/['gate']/['bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -539,6 +562,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -565,6 +589,7 @@ ".params/['params']/['decoder']/['layers']/['layers_0']/['GptOssMlp']/['gate']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -579,6 +604,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -614,6 +640,7 @@ ".params/['params']/['decoder']/['layers']/['layers_0']/['GptOssMlp']/['wi_0']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -628,6 +655,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -668,6 +696,7 @@ ".params/['params']/['decoder']/['layers']/['layers_0']/['GptOssMlp']/['wi_0_bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -682,6 +711,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -714,6 +744,7 @@ ".params/['params']/['decoder']/['layers']/['layers_0']/['GptOssMlp']/['wi_1']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -728,6 +759,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -768,6 +800,7 @@ ".params/['params']/['decoder']/['layers']/['layers_0']/['GptOssMlp']/['wi_1_bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -782,6 +815,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -814,6 +848,7 @@ ".params/['params']/['decoder']/['layers']/['layers_0']/['GptOssMlp']/['wo']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -828,6 +863,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -868,6 +904,7 @@ ".params/['params']/['decoder']/['layers']/['layers_0']/['GptOssMlp']/['wo_bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -882,6 +919,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -913,6 +951,7 @@ ".params/['params']/['decoder']/['layers']/['layers_0']/['post_self_attention_layer_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -927,6 +966,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -956,6 +996,7 @@ ".params/['params']/['decoder']/['layers']/['layers_0']/['pre_self_attention_layer_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -970,6 +1011,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -999,6 +1041,7 @@ ".params/['params']/['decoder']/['layers']/['layers_1']/['GptOssAttention']/['key']/['bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1013,6 +1056,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -1046,6 +1090,7 @@ ".params/['params']/['decoder']/['layers']/['layers_1']/['GptOssAttention']/['key']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1060,6 +1105,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -1101,6 +1147,7 @@ ".params/['params']/['decoder']/['layers']/['layers_1']/['GptOssAttention']/['out']/['bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1115,6 +1162,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -1148,6 +1196,7 @@ ".params/['params']/['decoder']/['layers']/['layers_1']/['GptOssAttention']/['out']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1162,6 +1211,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -1203,6 +1253,7 @@ ".params/['params']/['decoder']/['layers']/['layers_1']/['GptOssAttention']/['query']/['bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1217,6 +1268,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -1250,6 +1302,7 @@ ".params/['params']/['decoder']/['layers']/['layers_1']/['GptOssAttention']/['query']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1264,6 +1317,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -1305,6 +1359,7 @@ ".params/['params']/['decoder']/['layers']/['layers_1']/['GptOssAttention']/['sinks']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1319,6 +1374,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -1345,6 +1401,7 @@ ".params/['params']/['decoder']/['layers']/['layers_1']/['GptOssAttention']/['value']/['bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1359,6 +1416,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -1392,6 +1450,7 @@ ".params/['params']/['decoder']/['layers']/['layers_1']/['GptOssAttention']/['value']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1406,6 +1465,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -1447,6 +1507,7 @@ ".params/['params']/['decoder']/['layers']/['layers_1']/['GptOssMlp']/['gate']/['bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1461,6 +1522,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -1487,6 +1549,7 @@ ".params/['params']/['decoder']/['layers']/['layers_1']/['GptOssMlp']/['gate']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1501,6 +1564,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -1536,6 +1600,7 @@ ".params/['params']/['decoder']/['layers']/['layers_1']/['GptOssMlp']/['wi_0']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1550,6 +1615,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -1590,6 +1656,7 @@ ".params/['params']/['decoder']/['layers']/['layers_1']/['GptOssMlp']/['wi_0_bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1604,6 +1671,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -1636,6 +1704,7 @@ ".params/['params']/['decoder']/['layers']/['layers_1']/['GptOssMlp']/['wi_1']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1650,6 +1719,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -1690,6 +1760,7 @@ ".params/['params']/['decoder']/['layers']/['layers_1']/['GptOssMlp']/['wi_1_bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1704,6 +1775,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -1736,6 +1808,7 @@ ".params/['params']/['decoder']/['layers']/['layers_1']/['GptOssMlp']/['wo']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1750,6 +1823,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -1790,6 +1864,7 @@ ".params/['params']/['decoder']/['layers']/['layers_1']/['GptOssMlp']/['wo_bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1804,6 +1879,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -1835,6 +1911,7 @@ ".params/['params']/['decoder']/['layers']/['layers_1']/['post_self_attention_layer_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1849,6 +1926,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -1878,6 +1956,7 @@ ".params/['params']/['decoder']/['layers']/['layers_1']/['pre_self_attention_layer_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1892,6 +1971,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -1921,6 +2001,7 @@ ".params/['params']/['decoder']/['logits_dense']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1935,6 +2016,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -1972,6 +2054,7 @@ ".params/['params']/['token_embedder']/['embedding']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1986,6 +2069,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -2023,6 +2107,7 @@ ".opt_state/[0]/.count": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2037,6 +2122,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -2057,6 +2143,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['decoder_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2071,6 +2158,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -2098,6 +2186,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['layers_0']/['GptOssAttention']/['key']/['bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2112,6 +2201,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -2145,6 +2235,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['layers_0']/['GptOssAttention']/['key']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2159,6 +2250,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -2200,6 +2292,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['layers_0']/['GptOssAttention']/['out']/['bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2214,6 +2307,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -2247,6 +2341,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['layers_0']/['GptOssAttention']/['out']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2261,6 +2356,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -2302,6 +2398,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['layers_0']/['GptOssAttention']/['query']/['bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2316,6 +2413,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -2349,6 +2447,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['layers_0']/['GptOssAttention']/['query']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2363,6 +2462,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -2404,6 +2504,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['layers_0']/['GptOssAttention']/['sinks']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2418,6 +2519,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -2444,6 +2546,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['layers_0']/['GptOssAttention']/['value']/['bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2458,6 +2561,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -2491,6 +2595,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['layers_0']/['GptOssAttention']/['value']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2505,6 +2610,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -2546,6 +2652,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['layers_0']/['GptOssMlp']/['gate']/['bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2560,6 +2667,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -2586,6 +2694,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['layers_0']/['GptOssMlp']/['gate']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2600,6 +2709,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -2635,6 +2745,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['layers_0']/['GptOssMlp']/['wi_0']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2649,6 +2760,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -2689,6 +2801,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['layers_0']/['GptOssMlp']/['wi_0_bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2703,6 +2816,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -2735,6 +2849,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['layers_0']/['GptOssMlp']/['wi_1']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2749,6 +2864,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -2789,6 +2905,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['layers_0']/['GptOssMlp']/['wi_1_bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2803,6 +2920,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -2835,6 +2953,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['layers_0']/['GptOssMlp']/['wo']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2849,6 +2968,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -2889,6 +3009,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['layers_0']/['GptOssMlp']/['wo_bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2903,6 +3024,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -2934,6 +3056,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['layers_0']/['post_self_attention_layer_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2948,6 +3071,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -2977,6 +3101,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['layers_0']/['pre_self_attention_layer_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2991,6 +3116,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -3020,6 +3146,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['layers_1']/['GptOssAttention']/['key']/['bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3034,6 +3161,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -3067,6 +3195,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['layers_1']/['GptOssAttention']/['key']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3081,6 +3210,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -3122,6 +3252,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['layers_1']/['GptOssAttention']/['out']/['bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3136,6 +3267,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -3169,6 +3301,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['layers_1']/['GptOssAttention']/['out']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3183,6 +3316,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -3224,6 +3358,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['layers_1']/['GptOssAttention']/['query']/['bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3238,6 +3373,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -3271,6 +3407,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['layers_1']/['GptOssAttention']/['query']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3285,6 +3422,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -3326,6 +3464,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['layers_1']/['GptOssAttention']/['sinks']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3340,6 +3479,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -3366,6 +3506,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['layers_1']/['GptOssAttention']/['value']/['bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3380,6 +3521,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -3413,6 +3555,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['layers_1']/['GptOssAttention']/['value']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3427,6 +3570,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -3468,6 +3612,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['layers_1']/['GptOssMlp']/['gate']/['bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3482,6 +3627,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -3508,6 +3654,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['layers_1']/['GptOssMlp']/['gate']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3522,6 +3669,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -3557,6 +3705,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['layers_1']/['GptOssMlp']/['wi_0']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3571,6 +3720,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -3611,6 +3761,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['layers_1']/['GptOssMlp']/['wi_0_bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3625,6 +3776,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -3657,6 +3809,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['layers_1']/['GptOssMlp']/['wi_1']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3671,6 +3824,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -3711,6 +3865,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['layers_1']/['GptOssMlp']/['wi_1_bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3725,6 +3880,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -3757,6 +3913,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['layers_1']/['GptOssMlp']/['wo']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3771,6 +3928,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -3811,6 +3969,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['layers_1']/['GptOssMlp']/['wo_bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3825,6 +3984,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -3856,6 +4016,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['layers_1']/['post_self_attention_layer_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3870,6 +4031,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -3899,6 +4061,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['layers_1']/['pre_self_attention_layer_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3913,6 +4076,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -3942,6 +4106,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['logits_dense']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3956,6 +4121,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -3993,6 +4159,7 @@ ".opt_state/[0]/.mu/['params']/['token_embedder']/['embedding']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -4007,6 +4174,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -4044,6 +4212,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['decoder_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -4058,6 +4227,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -4085,6 +4255,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['layers_0']/['GptOssAttention']/['key']/['bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -4099,6 +4270,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -4132,6 +4304,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['layers_0']/['GptOssAttention']/['key']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -4146,6 +4319,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -4187,6 +4361,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['layers_0']/['GptOssAttention']/['out']/['bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -4201,6 +4376,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -4234,6 +4410,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['layers_0']/['GptOssAttention']/['out']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -4248,6 +4425,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -4289,6 +4467,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['layers_0']/['GptOssAttention']/['query']/['bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -4303,6 +4482,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -4336,6 +4516,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['layers_0']/['GptOssAttention']/['query']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -4350,6 +4531,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -4391,6 +4573,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['layers_0']/['GptOssAttention']/['sinks']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -4405,6 +4588,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -4431,6 +4615,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['layers_0']/['GptOssAttention']/['value']/['bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -4445,6 +4630,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -4478,6 +4664,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['layers_0']/['GptOssAttention']/['value']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -4492,6 +4679,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -4533,6 +4721,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['layers_0']/['GptOssMlp']/['gate']/['bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -4547,6 +4736,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -4573,6 +4763,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['layers_0']/['GptOssMlp']/['gate']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -4587,6 +4778,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -4622,6 +4814,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['layers_0']/['GptOssMlp']/['wi_0']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -4636,6 +4829,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -4676,6 +4870,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['layers_0']/['GptOssMlp']/['wi_0_bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -4690,6 +4885,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -4722,6 +4918,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['layers_0']/['GptOssMlp']/['wi_1']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -4736,6 +4933,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -4776,6 +4974,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['layers_0']/['GptOssMlp']/['wi_1_bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -4790,6 +4989,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -4822,6 +5022,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['layers_0']/['GptOssMlp']/['wo']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -4836,6 +5037,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -4876,6 +5078,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['layers_0']/['GptOssMlp']/['wo_bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -4890,6 +5093,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -4921,6 +5125,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['layers_0']/['post_self_attention_layer_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -4935,6 +5140,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -4964,6 +5170,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['layers_0']/['pre_self_attention_layer_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -4978,6 +5185,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -5007,6 +5215,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['layers_1']/['GptOssAttention']/['key']/['bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -5021,6 +5230,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -5054,6 +5264,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['layers_1']/['GptOssAttention']/['key']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -5068,6 +5279,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -5109,6 +5321,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['layers_1']/['GptOssAttention']/['out']/['bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -5123,6 +5336,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -5156,6 +5370,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['layers_1']/['GptOssAttention']/['out']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -5170,6 +5385,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -5211,6 +5427,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['layers_1']/['GptOssAttention']/['query']/['bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -5225,6 +5442,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -5258,6 +5476,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['layers_1']/['GptOssAttention']/['query']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -5272,6 +5491,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -5313,6 +5533,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['layers_1']/['GptOssAttention']/['sinks']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -5327,6 +5548,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -5353,6 +5575,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['layers_1']/['GptOssAttention']/['value']/['bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -5367,6 +5590,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -5400,6 +5624,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['layers_1']/['GptOssAttention']/['value']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -5414,6 +5639,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -5455,6 +5681,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['layers_1']/['GptOssMlp']/['gate']/['bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -5469,6 +5696,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -5495,6 +5723,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['layers_1']/['GptOssMlp']/['gate']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -5509,6 +5738,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -5544,6 +5774,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['layers_1']/['GptOssMlp']/['wi_0']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -5558,6 +5789,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -5598,6 +5830,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['layers_1']/['GptOssMlp']/['wi_0_bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -5612,6 +5845,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -5644,6 +5878,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['layers_1']/['GptOssMlp']/['wi_1']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -5658,6 +5893,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -5698,6 +5934,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['layers_1']/['GptOssMlp']/['wi_1_bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -5712,6 +5949,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -5744,6 +5982,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['layers_1']/['GptOssMlp']/['wo']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -5758,6 +5997,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -5798,6 +6038,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['layers_1']/['GptOssMlp']/['wo_bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -5812,6 +6053,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -5843,6 +6085,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['layers_1']/['post_self_attention_layer_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -5857,6 +6100,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -5886,6 +6130,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['layers_1']/['pre_self_attention_layer_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -5900,6 +6145,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -5929,6 +6175,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['logits_dense']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -5943,6 +6190,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -5980,6 +6228,7 @@ ".opt_state/[0]/.nu/['params']/['token_embedder']/['embedding']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -5994,6 +6243,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -6031,6 +6281,7 @@ ".opt_state/[2]/.count": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -6045,6 +6296,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, diff --git a/tests/utils/sharding_info/gpt-oss-20b/v5p-16/slice_1/named_shardings.json b/tests/utils/sharding_info/gpt-oss-20b/v5p-16/slice_1/named_shardings.json index a291ec09db..8d8089aac3 100644 --- a/tests/utils/sharding_info/gpt-oss-20b/v5p-16/slice_1/named_shardings.json +++ b/tests/utils/sharding_info/gpt-oss-20b/v5p-16/slice_1/named_shardings.json @@ -2,6 +2,7 @@ ".step": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -16,6 +17,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -36,6 +38,7 @@ ".params/['params']/['decoder']/['decoder_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -50,6 +53,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -77,6 +81,7 @@ ".params/['params']/['decoder']/['layers']/['layers_0']/['GptOssAttention']/['key']/['bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -91,6 +96,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -124,6 +130,7 @@ ".params/['params']/['decoder']/['layers']/['layers_0']/['GptOssAttention']/['key']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -138,6 +145,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -179,6 +187,7 @@ ".params/['params']/['decoder']/['layers']/['layers_0']/['GptOssAttention']/['out']/['bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -193,6 +202,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -226,6 +236,7 @@ ".params/['params']/['decoder']/['layers']/['layers_0']/['GptOssAttention']/['out']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -240,6 +251,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -281,6 +293,7 @@ ".params/['params']/['decoder']/['layers']/['layers_0']/['GptOssAttention']/['query']/['bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -295,6 +308,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -328,6 +342,7 @@ ".params/['params']/['decoder']/['layers']/['layers_0']/['GptOssAttention']/['query']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -342,6 +357,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -383,6 +399,7 @@ ".params/['params']/['decoder']/['layers']/['layers_0']/['GptOssAttention']/['sinks']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -397,6 +414,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -423,6 +441,7 @@ ".params/['params']/['decoder']/['layers']/['layers_0']/['GptOssAttention']/['value']/['bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -437,6 +456,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -470,6 +490,7 @@ ".params/['params']/['decoder']/['layers']/['layers_0']/['GptOssAttention']/['value']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -484,6 +505,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -525,6 +547,7 @@ ".params/['params']/['decoder']/['layers']/['layers_0']/['GptOssMlp']/['gate']/['bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -539,6 +562,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -565,6 +589,7 @@ ".params/['params']/['decoder']/['layers']/['layers_0']/['GptOssMlp']/['gate']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -579,6 +604,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -614,6 +640,7 @@ ".params/['params']/['decoder']/['layers']/['layers_0']/['GptOssMlp']/['wi_0']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -628,6 +655,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -668,6 +696,7 @@ ".params/['params']/['decoder']/['layers']/['layers_0']/['GptOssMlp']/['wi_0_bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -682,6 +711,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -714,6 +744,7 @@ ".params/['params']/['decoder']/['layers']/['layers_0']/['GptOssMlp']/['wi_1']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -728,6 +759,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -768,6 +800,7 @@ ".params/['params']/['decoder']/['layers']/['layers_0']/['GptOssMlp']/['wi_1_bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -782,6 +815,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -814,6 +848,7 @@ ".params/['params']/['decoder']/['layers']/['layers_0']/['GptOssMlp']/['wo']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -828,6 +863,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -868,6 +904,7 @@ ".params/['params']/['decoder']/['layers']/['layers_0']/['GptOssMlp']/['wo_bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -882,6 +919,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -913,6 +951,7 @@ ".params/['params']/['decoder']/['layers']/['layers_0']/['post_self_attention_layer_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -927,6 +966,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -956,6 +996,7 @@ ".params/['params']/['decoder']/['layers']/['layers_0']/['pre_self_attention_layer_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -970,6 +1011,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -999,6 +1041,7 @@ ".params/['params']/['decoder']/['layers']/['layers_1']/['GptOssAttention']/['key']/['bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1013,6 +1056,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -1046,6 +1090,7 @@ ".params/['params']/['decoder']/['layers']/['layers_1']/['GptOssAttention']/['key']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1060,6 +1105,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -1101,6 +1147,7 @@ ".params/['params']/['decoder']/['layers']/['layers_1']/['GptOssAttention']/['out']/['bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1115,6 +1162,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -1148,6 +1196,7 @@ ".params/['params']/['decoder']/['layers']/['layers_1']/['GptOssAttention']/['out']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1162,6 +1211,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -1203,6 +1253,7 @@ ".params/['params']/['decoder']/['layers']/['layers_1']/['GptOssAttention']/['query']/['bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1217,6 +1268,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -1250,6 +1302,7 @@ ".params/['params']/['decoder']/['layers']/['layers_1']/['GptOssAttention']/['query']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1264,6 +1317,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -1305,6 +1359,7 @@ ".params/['params']/['decoder']/['layers']/['layers_1']/['GptOssAttention']/['sinks']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1319,6 +1374,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -1345,6 +1401,7 @@ ".params/['params']/['decoder']/['layers']/['layers_1']/['GptOssAttention']/['value']/['bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1359,6 +1416,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -1392,6 +1450,7 @@ ".params/['params']/['decoder']/['layers']/['layers_1']/['GptOssAttention']/['value']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1406,6 +1465,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -1447,6 +1507,7 @@ ".params/['params']/['decoder']/['layers']/['layers_1']/['GptOssMlp']/['gate']/['bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1461,6 +1522,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -1487,6 +1549,7 @@ ".params/['params']/['decoder']/['layers']/['layers_1']/['GptOssMlp']/['gate']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1501,6 +1564,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -1536,6 +1600,7 @@ ".params/['params']/['decoder']/['layers']/['layers_1']/['GptOssMlp']/['wi_0']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1550,6 +1615,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -1590,6 +1656,7 @@ ".params/['params']/['decoder']/['layers']/['layers_1']/['GptOssMlp']/['wi_0_bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1604,6 +1671,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -1636,6 +1704,7 @@ ".params/['params']/['decoder']/['layers']/['layers_1']/['GptOssMlp']/['wi_1']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1650,6 +1719,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -1690,6 +1760,7 @@ ".params/['params']/['decoder']/['layers']/['layers_1']/['GptOssMlp']/['wi_1_bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1704,6 +1775,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -1736,6 +1808,7 @@ ".params/['params']/['decoder']/['layers']/['layers_1']/['GptOssMlp']/['wo']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1750,6 +1823,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -1790,6 +1864,7 @@ ".params/['params']/['decoder']/['layers']/['layers_1']/['GptOssMlp']/['wo_bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1804,6 +1879,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -1835,6 +1911,7 @@ ".params/['params']/['decoder']/['layers']/['layers_1']/['post_self_attention_layer_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1849,6 +1926,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -1878,6 +1956,7 @@ ".params/['params']/['decoder']/['layers']/['layers_1']/['pre_self_attention_layer_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1892,6 +1971,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -1921,6 +2001,7 @@ ".params/['params']/['decoder']/['logits_dense']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1935,6 +2016,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -1972,6 +2054,7 @@ ".params/['params']/['token_embedder']/['embedding']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1986,6 +2069,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -2023,6 +2107,7 @@ ".opt_state/[0]/.count": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2037,6 +2122,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -2057,6 +2143,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['decoder_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2071,6 +2158,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -2098,6 +2186,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['layers_0']/['GptOssAttention']/['key']/['bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2112,6 +2201,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -2145,6 +2235,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['layers_0']/['GptOssAttention']/['key']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2159,6 +2250,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -2200,6 +2292,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['layers_0']/['GptOssAttention']/['out']/['bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2214,6 +2307,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -2247,6 +2341,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['layers_0']/['GptOssAttention']/['out']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2261,6 +2356,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -2302,6 +2398,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['layers_0']/['GptOssAttention']/['query']/['bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2316,6 +2413,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -2349,6 +2447,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['layers_0']/['GptOssAttention']/['query']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2363,6 +2462,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -2404,6 +2504,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['layers_0']/['GptOssAttention']/['sinks']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2418,6 +2519,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -2444,6 +2546,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['layers_0']/['GptOssAttention']/['value']/['bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2458,6 +2561,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -2491,6 +2595,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['layers_0']/['GptOssAttention']/['value']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2505,6 +2610,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -2546,6 +2652,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['layers_0']/['GptOssMlp']/['gate']/['bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2560,6 +2667,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -2586,6 +2694,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['layers_0']/['GptOssMlp']/['gate']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2600,6 +2709,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -2635,6 +2745,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['layers_0']/['GptOssMlp']/['wi_0']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2649,6 +2760,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -2689,6 +2801,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['layers_0']/['GptOssMlp']/['wi_0_bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2703,6 +2816,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -2735,6 +2849,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['layers_0']/['GptOssMlp']/['wi_1']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2749,6 +2864,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -2789,6 +2905,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['layers_0']/['GptOssMlp']/['wi_1_bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2803,6 +2920,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -2835,6 +2953,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['layers_0']/['GptOssMlp']/['wo']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2849,6 +2968,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -2889,6 +3009,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['layers_0']/['GptOssMlp']/['wo_bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2903,6 +3024,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -2934,6 +3056,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['layers_0']/['post_self_attention_layer_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2948,6 +3071,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -2977,6 +3101,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['layers_0']/['pre_self_attention_layer_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2991,6 +3116,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -3020,6 +3146,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['layers_1']/['GptOssAttention']/['key']/['bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3034,6 +3161,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -3067,6 +3195,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['layers_1']/['GptOssAttention']/['key']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3081,6 +3210,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -3122,6 +3252,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['layers_1']/['GptOssAttention']/['out']/['bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3136,6 +3267,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -3169,6 +3301,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['layers_1']/['GptOssAttention']/['out']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3183,6 +3316,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -3224,6 +3358,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['layers_1']/['GptOssAttention']/['query']/['bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3238,6 +3373,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -3271,6 +3407,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['layers_1']/['GptOssAttention']/['query']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3285,6 +3422,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -3326,6 +3464,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['layers_1']/['GptOssAttention']/['sinks']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3340,6 +3479,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -3366,6 +3506,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['layers_1']/['GptOssAttention']/['value']/['bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3380,6 +3521,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -3413,6 +3555,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['layers_1']/['GptOssAttention']/['value']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3427,6 +3570,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -3468,6 +3612,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['layers_1']/['GptOssMlp']/['gate']/['bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3482,6 +3627,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -3508,6 +3654,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['layers_1']/['GptOssMlp']/['gate']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3522,6 +3669,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -3557,6 +3705,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['layers_1']/['GptOssMlp']/['wi_0']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3571,6 +3720,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -3611,6 +3761,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['layers_1']/['GptOssMlp']/['wi_0_bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3625,6 +3776,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -3657,6 +3809,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['layers_1']/['GptOssMlp']/['wi_1']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3671,6 +3824,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -3711,6 +3865,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['layers_1']/['GptOssMlp']/['wi_1_bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3725,6 +3880,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -3757,6 +3913,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['layers_1']/['GptOssMlp']/['wo']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3771,6 +3928,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -3811,6 +3969,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['layers_1']/['GptOssMlp']/['wo_bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3825,6 +3984,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -3856,6 +4016,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['layers_1']/['post_self_attention_layer_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3870,6 +4031,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -3899,6 +4061,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['layers_1']/['pre_self_attention_layer_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3913,6 +4076,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -3942,6 +4106,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['logits_dense']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3956,6 +4121,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -3993,6 +4159,7 @@ ".opt_state/[0]/.mu/['params']/['token_embedder']/['embedding']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -4007,6 +4174,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -4044,6 +4212,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['decoder_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -4058,6 +4227,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -4085,6 +4255,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['layers_0']/['GptOssAttention']/['key']/['bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -4099,6 +4270,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -4132,6 +4304,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['layers_0']/['GptOssAttention']/['key']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -4146,6 +4319,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -4187,6 +4361,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['layers_0']/['GptOssAttention']/['out']/['bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -4201,6 +4376,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -4234,6 +4410,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['layers_0']/['GptOssAttention']/['out']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -4248,6 +4425,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -4289,6 +4467,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['layers_0']/['GptOssAttention']/['query']/['bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -4303,6 +4482,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -4336,6 +4516,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['layers_0']/['GptOssAttention']/['query']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -4350,6 +4531,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -4391,6 +4573,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['layers_0']/['GptOssAttention']/['sinks']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -4405,6 +4588,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -4431,6 +4615,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['layers_0']/['GptOssAttention']/['value']/['bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -4445,6 +4630,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -4478,6 +4664,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['layers_0']/['GptOssAttention']/['value']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -4492,6 +4679,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -4533,6 +4721,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['layers_0']/['GptOssMlp']/['gate']/['bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -4547,6 +4736,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -4573,6 +4763,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['layers_0']/['GptOssMlp']/['gate']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -4587,6 +4778,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -4622,6 +4814,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['layers_0']/['GptOssMlp']/['wi_0']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -4636,6 +4829,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -4676,6 +4870,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['layers_0']/['GptOssMlp']/['wi_0_bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -4690,6 +4885,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -4722,6 +4918,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['layers_0']/['GptOssMlp']/['wi_1']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -4736,6 +4933,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -4776,6 +4974,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['layers_0']/['GptOssMlp']/['wi_1_bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -4790,6 +4989,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -4822,6 +5022,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['layers_0']/['GptOssMlp']/['wo']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -4836,6 +5037,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -4876,6 +5078,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['layers_0']/['GptOssMlp']/['wo_bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -4890,6 +5093,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -4921,6 +5125,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['layers_0']/['post_self_attention_layer_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -4935,6 +5140,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -4964,6 +5170,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['layers_0']/['pre_self_attention_layer_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -4978,6 +5185,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -5007,6 +5215,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['layers_1']/['GptOssAttention']/['key']/['bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -5021,6 +5230,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -5054,6 +5264,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['layers_1']/['GptOssAttention']/['key']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -5068,6 +5279,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -5109,6 +5321,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['layers_1']/['GptOssAttention']/['out']/['bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -5123,6 +5336,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -5156,6 +5370,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['layers_1']/['GptOssAttention']/['out']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -5170,6 +5385,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -5211,6 +5427,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['layers_1']/['GptOssAttention']/['query']/['bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -5225,6 +5442,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -5258,6 +5476,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['layers_1']/['GptOssAttention']/['query']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -5272,6 +5491,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -5313,6 +5533,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['layers_1']/['GptOssAttention']/['sinks']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -5327,6 +5548,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -5353,6 +5575,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['layers_1']/['GptOssAttention']/['value']/['bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -5367,6 +5590,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -5400,6 +5624,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['layers_1']/['GptOssAttention']/['value']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -5414,6 +5639,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -5455,6 +5681,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['layers_1']/['GptOssMlp']/['gate']/['bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -5469,6 +5696,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -5495,6 +5723,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['layers_1']/['GptOssMlp']/['gate']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -5509,6 +5738,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -5544,6 +5774,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['layers_1']/['GptOssMlp']/['wi_0']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -5558,6 +5789,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -5598,6 +5830,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['layers_1']/['GptOssMlp']/['wi_0_bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -5612,6 +5845,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -5644,6 +5878,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['layers_1']/['GptOssMlp']/['wi_1']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -5658,6 +5893,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -5698,6 +5934,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['layers_1']/['GptOssMlp']/['wi_1_bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -5712,6 +5949,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -5744,6 +5982,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['layers_1']/['GptOssMlp']/['wo']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -5758,6 +5997,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -5798,6 +6038,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['layers_1']/['GptOssMlp']/['wo_bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -5812,6 +6053,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -5843,6 +6085,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['layers_1']/['post_self_attention_layer_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -5857,6 +6100,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -5886,6 +6130,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['layers_1']/['pre_self_attention_layer_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -5900,6 +6145,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -5929,6 +6175,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['logits_dense']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -5943,6 +6190,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -5980,6 +6228,7 @@ ".opt_state/[0]/.nu/['params']/['token_embedder']/['embedding']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -5994,6 +6243,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -6031,6 +6281,7 @@ ".opt_state/[2]/.count": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -6045,6 +6296,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, diff --git a/tests/utils/sharding_info/gpt-oss-20b/v5p-16/slice_4/named_shardings.json b/tests/utils/sharding_info/gpt-oss-20b/v5p-16/slice_4/named_shardings.json index 1e20b637fe..a395dba2ea 100644 --- a/tests/utils/sharding_info/gpt-oss-20b/v5p-16/slice_4/named_shardings.json +++ b/tests/utils/sharding_info/gpt-oss-20b/v5p-16/slice_4/named_shardings.json @@ -2,6 +2,7 @@ ".step": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -16,6 +17,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -36,6 +38,7 @@ ".params/['params']/['decoder']/['decoder_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -50,6 +53,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -77,6 +81,7 @@ ".params/['params']/['decoder']/['layers']/['layers_0']/['GptOssAttention']/['key']/['bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -91,6 +96,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -124,6 +130,7 @@ ".params/['params']/['decoder']/['layers']/['layers_0']/['GptOssAttention']/['key']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -138,6 +145,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -179,6 +187,7 @@ ".params/['params']/['decoder']/['layers']/['layers_0']/['GptOssAttention']/['out']/['bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -193,6 +202,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -226,6 +236,7 @@ ".params/['params']/['decoder']/['layers']/['layers_0']/['GptOssAttention']/['out']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -240,6 +251,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -281,6 +293,7 @@ ".params/['params']/['decoder']/['layers']/['layers_0']/['GptOssAttention']/['query']/['bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -295,6 +308,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -328,6 +342,7 @@ ".params/['params']/['decoder']/['layers']/['layers_0']/['GptOssAttention']/['query']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -342,6 +357,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -383,6 +399,7 @@ ".params/['params']/['decoder']/['layers']/['layers_0']/['GptOssAttention']/['sinks']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -397,6 +414,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -423,6 +441,7 @@ ".params/['params']/['decoder']/['layers']/['layers_0']/['GptOssAttention']/['value']/['bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -437,6 +456,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -470,6 +490,7 @@ ".params/['params']/['decoder']/['layers']/['layers_0']/['GptOssAttention']/['value']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -484,6 +505,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -525,6 +547,7 @@ ".params/['params']/['decoder']/['layers']/['layers_0']/['GptOssMlp']/['gate']/['bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -539,6 +562,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -565,6 +589,7 @@ ".params/['params']/['decoder']/['layers']/['layers_0']/['GptOssMlp']/['gate']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -579,6 +604,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -614,6 +640,7 @@ ".params/['params']/['decoder']/['layers']/['layers_0']/['GptOssMlp']/['wi_0']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -628,6 +655,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -668,6 +696,7 @@ ".params/['params']/['decoder']/['layers']/['layers_0']/['GptOssMlp']/['wi_0_bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -682,6 +711,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -714,6 +744,7 @@ ".params/['params']/['decoder']/['layers']/['layers_0']/['GptOssMlp']/['wi_1']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -728,6 +759,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -768,6 +800,7 @@ ".params/['params']/['decoder']/['layers']/['layers_0']/['GptOssMlp']/['wi_1_bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -782,6 +815,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -814,6 +848,7 @@ ".params/['params']/['decoder']/['layers']/['layers_0']/['GptOssMlp']/['wo']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -828,6 +863,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -868,6 +904,7 @@ ".params/['params']/['decoder']/['layers']/['layers_0']/['GptOssMlp']/['wo_bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -882,6 +919,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -913,6 +951,7 @@ ".params/['params']/['decoder']/['layers']/['layers_0']/['post_self_attention_layer_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -927,6 +966,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -956,6 +996,7 @@ ".params/['params']/['decoder']/['layers']/['layers_0']/['pre_self_attention_layer_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -970,6 +1011,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -999,6 +1041,7 @@ ".params/['params']/['decoder']/['layers']/['layers_1']/['GptOssAttention']/['key']/['bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1013,6 +1056,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -1046,6 +1090,7 @@ ".params/['params']/['decoder']/['layers']/['layers_1']/['GptOssAttention']/['key']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1060,6 +1105,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -1101,6 +1147,7 @@ ".params/['params']/['decoder']/['layers']/['layers_1']/['GptOssAttention']/['out']/['bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1115,6 +1162,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -1148,6 +1196,7 @@ ".params/['params']/['decoder']/['layers']/['layers_1']/['GptOssAttention']/['out']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1162,6 +1211,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -1203,6 +1253,7 @@ ".params/['params']/['decoder']/['layers']/['layers_1']/['GptOssAttention']/['query']/['bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1217,6 +1268,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -1250,6 +1302,7 @@ ".params/['params']/['decoder']/['layers']/['layers_1']/['GptOssAttention']/['query']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1264,6 +1317,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -1305,6 +1359,7 @@ ".params/['params']/['decoder']/['layers']/['layers_1']/['GptOssAttention']/['sinks']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1319,6 +1374,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -1345,6 +1401,7 @@ ".params/['params']/['decoder']/['layers']/['layers_1']/['GptOssAttention']/['value']/['bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1359,6 +1416,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -1392,6 +1450,7 @@ ".params/['params']/['decoder']/['layers']/['layers_1']/['GptOssAttention']/['value']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1406,6 +1465,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -1447,6 +1507,7 @@ ".params/['params']/['decoder']/['layers']/['layers_1']/['GptOssMlp']/['gate']/['bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1461,6 +1522,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -1487,6 +1549,7 @@ ".params/['params']/['decoder']/['layers']/['layers_1']/['GptOssMlp']/['gate']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1501,6 +1564,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -1536,6 +1600,7 @@ ".params/['params']/['decoder']/['layers']/['layers_1']/['GptOssMlp']/['wi_0']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1550,6 +1615,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -1590,6 +1656,7 @@ ".params/['params']/['decoder']/['layers']/['layers_1']/['GptOssMlp']/['wi_0_bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1604,6 +1671,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -1636,6 +1704,7 @@ ".params/['params']/['decoder']/['layers']/['layers_1']/['GptOssMlp']/['wi_1']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1650,6 +1719,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -1690,6 +1760,7 @@ ".params/['params']/['decoder']/['layers']/['layers_1']/['GptOssMlp']/['wi_1_bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1704,6 +1775,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -1736,6 +1808,7 @@ ".params/['params']/['decoder']/['layers']/['layers_1']/['GptOssMlp']/['wo']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1750,6 +1823,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -1790,6 +1864,7 @@ ".params/['params']/['decoder']/['layers']/['layers_1']/['GptOssMlp']/['wo_bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1804,6 +1879,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -1835,6 +1911,7 @@ ".params/['params']/['decoder']/['layers']/['layers_1']/['post_self_attention_layer_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1849,6 +1926,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -1878,6 +1956,7 @@ ".params/['params']/['decoder']/['layers']/['layers_1']/['pre_self_attention_layer_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1892,6 +1971,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -1921,6 +2001,7 @@ ".params/['params']/['decoder']/['logits_dense']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1935,6 +2016,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -1972,6 +2054,7 @@ ".params/['params']/['token_embedder']/['embedding']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1986,6 +2069,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -2023,6 +2107,7 @@ ".opt_state/[0]/.count": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2037,6 +2122,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -2057,6 +2143,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['decoder_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2071,6 +2158,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -2098,6 +2186,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['layers_0']/['GptOssAttention']/['key']/['bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2112,6 +2201,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -2145,6 +2235,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['layers_0']/['GptOssAttention']/['key']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2159,6 +2250,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -2200,6 +2292,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['layers_0']/['GptOssAttention']/['out']/['bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2214,6 +2307,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -2247,6 +2341,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['layers_0']/['GptOssAttention']/['out']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2261,6 +2356,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -2302,6 +2398,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['layers_0']/['GptOssAttention']/['query']/['bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2316,6 +2413,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -2349,6 +2447,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['layers_0']/['GptOssAttention']/['query']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2363,6 +2462,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -2404,6 +2504,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['layers_0']/['GptOssAttention']/['sinks']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2418,6 +2519,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -2444,6 +2546,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['layers_0']/['GptOssAttention']/['value']/['bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2458,6 +2561,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -2491,6 +2595,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['layers_0']/['GptOssAttention']/['value']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2505,6 +2610,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -2546,6 +2652,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['layers_0']/['GptOssMlp']/['gate']/['bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2560,6 +2667,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -2586,6 +2694,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['layers_0']/['GptOssMlp']/['gate']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2600,6 +2709,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -2635,6 +2745,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['layers_0']/['GptOssMlp']/['wi_0']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2649,6 +2760,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -2689,6 +2801,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['layers_0']/['GptOssMlp']/['wi_0_bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2703,6 +2816,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -2735,6 +2849,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['layers_0']/['GptOssMlp']/['wi_1']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2749,6 +2864,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -2789,6 +2905,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['layers_0']/['GptOssMlp']/['wi_1_bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2803,6 +2920,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -2835,6 +2953,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['layers_0']/['GptOssMlp']/['wo']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2849,6 +2968,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -2889,6 +3009,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['layers_0']/['GptOssMlp']/['wo_bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2903,6 +3024,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -2934,6 +3056,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['layers_0']/['post_self_attention_layer_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2948,6 +3071,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -2977,6 +3101,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['layers_0']/['pre_self_attention_layer_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2991,6 +3116,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -3020,6 +3146,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['layers_1']/['GptOssAttention']/['key']/['bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3034,6 +3161,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -3067,6 +3195,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['layers_1']/['GptOssAttention']/['key']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3081,6 +3210,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -3122,6 +3252,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['layers_1']/['GptOssAttention']/['out']/['bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3136,6 +3267,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -3169,6 +3301,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['layers_1']/['GptOssAttention']/['out']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3183,6 +3316,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -3224,6 +3358,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['layers_1']/['GptOssAttention']/['query']/['bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3238,6 +3373,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -3271,6 +3407,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['layers_1']/['GptOssAttention']/['query']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3285,6 +3422,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -3326,6 +3464,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['layers_1']/['GptOssAttention']/['sinks']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3340,6 +3479,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -3366,6 +3506,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['layers_1']/['GptOssAttention']/['value']/['bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3380,6 +3521,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -3413,6 +3555,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['layers_1']/['GptOssAttention']/['value']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3427,6 +3570,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -3468,6 +3612,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['layers_1']/['GptOssMlp']/['gate']/['bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3482,6 +3627,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -3508,6 +3654,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['layers_1']/['GptOssMlp']/['gate']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3522,6 +3669,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -3557,6 +3705,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['layers_1']/['GptOssMlp']/['wi_0']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3571,6 +3720,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -3611,6 +3761,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['layers_1']/['GptOssMlp']/['wi_0_bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3625,6 +3776,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -3657,6 +3809,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['layers_1']/['GptOssMlp']/['wi_1']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3671,6 +3824,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -3711,6 +3865,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['layers_1']/['GptOssMlp']/['wi_1_bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3725,6 +3880,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -3757,6 +3913,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['layers_1']/['GptOssMlp']/['wo']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3771,6 +3928,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -3811,6 +3969,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['layers_1']/['GptOssMlp']/['wo_bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3825,6 +3984,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -3856,6 +4016,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['layers_1']/['post_self_attention_layer_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3870,6 +4031,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -3899,6 +4061,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['layers_1']/['pre_self_attention_layer_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3913,6 +4076,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -3942,6 +4106,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['logits_dense']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3956,6 +4121,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -3993,6 +4159,7 @@ ".opt_state/[0]/.mu/['params']/['token_embedder']/['embedding']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -4007,6 +4174,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -4044,6 +4212,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['decoder_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -4058,6 +4227,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -4085,6 +4255,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['layers_0']/['GptOssAttention']/['key']/['bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -4099,6 +4270,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -4132,6 +4304,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['layers_0']/['GptOssAttention']/['key']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -4146,6 +4319,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -4187,6 +4361,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['layers_0']/['GptOssAttention']/['out']/['bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -4201,6 +4376,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -4234,6 +4410,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['layers_0']/['GptOssAttention']/['out']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -4248,6 +4425,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -4289,6 +4467,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['layers_0']/['GptOssAttention']/['query']/['bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -4303,6 +4482,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -4336,6 +4516,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['layers_0']/['GptOssAttention']/['query']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -4350,6 +4531,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -4391,6 +4573,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['layers_0']/['GptOssAttention']/['sinks']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -4405,6 +4588,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -4431,6 +4615,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['layers_0']/['GptOssAttention']/['value']/['bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -4445,6 +4630,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -4478,6 +4664,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['layers_0']/['GptOssAttention']/['value']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -4492,6 +4679,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -4533,6 +4721,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['layers_0']/['GptOssMlp']/['gate']/['bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -4547,6 +4736,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -4573,6 +4763,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['layers_0']/['GptOssMlp']/['gate']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -4587,6 +4778,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -4622,6 +4814,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['layers_0']/['GptOssMlp']/['wi_0']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -4636,6 +4829,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -4676,6 +4870,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['layers_0']/['GptOssMlp']/['wi_0_bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -4690,6 +4885,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -4722,6 +4918,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['layers_0']/['GptOssMlp']/['wi_1']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -4736,6 +4933,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -4776,6 +4974,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['layers_0']/['GptOssMlp']/['wi_1_bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -4790,6 +4989,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -4822,6 +5022,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['layers_0']/['GptOssMlp']/['wo']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -4836,6 +5037,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -4876,6 +5078,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['layers_0']/['GptOssMlp']/['wo_bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -4890,6 +5093,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -4921,6 +5125,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['layers_0']/['post_self_attention_layer_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -4935,6 +5140,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -4964,6 +5170,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['layers_0']/['pre_self_attention_layer_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -4978,6 +5185,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -5007,6 +5215,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['layers_1']/['GptOssAttention']/['key']/['bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -5021,6 +5230,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -5054,6 +5264,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['layers_1']/['GptOssAttention']/['key']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -5068,6 +5279,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -5109,6 +5321,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['layers_1']/['GptOssAttention']/['out']/['bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -5123,6 +5336,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -5156,6 +5370,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['layers_1']/['GptOssAttention']/['out']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -5170,6 +5385,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -5211,6 +5427,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['layers_1']/['GptOssAttention']/['query']/['bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -5225,6 +5442,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -5258,6 +5476,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['layers_1']/['GptOssAttention']/['query']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -5272,6 +5491,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -5313,6 +5533,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['layers_1']/['GptOssAttention']/['sinks']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -5327,6 +5548,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -5353,6 +5575,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['layers_1']/['GptOssAttention']/['value']/['bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -5367,6 +5590,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -5400,6 +5624,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['layers_1']/['GptOssAttention']/['value']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -5414,6 +5639,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -5455,6 +5681,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['layers_1']/['GptOssMlp']/['gate']/['bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -5469,6 +5696,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -5495,6 +5723,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['layers_1']/['GptOssMlp']/['gate']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -5509,6 +5738,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -5544,6 +5774,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['layers_1']/['GptOssMlp']/['wi_0']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -5558,6 +5789,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -5598,6 +5830,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['layers_1']/['GptOssMlp']/['wi_0_bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -5612,6 +5845,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -5644,6 +5878,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['layers_1']/['GptOssMlp']/['wi_1']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -5658,6 +5893,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -5698,6 +5934,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['layers_1']/['GptOssMlp']/['wi_1_bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -5712,6 +5949,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -5744,6 +5982,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['layers_1']/['GptOssMlp']/['wo']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -5758,6 +5997,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -5798,6 +6038,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['layers_1']/['GptOssMlp']/['wo_bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -5812,6 +6053,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -5843,6 +6085,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['layers_1']/['post_self_attention_layer_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -5857,6 +6100,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -5886,6 +6130,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['layers_1']/['pre_self_attention_layer_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -5900,6 +6145,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -5929,6 +6175,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['logits_dense']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -5943,6 +6190,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -5980,6 +6228,7 @@ ".opt_state/[0]/.nu/['params']/['token_embedder']/['embedding']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -5994,6 +6243,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -6031,6 +6281,7 @@ ".opt_state/[2]/.count": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -6045,6 +6296,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, diff --git a/tests/utils/sharding_info/gpt-oss-20b/v6e-16/slice_1/named_shardings.json b/tests/utils/sharding_info/gpt-oss-20b/v6e-16/slice_1/named_shardings.json index 6a4eb12a10..78e42a8848 100644 --- a/tests/utils/sharding_info/gpt-oss-20b/v6e-16/slice_1/named_shardings.json +++ b/tests/utils/sharding_info/gpt-oss-20b/v6e-16/slice_1/named_shardings.json @@ -2,6 +2,7 @@ ".step": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -16,6 +17,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -36,6 +38,7 @@ ".params/['params']/['decoder']/['decoder_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -50,6 +53,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -77,6 +81,7 @@ ".params/['params']/['decoder']/['layers']/['layers_0']/['GptOssAttention']/['key']/['bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -91,6 +96,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -124,6 +130,7 @@ ".params/['params']/['decoder']/['layers']/['layers_0']/['GptOssAttention']/['key']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -138,6 +145,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -179,6 +187,7 @@ ".params/['params']/['decoder']/['layers']/['layers_0']/['GptOssAttention']/['out']/['bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -193,6 +202,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -226,6 +236,7 @@ ".params/['params']/['decoder']/['layers']/['layers_0']/['GptOssAttention']/['out']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -240,6 +251,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -281,6 +293,7 @@ ".params/['params']/['decoder']/['layers']/['layers_0']/['GptOssAttention']/['query']/['bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -295,6 +308,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -328,6 +342,7 @@ ".params/['params']/['decoder']/['layers']/['layers_0']/['GptOssAttention']/['query']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -342,6 +357,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -383,6 +399,7 @@ ".params/['params']/['decoder']/['layers']/['layers_0']/['GptOssAttention']/['sinks']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -397,6 +414,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -423,6 +441,7 @@ ".params/['params']/['decoder']/['layers']/['layers_0']/['GptOssAttention']/['value']/['bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -437,6 +456,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -470,6 +490,7 @@ ".params/['params']/['decoder']/['layers']/['layers_0']/['GptOssAttention']/['value']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -484,6 +505,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -525,6 +547,7 @@ ".params/['params']/['decoder']/['layers']/['layers_0']/['GptOssMlp']/['gate']/['bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -539,6 +562,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -565,6 +589,7 @@ ".params/['params']/['decoder']/['layers']/['layers_0']/['GptOssMlp']/['gate']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -579,6 +604,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -614,6 +640,7 @@ ".params/['params']/['decoder']/['layers']/['layers_0']/['GptOssMlp']/['wi_0']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -628,6 +655,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -668,6 +696,7 @@ ".params/['params']/['decoder']/['layers']/['layers_0']/['GptOssMlp']/['wi_0_bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -682,6 +711,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -714,6 +744,7 @@ ".params/['params']/['decoder']/['layers']/['layers_0']/['GptOssMlp']/['wi_1']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -728,6 +759,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -768,6 +800,7 @@ ".params/['params']/['decoder']/['layers']/['layers_0']/['GptOssMlp']/['wi_1_bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -782,6 +815,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -814,6 +848,7 @@ ".params/['params']/['decoder']/['layers']/['layers_0']/['GptOssMlp']/['wo']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -828,6 +863,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -868,6 +904,7 @@ ".params/['params']/['decoder']/['layers']/['layers_0']/['GptOssMlp']/['wo_bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -882,6 +919,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -913,6 +951,7 @@ ".params/['params']/['decoder']/['layers']/['layers_0']/['post_self_attention_layer_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -927,6 +966,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -956,6 +996,7 @@ ".params/['params']/['decoder']/['layers']/['layers_0']/['pre_self_attention_layer_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -970,6 +1011,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -999,6 +1041,7 @@ ".params/['params']/['decoder']/['layers']/['layers_1']/['GptOssAttention']/['key']/['bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1013,6 +1056,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -1046,6 +1090,7 @@ ".params/['params']/['decoder']/['layers']/['layers_1']/['GptOssAttention']/['key']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1060,6 +1105,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -1101,6 +1147,7 @@ ".params/['params']/['decoder']/['layers']/['layers_1']/['GptOssAttention']/['out']/['bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1115,6 +1162,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -1148,6 +1196,7 @@ ".params/['params']/['decoder']/['layers']/['layers_1']/['GptOssAttention']/['out']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1162,6 +1211,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -1203,6 +1253,7 @@ ".params/['params']/['decoder']/['layers']/['layers_1']/['GptOssAttention']/['query']/['bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1217,6 +1268,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -1250,6 +1302,7 @@ ".params/['params']/['decoder']/['layers']/['layers_1']/['GptOssAttention']/['query']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1264,6 +1317,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -1305,6 +1359,7 @@ ".params/['params']/['decoder']/['layers']/['layers_1']/['GptOssAttention']/['sinks']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1319,6 +1374,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -1345,6 +1401,7 @@ ".params/['params']/['decoder']/['layers']/['layers_1']/['GptOssAttention']/['value']/['bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1359,6 +1416,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -1392,6 +1450,7 @@ ".params/['params']/['decoder']/['layers']/['layers_1']/['GptOssAttention']/['value']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1406,6 +1465,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -1447,6 +1507,7 @@ ".params/['params']/['decoder']/['layers']/['layers_1']/['GptOssMlp']/['gate']/['bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1461,6 +1522,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -1487,6 +1549,7 @@ ".params/['params']/['decoder']/['layers']/['layers_1']/['GptOssMlp']/['gate']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1501,6 +1564,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -1536,6 +1600,7 @@ ".params/['params']/['decoder']/['layers']/['layers_1']/['GptOssMlp']/['wi_0']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1550,6 +1615,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -1590,6 +1656,7 @@ ".params/['params']/['decoder']/['layers']/['layers_1']/['GptOssMlp']/['wi_0_bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1604,6 +1671,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -1636,6 +1704,7 @@ ".params/['params']/['decoder']/['layers']/['layers_1']/['GptOssMlp']/['wi_1']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1650,6 +1719,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -1690,6 +1760,7 @@ ".params/['params']/['decoder']/['layers']/['layers_1']/['GptOssMlp']/['wi_1_bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1704,6 +1775,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -1736,6 +1808,7 @@ ".params/['params']/['decoder']/['layers']/['layers_1']/['GptOssMlp']/['wo']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1750,6 +1823,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -1790,6 +1864,7 @@ ".params/['params']/['decoder']/['layers']/['layers_1']/['GptOssMlp']/['wo_bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1804,6 +1879,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -1835,6 +1911,7 @@ ".params/['params']/['decoder']/['layers']/['layers_1']/['post_self_attention_layer_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1849,6 +1926,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -1878,6 +1956,7 @@ ".params/['params']/['decoder']/['layers']/['layers_1']/['pre_self_attention_layer_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1892,6 +1971,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -1921,6 +2001,7 @@ ".params/['params']/['decoder']/['logits_dense']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1935,6 +2016,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -1972,6 +2054,7 @@ ".params/['params']/['token_embedder']/['embedding']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1986,6 +2069,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -2023,6 +2107,7 @@ ".opt_state/[0]/.count": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2037,6 +2122,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -2057,6 +2143,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['decoder_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2071,6 +2158,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -2098,6 +2186,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['layers_0']/['GptOssAttention']/['key']/['bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2112,6 +2201,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -2145,6 +2235,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['layers_0']/['GptOssAttention']/['key']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2159,6 +2250,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -2200,6 +2292,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['layers_0']/['GptOssAttention']/['out']/['bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2214,6 +2307,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -2247,6 +2341,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['layers_0']/['GptOssAttention']/['out']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2261,6 +2356,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -2302,6 +2398,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['layers_0']/['GptOssAttention']/['query']/['bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2316,6 +2413,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -2349,6 +2447,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['layers_0']/['GptOssAttention']/['query']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2363,6 +2462,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -2404,6 +2504,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['layers_0']/['GptOssAttention']/['sinks']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2418,6 +2519,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -2444,6 +2546,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['layers_0']/['GptOssAttention']/['value']/['bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2458,6 +2561,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -2491,6 +2595,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['layers_0']/['GptOssAttention']/['value']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2505,6 +2610,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -2546,6 +2652,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['layers_0']/['GptOssMlp']/['gate']/['bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2560,6 +2667,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -2586,6 +2694,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['layers_0']/['GptOssMlp']/['gate']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2600,6 +2709,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -2635,6 +2745,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['layers_0']/['GptOssMlp']/['wi_0']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2649,6 +2760,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -2689,6 +2801,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['layers_0']/['GptOssMlp']/['wi_0_bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2703,6 +2816,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -2735,6 +2849,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['layers_0']/['GptOssMlp']/['wi_1']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2749,6 +2864,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -2789,6 +2905,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['layers_0']/['GptOssMlp']/['wi_1_bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2803,6 +2920,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -2835,6 +2953,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['layers_0']/['GptOssMlp']/['wo']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2849,6 +2968,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -2889,6 +3009,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['layers_0']/['GptOssMlp']/['wo_bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2903,6 +3024,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -2934,6 +3056,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['layers_0']/['post_self_attention_layer_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2948,6 +3071,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -2977,6 +3101,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['layers_0']/['pre_self_attention_layer_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2991,6 +3116,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -3020,6 +3146,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['layers_1']/['GptOssAttention']/['key']/['bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3034,6 +3161,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -3067,6 +3195,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['layers_1']/['GptOssAttention']/['key']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3081,6 +3210,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -3122,6 +3252,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['layers_1']/['GptOssAttention']/['out']/['bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3136,6 +3267,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -3169,6 +3301,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['layers_1']/['GptOssAttention']/['out']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3183,6 +3316,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -3224,6 +3358,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['layers_1']/['GptOssAttention']/['query']/['bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3238,6 +3373,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -3271,6 +3407,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['layers_1']/['GptOssAttention']/['query']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3285,6 +3422,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -3326,6 +3464,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['layers_1']/['GptOssAttention']/['sinks']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3340,6 +3479,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -3366,6 +3506,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['layers_1']/['GptOssAttention']/['value']/['bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3380,6 +3521,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -3413,6 +3555,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['layers_1']/['GptOssAttention']/['value']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3427,6 +3570,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -3468,6 +3612,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['layers_1']/['GptOssMlp']/['gate']/['bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3482,6 +3627,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -3508,6 +3654,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['layers_1']/['GptOssMlp']/['gate']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3522,6 +3669,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -3557,6 +3705,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['layers_1']/['GptOssMlp']/['wi_0']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3571,6 +3720,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -3611,6 +3761,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['layers_1']/['GptOssMlp']/['wi_0_bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3625,6 +3776,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -3657,6 +3809,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['layers_1']/['GptOssMlp']/['wi_1']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3671,6 +3824,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -3711,6 +3865,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['layers_1']/['GptOssMlp']/['wi_1_bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3725,6 +3880,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -3757,6 +3913,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['layers_1']/['GptOssMlp']/['wo']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3771,6 +3928,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -3811,6 +3969,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['layers_1']/['GptOssMlp']/['wo_bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3825,6 +3984,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -3856,6 +4016,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['layers_1']/['post_self_attention_layer_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3870,6 +4031,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -3899,6 +4061,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['layers_1']/['pre_self_attention_layer_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3913,6 +4076,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -3942,6 +4106,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['logits_dense']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3956,6 +4121,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -3993,6 +4159,7 @@ ".opt_state/[0]/.mu/['params']/['token_embedder']/['embedding']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -4007,6 +4174,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -4044,6 +4212,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['decoder_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -4058,6 +4227,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -4085,6 +4255,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['layers_0']/['GptOssAttention']/['key']/['bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -4099,6 +4270,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -4132,6 +4304,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['layers_0']/['GptOssAttention']/['key']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -4146,6 +4319,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -4187,6 +4361,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['layers_0']/['GptOssAttention']/['out']/['bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -4201,6 +4376,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -4234,6 +4410,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['layers_0']/['GptOssAttention']/['out']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -4248,6 +4425,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -4289,6 +4467,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['layers_0']/['GptOssAttention']/['query']/['bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -4303,6 +4482,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -4336,6 +4516,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['layers_0']/['GptOssAttention']/['query']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -4350,6 +4531,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -4391,6 +4573,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['layers_0']/['GptOssAttention']/['sinks']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -4405,6 +4588,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -4431,6 +4615,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['layers_0']/['GptOssAttention']/['value']/['bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -4445,6 +4630,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -4478,6 +4664,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['layers_0']/['GptOssAttention']/['value']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -4492,6 +4679,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -4533,6 +4721,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['layers_0']/['GptOssMlp']/['gate']/['bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -4547,6 +4736,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -4573,6 +4763,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['layers_0']/['GptOssMlp']/['gate']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -4587,6 +4778,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -4622,6 +4814,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['layers_0']/['GptOssMlp']/['wi_0']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -4636,6 +4829,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -4676,6 +4870,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['layers_0']/['GptOssMlp']/['wi_0_bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -4690,6 +4885,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -4722,6 +4918,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['layers_0']/['GptOssMlp']/['wi_1']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -4736,6 +4933,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -4776,6 +4974,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['layers_0']/['GptOssMlp']/['wi_1_bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -4790,6 +4989,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -4822,6 +5022,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['layers_0']/['GptOssMlp']/['wo']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -4836,6 +5037,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -4876,6 +5078,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['layers_0']/['GptOssMlp']/['wo_bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -4890,6 +5093,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -4921,6 +5125,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['layers_0']/['post_self_attention_layer_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -4935,6 +5140,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -4964,6 +5170,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['layers_0']/['pre_self_attention_layer_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -4978,6 +5185,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -5007,6 +5215,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['layers_1']/['GptOssAttention']/['key']/['bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -5021,6 +5230,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -5054,6 +5264,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['layers_1']/['GptOssAttention']/['key']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -5068,6 +5279,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -5109,6 +5321,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['layers_1']/['GptOssAttention']/['out']/['bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -5123,6 +5336,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -5156,6 +5370,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['layers_1']/['GptOssAttention']/['out']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -5170,6 +5385,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -5211,6 +5427,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['layers_1']/['GptOssAttention']/['query']/['bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -5225,6 +5442,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -5258,6 +5476,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['layers_1']/['GptOssAttention']/['query']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -5272,6 +5491,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -5313,6 +5533,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['layers_1']/['GptOssAttention']/['sinks']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -5327,6 +5548,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -5353,6 +5575,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['layers_1']/['GptOssAttention']/['value']/['bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -5367,6 +5590,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -5400,6 +5624,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['layers_1']/['GptOssAttention']/['value']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -5414,6 +5639,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -5455,6 +5681,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['layers_1']/['GptOssMlp']/['gate']/['bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -5469,6 +5696,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -5495,6 +5723,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['layers_1']/['GptOssMlp']/['gate']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -5509,6 +5738,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -5544,6 +5774,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['layers_1']/['GptOssMlp']/['wi_0']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -5558,6 +5789,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -5598,6 +5830,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['layers_1']/['GptOssMlp']/['wi_0_bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -5612,6 +5845,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -5644,6 +5878,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['layers_1']/['GptOssMlp']/['wi_1']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -5658,6 +5893,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -5698,6 +5934,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['layers_1']/['GptOssMlp']/['wi_1_bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -5712,6 +5949,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -5744,6 +5982,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['layers_1']/['GptOssMlp']/['wo']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -5758,6 +5997,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -5798,6 +6038,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['layers_1']/['GptOssMlp']/['wo_bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -5812,6 +6053,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -5843,6 +6085,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['layers_1']/['post_self_attention_layer_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -5857,6 +6100,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -5886,6 +6130,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['layers_1']/['pre_self_attention_layer_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -5900,6 +6145,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -5929,6 +6175,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['logits_dense']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -5943,6 +6190,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -5980,6 +6228,7 @@ ".opt_state/[0]/.nu/['params']/['token_embedder']/['embedding']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -5994,6 +6243,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -6031,6 +6281,7 @@ ".opt_state/[2]/.count": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -6045,6 +6296,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, diff --git a/tests/utils/sharding_info/gpt-oss-20b/v6e-16/slice_4/named_shardings.json b/tests/utils/sharding_info/gpt-oss-20b/v6e-16/slice_4/named_shardings.json index fffa91ebe5..ed765f1d18 100644 --- a/tests/utils/sharding_info/gpt-oss-20b/v6e-16/slice_4/named_shardings.json +++ b/tests/utils/sharding_info/gpt-oss-20b/v6e-16/slice_4/named_shardings.json @@ -2,6 +2,7 @@ ".step": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -16,6 +17,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -36,6 +38,7 @@ ".params/['params']/['decoder']/['decoder_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -50,6 +53,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -77,6 +81,7 @@ ".params/['params']/['decoder']/['layers']/['layers_0']/['GptOssAttention']/['key']/['bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -91,6 +96,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -124,6 +130,7 @@ ".params/['params']/['decoder']/['layers']/['layers_0']/['GptOssAttention']/['key']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -138,6 +145,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -179,6 +187,7 @@ ".params/['params']/['decoder']/['layers']/['layers_0']/['GptOssAttention']/['out']/['bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -193,6 +202,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -226,6 +236,7 @@ ".params/['params']/['decoder']/['layers']/['layers_0']/['GptOssAttention']/['out']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -240,6 +251,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -281,6 +293,7 @@ ".params/['params']/['decoder']/['layers']/['layers_0']/['GptOssAttention']/['query']/['bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -295,6 +308,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -328,6 +342,7 @@ ".params/['params']/['decoder']/['layers']/['layers_0']/['GptOssAttention']/['query']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -342,6 +357,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -383,6 +399,7 @@ ".params/['params']/['decoder']/['layers']/['layers_0']/['GptOssAttention']/['sinks']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -397,6 +414,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -423,6 +441,7 @@ ".params/['params']/['decoder']/['layers']/['layers_0']/['GptOssAttention']/['value']/['bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -437,6 +456,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -470,6 +490,7 @@ ".params/['params']/['decoder']/['layers']/['layers_0']/['GptOssAttention']/['value']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -484,6 +505,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -525,6 +547,7 @@ ".params/['params']/['decoder']/['layers']/['layers_0']/['GptOssMlp']/['gate']/['bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -539,6 +562,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -565,6 +589,7 @@ ".params/['params']/['decoder']/['layers']/['layers_0']/['GptOssMlp']/['gate']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -579,6 +604,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -614,6 +640,7 @@ ".params/['params']/['decoder']/['layers']/['layers_0']/['GptOssMlp']/['wi_0']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -628,6 +655,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -668,6 +696,7 @@ ".params/['params']/['decoder']/['layers']/['layers_0']/['GptOssMlp']/['wi_0_bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -682,6 +711,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -714,6 +744,7 @@ ".params/['params']/['decoder']/['layers']/['layers_0']/['GptOssMlp']/['wi_1']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -728,6 +759,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -768,6 +800,7 @@ ".params/['params']/['decoder']/['layers']/['layers_0']/['GptOssMlp']/['wi_1_bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -782,6 +815,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -814,6 +848,7 @@ ".params/['params']/['decoder']/['layers']/['layers_0']/['GptOssMlp']/['wo']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -828,6 +863,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -868,6 +904,7 @@ ".params/['params']/['decoder']/['layers']/['layers_0']/['GptOssMlp']/['wo_bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -882,6 +919,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -913,6 +951,7 @@ ".params/['params']/['decoder']/['layers']/['layers_0']/['post_self_attention_layer_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -927,6 +966,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -956,6 +996,7 @@ ".params/['params']/['decoder']/['layers']/['layers_0']/['pre_self_attention_layer_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -970,6 +1011,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -999,6 +1041,7 @@ ".params/['params']/['decoder']/['layers']/['layers_1']/['GptOssAttention']/['key']/['bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1013,6 +1056,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -1046,6 +1090,7 @@ ".params/['params']/['decoder']/['layers']/['layers_1']/['GptOssAttention']/['key']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1060,6 +1105,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -1101,6 +1147,7 @@ ".params/['params']/['decoder']/['layers']/['layers_1']/['GptOssAttention']/['out']/['bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1115,6 +1162,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -1148,6 +1196,7 @@ ".params/['params']/['decoder']/['layers']/['layers_1']/['GptOssAttention']/['out']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1162,6 +1211,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -1203,6 +1253,7 @@ ".params/['params']/['decoder']/['layers']/['layers_1']/['GptOssAttention']/['query']/['bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1217,6 +1268,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -1250,6 +1302,7 @@ ".params/['params']/['decoder']/['layers']/['layers_1']/['GptOssAttention']/['query']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1264,6 +1317,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -1305,6 +1359,7 @@ ".params/['params']/['decoder']/['layers']/['layers_1']/['GptOssAttention']/['sinks']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1319,6 +1374,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -1345,6 +1401,7 @@ ".params/['params']/['decoder']/['layers']/['layers_1']/['GptOssAttention']/['value']/['bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1359,6 +1416,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -1392,6 +1450,7 @@ ".params/['params']/['decoder']/['layers']/['layers_1']/['GptOssAttention']/['value']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1406,6 +1465,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -1447,6 +1507,7 @@ ".params/['params']/['decoder']/['layers']/['layers_1']/['GptOssMlp']/['gate']/['bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1461,6 +1522,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -1487,6 +1549,7 @@ ".params/['params']/['decoder']/['layers']/['layers_1']/['GptOssMlp']/['gate']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1501,6 +1564,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -1536,6 +1600,7 @@ ".params/['params']/['decoder']/['layers']/['layers_1']/['GptOssMlp']/['wi_0']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1550,6 +1615,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -1590,6 +1656,7 @@ ".params/['params']/['decoder']/['layers']/['layers_1']/['GptOssMlp']/['wi_0_bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1604,6 +1671,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -1636,6 +1704,7 @@ ".params/['params']/['decoder']/['layers']/['layers_1']/['GptOssMlp']/['wi_1']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1650,6 +1719,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -1690,6 +1760,7 @@ ".params/['params']/['decoder']/['layers']/['layers_1']/['GptOssMlp']/['wi_1_bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1704,6 +1775,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -1736,6 +1808,7 @@ ".params/['params']/['decoder']/['layers']/['layers_1']/['GptOssMlp']/['wo']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1750,6 +1823,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -1790,6 +1864,7 @@ ".params/['params']/['decoder']/['layers']/['layers_1']/['GptOssMlp']/['wo_bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1804,6 +1879,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -1835,6 +1911,7 @@ ".params/['params']/['decoder']/['layers']/['layers_1']/['post_self_attention_layer_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1849,6 +1926,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -1878,6 +1956,7 @@ ".params/['params']/['decoder']/['layers']/['layers_1']/['pre_self_attention_layer_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1892,6 +1971,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -1921,6 +2001,7 @@ ".params/['params']/['decoder']/['logits_dense']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1935,6 +2016,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -1972,6 +2054,7 @@ ".params/['params']/['token_embedder']/['embedding']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1986,6 +2069,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -2023,6 +2107,7 @@ ".opt_state/[0]/.count": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2037,6 +2122,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -2057,6 +2143,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['decoder_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2071,6 +2158,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -2098,6 +2186,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['layers_0']/['GptOssAttention']/['key']/['bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2112,6 +2201,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -2145,6 +2235,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['layers_0']/['GptOssAttention']/['key']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2159,6 +2250,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -2200,6 +2292,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['layers_0']/['GptOssAttention']/['out']/['bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2214,6 +2307,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -2247,6 +2341,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['layers_0']/['GptOssAttention']/['out']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2261,6 +2356,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -2302,6 +2398,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['layers_0']/['GptOssAttention']/['query']/['bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2316,6 +2413,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -2349,6 +2447,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['layers_0']/['GptOssAttention']/['query']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2363,6 +2462,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -2404,6 +2504,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['layers_0']/['GptOssAttention']/['sinks']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2418,6 +2519,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -2444,6 +2546,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['layers_0']/['GptOssAttention']/['value']/['bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2458,6 +2561,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -2491,6 +2595,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['layers_0']/['GptOssAttention']/['value']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2505,6 +2610,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -2546,6 +2652,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['layers_0']/['GptOssMlp']/['gate']/['bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2560,6 +2667,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -2586,6 +2694,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['layers_0']/['GptOssMlp']/['gate']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2600,6 +2709,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -2635,6 +2745,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['layers_0']/['GptOssMlp']/['wi_0']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2649,6 +2760,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -2689,6 +2801,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['layers_0']/['GptOssMlp']/['wi_0_bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2703,6 +2816,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -2735,6 +2849,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['layers_0']/['GptOssMlp']/['wi_1']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2749,6 +2864,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -2789,6 +2905,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['layers_0']/['GptOssMlp']/['wi_1_bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2803,6 +2920,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -2835,6 +2953,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['layers_0']/['GptOssMlp']/['wo']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2849,6 +2968,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -2889,6 +3009,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['layers_0']/['GptOssMlp']/['wo_bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2903,6 +3024,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -2934,6 +3056,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['layers_0']/['post_self_attention_layer_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2948,6 +3071,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -2977,6 +3101,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['layers_0']/['pre_self_attention_layer_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2991,6 +3116,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -3020,6 +3146,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['layers_1']/['GptOssAttention']/['key']/['bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3034,6 +3161,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -3067,6 +3195,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['layers_1']/['GptOssAttention']/['key']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3081,6 +3210,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -3122,6 +3252,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['layers_1']/['GptOssAttention']/['out']/['bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3136,6 +3267,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -3169,6 +3301,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['layers_1']/['GptOssAttention']/['out']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3183,6 +3316,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -3224,6 +3358,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['layers_1']/['GptOssAttention']/['query']/['bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3238,6 +3373,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -3271,6 +3407,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['layers_1']/['GptOssAttention']/['query']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3285,6 +3422,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -3326,6 +3464,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['layers_1']/['GptOssAttention']/['sinks']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3340,6 +3479,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -3366,6 +3506,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['layers_1']/['GptOssAttention']/['value']/['bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3380,6 +3521,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -3413,6 +3555,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['layers_1']/['GptOssAttention']/['value']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3427,6 +3570,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -3468,6 +3612,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['layers_1']/['GptOssMlp']/['gate']/['bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3482,6 +3627,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -3508,6 +3654,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['layers_1']/['GptOssMlp']/['gate']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3522,6 +3669,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -3557,6 +3705,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['layers_1']/['GptOssMlp']/['wi_0']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3571,6 +3720,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -3611,6 +3761,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['layers_1']/['GptOssMlp']/['wi_0_bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3625,6 +3776,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -3657,6 +3809,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['layers_1']/['GptOssMlp']/['wi_1']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3671,6 +3824,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -3711,6 +3865,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['layers_1']/['GptOssMlp']/['wi_1_bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3725,6 +3880,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -3757,6 +3913,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['layers_1']/['GptOssMlp']/['wo']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3771,6 +3928,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -3811,6 +3969,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['layers_1']/['GptOssMlp']/['wo_bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3825,6 +3984,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -3856,6 +4016,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['layers_1']/['post_self_attention_layer_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3870,6 +4031,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -3899,6 +4061,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['layers_1']/['pre_self_attention_layer_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3913,6 +4076,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -3942,6 +4106,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['logits_dense']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -3956,6 +4121,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -3993,6 +4159,7 @@ ".opt_state/[0]/.mu/['params']/['token_embedder']/['embedding']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -4007,6 +4174,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -4044,6 +4212,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['decoder_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -4058,6 +4227,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -4085,6 +4255,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['layers_0']/['GptOssAttention']/['key']/['bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -4099,6 +4270,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -4132,6 +4304,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['layers_0']/['GptOssAttention']/['key']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -4146,6 +4319,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -4187,6 +4361,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['layers_0']/['GptOssAttention']/['out']/['bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -4201,6 +4376,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -4234,6 +4410,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['layers_0']/['GptOssAttention']/['out']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -4248,6 +4425,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -4289,6 +4467,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['layers_0']/['GptOssAttention']/['query']/['bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -4303,6 +4482,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -4336,6 +4516,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['layers_0']/['GptOssAttention']/['query']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -4350,6 +4531,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -4391,6 +4573,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['layers_0']/['GptOssAttention']/['sinks']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -4405,6 +4588,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -4431,6 +4615,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['layers_0']/['GptOssAttention']/['value']/['bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -4445,6 +4630,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -4478,6 +4664,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['layers_0']/['GptOssAttention']/['value']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -4492,6 +4679,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -4533,6 +4721,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['layers_0']/['GptOssMlp']/['gate']/['bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -4547,6 +4736,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -4573,6 +4763,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['layers_0']/['GptOssMlp']/['gate']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -4587,6 +4778,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -4622,6 +4814,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['layers_0']/['GptOssMlp']/['wi_0']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -4636,6 +4829,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -4676,6 +4870,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['layers_0']/['GptOssMlp']/['wi_0_bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -4690,6 +4885,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -4722,6 +4918,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['layers_0']/['GptOssMlp']/['wi_1']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -4736,6 +4933,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -4776,6 +4974,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['layers_0']/['GptOssMlp']/['wi_1_bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -4790,6 +4989,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -4822,6 +5022,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['layers_0']/['GptOssMlp']/['wo']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -4836,6 +5037,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -4876,6 +5078,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['layers_0']/['GptOssMlp']/['wo_bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -4890,6 +5093,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -4921,6 +5125,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['layers_0']/['post_self_attention_layer_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -4935,6 +5140,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -4964,6 +5170,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['layers_0']/['pre_self_attention_layer_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -4978,6 +5185,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -5007,6 +5215,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['layers_1']/['GptOssAttention']/['key']/['bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -5021,6 +5230,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -5054,6 +5264,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['layers_1']/['GptOssAttention']/['key']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -5068,6 +5279,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -5109,6 +5321,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['layers_1']/['GptOssAttention']/['out']/['bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -5123,6 +5336,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -5156,6 +5370,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['layers_1']/['GptOssAttention']/['out']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -5170,6 +5385,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -5211,6 +5427,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['layers_1']/['GptOssAttention']/['query']/['bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -5225,6 +5442,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -5258,6 +5476,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['layers_1']/['GptOssAttention']/['query']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -5272,6 +5491,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -5313,6 +5533,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['layers_1']/['GptOssAttention']/['sinks']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -5327,6 +5548,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -5353,6 +5575,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['layers_1']/['GptOssAttention']/['value']/['bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -5367,6 +5590,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -5400,6 +5624,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['layers_1']/['GptOssAttention']/['value']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -5414,6 +5639,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -5455,6 +5681,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['layers_1']/['GptOssMlp']/['gate']/['bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -5469,6 +5696,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -5495,6 +5723,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['layers_1']/['GptOssMlp']/['gate']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -5509,6 +5738,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -5544,6 +5774,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['layers_1']/['GptOssMlp']/['wi_0']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -5558,6 +5789,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -5598,6 +5830,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['layers_1']/['GptOssMlp']/['wi_0_bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -5612,6 +5845,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -5644,6 +5878,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['layers_1']/['GptOssMlp']/['wi_1']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -5658,6 +5893,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -5698,6 +5934,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['layers_1']/['GptOssMlp']/['wi_1_bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -5712,6 +5949,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -5744,6 +5982,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['layers_1']/['GptOssMlp']/['wo']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -5758,6 +5997,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -5798,6 +6038,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['layers_1']/['GptOssMlp']/['wo_bias']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -5812,6 +6053,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -5843,6 +6085,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['layers_1']/['post_self_attention_layer_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -5857,6 +6100,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -5886,6 +6130,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['layers_1']/['pre_self_attention_layer_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -5900,6 +6145,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -5929,6 +6175,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['logits_dense']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -5943,6 +6190,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -5980,6 +6228,7 @@ ".opt_state/[0]/.nu/['params']/['token_embedder']/['embedding']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -5994,6 +6243,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -6031,6 +6281,7 @@ ".opt_state/[2]/.count": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -6045,6 +6296,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, diff --git a/tests/utils/sharding_info/qwen3-0.6b/tpu7x-16/slice_1/named_shardings.json b/tests/utils/sharding_info/qwen3-0.6b/tpu7x-16/slice_1/named_shardings.json index 0ad9713479..6208b4ba80 100644 --- a/tests/utils/sharding_info/qwen3-0.6b/tpu7x-16/slice_1/named_shardings.json +++ b/tests/utils/sharding_info/qwen3-0.6b/tpu7x-16/slice_1/named_shardings.json @@ -2,6 +2,7 @@ ".step": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -16,6 +17,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -36,6 +38,7 @@ ".params/['params']/['decoder']/['decoder_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -50,6 +53,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -77,6 +81,7 @@ ".params/['params']/['decoder']/['layers']/['mlp']/['wi_0']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -91,6 +96,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -130,6 +136,7 @@ ".params/['params']/['decoder']/['layers']/['mlp']/['wi_1']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -144,6 +151,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -183,6 +191,7 @@ ".params/['params']/['decoder']/['layers']/['mlp']/['wo']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -197,6 +206,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -236,6 +246,7 @@ ".params/['params']/['decoder']/['layers']/['post_self_attention_layer_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -250,6 +261,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -279,6 +291,7 @@ ".params/['params']/['decoder']/['layers']/['pre_self_attention_layer_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -293,6 +306,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -322,6 +336,7 @@ ".params/['params']/['decoder']/['layers']/['self_attention']/['key']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -336,6 +351,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -377,6 +393,7 @@ ".params/['params']/['decoder']/['layers']/['self_attention']/['key_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -391,6 +408,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -420,6 +438,7 @@ ".params/['params']/['decoder']/['layers']/['self_attention']/['out']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -434,6 +453,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -475,6 +495,7 @@ ".params/['params']/['decoder']/['layers']/['self_attention']/['query']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -489,6 +510,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -530,6 +552,7 @@ ".params/['params']/['decoder']/['layers']/['self_attention']/['query_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -544,6 +567,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -573,6 +597,7 @@ ".params/['params']/['decoder']/['layers']/['self_attention']/['value']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -587,6 +612,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -628,6 +654,7 @@ ".params/['params']/['token_embedder']/['embedding']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -642,6 +669,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -679,6 +707,7 @@ ".opt_state/[0]/.count": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -693,6 +722,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -713,6 +743,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['decoder_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -727,6 +758,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -754,6 +786,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['mlp']/['wi_0']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -768,6 +801,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -807,6 +841,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['mlp']/['wi_1']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -821,6 +856,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -860,6 +896,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['mlp']/['wo']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -874,6 +911,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -913,6 +951,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['post_self_attention_layer_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -927,6 +966,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -956,6 +996,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['pre_self_attention_layer_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -970,6 +1011,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -999,6 +1041,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['self_attention']/['key']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1013,6 +1056,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -1054,6 +1098,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['self_attention']/['key_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1068,6 +1113,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -1097,6 +1143,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['self_attention']/['out']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1111,6 +1158,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -1152,6 +1200,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['self_attention']/['query']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1166,6 +1215,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -1207,6 +1257,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['self_attention']/['query_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1221,6 +1272,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -1250,6 +1302,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['self_attention']/['value']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1264,6 +1317,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -1305,6 +1359,7 @@ ".opt_state/[0]/.mu/['params']/['token_embedder']/['embedding']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1319,6 +1374,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -1356,6 +1412,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['decoder_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1370,6 +1427,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -1397,6 +1455,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['mlp']/['wi_0']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1411,6 +1470,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -1450,6 +1510,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['mlp']/['wi_1']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1464,6 +1525,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -1503,6 +1565,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['mlp']/['wo']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1517,6 +1580,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -1556,6 +1620,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['post_self_attention_layer_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1570,6 +1635,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -1599,6 +1665,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['pre_self_attention_layer_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1613,6 +1680,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -1642,6 +1710,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['self_attention']/['key']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1656,6 +1725,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -1697,6 +1767,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['self_attention']/['key_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1711,6 +1782,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -1740,6 +1812,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['self_attention']/['out']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1754,6 +1827,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -1795,6 +1869,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['self_attention']/['query']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1809,6 +1884,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -1850,6 +1926,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['self_attention']/['query_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1864,6 +1941,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -1893,6 +1971,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['self_attention']/['value']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1907,6 +1986,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -1948,6 +2028,7 @@ ".opt_state/[0]/.nu/['params']/['token_embedder']/['embedding']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1962,6 +2043,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -1999,6 +2081,7 @@ ".opt_state/[2]/.count": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2013,6 +2096,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, diff --git a/tests/utils/sharding_info/qwen3-0.6b/tpu7x-16/slice_4/named_shardings.json b/tests/utils/sharding_info/qwen3-0.6b/tpu7x-16/slice_4/named_shardings.json index 8e13360273..31499e643e 100644 --- a/tests/utils/sharding_info/qwen3-0.6b/tpu7x-16/slice_4/named_shardings.json +++ b/tests/utils/sharding_info/qwen3-0.6b/tpu7x-16/slice_4/named_shardings.json @@ -2,6 +2,7 @@ ".step": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -16,6 +17,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -36,6 +38,7 @@ ".params/['params']/['decoder']/['decoder_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -50,6 +53,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -77,6 +81,7 @@ ".params/['params']/['decoder']/['layers']/['mlp']/['wi_0']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -91,6 +96,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -130,6 +136,7 @@ ".params/['params']/['decoder']/['layers']/['mlp']/['wi_1']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -144,6 +151,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -183,6 +191,7 @@ ".params/['params']/['decoder']/['layers']/['mlp']/['wo']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -197,6 +206,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -236,6 +246,7 @@ ".params/['params']/['decoder']/['layers']/['post_self_attention_layer_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -250,6 +261,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -279,6 +291,7 @@ ".params/['params']/['decoder']/['layers']/['pre_self_attention_layer_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -293,6 +306,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -322,6 +336,7 @@ ".params/['params']/['decoder']/['layers']/['self_attention']/['key']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -336,6 +351,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -377,6 +393,7 @@ ".params/['params']/['decoder']/['layers']/['self_attention']/['key_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -391,6 +408,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -420,6 +438,7 @@ ".params/['params']/['decoder']/['layers']/['self_attention']/['out']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -434,6 +453,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -475,6 +495,7 @@ ".params/['params']/['decoder']/['layers']/['self_attention']/['query']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -489,6 +510,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -530,6 +552,7 @@ ".params/['params']/['decoder']/['layers']/['self_attention']/['query_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -544,6 +567,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -573,6 +597,7 @@ ".params/['params']/['decoder']/['layers']/['self_attention']/['value']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -587,6 +612,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -628,6 +654,7 @@ ".params/['params']/['token_embedder']/['embedding']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -642,6 +669,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -679,6 +707,7 @@ ".opt_state/[0]/.count": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -693,6 +722,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -713,6 +743,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['decoder_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -727,6 +758,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -754,6 +786,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['mlp']/['wi_0']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -768,6 +801,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -807,6 +841,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['mlp']/['wi_1']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -821,6 +856,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -860,6 +896,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['mlp']/['wo']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -874,6 +911,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -913,6 +951,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['post_self_attention_layer_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -927,6 +966,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -956,6 +996,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['pre_self_attention_layer_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -970,6 +1011,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -999,6 +1041,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['self_attention']/['key']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1013,6 +1056,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -1054,6 +1098,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['self_attention']/['key_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1068,6 +1113,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -1097,6 +1143,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['self_attention']/['out']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1111,6 +1158,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -1152,6 +1200,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['self_attention']/['query']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1166,6 +1215,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -1207,6 +1257,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['self_attention']/['query_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1221,6 +1272,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -1250,6 +1302,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['self_attention']/['value']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1264,6 +1317,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -1305,6 +1359,7 @@ ".opt_state/[0]/.mu/['params']/['token_embedder']/['embedding']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1319,6 +1374,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -1356,6 +1412,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['decoder_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1370,6 +1427,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -1397,6 +1455,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['mlp']/['wi_0']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1411,6 +1470,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -1450,6 +1510,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['mlp']/['wi_1']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1464,6 +1525,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -1503,6 +1565,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['mlp']/['wo']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1517,6 +1580,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -1556,6 +1620,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['post_self_attention_layer_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1570,6 +1635,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -1599,6 +1665,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['pre_self_attention_layer_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1613,6 +1680,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -1642,6 +1710,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['self_attention']/['key']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1656,6 +1725,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -1697,6 +1767,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['self_attention']/['key_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1711,6 +1782,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -1740,6 +1812,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['self_attention']/['out']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1754,6 +1827,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -1795,6 +1869,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['self_attention']/['query']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1809,6 +1884,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -1850,6 +1926,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['self_attention']/['query_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1864,6 +1941,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -1893,6 +1971,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['self_attention']/['value']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1907,6 +1986,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -1948,6 +2028,7 @@ ".opt_state/[0]/.nu/['params']/['token_embedder']/['embedding']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1962,6 +2043,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -1999,6 +2081,7 @@ ".opt_state/[2]/.count": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2013,6 +2096,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, diff --git a/tests/utils/sharding_info/qwen3-0.6b/v5p-16/slice_1/named_shardings.json b/tests/utils/sharding_info/qwen3-0.6b/v5p-16/slice_1/named_shardings.json index 40d1315185..2cce1577f2 100644 --- a/tests/utils/sharding_info/qwen3-0.6b/v5p-16/slice_1/named_shardings.json +++ b/tests/utils/sharding_info/qwen3-0.6b/v5p-16/slice_1/named_shardings.json @@ -2,6 +2,7 @@ ".step": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -16,6 +17,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -36,6 +38,7 @@ ".params/['params']/['decoder']/['decoder_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -50,6 +53,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -77,6 +81,7 @@ ".params/['params']/['decoder']/['layers']/['mlp']/['wi_0']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -91,6 +96,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -130,6 +136,7 @@ ".params/['params']/['decoder']/['layers']/['mlp']/['wi_1']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -144,6 +151,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -183,6 +191,7 @@ ".params/['params']/['decoder']/['layers']/['mlp']/['wo']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -197,6 +206,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -236,6 +246,7 @@ ".params/['params']/['decoder']/['layers']/['post_self_attention_layer_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -250,6 +261,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -279,6 +291,7 @@ ".params/['params']/['decoder']/['layers']/['pre_self_attention_layer_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -293,6 +306,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -322,6 +336,7 @@ ".params/['params']/['decoder']/['layers']/['self_attention']/['key']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -336,6 +351,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -377,6 +393,7 @@ ".params/['params']/['decoder']/['layers']/['self_attention']/['key_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -391,6 +408,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -420,6 +438,7 @@ ".params/['params']/['decoder']/['layers']/['self_attention']/['out']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -434,6 +453,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -475,6 +495,7 @@ ".params/['params']/['decoder']/['layers']/['self_attention']/['query']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -489,6 +510,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -530,6 +552,7 @@ ".params/['params']/['decoder']/['layers']/['self_attention']/['query_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -544,6 +567,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -573,6 +597,7 @@ ".params/['params']/['decoder']/['layers']/['self_attention']/['value']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -587,6 +612,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -628,6 +654,7 @@ ".params/['params']/['token_embedder']/['embedding']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -642,6 +669,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -679,6 +707,7 @@ ".opt_state/[0]/.count": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -693,6 +722,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -713,6 +743,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['decoder_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -727,6 +758,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -754,6 +786,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['mlp']/['wi_0']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -768,6 +801,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -807,6 +841,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['mlp']/['wi_1']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -821,6 +856,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -860,6 +896,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['mlp']/['wo']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -874,6 +911,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -913,6 +951,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['post_self_attention_layer_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -927,6 +966,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -956,6 +996,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['pre_self_attention_layer_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -970,6 +1011,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -999,6 +1041,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['self_attention']/['key']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1013,6 +1056,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -1054,6 +1098,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['self_attention']/['key_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1068,6 +1113,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -1097,6 +1143,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['self_attention']/['out']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1111,6 +1158,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -1152,6 +1200,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['self_attention']/['query']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1166,6 +1215,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -1207,6 +1257,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['self_attention']/['query_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1221,6 +1272,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -1250,6 +1302,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['self_attention']/['value']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1264,6 +1317,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -1305,6 +1359,7 @@ ".opt_state/[0]/.mu/['params']/['token_embedder']/['embedding']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1319,6 +1374,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -1356,6 +1412,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['decoder_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1370,6 +1427,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -1397,6 +1455,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['mlp']/['wi_0']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1411,6 +1470,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -1450,6 +1510,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['mlp']/['wi_1']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1464,6 +1525,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -1503,6 +1565,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['mlp']/['wo']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1517,6 +1580,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -1556,6 +1620,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['post_self_attention_layer_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1570,6 +1635,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -1599,6 +1665,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['pre_self_attention_layer_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1613,6 +1680,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -1642,6 +1710,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['self_attention']/['key']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1656,6 +1725,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -1697,6 +1767,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['self_attention']/['key_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1711,6 +1782,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -1740,6 +1812,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['self_attention']/['out']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1754,6 +1827,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -1795,6 +1869,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['self_attention']/['query']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1809,6 +1884,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -1850,6 +1926,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['self_attention']/['query_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1864,6 +1941,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -1893,6 +1971,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['self_attention']/['value']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1907,6 +1986,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -1948,6 +2028,7 @@ ".opt_state/[0]/.nu/['params']/['token_embedder']/['embedding']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1962,6 +2043,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, @@ -1999,6 +2081,7 @@ ".opt_state/[2]/.count": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2013,6 +2096,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 8, diff --git a/tests/utils/sharding_info/qwen3-0.6b/v5p-16/slice_4/named_shardings.json b/tests/utils/sharding_info/qwen3-0.6b/v5p-16/slice_4/named_shardings.json index 5fc1a68eed..b9512d15f0 100644 --- a/tests/utils/sharding_info/qwen3-0.6b/v5p-16/slice_4/named_shardings.json +++ b/tests/utils/sharding_info/qwen3-0.6b/v5p-16/slice_4/named_shardings.json @@ -2,6 +2,7 @@ ".step": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -16,6 +17,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -36,6 +38,7 @@ ".params/['params']/['decoder']/['decoder_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -50,6 +53,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -77,6 +81,7 @@ ".params/['params']/['decoder']/['layers']/['mlp']/['wi_0']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -91,6 +96,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -130,6 +136,7 @@ ".params/['params']/['decoder']/['layers']/['mlp']/['wi_1']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -144,6 +151,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -183,6 +191,7 @@ ".params/['params']/['decoder']/['layers']/['mlp']/['wo']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -197,6 +206,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -236,6 +246,7 @@ ".params/['params']/['decoder']/['layers']/['post_self_attention_layer_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -250,6 +261,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -279,6 +291,7 @@ ".params/['params']/['decoder']/['layers']/['pre_self_attention_layer_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -293,6 +306,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -322,6 +336,7 @@ ".params/['params']/['decoder']/['layers']/['self_attention']/['key']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -336,6 +351,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -377,6 +393,7 @@ ".params/['params']/['decoder']/['layers']/['self_attention']/['key_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -391,6 +408,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -420,6 +438,7 @@ ".params/['params']/['decoder']/['layers']/['self_attention']/['out']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -434,6 +453,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -475,6 +495,7 @@ ".params/['params']/['decoder']/['layers']/['self_attention']/['query']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -489,6 +510,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -530,6 +552,7 @@ ".params/['params']/['decoder']/['layers']/['self_attention']/['query_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -544,6 +567,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -573,6 +597,7 @@ ".params/['params']/['decoder']/['layers']/['self_attention']/['value']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -587,6 +612,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -628,6 +654,7 @@ ".params/['params']/['token_embedder']/['embedding']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -642,6 +669,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -679,6 +707,7 @@ ".opt_state/[0]/.count": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -693,6 +722,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -713,6 +743,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['decoder_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -727,6 +758,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -754,6 +786,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['mlp']/['wi_0']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -768,6 +801,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -807,6 +841,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['mlp']/['wi_1']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -821,6 +856,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -860,6 +896,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['mlp']/['wo']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -874,6 +911,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -913,6 +951,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['post_self_attention_layer_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -927,6 +966,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -956,6 +996,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['pre_self_attention_layer_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -970,6 +1011,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -999,6 +1041,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['self_attention']/['key']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1013,6 +1056,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -1054,6 +1098,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['self_attention']/['key_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1068,6 +1113,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -1097,6 +1143,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['self_attention']/['out']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1111,6 +1158,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -1152,6 +1200,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['self_attention']/['query']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1166,6 +1215,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -1207,6 +1257,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['self_attention']/['query_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1221,6 +1272,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -1250,6 +1302,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['self_attention']/['value']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1264,6 +1317,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -1305,6 +1359,7 @@ ".opt_state/[0]/.mu/['params']/['token_embedder']/['embedding']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1319,6 +1374,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -1356,6 +1412,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['decoder_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1370,6 +1427,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -1397,6 +1455,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['mlp']/['wi_0']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1411,6 +1470,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -1450,6 +1510,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['mlp']/['wi_1']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1464,6 +1525,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -1503,6 +1565,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['mlp']/['wo']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1517,6 +1580,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -1556,6 +1620,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['post_self_attention_layer_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1570,6 +1635,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -1599,6 +1665,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['pre_self_attention_layer_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1613,6 +1680,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -1642,6 +1710,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['self_attention']/['key']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1656,6 +1725,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -1697,6 +1767,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['self_attention']/['key_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1711,6 +1782,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -1740,6 +1812,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['self_attention']/['out']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1754,6 +1827,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -1795,6 +1869,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['self_attention']/['query']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1809,6 +1884,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -1850,6 +1926,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['self_attention']/['query_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1864,6 +1941,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -1893,6 +1971,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['self_attention']/['value']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1907,6 +1986,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -1948,6 +2028,7 @@ ".opt_state/[0]/.nu/['params']/['token_embedder']/['embedding']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1962,6 +2043,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, @@ -1999,6 +2081,7 @@ ".opt_state/[2]/.count": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2013,6 +2096,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 8, diff --git a/tests/utils/sharding_info/qwen3-0.6b/v6e-16/slice_1/named_shardings.json b/tests/utils/sharding_info/qwen3-0.6b/v6e-16/slice_1/named_shardings.json index 0ad9713479..6208b4ba80 100644 --- a/tests/utils/sharding_info/qwen3-0.6b/v6e-16/slice_1/named_shardings.json +++ b/tests/utils/sharding_info/qwen3-0.6b/v6e-16/slice_1/named_shardings.json @@ -2,6 +2,7 @@ ".step": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -16,6 +17,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -36,6 +38,7 @@ ".params/['params']/['decoder']/['decoder_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -50,6 +53,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -77,6 +81,7 @@ ".params/['params']/['decoder']/['layers']/['mlp']/['wi_0']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -91,6 +96,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -130,6 +136,7 @@ ".params/['params']/['decoder']/['layers']/['mlp']/['wi_1']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -144,6 +151,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -183,6 +191,7 @@ ".params/['params']/['decoder']/['layers']/['mlp']/['wo']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -197,6 +206,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -236,6 +246,7 @@ ".params/['params']/['decoder']/['layers']/['post_self_attention_layer_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -250,6 +261,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -279,6 +291,7 @@ ".params/['params']/['decoder']/['layers']/['pre_self_attention_layer_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -293,6 +306,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -322,6 +336,7 @@ ".params/['params']/['decoder']/['layers']/['self_attention']/['key']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -336,6 +351,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -377,6 +393,7 @@ ".params/['params']/['decoder']/['layers']/['self_attention']/['key_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -391,6 +408,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -420,6 +438,7 @@ ".params/['params']/['decoder']/['layers']/['self_attention']/['out']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -434,6 +453,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -475,6 +495,7 @@ ".params/['params']/['decoder']/['layers']/['self_attention']/['query']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -489,6 +510,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -530,6 +552,7 @@ ".params/['params']/['decoder']/['layers']/['self_attention']/['query_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -544,6 +567,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -573,6 +597,7 @@ ".params/['params']/['decoder']/['layers']/['self_attention']/['value']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -587,6 +612,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -628,6 +654,7 @@ ".params/['params']/['token_embedder']/['embedding']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -642,6 +669,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -679,6 +707,7 @@ ".opt_state/[0]/.count": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -693,6 +722,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -713,6 +743,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['decoder_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -727,6 +758,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -754,6 +786,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['mlp']/['wi_0']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -768,6 +801,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -807,6 +841,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['mlp']/['wi_1']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -821,6 +856,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -860,6 +896,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['mlp']/['wo']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -874,6 +911,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -913,6 +951,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['post_self_attention_layer_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -927,6 +966,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -956,6 +996,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['pre_self_attention_layer_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -970,6 +1011,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -999,6 +1041,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['self_attention']/['key']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1013,6 +1056,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -1054,6 +1098,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['self_attention']/['key_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1068,6 +1113,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -1097,6 +1143,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['self_attention']/['out']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1111,6 +1158,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -1152,6 +1200,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['self_attention']/['query']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1166,6 +1215,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -1207,6 +1257,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['self_attention']/['query_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1221,6 +1272,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -1250,6 +1302,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['self_attention']/['value']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1264,6 +1317,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -1305,6 +1359,7 @@ ".opt_state/[0]/.mu/['params']/['token_embedder']/['embedding']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1319,6 +1374,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -1356,6 +1412,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['decoder_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1370,6 +1427,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -1397,6 +1455,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['mlp']/['wi_0']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1411,6 +1470,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -1450,6 +1510,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['mlp']/['wi_1']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1464,6 +1525,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -1503,6 +1565,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['mlp']/['wo']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1517,6 +1580,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -1556,6 +1620,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['post_self_attention_layer_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1570,6 +1635,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -1599,6 +1665,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['pre_self_attention_layer_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1613,6 +1680,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -1642,6 +1710,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['self_attention']/['key']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1656,6 +1725,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -1697,6 +1767,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['self_attention']/['key_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1711,6 +1782,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -1740,6 +1812,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['self_attention']/['out']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1754,6 +1827,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -1795,6 +1869,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['self_attention']/['query']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1809,6 +1884,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -1850,6 +1926,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['self_attention']/['query_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1864,6 +1941,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -1893,6 +1971,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['self_attention']/['value']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1907,6 +1986,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -1948,6 +2028,7 @@ ".opt_state/[0]/.nu/['params']/['token_embedder']/['embedding']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1962,6 +2043,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, @@ -1999,6 +2081,7 @@ ".opt_state/[2]/.count": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2013,6 +2096,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 1, "stage": 1, "fsdp": 16, diff --git a/tests/utils/sharding_info/qwen3-0.6b/v6e-16/slice_4/named_shardings.json b/tests/utils/sharding_info/qwen3-0.6b/v6e-16/slice_4/named_shardings.json index 8e13360273..31499e643e 100644 --- a/tests/utils/sharding_info/qwen3-0.6b/v6e-16/slice_4/named_shardings.json +++ b/tests/utils/sharding_info/qwen3-0.6b/v6e-16/slice_4/named_shardings.json @@ -2,6 +2,7 @@ ".step": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -16,6 +17,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -36,6 +38,7 @@ ".params/['params']/['decoder']/['decoder_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -50,6 +53,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -77,6 +81,7 @@ ".params/['params']/['decoder']/['layers']/['mlp']/['wi_0']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -91,6 +96,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -130,6 +136,7 @@ ".params/['params']/['decoder']/['layers']/['mlp']/['wi_1']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -144,6 +151,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -183,6 +191,7 @@ ".params/['params']/['decoder']/['layers']/['mlp']/['wo']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -197,6 +206,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -236,6 +246,7 @@ ".params/['params']/['decoder']/['layers']/['post_self_attention_layer_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -250,6 +261,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -279,6 +291,7 @@ ".params/['params']/['decoder']/['layers']/['pre_self_attention_layer_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -293,6 +306,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -322,6 +336,7 @@ ".params/['params']/['decoder']/['layers']/['self_attention']/['key']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -336,6 +351,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -377,6 +393,7 @@ ".params/['params']/['decoder']/['layers']/['self_attention']/['key_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -391,6 +408,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -420,6 +438,7 @@ ".params/['params']/['decoder']/['layers']/['self_attention']/['out']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -434,6 +453,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -475,6 +495,7 @@ ".params/['params']/['decoder']/['layers']/['self_attention']/['query']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -489,6 +510,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -530,6 +552,7 @@ ".params/['params']/['decoder']/['layers']/['self_attention']/['query_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -544,6 +567,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -573,6 +597,7 @@ ".params/['params']/['decoder']/['layers']/['self_attention']/['value']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -587,6 +612,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -628,6 +654,7 @@ ".params/['params']/['token_embedder']/['embedding']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -642,6 +669,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -679,6 +707,7 @@ ".opt_state/[0]/.count": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -693,6 +722,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -713,6 +743,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['decoder_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -727,6 +758,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -754,6 +786,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['mlp']/['wi_0']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -768,6 +801,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -807,6 +841,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['mlp']/['wi_1']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -821,6 +856,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -860,6 +896,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['mlp']/['wo']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -874,6 +911,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -913,6 +951,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['post_self_attention_layer_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -927,6 +966,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -956,6 +996,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['pre_self_attention_layer_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -970,6 +1011,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -999,6 +1041,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['self_attention']/['key']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1013,6 +1056,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -1054,6 +1098,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['self_attention']/['key_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1068,6 +1113,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -1097,6 +1143,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['self_attention']/['out']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1111,6 +1158,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -1152,6 +1200,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['self_attention']/['query']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1166,6 +1215,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -1207,6 +1257,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['self_attention']/['query_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1221,6 +1272,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -1250,6 +1302,7 @@ ".opt_state/[0]/.mu/['params']/['decoder']/['layers']/['self_attention']/['value']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1264,6 +1317,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -1305,6 +1359,7 @@ ".opt_state/[0]/.mu/['params']/['token_embedder']/['embedding']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1319,6 +1374,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -1356,6 +1412,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['decoder_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1370,6 +1427,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -1397,6 +1455,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['mlp']/['wi_0']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1411,6 +1470,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -1450,6 +1510,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['mlp']/['wi_1']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1464,6 +1525,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -1503,6 +1565,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['mlp']/['wo']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1517,6 +1580,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -1556,6 +1620,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['post_self_attention_layer_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1570,6 +1635,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -1599,6 +1665,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['pre_self_attention_layer_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1613,6 +1680,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -1642,6 +1710,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['self_attention']/['key']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1656,6 +1725,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -1697,6 +1767,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['self_attention']/['key_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1711,6 +1782,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -1740,6 +1812,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['self_attention']/['out']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1754,6 +1827,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -1795,6 +1869,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['self_attention']/['query']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1809,6 +1884,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -1850,6 +1926,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['self_attention']/['query_norm']/['scale']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1864,6 +1941,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -1893,6 +1971,7 @@ ".opt_state/[0]/.nu/['params']/['decoder']/['layers']/['self_attention']/['value']/['kernel']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1907,6 +1986,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -1948,6 +2028,7 @@ ".opt_state/[0]/.nu/['params']/['token_embedder']/['embedding']": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -1962,6 +2043,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16, @@ -1999,6 +2081,7 @@ ".opt_state/[2]/.count": { "mesh": { "axis_names": [ + "diloco", "data", "stage", "fsdp", @@ -2013,6 +2096,7 @@ "autoregressive" ], "shape": { + "diloco": 1, "data": 4, "stage": 1, "fsdp": 16,