AReaL's Megatron training is much slower than FSDP

[spencergotowork]

I've been comparing the training efficiency of AReaL and verl in text generation tasks, and encountered a confusion that I hope to get some insights on. Here's the detailed context:

Experiment Setup

• Task: Text generation training (comparing AReaL and verl)
• Reward Function: Simplified to directly return 1 (no complex evaluation logic)
• Dataset: Identical dataset used for both frameworks
• Text Length: All sequences (including prompt and response) are within 8k tokens
• Inference Acceleration: Both use SGLang for inference speedup
• Training Algorithm: GRPO

Observation

Under the above identical settings, the training duration of AReaL is significantly longer than that of verl.

After profiling, I found that the most time-consuming part in AReaL's model_worker is the loss calculation and model update step, specifically the code around this line in ppo_interface.py.

Question

Could anyone help explain why AReaL takes longer in this scenario, especially regarding the performance bottleneck in the loss calculation and model update stage? Are there any configuration adjustments or optimizations I might be missing for AReaL under such experimental settings?

Thanks in advance for your insights!

[spencergotowork]

The command is:

python3 training/main_async_ppo.py \
    n_nodes=1 n_gpus_per_node=8 \
    allocation_mode=sglang.d4p1m1+d2p2m1 \  # somtimes allocation_mode=sglang.d4p1m1+d4p1m1 \
    cluster.fileroot=experiments \
    actor.type._class=qwen3 \
    actor.path=Qwen/Qwen3-1.7B \
    ref.type._class=qwen3 \
    ref.path=Qwen/Qwen3-1.7B \
    dataset.path=<my-datasets> \
    dataset.train_bs_n_seqs=32 \
    group_size=8 \
    ppo.gen.max_new_tokens=4096 \
    ppo.ppo_n_minibatches=4 \
    actor_train.mb_spec.max_tokens_per_mb=32768 \
    actor_inf.mb_spec.max_tokens_per_mb=32768 \
    max_concurrent_rollouts=16 \
    max_head_offpolicyness=4 \
    group_adv_norm=True \
    ppo.adv_norm=False

[garrett4wade]

Hi, sorry for the late reply. The title and the issue content seem to mismatch. I assume that the issue is AReaL's slow training rather than generation.

Your link to ppo_interface.py marks an if condition before reward calculation. The reward calculation (i.e., get_packed_rewards) should not be slow as it is composed of simple pytorch operations, but the advantage calculation (i.e., get_packed_advantages_and_returns) can indeed be slow due to the lack of the cugae package. Can you check whether you have this package installed?

[spencergotowork]

Hi, sorry for the late reply. The title and the issue content seem to mismatch. I assume that the issue is AReaL's slow training rather than generation.

Your link to ppo_interface.py marks an if condition before reward calculation. The reward calculation (i.e., get_packed_rewards) should not be slow as it is composed of simple pytorch operations, but the advantage calculation (i.e., get_packed_advantages_and_returns) can indeed be slow due to the lack of the cugae package. Can you check whether you have this package installed?

Thank you very much for your reply！
After further localization, I found that the part where the loss calculation begins is the most time-consuming. At https://github.com/inclusionAI/AReaL/blob/main/realhf/impl/model/interface/ppo_interface.py#L788

[garrett4wade]

Actually I don't have no idea why this would happen. Could it be related to using a local function, and that python pickling is slow?

I'll try to reproduce your results recently. cc @nuzant

[spencergotowork]

Actually I don't have no idea why this would happen. Could it be related to using a local function, and that python pickling is slow?

I'll try to reproduce your results recently. cc @nuzant

Thank you for your reply! I've noticed that one potential influencing factor might be the max_tokens_per_mb parameter, which seems to limit the size of the micro batch. In the AReal model, this value is around 30720. Assuming the total length of each prompt and response is 10k, the micro batch size would only be around 3.

However, in verl, I can set the micro batch to 2 or even lower in stages like "ref" and "train", which might be one of the main factors affecting the train_step. When I tried doubling AReal's max_tokens_per_mb to 61440, a significant OutOfMemory (OOM) issue occurred. Further observation revealed that AReal's memory consumption is noticeably higher than that of verl, and I haven't yet pinpointed the exact cause.

