NVIDIA-NeMo · guyueh1 · Jun 11, 2025 · Jun 11, 2025 · Jun 12, 2025 · Jun 12, 2025
diff --git a/nemo_rl/algorithms/grpo.py b/nemo_rl/algorithms/grpo.py
@@ -290,38 +290,44 @@ def refit_policy_generation(
     policy: ColocatablePolicyInterface,
     policy_generation: GenerationInterface,
     refit_buffer_size_gb: int,  # GB
+    timer: Timer,
 ) -> None:
     """Refit the policy generation interface with the latest policy weights."""
+
     policy.offload_before_refit()
-    policy_generation.prepare_for_generation(tags=["weights"])
-    # Streaming update weights to save memory
-    state_dict_info: list[tuple[str, int]] = policy.prepare_weights_for_ipc()
-    # group keys to save time
-    available_bytes = refit_buffer_size_gb * (1024**3)
-    split_keys: list[list[str]] = []
-    keys: list[str] = []
-    for key, size_in_bytes in state_dict_info:
-        if size_in_bytes > available_bytes:
-            if keys:
-                split_keys.append(keys)
-                keys = []
-            available_bytes = refit_buffer_size_gb * (1024**3)
-
-        keys.append(key)
-        available_bytes -= size_in_bytes
-
-    if len(keys) > 0:
-        split_keys.append(keys)
-    # do update
-    for keys in split_keys:
-        ipc_handles = policy.get_weights_ipc_handles(keys)
-        if not policy_generation.update_weights(ipc_handles):
-            error_message = (
-                "❌ Error: Updating weights for the generation policy failed during refit.\n"
-                "This often indicates an issue with cuda-ipc or "
-                "a problem within the generation backend (e.g., vLLM worker).\n"
-            )
-            raise RuntimeError(error_message)
+
+    with timer.time("prepare_for_generation/reshard"):
+        policy_generation.prepare_for_generation(tags=["weights"])
+        # Streaming update weights to save memory
+        state_dict_info: list[tuple[str, int]] = policy.prepare_weights_for_ipc()
+        # group keys to save time
+        available_bytes = refit_buffer_size_gb * (1024**3)
+        split_keys: list[list[str]] = []
+        keys: list[str] = []
+        for key, size_in_bytes in state_dict_info:
+            if size_in_bytes > available_bytes:
+                if keys:
+                    split_keys.append(keys)
+                    keys = []
+                available_bytes = refit_buffer_size_gb * (1024**3)
+
+            keys.append(key)
+            available_bytes -= size_in_bytes
+
+        if len(keys) > 0:
+            split_keys.append(keys)
+        # do update
+        for keys in split_keys:
+            with timer.time("prepare_for_generation/reshard/get_ipc_handles"):
+                ipc_handles = policy.get_weights_ipc_handles(keys)
+            with timer.time("prepare_for_generation/reshard/update_weights"):
+                if not policy_generation.update_weights(ipc_handles):
+                    error_message = (
+                        "❌ Error: Updating weights for the generation policy failed during refit.\n"
+                        "This often indicates an issue with cuda-ipc or "
+                        "a problem within the generation backend (e.g., vLLM worker).\n"
+                    )
+                    raise RuntimeError(error_message)
     policy.offload_after_refit()
     policy_generation.prepare_for_generation(tags=["kv_cache"])
 
