From 54978f19c062fe745c254cc58c0d47ccc33732c6 Mon Sep 17 00:00:00 2001 From: chaofengw Date: Thu, 10 Sep 2026 05:05:29 +0000 Subject: [PATCH 1/6] fix(eagle_vlm): shard SwiGLU weights for TP Slice gate and up projections along their output dimension and the down projection along its input dimension so rank-local weights match the reduced MLP width. Add family-owned numerical regression coverage for TP 1, 2, 4, and 8 without changing E2E passing criteria. Signed-off-by: chaofengw --- families/eagle_vlm/parallel.py | 4 +- .../eagle_vlm/tests/test_parallel_swiglu.py | 82 +++++++++++++++++++ 2 files changed, 84 insertions(+), 2 deletions(-) create mode 100644 families/eagle_vlm/tests/test_parallel_swiglu.py 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, + ) From 1210731a82bd83a3b18e3224be1b8a65d95c61e3 Mon Sep 17 00:00:00 2001 From: chaofengw Date: Thu, 10 Sep 2026 05:05:29 +0000 Subject: [PATCH 2/6] fix(glm): shard SwiGLU weights for TP Slice gate and up projections along their output dimension and the down projection along its input dimension so rank-local weights match the reduced MLP width. Add family-owned numerical regression coverage for TP 1, 2, 4, and 8 without changing E2E passing criteria. Signed-off-by: chaofengw --- families/glm/parallel.py | 4 +- families/glm/tests/test_parallel_swiglu.py | 82 ++++++++++++++++++++++ 2 files changed, 84 insertions(+), 2 deletions(-) create mode 100644 families/glm/tests/test_parallel_swiglu.py 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, + ) From 04051e6ec5f7dff7811e7fa71c3bbdfe59582f2e Mon Sep 17 00:00:00 2001 From: chaofengw Date: Thu, 10 Sep 2026 05:05:30 +0000 Subject: [PATCH 3/6] fix(granite): shard SwiGLU weights for TP Slice gate and up projections along their output dimension and the down projection along its input dimension so rank-local weights match the reduced MLP width. Add family-owned numerical regression coverage for TP 1, 2, 4, and 8 without changing E2E passing criteria. Signed-off-by: chaofengw --- families/granite/parallel.py | 4 +- .../granite/tests/test_parallel_swiglu.py | 82 +++++++++++++++++++ 2 files changed, 84 insertions(+), 2 deletions(-) create mode 100644 families/granite/tests/test_parallel_swiglu.py 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, + ) From a886cf83e8448a20b55e3e9b998f657c7702eb4f Mon Sep 17 00:00:00 2001 From: chaofengw Date: Thu, 10 Sep 2026 05:05:30 +0000 Subject: [PATCH 4/6] fix(olmo): shard SwiGLU weights for TP Slice gate and up projections along their output dimension and the down projection along its input dimension so rank-local weights match the reduced MLP width. Add family-owned numerical regression coverage for TP 1, 2, 4, and 8 without changing E2E passing criteria. Signed-off-by: chaofengw --- families/olmo/parallel.py | 4 +- families/olmo/tests/test_parallel_swiglu.py | 82 +++++++++++++++++++++ 2 files changed, 84 insertions(+), 2 deletions(-) create mode 100644 families/olmo/tests/test_parallel_swiglu.py 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, + ) From 1ded5d9a5a7803bf97360341d38376680a64dd8b Mon Sep 17 00:00:00 2001 From: chaofengw Date: Thu, 10 Sep 2026 05:05:30 +0000 Subject: [PATCH 5/6] fix(qwen_vl): shard SwiGLU weights for TP Slice gate and up projections along their output dimension and the down projection along its input dimension so rank-local weights match the reduced MLP width. Add family-owned numerical regression coverage for TP 1, 2, 4, and 8 without changing E2E passing criteria. Signed-off-by: chaofengw --- families/qwen_vl/parallel.py | 4 +- .../qwen_vl/tests/test_parallel_swiglu.py | 82 +++++++++++++++++++ 2 files changed, 84 insertions(+), 2 deletions(-) create mode 100644 families/qwen_vl/tests/test_parallel_swiglu.py 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, + ) From cbbe5c24324af75038e8090b036865696fa4221e Mon Sep 17 00:00:00 2001 From: chaofengw Date: Thu, 10 Sep 2026 05:05:30 +0000 Subject: [PATCH 6/6] fix(stablelm): shard SwiGLU weights for TP Slice gate and up projections along their output dimension and the down projection along its input dimension so rank-local weights match the reduced MLP width. Add family-owned numerical regression coverage for TP 1, 2, 4, and 8 without changing E2E passing criteria. Signed-off-by: chaofengw --- families/stablelm/parallel.py | 4 +- .../stablelm/tests/test_parallel_swiglu.py | 82 +++++++++++++++++++ 2 files changed, 84 insertions(+), 2 deletions(-) create mode 100644 families/stablelm/tests/test_parallel_swiglu.py 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, + )