For your reference, here are the verl parameters I'm using:

algorithm.adv_estimator=grpo \
data.train_batch_size=32 \
data.max_prompt_length=5000 \
data.max_response_length=5000 \
actor_rollout_ref.model.path=Qwen/Qwen3-8B \
actor_rollout_ref.actor.optim.lr=1e-6 \
actor_rollout_ref.actor.ppo_mini_batch_size=32 \
actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=4 \
actor_rollout_ref.actor.use_kl_loss=True \
actor_rollout_ref.actor.kl_loss_coef=0.001 \
actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=4 \
actor_rollout_ref.rollout.tensor_model_parallel_size=2 \
actor_rollout_ref.rollout.name=sglang \
actor_rollout_ref.rollout.gpu_memory_utilization=0.3 \
actor_rollout_ref.rollout.n=8 \
actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=4 \
actor_rollout_ref.ref.fsdp_config.param_offload=True \
trainer.n_gpus_per_node=8 \
trainer.nnodes=1 

Could you please provide some recommendations for AReal's parameter configuration based on this information? Thank you very much for your help!

[garrett4wade]

Here are some tips for AReaL:

1. The max_tokens_per_mb parameter can be conservatively increased, e.g., to 33000 or 35000. The value doesn't have to be a two's exponential. This value can also be adjusted for different stages, e.g., actor_train.max_tokens_per_mb, actor_train.max_tokens_per_mb, and actor_train.max_tokens_per_mb.
2. If you still encounter OOM, try using pipeline parallel with allocation_mode=sglang.d4p1m1+d1p4m1. AReaL will do micro-batching across stages so it will not be significantly slower, but it indeed decreases the GPU memory usage so that you can set a larger max_tokens_per_mb.

I also noticed that you enabled parameter offload for verl FSDP. This is why you can use a large batch size. Parameter offloading is not supported for AReaL now.

In our internal experiments, we also found that FSDP is faster than the Megatron intergration of AReaL. We are still trying to figure out the reason.

[spencergotowork]

Here are some tips for AReaL:

1. The max_tokens_per_mb parameter can be conservatively increased, e.g., to 33000 or 35000. The value doesn't have to be a two's exponential. This value can also be adjusted for different stages, e.g., actor_train.max_tokens_per_mb, actor_train.max_tokens_per_mb, and actor_train.max_tokens_per_mb.
2. If you still encounter OOM, try using pipeline parallel with allocation_mode=sglang.d4p1m1+d1p4m1. AReaL will do micro-batching across stages so it will not be significantly slower, but it indeed decreases the GPU memory usage so that you can set a larger max_tokens_per_mb.

I also noticed that you enabled parameter offload for verl FSDP. This is why you can use a large batch size. Parameter offloading is not supported for AReaL now.

In our internal experiments, we also found that FSDP is faster than the Megatron intergration of AReaL. We are still trying to figure out the reason.

Thank you so much for the incredible work the Areal community has put in! I’ve noticed a significant improvement in training speed after switching to FSDP—this is really impressive.

I’m curious to learn more about what’s driving this performance boost. Previously, Areal had relatively slow training speeds, and I’m wondering if the shift to FSDP as the training backend is the primary reason behind the improvement? Or are there other optimizations that contributed to this noticeable speedup?

[garrett4wade]

Hi @spencergotowork , it seems that you have tried out the FSDP in AReaL lite, right? We are really glad to see that works for you!

if the shift to FSDP as the training backend is the primary reason behind the improvement?

I'm afraid the answer is yes. We reckon Megatron introduces some overhead in small-scale training scenarios, but we still don't have bandwidth to do profiling and figure it out in details.

Could you please edit the title to something like "AReaL's Megatron training is much slower than FSDP"? Let's keep it accurate and remain it open for now.

[github-actions]

This issue has been automatically marked as stale because it has not had recent activity within the last 14 days.

If this issue is still relevant, please add a comment to keep it open. Otherwise, it will be automatically closed in 16 days.

Thank you for your contribution!

[rchardx]

I'm glad to inform you that #680 supports parameter offloading in AReaL's FSDP.
Meanwhile, AReaL's overall performance has been improved over the last couple of months. Try it out!