[Tracing] Tracing is slower than the execution

[baoleai]

❓ Questions and Help

Hi, I am currently experiencing slower performance with XLA tracing compared to execution and I'd like to get some help on how to investigate this issue @JackCaoG .
The timeline is like:

My test model is a 7B llm based on xla SPMD.

[JackCaoG]

It is a bit hard for me to read what's going on here. Is thread 971120(fwd) and thread 972534(bwd) doing tracing? What model is it and what parameter size? Tracing should be overlapped with execution, for most transformer models they are execution bound, we did observe that some models like stable diffusion can be tracing bound sometime.

The reason you see memory HtoD happened after the tracing is done is because virtual device. The idea is that we don't know the sharding annotation of the tensor until user tells us the annotation. In order to avoid replicate every tensor(parameter or data) we delay the transfer until tracing is done. By the time execution happens, all of the non-annotaed tensor will be replicated. I think this part can be improved. I plan to work on this later.

[JackCaoG]

If this model is tracing bound, I think maybe we should give dynamo a try. We supported dynamo+spmd for inference a month ago but it is currently broken on head. @wonjoolee95 is looking into the issue right now. We never tried dynamo + training but you can try that after we fixed current inference issue.

[ManfeiBai]

Hi, @JackCaoG, is it ok to assign this issue to you?

[JackCaoG]

@baoleai Is this still a valid issue, or you are OK with it after your pr to removed the unnecessary copy of the HLO

[baoleai]

Although the impact may have diminished after eliminating the copied data, the blank space still remains in the timeline. In my opinion, the primary contributing factors are likely the for loop process in PrepareOutputShardingPropagation and the issue with slow trace.

[JackCaoG]

sg, I will keep this issue open and take a look when I get some cycle.