Question about pipeline

[YYue000]

Hi,

NEO is an excellent research work and I'm very glad to have access to its open-source implementation. I have a question.

How does NEO achieve parallelism between the two sub-batches? Here two _forward_pipeline_stage execute sequentially and I didn't find cuda stream or other async operations.

I'm looking forward to your reply.

[SwiftSeal03]

Thanks for pointing this out. The "stages" execute sequentially, while the two batches execute in parallel within each stage. For synchronization details, you can refer to calls of self._compute_wait_comm() and self._comm_wait_compute() in the LlamaTransformerLayer class.

In a single "stage", NEO launches all CUDA kernels before computing decoding attention on the CPU. The kernel execution is asynchronous, so GPU kernels and CPU kernels are actually executed in parallel.