alibaba / alibaba/ROLL

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

Open
#484 0 comments 0 reactions 0 assignees View on GitHub
Dominant language
Python
Stars
3.4k
Forks
312
Avg merge
1h 2m
Merged PRs (30d)
2

Description

## 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:

```text
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:

```text
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:

```text
624 tensors
169,417,728 bytes
```

## Scope
The workaround is strictly limited to:

```text
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.

Contributor guide

No contributing guide indexed for this repository

Research direction

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.

Written by the indexing model from the issue text.

Assessment

Tech stack
python, pytorch
Domain
distributed-systems, machine-learning
Issue type
Feature
Difficulty
5/5
Estimated time
Over a week
Activity status
Active
Clarity
Needs clarification
Newbie friendliness
35/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.