refactor(nixl): pick the wire dtype with resolve_wire_dtype

src/prime_rl/transports/weights/nixl/nixl.py

@@ -21,7 +21,6 @@
 
 from prime_rl.configs.shared import NIXLWeightBroadcastConfig
 from prime_rl.orchestrator.clients import init_nixl_broadcast
-from prime_rl.trainer.models.base import PreTrainedModelPrimeRL
 from prime_rl.trainer.parallel_dims import ParallelDims
 from prime_rl.trainer.world import wait_for_master
 from prime_rl.transports.weights.base import WeightReceiver, WeightSender
@@ -42,6 +41,7 @@
     TrainerTensor,
     TrainerTensorTable,
 )
+from prime_rl.utils.weights import resolve_wire_dtype
 
 LAYER_RE = re.compile(r"(?:^|\.)layers\.(\d+)(?=\.|$)")
 
@@ -123,7 +123,7 @@ def collect_local_tensor_shards(
         self,
         state_dict: dict[str, torch.Tensor],
         transfer_groups: TransferGroupIndex,
-        keep_in_fp32: Callable[[str], bool],
+        keep_in_fp32: Callable[[str], bool] | None,
     ) -> list[StagedTensorShard]:
         local_shards: list[StagedTensorShard] = []
         for name, value in state_dict.items():
@@ -132,7 +132,7 @@ def collect_local_tensor_shards(
                 continue
             full_shape = tuple(value.shape)
             group_index = self.find_transfer_group_index(name, transfer_groups)
-            wire_dtype = torch.float32 if keep_in_fp32(name) else torch.bfloat16
+            wire_dtype = resolve_wire_dtype(keep_in_fp32, name, torch.bfloat16)
 
             # Unsharded tensors are identical on every rank, so rank 0 serves the only copy.
             if not isinstance(value, DTensor):
@@ -327,15 +327,14 @@ def merge_trainer_table_fragments(self, table_fragments: list[bytes]) -> Trainer
     def initialize_transfer(self, model: nn.Module) -> None:
         if self.initialized:
             return
-        model = cast(PreTrainedModelPrimeRL, model)
         state_dict = model.state_dict()
         transfer_groups = self.build_transfer_group_index(state_dict)
         self.transfer_group_names = transfer_groups.group_names
         if self.is_serving_rank:
             self.staged_shards = self.collect_local_tensor_shards(
                 state_dict,
                 transfer_groups,
-                model.keep_in_fp32_for_weight_transfer,
+                getattr(model, "keep_in_fp32_for_weight_transfer", None),
             )
         self.prepare_staging_buffers()
         table_fragments = self.gather_trainer_table_fragments()

src/prime_rl/utils/weights.py

@@ -152,10 +152,7 @@ def partition_weights(state_dict: dict[str, Tensor], world_size: int, dtype: tor
 
 
 def resolve_wire_dtype(keep_in_fp32: Callable[[str], bool] | None, key: str, default: torch.dtype) -> torch.dtype:
-    """The dtype a tensor must travel to the inference engine in.
-
-    TODO: NIXL keeps its own inline copy of this rule; unify the handling.
-    """
+    """The dtype a tensor must travel to the inference engine in."""
     return torch.float32 if keep_in_fp32 is not None and keep_in_fp32(key) else default