mirror of https://github.com/hpcaitech/ColossalAI
Browse Source
* add SimPO
* fix dataloader
* remove debug code
* add orpo
* fix style
* fix colossalai, transformers version
* fix colossalai, transformers version
* fix colossalai, transformers version
* fix torch colossalai version
* update transformers version
* [shardformer] DeepseekMoE support (#5871)
* [Feature] deepseek moe expert parallel implement
* [misc] fix typo, remove redundant file (#5867)
* [misc] fix typo
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
---------
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
* [Feature] deepseek support & unit test
* [misc] remove debug code & useless print
* [misc] fix typos (#5872)
* [Feature] remove modeling file, use auto config. (#5884)
* [misc] fix typos
* [Feature] deepseek support via auto model, remove modeling file
* [misc] delete useless file
* [misc] fix typos
* [Deepseek] remove redundant code (#5888)
* [misc] fix typos
* [Feature] deepseek support via auto model, remove modeling file
* [misc] delete useless file
* [misc] fix typos
* [misc] remove redundant code
* [Feature/deepseek] resolve comment. (#5889)
* [misc] fix typos
* [Feature] deepseek support via auto model, remove modeling file
* [misc] delete useless file
* [misc] fix typos
* [misc] remove redundant code
* [misc] mv module replacement into if branch
* [misc] add some warning message and modify some code in unit test
* [misc] fix typos
---------
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
* [Hoxfix] Fix CUDA_DEVICE_MAX_CONNECTIONS for comm overlap
Co-authored-by: Edenzzzz <wtan45@wisc.edu>
* [Feat] Diffusion Model(PixArtAlpha/StableDiffusion3) Support (#5838)
* Diffusion Model Inference support
* Stable Diffusion 3 Support
* pixartalpha support
* [HotFix] CI,import,requirements-test for #5838 (#5892)
* [Hot Fix] CI,import,requirements-test
---------
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
* [Feature] Enable PP + SP for llama (#5868)
* fix cross-PP-stage position id length diff bug
* fix typo
* fix typo
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* use a one cross entropy func for all shardformer models
---------
Co-authored-by: Edenzzzz <wtan45@wisc.edu>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
* [ShardFormer] Add Ulysses Sequence Parallelism support for Command-R, Qwen2 and ChatGLM (#5897)
* add benchmark for sft, dpo, simpo, orpo. Add benchmarking result. Support lora with gradient checkpoint
* fix style
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* fix eval
* hotfix citation
* [zero] support all-gather overlap (#5898)
* [zero] support all-gather overlap
* [zero] add overlap all-gather flag
* [misc] fix typo
* [zero] update api
* fix orpo cross entropy loss
* [Auto Parallel]: Speed up intra-op plan generation by 44% (#5446)
* Remove unnecessary calls to deepcopy
* Build DimSpec's difference dict only once
This change considerably speeds up construction speed of DimSpec objects. The difference_dict is the same for each DimSpec object, so a single copy of it is enough.
* Fix documentation of DimSpec's difference method
* [ShardFormer] fix qwen2 sp (#5903)
* [compatibility] support torch 2.2 (#5875)
* Support Pytorch 2.2.2
* keep build_on_pr file and update .compatibility
* fix object_to_tensor usage when torch>=2.3.0 (#5820)
* [misc] support torch2.3 (#5893)
* [misc] support torch2.3
* [devops] update compatibility ci
* [devops] update compatibility ci
* [devops] add debug
* [devops] add debug
* [devops] add debug
* [devops] add debug
* [devops] remove debug
* [devops] remove debug
* [release] update version (#5912)
* [plugin] support all-gather overlap for hybrid parallel (#5919)
* [plugin] fixed all-gather overlap support for hybrid parallel
* add kto
* fix style, add kto data sample
* [Examples] Add lazy init to OPT and GPT examples (#5924)
Co-authored-by: Edenzzzz <wtan45@wisc.edu>
* [ColossalChat] Hotfix for ColossalChat (#5910)
* add ignore and tiny llama
* fix path issue
* run style
* fix issue
* update bash
* add ignore and tiny llama
* fix path issue
* run style
* fix issue
* update bash
* fix ddp issue
* add Qwen 1.5 32B
* refactor tokenization
* [FIX BUG] UnboundLocalError: cannot access local variable 'default_conversation' where it is not associated with a value (#5931)
* cannot access local variable 'default_conversation' where it is not associated with a value
set default value for 'default_conversation'
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
---------
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
* fix test data
* refactor evaluation
* remove real data path
* remove real data path
* Add n_fused as an input from native_module (#5894)
* [FIX BUG] convert env param to int in (#5934)
* [Hotfix] Fix ZeRO typo #5936
Co-authored-by: Edenzzzz <wtan45@wisc.edu>
* [Feature] Add a switch to control whether the model checkpoint needs to be saved after each epoch ends (#5941)
* Add a switch to control whether the model checkpoint needs to be saved after each epoch ends
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
---------
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
* fix style
* fix style
* fix style
* [shardformer] hotfix attn mask (#5945)
* [shardformer] hotfix attn mask (#5947)
* [Feat] Distrifusion Acceleration Support for Diffusion Inference (#5895)
* Distrifusion Support source
* comp comm overlap optimization
* sd3 benchmark
* pixart distrifusion bug fix
* sd3 bug fix and benchmark
* generation bug fix
* naming fix
* add docstring, fix counter and shape error
* add reference
* readme and requirement
* [zero] hotfix update master params (#5951)
* [release] update version (#5952)
* [Chat] Fix lora (#5946)
* fix merging
* remove filepath
* fix style
* Update README.md (#5958)
* [hotfix] Remove unused plan section (#5957)
* remove readme
* fix readme
* update
* [test] add mixtral for sequence classification
* [test] add mixtral transformer test
* [moe] fix plugin
* [test] mixtra pp shard test
* [chore] handle non member group
* [zero] solve hang
* [test] pass mixtral shardformer test
* [moe] implement transit between non moe tp and ep
* [zero] solve hang
* [misc] solve booster hang by rename the variable
* solve hang when parallel mode = pp + dp
* [moe] implement submesh initialization
* [moe] add mixtral dp grad scaling when not all experts are activated
* [chore] manually revert unintended commit
* [chore] trivial fix
* [chore] arg pass & remove drop token
* [test] add mixtral modelling test
* [moe] implement tp
* [moe] test deepseek
* [moe] clean legacy code
* [Feature] MoE Ulysses Support (#5918)
* moe sp support
* moe sp bug solve
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
---------
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
* [chore] minor fix
* [moe] init moe plugin comm setting with sp
* moe sp + ep bug fix
* [moe] finalize test (no pp)
* [moe] full test for deepseek and mixtral (pp + sp to fix)
* [chore] minor fix after rebase
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* [chore] solve moe ckpt test failure and some other arg pass failure
* [moe] remove ops
* [test] fix test: test_zero1_2
* [bug] fix: somehow logger hangs the program
* [moe] deepseek moe sp support
* [test] add check
* [deepseek] replace attn (a workaround for bug in transformers)
* [misc] skip redunant test
* [misc] remove debug/print code
* [moe] refactor mesh assignment
* Revert "[moe] implement submesh initialization"
This reverts commit 2f9bce6686.
* [chore] change moe_pg_mesh to private
* [misc] remove incompatible test config
* [misc] fix ci failure: change default value to false in moe plugin
* [misc] remove useless condition
* [chore] docstring
* [moe] remove force_overlap_comm flag and add warning instead
* [doc] add MoeHybridParallelPlugin docstring
* [moe] solve dp axis issue
* [chore] remove redundant test case, print string & reduce test tokens
* [feat] Dist Loader for Eval (#5950)
* support auto distributed data loader
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* support auto distributed data loader
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* fix tp error
* remove unused parameters
* remove unused
* update inference
* update docs
* update inference
---------
Co-authored-by: Michelle <qianranma8@gmail.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
* [lora] lora support hybrid parallel plugin (#5956)
* lora support hybrid plugin
* fix
* fix
* fix
* fix
* Support overall loss, update KTO logging
* [Docs] clarify launch port
Co-authored-by: Edenzzzz <wtan45@wisc.edu>
* [Hotfix] README link (#5966)
* update ignore
* update readme
* run style
* update readme
* [Hotfix] Avoid fused RMSnorm import error without apex (#5985)
Co-authored-by: Edenzzzz <wtan45@wisc.edu>
* [Chat] fix readme (#5989)
* fix readme
* fix readme, tokenization fully tested
* fix readme, tokenization fully tested
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
---------
Co-authored-by: root <root@notebook-8f919155-6035-47b4-9c6f-1be133b9e2c9-0.notebook-8f919155-6035-47b4-9c6f-1be133b9e2c9.colossal-ai.svc.cluster.local>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
* fix sync condition (#6000)
* [plugin] add cast inputs option for zero (#6003)
* [pre-commit.ci] pre-commit autoupdate (#5995)
updates:
- [github.com/psf/black-pre-commit-mirror: 24.4.2 → 24.8.0](https://github.com/psf/black-pre-commit-mirror/compare/24.4.2...24.8.0)
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
* [misc] Bypass the huggingface bug to solve the mask mismatch problem (#5991)
* [Feature] Zigzag Ring attention (#5905)
* halfway
* fix cross-PP-stage position id length diff bug
* fix typo
* fix typo
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* unified cross entropy func for all shardformer models
* remove redundant lines
* add basic ring attn; debug cross entropy
* fwd bwd logic complete
* fwd bwd logic complete; add experimental triton rescale
* precision tests passed
* precision tests passed
* fix typos and remove misc files
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* add sp_mode to benchmark; fix varlen interface
* update softmax_lse shape by new interface
* change tester name
* remove buffer clone; support packed seq layout
* add varlen tests
* fix typo
* all tests passed
* add dkv_group; fix mask
* remove debug statements
---------
Co-authored-by: Edenzzzz <wtan45@wisc.edu>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
* [misc] update compatibility (#6008)
* [misc] update compatibility
* [misc] update requirements
* [devops] disable requirements cache
* [test] fix torch ddp test
* [test] fix rerun on address in use
* [test] fix lazy init
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* fix the merge
* fix the merge
* overlap kv comm with output rescale (#6017)
Co-authored-by: Edenzzzz <wtan45@wisc.edu>
* fix the merge
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* fix the merge
* fix
* fix
* fix the merge
* fix
* [misc] Use dist logger in plugins (#6011)
* use dist logger in plugins
* remove trash
* print on rank 0
---------
Co-authored-by: Edenzzzz <wtan45@wisc.edu>
* fix
* fix
* fix
* fix
* fix the merge
* fix
* fix
* fix
* fix
---------
Co-authored-by: YeAnbang <anbangy2@outlook.com>
Co-authored-by: Haze188 <haze188@qq.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Edenzzzz <wenxuan.tan@wisc.edu>
Co-authored-by: Edenzzzz <wtan45@wisc.edu>
Co-authored-by: Runyu Lu <77330637+LRY89757@users.noreply.github.com>
Co-authored-by: Guangyao Zhang <xjtu521@qq.com>
Co-authored-by: YeAnbang <44796419+YeAnbang@users.noreply.github.com>
Co-authored-by: Hongxin Liu <lhx0217@gmail.com>
Co-authored-by: Stephan Kö <stephankoe@users.noreply.github.com>
Co-authored-by: アマデウス <kurisusnowdeng@users.noreply.github.com>
Co-authored-by: Tong Li <tong.li352711588@gmail.com>
Co-authored-by: zhurunhua <1281592874@qq.com>
Co-authored-by: Insu Jang <insujang@umich.edu>
Co-authored-by: Gao, Ruiyuan <905370712@qq.com>
Co-authored-by: hxwang <wang1570@e.ntu.edu.sg>
Co-authored-by: Michelle <qianranma8@gmail.com>
Co-authored-by: root <root@notebook-8f919155-6035-47b4-9c6f-1be133b9e2c9-0.notebook-8f919155-6035-47b4-9c6f-1be133b9e2c9.colossal-ai.svc.cluster.local>
pull/6024/head
Wang Binluo
3 months ago
committed by
GitHub
92 changed files with 2222 additions and 463 deletions
@ -1,3 +1,4 @@
|
||||
2.1.0-12.1.0 |
||||
2.2.2-12.1.0 |
||||
2.3.0-12.1.0 |
||||
2.4.0-12.4.1 |
||||
|
||||
File diff suppressed because it is too large
Load Diff
@ -0,0 +1,186 @@
|
||||
import torch |
||||
import torch.distributed as dist |
||||
import torch.nn.functional as F |
||||
from flash_attn import flash_attn_qkvpacked_func, flash_attn_varlen_qkvpacked_func |
||||
from torch.testing import assert_close |
||||
|
||||
import colossalai |
||||
from colossalai.shardformer.layer import AttnMaskType |
||||
from colossalai.shardformer.layer.attn import AttnMaskType, RingAttention |
||||
from colossalai.shardformer.layer.utils import split_batch_zigzag, split_varlen_zigzag |
||||
from colossalai.testing import parameterize, rerun_if_address_is_in_use, spawn |
||||
from colossalai.utils import get_current_device |
||||
|
||||
|
||||
@parameterize("seq_len", [4096]) |
||||
@parameterize("bs", [2]) |
||||
@parameterize("nheads", [5]) |
||||
@parameterize("d", [128]) |
||||
@parameterize("dtype", [torch.bfloat16, torch.float16]) |
||||
def check_ring_attn(seq_len, bs, nheads, d, dtype): |
||||
torch.cuda.manual_seed(2) |
||||
device = get_current_device() |
||||
sp_group = dist.group.WORLD |
||||
sp_size = dist.get_world_size() |
||||
# Some outliers may seem large, but our errors are still lower than |
||||
# than Megatron-LM context parallel's |
||||
# (https://github.com/NVIDIA/TransformerEngine/blob/33a3d02f81c56e6f7b542c09bfa86657078d57fb/tests/pytorch/fused_attn/run_fused_attn_with_cp.py#L215) |
||||
# and the original zigzag implementation's (https://github.com/zhuzilin/ring-flash-attention/tree/main) |
||||
atol = rtol = 7e-3 |
||||
|
||||
# Setup inputs |
||||
qkv = torch.randn(bs, seq_len, 3, nheads, d, device=device, dtype=dtype, requires_grad=True) |
||||
local_qkv = split_batch_zigzag(qkv, sp_group) |
||||
q, k, v = local_qkv.unbind(dim=-3) |
||||
q, k, v = [x.squeeze(2).detach().clone().transpose(1, 2) for x in (q, k, v)] # (B, nHeads, Sq, D) |
||||
q.requires_grad = k.requires_grad = v.requires_grad = True |
||||
|
||||
# Ring attention vs single GPU |
||||
ring_out, ring_lse = RingAttention.attention( |
||||
q, |
||||
k, |
||||
v, |
||||
sp_group, |
||||
AttnMaskType.CAUSAL, |
||||
return_softmax=True, |
||||
inner_ring_size=max(2, sp_size // 2), |
||||
# inner_ring_size=4 |
||||
) |
||||
ring_out = ring_out.transpose(1, 2) |
||||
out, lse, _ = flash_attn_qkvpacked_func( |
||||
qkv, dropout_p=0.0, causal=True, window_size=(-1, -1), alibi_slopes=None, return_attn_probs=True |
||||
) |
||||
|
||||
# Checkout out and softmax denominator |
||||
local_out = split_batch_zigzag(out, sp_group) |
||||
local_lse = split_batch_zigzag(lse, sp_group, seq_dim=-1) |
||||
local_lse = local_lse.transpose(1, 2).contiguous().view(-1, ring_lse.shape[-1]) # (B, nHeads, Sq) -> (T, nHeads) |
||||
assert_close(ring_lse, local_lse, atol=atol, rtol=rtol) |
||||
assert_close(ring_out, local_out, atol=atol, rtol=rtol) |
||||
|
||||
# Check grads |
||||
ring_out.sum().backward() |
||||
out.sum().backward() |
||||
ring_dq, ring_dk, ring_dv = [x.transpose(1, 2) for x in (q.grad, k.grad, v.grad)] |
||||
dqkv = qkv.grad |
||||
local_dqkv = split_batch_zigzag(dqkv, sp_group) |
||||
|
||||
assert_close(ring_dq, local_dqkv[:, :, 0], atol=atol, rtol=rtol) |
||||
assert_close(ring_dk, local_dqkv[:, :, 1], atol=atol, rtol=rtol) |
||||
assert_close(ring_dv, local_dqkv[:, :, 2], atol=atol, rtol=rtol) |
||||
if dist.get_rank() == 0: |
||||
print( |
||||
f"sp_size {dist.get_world_size()}, inner ring size {dist.get_world_size(RingAttention.INNER_RING_GROUP)} passed." |
||||
) |
||||
|
||||
|
||||
@parameterize("seqlen", [4096]) |
||||
@parameterize("bs", [2]) |
||||
@parameterize("nheads", [5]) |
||||
@parameterize("d", [128]) |
||||
@parameterize("dtype", [torch.bfloat16, torch.float16]) |
||||
def check_packed_seq(seqlen, bs, nheads, d, dtype): |
||||
device = get_current_device() |
||||
sp_group = dist.group.WORLD |
||||
sp_size = dist.get_world_size() |
||||
atol = rtol = 7e-3 |
||||
torch.cuda.manual_seed(2) |
||||
# Prepare varlen attention mask |
||||
padding_mask = torch.ones((bs, seqlen), dtype=torch.int, device=device) |
||||
padding_mask[: bs // 2, (seqlen // 4) * 3 :] = 0 |
||||
padding_mask[:, seqlen // 2 :] = 0 |
||||
|
||||
input_embeds = torch.randn(bs, seqlen, nheads, d, device=device, dtype=dtype, requires_grad=True) |
||||
|
||||
# Forward |
||||
# out = ColoAttention.attention(q, k, v, **mask_info) |
||||
flat_input = input_embeds.view(-1, nheads, d)[padding_mask.flatten().nonzero().squeeze()] |
||||
qkv = torch.stack([flat_input] * 3, dim=1) |
||||
qkv.retain_grad() |
||||
|
||||
input_embeds, mask_info, _ = RingAttention.prepare_varlen_batch(padding_mask, sp_group, input_embeds) |
||||
out, lse, _ = flash_attn_varlen_qkvpacked_func( |
||||
qkv, |
||||
mask_info["cu_seqlens"] * sp_size, |
||||
mask_info["max_seqlen"] * sp_size, |
||||
return_attn_probs=True, |
||||
causal=True, |
||||
# deterministic=True |
||||
) |
||||
# Test the splitting function |
||||
local_input = split_varlen_zigzag( |
||||
flat_input, mask_info["cu_seqlens"] * sp_size, sp_group, mask_info["max_seqlen"] * sp_size |
||||
) |
||||
assert (local_input == input_embeds.view(-1, nheads, d)[mask_info["valid_indices"]]).all() |
||||
del local_input, flat_input |
||||
|
||||
q_ring, k_ring, v_ring = [input_embeds.clone().transpose(1, 2) for _ in range(3)] |
||||
q_ring.retain_grad() |
||||
k_ring.retain_grad() |
||||
v_ring.retain_grad() |
||||
|
||||
ring_out, ring_lse = RingAttention.attention( |
||||
q_ring, |
||||
k_ring, |
||||
v_ring, |
||||
sp_group, |
||||
**mask_info, |
||||
pad_output=False, |
||||
return_softmax=True, |
||||
# deterministic=True |
||||
) |
||||
ring_out = ring_out.transpose(1, 2).reshape(-1, nheads, d) |
||||
# Check output |
||||
lse = lse.transpose(0, 1) |
||||
out, lse = split_varlen_zigzag( |
||||
[out, lse], mask_info["cu_seqlens"] * sp_size, sp_group, mask_info["max_seqlen"] * sp_size |
||||
) |
||||
assert_close(lse, ring_lse, atol=atol, rtol=rtol) |
||||
assert_close(out, ring_out, atol=atol, rtol=rtol) |
||||
|
||||
# Check grads |
||||
labels = torch.ones(out.shape[0], dtype=dtype, device=device) |
||||
F.mse_loss(out.sum((-2, -1)), labels).backward() |
||||
F.mse_loss(ring_out.sum((-2, -1)), labels[: ring_out.shape[0]]).backward() |
||||
dq, dk, dv = [ |
||||
split_varlen_zigzag( |
||||
qkv.grad[:, i], mask_info["cu_seqlens"] * sp_size, sp_group, mask_info["max_seqlen"] * sp_size |
||||
) |
||||
for i in range(3) |
||||
] |
||||
dq_ring, dk_ring, dv_ring = [ |
||||
x.transpose(1, 2).reshape(-1, nheads, d)[mask_info["valid_indices"]] |
||||
for x in (q_ring.grad, k_ring.grad, v_ring.grad) |
||||
] |
||||
|
||||
assert_close(dq, dq_ring, atol=atol, rtol=rtol) |
||||
assert_close(dk, dk_ring, atol=atol, rtol=rtol) |
||||
assert_close(dv, dv_ring, atol=atol, rtol=rtol) |
||||
|
||||
|
||||
def launch_single_ring(rank, world_size, port): |
||||
colossalai.launch(rank, world_size, "localhost", port) |
||||
check_packed_seq() |
||||
check_ring_attn() |
||||
|
||||
|
||||
def launch_double_ring(rank, world_size, port): |
||||
colossalai.launch(rank, world_size, "localhost", port) |
||||
check_ring_attn() |
||||
|
||||
|
||||
@rerun_if_address_is_in_use() |
||||
@parameterize("world_size", [2]) |
||||
def test_ring_attn(world_size): |
||||
spawn(launch_single_ring, nprocs=world_size) |
||||
|
||||
|
||||
@rerun_if_address_is_in_use() |
||||
@parameterize("world_size", [4]) |
||||
def test_double_ring(world_size): |
||||
spawn(launch_double_ring, nprocs=world_size) |
||||
|
||||
|
||||
if __name__ == "__main__": |
||||
test_ring_attn() |
||||
test_double_ring() |
||||
Loading…
Reference in new issue