Megatron colocated LoRA weight sync fails on CUDA IPC and needs CPU staging fallback

No file paths or tests are named. Start by locating the Megatron colocated LoRA synchronization path and the reported serialize_named_weights → ForkingPickler call chain, then review how lora_rank and converter kwargs flow through it. Done should be defined separately for the optional CPU-staging transport, converter propagation fixes, and manifest prepare/commit behavior.

Problem

We encountered a weight synchronization failure when running Qwen3.6-27B LoRA PPO with:

• Megatron actor training
• vLLM actor inference
• colocated train/infer workers
• TP=4, PP=4
• LoRA weight updates

The original colocated synchronization path serializes GPU tensors through MultiprocessingSerializer, which eventually calls:

serialize_named_weights
-> ForkingPickler
-> torch.multiprocessing.reductions.reduce_tensor
-> storage._share_cuda_()
-> CUDA driver error: unknown error

This happens on our accelerator environment where CUDA IPC is unavailable or incompatible.

Workaround that succeeded

We implemented an optional CPU staging path for colocated LoRA updates:

GPU LoRA tensors
-> flattened bucket
-> bucket.to("cpu")
-> standard pickle bytes
-> Gloo gather_object
-> Ray RPC
-> one bucket.to(device) on the vLLM worker
-> reconstruct tensor views
-> prepare and commit LoRA

With this workaround, the PPO task trained continuously for more than 91 steps.The synchronized payload was LoRA-only:

624 tensors
169,417,728 bytes

Scope

The workaround is strictly limited to:

Megatron + colocated + LoRA

We have not validated it for:

• Full-weight synchronization
• Megatron separated mode
• FSDP2
• SGLang

We also did not add or register GDN LoRA support. GDN in_proj LoRA modules were excluded because the current converter expects special StackedTensors semantics.

Related fixes

During this investigation, we also found:

1. lora_rank was not propagated through the complete Megatron-to-HF converter call chain.
2. **kwargs could be lost during an intermediate weight-buffer flush.
3. Dynamic LoRA updates benefit from an update_id, manifest validation, and prepare/commit semantics to prevent partial updates.

Proposal

Would the maintainers be interested in:

1. A generic optional CPU-staging transport for colocated LoRA synchronization?
2. A separate fix for lora_rank and converter kwargs propagation?
3. LoRA update manifest validation and prepare/commit support?

If this direction is acceptable, I can prepare the changes as separate pull requests with tests.