perf(ckpt): build remaining split indices on the input device (#6347)

Signed-off-by: Eugenio "Jay" Zuccarelli <11176606+jayzuccarelli@users.noreply.github.com>

Author: jayzuccarelli

src/megatron/bridge/models/conversion/param_mapping.py

@@ -2575,11 +2575,13 @@ def hf_to_megatron(
             d_tot_ssm = config.mamba_state_dim * config.mamba_num_groups
 
             # Define component indices in the concatenated tensor
-            z_shard_idx = torch.arange(d_inner)
-            x_shard_idx = torch.arange(d_inner, 2 * d_inner)
-            B_shard_idx = torch.arange(2 * d_inner, 2 * d_inner + d_tot_ssm)
-            C_shard_idx = torch.arange(2 * d_inner + d_tot_ssm, 2 * d_inner + 2 * d_tot_ssm)
-            dt_shard_idx = torch.arange(2 * (d_inner + d_tot_ssm), 2 * (d_inner + d_tot_ssm) + config.mamba_num_heads)
+            z_shard_idx = torch.arange(d_inner, device=hf_weights.device)
+            x_shard_idx = torch.arange(d_inner, 2 * d_inner, device=hf_weights.device)
+            B_shard_idx = torch.arange(2 * d_inner, 2 * d_inner + d_tot_ssm, device=hf_weights.device)
+            C_shard_idx = torch.arange(2 * d_inner + d_tot_ssm, 2 * d_inner + 2 * d_tot_ssm, device=hf_weights.device)
+            dt_shard_idx = torch.arange(
+                2 * (d_inner + d_tot_ssm), 2 * (d_inner + d_tot_ssm) + config.mamba_num_heads, device=hf_weights.device
+            )
 
             # Reshape for tensor parallel distribution
             target_shape = (self.tp_size, -1, config.hidden_size)
@@ -2625,11 +2627,15 @@ def megatron_to_hf(
         n_heads_local = config.mamba_num_heads // self.tp_size
 
         # Extract local components
-        z_shard_idx = torch.arange(d_inner_local)
-        x_shard_idx = torch.arange(d_inner_local) + d_inner_local
-        B_shard_idx = torch.arange(d_tot_ssm_local) + 2 * d_inner_local
-        C_shard_idx = torch.arange(d_tot_ssm_local) + 2 * d_inner_local + d_tot_ssm_local
-        dt_shard_idx = torch.arange(n_heads_local) + 2 * (d_inner_local + d_tot_ssm_local)
+        z_shard_idx = torch.arange(d_inner_local, device=megatron_weights.device)
+        x_shard_idx = torch.arange(d_inner_local, device=megatron_weights.device) + d_inner_local
+        B_shard_idx = torch.arange(d_tot_ssm_local, device=megatron_weights.device) + 2 * d_inner_local
+        C_shard_idx = (
+            torch.arange(d_tot_ssm_local, device=megatron_weights.device) + 2 * d_inner_local + d_tot_ssm_local
+        )
+        dt_shard_idx = torch.arange(n_heads_local, device=megatron_weights.device) + 2 * (
+            d_inner_local + d_tot_ssm_local
+        )
 
         local_components = [
             megatron_weights[z_shard_idx],
@@ -2668,7 +2674,9 @@ def __init__(self, megatron_param: str, hf_param: str):
         self._tp_mapping = ColumnParallelMapping(megatron_param, megatron_param)
 
     @abstractmethod
-    def get_shard_idx(self, config: TransformerConfig, local_tp: bool) -> List[int]:
+    def get_shard_idx(
+        self, config: TransformerConfig, local_tp: bool, device: torch.device | None = None
+    ) -> List[int]:
         """Get shard indices for the given config."""
         ...
 
@@ -2688,7 +2696,7 @@ def hf_to_megatron(
                 assert "bias" in self.megatron_param, "Only bias and weight are supported for conv1d"
                 target_shape = (self.tp_size, -1)
 
-            shard_idx = self.get_shard_idx(config, local_tp=False)
+            shard_idx = self.get_shard_idx(config, local_tp=False, device=hf_weights.device)
 
             # Extract and reshape components
             sharded_weights = [hf_weights[idx].reshape(target_shape) for idx in shard_idx]
@@ -2723,7 +2731,7 @@ def megatron_to_hf(
             config = remove_non_pickleables(config, max_depth=3)
             config = self.broadcast_obj_from_pp_rank(config, cache_key="config")
 
-        shard_idx = self.get_shard_idx(config, local_tp=True)
+        shard_idx = self.get_shard_idx(config, local_tp=True, device=megatron_weights.device)
 
         local_components = [megatron_weights[idx] for idx in shard_idx]
 
@@ -2747,17 +2755,19 @@ class GDNConv1dMapping(ChunkedMapping):
     tensor-parallel distributed format for GDN SSM layers.
     """
 
-    def get_shard_idx(self, config: TransformerConfig, local_tp: bool) -> List[int]:
+    def get_shard_idx(
+        self, config: TransformerConfig, local_tp: bool, device: torch.device | None = None
+    ) -> List[int]:
         """Get shard indices for the given config."""
         qk_dim = config.linear_key_head_dim * config.linear_num_key_heads
         v_dim = config.linear_value_head_dim * config.linear_num_value_heads
         if local_tp:
             qk_dim = qk_dim // self.tp_size
             v_dim = v_dim // self.tp_size
 
-        q_shard_idx = torch.arange(qk_dim)
-        k_shard_idx = torch.arange(qk_dim) + qk_dim
-        v_shard_idx = torch.arange(v_dim) + qk_dim * 2
+        q_shard_idx = torch.arange(qk_dim, device=device)
+        k_shard_idx = torch.arange(qk_dim, device=device) + qk_dim
+        v_shard_idx = torch.arange(v_dim, device=device) + qk_dim * 2
 
         return [q_shard_idx, k_shard_idx, v_shard_idx]
 
@@ -2769,7 +2779,9 @@ class MambaConv1dMapping(ChunkedMapping):
     tensor-parallel distributed format for Mamba SSM layers.
     """
 
-    def get_shard_idx(self, config: TransformerConfig, local_tp: bool) -> List[int]:
+    def get_shard_idx(
+        self, config: TransformerConfig, local_tp: bool, device: torch.device | None = None
+    ) -> List[int]:
         """Get shard indices for the given config."""
         d_inner = config.mamba_num_heads * config.mamba_head_dim
         d_tot_ssm = config.mamba_state_dim * config.mamba_num_groups
@@ -2778,9 +2790,9 @@ def get_shard_idx(self, config: TransformerConfig, local_tp: bool) -> List[int]:
             d_tot_ssm = d_tot_ssm // self.tp_size
 
         # Extract local components
-        x_shard_idx = torch.arange(d_inner)
-        B_shard_idx = torch.arange(d_tot_ssm) + d_inner
-        C_shard_idx = torch.arange(d_tot_ssm) + d_inner + d_tot_ssm
+        x_shard_idx = torch.arange(d_inner, device=device)
+        B_shard_idx = torch.arange(d_tot_ssm, device=device) + d_inner
+        C_shard_idx = torch.arange(d_tot_ssm, device=device) + d_inner + d_tot_ssm
 
         return [x_shard_idx, B_shard_idx, C_shard_idx]
 

tests/unit_tests/models/test_param_mapping.py

@@ -29,7 +29,10 @@
     FusedExpertMapping,
     FusedGatedExpertMapping,
     GatedMLPMapping,
+    GDNConv1dMapping,
     KVMapping,
+    MambaConv1dMapping,
+    MambaInProjMapping,
     MegatronParamMapping,
     QKVMapping,
     ReplicatedMapping,
@@ -802,6 +805,45 @@ def test_megatron_to_hf(self, mock_distributed_env, transformer_config):
         assert torch.equal(result["v.weight"], v)
 
 
+class TestSSMMappings:
+    config = SimpleNamespace(
+        hidden_size=8,
+        mamba_num_heads=4,
+        mamba_head_dim=2,
+        mamba_state_dim=2,
+        mamba_num_groups=1,
+        linear_key_head_dim=2,
+        linear_num_key_heads=2,
+        linear_value_head_dim=2,
+        linear_num_value_heads=2,
+    )
+
+    @pytest.mark.parametrize(
+        "mapping, hf_weight",
+        [
+            (MambaInProjMapping("mixer.in_proj.weight", "hf.in_proj.weight"), torch.randn(24, 8)),
+            (MambaConv1dMapping("mixer.conv1d.weight", "hf.conv1d.weight"), torch.randn(12, 1, 4)),
+            (MambaConv1dMapping("mixer.conv1d.bias", "hf.conv1d.bias"), torch.randn(12)),
+            (GDNConv1dMapping("mixer.conv1d.weight", "hf.conv1d.weight"), torch.randn(12, 1, 4)),
+        ],
+    )
+    def test_round_trip_builds_shard_idx_on_input_device(self, mock_distributed_env, mapping, hf_weight):
+        mock_distributed_env()
+        megatron_module = MockModule(self.config, weight_shape=hf_weight.shape)
+
+        with (
+            patch.object(mapping._tp_mapping, "hf_to_megatron", side_effect=lambda weight, module: weight),
+            patch("torch.arange", wraps=torch.arange) as mock_arange,
+        ):
+            megatron_weight = mapping.hf_to_megatron(hf_weight, megatron_module)
+            result = mapping.megatron_to_hf(megatron_weight, megatron_module)
+
+        assert torch.equal(result[mapping.hf_param], hf_weight)
+        assert mock_arange.call_count > 0
+        for call in mock_arange.call_args_list:
+            assert call.kwargs.get("device") == hf_weight.device
+
+
 class TestGatedMLPMapping:
     def test_hf_to_megatron_single_tp(self, mock_distributed_env, transformer_config):
         """Test gate+up merging with single TP rank."""