[sglang] fix: Bug in megatron+sglang TP16 update_weights.

[SuperCB]

What does this PR do?

We observe the following when using Megatron + Sglang + TP16:

After investigation, we found that this was caused by the cudaipc mechanism not supporting cross-machine access. We have resolved and fixed this bug.

Checklist Before Starting

• Search for similar PRs. Paste at least one query link here: ...
• Format the PR title as [{modules}] {type}: {description} (This will be checked by the CI)
  • {modules} include fsdp, megatron, sglang, vllm, rollout, trainer, ci, training_utils, recipe, hardware, deployment, ray, worker, single_controller, misc, perf, model, algo, env, tool, ckpt, doc, data
  • If this PR involves multiple modules, separate them with , like [megatron, fsdp, doc]
  • {type} is in feat, fix, refactor, chore, test
  • If this PR breaks any API (CLI arguments, config, function signature, etc.), add [BREAKING] to the beginning of the title.
  • Example: [BREAKING][fsdp, megatron] feat: dynamic batching

High-Level Design

Specific Changes

List the specific changes.

Checklist Before Submitting

Important

Please check all the following items before requesting a review, otherwise the reviewer might deprioritize this PR for review.

• Read the Contribute Guide.
• Apply pre-commit checks: pre-commit install && pre-commit run --all-files --show-diff-on-failure --color=always
• Add / Update the documentation.
• Add unit or end-to-end test(s) to the CI workflow to cover all the code. If not feasible, explain why: ...
• Once your PR is ready for CI, send a message in the ci-request channel in the verl Slack workspace.

[CLAassistant]

All committers have signed the CLA.

[zhaochenyang20]

Could you give me an end to end docs to reproduce your fix and verify, like this:

https://github.com/zhaochenyang20/Awesome-ML-SYS-Tutorial/blob/main/rlhf/verl/multi-turn/tool_examples/debug.md

[SuperCB]

Could you give me an end to end docs to reproduce your fix and verify, like this:

https://github.com/zhaochenyang20/Awesome-ML-SYS-Tutorial/blob/main/rlhf/verl/multi-turn/tool_examples/debug.md

This might take some time because our code and configurations are for training internal, self-developed models. We need to prepare reproducible configurations for training on open-source models.

[ETOgaosion]

I think that this logic should be moved to SGLang's processing input weights stage. In verl's sharding manager, we only provide gathered training full tensors and directly send into inference engine. It's better let inference engine to deal with these tensors as abstraction and generalization.

@zhaochenyang20 What do you think of integrating this in SGLang?

[zhaochenyang20]

I think that this logic should be moved to SGLang's processing input weights stage. In verl's sharding manager, we only provide gathered training full tensors and directly send into inference engine. It's better let inference engine to deal with these tensors as abstraction and generalization.

@zhaochenyang20 What do you think of integrating this in SGLang?

Let me discuss with our team

[Hecate0821]

I think that this logic should be moved to SGLang's processing input weights stage. In verl's sharding manager, we only provide gathered training full tensors and directly send into inference engine. It's better let inference engine to deal with these tensors as abstraction and generalization.

@zhaochenyang20 What do you think of integrating this in SGLang?

When passing a tensor across processes, it should be serialized first. You can find a similar implementation in fsdp_sglang.py.

[SuperCB]

I think that this logic should be moved to SGLang's processing input weights stage. In verl's sharding manager, we only provide gathered training full tensors and directly send into inference engine. It's better let inference engine to deal with these tensors as abstraction and generalization.

@zhaochenyang20 What do you think of integrating this in SGLang?

It's a good idea.

[zhaochenyang20]

I think that this logic should be moved to SGLang's processing input weights stage. In verl's sharding manager, we only provide gathered training full tensors and directly send into inference engine. It's better let inference engine to deal with these tensors as abstraction and generalization.
@zhaochenyang20 What do you think of integrating this in SGLang?

It's a good idea.

I think it's better to be serialized in verl and only pass handle tuple to SGLang. Like FSDP does:

def _preprocess_tensor_for_update_weights(tensor: torch.Tensor):
    if isinstance(tensor, DTensor):
        return tensor.full_tensor()
    return tensor

