feat(weight_sync): Raiden weight synchronization with progressive memory reclamation

[SurbhiJainUSC]

This PR enhances the Raiden weight synchronization system with progressive host memory reclamation during parameter staging (preventing host OOMs on 30B+ MoE models), serializes per-source chunk transfers with retry backoff to prevent controller segfaults, and introduces a canonical dict_to_metadata deserializer for RPC workflows.

Motivation & Problem Statement

1. Host OOM During Model Staging: Staging 30B+ MoE parameters (e.g. padded MoE MLP weights requiring ~80GB of host-resident bf16 arrays) previously held the entire JAX PyTree and NumPy arrays simultaneously in RAM, causing peak memory to double and triggering host OOMKills.
2. Multi-Source Transfer Segfaults: Attempting to transfer multiple source work units in a single transfer() call triggered native Raiden controller multi-scheduling crashes at runtime.
3. RPC Metadata Deserialization: Remote inference samplers (such as RLVllmSampler) emit transport metadata as plain dictionary objects over RPC/IPC, requiring a robust deserializer to convert them into strongly-typed WorkUnitMetadata instances.

Checklist

• I have added all the necessary unit tests for my change.
• I have verified that my change does not break existing code and all unit tests pass.
• I have added all appropriate doc-strings/documentation.
• My PR is based on the latest changes of the main branch (if unsure, rebase the code).
• I have signed the Contributor License Agreement.
• I have followed Contribution Guidelines.

[jiangyangmu]

Do you have any test results on this ? like some before/after screenshots to show the improvement.

[SurbhiJainUSC]

@khatwanimohit - do you have any results on this?

[tianshub]

what much performance impact this has? every 4 leaves seem like a very aggressive gc schedule.
Also are we allocating all RAM (e.g. 960B per node for v7x)? curious why we see kubelet eviction, since RAM pressure should be mild right?

[SurbhiJainUSC]

gc.collect() takes ~1ms per invocation. For a model tree with ~100–200 tensor leaves, running gc.collect() every 4 leaves results in ~25–50 collections, adding only ~25–50ms across a weight-sync operation. We can relax this frequency of collection, if preferred.

[datenglin]

Is this trying to unblock some urgent testing?

If not, I think we should fix this by using FFI for the Pathways trainer for weights sync, instead of progressively deleting and garbage-collecting leaves one-by-one during staging.

[tianshub]

are we changing the obj type? how come we need to do explicit dict to metadata conversion now but not before? also is current unit test covering this?