[Tiny] Fix docstring in CP abstractions (#30579)
This commit is contained in:
@@ -16,10 +16,10 @@
|
||||
|
||||
For ``cp_size = 4``, each rank owns every fourth token:
|
||||
|
||||
dp_attn_tp0: token0, token4, token8, token12, token16, ...
|
||||
dp_attn_tp1: token1, token5, token9, token13, token17, ...
|
||||
dp_attn_tp2: token2, token6, token10, token14, token18, ...
|
||||
dp_attn_tp3: token3, token7, token11, token15, token19, ...
|
||||
cp0: token0, token4, token8, token12, token16, ...
|
||||
cp1: token1, token5, token9, token13, token17, ...
|
||||
cp2: token2, token6, token10, token14, token18, ...
|
||||
cp3: token3, token7, token11, token15, token19, ...
|
||||
|
||||
After all-gather, tokens are restored to the original order:
|
||||
|
||||
|
||||
@@ -17,10 +17,10 @@
|
||||
For ``cp_size = 4``, each sequence is split into ``2 * cp_size`` blocks. Each
|
||||
rank owns one early block and one late block:
|
||||
|
||||
dp_attn_tp0: block0, block7
|
||||
dp_attn_tp1: block1, block6
|
||||
dp_attn_tp2: block2, block5
|
||||
dp_attn_tp3: block3, block4
|
||||
cp0: block0, block7
|
||||
cp1: block1, block6
|
||||
cp2: block2, block5
|
||||
cp3: block3, block4
|
||||
|
||||
After all-gather, the blocks are reranged back to their original order:
|
||||
|
||||
|
||||
Reference in New Issue
Block a user