[Draft] TP Overlap of Micro Batches

[inkcherry]

Motivation

Draft of TP Overlap of Micro-Batches (LLaMA)
This is the forward part as described in arXiv:2409.15241.
Currently, it is implemented with a rough code hack for the POC stage. Functional and mmlu tests have been passed on the torch kernel backend and torch comm backend

Modifications

Checklist

• Format your code according to the Code Formatting with Pre-Commit.
• Add unit tests as outlined in the Running Unit Tests.
• Update documentation / docstrings / example tutorials as needed, according to Writing Documentation.
• Provide throughput / latency benchmark results and accuracy evaluation results as needed, according to Benchmark and Profiling and Accuracy Results.
• For reviewers: If you haven't made any contributions to this PR and are only assisting with merging the main branch, please remove yourself as a co-author when merging the PR.
• Please feel free to join our Slack channel at https://slack.sglang.ai to discuss your PR.

[fzyzcjy]

Wondering whether code in #4068 may be a bit helpful. For example,

• Code to split a batch into two micro batches
• Code to allow users to write two-batch-overlap code as if one-batch code

[zhyncs]

Amazing!

[fzyzcjy]

Extracting: #4965

[inkcherry]

Wondering whether code in #4068 may be a bit helpful. For example,

• Code to split a batch into two micro batches
• Code to allow users to write two-batch-overlap code as if one-batch code

Thanks! this PR is only at the POC stage and hasn't considered the design aspect.
#4068 has a good design, let me understand how to use it.

[fzyzcjy]

@inkcherry Btw many things in diff of #4068 are unrelated to it, but are actually other not-yet-merged PRs that 4068 utilizes , so you can ignore them when reading.

[inkcherry]

Hi, @fzyzcjy,
I’m able to adopt your designed API and have verified mmlu with local tests, including:

• split a batch into two micro batches
• Two-batch stage pipeline layout design

That’s very helpful, once the related code is merged, I’ll continue with further updates : )

[fzyzcjy]

Looks great! I hope my PRs can be merged soon.