@@ -366,7 +372,7 @@ def grpo_train(
     if val_at_start and step == 0:
         print("\n🔍 Running initial validation...")
         if NEED_REFIT and POLICY_GENERATION_STALE:
-            refit_policy_generation(policy, policy_generation, refit_buffer_size_gb)
+            refit_policy_generation(policy, policy_generation, refit_buffer_size_gb, timer=timer)
             POLICY_GENERATION_STALE = False
         else:
             policy_generation.prepare_for_generation()
@@ -413,6 +419,7 @@ def grpo_train(
                         policy,
                         policy_generation,
                         refit_buffer_size_gb,
+                        timer=timer,
                     )
                     POLICY_GENERATION_STALE = False
                 else:
@@ -529,6 +536,7 @@ def grpo_train(
                         policy,
                         policy_generation,
                         refit_buffer_size_gb,
+                        timer=timer,
                     )
                     POLICY_GENERATION_STALE = False
                 else:

diff --git a/nemo_rl/distributed/batched_data_dict.py b/nemo_rl/distributed/batched_data_dict.py
@@ -431,7 +431,10 @@ def _get_padded_seqlen(seqlen: int) -> int:
                 for shard_idx in range(shards):
                     sharded_micro_indices[shard_idx].append(chunk_sharded_micro_indices[shard_idx])
                     sharded_micro_lengths[shard_idx].append(chunk_sharded_micro_lengths[shard_idx])
-                    sharded_elem_counts_per_gb[shard_idx].append(chunk_sharded_micro_indices[shard_idx][-1][1])
+                    if chunk_sharded_micro_indices[shard_idx] == []:
+                        sharded_elem_counts_per_gb[shard_idx].append(0)
+                    else:
+                        sharded_elem_counts_per_gb[shard_idx].append(chunk_sharded_micro_indices[shard_idx][-1][1])
 
             # flatten global_indices_per_shard
             batch_sorted_indices = []

diff --git a/nemo_rl/models/generation/vllm_backend.py b/nemo_rl/models/generation/vllm_backend.py
@@ -41,22 +41,26 @@ def update_weights_from_ipc_handles(self, ipc_handles):
         try:
             # Get handles for this device
             device_uuid = self.report_device_id()
-            handles = ipc_handles[device_uuid]
+            deserialized = ipc_handles[device_uuid]
             device_id = self.device.index
             weights = []
 
-            # Process each handle to get the tensor
-            for name, handle in handles:
-                func, args = handle
+            all_handles, key_to_type_and_offset_and_size_in_big_tensor = deserialized
+8000
            type_to_packed_big_tensor_size = {}
+            for k, tensor_handle in all_handles:
+                func, args = tensor_handle
                 list_args = list(args)
-                # Update device ID to match the current device
                 list_args[6] = device_id
                 tensor = func(*list_args)
-                weights.append((name, tensor))
+                type_to_packed_big_tensor_size[k] = tensor
+
+            for key, shape, type, offset, size in key_to_type_and_offset_and_size_in_big_tensor:
+                assert offset+size <= type_to_packed_big_tensor_size[type].numel()
+                tensor = type_to_packed_big_tensor_size[type][offset:offset+size].clone().reshape(shape)
+                weights.append((key, tensor))
 
             # Load weights into the model
             self.model_runner.model.load_weights(weights=weights)
-            torch.cuda.synchronize()
             return True
         except Exception as e:
             print(

diff --git a/nemo_rl/models/megatron/refit_utils.py b/nemo_rl/models/megatron/refit_utils.py
@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import re
-from typing import Dict, List, Tuple
+from typing import Dict, List, Tuple, Optional
 
 import torch
 from megatron.core import parallel_state
@@ -79,10 +79,8 @@ def get_tp_dim(model, param_name, named_modules_dict):
 def gather_params(
     model,
     keys,
+    key_to_global_keys: Optional[Dict[str, List[str]]] = None,
 ):
-    import time
-
-    st = time.time()
 
     tp_group = parallel_state.get_tensor_model_parallel_group()
     tp_world_size = torch.distributed.get_world_size(tp_group)
@@ -95,6 +93,7 @@ def gather_params(
     state_dict = model.state_dict()
     gathered_params = {}
     ep_pattern = re.compile(r"mlp\.experts.*\.weight\d*$")
+    use_cached_key_to_global_key_map = key_to_global_keys is not None
     for local_key, shape, dtype in sorted(keys):
         if local_key in state_dict:
             param = state_dict[local_key]
@@ -113,21 +112,23 @@ def gather_params(
             else:
                 # TODO: why do we need to clone?
                 full_param = param
-            global_key = get_global_key_from_local_key(local_key, model.config)
+            if not use_cached_key_to_global_key_map:
+                global_key = get_global_key_from_local_key(local_key, model.config)
         else:
             #  params that may not be on every rank, e.g. the embedding layer
-            global_key = None
-            full_param = torch.empty(
-                *shape, dtype=dtype, device=torch.cuda.current_device()
-            )
+            if not use_cached_key_to_global_key_map:
+                global_key = None
+            full_param = torch.empty(*shape, dtype=dtype, device=torch.cuda.current_device())
 
         # gather across PP group
-        pp_gathered_global_keys = [None] * pp_world_size
-        torch.distributed.all_gather_object(
-            pp_gathered_global_keys, global_key, group=pp_group
-        )
-        # To test no gather:
-        # pp_gathered_global_keys = [global_key] * pp_world_size
+        if not use_cached_key_to_global_key_map:
+            pp_gathered_global_keys = [None] * pp_world_size
+            torch.distributed.all_gather_object(
+                pp_gathered_global_keys, global_key, group=pp_group
+            )
+            # To test no gather:
+            # pp_gathered_global_keys = [global_key] * pp_world_size
+
 
         pp_gathered_params = [
             torch.empty(*shape, dtype=dtype, device=torch.cuda.current_device())
@@ -137,12 +138,13 @@ def gather_params(
 
         # gather across EP group
         if ep_pattern.search(local_key):
-            ep_gathered_global_keys = [None] * ep_world_size
-            torch.distributed.all_gather_object(
-                ep_gathered_global_keys, pp_gathered_global_keys, group=ep_group
-            )
-            # To test no gather:
-            # ep_gathered_global_keys = [pp_gathered_global_keys] * ep_world_size
+            if not use_cached_key_to_global_key_map:
+                ep_gathered_global_keys = [None] * ep_world_size
+                torch.distributed.all_gather_object(
+                    ep_gathered_global_keys, pp_gathered_global_keys, group=ep_group
+                )
+                # To test no gather:
+                # ep_gathered_global_keys = [pp_gathered_global_keys] * ep_world_size
 
             stacked_pp_gathered_params = torch.stack(pp_gathered_params)
             ep_gathered_params = [
@@ -157,19 +159,20 @@ def gather_params(
                 ep_gathered_params, stacked_pp_gathered_params, group=ep_group
             )
 
-            flat_gathered_global_keys = [x for y in ep_gathered_global_keys for x in y]
-            flat_gathered_params = [
-                x for y in ep_gathered_params for x in torch.unbind(y)
-            ]
+            if not use_cached_key_to_global_key_map:
+                flat_gathered_global_keys = [x for y in ep_gathered_global_keys for x in y]
+            flat_gathered_params = [x for y in ep_gathered_params for x in torch.unbind(y)]
+
         else:
-            flat_gathered_global_keys = pp_gathered_global_keys
+            if not use_cached_key_to_global_key_map:
+                flat_gathered_global_keys = pp_gathered_global_keys
             flat_gathered_params = pp_gathered_params
 
+        if use_cached_key_to_global_key_map:
+            flat_gathered_global_keys = key_to_global_keys[local_key]
+
         for k, p in zip(flat_gathered_global_keys, flat_gathered_params):
             if k is not None:
                 gathered_params[k] = p
 
-    torch.cuda.empty_cache()
-    torch.cuda.synchronize()
-    print(f"Time taken to gather params: {time.time() - st}")
     return gathered_params