[BUG]The time consumption of these two functions, backward_step and get_grad_norm, has increased

[zwc163]

Describe the bug
When I trained a deepseek 16B model using GRPO+sglang+megatron in verl, I found the time taken for each step is steadily increasing, specifically the time for update_policy is continuously growing. Is there any bug？

Stack trace/logs
I found that the time consumption of these two functions in megatron-core, (backward_step and get_grad_norm) has increased.
stack trace 1:
forward_backward_pipelining_without_interleaving (core/pipeline_parallel/schedules.py:1959
backward_step (core/pipeline_parallel/schedules.py:400)
backward (torch/autograd/init.py:347)
_engine_run_backward (torch/autograd/graph.py:823)
stack trace 2:
get_grad_norm (core/optimizer/optimizer.py:192)
get_grad_norm_fp32 (core/optimizer/clip_grads.py:137)

Environment (please complete the following information):

• Megatron-LM : megatron-core==0.12.0
• PyTorch 2.6.0
• CUDA version:12.4
• NCCL version:2.21.5

Additional context
Add any other context about the problem here.

[UnpureRationalist]

We encountered a similar issue when training model using GRPO+sglang+megatron in verl. After training for about 10 hours, the time of backward of a step increase obviously. The nsys profiling results of early steps and steps after 10 hours are shown as follows:

The bwd step in the figure are marked in Megatron source code(https://github.com/NVIDIA/Megatron-LM/blob/core_v0.12.0/megatron/core/pipeline_parallel/schedules.py#L391) using nvtx:

    torch.cuda.nvtx.range_push("bwd step")
    if output_tensor[0].requires_grad:
        if config.deallocate_pipeline_outputs:
            custom_backward(output_tensor[0], output_tensor_grad[0])
        else:
            torch.autograd.backward(output_tensor[0], grad_tensors=output_tensor_grad[0])
    torch.cuda.nvtx.range_pop()

And there are GPU bubbles in the late stage of backward. Is there any solutions? Looking forward to your insights.

[erictang000]

also having this problem with megatron 13.0 while trying to add a megatron backend for SkyRL to do GRPO! Has anyone been able to diagnose or fix this issue? When profiling, I'm seeing the backward pass time is the main culprit for the increase (at line 368 in pipeline_parallel/schedules.py)

For context, here's a 2x increase in training time over the course of a GRPO training run for Qwen-30B-A3B on 4 8xH100 nodes:

Input length and time for a forward pass (fwd_logprobs) is uncorrelated with the increase in training time:

I ran a more minimal reproduction of this for training with Qwen2.5-3B on 1 8xH100 with ~4k max context, where generation length is similarly stable, and logprobs fwd pass time doesn't increase too much either.

Step 0

Fwd pass time (~240 ms):

Backward pass time (~200 ms):

Step 160

Fwd pass time is the same (~240 ms)

Backward pass time increases (~310 ms):

[ISEEKYAN]

Hi, I guess it might be caused by imbalanced data distribution, would you please share more details about the response length curve and the MFU curve?

[erictang000]

Hey @ISEEKYAN! Thanks for your work on mbridge by the way, it's been very useful in adding megatron support to SkyRL!

The response length curve corresponding to the above profiling looks like this - it increases over time, but the increases in training time do not seem directly correlated to the increase in response length (and the fwd pass time remains constant):

I don't have exact numbers on MFU unfortunately. Could you say more about what you mean by imbalanced data distribution? Do you mean within a mini-batch, different GPU micro batches could be imbalanced (our train batch size is fairly low right now, only 2, but using sequence packing)? Or imbalance at the mini batch/batch level?

[ISEEKYAN]

hi @erictang000, very glad that my efforts could be helpful.
By saying imbalance data distribution, I mean different DP group has difference workload, but the DP sync will make all DP group wait for the slowest one. I have a PR in verl to alleviate this see volcengine/verl#3605 for details.

After seeing the response length curve (mean and std) and the increasing time curve of updating policy, I don't think the imbalanced data is the root cause. But it will be helpful if we have an experiment with DP_size=1 and find out what is going on with the backward in profiling timeline.

[robinren03]

also having this problem with megatron 13.0 while trying to add a megatron backend for SkyRL to do GRPO! Has anyone been able to diagnose or fix this issue? When profiling, I'm seeing the backward pass time is the main culprit for the increase (at line 368 in pipeline_parallel/schedules.py)

For context, here's a 2x increase in training time over the course of a GRPO training run for Qwen-30B-A3B on 4 8xH100 nodes:

Input length and time for a forward pass (fwd_logprobs) is uncorrelated with the increase in training time:

I ran a more minimal reproduction of this for training with Qwen2.5-3B on 1 8xH100 with ~4k max context, where generation length is similarly stable, and logprobs fwd pass time doesn't increase too much either.

Step 0

Fwd pass time (~240 ms):

Backward pass time (~200 ms):

## Step 160

Fwd pass time is the same (~240 ms)

Backward pass time increases (~310 ms):

Hi @erictang000 , thanks for your details. I am currently training Qwen-2.5-32B on Megatron. A similar problem emerges to me. I also find it weird that in this case, the forward pass takes similar or a little bit longer than backward pass. And it is also very strange that FSDP performs better than Megatron in this scale. I wonder if all these are normal from your insight. Thanks!

[UnpureRationalist]

This problem in verl can be solved according to link , hope this helpful to you. I think this is related to the communcation in PyTorch affects the backward pass, but I have few knowledge about the details.