diff --git a/families/eagle_vlm/parallel.py b/families/eagle_vlm/parallel.py index 77f29d0a80..d639386930 100644 --- a/families/eagle_vlm/parallel.py +++ b/families/eagle_vlm/parallel.py @@ -82,9 +82,9 @@ def shard_standard_decoder_weights( sharded[key] = value elif key.endswith((".w_q", ".w_k", ".w_v", ".q_bias", ".k_bias", ".v_bias")): sharded[key] = _slice_last_dim(value, rank, tp_size) - elif key.endswith((".w_o", ".w_fc2")): + elif key.endswith((".w_o", ".w_fc2", ".w_down")): sharded[key] = _slice_first_dim(value, rank, tp_size) - elif key.endswith((".w_fc1", ".fc1_bias")): + elif key.endswith((".w_fc1", ".fc1_bias", ".w_gate", ".w_up")): sharded[key] = _slice_last_dim(value, rank, tp_size) else: sharded[key] = value diff --git a/families/eagle_vlm/tests/test_parallel_swiglu.py b/families/eagle_vlm/tests/test_parallel_swiglu.py new file mode 100644 index 0000000000..6fd6e796cd --- /dev/null +++ b/families/eagle_vlm/tests/test_parallel_swiglu.py @@ -0,0 +1,82 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""eagle_vlm SwiGLU weight shards must reconstruct the unsharded computation.""" + +import numpy as np +import pytest + +from ..checkpoint_mapper import WeightDict +from ..config import ModelConfig +from ..parallel import ParallelConfig, shard_standard_decoder_weights + + +@pytest.mark.parametrize("tp_size", [1, 2, 4, 8]) +@pytest.mark.parametrize("dtype", [np.float32, np.float16]) +def test_swiglu_rank_shards_reconstruct_weights_and_output(tp_size, dtype): + config = ModelConfig( + hidden_size=8, + intermediate_size=24, + num_attention_heads=8, + num_key_value_heads=8, + ) + rng = np.random.default_rng(42) + weights = WeightDict( + { + "_attention_size": 8, + "_kv_attention_size": 8, + "_mlp_size": 24, + "embedding": rng.normal(size=(16, 8)).astype(dtype), + "layer.0.w_gate": rng.normal(scale=0.2, size=(8, 24)).astype(dtype), + "layer.0.w_up": rng.normal(scale=0.2, size=(8, 24)).astype(dtype), + "layer.0.w_down": rng.normal(scale=0.2, size=(24, 8)).astype(dtype), + "layer.0.post_attn_norm": np.ones(8, dtype=dtype), + } + ) + original = { + key: value.copy() for key, value in weights.items() if isinstance(value, np.ndarray) + } + shards = [ + shard_standard_decoder_weights(config, weights, ParallelConfig(tp_size, rank)) + for rank in range(tp_size) + ] + + for key, axis in (("w_gate", 1), ("w_up", 1), ("w_down", 0)): + full = weights[f"layer.0.{key}"] + expected_shape = list(full.shape) + expected_shape[axis] //= tp_size + for shard in shards: + part = shard[f"layer.0.{key}"] + assert part.shape == tuple(expected_shape) + assert part.dtype == dtype + assert part.flags.c_contiguous + np.testing.assert_array_equal( + np.concatenate([shard[f"layer.0.{key}"] for shard in shards], axis=axis), + full, + ) + + for rank, shard in enumerate(shards): + assert shard["_mlp_size"] == 24 // tp_size + assert isinstance(shard, WeightDict) + for key in ("embedding", "layer.0.post_attn_norm"): + np.testing.assert_array_equal(shard[key], weights[key]) + if tp_size > 1: + assert shard["_tensor_parallel_rank"] == rank + for key, value in original.items(): + np.testing.assert_array_equal(weights[key], value) + assert weights["_mlp_size"] == 24 + + inputs = rng.normal(size=(3, 8)) + + def evaluate(rank_weights): + gate = inputs @ rank_weights["layer.0.w_gate"].astype(np.float64) + up = inputs @ rank_weights["layer.0.w_up"].astype(np.float64) + return (gate / (1 + np.exp(-gate)) * up) @ rank_weights["layer.0.w_down"].astype(np.float64) + + # Row-parallel down projections are summed by the runtime ALL_REDUCE. + np.testing.assert_allclose( + sum(evaluate(shard) for shard in shards), + evaluate(weights), + rtol=1e-12, + atol=1e-12, + ) diff --git a/families/glm/parallel.py b/families/glm/parallel.py index 4d88550d49..9bfc5a21c3 100644 --- a/families/glm/parallel.py +++ b/families/glm/parallel.py @@ -82,9 +82,9 @@ def shard_standard_decoder_weights( sharded[key] = value elif key.endswith((".w_q", ".w_k", ".w_v", ".q_bias", ".k_bias", ".v_bias")): sharded[key] = _slice_last_dim(value, rank, tp_size) - elif key.endswith((".w_o", ".w_fc2")): + elif key.endswith((".w_o", ".w_fc2", ".w_down")): sharded[key] = _slice_first_dim(value, rank, tp_size) - elif key.endswith((".w_fc1", ".fc1_bias")): + elif key.endswith((".w_fc1", ".fc1_bias", ".w_gate", ".w_up")): sharded[key] = _slice_last_dim(value, rank, tp_size) else: sharded[key] = value diff --git a/families/glm/tests/test_parallel_swiglu.py b/families/glm/tests/test_parallel_swiglu.py new file mode 100644 index 0000000000..98c3a04a27 --- /dev/null +++ b/families/glm/tests/test_parallel_swiglu.py @@ -0,0 +1,82 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""glm SwiGLU weight shards must reconstruct the unsharded computation.""" + +import numpy as np +import pytest + +from ..checkpoint_mapper import WeightDict +from ..config import ModelConfig +from ..parallel import ParallelConfig, shard_standard_decoder_weights + + +@pytest.mark.parametrize("tp_size", [1, 2, 4, 8]) +@pytest.mark.parametrize("dtype", [np.float32, np.float16]) +def test_swiglu_rank_shards_reconstruct_weights_and_output(tp_size, dtype): + config = ModelConfig( + hidden_size=8, + intermediate_size=24, + num_attention_heads=8, + num_key_value_heads=8, + ) + rng = np.random.default_rng(42) + weights = WeightDict( + { + "_attention_size": 8, + "_kv_attention_size": 8, + "_mlp_size": 24, + "embedding": rng.normal(size=(16, 8)).astype(dtype), + "layer.0.w_gate": rng.normal(scale=0.2, size=(8, 24)).astype(dtype), + "layer.0.w_up": rng.normal(scale=0.2, size=(8, 24)).astype(dtype), + "layer.0.w_down": rng.normal(scale=0.2, size=(24, 8)).astype(dtype), + "layer.0.post_attn_norm": np.ones(8, dtype=dtype), + } + ) + original = { + key: value.copy() for key, value in weights.items() if isinstance(value, np.ndarray) + } + shards = [ + shard_standard_decoder_weights(config, weights, ParallelConfig(tp_size, rank)) + for rank in range(tp_size) + ] + + for key, axis in (("w_gate", 1), ("w_up", 1), ("w_down", 0)): + full = weights[f"layer.0.{key}"] + expected_shape = list(full.shape) + expected_shape[axis] //= tp_size + for shard in shards: + part = shard[f"layer.0.{key}"] + assert part.shape == tuple(expected_shape) + assert part.dtype == dtype + assert part.flags.c_contiguous + np.testing.assert_array_equal( + np.concatenate([shard[f"layer.0.{key}"] for shard in shards], axis=axis), + full, + ) + + for rank, shard in enumerate(shards): + assert shard["_mlp_size"] == 24 // tp_size + assert isinstance(shard, WeightDict) + for key in ("embedding", "layer.0.post_attn_norm"): + np.testing.assert_array_equal(shard[key], weights[key]) + if tp_size > 1: + assert shard["_tensor_parallel_rank"] == rank + for key, value in original.items(): + np.testing.assert_array_equal(weights[key], value) + assert weights["_mlp_size"] == 24 + + inputs = rng.normal(size=(3, 8)) + + def evaluate(rank_weights): + gate = inputs @ rank_weights["layer.0.w_gate"].astype(np.float64) + up = inputs @ rank_weights["layer.0.w_up"].astype(np.float64) + return (gate / (1 + np.exp(-gate)) * up) @ rank_weights["layer.0.w_down"].astype(np.float64) + + # Row-parallel down projections are summed by the runtime ALL_REDUCE. + np.testing.assert_allclose( + sum(evaluate(shard) for shard in shards), + evaluate(weights), + rtol=1e-12, + atol=1e-12, + ) diff --git a/families/granite/parallel.py b/families/granite/parallel.py index 998427daf8..d7dc612f10 100644 --- a/families/granite/parallel.py +++ b/families/granite/parallel.py @@ -82,9 +82,9 @@ def shard_standard_decoder_weights( sharded[key] = value elif key.endswith((".w_q", ".w_k", ".w_v", ".q_bias", ".k_bias", ".v_bias")): sharded[key] = _slice_last_dim(value, rank, tp_size) - elif key.endswith((".w_o", ".w_fc2")): + elif key.endswith((".w_o", ".w_fc2", ".w_down")): sharded[key] = _slice_first_dim(value, rank, tp_size) - elif key.endswith((".w_fc1", ".fc1_bias")): + elif key.endswith((".w_fc1", ".fc1_bias", ".w_gate", ".w_up")): sharded[key] = _slice_last_dim(value, rank, tp_size) else: sharded[key] = value diff --git a/families/granite/tests/test_parallel_swiglu.py b/families/granite/tests/test_parallel_swiglu.py new file mode 100644 index 0000000000..4cdf6a8b7f --- /dev/null +++ b/families/granite/tests/test_parallel_swiglu.py @@ -0,0 +1,82 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""granite SwiGLU weight shards must reconstruct the unsharded computation.""" + +import numpy as np +import pytest + +from ..checkpoint_mapper import WeightDict +from ..config import ModelConfig +from ..parallel import ParallelConfig, shard_standard_decoder_weights + + +@pytest.mark.parametrize("tp_size", [1, 2, 4, 8]) +@pytest.mark.parametrize("dtype", [np.float32, np.float16]) +def test_swiglu_rank_shards_reconstruct_weights_and_output(tp_size, dtype): + config = ModelConfig( + hidden_size=8, + intermediate_size=24, + num_attention_heads=8, + num_key_value_heads=8, + ) + rng = np.random.default_rng(42) + weights = WeightDict( + { + "_attention_size": 8, + "_kv_attention_size": 8, + "_mlp_size": 24, + "embedding": rng.normal(size=(16, 8)).astype(dtype), + "layer.0.w_gate": rng.normal(scale=0.2, size=(8, 24)).astype(dtype), + "layer.0.w_up": rng.normal(scale=0.2, size=(8, 24)).astype(dtype), + "layer.0.w_down": rng.normal(scale=0.2, size=(24, 8)).astype(dtype), + "layer.0.post_attn_norm": np.ones(8, dtype=dtype), + } + ) + original = { + key: value.copy() for key, value in weights.items() if isinstance(value, np.ndarray) + } + shards = [ + shard_standard_decoder_weights(config, weights, ParallelConfig(tp_size, rank)) + for rank in range(tp_size) + ] + + for key, axis in (("w_gate", 1), ("w_up", 1), ("w_down", 0)): + full = weights[f"layer.0.{key}"] + expected_shape = list(full.shape) + expected_shape[axis] //= tp_size + for shard in shards: + part = shard[f"layer.0.{key}"] + assert part.shape == tuple(expected_shape) + assert part.dtype == dtype + assert part.flags.c_contiguous + np.testing.assert_array_equal( + np.concatenate([shard[f"layer.0.{key}"] for shard in shards], axis=axis), + full, + ) + + for rank, shard in enumerate(shards): + assert shard["_mlp_size"] == 24 // tp_size + assert isinstance(shard, WeightDict) + for key in ("embedding", "layer.0.post_attn_norm"): + np.testing.assert_array_equal(shard[key], weights[key]) + if tp_size > 1: + assert shard["_tensor_parallel_rank"] == rank + for key, value in original.items(): + np.testing.assert_array_equal(weights[key], value) + assert weights["_mlp_size"] == 24 + + inputs = rng.normal(size=(3, 8)) + + def evaluate(rank_weights): + gate = inputs @ rank_weights["layer.0.w_gate"].astype(np.float64) + up = inputs @ rank_weights["layer.0.w_up"].astype(np.float64) + return (gate / (1 + np.exp(-gate)) * up) @ rank_weights["layer.0.w_down"].astype(np.float64) + + # Row-parallel down projections are summed by the runtime ALL_REDUCE. + np.testing.assert_allclose( + sum(evaluate(shard) for shard in shards), + evaluate(weights), + rtol=1e-12, + atol=1e-12, + ) diff --git a/families/olmo/parallel.py b/families/olmo/parallel.py index f4de318f93..9484cf4c8d 100644 --- a/families/olmo/parallel.py +++ b/families/olmo/parallel.py @@ -82,9 +82,9 @@ def shard_standard_decoder_weights( sharded[key] = value elif key.endswith((".w_q", ".w_k", ".w_v", ".q_bias", ".k_bias", ".v_bias")): sharded[key] = _slice_last_dim(value, rank, tp_size) - elif key.endswith((".w_o", ".w_fc2")): + elif key.endswith((".w_o", ".w_fc2", ".w_down")): sharded[key] = _slice_first_dim(value, rank, tp_size) - elif key.endswith((".w_fc1", ".fc1_bias")): + elif key.endswith((".w_fc1", ".fc1_bias", ".w_gate", ".w_up")): sharded[key] = _slice_last_dim(value, rank, tp_size) else: sharded[key] = value diff --git a/families/olmo/tests/test_parallel_swiglu.py b/families/olmo/tests/test_parallel_swiglu.py new file mode 100644 index 0000000000..ee77a49adb --- /dev/null +++ b/families/olmo/tests/test_parallel_swiglu.py @@ -0,0 +1,82 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""olmo SwiGLU weight shards must reconstruct the unsharded computation.""" + +import numpy as np +import pytest + +from ..checkpoint_mapper import WeightDict +from ..config import ModelConfig +from ..parallel import ParallelConfig, shard_standard_decoder_weights + + +@pytest.mark.parametrize("tp_size", [1, 2, 4, 8]) +@pytest.mark.parametrize("dtype", [np.float32, np.float16]) +def test_swiglu_rank_shards_reconstruct_weights_and_output(tp_size, dtype): + config = ModelConfig( + hidden_size=8, + intermediate_size=24, + num_attention_heads=8, + num_key_value_heads=8, + ) + rng = np.random.default_rng(42) + weights = WeightDict( + { + "_attention_size": 8, + "_kv_attention_size": 8, + "_mlp_size": 24, + "embedding": rng.normal(size=(16, 8)).astype(dtype), + "layer.0.w_gate": rng.normal(scale=0.2, size=(8, 24)).astype(dtype), + "layer.0.w_up": rng.normal(scale=0.2, size=(8, 24)).astype(dtype), + "layer.0.w_down": rng.normal(scale=0.2, size=(24, 8)).astype(dtype), + "layer.0.post_attn_norm": np.ones(8, dtype=dtype), + } + ) + original = { + key: value.copy() for key, value in weights.items() if isinstance(value, np.ndarray) + } + shards = [ + shard_standard_decoder_weights(config, weights, ParallelConfig(tp_size, rank)) + for rank in range(tp_size) + ] + + for key, axis in (("w_gate", 1), ("w_up", 1), ("w_down", 0)): + full = weights[f"layer.0.{key}"] + expected_shape = list(full.shape) + expected_shape[axis] //= tp_size + for shard in shards: + part = shard[f"layer.0.{key}"] + assert part.shape == tuple(expected_shape) + assert part.dtype == dtype + assert part.flags.c_contiguous + np.testing.assert_array_equal( + np.concatenate([shard[f"layer.0.{key}"] for shard in shards], axis=axis), + full, + ) + + for rank, shard in enumerate(shards): + assert shard["_mlp_size"] == 24 // tp_size + assert isinstance(shard, WeightDict) + for key in ("embedding", "layer.0.post_attn_norm"): + np.testing.assert_array_equal(shard[key], weights[key]) + if tp_size > 1: + assert shard["_tensor_parallel_rank"] == rank + for key, value in original.items(): + np.testing.assert_array_equal(weights[key], value) + assert weights["_mlp_size"] == 24 + + inputs = rng.normal(size=(3, 8)) + + def evaluate(rank_weights): + gate = inputs @ rank_weights["layer.0.w_gate"].astype(np.float64) + up = inputs @ rank_weights["layer.0.w_up"].astype(np.float64) + return (gate / (1 + np.exp(-gate)) * up) @ rank_weights["layer.0.w_down"].astype(np.float64) + + # Row-parallel down projections are summed by the runtime ALL_REDUCE. + np.testing.assert_allclose( + sum(evaluate(shard) for shard in shards), + evaluate(weights), + rtol=1e-12, + atol=1e-12, + ) diff --git a/families/qwen_vl/parallel.py b/families/qwen_vl/parallel.py index 85bff94b95..8e6f6a0d31 100644 --- a/families/qwen_vl/parallel.py +++ b/families/qwen_vl/parallel.py @@ -82,9 +82,9 @@ def shard_standard_decoder_weights( sharded[key] = value elif key.endswith((".w_q", ".w_k", ".w_v", ".q_bias", ".k_bias", ".v_bias")): sharded[key] = _slice_last_dim(value, rank, tp_size) - elif key.endswith((".w_o", ".w_fc2")): + elif key.endswith((".w_o", ".w_fc2", ".w_down")): sharded[key] = _slice_first_dim(value, rank, tp_size) - elif key.endswith((".w_fc1", ".fc1_bias")): + elif key.endswith((".w_fc1", ".fc1_bias", ".w_gate", ".w_up")): sharded[key] = _slice_last_dim(value, rank, tp_size) else: sharded[key] = value diff --git a/families/qwen_vl/tests/test_parallel_swiglu.py b/families/qwen_vl/tests/test_parallel_swiglu.py new file mode 100644 index 0000000000..b20d7ec30b --- /dev/null +++ b/families/qwen_vl/tests/test_parallel_swiglu.py @@ -0,0 +1,82 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""qwen_vl SwiGLU weight shards must reconstruct the unsharded computation.""" + +import numpy as np +import pytest + +from ..checkpoint_mapper import WeightDict +from ..config import ModelConfig +from ..parallel import ParallelConfig, shard_standard_decoder_weights + + +@pytest.mark.parametrize("tp_size", [1, 2, 4, 8]) +@pytest.mark.parametrize("dtype", [np.float32, np.float16]) +def test_swiglu_rank_shards_reconstruct_weights_and_output(tp_size, dtype): + config = ModelConfig( + hidden_size=8, + intermediate_size=24, + num_attention_heads=8, + num_key_value_heads=8, + ) + rng = np.random.default_rng(42) + weights = WeightDict( + { + "_attention_size": 8, + "_kv_attention_size": 8, + "_mlp_size": 24, + "embedding": rng.normal(size=(16, 8)).astype(dtype), + "layer.0.w_gate": rng.normal(scale=0.2, size=(8, 24)).astype(dtype), + "layer.0.w_up": rng.normal(scale=0.2, size=(8, 24)).astype(dtype), + "layer.0.w_down": rng.normal(scale=0.2, size=(24, 8)).astype(dtype), + "layer.0.post_attn_norm": np.ones(8, dtype=dtype), + } + ) + original = { + key: value.copy() for key, value in weights.items() if isinstance(value, np.ndarray) + } + shards = [ + shard_standard_decoder_weights(config, weights, ParallelConfig(tp_size, rank)) + for rank in range(tp_size) + ] + + for key, axis in (("w_gate", 1), ("w_up", 1), ("w_down", 0)): + full = weights[f"layer.0.{key}"] + expected_shape = list(full.shape) + expected_shape[axis] //= tp_size + for shard in shards: + part = shard[f"layer.0.{key}"] + assert part.shape == tuple(expected_shape) + assert part.dtype == dtype + assert part.flags.c_contiguous + np.testing.assert_array_equal( + np.concatenate([shard[f"layer.0.{key}"] for shard in shards], axis=axis), + full, + ) + + for rank, shard in enumerate(shards): + assert shard["_mlp_size"] == 24 // tp_size + assert isinstance(shard, WeightDict) + for key in ("embedding", "layer.0.post_attn_norm"): + np.testing.assert_array_equal(shard[key], weights[key]) + if tp_size > 1: + assert shard["_tensor_parallel_rank"] == rank + for key, value in original.items(): + np.testing.assert_array_equal(weights[key], value) + assert weights["_mlp_size"] == 24 + + inputs = rng.normal(size=(3, 8)) + + def evaluate(rank_weights): + gate = inputs @ rank_weights["layer.0.w_gate"].astype(np.float64) + up = inputs @ rank_weights["layer.0.w_up"].astype(np.float64) + return (gate / (1 + np.exp(-gate)) * up) @ rank_weights["layer.0.w_down"].astype(np.float64) + + # Row-parallel down projections are summed by the runtime ALL_REDUCE. + np.testing.assert_allclose( + sum(evaluate(shard) for shard in shards), + evaluate(weights), + rtol=1e-12, + atol=1e-12, + ) diff --git a/families/stablelm/parallel.py b/families/stablelm/parallel.py index 160f27582a..a5fd2e60f9 100644 --- a/families/stablelm/parallel.py +++ b/families/stablelm/parallel.py @@ -82,9 +82,9 @@ def shard_standard_decoder_weights( sharded[key] = value elif key.endswith((".w_q", ".w_k", ".w_v", ".q_bias", ".k_bias", ".v_bias")): sharded[key] = _slice_last_dim(value, rank, tp_size) - elif key.endswith((".w_o", ".w_fc2")): + elif key.endswith((".w_o", ".w_fc2", ".w_down")): sharded[key] = _slice_first_dim(value, rank, tp_size) - elif key.endswith((".w_fc1", ".fc1_bias")): + elif key.endswith((".w_fc1", ".fc1_bias", ".w_gate", ".w_up")): sharded[key] = _slice_last_dim(value, rank, tp_size) else: sharded[key] = value diff --git a/families/stablelm/tests/test_parallel_swiglu.py b/families/stablelm/tests/test_parallel_swiglu.py new file mode 100644 index 0000000000..e60483e0b3 --- /dev/null +++ b/families/stablelm/tests/test_parallel_swiglu.py @@ -0,0 +1,82 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""stablelm SwiGLU weight shards must reconstruct the unsharded computation.""" + +import numpy as np +import pytest + +from ..checkpoint_mapper import WeightDict +from ..config import ModelConfig +from ..parallel import ParallelConfig, shard_standard_decoder_weights + + +@pytest.mark.parametrize("tp_size", [1, 2, 4, 8]) +@pytest.mark.parametrize("dtype", [np.float32, np.float16]) +def test_swiglu_rank_shards_reconstruct_weights_and_output(tp_size, dtype): + config = ModelConfig( + hidden_size=8, + intermediate_size=24, + num_attention_heads=8, + num_key_value_heads=8, + ) + rng = np.random.default_rng(42) + weights = WeightDict( + { + "_attention_size": 8, + "_kv_attention_size": 8, + "_mlp_size": 24, + "embedding": rng.normal(size=(16, 8)).astype(dtype), + "layer.0.w_gate": rng.normal(scale=0.2, size=(8, 24)).astype(dtype), + "layer.0.w_up": rng.normal(scale=0.2, size=(8, 24)).astype(dtype), + "layer.0.w_down": rng.normal(scale=0.2, size=(24, 8)).astype(dtype), + "layer.0.post_attn_norm": np.ones(8, dtype=dtype), + } + ) + original = { + key: value.copy() for key, value in weights.items() if isinstance(value, np.ndarray) + } + shards = [ + shard_standard_decoder_weights(config, weights, ParallelConfig(tp_size, rank)) + for rank in range(tp_size) + ] + + for key, axis in (("w_gate", 1), ("w_up", 1), ("w_down", 0)): + full = weights[f"layer.0.{key}"] + expected_shape = list(full.shape) + expected_shape[axis] //= tp_size + for shard in shards: + part = shard[f"layer.0.{key}"] + assert part.shape == tuple(expected_shape) + assert part.dtype == dtype + assert part.flags.c_contiguous + np.testing.assert_array_equal( + np.concatenate([shard[f"layer.0.{key}"] for shard in shards], axis=axis), + full, + ) + + for rank, shard in enumerate(shards): + assert shard["_mlp_size"] == 24 // tp_size + assert isinstance(shard, WeightDict) + for key in ("embedding", "layer.0.post_attn_norm"): + np.testing.assert_array_equal(shard[key], weights[key]) + if tp_size > 1: + assert shard["_tensor_parallel_rank"] == rank + for key, value in original.items(): + np.testing.assert_array_equal(weights[key], value) + assert weights["_mlp_size"] == 24 + + inputs = rng.normal(size=(3, 8)) + + def evaluate(rank_weights): + gate = inputs @ rank_weights["layer.0.w_gate"].astype(np.float64) + up = inputs @ rank_weights["layer.0.w_up"].astype(np.float64) + return (gate / (1 + np.exp(-gate)) * up) @ rank_weights["layer.0.w_down"].astype(np.float64) + + # Row-parallel down projections are summed by the runtime ALL_REDUCE. + np.testing.assert_allclose( + sum(evaluate(shard) for shard in shards), + evaluate(weights), + rtol=1e-12, + atol=1e-12, + )