async def update_weights(self, params):
    named_tensors = [(k, v) for k, v in params.items()]
    load_format = None
    for tensor_index, (name, tensor) in enumerate(named_tensors):
        serialized_tensor = MultiprocessingSerializer.serialize(_preprocess_tensor_for_update_weights(tensor))

        if self.device_mesh["infer_tp"].get_local_rank() == 0:
            gathered_serialized_tensors = [None for _ in range(self.device_mesh["infer_tp"].mesh.size()[0])]
        else:
            gathered_serialized_tensors = None
        dist.gather_object(
            obj=serialized_tensor,
            object_gather_list=gathered_serialized_tensors,
            dst=self.device_mesh["infer_tp"].mesh.tolist()[0],
            group=self.device_mesh["infer_tp"].get_group(),
        )

        if self.device_mesh["infer_tp"].get_local_rank() == 0:
            await self.inference_engine.update_weights_from_tensor(
                named_tensors=[
                    (
                        name,
                        LocalSerializedTensor(values=gathered_serialized_tensors),
                    )
                ],
                load_format=load_format,
                flush_cache=tensor_index == len(named_tensors) - 1,
            )

[SuperCB]

if self.device_mesh["infer_tp"].get_local_rank() == 0:
            gathered_serialized_tensors = [None for _ in range(self.device_mesh["infer_tp"].mesh.size()[0])]
        else:
            gathered_serialized_tensors = None
        dist.gather_object(
            obj=serialized_tensor,
            object_gather_list=gathered_serialized_tensors,
            dst=self.device_mesh["infer_tp"].mesh.tolist()[0],
            group=self.device_mesh["infer_tp"].get_group(),
        )

This functionality can be integrated into the SglangRollout class to avoid similar code appearing in FSDPShard Manager and MegatronShard Manager.

[Hecate0821]

if self.device_mesh["infer_tp"].get_local_rank() == 0:
            gathered_serialized_tensors = [None for _ in range(self.device_mesh["infer_tp"].mesh.size()[0])]
        else:
            gathered_serialized_tensors = None
        dist.gather_object(
            obj=serialized_tensor,
            object_gather_list=gathered_serialized_tensors,
            dst=self.device_mesh["infer_tp"].mesh.tolist()[0],
            group=self.device_mesh["infer_tp"].get_group(),
        )

This functionality can be integrated into the SglangRollout class to avoid similar code appearing in FSDPShard Manager and MegatronShard Manager.

if self.device_mesh["infer_tp"].get_local_rank() == 0:
            gathered_serialized_tensors = [None for _ in range(self.device_mesh["infer_tp"].mesh.size()[0])]
        else:
            gathered_serialized_tensors = None
        dist.gather_object(
            obj=serialized_tensor,
            object_gather_list=gathered_serialized_tensors,
            dst=self.device_mesh["infer_tp"].mesh.tolist()[0],
            group=self.device_mesh["infer_tp"].get_group(),
        )

This functionality can be integrated into the SglangRollout class to avoid similar code appearing in FSDPShard Manager and MegatronShard Manager.

You mean add a helper function to SGLang, and call it like this?

for idx, (name, tensor) in enumerate(named_tensors):
    serialized_tensor = MultiprocessingSerializer.serialize(_preprocess_tensor_for_update_weights(tensor))

    gathered_serialized_tensors = SGLangRollout.gather_serialized_tensor(serialized_tensor, self.device_mesh)

    if self.device_mesh["infer_tp"].get_local_rank() == 0:
        await self.inference_engine.update_weights_from_tensor(
            named_tensors=[(name, LocalSerializedTensor(values=gathered_serialized_tensors))],
            ...
        )

[zhaochenyang20]

if self.device_mesh["infer_tp"].get_local_rank() == 0:
            gathered_serialized_tensors = [None for _ in range(self.device_mesh["infer_tp"].mesh.size()[0])]
        else:
            gathered_serialized_tensors = None
        dist.gather_object(
            obj=serialized_tensor,
            object_gather_list=gathered_serialized_tensors,
            dst=self.device_mesh["infer_tp"].mesh.tolist()[0],
            group=self.device_mesh["infer_tp"].get_group(),
        )

This functionality can be integrated into the SglangRollout class to avoid similar code appearing in FSDPShard Manager and MegatronShard Manager.

if self.device_mesh["infer_tp"].get_local_rank() == 0:
            gathered_serialized_tensors = [None for _ in range(self.device_mesh["infer_tp"].mesh.size()[0])]
        else:
            gathered_serialized_tensors = None
        dist.gather_object(
            obj=serialized_tensor,
            object_gather_list=gathered_serialized_tensors,
            dst=self.device_mesh["infer_tp"].mesh.tolist()[0],
            group=self.device_mesh["infer_tp"].get_group(),
        )

This functionality can be integrated into the SglangRollout class to avoid similar code appearing in FSDPShard Manager and MegatronShard Manager.

