Performance regression on NPU after #154

[A1waysBeenHere]

Qwen3-32B SFT performance regressed by 6x on NPU
Take Qwen3 as a example.
#154 patched some code on Qwen3ForCausalLM.forward to keep cu_seq_lens_x on CPU in NPU case, but which lead to performance regression. A lot communication cannot be overlapped base on profiling.
I am not sure if .cpu() Qwen3ForCausalLM.forward lead to some sync op. And it is all good once I moved the .cpu() to Qwen3Attention.forward.

And I get #199 to patch the .cpu() code in Attention class only.

PS:
We noticed main branch code now, in my case, the memory usage will increase about 4 gigabytes, and which also leads a huge time-consuming in empty_tensor, we canexport MULTI_STREAM_MEMORY_REUSE=2 export PYTORCH_NPU_ALLOC_CONF=expandable_segments:True to optimize memory usage and avoid this problem. But still confused about why #154 get extra 4 gigabytes in memory usage

[Luosuu]

Thank you for your investigation. cc @wey-code @piyifan123

[Crystal-jiang]

ascend

[wey-code]

@A1waysBeenHere
Thanks for the contribution～

In our previous experiments on Qwen models, we did not observe similar extra synchronization being triggered by this change. From a theoretical perspective, the modification introduced in the referenced MR should also be unrelated to the observed increase in NPU memory usage.

To help us better understand and analyze the performance regression and memory behavior, could you please provide the relevant profiling data (e.g., profiler traces or logs)? That would be very helpful for further investigation.

Thanks again for your efforts and support.

[piyifan123]

+1 it would be nice to see a trace (you can use veomni's enable_profiling function). We don't observe this issue in our cases when training qwen3 dense models. The particular PR you refers to actually reduce the per layer attention cu seqlens GPU to CPU sync which was a huge performance degrader.

[A1waysBeenHere]

@A1waysBeenHere Thanks for the contribution～

In our previous experiments on Qwen models, we did not observe similar extra synchronization being triggered by this change. From a theoretical perspective, the modification introduced in the referenced MR should also be unrelated to the observed increase in NPU memory usage.

To help us better understand and analyze the performance regression and memory behavior, could you please provide the relevant profiling data (e.g., profiler traces or logs)? That would be very helpful for further investigation.

Thanks again for your efforts and support.

Will get profiling soon

[A1waysBeenHere]

Cannot get profiling file for some reasons, but I got torch memory snapshot. I noticed that when we have cu_seq_len_q.cpu(), compete layers accumulated and cannot be resharded in time, it's where the extra memory usage from.

Maybe synchronization triggered by .cpu() in def forward collapse the workflow (layer X unshard -> layer X compute/layer X + 1 unshard -> layer X reshard)

abnormal case: There are 11 layers at the end of the backward phase at the same time. Similar situations exist in both forward and backward phase.

normal case:

@piyifan123 @wey-code

[onehaitao]

I think .cpu() operation that caused stream synchronization, so the memory problem was alleviated. If there is no .cpu() operation, delayed memory reuse can lead to memory stacking and push up memory.

[wey-code]

Cannot get profiling file for some reasons, but I got torch memory snapshot. I noticed that when we have cu_seq_len_q.cpu(), compete layers accumulated and cannot be resharded in time, it's where the extra memory usage from.

Maybe synchronization triggered by .cpu() in def forward collapse the workflow (layer X unshard -> layer X compute/layer X + 1 unshard -> layer X reshard)

abnormal case: There are 11 layers at the end of the backward phase at the same time. Similar situations exist in both forward and backward phase.

normal case: [@piyifan123](https://github.com/piyifan123) [@wey-code](https://github.com/wey-code)

This seems related to the FSDP implementation. Could you provide the versions of torch, torch_npu, CANN etc. that you're using?

[A1waysBeenHere]

Cannot get profiling file for some reasons, but I got torch memory snapshot. I noticed that when we have cu_seq_len_q.cpu(), compete layers accumulated and cannot be resharded in time, it's where the extra memory usage from.
Maybe synchronization triggered by .cpu() in def forward collapse the workflow (layer X unshard -> layer X compute/layer X + 1 unshard -> layer X reshard)
abnormal case: There are 11 layers at the end of the backward phase at the same time. Similar situations exist in both forward and backward phase.

normal case:

[@piyifan123](https://github.com/piyifan123) [@wey-code](https://github.com/wey-code)

This seems related to the FSDP implementation. Could you provide the versions of torch, torch_npu, CANN etc. that you're using?

torch 2.7.1
torch_npu 2.7.1.dev20251023
CANN 8.2RC1