Refit Performance for Large MoE Models

[guyueh1]

[guyueh1]

Is your feature request related to a problem? Please describe.

Currently when running GRPO in large number of GPUs (>= 512), the Ray overhead during refit is large. It is mainly because the single controller sits in the middle of training workers and generation workers and becomes the bottleneck of IPC passing and job submission, as shown in the following figure:

Current refit implementation (code) is a loop over groups of model parameters, and for each group we perform:

• policy.get_weights_ipc_handles: in this step, all policy workers collect parameters to their own GPUs, prepare IPC handles and other metadata, and return them to the main actor
• policy_generation.update_weights_from_ipc_handles: in this step, the main actor sends the IPC handles and metadata to VllmGenerationWorkers, who consequently send them to Vllm engines;

The Ray overhead during this process is that the following parts are carried out by main actor sequantially:

1. job submission: submission of jobs of all policy workers and VllmGenerationWorkers are serial
2. argument serialization and result deserialization
   Adding these overheads up, it becomes a big bottleneck in refit when number of GPUs is large. For instance, in 512 GPU GRPO job for deepseek v3, we spend ~33s out of 48s refit time on passing IPC handles and metadata.

Describe the solution you'd like

The target is to create a way for train worker to directly submit job to vllm engine. As shown in the graph, in co-located case all GPUs have one train worker and one vllm engine; if each train worker can submit job to the vllm engine on the same GPU, an ideal refit process will look like this pseudocode:

# main actor code
policy.refit()

# train worker code
class MegatronPolicyWorker:
    ...
    def refit(self):
        for keys in grouped_param_keys:
            ipc_handles = self.get_weight_ipc_handles(keys)
            self.colocated_vllm_engine.execute_method(
                "update_weights_from_ipc_handles",
                args=ipc_handles
            )

This way, each train worker can call method on vllm engine and all the job submission is distributed/parallelized instead of centralized/serialized in the main actor. This is expected to significantly reduce the 33s part of refit time for DeepSeek.

The overall goal is to remove main actor from the critical path. A possible implementation path is via zero-mq by adding the following components in our code:

• Enable other actors to submit job to vllm engine via zeromq
• Expose the zeromq address for each vllm engine to the corresponding train worker before training starts
• Implement a new refit method based as described above.
  This plan is subject to change when we figure out more details.

Describe alternatives you've considered

To reduce Ray overhead, instead of this way we can also maybe leverage Ray compiled graph.

Additional context

N/A

[guyueh1]

The performance of deepseek refit should be much improved with the following two PRs:

• fix: Fix policy worker placement when using unified placement group #1341 Improve the NCCL speed at trainer side
• feat: refit refactoring with zmq and overlapping #1267 Reduce the overhead of Cuda IPC
  Also, using larger EP and smaller PP helps refit performance (makes total NCCL time shorter). Because smaller PP also helps reduce bubbles, I recommend using EP64PP4 from now on.

[guyueh1]

Some results (Vllm config is fixed to async TP32)

Vllm config	Megatron config	Branch	Note	Refit time
sync TP64	EP16PP16	main		67s
async TP32	EP16PP16	#1341		64s
async TP32	EP16PP16	combining 1267 and 1341		24s
async TP32	EP16PP16	#1341	only run megatron->hf, no vllm update. This is SOL.	9s
async TP32	EP64PP4	#1341	only run megatron->hf, no vllm update. This is SOL.	~7s

Some conclusions:

1. 1267 can significantly reduce the overhead of vllm (64s to 24s)
2. still ~15s away from SOL, need to understand the actual vllm update time to figure out if we can bridge this 15s
3. Using a larger EP can help with all-gather performance, but if we can't reduce the vllm update overhead, this impact is very small.