You mean add a helper function to SGLang, and call it like this?

for idx, (name, tensor) in enumerate(named_tensors):
    serialized_tensor = MultiprocessingSerializer.serialize(_preprocess_tensor_for_update_weights(tensor))

    gathered_serialized_tensors = SGLangRollout.gather_serialized_tensor(serialized_tensor, self.device_mesh)

    if self.device_mesh["infer_tp"].get_local_rank() == 0:
        await self.inference_engine.update_weights_from_tensor(
            named_tensors=[(name, LocalSerializedTensor(values=gathered_serialized_tensors))],
            ...
        )

I mean, just leave this part in verl. If you can make it a utils function, that would be better. No need to change SGLang side.

[Hecate0821]

if self.device_mesh["infer_tp"].get_local_rank() == 0:
            gathered_serialized_tensors = [None for _ in range(self.device_mesh["infer_tp"].mesh.size()[0])]
        else:
            gathered_serialized_tensors = None
        dist.gather_object(
            obj=serialized_tensor,
            object_gather_list=gathered_serialized_tensors,
            dst=self.device_mesh["infer_tp"].mesh.tolist()[0],
            group=self.device_mesh["infer_tp"].get_group(),
        )

This functionality can be integrated into the SglangRollout class to avoid similar code appearing in FSDPShard Manager and MegatronShard Manager.

You mean add a helper function to SGLang, and call it like this?

for idx, (name, tensor) in enumerate(named_tensors):
    serialized_tensor = MultiprocessingSerializer.serialize(_preprocess_tensor_for_update_weights(tensor))

    gathered_serialized_tensors = SGLangRollout.gather_serialized_tensor(serialized_tensor, self.device_mesh)

    if self.device_mesh["infer_tp"].get_local_rank() == 0:
        await self.inference_engine.update_weights_from_tensor(
            named_tensors=[(name, LocalSerializedTensor(values=gathered_serialized_tensors))],
            ...
        )

I mean, just leave this part in verl. If you can make it a utils function, that would be better. No need to change SGLang side.

You're right, that's a better approach. What do you think, @SuperCB?

[SuperCB]

if self.device_mesh["infer_tp"].get_local_rank() == 0:
            gathered_serialized_tensors = [None for _ in range(self.device_mesh["infer_tp"].mesh.size()[0])]
        else:
            gathered_serialized_tensors = None
        dist.gather_object(
            obj=serialized_tensor,
            object_gather_list=gathered_serialized_tensors,
            dst=self.device_mesh["infer_tp"].mesh.tolist()[0],
            group=self.device_mesh["infer_tp"].get_group(),
        )

This functionality can be integrated into the SglangRollout class to avoid similar code appearing in FSDPShard Manager and MegatronShard Manager.

You mean add a helper function to SGLang, and call it like this?

for idx, (name, tensor) in enumerate(named_tensors):
    serialized_tensor = MultiprocessingSerializer.serialize(_preprocess_tensor_for_update_weights(tensor))

    gathered_serialized_tensors = SGLangRollout.gather_serialized_tensor(serialized_tensor, self.device_mesh)

    if self.device_mesh["infer_tp"].get_local_rank() == 0:
        await self.inference_engine.update_weights_from_tensor(
            named_tensors=[(name, LocalSerializedTensor(values=gathered_serialized_tensors))],
            ...
        )

I mean, just leave this part in verl. If you can make it a utils function, that would be better. No need to change SGLang side.

You're right, that's a better approach. What do you think, @SuperCB?

I agree

[ETOgaosion]

The key point is that whether RL frameworks provide training side collected tensors or serialized tensors. Current implementation is OK since SGLang's IPC weight sync design.

I leave 1 questions here:

For extendability, when SGLang can use EP/PP in verl, how to design the sharding process?

[ETOgaosion]

@SuperCB Could you sign the CLA and register your email in this commit to your GitHub account?

[SuperCB]

The key point is that whether RL frameworks provide training side collected tensors or serialized tensors. Current implementation is OK since SGLang's IPC weight sync design.

I leave 1 questions here:

For extendability, when SGLang can use EP/PP in verl, how to design the sharding process?

This is a great question. Our team is also thinking about this issue. If you have any ideas, we can collaborate.

[SuperCB]

I accidentally messed up the commit history,

[SuperCB]

Why would my changes cause this problem? That looks like really unbelievable.

[zhaochenyang20]

Why would my changes cause this problem? That looks like really unbelievable.

could you directly ping me in wechat?