mirror of https://github.com/hpcaitech/ColossalAI
Browse Source
* [feat] add zerobubble pp (just a frame now); add POC test for dx_dw; add test for zerobubble; * [feat] add dw test; * [fix] fix weight not close; * [update] update text; * [feat] add test run_fwd_bwd automatic scheduling; * [feat] split communication and calculation; fix pop empty send_bwd_buffer error; * [feat] add test for p & p grad; * [feat] add comments for ZBV func; * [fix] rm useless assign and comments; * [fix] fix ci test; add pytest; * [feat] add run_fwd_bwd_with_microbatch (replace input) & test; add p&p.grad assert close test & all pass; * [feat] add apply v_schedule graph; p & p.grad assert err exist; * [fix] update * [feat] fix ci; add assert; * [feat] fix poc format * [feat] fix func name & ci; add comments; * [fix] fix poc test; add comments in poc; * [feat] add optim backward_b_by_grad * [feat] fix optimizer bwd b & w; support return accum loss & output * [feat] add fwd_bwd_step, run_fwd_only; * [fix] fix optim bwd; add license for v_schedule; remove redundant attributes; fix schedule loop "while"--> "for"; add communication dict; * [fix] fix communication_map; * [feat] update test; rm comments; * [fix] rm zbv in hybridplugin * [fix] fix optim bwd; * [fix] fix optim bwd; * [fix] rm output.data after send fwd; * [fix] fix bwd step if condition; remove useless comments and format info; * [fix] fix detach output & release output; * [fix] rm requir_grad for output; * [fix] fix requir grad position and detach position and input&output local buffer append position; * [feat] add memory assertation; * [fix] fix mem check; * [fix] mem assertation' * [fix] fix mem assertation * [fix] fix mem; use a new model shape; only assert mem less and equal than theo; * [fix] fix model zoo import; * [fix] fix redundant detach & clone; add buffer assertation in the end; * [fix] add output_obj_grad assert None at bwd b step; replace input_obj.require_grad_ with treemap; * [fix] update optim state dict assert (include param group & state); fix mem assert after add optim; * [fix] add testcase with microbatch 4; * [zerobubble]Support ZeroBubble Pipeline (#6034) * [feat] add zerobubble pp (just a frame now); add POC test for dx_dw; add test for zerobubble; * [feat] add dw test; * [fix] fix weight not close; * [update] update text; * [feat] add test run_fwd_bwd automatic scheduling; * [feat] split communication and calculation; fix pop empty send_bwd_buffer error; * [feat] add test for p & p grad; * [feat] add comments for ZBV func; * [fix] rm useless assign and comments; * [fix] fix ci test; add pytest; * [feat] add run_fwd_bwd_with_microbatch (replace input) & test; add p&p.grad assert close test & all pass; * [feat] add apply v_schedule graph; p & p.grad assert err exist; * [fix] update * [feat] fix ci; add assert; * [feat] fix poc format * [feat] fix func name & ci; add comments; * [fix] fix poc test; add comments in poc; * [feat] add optim backward_b_by_grad * [feat] fix optimizer bwd b & w; support return accum loss & output * [feat] add fwd_bwd_step, run_fwd_only; * [fix] fix optim bwd; add license for v_schedule; remove redundant attributes; fix schedule loop "while"--> "for"; add communication dict; * [fix] fix communication_map; * [feat] update test; rm comments; * [fix] rm zbv in hybridplugin * [fix] fix optim bwd; * [fix] fix optim bwd; * [fix] rm output.data after send fwd; * [fix] fix bwd step if condition; remove useless comments and format info; * [fix] fix detach output & release output; * [fix] rm requir_grad for output; * [fix] fix requir grad position and detach position and input&output local buffer append position; * [feat] add memory assertation; * [fix] fix mem check; * [fix] mem assertation' * [fix] fix mem assertation * [fix] fix mem; use a new model shape; only assert mem less and equal than theo; * [fix] fix model zoo import; * [fix] fix redundant detach & clone; add buffer assertation in the end; * [fix] add output_obj_grad assert None at bwd b step; replace input_obj.require_grad_ with treemap; * [fix] update optim state dict assert (include param group & state); fix mem assert after add optim; * [fix] add testcase with microbatch 4; * [feat] moehybrid support zerobubble; * [fix] fix zerobubble pp for shardformer type input; * [feat] add more test; * [fix] fix require_grad & deallocate call; * [fix] updatw bwd b&w input; dict --> list[torch.Tensor] * [fix] fix bwd w input; * [fix] fix mem assert; * [fix] fix input_tensors buffer append input_obj(dict) --> Tuple (microbatch, input_obj) , and all bwd b related cal logic; * [fix] use tree_flatten replace dict traverse; * [fix] rm comments; * [fix] fix fwd branch, fwd pass both micro_batch & internal_inputs' * [fix] fix pipeline util func deallocate --> release_tensor_data; fix bwd_b loss bwd branch; * [fix] fix detach clone release order; * [fix] fix ci --> oom in 4096 hidden dim; * [fix] fix dumb clone; * [fix] fix detach_output_obj clone; * [fix] fix stage_indices; * [fix] fix traverse; traverse dict --> traverse tensor List; * [fix] fix zerobubble; support shardformer model type; * [fix] rm comments; * [fix] fix test_pipeline_utils ci; * [fix] remove duplicate arg; rm comments; * [fix] remove chunk 0 stage 0 bwd b; u don't have to cal micrbatch's dx; * [fix] rm print & comments; * [plugin] hybrid support zero bubble pipeline (#6060) * hybrid support zbv * fix fix * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix * Update zero_bubble_pp.py * fix * fix-ci * fix [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci fix * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix * fix * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix * fix * fix * fix * [zerobubble]Support ZeroBubble Pipeline (#6034) * [feat] add zerobubble pp (just a frame now); add POC test for dx_dw; add test for zerobubble; * [feat] add dw test; * [fix] fix weight not close; * [update] update text; * [feat] add test run_fwd_bwd automatic scheduling; * [feat] split communication and calculation; fix pop empty send_bwd_buffer error; * [feat] add test for p & p grad; * [feat] add comments for ZBV func; * [fix] rm useless assign and comments; * [fix] fix ci test; add pytest; * [feat] add run_fwd_bwd_with_microbatch (replace input) & test; add p&p.grad assert close test & all pass; * [feat] add apply v_schedule graph; p & p.grad assert err exist; * [fix] update * [feat] fix ci; add assert; * [feat] fix poc format * [feat] fix func name & ci; add comments; * [fix] fix poc test; add comments in poc; * [feat] add optim backward_b_by_grad * [feat] fix optimizer bwd b & w; support return accum loss & output * [feat] add fwd_bwd_step, run_fwd_only; * [fix] fix optim bwd; add license for v_schedule; remove redundant attributes; fix schedule loop "while"--> "for"; add communication dict; * [fix] fix communication_map; * [feat] update test; rm comments; * [fix] rm zbv in hybridplugin * [fix] fix optim bwd; * [fix] fix optim bwd; * [fix] rm output.data after send fwd; * [fix] fix bwd step if condition; remove useless comments and format info; * [fix] fix detach output & release output; * [fix] rm requir_grad for output; * [fix] fix requir grad position and detach position and input&output local buffer append position; * [feat] add memory assertation; * [fix] fix mem check; * [fix] mem assertation' * [fix] fix mem assertation * [fix] fix mem; use a new model shape; only assert mem less and equal than theo; * [fix] fix model zoo import; * [fix] fix redundant detach & clone; add buffer assertation in the end; * [fix] add output_obj_grad assert None at bwd b step; replace input_obj.require_grad_ with treemap; * [fix] update optim state dict assert (include param group & state); fix mem assert after add optim; * [fix] add testcase with microbatch 4; * hybrid support zbv * fix fix * fix * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Update zero_bubble_pp.py * fix * fix-ci * fix [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci fix * fix * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix * fix * fix * fix * fix * fix * fix * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix * fix * fix * fix --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: duanjunwen <935724073@qq.com> * [feat] zerobubble support moehybridplugin; * [feat] update optimizer bwd; ä¸ * [fix] fix build ci; * [zerobubble] rebase main (#6075) * fp8 operators for compressed communication cast_to_fp8, cast_from_fp8, all_reduce_fp8 * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix typo * fix scaling algorithm in FP8 casting * support fp8 communication in pipeline parallelism * add fp8_communication flag in the script * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * shardformer fp8 * fix rebase * remove all to all * fix shardformer fp8 communication training degradation * [fp8] support all-gather flat tensor (#5932) * [fp8] add fp8 comm for low level zero * [test] add zero fp8 test case * [Feature] llama shardformer fp8 support (#5938) * add llama shardformer fp8 * Llama Shardformer Parity * fix typo * fix all reduce * fix pytest failure * fix reduce op and move function to fp8.py * fix typo * [FP8] rebase main (#5963) * 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 commitpull/6147/head2f9bce6686. * [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 * fp8 operators for compressed communication cast_to_fp8, cast_from_fp8, all_reduce_fp8 * fix scaling algorithm in FP8 casting * support fp8 communication in pipeline parallelism * add fp8_communication flag in the script * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix typo * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * shardformer fp8 * fix rebase * remove all to all * fix shardformer fp8 communication training degradation * [fp8] support all-gather flat tensor (#5932) * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix * Update low_level_optim.py --------- 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: Wang Binluo <32676639+wangbluo@users.noreply.github.com> Co-authored-by: HangXu <hangxu0304@gmail.com> * [fp8]support all2all fp8 (#5953) * support all2all fp8 * fix * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix * fix * fix * [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> * [fp8] add fp8 linear (#5967) * [fp8] add fp8 linear * [test] fix fp8 linear test condition * [test] fix fp8 linear test condition * [test] fix fp8 linear test condition * [fp8] support fp8 amp for hybrid parallel plugin (#5975) * [fp8] support fp8 amp for hybrid parallel plugin * [test] add fp8 hook test * [fp8] fix fp8 linear compatibility * fix (#5976) * [Feature]: support FP8 communication in DDP, FSDP, Gemini (#5928) * support fp8_communication in the Torch DDP grad comm, FSDP grad comm, and FSDP params comm * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * implement communication hook for FSDP params all-gather * added unit test for fp8 operators * support fp8 communication in GeminiPlugin * update training scripts to support fsdp and fp8 communication * fixed some minor bugs observed in unit test * add all_gather_into_tensor_flat_fp8 * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * add skip the test if torch < 2.2.0 * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * add skip the test if torch < 2.2.0 * add skip the test if torch < 2.2.0 * add fp8_comm flag * rebase latest fp8 operators * rebase latest fp8 operators * [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> * [test ci]Feature/fp8 comm (#5981) * fix * fix * fix * [fp8] support gemini plugin (#5978) * [fp8] refactor hook * [fp8] support gemini plugin * [example] add fp8 option for llama benchmark * [fp8] use torch compile (torch >= 2.3.0) (#5979) * [fp8] use torch compile (torch >= 2.4.0) * [fp8] set use_fast_accum in linear * [chore] formal version check * [chore] fix sig * [fp8]Moe support fp8 communication (#5977) * fix * support moe fp8 * fix * fix * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix * fix * fix * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix * fix * fix fix fi * fix * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> * [fp8] support hybrid parallel plugin (#5982) * support fp8 comm for qwen2 model * support fp8 comm for qwen2 model * support fp8 comm for qwen2 model * fp8 * fix * bert and bloom * chatglm and command * gpt2,gptj,bert, falcon,blip2 * mistral,opy,sam,t5,vit,whisper * fix * fix * fix * [fp8] refactor fp8 linear with compile (#5993) * [fp8] refactor fp8 linear with compile * [fp8] fix linear test * [fp8] fix linear test * [fp8] support asynchronous FP8 communication (#5997) * fix * fix * fix * support async all2all * support async op for all gather * fix * fix * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> * [fp8] update torch.compile for linear_fp8 to >= 2.4.0 (#6004) * [fp8] linear perf enhancement * [fp8]update reduce-scatter test (#6002) * fix * fix * fix * fix * [fp8] add use_fp8 option for MoeHybridParallelPlugin (#6009) * [fp8] zero support fp8 linear. (#6006) * fix * fix * fix * zero fp8 * zero fp8 * Update requirements.txt * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix the merge * fix the merge * 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 * fix * fix * fix * fix * fix the merge * fix * fix * fix * fix * [fp8] Merge feature/fp8_comm to main branch of Colossalai (#6016) * 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 commit2f9bce6686. * [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> * fix * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Update train_dpo.py * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Update low_level_zero_plugin.py * fix * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * [CI] Remove triton version for compatibility bug; update req torch >=2.2 (#6018) * remove triton version * remove torch 2.2 * remove torch 2.1 * debug * remove 2.1 build tests * require torch >=2.2 --------- Co-authored-by: Edenzzzz <wtan45@wisc.edu> * [plugin] hotfix zero plugin (#6036) * [plugin] hotfix zero plugin * [plugin] hotfix zero plugin * [Colossal-LLaMA] Refactor latest APIs (#6030) * refactor latest code * update api * add dummy dataset * update Readme * add setup * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * update files * add PP support * update arguments * update argument * reorg folder * update version * remove IB infor * update utils * update readme * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * update save for zero * update save * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * add apex * update --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> * add fused norm (#6038) * [FP8] unsqueeze scale to make it compatible with torch.compile (#6040) * [colossalai/checkpoint_io/...] fix bug in load_state_dict_into_model; format error msg (#6020) * fix bug in load_state_dict_into_model; format error msg * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Update utils.py to support checking missing_keys * Update general_checkpoint_io.py fix bug in missing_keys error message * retrigger tests --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> * [Hotfix] Remove deprecated install (#6042) * remove deprecated install * remove unused folder * [fp8] optimize all-gather (#6043) * [fp8] optimize all-gather * [fp8] fix all gather fp8 ring * [fp8] enable compile * [fp8] fix all gather fp8 ring * [fp8] fix linear hook (#6046) * [fp8] disable all_to_all_fp8 in intranode (#6045) * enhance all_to_all_fp8 with internode comm control * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * disable some fp8 ops due to performance issue * [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> * [release] update version (#6041) * [release] update version * [devops] update comp test * [devops] update comp test debug * [devops] debug comp test * [devops] debug comp test * [devops] debug comp test * [devops] debug comp test * [devops] debug comp test * [Feature] Split cross-entropy computation in SP (#5959) * halfway * fix cross-PP-stage position id length diff bug * 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 * 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 * adapt chatglm, command-R, qwen * debug * 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 * add sp_mode to benchmark; fix varlen interface * update softmax_lse shape by new interface * add varlen tests * fix typo * all tests passed * add dkv_group; fix mask * remove debug statements * add comments * q1 index only once * remove events to simplify stream sync * simplify forward/backward logic * 2d ring forward passed * 2d ring backward passed * fixes * fix ring attn loss * 2D ring backward + llama passed * merge * update logger * fix typo * rebase * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix typo * remove typos * fixes * support GPT --------- Co-authored-by: Edenzzzz <wtan45@wisc.edu> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> * [hotfix] moe hybrid parallelism benchmark & follow-up fix (#6048) * [example] pass use_fp8_comm flag to all plugins * [example] add mixtral benchmark * [moe] refine assertion and check * [moe] fix mixtral & add more tests * [moe] consider checking dp * sp group and moe_dp_group * [mixtral] remove gate tp & add more tests * [deepseek] fix tp & sp for deepseek * [mixtral] minor fix * [deepseek] add deepseek benchmark * [fp8] hotfix backward hook (#6053) * [fp8] hotfix backward hook * [fp8] hotfix pipeline loss accumulation * [doc] update sp doc (#6055) * update sp doc * fix * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix * fix * fix --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> * fix the sp * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix the attn * fix * fix * fix * fix * [zerobubble]Support ZeroBubble Pipeline (#6034) * [feat] add zerobubble pp (just a frame now); add POC test for dx_dw; add test for zerobubble; * [feat] add dw test; * [fix] fix weight not close; * [update] update text; * [feat] add test run_fwd_bwd automatic scheduling; * [feat] split communication and calculation; fix pop empty send_bwd_buffer error; * [feat] add test for p & p grad; * [feat] add comments for ZBV func; * [fix] rm useless assign and comments; * [fix] fix ci test; add pytest; * [feat] add run_fwd_bwd_with_microbatch (replace input) & test; add p&p.grad assert close test & all pass; * [feat] add apply v_schedule graph; p & p.grad assert err exist; * [fix] update * [feat] fix ci; add assert; * [feat] fix poc format * [feat] fix func name & ci; add comments; * [fix] fix poc test; add comments in poc; * [feat] add optim backward_b_by_grad * [feat] fix optimizer bwd b & w; support return accum loss & output * [feat] add fwd_bwd_step, run_fwd_only; * [fix] fix optim bwd; add license for v_schedule; remove redundant attributes; fix schedule loop "while"--> "for"; add communication dict; * [fix] fix communication_map; * [feat] update test; rm comments; * [fix] rm zbv in hybridplugin * [fix] fix optim bwd; * [fix] fix optim bwd; * [fix] rm output.data after send fwd; * [fix] fix bwd step if condition; remove useless comments and format info; * [fix] fix detach output & release output; * [fix] rm requir_grad for output; * [fix] fix requir grad position and detach position and input&output local buffer append position; * [feat] add memory assertation; * [fix] fix mem check; * [fix] mem assertation' * [fix] fix mem assertation * [fix] fix mem; use a new model shape; only assert mem less and equal than theo; * [fix] fix model zoo import; * [fix] fix redundant detach & clone; add buffer assertation in the end; * [fix] add output_obj_grad assert None at bwd b step; replace input_obj.require_grad_ with treemap; * [fix] update optim state dict assert (include param group & state); fix mem assert after add optim; * [fix] add testcase with microbatch 4; * [fp8] fix missing fp8_comm flag in mixtral (#6057) * fix * fix * fix * [fp8] Disable all_gather intranode. Disable Redundant all_gather fp8 (#6059) * all_gather only internode, fix pytest * fix cuda arch <89 compile pytest error * fix pytest failure * disable all_gather_into_tensor_flat_fp8 * fix fp8 format * fix pytest * fix conversations * fix chunk tuple to list * [doc] FP8 training and communication document (#6050) * Add FP8 training and communication document * add fp8 docstring for plugins * fix typo * fix typo * fix * fix * [moe] add parallel strategy for shared_expert && fix test for deepseek (#6063) * [ColossalEval] support for vllm (#6056) * support vllm * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * modify vllm and update readme * run pre-commit * remove dupilicated lines and refine code * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * update param name * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * refine code * update readme * refine code * [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> * [release] update version (#6062) * [feat] add zerobubble pp (just a frame now); add POC test for dx_dw; add test for zerobubble; * [update] update text; * [feat] add test run_fwd_bwd automatic scheduling; * [feat] fix poc format * [fix] fix poc test; add comments in poc; * [feat] add optim backward_b_by_grad * [feat] fix optimizer bwd b & w; support return accum loss & output * [fix] fix optim bwd; add license for v_schedule; remove redundant attributes; fix schedule loop "while"--> "for"; add communication dict; * [feat] update test; rm comments; * [fix] rm zbv in hybridplugin * [fix] fix optim bwd; * [fix] fix optim bwd; * [fix] rm output.data after send fwd; * [fix] fix bwd step if condition; remove useless comments and format info; * [fix] fix mem check; * [fix] fix mem assertation * [fix] fix mem; use a new model shape; only assert mem less and equal than theo; * [fix] fix model zoo import; * [feat] moehybrid support zerobubble; * [fix] fix zerobubble pp for shardformer type input; * [fix] fix require_grad & deallocate call; * [fix] fix mem assert; * [fix] fix fwd branch, fwd pass both micro_batch & internal_inputs' * [fix] fix pipeline util func deallocate --> release_tensor_data; fix bwd_b loss bwd branch; * [fix] fix zerobubble; support shardformer model type; * [fix] fix test_pipeline_utils ci; * [plugin] hybrid support zero bubble pipeline (#6060) * hybrid support zbv * fix fix * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix * Update zero_bubble_pp.py * fix * fix-ci * fix [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci fix * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix * fix * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix * fix * fix * fix * [zerobubble]Support ZeroBubble Pipeline (#6034) * [feat] add zerobubble pp (just a frame now); add POC test for dx_dw; add test for zerobubble; * [feat] add dw test; * [fix] fix weight not close; * [update] update text; * [feat] add test run_fwd_bwd automatic scheduling; * [feat] split communication and calculation; fix pop empty send_bwd_buffer error; * [feat] add test for p & p grad; * [feat] add comments for ZBV func; * [fix] rm useless assign and comments; * [fix] fix ci test; add pytest; * [feat] add run_fwd_bwd_with_microbatch (replace input) & test; add p&p.grad assert close test & all pass; * [feat] add apply v_schedule graph; p & p.grad assert err exist; * [fix] update * [feat] fix ci; add assert; * [feat] fix poc format * [feat] fix func name & ci; add comments; * [fix] fix poc test; add comments in poc; * [feat] add optim backward_b_by_grad * [feat] fix optimizer bwd b & w; support return accum loss & output * [feat] add fwd_bwd_step, run_fwd_only; * [fix] fix optim bwd; add license for v_schedule; remove redundant attributes; fix schedule loop "while"--> "for"; add communication dict; * [fix] fix communication_map; * [feat] update test; rm comments; * [fix] rm zbv in hybridplugin * [fix] fix optim bwd; * [fix] fix optim bwd; * [fix] rm output.data after send fwd; * [fix] fix bwd step if condition; remove useless comments and format info; * [fix] fix detach output & release output; * [fix] rm requir_grad for output; * [fix] fix requir grad position and detach position and input&output local buffer append position; * [feat] add memory assertation; * [fix] fix mem check; * [fix] mem assertation' * [fix] fix mem assertation * [fix] fix mem; use a new model shape; only assert mem less and equal than theo; * [fix] fix model zoo import; * [fix] fix redundant detach & clone; add buffer assertation in the end; * [fix] add output_obj_grad assert None at bwd b step; replace input_obj.require_grad_ with treemap; * [fix] update optim state dict assert (include param group & state); fix mem assert after add optim; * [fix] add testcase with microbatch 4; * hybrid support zbv * fix fix * fix * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Update zero_bubble_pp.py * fix * fix-ci * fix [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci fix * fix * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix * fix * fix * fix * fix * fix * fix * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix * fix * fix * fix --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: duanjunwen <935724073@qq.com> * [feat] add zerobubble pp (just a frame now); add POC test for dx_dw; add test for zerobubble; * [update] update text; * [feat] add test run_fwd_bwd automatic scheduling; * [feat] fix poc format * [fix] fix poc test; add comments in poc; * [feat] add optim backward_b_by_grad * [feat] fix optimizer bwd b & w; support return accum loss & output * [fix] fix optim bwd; add license for v_schedule; remove redundant attributes; fix schedule loop "while"--> "for"; add communication dict; * [feat] update test; rm comments; * [fix] fix optim bwd; * [fix] fix optim bwd; * [fix] rm output.data after send fwd; * [fix] fix bwd step if condition; remove useless comments and format info; * [fix] fix mem check; * [fix] fix mem assertation * [fix] fix mem; use a new model shape; only assert mem less and equal than theo; * [fix] fix model zoo import; * [fix] fix mem assert; * [fix] fix fwd branch, fwd pass both micro_batch & internal_inputs' * [plugin] hybrid support zero bubble pipeline (#6060) * hybrid support zbv * fix fix * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix * Update zero_bubble_pp.py * fix * fix-ci * fix [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci fix * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix * fix * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix * fix * fix * fix * [zerobubble]Support ZeroBubble Pipeline (#6034) * [feat] add zerobubble pp (just a frame now); add POC test for dx_dw; add test for zerobubble; * [feat] add dw test; * [fix] fix weight not close; * [update] update text; * [feat] add test run_fwd_bwd automatic scheduling; * [feat] split communication and calculation; fix pop empty send_bwd_buffer error; * [feat] add test for p & p grad; * [feat] add comments for ZBV func; * [fix] rm useless assign and comments; * [fix] fix ci test; add pytest; * [feat] add run_fwd_bwd_with_microbatch (replace input) & test; add p&p.grad assert close test & all pass; * [feat] add apply v_schedule graph; p & p.grad assert err exist; * [fix] update * [feat] fix ci; add assert; * [feat] fix poc format * [feat] fix func name & ci; add comments; * [fix] fix poc test; add comments in poc; * [feat] add optim backward_b_by_grad * [feat] fix optimizer bwd b & w; support return accum loss & output * [feat] add fwd_bwd_step, run_fwd_only; * [fix] fix optim bwd; add license for v_schedule; remove redundant attributes; fix schedule loop "while"--> "for"; add communication dict; * [fix] fix communication_map; * [feat] update test; rm comments; * [fix] rm zbv in hybridplugin * [fix] fix optim bwd; * [fix] fix optim bwd; * [fix] rm output.data after send fwd; * [fix] fix bwd step if condition; remove useless comments and format info; * [fix] fix detach output & release output; * [fix] rm requir_grad for output; * [fix] fix requir grad position and detach position and input&output local buffer append position; * [feat] add memory assertation; * [fix] fix mem check; * [fix] mem assertation' * [fix] fix mem assertation * [fix] fix mem; use a new model shape; only assert mem less and equal than theo; * [fix] fix model zoo import; * [fix] fix redundant detach & clone; add buffer assertation in the end; * [fix] add output_obj_grad assert None at bwd b step; replace input_obj.require_grad_ with treemap; * [fix] update optim state dict assert (include param group & state); fix mem assert after add optim; * [fix] add testcase with microbatch 4; * hybrid support zbv * fix fix * fix * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Update zero_bubble_pp.py * fix * fix-ci * fix [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci fix * fix * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix * fix * fix * fix * fix * fix * fix * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix * fix * fix * fix --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: duanjunwen <935724073@qq.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Co-authored-by: HangXu <hangxu0304@gmail.com> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: GuangyaoZhang <xjtu521@qq.com> Co-authored-by: Hongxin Liu <lhx0217@gmail.com> Co-authored-by: YeAnbang <anbangy2@outlook.com> Co-authored-by: Haze188 <haze188@qq.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: YeAnbang <44796419+YeAnbang@users.noreply.github.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: Wang Binluo <32676639+wangbluo@users.noreply.github.com> Co-authored-by: wangbluo <2538539015@qq.com> 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: duanjunwen <935724073@qq.com> Co-authored-by: Camille Zhong <44392324+Camille7777@users.noreply.github.com> * [fix] fix mixtral policy; * [fix] fix mixtral policy; * [feat] support zbv in mixtral benchmark; * [fix] MixtralForCausalLMPolicy get_held_layer support zbv; * [feat] update MixtralPipelineForwards --> mixtral_model_forward; support zbv; * [feat] support MixtralPipelineForwards--> mixtral_for_causal_lm_forward for zbv * [zero bubble] support zero (#6080) * fp8 operators for compressed communication cast_to_fp8, cast_from_fp8, all_reduce_fp8 * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix typo * fix scaling algorithm in FP8 casting * support fp8 communication in pipeline parallelism * add fp8_communication flag in the script * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * shardformer fp8 * fix rebase * remove all to all * fix shardformer fp8 communication training degradation * [fp8] support all-gather flat tensor (#5932) * [fp8] add fp8 comm for low level zero * [test] add zero fp8 test case * [Feature] llama shardformer fp8 support (#5938) * add llama shardformer fp8 * Llama Shardformer Parity * fix typo * fix all reduce * fix pytest failure * fix reduce op and move function to fp8.py * fix typo * [FP8] rebase main (#5963) * 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 commit2f9bce6686. * [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 * fp8 operators for compressed communication cast_to_fp8, cast_from_fp8, all_reduce_fp8 * fix scaling algorithm in FP8 casting * support fp8 communication in pipeline parallelism * add fp8_communication flag in the script * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix typo * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * shardformer fp8 * fix rebase * remove all to all * fix shardformer fp8 communication training degradation * [fp8] support all-gather flat tensor (#5932) * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix * Update low_level_optim.py --------- 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: Wang Binluo <32676639+wangbluo@users.noreply.github.com> Co-authored-by: HangXu <hangxu0304@gmail.com> * [fp8]support all2all fp8 (#5953) * support all2all fp8 * fix * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix * fix * fix * [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> * [fp8] add fp8 linear (#5967) * [fp8] add fp8 linear * [test] fix fp8 linear test condition * [test] fix fp8 linear test condition * [test] fix fp8 linear test condition * [fp8] support fp8 amp for hybrid parallel plugin (#5975) * [fp8] support fp8 amp for hybrid parallel plugin * [test] add fp8 hook test * [fp8] fix fp8 linear compatibility * fix (#5976) * [Feature]: support FP8 communication in DDP, FSDP, Gemini (#5928) * support fp8_communication in the Torch DDP grad comm, FSDP grad comm, and FSDP params comm * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * implement communication hook for FSDP params all-gather * added unit test for fp8 operators * support fp8 communication in GeminiPlugin * update training scripts to support fsdp and fp8 communication * fixed some minor bugs observed in unit test * add all_gather_into_tensor_flat_fp8 * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * add skip the test if torch < 2.2.0 * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * add skip the test if torch < 2.2.0 * add skip the test if torch < 2.2.0 * add fp8_comm flag * rebase latest fp8 operators * rebase latest fp8 operators * [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> * [test ci]Feature/fp8 comm (#5981) * fix * fix * fix * [fp8] support gemini plugin (#5978) * [fp8] refactor hook * [fp8] support gemini plugin * [example] add fp8 option for llama benchmark * [fp8] use torch compile (torch >= 2.3.0) (#5979) * [fp8] use torch compile (torch >= 2.4.0) * [fp8] set use_fast_accum in linear * [chore] formal version check * [chore] fix sig * [fp8]Moe support fp8 communication (#5977) * fix * support moe fp8 * fix * fix * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix * fix * fix * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix * fix * fix fix fi * fix * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> * [fp8] support hybrid parallel plugin (#5982) * support fp8 comm for qwen2 model * support fp8 comm for qwen2 model * support fp8 comm for qwen2 model * fp8 * fix * bert and bloom * chatglm and command * gpt2,gptj,bert, falcon,blip2 * mistral,opy,sam,t5,vit,whisper * fix * fix * fix * [fp8] refactor fp8 linear with compile (#5993) * [fp8] refactor fp8 linear with compile * [fp8] fix linear test * [fp8] fix linear test * [fp8] support asynchronous FP8 communication (#5997) * fix * fix * fix * support async all2all * support async op for all gather * fix * fix * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> * [fp8] update torch.compile for linear_fp8 to >= 2.4.0 (#6004) * [fp8] linear perf enhancement * [fp8]update reduce-scatter test (#6002) * fix * fix * fix * fix * [fp8] add use_fp8 option for MoeHybridParallelPlugin (#6009) * [fp8] zero support fp8 linear. (#6006) * fix * fix * fix * zero fp8 * zero fp8 * Update requirements.txt * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix the merge * fix the merge * 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 * fix * fix * fix * fix * fix the merge * fix * fix * fix * fix * [fp8] Merge feature/fp8_comm to main branch of Colossalai (#6016) * 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 commit2f9bce6686. * [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> * fix * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Update train_dpo.py * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Update low_level_zero_plugin.py * fix * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * [CI] Remove triton version for compatibility bug; update req torch >=2.2 (#6018) * remove triton version * remove torch 2.2 * remove torch 2.1 * debug * remove 2.1 build tests * require torch >=2.2 --------- Co-authored-by: Edenzzzz <wtan45@wisc.edu> * [plugin] hotfix zero plugin (#6036) * [plugin] hotfix zero plugin * [plugin] hotfix zero plugin * [Colossal-LLaMA] Refactor latest APIs (#6030) * refactor latest code * update api * add dummy dataset * update Readme * add setup * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * update files * add PP support * update arguments * update argument * reorg folder * update version * remove IB infor * update utils * update readme * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * update save for zero * update save * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * add apex * update --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> * add fused norm (#6038) * [FP8] unsqueeze scale to make it compatible with torch.compile (#6040) * [colossalai/checkpoint_io/...] fix bug in load_state_dict_into_model; format error msg (#6020) * fix bug in load_state_dict_into_model; format error msg * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Update utils.py to support checking missing_keys * Update general_checkpoint_io.py fix bug in missing_keys error message * retrigger tests --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> * [Hotfix] Remove deprecated install (#6042) * remove deprecated install * remove unused folder * [fp8] optimize all-gather (#6043) * [fp8] optimize all-gather * [fp8] fix all gather fp8 ring * [fp8] enable compile * [fp8] fix all gather fp8 ring * [fp8] fix linear hook (#6046) * [fp8] disable all_to_all_fp8 in intranode (#6045) * enhance all_to_all_fp8 with internode comm control * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * disable some fp8 ops due to performance issue * [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> * [release] update version (#6041) * [release] update version * [devops] update comp test * [devops] update comp test debug * [devops] debug comp test * [devops] debug comp test * [devops] debug comp test * [devops] debug comp test * [devops] debug comp test * [Feature] Split cross-entropy computation in SP (#5959) * halfway * fix cross-PP-stage position id length diff bug * 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 * 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 * adapt chatglm, command-R, qwen * debug * 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 * add sp_mode to benchmark; fix varlen interface * update softmax_lse shape by new interface * add varlen tests * fix typo * all tests passed * add dkv_group; fix mask * remove debug statements * add comments * q1 index only once * remove events to simplify stream sync * simplify forward/backward logic * 2d ring forward passed * 2d ring backward passed * fixes * fix ring attn loss * 2D ring backward + llama passed * merge * update logger * fix typo * rebase * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix typo * remove typos * fixes * support GPT --------- Co-authored-by: Edenzzzz <wtan45@wisc.edu> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> * [hotfix] moe hybrid parallelism benchmark & follow-up fix (#6048) * [example] pass use_fp8_comm flag to all plugins * [example] add mixtral benchmark * [moe] refine assertion and check * [moe] fix mixtral & add more tests * [moe] consider checking dp * sp group and moe_dp_group * [mixtral] remove gate tp & add more tests * [deepseek] fix tp & sp for deepseek * [mixtral] minor fix * [deepseek] add deepseek benchmark * [fp8] hotfix backward hook (#6053) * [fp8] hotfix backward hook * [fp8] hotfix pipeline loss accumulation * [doc] update sp doc (#6055) * update sp doc * fix * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix * fix * fix --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> * fix the sp * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix the attn * fix * fix * fix * fix * [zerobubble]Support ZeroBubble Pipeline (#6034) * [feat] add zerobubble pp (just a frame now); add POC test for dx_dw; add test for zerobubble; * [feat] add dw test; * [fix] fix weight not close; * [update] update text; * [feat] add test run_fwd_bwd automatic scheduling; * [feat] split communication and calculation; fix pop empty send_bwd_buffer error; * [feat] add test for p & p grad; * [feat] add comments for ZBV func; * [fix] rm useless assign and comments; * [fix] fix ci test; add pytest; * [feat] add run_fwd_bwd_with_microbatch (replace input) & test; add p&p.grad assert close test & all pass; * [feat] add apply v_schedule graph; p & p.grad assert err exist; * [fix] update * [feat] fix ci; add assert; * [feat] fix poc format * [feat] fix func name & ci; add comments; * [fix] fix poc test; add comments in poc; * [feat] add optim backward_b_by_grad * [feat] fix optimizer bwd b & w; support return accum loss & output * [feat] add fwd_bwd_step, run_fwd_only; * [fix] fix optim bwd; add license for v_schedule; remove redundant attributes; fix schedule loop "while"--> "for"; add communication dict; * [fix] fix communication_map; * [feat] update test; rm comments; * [fix] rm zbv in hybridplugin * [fix] fix optim bwd; * [fix] fix optim bwd; * [fix] rm output.data after send fwd; * [fix] fix bwd step if condition; remove useless comments and format info; * [fix] fix detach output & release output; * [fix] rm requir_grad for output; * [fix] fix requir grad position and detach position and input&output local buffer append position; * [feat] add memory assertation; * [fix] fix mem check; * [fix] mem assertation' * [fix] fix mem assertation * [fix] fix mem; use a new model shape; only assert mem less and equal than theo; * [fix] fix model zoo import; * [fix] fix redundant detach & clone; add buffer assertation in the end; * [fix] add output_obj_grad assert None at bwd b step; replace input_obj.require_grad_ with treemap; * [fix] update optim state dict assert (include param group & state); fix mem assert after add optim; * [fix] add testcase with microbatch 4; * [fp8] fix missing fp8_comm flag in mixtral (#6057) * fix * fix * fix * [fp8] Disable all_gather intranode. Disable Redundant all_gather fp8 (#6059) * all_gather only internode, fix pytest * fix cuda arch <89 compile pytest error * fix pytest failure * disable all_gather_into_tensor_flat_fp8 * fix fp8 format * fix pytest * fix conversations * fix chunk tuple to list * [doc] FP8 training and communication document (#6050) * Add FP8 training and communication document * add fp8 docstring for plugins * fix typo * fix typo * fix * fix * [moe] add parallel strategy for shared_expert && fix test for deepseek (#6063) * [ColossalEval] support for vllm (#6056) * support vllm * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * modify vllm and update readme * run pre-commit * remove dupilicated lines and refine code * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * update param name * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * refine code * update readme * refine code * [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> * [release] update version (#6062) * [feat] add zerobubble pp (just a frame now); add POC test for dx_dw; add test for zerobubble; * [update] update text; * [feat] add test run_fwd_bwd automatic scheduling; * [feat] fix poc format * [fix] fix poc test; add comments in poc; * [feat] add optim backward_b_by_grad * [feat] fix optimizer bwd b & w; support return accum loss & output * [fix] fix optim bwd; add license for v_schedule; remove redundant attributes; fix schedule loop "while"--> "for"; add communication dict; * [feat] update test; rm comments; * [fix] rm zbv in hybridplugin * [fix] fix optim bwd; * [fix] fix optim bwd; * [fix] rm output.data after send fwd; * [fix] fix bwd step if condition; remove useless comments and format info; * [fix] fix mem check; * [fix] fix mem assertation * [fix] fix mem; use a new model shape; only assert mem less and equal than theo; * [fix] fix model zoo import; * [feat] moehybrid support zerobubble; * [fix] fix zerobubble pp for shardformer type input; * [fix] fix require_grad & deallocate call; * [fix] fix mem assert; * [fix] fix fwd branch, fwd pass both micro_batch & internal_inputs' * [fix] fix pipeline util func deallocate --> release_tensor_data; fix bwd_b loss bwd branch; * [fix] fix zerobubble; support shardformer model type; * [fix] fix test_pipeline_utils ci; * [plugin] hybrid support zero bubble pipeline (#6060) * hybrid support zbv * fix fix * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix * Update zero_bubble_pp.py * fix * fix-ci * fix [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci fix * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix * fix * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix * fix * fix * fix * [zerobubble]Support ZeroBubble Pipeline (#6034) * [feat] add zerobubble pp (just a frame now); add POC test for dx_dw; add test for zerobubble; * [feat] add dw test; * [fix] fix weight not close; * [update] update text; * [feat] add test run_fwd_bwd automatic scheduling; * [feat] split communication and calculation; fix pop empty send_bwd_buffer error; * [feat] add test for p & p grad; * [feat] add comments for ZBV func; * [fix] rm useless assign and comments; * [fix] fix ci test; add pytest; * [feat] add run_fwd_bwd_with_microbatch (replace input) & test; add p&p.grad assert close test & all pass; * [feat] add apply v_schedule graph; p & p.grad assert err exist; * [fix] update * [feat] fix ci; add assert; * [feat] fix poc format * [feat] fix func name & ci; add comments; * [fix] fix poc test; add comments in poc; * [feat] add optim backward_b_by_grad * [feat] fix optimizer bwd b & w; support return accum loss & output * [feat] add fwd_bwd_step, run_fwd_only; * [fix] fix optim bwd; add license for v_schedule; remove redundant attributes; fix schedule loop "while"--> "for"; add communication dict; * [fix] fix communication_map; * [feat] update test; rm comments; * [fix] rm zbv in hybridplugin * [fix] fix optim bwd; * [fix] fix optim bwd; * [fix] rm output.data after send fwd; * [fix] fix bwd step if condition; remove useless comments and format info; * [fix] fix detach output & release output; * [fix] rm requir_grad for output; * [fix] fix requir grad position and detach position and input&output local buffer append position; * [feat] add memory assertation; * [fix] fix mem check; * [fix] mem assertation' * [fix] fix mem assertation * [fix] fix mem; use a new model shape; only assert mem less and equal than theo; * [fix] fix model zoo import; * [fix] fix redundant detach & clone; add buffer assertation in the end; * [fix] add output_obj_grad assert None at bwd b step; replace input_obj.require_grad_ with treemap; * [fix] update optim state dict assert (include param group & state); fix mem assert after add optim; * [fix] add testcase with microbatch 4; * hybrid support zbv * fix fix * fix * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Update zero_bubble_pp.py * fix * fix-ci * fix [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci fix * fix * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix * fix * fix * fix * fix * fix * fix * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix * fix * fix * fix --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: duanjunwen <935724073@qq.com> * [feat] add zerobubble pp (just a frame now); add POC test for dx_dw; add test for zerobubble; * [update] update text; * [feat] add test run_fwd_bwd automatic scheduling; * [feat] fix poc format * [fix] fix poc test; add comments in poc; * [feat] add optim backward_b_by_grad * [feat] fix optimizer bwd b & w; support return accum loss & output * [fix] fix optim bwd; add license for v_schedule; remove redundant attributes; fix schedule loop "while"--> "for"; add communication dict; * [feat] update test; rm comments; * [fix] fix optim bwd; * [fix] fix optim bwd; * [fix] rm output.data after send fwd; * [fix] fix bwd step if condition; remove useless comments and format info; * [fix] fix mem check; * [fix] fix mem assertation * [fix] fix mem; use a new model shape; only assert mem less and equal than theo; * [fix] fix model zoo import; * [fix] fix mem assert; * [fix] fix fwd branch, fwd pass both micro_batch & internal_inputs' * [plugin] hybrid support zero bubble pipeline (#6060) * hybrid support zbv * fix fix * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix * Update zero_bubble_pp.py * fix * fix-ci * fix [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci fix * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix * fix * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix * fix * fix * fix * [zerobubble]Support ZeroBubble Pipeline (#6034) * [feat] add zerobubble pp (just a frame now); add POC test for dx_dw; add test for zerobubble; * [feat] add dw test; * [fix] fix weight not close; * [update] update text; * [feat] add test run_fwd_bwd automatic scheduling; * [feat] split communication and calculation; fix pop empty send_bwd_buffer error; * [feat] add test for p & p grad; * [feat] add comments for ZBV func; * [fix] rm useless assign and comments; * [fix] fix ci test; add pytest; * [feat] add run_fwd_bwd_with_microbatch (replace input) & test; add p&p.grad assert close test & all pass; * [feat] add apply v_schedule graph; p & p.grad assert err exist; * [fix] update * [feat] fix ci; add assert; * [feat] fix poc format * [feat] fix func name & ci; add comments; * [fix] fix poc test; add comments in poc; * [feat] add optim backward_b_by_grad * [feat] fix optimizer bwd b & w; support return accum loss & output * [feat] add fwd_bwd_step, run_fwd_only; * [fix] fix optim bwd; add license for v_schedule; remove redundant attributes; fix schedule loop "while"--> "for"; add communication dict; * [fix] fix communication_map; * [feat] update test; rm comments; * [fix] rm zbv in hybridplugin * [fix] fix optim bwd; * [fix] fix optim bwd; * [fix] rm output.data after send fwd; * [fix] fix bwd step if condition; remove useless comments and format info; * [fix] fix detach output & release output; * [fix] rm requir_grad for output; * [fix] fix requir grad position and detach position and input&output local buffer append position; * [feat] add memory assertation; * [fix] fix mem check; * [fix] mem assertation' * [fix] fix mem assertation * [fix] fix mem; use a new model shape; only assert mem less and equal than theo; * [fix] fix model zoo import; * [fix] fix redundant detach & clone; add buffer assertation in the end; * [fix] add output_obj_grad assert None at bwd b step; replace input_obj.require_grad_ with treemap; * [fix] update optim state dict assert (include param group & state); fix mem assert after add optim; * [fix] add testcase with microbatch 4; * hybrid support zbv * fix fix * fix * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Update zero_bubble_pp.py * fix * fix-ci * fix [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci fix * fix * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix * fix * fix * fix * fix * fix * fix * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix * fix * fix * fix --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: duanjunwen <935724073@qq.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * zbv support zero * fix * fix * fix --------- Co-authored-by: HangXu <hangxu0304@gmail.com> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: GuangyaoZhang <xjtu521@qq.com> Co-authored-by: Hongxin Liu <lhx0217@gmail.com> Co-authored-by: YeAnbang <anbangy2@outlook.com> Co-authored-by: Haze188 <haze188@qq.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: YeAnbang <44796419+YeAnbang@users.noreply.github.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: Wang Binluo <32676639+wangbluo@users.noreply.github.com> Co-authored-by: wangbluo <2538539015@qq.com> 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: duanjunwen <935724073@qq.com> Co-authored-by: Camille Zhong <44392324+Camille7777@users.noreply.github.com> * [fix] fix llama, mixtral benchmark zbv loss none bug; update mixtral & llama policy and modeling; * [feat] Linear1D_COL/ROW support zbv WeightGradStore; * [feat] support use_zbv in llama, mixtral modeling; only replace Linear1D_Col/Row policy; * [fix] fix test case; moe error in second iter * [feat]EPMixtralSparseMoeBlock (op in MOE) support zbv; * [fix] fix bwd b; now bwd w only for Layer replaced by Linear1D_Col/Row; other layer perform a fully bwd; * [fix] debug zbv llama test; * [fix] rm use_zbv flag in Shardconfig; rm debug info; * [fix] add & fix llama test * [feat] support meta cache, meta_grad_send, meta_tensor_send; fix runtime too long in Recv Bwd; benchmark for llama + Hybrid(tp+pp); * [fix\ fix fail case test_shard_llama * [fix] fix test_shard_llama * [fix] fix llama modeling policy; * [fix] fix test_shard_llama ci; * [fix] fix test zerobubble * [fix] fix handle name; rm useless comments; * [fix] fix send recv signature; * [fix] fix comment in llama & benchmark * [feat] support no tensor parallel Linear in shardformer; Add test for use weightGradStore and not use WeightGradStore * [fix] fix linear (no tp) ops func name; * [feat] support zbv in mixtral benchmark; (#6083) * [feat] support zbv in mixtral benchmark; * [fix] MixtralForCausalLMPolicy get_held_layer support zbv; * [feat] update MixtralPipelineForwards --> mixtral_model_forward; support zbv; * [feat] support MixtralPipelineForwards--> mixtral_for_causal_lm_forward for zbv * [fix] fix llama, mixtral benchmark zbv loss none bug; update mixtral & llama policy and modeling; * [feat] Linear1D_COL/ROW support zbv WeightGradStore; * [feat] support use_zbv in llama, mixtral modeling; only replace Linear1D_Col/Row policy; * [fix] fix test case; moe error in second iter * [feat]EPMixtralSparseMoeBlock (op in MOE) support zbv; * [fix] fix bwd b; now bwd w only for Layer replaced by Linear1D_Col/Row; other layer perform a fully bwd; * [fix] debug zbv llama test; * [fix] rm use_zbv flag in Shardconfig; rm debug info; * [fix] add & fix llama test * [feat] support meta cache, meta_grad_send, meta_tensor_send; fix runtime too long in Recv Bwd; benchmark for llama + Hybrid(tp+pp); * [fix\ fix fail case test_shard_llama * [fix] fix test_shard_llama * [fix] fix llama modeling policy; * [fix] fix test_shard_llama ci; * [fix] fix test zerobubble * [fix] fix handle name; rm useless comments; * [fix] fix send recv signature; * [fix] fix comment in llama & benchmark * [feat] support no tensor parallel Linear in shardformer; Add test for use weightGradStore and not use WeightGradStore * [fix] fix linear (no tp) ops func name; * [fix] fix fp8 args in HybridParallel * [fix] fix hybridparall use_fp8 config * [fix] fix use_fp8 flag * [fix] fix model zoo init * [feat] support no_tp Linear for sharderformer.llama * [fix] fix zbv llama pp4 * [fix] fix send_tensor_metadata & send_grad_metadata; * [feat] fix testcase; * [feat] support mixtral policy with zbv tp_Linear & non_tp_Linear * [feat] update mixtral policy & bert policy for zerobubble * [fix] fix p2p error in zbv * [fix] fix attn * [fix] fix mixtral modeling & policy; update wait handles; doing benchmarking for llama hybrid; * [fix] fix zbv wait_handle * [fix] rm debug info; update llama policy; update wait handle * [fix] fix test_lora * [fix] fix test_lora in llama policy * [fix] fix wait handle in run_fwd_bwd * [fix] remove debug info; * [fix] rm unused comments * [fix] fix fp8 overlap code * [fix] fix yml file & v_schedule comments * [fix] rm fwd only meta cache comments; --------- Co-authored-by: flybird11111 <1829166702@qq.com> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: HangXu <hangxu0304@gmail.com> Co-authored-by: GuangyaoZhang <xjtu521@qq.com> Co-authored-by: Hongxin Liu <lhx0217@gmail.com> Co-authored-by: YeAnbang <anbangy2@outlook.com> Co-authored-by: Haze188 <haze188@qq.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: YeAnbang <44796419+YeAnbang@users.noreply.github.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: Wang Binluo <32676639+wangbluo@users.noreply.github.com> Co-authored-by: wangbluo <2538539015@qq.com> 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: Camille Zhong <44392324+Camille7777@users.noreply.github.com>
duanjunwen
2 days ago
committed by
GitHub
35 changed files with 3682 additions and 123 deletions
@ -1,11 +1,12 @@
|
||||
from .p2p import PipelineP2PCommunication |
||||
from .schedule import InterleavedSchedule, OneForwardOneBackwardSchedule, PipelineSchedule |
||||
from .schedule import InterleavedSchedule, OneForwardOneBackwardSchedule, PipelineSchedule, ZeroBubbleVPipeScheduler |
||||
from .stage_manager import PipelineStageManager |
||||
|
||||
__all__ = [ |
||||
"PipelineSchedule", |
||||
"OneForwardOneBackwardSchedule", |
||||
"InterleavedSchedule", |
||||
"ZeroBubbleVPipeScheduler", |
||||
"PipelineP2PCommunication", |
||||
"PipelineStageManager", |
||||
] |
||||
|
||||
@ -1,9 +1,11 @@
|
||||
from .base import PipelineSchedule |
||||
from .interleaved_pp import InterleavedSchedule |
||||
from .one_f_one_b import OneForwardOneBackwardSchedule |
||||
from .zero_bubble_pp import ZeroBubbleVPipeScheduler |
||||
|
||||
__all__ = [ |
||||
"PipelineSchedule", |
||||
"OneForwardOneBackwardSchedule", |
||||
"InterleavedSchedule", |
||||
"ZeroBubbleVPipeScheduler", |
||||
] |
||||
|
||||
@ -0,0 +1,449 @@
|
||||
# Refer from Zero Bubble Pipeline Parallelism. |
||||
# Github: https://github.com/sail-sg/zero-bubble-pipeline-parallelism |
||||
# Paper: https://arxiv.org/abs/2401.10241 |
||||
# The following applies to all files unless otherwise noted: |
||||
# Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved. |
||||
# |
||||
# Redistribution and use in source and binary forms, with or without |
||||
# modification, are permitted provided that the following conditions |
||||
# are met: |
||||
# * Redistributions of source code must retain the above copyright |
||||
# notice, this list of conditions and the following disclaimer. |
||||
# * Redistributions in binary form must reproduce the above copyright |
||||
# notice, this list of conditions and the following disclaimer in the |
||||
# documentation and/or other materials provided with the distribution. |
||||
# * Neither the name of NVIDIA CORPORATION nor the names of its |
||||
# contributors may be used to endorse or promote products derived |
||||
# from this software without specific prior written permission. |
||||
# |
||||
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY |
||||
# EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE |
||||
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR |
||||
# PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR |
||||
# CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, |
||||
# EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, |
||||
# PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR |
||||
# PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY |
||||
# OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT |
||||
# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE |
||||
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. |
||||
|
||||
from collections import deque |
||||
from dataclasses import dataclass |
||||
|
||||
|
||||
@dataclass(eq=True, frozen=True) |
||||
class ScheduledNode: |
||||
type: str |
||||
chunk: int |
||||
stage: int |
||||
minibatch: int |
||||
start_time: int = 0 |
||||
completion_time: int = 0 |
||||
rollback: bool = False |
||||
|
||||
|
||||
class PipelineGraph(object): |
||||
"""PipelineGraph""" |
||||
|
||||
def __init__( |
||||
self, |
||||
n_stage, |
||||
n_micro, |
||||
f_cost, |
||||
b_cost, |
||||
w_cost, |
||||
c_cost, |
||||
f_mem, |
||||
b_mem, |
||||
w_mem, |
||||
max_mem=None, |
||||
): |
||||
self.n_node = 6 * n_stage * n_micro |
||||
self.n_stage = n_stage |
||||
self.n_micro = n_micro |
||||
self.f_cost = f_cost |
||||
self.b_cost = b_cost |
||||
self.w_cost = w_cost |
||||
self.c_cost = c_cost |
||||
self.f_mem = f_mem |
||||
self.b_mem = b_mem |
||||
self.w_mem = w_mem |
||||
self.fbw_cost = [f_cost, b_cost, w_cost] |
||||
self.fbw_mem = [f_mem, b_mem, w_mem] |
||||
self.max_mem = max_mem or f_mem * self.n_stage * 2 |
||||
|
||||
def get_id(self, cat, chunk, stage, micro): |
||||
return ( |
||||
cat * 2 * self.n_stage * self.n_micro + chunk * self.n_stage * self.n_micro + stage * self.n_micro + micro |
||||
) |
||||
|
||||
def try_v_schedule(self, fill_f=True, fill_b=True, approved_bubble=None): |
||||
count = [] |
||||
for i in range(self.n_stage): |
||||
count.append([0] * 6) |
||||
|
||||
end_time = [-1] * self.n_node |
||||
cur_time = [0] * self.n_stage |
||||
mem = [0] * self.n_stage |
||||
stage_bubble = [0] * self.n_stage |
||||
pending_w = [deque() for _ in range(self.n_stage)] |
||||
schedule = [[] for _ in range(self.n_stage)] |
||||
stage_str = [" " * i for i in range(self.n_stage)] |
||||
|
||||
if approved_bubble is None: |
||||
approved_bubble = [-1] * self.n_stage |
||||
max_approved_bubble = max(approved_bubble) |
||||
|
||||
def get_max_stage_bubble(stage=-1): |
||||
max_stage_bubble = 0 |
||||
for bb in stage_bubble: |
||||
max_stage_bubble = max(max_stage_bubble, bb) |
||||
if stage >= 0: |
||||
max_stage_bubble = max(max_stage_bubble, max_approved_bubble - approved_bubble[stage]) |
||||
return max_stage_bubble |
||||
|
||||
def put_w(stage): |
||||
assert len(pending_w[stage]) > 0 |
||||
_, chunk_, _ = pending_w[stage].popleft() |
||||
put(2, chunk_, stage) |
||||
|
||||
def put(cat, chunk, stage, assert_cnt=True): |
||||
_tmp = _no_bubble = cur_time[stage] + self.fbw_cost[cat] |
||||
_cnt = count[stage][cat * 2 + chunk] |
||||
# assert _cnt < self.n_micro |
||||
if _cnt >= self.n_micro: |
||||
if not assert_cnt: |
||||
stage_str[stage] += " " |
||||
cur_time[stage] = _tmp # TODO |
||||
return |
||||
assert False |
||||
assert mem[stage] + self.fbw_mem[cat] <= self.max_mem |
||||
stage_str[stage] += "FfBbWw"[cat * 2 + chunk] + str(_cnt + 1) + " " * (3 - len(str(_cnt + 1))) |
||||
if cat > 0 or chunk > 0: |
||||
last_id = cat * 2 + chunk - 1 |
||||
if cat < 2: |
||||
assert end_time[self.get_id(last_id // 2, last_id % 2, stage, _cnt)] >= 0 |
||||
else: |
||||
assert end_time[self.get_id(1, chunk, stage, _cnt)] >= 0 |
||||
if chunk == 1 and cat < 2: |
||||
if stage < self.n_stage - 1: |
||||
_fa_id = self.get_id(cat, chunk, stage + 1, _cnt) |
||||
assert end_time[_fa_id] >= 0 |
||||
_tmp = max(_tmp, end_time[_fa_id] + self.c_cost + self.fbw_cost[cat]) |
||||
if chunk == 0 and cat < 2: |
||||
if stage > 0: |
||||
_fa_id = self.get_id(cat, chunk, stage - 1, _cnt) |
||||
assert end_time[_fa_id] >= 0, f"{cat}, {chunk}, {stage}, {_cnt}" |
||||
_tmp = max(_tmp, end_time[_fa_id] + self.c_cost + self.fbw_cost[cat]) |
||||
_id = self.get_id(cat, chunk, stage, _cnt) |
||||
if count[stage][0] > 0: |
||||
stage_bubble[stage] += _tmp - _no_bubble |
||||
end_time[_id] = _tmp |
||||
cur_time[stage] = _tmp |
||||
mem[stage] += self.fbw_mem[cat] |
||||
# noinspection PyTypeChecker |
||||
schedule[stage].append((cat, chunk, _cnt)) |
||||
if cat == 1: |
||||
pending_w[stage].append((2, chunk, _cnt)) |
||||
count[stage][cat * 2 + chunk] += 1 |
||||
|
||||
for i in range(self.n_stage): |
||||
put(0, 0, i) |
||||
for i in range(self.n_stage - 1, -1, -1): |
||||
if i == self.n_stage - 1: |
||||
put(0, 1, i) |
||||
continue |
||||
tmp = end_time[self.get_id(0, 1, i + 1, 0)] + self.c_cost |
||||
while ( |
||||
mem[i] + self.fbw_mem[0] * (2 + i * 2) <= self.max_mem |
||||
and cur_time[i] + self.fbw_cost[0] <= tmp |
||||
and count[i][0] < self.n_micro |
||||
): |
||||
for j in range(i + 1): |
||||
put(0, 0, j) |
||||
put(0, 1, i) |
||||
iter_chunk_ = 0 |
||||
end_tmp = 0 |
||||
for i in range(self.n_stage): |
||||
if i == 0: |
||||
end_tmp = cur_time[0] + self.fbw_cost[1] |
||||
continue |
||||
tmp = end_tmp + self.c_cost |
||||
while ( |
||||
count[i][0] + count[i][1] < count[i - 1][0] + count[i - 1][1] |
||||
or count[i][1] <= count[i - 1][1] < self.n_micro |
||||
): |
||||
for j in range(self.n_stage - 1, i - 1, -1): |
||||
if count[j][iter_chunk_] < self.n_micro: |
||||
put(0, iter_chunk_, j) |
||||
iter_chunk_ = 1 - iter_chunk_ |
||||
|
||||
for _ in range(2 * self.n_micro): |
||||
# check mem before putting b |
||||
for i in range(self.n_stage): |
||||
while mem[i] + self.fbw_mem[1] > self.max_mem: |
||||
assert len(pending_w[i]) > 0 |
||||
put_w(i) |
||||
b0_ranks, b1_ranks = [], [] |
||||
for i in range(self.n_stage): |
||||
if count[i][3] >= count[i][2]: |
||||
b0_ranks.append(i) |
||||
elif i == self.n_stage - 1: |
||||
b1_ranks.append(i) |
||||
else: |
||||
fa_id = self.get_id(1, 1, i + 1, count[i][3]) |
||||
if end_time[fa_id] >= 0 or count[i][2] >= self.n_micro: |
||||
b1_ranks.append(i) |
||||
else: |
||||
b0_ranks.append(i) |
||||
b_ranks = [] |
||||
# put b1 |
||||
for i in reversed(b1_ranks): |
||||
b_ranks.append((i, 1)) |
||||
# put b0 |
||||
for i in b0_ranks: |
||||
b_ranks.append((i, 0)) |
||||
for i, _chunk_ in b_ranks: |
||||
fa_id = -1 |
||||
if _chunk_ == 1 and i < self.n_stage - 1: |
||||
fa_id = self.get_id(1, 1, i + 1, count[i][3]) |
||||
if _chunk_ == 0 and i > 0: |
||||
fa_id = self.get_id(1, 0, i - 1, count[i][2]) |
||||
while ( |
||||
len(pending_w[i]) > 0 |
||||
and fa_id >= 0 |
||||
and end_time[fa_id] + self.c_cost >= cur_time[i] + self.fbw_cost[2] |
||||
): |
||||
# fill the bubble |
||||
put_w(i) |
||||
if ( |
||||
len(pending_w[i]) > 0 |
||||
and end_time[fa_id] + self.c_cost - cur_time[i] > get_max_stage_bubble(i) - stage_bubble[i] |
||||
): |
||||
if _chunk_ == 1: |
||||
put_w(i) |
||||
elif fill_b: |
||||
put_w(i) |
||||
put(1, _chunk_, i) |
||||
|
||||
# put f |
||||
for i in range(self.n_stage): |
||||
if count[i][1] >= self.n_micro: |
||||
continue |
||||
put_item = None |
||||
if count[i][1] >= count[i][0]: |
||||
put_item = 0 |
||||
elif i == self.n_stage - 1: |
||||
put_item = 1 |
||||
else: |
||||
if end_time[self.get_id(0, 1, i + 1, count[i][1])] >= 0: |
||||
put_item = 1 |
||||
elif count[i][0] < self.n_micro: |
||||
if i == 0: |
||||
put_item = 0 |
||||
elif end_time[self.get_id(0, 0, i - 1, count[i][0])] >= 0: |
||||
put_item = 0 |
||||
if put_item is None: |
||||
continue |
||||
# check mem before putting f |
||||
while mem[i] + self.fbw_mem[0] > self.max_mem: |
||||
assert len(pending_w[i]) > 0 |
||||
put_w(i) |
||||
fa_id = -1 |
||||
if put_item == 0 and i > 0: |
||||
fa_id = self.get_id(0, 0, i - 1, count[i][0]) |
||||
if put_item == 1 and i < self.n_stage - 1: |
||||
fa_id = self.get_id(0, 1, i + 1, count[i][1]) |
||||
while ( |
||||
len(pending_w[i]) > 0 |
||||
and fa_id >= 0 |
||||
and end_time[fa_id] + self.c_cost >= cur_time[i] + self.fbw_cost[2] |
||||
): |
||||
# fill the bubble |
||||
put_w(i) |
||||
if ( |
||||
len(pending_w[i]) > 0 |
||||
and end_time[fa_id] + self.c_cost - cur_time[i] > get_max_stage_bubble(i) - stage_bubble[i] |
||||
): |
||||
if fill_f: |
||||
put_w(i) |
||||
put(0, put_item, i) |
||||
|
||||
for i in range(self.n_stage): |
||||
while len(pending_w[i]) > 0: |
||||
put_w(i) |
||||
|
||||
max_bubble = get_max_stage_bubble() |
||||
expected_time = sum(self.fbw_cost) * self.n_micro * 2 |
||||
max_bubble / expected_time |
||||
if max_approved_bubble < 0 or max_bubble < max_approved_bubble: |
||||
_schedule, _end_time, _max_bubble = self.try_v_schedule( |
||||
fill_f=fill_f, |
||||
fill_b=fill_b, |
||||
approved_bubble=stage_bubble, |
||||
) |
||||
if _max_bubble < max_bubble: |
||||
return _schedule, _end_time, _max_bubble |
||||
return schedule, end_time, max_bubble |
||||
|
||||
def print_details(self, end_time, print_scaling=1): |
||||
for stage in range(self.n_stage): |
||||
stage_str = ["."] * int(max(end_time) / print_scaling) |
||||
for _cat in range(3): |
||||
for _chunk in range(2): |
||||
for _micro in range(self.n_micro): |
||||
_id = self.get_id(_cat, _chunk, stage, _micro) |
||||
if end_time[_id] < 0: |
||||
continue |
||||
end = int(end_time[_id] / print_scaling) |
||||
start = int((end_time[_id] - self.fbw_cost[_cat]) / print_scaling) |
||||
for j in range(start, end): |
||||
if j == start or j == end - 1: |
||||
stage_str[j] = "FfBbWw"[_cat * 2 + _chunk] |
||||
elif j == start + 1: |
||||
if _micro >= 10: |
||||
stage_str[j] = str(_micro // 10) |
||||
else: |
||||
stage_str[j] = str(_micro) |
||||
elif j == start + 2 and _micro >= 10: |
||||
stage_str[j] = str(_micro % 10) |
||||
else: |
||||
stage_str[j] = "-" |
||||
_str = "" |
||||
for _c in stage_str: |
||||
_str += _c |
||||
print(_str) |
||||
|
||||
def get_v_schedule(self, only_run_time=False): |
||||
schedule, end_time, max_bubble = None, None, None |
||||
expected_time = sum(self.fbw_cost) * self.n_micro * 2 |
||||
for fill_b in [True, False]: |
||||
for fill_f in [True, False]: |
||||
_schedule, _end_time, _max_bubble = self.try_v_schedule(fill_b=fill_b, fill_f=fill_f) |
||||
if max_bubble is None or _max_bubble < max_bubble: |
||||
max_bubble = _max_bubble |
||||
schedule = _schedule |
||||
end_time = _end_time |
||||
if only_run_time: |
||||
return max_bubble + expected_time |
||||
max_bubble / (expected_time + max_bubble) |
||||
local_order = [[] for _ in range(self.n_stage)] |
||||
comm_id = {} |
||||
comm_id_counter = 0 |
||||
post_validation_time = 0 |
||||
for i in range(self.n_stage - 1, -1, -1): |
||||
pv_id = min(2 * (self.n_stage - 1 - i), self.n_micro - 1) |
||||
post_validation_time = max( |
||||
post_validation_time, end_time[self.get_id(0, 0, i, pv_id)] - self.fbw_cost[0] - self.c_cost |
||||
) |
||||
# post_validation_time = 0 |
||||
for it in ["RECV_", "SEND_", ""]: |
||||
if i == 0 and it == "SEND_": |
||||
continue |
||||
if i == self.n_stage - 1 and it == "RECV_": |
||||
continue |
||||
# stage_ = i - 1 if it == "RECV_" else i |
||||
stage_ = i |
||||
local_order[stage_].append( |
||||
ScheduledNode( |
||||
type=it + "POST_VALIDATION", |
||||
chunk=0, |
||||
stage=stage_, |
||||
minibatch=0, |
||||
start_time=post_validation_time, |
||||
completion_time=post_validation_time, |
||||
) |
||||
) |
||||
comm_id[local_order[stage_][-1]] = comm_id_counter |
||||
comm_id_counter += 1 |
||||
for i in range(self.n_stage): |
||||
for _cat_, _chunk_, _micro_ in schedule[i]: |
||||
complete_time = end_time[self.get_id(_cat_, _chunk_, i, _micro_)] |
||||
local_order[i].append( |
||||
ScheduledNode( |
||||
type="FBW"[_cat_], |
||||
chunk=_chunk_ if _cat_ == 0 else 1 - _chunk_, |
||||
stage=i, |
||||
minibatch=_micro_, |
||||
start_time=complete_time - self.fbw_cost[_cat_], |
||||
completion_time=complete_time, |
||||
) |
||||
) |
||||
if _cat_ == 2: # no communication for W |
||||
continue |
||||
cat_str = "FORWARD" if _cat_ == 0 else "BACKWARD" |
||||
|
||||
def communicate(send_recv, stage_): |
||||
# noinspection PyTypeChecker |
||||
local_order[stage_].append( |
||||
ScheduledNode( |
||||
type=send_recv + cat_str, |
||||
chunk=_chunk_ if _cat_ == 0 else 1 - _chunk_, |
||||
stage=stage_, |
||||
minibatch=_micro_, |
||||
start_time=complete_time, |
||||
completion_time=complete_time, |
||||
) |
||||
) |
||||
comm_id[local_order[stage_][-1]] = comm_id_counter |
||||
|
||||
if _chunk_ == 1 and i > 0: |
||||
communicate("SEND_", i) |
||||
communicate("RECV_", i - 1) |
||||
if _chunk_ == 0 and i < self.n_stage - 1: |
||||
communicate("SEND_", i) |
||||
communicate("RECV_", i + 1) |
||||
comm_id_counter += 1 |
||||
for rank in range(self.n_stage): |
||||
# For nodes with the same timestamp on the same stage, communication will be prioritized. |
||||
def even_breaker(x: ScheduledNode): |
||||
# Compute nodes are always delayed. |
||||
if x.type in ["F", "B", "W"]: |
||||
return comm_id_counter |
||||
# For comm nodes, order by their unique comm id |
||||
return comm_id[x] |
||||
|
||||
local_order[rank] = list(sorted(local_order[rank], key=lambda x: (x.start_time, even_breaker(x)))) |
||||
# If a recv with intersects with previous computation, reorder them so that recv |
||||
# is executed before computation and hence can be overlapped. |
||||
for i in range(len(local_order[rank])): |
||||
if ( |
||||
i > 0 |
||||
and local_order[rank][i - 1].type in {"F", "B", "W"} |
||||
and local_order[rank][i].type.startswith("RECV") |
||||
and "POST_VALIDATION" not in local_order[rank][i].type |
||||
and local_order[rank][i].start_time <= local_order[rank][i - 1].completion_time |
||||
): |
||||
local_order[rank][i], local_order[rank][i - 1] = local_order[rank][i - 1], local_order[rank][i] |
||||
|
||||
local_order_with_rollback = [[] for _ in range(self.n_stage)] |
||||
for rank in range(self.n_stage): |
||||
rollback_comm = set() |
||||
if rank > 0: |
||||
for node in local_order[rank - 1]: |
||||
if node.type == "POST_VALIDATION": |
||||
break |
||||
if node.type == "SEND_FORWARD": |
||||
assert node.chunk == 0 |
||||
rollback_comm.add(node.minibatch) |
||||
for node in local_order[rank]: |
||||
if node.type == "RECV_FORWARD" and node.chunk == 0 and node.minibatch in rollback_comm: |
||||
rollback = True |
||||
rollback_comm.remove(node.minibatch) |
||||
else: |
||||
rollback = False |
||||
local_order_with_rollback[rank].append( |
||||
ScheduledNode( |
||||
type=node.type, |
||||
chunk=node.chunk, |
||||
stage=node.stage, |
||||
minibatch=node.minibatch, |
||||
start_time=node.start_time, |
||||
completion_time=node.completion_time, |
||||
rollback=rollback, |
||||
) |
||||
) |
||||
assert len(rollback_comm) == 0 |
||||
|
||||
return local_order_with_rollback |
||||
@ -0,0 +1,958 @@
|
||||
from functools import partial |
||||
from typing import Any, Callable, Dict, Iterable, List, Optional, Union |
||||
|
||||
import torch |
||||
import torch.cuda |
||||
import torch.distributed |
||||
from torch.nn import Module, ModuleList |
||||
from torch.utils._pytree import tree_flatten, tree_map |
||||
|
||||
from colossalai.accelerator import get_accelerator |
||||
from colossalai.interface import OptimizerWrapper |
||||
from colossalai.pipeline.p2p import PipelineP2PCommunication, create_send_metadata |
||||
from colossalai.pipeline.schedule.v_schedule import ScheduledNode |
||||
from colossalai.pipeline.stage_manager import PipelineStageManager |
||||
from colossalai.pipeline.weight_grad_store import WeightGradStore |
||||
|
||||
from ._utils import ( |
||||
clone, |
||||
detach, |
||||
get_batch_size, |
||||
get_micro_batch, |
||||
merge_batch, |
||||
model_forward, |
||||
release_tensor_data, |
||||
require_grad, |
||||
retain_grad, |
||||
to_device, |
||||
) |
||||
from .base import PipelineSchedule |
||||
|
||||
AUTO_SCHEDULE_COMMUNICATION_TYPES = {"RECV_FORWARD", "RECV_BACKWARD", "SEND_FORWARD", "SEND_BACKWARD"} |
||||
|
||||
|
||||
def _wait_p2p(wait_handles: List[torch.cuda.Event]) -> None: |
||||
if wait_handles is not None: |
||||
for req in wait_handles: |
||||
req.wait() |
||||
|
||||
|
||||
class ZeroBubbleVPipeScheduler(PipelineSchedule): |
||||
def __init__( |
||||
self, |
||||
stage_manager: PipelineStageManager, |
||||
schedule: List[ScheduledNode], |
||||
num_model_chunks: int, |
||||
num_microbatch: Optional[int] = None, |
||||
microbatch_size: Optional[int] = None, |
||||
enable_metadata_cache: bool = True, |
||||
overlap_p2p: bool = True, |
||||
): |
||||
super().__init__(stage_manager) |
||||
# batch info |
||||
self.num_microbatch = num_microbatch |
||||
self.microbatch_size = microbatch_size |
||||
self.num_model_chunks = num_model_chunks |
||||
self.batch: Any |
||||
self.batch_size: int |
||||
self.last_batch_size: Optional[int] = None |
||||
self.microbatch_offset: List[int] |
||||
|
||||
self.schedules = schedule |
||||
# TODO: optim post valid |
||||
self.do_post_validation = False |
||||
|
||||
# P2PMeta cache |
||||
self.enable_metadata_cache = enable_metadata_cache |
||||
|
||||
# check send_tensor_metadata, send_grad_metadata |
||||
# pp4 as sample, we should follow this meta strategy |
||||
# send_tensor_meta(fwd) send_grad_meta(bwd) |
||||
# chunk0 | chunk1 chunk0 | chunk 1 |
||||
# stage 0 T | F F | T |
||||
# stage 1 T | T T | T |
||||
# stage 2 T | T T | T |
||||
# stage 3 F | T F | T |
||||
if stage_manager.is_first_stage(ignore_chunk=True): |
||||
self.send_tensor_metadata = [True, False] |
||||
self.send_grad_metadata = [False, True] |
||||
elif stage_manager.is_last_stage(ignore_chunk=True): |
||||
self.send_tensor_metadata = [False, True] |
||||
self.send_grad_metadata = [True, False] |
||||
else: |
||||
self.send_tensor_metadata = [True, True] |
||||
self.send_grad_metadata = [True, True] |
||||
|
||||
# meta cache buffer |
||||
self.tensor_metadata_recv = [None, None] # [chunk 0 meta, chunk 1 meta] |
||||
self.grad_metadata_recv = [None, None] |
||||
|
||||
# P2P communication |
||||
self.comm = PipelineP2PCommunication(stage_manager, overlap_p2p=overlap_p2p) |
||||
|
||||
# init communication map |
||||
self.communication_map = { |
||||
"SEND_FORWARD": self.send_forward, |
||||
"RECV_FORWARD": self.recv_forward, |
||||
"SEND_BACKWARD": self.send_backward, |
||||
"RECV_BACKWARD": self.recv_backward, |
||||
} |
||||
|
||||
# init buffer |
||||
self._free_buffers() |
||||
|
||||
def _free_buffers(self): |
||||
# free local buffer |
||||
# two dim array, first dim is the model chunk, second dim is the microbatch queue |
||||
|
||||
# x & y buffer for schedule b |
||||
self.input_tensors = [[], []] |
||||
self.output_tensors = [[], []] |
||||
|
||||
# y & dy buffer for schedule w |
||||
self.output_tensors_dw = [[], []] |
||||
self.output_tensors_grad_dw = [[], []] |
||||
|
||||
# buffer for communication |
||||
self.send_forward_buffer = [[], []] # [chunk0:[torch.Tensor], chunk1:[torch.Tensor]] |
||||
self.recv_forward_buffer = [ |
||||
[], |
||||
[], |
||||
] # [chunk0:[(torch.Tensor, wait_handle)], chunk1:[(torch.Tensor, wait_handle)]] |
||||
self.send_backward_buffer = [[], []] # [chunk0:[torch.Tensor], chunk1:[torch.Tensor]] |
||||
self.recv_backward_buffer = [ |
||||
[], |
||||
[], |
||||
] # [chunk0:[(torch.Tensor, wait_handle)], chunk1:[(torch.Tensor, wait_handle)]] |
||||
|
||||
# y buffer for local send fwd |
||||
self.local_send_forward_buffer = [] |
||||
# dy buffer for local send bwd |
||||
self.local_send_backward_buffer = [] |
||||
|
||||
# wait pp buffer |
||||
self.wait_handles = [] |
||||
|
||||
def assert_buffer_empty(self): |
||||
# assert buffer is empty at end |
||||
assert len(self.input_tensors[0]) == 0 |
||||
assert len(self.input_tensors[1]) == 0 |
||||
assert len(self.output_tensors[0]) == 0 |
||||
assert len(self.output_tensors[1]) == 0 |
||||
assert len(self.output_tensors_dw[0]) == 0 |
||||
assert len(self.output_tensors_dw[1]) == 0 |
||||
assert len(self.output_tensors_grad_dw[0]) == 0 |
||||
assert len(self.output_tensors_grad_dw[1]) == 0 |
||||
assert len(self.send_forward_buffer[0]) == 0 |
||||
assert len(self.send_forward_buffer[1]) == 0 |
||||
assert len(self.recv_forward_buffer[0]) == 0 |
||||
assert len(self.recv_forward_buffer[1]) == 0 |
||||
assert len(self.send_backward_buffer[0]) == 0 |
||||
assert len(self.send_backward_buffer[1]) == 0 |
||||
assert len(self.recv_backward_buffer[0]) == 0 |
||||
assert len(self.recv_backward_buffer[1]) == 0 |
||||
assert len(self.local_send_forward_buffer) == 0 |
||||
assert len(self.local_send_backward_buffer) == 0 |
||||
|
||||
def load_batch(self, data_iter: Iterable, device: Optional[torch.device] = None) -> None: |
||||
"""Load a batch from data iterator. |
||||
|
||||
Args: |
||||
data_iter (Iterable): Data iterator. |
||||
device (Optional[torch.device], optional): Target device. Defaults to None. |
||||
""" |
||||
batch = next(data_iter) |
||||
if device is not None: |
||||
batch = tree_map(partial(to_device, device=device), batch) |
||||
|
||||
self.microbatch_offset = [0 for _ in range(self.num_model_chunks)] |
||||
self.batch = batch |
||||
self.batch_size = get_batch_size(batch) |
||||
|
||||
if self.microbatch_size is None: |
||||
assert self.batch_size % self.num_microbatch == 0, "Batch size should divided by the number of microbatch" |
||||
self.microbatch_size = self.batch_size // self.num_microbatch |
||||
if self.num_microbatch is None: |
||||
assert self.batch_size % self.microbatch_size == 0, "Batch size should divided by the microbatch size" |
||||
self.num_microbatch = self.batch_size // self.microbatch_size |
||||
|
||||
if not self.forward_only: |
||||
assert self.last_batch_size is None or self.last_batch_size == self.batch_size |
||||
assert self.batch_size == self.microbatch_size * self.num_microbatch |
||||
|
||||
assert ( |
||||
self.num_microbatch % self.stage_manager.num_stages == 0 |
||||
), "Number of microbatch should be an integer multiple of number of pipeline parallel devices" |
||||
|
||||
if self.forward_only: |
||||
self.num_microbatch = (self.batch_size - 1) // self.microbatch_size + 1 |
||||
|
||||
self.last_batch_size = self.batch_size |
||||
|
||||
def load_micro_batch(self, model_chunk_id: int) -> Any: |
||||
"""Load a micro batch from the current batch. |
||||
|
||||
Args: |
||||
microbatch_id (int): the current model chunk idx. |
||||
|
||||
Returns: |
||||
Any: Micro batch. |
||||
""" |
||||
assert self.microbatch_offset[model_chunk_id] <= self.batch_size, "Microbatches exhausted" |
||||
micro_batch = get_micro_batch(self.batch, self.microbatch_offset[model_chunk_id], self.microbatch_size) |
||||
self.microbatch_offset[model_chunk_id] += self.microbatch_size |
||||
return tree_map(partial(to_device, device=get_accelerator().get_current_device()), micro_batch) |
||||
|
||||
def get_model_chunk_id(self, microbatch_id: int, is_forward: bool) -> int: |
||||
"""Helper method to get the model chunk ID given the iteration number. |
||||
|
||||
Args: |
||||
microbatch_id (int): the current microbatch idx |
||||
forward (bool): if is the forward process |
||||
|
||||
Returns: |
||||
int: The model chunk idx of the input microbatch_id |
||||
""" |
||||
assert ( |
||||
microbatch_id < self.num_microbatch * self.num_model_chunks |
||||
), f"microbatch_id {microbatch_id} is out of range ({self.num_microbatch * self.num_model_chunks})" |
||||
microbatch_id_in_group = microbatch_id % (self.stage_manager.num_stages * self.num_model_chunks) |
||||
model_chunk_id = microbatch_id_in_group // self.stage_manager.num_stages |
||||
if not is_forward: |
||||
# Reverse order |
||||
model_chunk_id = self.num_model_chunks - model_chunk_id - 1 |
||||
return model_chunk_id |
||||
|
||||
def recv_forward(self, model_chunk_id: int, prev_rank: int = None) -> List: |
||||
"""Copy the forward output from the previous stage in pipeline as the input tensor of this stage. |
||||
For ZBV. |
||||
|
||||
Args: |
||||
model_chunk_id (int): The current model chunk idx. |
||||
prev_rank (int, optional): The rank of the source of the tensor. |
||||
|
||||
Returns: |
||||
Any: The input tensor or input tensor list. |
||||
Any: The wait handles for the communication. |
||||
""" |
||||
with self.stage_manager.switch_model_chunk_id(model_chunk_id): |
||||
if model_chunk_id == 0: |
||||
################ |
||||
# chunk = 0 & is_first_stage |
||||
# do nothing; cause u are chunk 0 in first rank, u have no prev rank; |
||||
################# |
||||
if self.stage_manager.is_first_stage(ignore_chunk=True): |
||||
return [] |
||||
|
||||
################ |
||||
# chunk = 0 & not is_first_stage |
||||
# Recv y from PREV_rank as input |
||||
################# |
||||
else: |
||||
prev_rank = self.stage_manager.get_prev_rank() |
||||
input_tensor, wait_handles = self.comm.recv_forward( |
||||
prev_rank=prev_rank, metadata_recv=self.tensor_metadata_recv[model_chunk_id] |
||||
) |
||||
if self.enable_metadata_cache and self.tensor_metadata_recv[model_chunk_id] is None: |
||||
self.tensor_metadata_recv[model_chunk_id] = create_send_metadata(input_tensor) |
||||
self.recv_forward_buffer[model_chunk_id].append((input_tensor, wait_handles)) |
||||
return wait_handles |
||||
|
||||
else: |
||||
################ |
||||
# chunk = 1 & is_last_stage |
||||
# do nothing; cause u get y from local_send_forward_buffer in schedule f |
||||
################ |
||||
if self.stage_manager.is_last_stage(ignore_chunk=True): |
||||
# return None, [] |
||||
return [] |
||||
|
||||
################ |
||||
# chunk = 1 & not is_last_stage |
||||
# recv y from NEXT_rank as input |
||||
################ |
||||
else: |
||||
next_rank = self.stage_manager.get_next_rank() |
||||
input_tensor, wait_handles = self.comm.recv_forward( |
||||
next_rank, metadata_recv=self.tensor_metadata_recv[model_chunk_id] |
||||
) |
||||
if self.enable_metadata_cache and self.tensor_metadata_recv[model_chunk_id] is None: |
||||
self.tensor_metadata_recv[model_chunk_id] = create_send_metadata(input_tensor) |
||||
self.recv_forward_buffer[model_chunk_id].append((input_tensor, wait_handles)) |
||||
return wait_handles |
||||
|
||||
def recv_backward(self, model_chunk_id: int, next_rank: int = None) -> List: |
||||
"""Copy the gradient tensor from the next stage in pipeline as the input gradient of this stage. |
||||
For ZBV. |
||||
|
||||
Args: |
||||
model_chunk_id (int): The current model chunk idx. |
||||
next_rank (int, optional): The rank of the source of the tensor. |
||||
|
||||
Returns: |
||||
Any: The input gradient tensor or gradient tensor list. |
||||
Any: The wait handles for the communication. |
||||
""" |
||||
with self.stage_manager.switch_model_chunk_id(model_chunk_id): |
||||
if model_chunk_id == 0: |
||||
# bwd chunk0 is right V; |
||||
################ |
||||
# chunk = 0 & is_last_stage |
||||
# do nothing; Already get dy from local_send_backward_buffer in schedule b |
||||
################ |
||||
if self.stage_manager.is_last_stage(ignore_chunk=True): |
||||
return [] |
||||
|
||||
################ |
||||
# chunk = 0 & not is_last_stage |
||||
# Recv bwd from next stage; |
||||
################ |
||||
else: |
||||
next_rank = self.stage_manager.get_next_rank() |
||||
output_tensor_grad, wait_handles = self.comm.recv_backward( |
||||
next_rank, metadata_recv=self.grad_metadata_recv[model_chunk_id] |
||||
) |
||||
if self.enable_metadata_cache and self.grad_metadata_recv[model_chunk_id] is None: |
||||
self.grad_metadata_recv[model_chunk_id] = create_send_metadata(output_tensor_grad) |
||||
self.recv_backward_buffer[model_chunk_id].append((output_tensor_grad, wait_handles)) |
||||
return wait_handles |
||||
|
||||
else: |
||||
# bwd chunk1 is left V; |
||||
################ |
||||
# chunk = 1 & is_first_stage |
||||
# do nothing; get loss from local |
||||
################ |
||||
if self.stage_manager.is_first_stage(ignore_chunk=True): |
||||
return [] |
||||
|
||||
################ |
||||
# chunk = 1 & not first stage |
||||
# recv_backward recv bwd from prev stage; |
||||
################ |
||||
else: |
||||
prev_rank = self.stage_manager.get_prev_rank() |
||||
output_tensor_grad, wait_handles = self.comm.recv_backward( |
||||
next_rank=prev_rank, metadata_recv=self.grad_metadata_recv[model_chunk_id] |
||||
) |
||||
if self.enable_metadata_cache and self.grad_metadata_recv[model_chunk_id] is None: |
||||
self.grad_metadata_recv[model_chunk_id] = create_send_metadata(output_tensor_grad) |
||||
self.recv_backward_buffer[model_chunk_id].append((output_tensor_grad, wait_handles)) |
||||
return wait_handles |
||||
|
||||
def send_forward(self, model_chunk_id: int, next_rank: int = None) -> List: |
||||
"""Sends the input tensor to the next stage in pipeline. |
||||
For ZBV. |
||||
|
||||
Args: |
||||
model_chunk_id (int): The current model chunk idx. |
||||
next_rank (int, optional): The rank of the recipient of the tensor. |
||||
|
||||
Returns: |
||||
Any: The wait handles for the communication. |
||||
""" |
||||
|
||||
with self.stage_manager.switch_model_chunk_id(model_chunk_id): |
||||
if model_chunk_id == 0: |
||||
################ |
||||
# chunk = 0 && is_last_stage |
||||
# do nothing; hold y on local_send_forward_buffer |
||||
################ |
||||
if self.stage_manager.is_last_stage(ignore_chunk=True): |
||||
self.send_tensor_metadata[model_chunk_id] = not self.enable_metadata_cache |
||||
return [] |
||||
|
||||
################ |
||||
# chunk = 0 && not is_last_stage |
||||
# self.comm.send_forward send y to NEXT stage |
||||
################ |
||||
else: |
||||
next_rank = self.stage_manager.get_next_rank() |
||||
output_tensor = self.send_forward_buffer[model_chunk_id].pop(0) |
||||
send_handles = self.comm.send_forward( |
||||
output_object=output_tensor, |
||||
next_rank=next_rank, |
||||
send_metadata=self.send_tensor_metadata[model_chunk_id], |
||||
) |
||||
self.send_tensor_metadata[model_chunk_id] = not self.enable_metadata_cache |
||||
return send_handles |
||||
|
||||
else: |
||||
################ |
||||
# chunk = 1 && is_first_stage |
||||
# do nothing; Already send LOSS to local_send_backward_buffer in schedule f send part |
||||
################ |
||||
if self.stage_manager.is_first_stage(ignore_chunk=True): |
||||
self.send_tensor_metadata[model_chunk_id] = not self.enable_metadata_cache |
||||
return [] |
||||
|
||||
################ |
||||
# chunk = 1 && not is_first_stage |
||||
# self.comm.send_forward send y to PREV stage |
||||
################ |
||||
else: |
||||
prev_rank = self.stage_manager.get_prev_rank() |
||||
output_tensor = self.send_forward_buffer[model_chunk_id].pop(0) |
||||
send_handles = self.comm.send_forward( |
||||
output_tensor, prev_rank, send_metadata=self.send_tensor_metadata[model_chunk_id] |
||||
) |
||||
self.send_tensor_metadata[model_chunk_id] = not self.enable_metadata_cache |
||||
return send_handles |
||||
|
||||
def send_backward(self, model_chunk_id: int, prev_rank: int = None) -> List: |
||||
"""Sends the gradient tensor to the previous stage in pipeline. |
||||
For ZBV. |
||||
|
||||
Args: |
||||
model_chunk_id (int): The current model chunk idx. |
||||
prev_rank (int, optional): The rank of the recipient of the tensor |
||||
|
||||
Returns: |
||||
Any: The wait handles for the communication. |
||||
""" |
||||
|
||||
with self.stage_manager.switch_model_chunk_id(model_chunk_id): |
||||
if model_chunk_id == 0: |
||||
# bwd chunk0 is right V; |
||||
################ |
||||
# chunk = 0 && is_first_stage |
||||
# do nothing; cause u are the first chunk in first stage; bwd end |
||||
################ |
||||
if self.stage_manager.is_first_stage(ignore_chunk=True): |
||||
self.send_grad_metadata[model_chunk_id] = not self.enable_metadata_cache |
||||
return [] |
||||
|
||||
################ |
||||
# chunk = 0 && not is_first_stage |
||||
# Send dx to PREV stage; |
||||
################ |
||||
else: |
||||
prev_rank = self.stage_manager.get_prev_rank() |
||||
input_tensor_grad = self.send_backward_buffer[model_chunk_id].pop(0) |
||||
send_handles = self.comm.send_backward( |
||||
input_tensor_grad, prev_rank, send_metadata=self.send_grad_metadata[model_chunk_id] |
||||
) |
||||
self.send_grad_metadata[model_chunk_id] = not self.enable_metadata_cache |
||||
return send_handles |
||||
|
||||
# bwd chunk1 is left V; |
||||
else: |
||||
################ |
||||
# chunk = 1 && is_last_stage |
||||
# do nothing; Already send input_tensor_grad to local_send_bwd_buffer in schedule b; |
||||
################ |
||||
if self.stage_manager.is_last_stage(ignore_chunk=True): |
||||
self.send_grad_metadata[model_chunk_id] = not self.enable_metadata_cache |
||||
return [] |
||||
|
||||
################ |
||||
# chunk = 1 && not is_last_stage |
||||
# Send dx to NEXT stage; |
||||
################ |
||||
else: |
||||
next_rank = self.stage_manager.get_next_rank() |
||||
input_tensor_grad = self.send_backward_buffer[model_chunk_id].pop(0) |
||||
send_handles = self.comm.send_backward( |
||||
input_tensor_grad, next_rank, send_metadata=self.send_grad_metadata[model_chunk_id] |
||||
) |
||||
self.send_grad_metadata[model_chunk_id] = not self.enable_metadata_cache |
||||
return send_handles |
||||
|
||||
def forward_step( |
||||
self, |
||||
model_chunk: Union[ModuleList, Module], |
||||
model_chunk_id: int, |
||||
micro_batch: Optional[dict], |
||||
input_obj: Optional[dict], |
||||
criterion: Callable, |
||||
accum_loss: Optional[torch.Tensor] = None, |
||||
outputs: Optional[List[Any]] = None, |
||||
) -> Union[torch.Tensor, dict]: |
||||
"""Forward one step of the pipeline |
||||
Args: |
||||
model_chunk (ModuleList or Module): Model Chunk to be run; |
||||
model_chunk_id (int): The current model chunk idx; |
||||
input_obj (Optional[dict]): x; |
||||
criterion (Callable): loss function; |
||||
accum_loss (Optional[torch.Tensor], optional): Accumulated loss. Defaults to None. |
||||
outputs (Optional[List[Any]], optional): List to store the output of the last stage (final output). Defaults to None. |
||||
|
||||
Returns: |
||||
Union[torch.Tensor, dict]: The intermediate output (dict) of the current stage. If it is the last stage, the output is the loss (Tensor). |
||||
""" |
||||
# Load input ids, attention mask and labels |
||||
# for the first stage, input_obj is None; So,we use micro_batch as input_obj |
||||
# for other stages, input_obj is the output of the previous/next stage containing hidden_states etc. |
||||
# Only attention_mask from micro_batch is used |
||||
with self.stage_manager.switch_model_chunk_id(model_chunk_id): |
||||
# fwd calculate |
||||
internal_inputs = {} if input_obj is None else input_obj |
||||
internal_inputs["stage_index"] = self.stage_manager.stage_indices[model_chunk_id] |
||||
output_obj = model_forward(model_chunk, micro_batch, internal_inputs) |
||||
# last layer in model |
||||
if model_chunk_id == 1 and self.stage_manager.is_first_stage(ignore_chunk=True): |
||||
loss = criterion(output_obj, micro_batch) / self.num_microbatch |
||||
if accum_loss is not None: |
||||
accum_loss.add_(loss.detach()) |
||||
if outputs is not None: |
||||
outputs.append(tree_map(detach, output_obj)) |
||||
return loss |
||||
else: |
||||
return output_obj |
||||
|
||||
def backward_b_step( |
||||
self, |
||||
model_chunk: Union[ModuleList, Module], |
||||
model_chunk_id: int, |
||||
optimizer: OptimizerWrapper, |
||||
# micro_batch: Optional[dict], |
||||
input_obj: Optional[dict], |
||||
output_obj: Union[dict, torch.Tensor], |
||||
output_obj_grad: Optional[dict], |
||||
) -> Optional[dict]: |
||||
"""Backward dx step of the pipeline; we calculate "dx = w*dy" here; |
||||
|
||||
Args: |
||||
model_chunk (ModuleList or Module): Model Chunk to be run; |
||||
model_chunk_id (int): The current model chunk idx; |
||||
optimizer (OptimizerWrapper): Optimizer to update the model |
||||
input_obj (Optional[Tuple(dict)]): x. (microbatch, input_obj) |
||||
output_obj (Union[dict, torch.Tensor]): y. |
||||
output_obj_grad (dict): dy. |
||||
|
||||
Returns: |
||||
Optional[dict]: dx. |
||||
""" |
||||
# calculate bwd b step ; only dx = w*dy; |
||||
|
||||
# Retain the grad on the input_obj. No need retain_grad microbatch |
||||
if input_obj is not None: |
||||
tree_map(retain_grad, input_obj) |
||||
|
||||
# x, y, dy list for backward_by_grad; Type: list[tensor]; |
||||
input_obj_ = [] |
||||
output_obj_ = [] |
||||
output_obj_grad_ = [] |
||||
|
||||
# For chunk 0 stage 0, use micro_batch as input_obj_; and we don't have to cal microbatch dx. |
||||
|
||||
# For loss backward; output_obj is loss; output_obj_grad should be None |
||||
if model_chunk_id == 1 and self.stage_manager.is_first_stage(ignore_chunk=True): |
||||
assert output_obj_grad is None |
||||
input_obj_, _ = tree_flatten(input_obj) |
||||
output_obj_.append(output_obj) # LOSS |
||||
output_obj_grad_.append(output_obj_grad) # None |
||||
|
||||
# For other chunk stage, use input_obj as input_obj_; |
||||
else: |
||||
input_obj_, _ = tree_flatten(input_obj) |
||||
output_obj_, _ = tree_flatten(output_obj) # y |
||||
output_obj_grad_, _ = tree_flatten(output_obj_grad) # dy |
||||
|
||||
# filter item which is not torch.Tensor |
||||
input_obj_ = [v for v in input_obj_ if isinstance(v, torch.Tensor) or v is None] |
||||
output_obj_ = [v for v in output_obj_ if isinstance(v, torch.Tensor) or v is None] |
||||
output_obj_grad_ = [v for v in output_obj_grad_ if isinstance(v, torch.Tensor) or v is None] |
||||
|
||||
try: |
||||
ctx = optimizer.no_sync() |
||||
except AttributeError: |
||||
ctx = model_chunk.no_sync() |
||||
with ctx: |
||||
optimizer.backward_by_grad( |
||||
tensor=output_obj_, |
||||
grad=output_obj_grad_, |
||||
# inputs=input_obj_, |
||||
retain_graph=False, |
||||
) |
||||
# Format output_obj_grad |
||||
input_obj_grad = dict() |
||||
if model_chunk_id == 0 and self.stage_manager.is_first_stage(ignore_chunk=True): |
||||
pass |
||||
else: |
||||
for k, v in input_obj.items(): |
||||
if isinstance(v, torch.Tensor) and v.grad is not None: |
||||
input_obj_grad[k] = v.grad |
||||
return input_obj_grad |
||||
|
||||
def backward_w_step( |
||||
self, |
||||
model_chunk: Union[ModuleList, Module], |
||||
model_chunk_id: int, |
||||
optimizer: OptimizerWrapper, |
||||
output_obj: Union[dict, torch.Tensor], |
||||
output_obj_grad: Optional[dict], |
||||
): |
||||
"""Backward dw step of the pipeline; we calculate "dw = x*dy" here; |
||||
|
||||
Args: |
||||
model_chunk (ModuleList or Module): Model Chunk to be run; |
||||
model_chunk_id (int): The current model chunk idx; |
||||
optimizer (OptimizerWrapper): Optimizer to update the model |
||||
output_obj (Union[dict, torch.Tensor]): y. |
||||
output_obj_grad (dict): dy. |
||||
|
||||
Returns: |
||||
Nothing need to return; we only calculate dw then update w; |
||||
""" |
||||
# calculate bwd w step ; only dw = x*dy; |
||||
|
||||
# y, dy list for w backward_by_grad; Type: list[tensor]; |
||||
output_obj_ = [] |
||||
output_obj_grad_ = [] |
||||
|
||||
if model_chunk_id == 1 and self.stage_manager.is_first_stage(ignore_chunk=True): |
||||
# loss backward; output_obj is loss; |
||||
output_obj_.append(output_obj) # LOSS |
||||
output_obj_grad_.append(None) # None |
||||
else: |
||||
output_obj_, _ = tree_flatten(output_obj) # y |
||||
output_obj_grad_, _ = tree_flatten(output_obj_grad) # dy |
||||
|
||||
# filter item which is not torch.Tensor |
||||
output_obj_ = [v for v in output_obj_ if isinstance(v, torch.Tensor) or v is None] |
||||
output_obj_grad_ = [v for v in output_obj_grad_ if isinstance(v, torch.Tensor) or v is None] |
||||
|
||||
optimizer.backward_by_grad( |
||||
tensor=output_obj_, |
||||
grad=output_obj_grad_, |
||||
inputs=list(model_chunk.parameters()), |
||||
retain_graph=False, |
||||
) |
||||
|
||||
def schedule_f( |
||||
self, |
||||
scheduled_node, |
||||
model_chunk: torch.nn.ModuleList, |
||||
model_chunk_id: int, |
||||
criterion: Callable, |
||||
accum_loss: Optional[torch.Tensor] = None, |
||||
outputs: Optional[List[Any]] = None, |
||||
): |
||||
"""A complete forward schedule; Include recv fwd --> cal fwd --> send fwd; |
||||
|
||||
Args: |
||||
scheduled_node: |
||||
model_chunk (ModuleList or Module): Model Chunk to be run; |
||||
model_chunk_id (int): The current model chunk idx; |
||||
criterion (Callable): loss function; |
||||
accum_loss (Optional[torch.Tensor], optional): Accumulated loss. Defaults to None. |
||||
outputs (Optional[List[Any]], optional): List to store the output of the last stage (final output). Defaults to None. |
||||
|
||||
Returns: |
||||
Nothing. |
||||
""" |
||||
micro_batch = self.load_micro_batch(model_chunk_id=model_chunk_id) |
||||
# Step1: recv fwd |
||||
if model_chunk_id == 0: |
||||
# is first stage; get input from microbatch |
||||
if self.stage_manager.is_first_stage(ignore_chunk=True): |
||||
input_obj = None # (tensor, wait_handle) |
||||
else: |
||||
input_obj = self.recv_forward_buffer[model_chunk_id].pop(0) |
||||
for h in input_obj[1]: |
||||
h.wait() |
||||
input_obj = input_obj[0] |
||||
else: |
||||
# is last stage; recv from local |
||||
if self.stage_manager.is_last_stage(ignore_chunk=True): |
||||
input_obj = self.local_send_forward_buffer.pop(0) |
||||
# not last stage; recv from next |
||||
else: |
||||
input_obj = self.recv_forward_buffer[model_chunk_id].pop(0) |
||||
for h in input_obj[1]: |
||||
h.wait() |
||||
input_obj = input_obj[0] |
||||
# Here, let input_obj.requires_grad_() |
||||
# if input_obj is not None: |
||||
if not isinstance(input_obj, torch.Tensor): |
||||
tree_map(require_grad, input_obj) |
||||
|
||||
# Also requires_grad_ for micro_batch in stage 0 chunk 0 fwd, |
||||
# tree_map(torch.Tensor.requires_grad_, micro_batch) |
||||
|
||||
# Step2: fwd step |
||||
output_obj = self.forward_step( |
||||
model_chunk=model_chunk, |
||||
model_chunk_id=model_chunk_id, |
||||
micro_batch=micro_batch, |
||||
input_obj=input_obj, |
||||
criterion=criterion, |
||||
accum_loss=accum_loss, |
||||
outputs=outputs, |
||||
) |
||||
|
||||
# Step3: |
||||
# 3-1:detach output; detach output for send fwd; |
||||
if model_chunk_id == 1 and self.stage_manager.is_first_stage(ignore_chunk=True): |
||||
# We should not detach bwd LOSS |
||||
pass |
||||
else: |
||||
# detach output |
||||
detached_output_obj = tree_map(detach, output_obj) |
||||
# 3-2 clone detached_output_obj |
||||
detached_output_obj = tree_map(clone, detached_output_obj) |
||||
|
||||
# 3-3 release cloned output.data; release_tensor_data output for bwd b & w; (do not detach output) |
||||
if model_chunk_id == 1 and self.stage_manager.is_first_stage(ignore_chunk=True): |
||||
# We should not release_tensor_data bwd LOSS |
||||
pass |
||||
else: |
||||
# release_tensor_data output |
||||
tree_map(release_tensor_data, output_obj) |
||||
|
||||
# add input and output object for backward b |
||||
self.input_tensors[model_chunk_id].append(input_obj) |
||||
|
||||
# for bwd b&w, we only need the graph(grad_fn) of output_obj |
||||
# Do not release_tensor_data loss, release_tensor_data other output_obj; |
||||
if model_chunk_id == 1 and self.stage_manager.is_first_stage(ignore_chunk=True): |
||||
self.output_tensors[model_chunk_id].append(output_obj) |
||||
else: |
||||
self.output_tensors[model_chunk_id].append(output_obj) |
||||
|
||||
# add output to send_fwd_buffer |
||||
if model_chunk_id == 0: # chunk 0 |
||||
# is last stage; send to local_send_forward_buffer |
||||
if self.stage_manager.is_last_stage(ignore_chunk=True): |
||||
self.local_send_forward_buffer.append(detached_output_obj) |
||||
else: |
||||
self.send_forward_buffer[model_chunk_id].append(detached_output_obj) |
||||
else: # chunk 1 |
||||
# is first stage; end of fwd; do nothing |
||||
if self.stage_manager.is_first_stage(ignore_chunk=True): |
||||
pass |
||||
else: |
||||
self.send_forward_buffer[model_chunk_id].append(detached_output_obj) |
||||
|
||||
def schedule_b( |
||||
self, |
||||
scheduled_node, |
||||
model_chunk: Union[ModuleList, Module], |
||||
model_chunk_id: int, |
||||
optimizer: OptimizerWrapper, |
||||
): |
||||
"""A complete backward b schedule; Include recv bwd --> cal bwd step --> send bwd; |
||||
|
||||
Args: |
||||
scheduled_node: |
||||
model_chunk (ModuleList or Module): Model Chunk to be run; |
||||
model_chunk_id (int): The current model chunk idx; |
||||
Returns: |
||||
Nothing. |
||||
""" |
||||
# Step1: recv bwd |
||||
if model_chunk_id == 0: |
||||
# chunk0 is last stage; recv output_grad from local_send_backward_buffer |
||||
if self.stage_manager.is_last_stage(ignore_chunk=True): |
||||
output_tensor_grad = self.local_send_backward_buffer.pop(0) |
||||
# chunk0 not last stage; recv output_grad from recv_backward_buffer |
||||
else: |
||||
output_tensor_grad = self.recv_backward_buffer[model_chunk_id].pop(0) |
||||
for h in output_tensor_grad[1]: |
||||
h.wait() |
||||
output_tensor_grad = output_tensor_grad[0] |
||||
else: |
||||
# chunk1, is first stage; recv LOSS from local send bwd buffer |
||||
if self.stage_manager.is_first_stage(ignore_chunk=True): |
||||
output_tensor_grad = None |
||||
# chunk1, not first stage; recv output_grad from recv_backward_buffer |
||||
else: |
||||
output_tensor_grad = self.recv_backward_buffer[model_chunk_id].pop(0) |
||||
for h in output_tensor_grad[1]: |
||||
h.wait() |
||||
output_tensor_grad = output_tensor_grad[0] |
||||
|
||||
# get input and output object from buffer; |
||||
input_obj = self.input_tensors[model_chunk_id].pop(0) |
||||
output_obj = self.output_tensors[model_chunk_id].pop(0) |
||||
|
||||
input_object_grad = self.backward_b_step( |
||||
model_chunk=model_chunk, |
||||
model_chunk_id=model_chunk_id, |
||||
optimizer=optimizer, |
||||
input_obj=input_obj, |
||||
output_obj=output_obj, |
||||
output_obj_grad=output_tensor_grad, |
||||
) |
||||
|
||||
# Step3: send bwd |
||||
if model_chunk_id == 0: |
||||
# do nothing; end of bwd; |
||||
if self.stage_manager.is_first_stage(ignore_chunk=True): |
||||
pass |
||||
# save input_object_grad to send_backward_buffer |
||||
else: |
||||
self.send_backward_buffer[model_chunk_id].append(input_object_grad) |
||||
else: |
||||
# send to local_send_backward_buffer |
||||
if self.stage_manager.is_last_stage(ignore_chunk=True): |
||||
self.local_send_backward_buffer.append(input_object_grad) |
||||
# send to next |
||||
else: |
||||
self.send_backward_buffer[model_chunk_id].append(input_object_grad) |
||||
WeightGradStore.flush(chunk=model_chunk_id) |
||||
|
||||
def schedule_w( |
||||
self, |
||||
scheduled_node, |
||||
model_chunk: Union[ModuleList, Module], |
||||
model_chunk_id: int, |
||||
optimizer: OptimizerWrapper, |
||||
): |
||||
"""A complete backward w schedule; Include get y & dy from buffer --> cal bwd w step(cal dw & update w); |
||||
|
||||
Args: |
||||
scheduled_node: |
||||
model_chunk (ModuleList or Module): Model Chunk to be run; |
||||
model_chunk_id (int): The current model chunk idx; |
||||
Returns: |
||||
Nothing. |
||||
""" |
||||
WeightGradStore.pop(chunk=model_chunk_id) |
||||
|
||||
def run_forward_only( |
||||
self, |
||||
model_chunk: Union[ModuleList, Module], |
||||
data_iter: Iterable, |
||||
criterion: Callable[..., Any], |
||||
return_loss: bool = False, |
||||
return_outputs: bool = False, |
||||
) -> Dict: |
||||
assert self.forward_only |
||||
|
||||
# prepare batch |
||||
self.load_batch(data_iter) |
||||
|
||||
# prepare accum loss & output |
||||
accum_loss = None |
||||
|
||||
# reset accum loss at fwd end; |
||||
if return_loss and self.stage_manager.is_first_stage(ignore_chunk=True): |
||||
accum_loss = torch.scalar_tensor(0, device=get_accelerator().get_current_device()) |
||||
|
||||
outputs = [] if return_outputs and self.stage_manager.is_first_stage(ignore_chunk=True) else None |
||||
|
||||
# while we still have schedules_node in self.schedules |
||||
for it in range(len(self.schedules)): |
||||
scheduled_node = self.schedules[it] |
||||
|
||||
if scheduled_node.type in {"RECV_FORWARD", "SEND_FORWARD"}: |
||||
# communication |
||||
communication_func = self.communication_map[scheduled_node.type] |
||||
communication_func(scheduled_node.chunk) |
||||
if scheduled_node.type == "F": |
||||
self.schedule_f( |
||||
scheduled_node=scheduled_node, |
||||
model_chunk=model_chunk, |
||||
model_chunk_id=scheduled_node.chunk, |
||||
criterion=criterion, |
||||
accum_loss=accum_loss, |
||||
outputs=outputs, |
||||
) |
||||
# return loss & output |
||||
if outputs is not None: |
||||
outputs = merge_batch(outputs) |
||||
return {"loss": accum_loss, "outputs": outputs} |
||||
|
||||
def run_forward_backward( |
||||
self, |
||||
model_chunk: Union[ModuleList, Module], |
||||
data_iter: Iterable, |
||||
criterion: Callable[..., Any], |
||||
optimizer: Optional[OptimizerWrapper] = None, |
||||
return_loss: bool = False, |
||||
return_outputs: bool = False, |
||||
) -> Dict: |
||||
""" |
||||
Runs Zerobubble schedule, with communication between pipeline stages. |
||||
""" |
||||
# prepare batch |
||||
self.load_batch(data_iter) |
||||
|
||||
# prepare accum loss & output |
||||
accum_loss = None |
||||
|
||||
# reset accum loss at fwd end; |
||||
if return_loss and self.stage_manager.is_first_stage(ignore_chunk=True): |
||||
accum_loss = torch.scalar_tensor(0, device=get_accelerator().get_current_device()) |
||||
|
||||
outputs = [] if return_outputs and self.stage_manager.is_first_stage(ignore_chunk=True) else None |
||||
|
||||
# while we still have schedules_node in self.schedules |
||||
schedule = self.schedules[self.stage_manager.stage] # get schedule by stage (rank) |
||||
for it in range(len(schedule)): |
||||
scheduled_node = schedule[it] |
||||
if scheduled_node.type in AUTO_SCHEDULE_COMMUNICATION_TYPES: |
||||
# communication |
||||
communication_func = self.communication_map[scheduled_node.type] |
||||
wait_handle = communication_func(scheduled_node.chunk) |
||||
# We wait recv handle in fwd step and bwd step. Here only need to wait for send handle |
||||
if scheduled_node.type in {"SEND_FORWARD", "SEND_BACKWARD"}: |
||||
self.wait_handles.append(wait_handle) |
||||
elif scheduled_node.type == "F": |
||||
self.schedule_f( |
||||
scheduled_node=scheduled_node, |
||||
model_chunk=model_chunk, |
||||
model_chunk_id=scheduled_node.chunk, |
||||
criterion=criterion, |
||||
accum_loss=accum_loss, |
||||
outputs=outputs, |
||||
) |
||||
elif scheduled_node.type == "B": |
||||
self.schedule_b( |
||||
scheduled_node=scheduled_node, |
||||
model_chunk=model_chunk, |
||||
model_chunk_id=scheduled_node.chunk, |
||||
optimizer=optimizer, |
||||
) |
||||
elif scheduled_node.type == "W": |
||||
self.schedule_w( |
||||
scheduled_node=scheduled_node, |
||||
model_chunk=model_chunk, |
||||
model_chunk_id=scheduled_node.chunk, |
||||
optimizer=optimizer, |
||||
) |
||||
# wait here to ensure all communication is done |
||||
for h in self.wait_handles: |
||||
for hh in h: |
||||
hh.wait() |
||||
# return loss & output |
||||
if outputs is not None: |
||||
outputs = merge_batch(outputs) |
||||
return {"loss": accum_loss, "outputs": outputs} |
||||
|
||||
def forward_backward_step( |
||||
self, |
||||
model_chunk: Union[ModuleList, Module], |
||||
data_iter: Iterable, |
||||
criterion: Callable[..., Any], |
||||
optimizer: Optional[OptimizerWrapper] = None, |
||||
return_loss: bool = False, |
||||
return_outputs: bool = False, |
||||
) -> dict: |
||||
""" |
||||
Args: |
||||
model_chunk (ModuleList or Module): Model Chunk to be trained. Original interleaved uses a module list whereas shardformer uses entire model + layer specification |
||||
data_iter (Iterable): Data iterator. |
||||
criterion (Callable[[Any, Any], Tensor]): Criterion to be used. It should take two arguments: model outputs and inputs, and returns loss tensor. |
||||
optimizer (OptimizerWrapper, optional): Optimizer to be used. Can be None when only forward is executed. Defaults to None. |
||||
return_loss (bool, optional): Whether to return loss. Defaults to False. Whether to return loss. |
||||
return_outputs (bool, optional): Whether to return model outputs. Defaults to False. Whether to return model outputs. |
||||
|
||||
Returns: |
||||
dict: A dict with keys: 'loss' and 'outputs'. |
||||
""" |
||||
self.forward_only = not torch.is_grad_enabled() |
||||
if optimizer is None: |
||||
assert self.forward_only, "Optimizer should be passed when doing backward." |
||||
|
||||
if self.forward_only: |
||||
result = self.run_forward_only(model_chunk, data_iter, criterion, return_loss, return_outputs) |
||||
else: |
||||
result = self.run_forward_backward( |
||||
model_chunk, data_iter, criterion, optimizer, return_loss, return_outputs |
||||
) |
||||
|
||||
self.assert_buffer_empty() |
||||
return result |
||||
@ -0,0 +1,32 @@
|
||||
import queue |
||||
|
||||
|
||||
class WeightGradStore: |
||||
|
||||
cache = [] |
||||
weight_grad_queue = [queue.Queue(), queue.Queue()] |
||||
|
||||
@classmethod |
||||
def put(cls, total_input, grad_output, weight, func): |
||||
# func(total_input, grad_output, weight.main_grad) |
||||
cls.cache.append((total_input, grad_output, weight, func)) |
||||
|
||||
@classmethod |
||||
def flush(cls, chunk=0): |
||||
cls.weight_grad_queue[chunk].put(cls.cache) |
||||
cls.cache = [] |
||||
|
||||
@classmethod |
||||
def pop(cls, chunk=0): |
||||
# print(f"chunk id {chunk} queue size {cls.weight_grad_queue[chunk].qsize()}") |
||||
if cls.weight_grad_queue[chunk].qsize() > 0: |
||||
stored_grads = cls.weight_grad_queue[chunk].get() |
||||
for total_input, grad_output, weight, func in stored_grads: |
||||
if weight.grad is not None: |
||||
func(total_input, grad_output, weight.grad) |
||||
# for first bwd; weight.grad is None, assign grad_weight to weight.grad |
||||
else: |
||||
grad_weight = func(total_input, grad_output) |
||||
weight.grad = grad_weight |
||||
else: |
||||
raise Exception("Pop empty queue.") |
||||
File diff suppressed because it is too large
Load Diff
Loading…
Reference in new issue