perf(ckpt): cache config broadcasts for Mamba and GDN mappings

Signed-off-by: Begunner <went@bytedance.com>

Author: Begunner

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

@@ -2094,12 +2094,12 @@ def megatron_to_hf(
 
         # Broadcast config to all PP ranks for collective communication
         if megatron_module is None:
-            config = self.broadcast_obj_from_pp_rank(None)
+            config = self.broadcast_obj_from_pp_rank(None, cache_key="config")
         else:
             config = self._get_config(megatron_module)
             # create shallow copy and remove non-picklable objects with max depth=3
             config = remove_non_pickleables(config, max_depth=3)
-            config = self.broadcast_obj_from_pp_rank(config)
+            config = self.broadcast_obj_from_pp_rank(config, cache_key="config")
 
         d_inner_local = (config.mamba_num_heads * config.mamba_head_dim) // self.tp_size
         d_tot_ssm_local = (config.mamba_state_dim * config.mamba_num_groups) // self.tp_size
@@ -2197,12 +2197,12 @@ def megatron_to_hf(
 
         # Broadcast config to all PP ranks for collective communication
         if megatron_module is None:
-            config = self.broadcast_obj_from_pp_rank(None)
+            config = self.broadcast_obj_from_pp_rank(None, cache_key="config")
         else:
             config = self._get_config(megatron_module)
             # create shallow copy and remove non-picklable objects with max depth=3
             config = remove_non_pickleables(config, max_depth=3)
-            config = self.broadcast_obj_from_pp_rank(config)
+            config = self.broadcast_obj_from_pp_rank(config, cache_key="config")
 
         shard_idx = self.get_shard_idx(config, local_tp=True)
 
@@ -2316,12 +2316,12 @@ def megatron_to_hf(
         # collective communication.
         # ------------------------------------------------------------------
         if megatron_module is None:
-            config = self.broadcast_obj_from_pp_rank(None)
+            config = self.broadcast_obj_from_pp_rank(None, cache_key="config")
         else:
             config = self._get_config(megatron_module)
             # create shallow copy and remove non-picklable objects with max depth=3
             config = remove_non_pickleables(config, max_depth=3)
-            config = self.broadcast_obj_from_pp_rank(config)
+            config = self.broadcast_obj_from_pp_rank(config, cache_key="config")
 
         # Delegate TP/PP gathering.
         packed_dict = self._tp_mapping.megatron_to_hf(megatron_weights, megatron_module)
@@ -2416,11 +2416,11 @@ def megatron_to_hf(
 
         # Broadcast config across PP ranks (mirrors GDNLinearMapping).
         if megatron_module is None:
-            config = self.broadcast_obj_from_pp_rank(None)
+            config = self.broadcast_obj_from_pp_rank(None, cache_key="config")
         else:
             config = self._get_config(megatron_module)
             config = remove_non_pickleables(config, max_depth=3)
-            config = self.broadcast_obj_from_pp_rank(config)
+            config = self.broadcast_obj_from_pp_rank(config, cache_key="config")
 
         packed_dict = self._tp_mapping.megatron_to_hf(megatron_weights, megatron_module)
         if not packed_dict: