1"use strict";(self.webpackChunkdemo=self.webpackChunkdemo||[]).push([[2675],{6999:(e,t,n)=>{n.d(t,{Cl:()=>o,Dx:()=>s,Pc:()=>l,aE:()=>p,e_:()=>u,iz:()=>r,nT:()=>d});var i=n(7294),a=n(398);n(814);function o(e){return i.createElement("div",{className:"docstring-container"},e.children)}function l(e){return i.createElement("div",{className:"signature"},"(",e.children,")")}function r(e){return i.createElement("div",{class:"divider"},i.createElement("span",{class:"divider-text"},e.name))}function p(e){return i.createElement("div",null,i.createElement(r,{name:"Parameters"}),i.createElement(a.D,null,e.children))}function d(e){return i.createElement("div",null,i.createElement(r,{name:"Returns"}),i.createElement(a.D,null,`${e.name}: ${e.desc}`))}function s(e){return i.createElement("div",{className:"title-container"},i.createElement("div",{className:"title-module"},i.createElement("h5",null,e.type),"\xa0 ",i.createElement("h3",null,e.name)),i.createElement("div",{className:"title-source"},"<",i.createElement("a",{href:e.source,className:"title-source"},"source"),">"))}function u(e){return i.createElement("div",null,i.createElement(r,{name:"Example"}),i.createElement(a.D,null,e.code))}},8076:(e,t,n)=>{n.r(t),n.d(t,{assets:()=>d,contentTitle:()=>r,default:()=>m,frontMatter:()=>l,metadata:()=>p,toc:()=>s});var i=n(7462),a=(n(7294),n(3905)),o=n(6999);const l={},r="ZeroBubble Pipeline Parallelism",p={unversionedId:"features/zerobubble_pipeline_parallelism",id:"features/zerobubble_pipeline_parallelism",title:"ZeroBubble Pipeline Parallelism",description:"Author: Junwen Duan, Hongxin Liu",source:"@site/i18n/en/docusaurus-plugin-content-docs/current/features/zerobubble_pipeline_parallelism.md",sourceDirName:"features",slug:"/features/zerobubble_pipeline_parallelism",permalink:"/docs/features/zerobubble_pipeline_parallelism",draft:!1,editUrl:"https://github.com/facebook/docusaurus/tree/main/packages/create-docusaurus/templates/shared/docs/features/zerobubble_pipeline_parallelism.md",tags:[],version:"current",frontMatter:{}},d={},s=[{value:"Introduction",id:"introduction",level:2},{value:"Hands-On Practice",id:"hands-on-practice",level:2},{value:"step 1. Import libraries",id:"step-1-import-libraries",level:3},{value:"step 2. Initialize Distributed Environment and Parallism Group",id:"step-2-initialize-distributed-environment-and-parallism-group",level:3},{value:"step 3. Initialize Module, Optimizer, and Pipeline Schedule",id:"step-3-initialize-module-optimizer-and-pipeline-schedule",level:3},{value:"step 4. Initialize Module, Optimizer, and Pipeline Schedul",id:"step-4-initialize-module-optimizer-and-pipeline-schedul",level:3},{value:"step 5.Init Booster",id:"step-5init-booster",level:3},{value:"step 6.Train Your Model",id:"step-6train-your-model",level:3},{value:"Advanced Practice",id:"advanced-practice",level:2},{value:"1.Use MetaCache with ZeroBubble",id:"1use-metacache-with-zerobubble",level:3},{value:"2.HybridParallel with ZeroBubble",id:"2hybridparallel-with-zerobubble",level:3},{value:"3.Fine-tuning Scheduler parameters",id:"3fine-tuning-scheduler-parameters",level:3},{value:"Model compatibility",id:"model-compatibility",level:2},{value:"API Reference",id:"api-reference",level:2}],u={toc:s},c="wrapper";
1function m(e){let{components:t,...n}=e;return(0,a.kt)(c,(0,i.Z)({},u,n,{components:t,mdxType:"MDXLayout"}),(0,a.kt)("h1",{id:"zerobubble-pipeline-parallelism"},"ZeroBubble Pipeline Parallelism"),(0,a.kt)("p",null,"Author: ",(0,a.kt)("a",{parentName:"p",href:"https://github.com/duanjunwen"},"Junwen Duan"),", ",(0,a.kt)("a",{parentName:"p",href:"https://github.com/ver217"},"Hongxin Liu")),(0,a.kt)("p",null,(0,a.kt)("strong",{parentName:"p"},"Related Paper")),(0,a.kt)("ul",null,(0,a.kt)("li",{parentName:"ul"},(0,a.kt)("a",{parentName:"li",href:"https://arxiv.org/abs/2401.10241"},"Zero Bubble Pipeline Parallelism"))),(0,a.kt)("h2",{id:"introduction"},"Introduction"),(0,a.kt)("p",null,"ZeroBubble (V Schedule):\nCrucially, splitting B into two stages (also known as an activation gradient and a weight gradient) and a scheme like 1F1B1W can further reduce the bubble compared to the 1F1B scheme in earlier work."),(0,a.kt)("h2",{id:"hands-on-practice"},"Hands-On Practice"),(0,a.kt)("p",null,"We now demonstrate how to use ZeroBubble with booster API with 4 GPUs."),(0,a.kt)("h3",{id:"step-1-import-libraries"},"step 1. Import libraries"),(0,a.kt)("pre",null,(0,a.kt)("code",{parentName:"pre",className:"language-python"},"import torch\nimport torch.distributed as dist\nimport torch.nn as nn\nfrom torch.testing import assert_close\nfrom transformers.models.llama.configuration_llama import LlamaConfig\nfrom transformers.models.llama.modeling_llama import LlamaModel\n\nimport colossalai\nfrom colossalai.booster.booster import Booster\nfrom colossalai.booster.plugin.moe_hybrid_parallel_plugin import HybridParallelPlugin\nfrom colossalai.pipeline.schedule.zero_bubble_pp import ZeroBubbleVPipeScheduler\n")),(0,a.kt)("h3",{id:"step-2-initialize-distributed-environment-and-parallism-group"},"step 2. Initialize Distributed Environment and Parallism Group"),(0,a.kt)("pre",null,(0,a.kt)("code",{parentName:"pre",className:"language-python"},'colossalai.launch(rank=rank, world_size=world_size, host="localhost", port=port, backend="nccl")\n')),(0,a.kt)("h3",{id:"step-3-initialize-module-optimizer-and-pipeline-schedule"},"step 3. Initialize Module, Optimizer, and Pipeline Schedule"),(0,a.kt)("p",null,"Build our model and Optimizer. We created a Llama with 8 Decoder-Layer. Then, inite the PipelineGraph and Pipeline schedule by get_v_schedule() function."),(0,a.kt)("pre",null,(0,a.kt)("code",{parentName:"pre",className:"language-python"},'# Global Param\nNUM_BATCH = 8\nNUM_TOK_PER_BATCH = 4\nNUM_LAYERS = 8\nHIDDEN_SIZE_PER_HEAD = 4\nNUM_HEADS = 4\n# Init Llama from huggingface\nconfiguration = LlamaConfig(\n hidden_size=HIDDEN_SIZE_PER_HEAD * NUM_HEADS,\n intermediate_size=HIDDEN_SIZE_PER_HEAD * NUM_HEADS * 2,\n num_hidden_layers=NUM_LAYERS,\n num_attention_heads=NUM_HEADS,\n num_key_value_heads=NUM_HEADS,\n attn_implementation="flash_attention_2",\n)\nmodel = LlamaModel(configuration).cuda()\noptimizer = torch.optim.Adam(torch_model.parameters(), lr=1)\n')),(0,a.kt)("h3",{id:"step-4-initialize-module-optimizer-and-pipeline-schedul"},"step 4. Initialize Module, Optimizer, and Pipeline Schedul"),(0,a.kt)("p",null,"Then, we need to create the PipelineGraph and PipelineSchedule using the get_v_schedule() function. We need to initialise the PipelineGraph with the following parameters.\nx_cost represents the runtime consumed by operation x of each model chunk.\nx_mem represents the amount of memory consumed by the operation x of each model chunk.\nThese parameters are estimated and filled in before the pipeline starts. In fact, better results can be obtained based on the runtime and memory cost during the real computation of the model.\nIn the following example, we assume that the computation times for the model's forward, reverse B, and reverse W are 1, 1, 1, respectively, and the p2p communication time is 1."),(0,a.kt)("pre",null,(0,a.kt)("code",{parentName:"pre",className:"language-python"},"# Init schedule\nh, a, s = config.hidden_size, config.num_attention_heads, 1024\nmem_f = 34 * h + 5 * a * s\nmem_w = -32 * h\nmem_b = -mem_w - mem_f\ngraph = PipelineGraph(\n n_stage=pp_size,\n n_micro=num_microbatches,\n f_cost=1,\n b_cost=1,\n w_cost=1,\n c_cost=1,\n f_mem=mem_f,\n b_mem=mem_b,\n w_mem=mem_w,\n)\nzbv_schedule = graph.get_v_schedule()\n")),(0,a.kt)("h3",{id:"step-5init-booster"},"step 5.Init Booster"),(0,a.kt)("p",null,'Pass pp_style="zbv" when initialising the Plugin to use the ZeroBubble Pipeline.'),(0,a.kt)("pre",null,(0,a.kt)("code",{parentName:"pre",className:"language-python"},'plugin = HybridParallelPlugin(\n pp_size=4,\n num_microbatches=4,\n tp_size=1,\n sp_size=1,\n zero_stage=1,\n initial_scale=1,\n find_unused_parameters=True,\n pp_style="zbv",\n scheduler_nodes=zbv_schedule,\n num_model_chunks=2,\n)\n\ndp_size = plugin.dp_size\nbooster = Booster(plugin=plugin)\n')),(0,a.kt)("h3",{id:"step-6train-your-model"},"step 6.Train Your Model"),(0,a.kt)("pre",null,(0,a.kt)("code",{parentName:"pre",className:"language-python"},'steps = 10\nfor step in range(steps):\n input_embeddings = torch.rand(\n NUM_BATCH, NUM_TOK_PER_BATCH, HIDDEN_SIZE_PER_HEAD * NUM_HEADS, requires_grad=True\n ).cuda()\n dist.all_reduce(\n input_embeddings, group=plugin.pp_group\n )\n data_iter = iter([{"input
1s_embeds": input_embeddings}])\n output = booster.execute_pipeline(\n data_iter,\n model,\n lambda x, y: x.last_hidden_state.mean(),\n optimizer,\n return_loss=True,\n return_outputs=True,\n )\n optimizer.step()\n optimizer.zero_grad()\n')),(0,a.kt)("h2",{id:"advanced-practice"},"Advanced Practice"),(0,a.kt)("p",null,"In ColossalAI, you can get better training performance by using MetaCache and HybridParallel with ZeroBubble."),(0,a.kt)("h3",{id:"1use-metacache-with-zerobubble"},"1.Use MetaCache with ZeroBubble"),(0,a.kt)("p",null,'Pass "enable_metadata_cache=True" when initialising the Plugin to use the Meta Cache with ZeroBubble Pipeline.'),(0,a.kt)("pre",null,(0,a.kt)("code",{parentName:"pre",className:"language-python"},'plugin = HybridParallelPlugin(\n pp_size=2,\n num_microbatches=4,\n tp_size=2,\n sp_size=2,\n zero_stage=1,\n initial_scale=1,\n enable_metadata_cache=True,\n find_unused_parameters=True,\n pp_style="zbv",\n scheduler_nodes=zbv_schedule,\n num_model_chunks=2,\n)\n')),(0,a.kt)("h3",{id:"2hybridparallel-with-zerobubble"},"2.HybridParallel with ZeroBubble"),(0,a.kt)("p",null,"Pass pp_size, tp_size, sp_size when initialising the Plugin to use the HybridParallel with ZeroBubble Pipeline."),(0,a.kt)("pre",null,(0,a.kt)("code",{parentName:"pre",className:"language-python"},'plugin = HybridParallelPlugin(\n pp_size=2,\n num_microbatches=2,\n tp_size=2,\n sp_size=2,\n zero_stage=1,\n initial_scale=1,\n find_unused_parameters=True,\n pp_style="zbv",\n scheduler_nodes=zbv_schedule,\n num_model_chunks=2,\n)\n')),(0,a.kt)("p",null,"Performance Benchmark"),(0,a.kt)("table",null,(0,a.kt)("tr",null,(0,a.kt)("th",{nowrap:"nowrap"},"HybridParallel Strategy"),(0,a.kt)("th",{nowrap:"nowrap",align:"center"},"Pipeline Parallel"),(0,a.kt)("th",{nowrap:"nowrap",align:"center"},"Sequence Parallel + Pipeline Parallel"),(0,a.kt)("th",{nowrap:"nowrap",align:"center"},"Data Parallel + Pipeline Parallel")),(0,a.kt)("tr",null,(0,a.kt)("td",{nowrap:"nowrap",align:"center",title:"1F1B"},"With 1F1B"),(0,a.kt)("td",{nowrap:"nowrap",align:"center"},"15.27 samples/sec"),(0,a.kt)("td",{nowrap:"nowrap",align:"center"},"17.22 samples/sec"),(0,a.kt)("td",{nowrap:"nowrap",align:"center"},"14.06 samples/sec")),(0,a.kt)("tr",null,(0,a.kt)("td",{nowrap:"nowrap",align:"center",title:"Zero Bubble"},"With Zero Bubble"),(0,a.kt)("td",{nowrap:"nowrap",align:"center"},"17.36 samples/sec"),(0,a.kt)("td",{nowrap:"nowrap",align:"center"},"18.38 samples/sec"),(0,a.kt)("td",{nowrap:"nowrap",align:"center"},"14.44 samples/sec")),(0,a.kt)("tr",null,(0,a.kt)("td",{colspan:"39"}))),(0,a.kt)("h3",{id:"3fine-tuning-scheduler-parameters"},"3.Fine-tuning Scheduler parameters"),(0,a.kt)("pre",null,(0,a.kt)("code",{parentName:"pre",className:"language-python"},"")),(0,a.kt)("h2",{id:"model-compatibility"},"Model compatibility"),(0,a.kt)("table",null,(0,a.kt)("tr",null,(0,a.kt)("th",{nowrap:"nowrap"},"Shardformer/Model"),(0,a.kt)("th",{nowrap:"nowrap",align:"center"},"Bert"),(0,a.kt)("th",{nowrap:"nowrap",align:"center"},"Blip2"),(0,a.kt)("th",{nowrap:"nowrap",align:"center"},"Bloom"),(0,a.kt)("th",{nowrap:"nowrap",align:"center"},"Chatglm2"),(0,a.kt)("th",{nowrap:"nowrap",align:"center"},"Command"),(0,a.kt)("th",{nowrap:"nowrap",align:"center"},"Deepseek"),(0,a.kt)("th",{nowrap:"nowrap",align:"center"},"Falcon"),(0,a.kt)("th",{nowrap:"nowrap",align:"center"},"GPT2"),(0,a.kt)("th",{nowrap:"nowrap",align:"center"},"Gptj"),(0,a.kt)("th",{nowrap:"nowrap",align:"center"},"Llama"),(0,a.kt)("th",{nowrap:"nowrap",align:"center"},"Mistral"),(0,a.kt)("th",{nowrap:"nowrap",align:"center"},"Opt"),(0,a.kt)("th",{nowrap:"nowrap",align:"center"},"Qwen2"),(0,a.kt)("th",{nowrap:"nowrap",align:"center"},"Sam"),(0,a.kt)("th",{nowrap:"nowrap",align:"center"},"T5"),(0,a.kt)("th",{nowrap:"nowrap",align:"center"},"Vit"),(0,a.kt)("th",{nowrap:"nowrap",align:"center"},"Whisper")),(0,a.kt)("tr",null,(0,a.kt)("td",{nowrap:"nowrap",align:"center",title:"ZeroBubble"},"ZeroBubble"),(0,a.kt)("td",{nowrap:"nowrap",align:"center"},"\u2714\ufe0f"),(0,a.kt)("td",{nowrap:"nowrap",align:"center"},"\u2714\ufe0f"),(0,a.kt)("td",{nowrap:"nowrap",align:"center"},"\u2714\ufe0f"),(0,a.kt)("td",{nowrap:"nowrap",align:"center"},"\u2714\ufe0f"),(0,a.kt)("td",{nowrap:"nowrap",align:"center"},"\u2714\ufe0f"),(0,a.kt)("td",{nowrap:"nowrap",align:"center"},"\u2714\ufe0f"),(0,a.kt)("td",{nowrap:"nowrap",align:"center"},"\u2714\ufe0f"),(0,a.kt)("td",{nowrap:"nowrap",align:"center"},"\u2714\ufe0f"),(0,a.kt)("td",{nowrap:"nowrap",align:"center"},"\u2714\ufe0f"),(0,a.kt)("td",{nowrap:"nowrap",align:"center"},"\u2714\ufe0f"),(0,a.kt)("td",{nowrap:"nowrap",align:"center"},"\u2714\ufe0f"),(0,a.kt)("td",{nowrap:"nowrap",align:"center"},"\u2714\ufe0f"),(0,a.kt)("td",{nowrap:"nowrap",align:"center"},"\u2714\ufe0f"),(0,a.kt)("td",{nowrap:"nowrap",align:"center"},"\u2714\ufe0f"),(0,a.kt)("td",{nowrap:"nowrap",align:"center"},"\u2714\ufe0f"),(0,a.kt)("td",{nowrap:"nowrap",align:"center"},"\u2714\ufe0f")),(0,a.kt)("tr",null,(0,a.kt)("td",{colspan:"39"}))),(0,a.kt)("h2",{id:"api-reference"},"API Reference"),(0,a.kt)(o.Cl,{mdxType:"DocStringContainer"},(0,a.kt)("div",null,(0,a.kt)(o.Dx,{type:"class",name:"colossalai.pipeline.ZeroBubbleVPipeScheduler",source:"https://github.com/hpcaitech/ColossalAI/blob/main/colossalai/pipeline/schedule/zero_bubble_pp.py#L40",mdxType:"Title"}),(0,a.kt)(o.Pc,{mdxType:"Signature"},"stage_manager: PipelineStageManager, schedule: typing.List[colossalai.pipeline.sc
1hedule.v_schedule.ScheduledNode], num_model_chunks: int, num_microbatch: typing.Optional[int] = None, microbatch_size: typing.Optional[int] = None, enable_metadata_cache: bool = True, overlap_p2p: bool = True"),(0,a.kt)(o.aE,{mdxType:"Parameters"},"- **stage_manager** (PipelineStageManager) -- If using pipeline parallelism, it's necessary to specify a pipeline stage manager for inter-process communication in pipeline parallelism. Defaults to None, which means not using pipeline parallelism.\n- **schedule** (List[ScheduledNode]) -- Schedule for ZeroBubbleVPipe.\n- **num_model_chunks** (int) -- The number of model chunk in a device.\n- **num_microbatch** (Optional[int]) -- The number of microbatch.\n- **microbatch_size** (Optional[int]) -- The size per microbatch.\n- **enable_metadata_cache** (bool) -- whether to enable metadata cache to acclerate communication.\n- **overlap_p2p** (bool) -- whether to use overlap_p2p.")),(0,a.kt)("div",null,(0,a.kt)(o.iz,{name:"Description",mdxType:"Divider"}),(0,a.kt)("p",null,"ZeroBubbleVPipeScheduler")),(0,a.kt)(o.Cl,{mdxType:"DocStringContainer"},(0,a.kt)("div",null,(0,a.kt)(o.Dx,{type:"function",name:"backward_b_step",source:"https://github.com/hpcaitech/ColossalAI/blob/main/colossalai/pipeline/schedule/zero_bubble_pp.py#L516",mdxType:"Title"}),(0,a.kt)(o.Pc,{mdxType:"Signature"},"model_chunk: typing.Union[torch.nn.modules.container.ModuleList, torch.nn.modules.module.Module], model_chunk_id: int, optimizer: OptimizerWrapper, input_obj: typing.Optional[dict], output_obj: typing.Union[dict, torch.Tensor], output_obj_grad: typing.Optional[dict]"),(0,a.kt)(o.aE,{mdxType:"Parameters"},"- **model_chunk** (ModuleList or Module) -- Model Chunk to be run;\n- **model_chunk_id** (int) -- The current model chunk idx;\n- **optimizer** (OptimizerWrapper) -- Optimizer to update the model\n- **input_obj** (Optional[Tuple(dict)]) -- x. (microbatch, input_obj)\n- **output_obj** (Union[dict, torch.Tensor]) -- y.\n- **output_obj_grad** (dict) -- dy."),(0,a.kt)(o.nT,{name:"Optional[dict]",desc:"dx.",mdxType:"Returns"})),(0,a.kt)("div",null,(0,a.kt)(o.iz,{name:"Description",mdxType:"Divider"}),'Backward dx step of the pipeline; we calculate "dx = w*dy" here;')),(0,a.kt)(o.Cl,{mdxType:"DocStringContainer"},(0,a.kt)("div",null,(0,a.kt)(o.Dx,{type:"function",name:"backward_w_step",source:"https://github.com/hpcaitech/ColossalAI/blob/main/colossalai/pipeline/schedule/zero_bubble_pp.py#L591",mdxType:"Title"}),(0,a.kt)(o.Pc,{mdxType:"Signature"},"model_chunk: typing.Union[torch.nn.modules.container.ModuleList, torch.nn.modules.module.Module], model_chunk_id: int, optimizer: OptimizerWrapper, output_obj: typing.Union[dict, torch.Tensor], output_obj_grad: typing.Optional[dict]"),(0,a.kt)(o.aE,{mdxType:"Parameters"},"- **model_chunk** (ModuleList or Module) -- Model Chunk to be run;\n- **model_chunk_id** (int) -- The current model chunk idx;\n- **optimizer** (OptimizerWrapper) -- Optimizer to update the model\n- **output_obj** (Union[dict, torch.Tensor]) -- y.\n- **output_obj_grad** (dict) -- dy."),(0,a.kt)(o.nT,{name:"",desc:"Nothing need to return; we only calculate dw then update w;",mdxType:"Returns"})),(0,a.kt)("div",null,(0,a.kt)(o.iz,{name:"Description",mdxType:"Divider"}),'Backward dw step of the pipeline; we calculate "dw = x*dy" here;')),(0,a.kt)(o.Cl,{mdxType:"DocStringContainer"},(0,a.kt)("div",null,(0,a.kt)(o.Dx,{type:"function",name:"forward_backward_step",source:"https://github.com/hpcaitech/ColossalAI/blob/main/colossalai/pipeline/schedule/zero_bubble_pp.py#L938",mdxType:"Title"}),(0,a.kt)(o.Pc,{mdxType:"Signature"},"model_chunk: typing.Union[torch.nn.modules.container.ModuleList, torch.nn.modules.module.Module], data_iter: typing.Iterable, criterion: typing.Callable[..., typing.Any], optimizer: typing.Optional[colossalai.interface.optimizer.OptimizerWrapper] = None, return_loss: bool = False, return_outputs: bool = False"),(0,a.kt)(o.aE,{mdxType:"Parameters"},"- **model_chunk** (ModuleList or Module) -- Model Chunk to be trained. Original interleaved uses a module list whereas shardformer uses entire model + layer specification\n- **d
1ata_iter** (Iterable) -- Data iterator.\n- **criterion** (Callable[[Any, Any], Tensor]) -- Criterion to be used. It should take two arguments: model outputs and inputs, and returns loss tensor.\n- **optimizer** (OptimizerWrapper, optional) -- Optimizer to be used. Can be None when only forward is executed. Defaults to None.\n- **return_loss** (bool, optional) -- Whether to return loss. Defaults to False. Whether to return loss.\n- **return_outputs** (bool, optional) -- Whether to return model outputs. Defaults to False. Whether to return model outputs."),(0,a.kt)(o.nT,{name:"dict",desc:"A dict with keys: 'loss' and 'outputs'.",mdxType:"Returns"})),(0,a.kt)("div",null,(0,a.kt)(o.iz,{name:"Description",mdxType:"Divider"}))),(0,a.kt)(o.Cl,{mdxType:"DocStringContainer"},(0,a.kt)("div",null,(0,a.kt)(o.Dx,{type:"function",name:"forward_step",source:"https://github.com/hpcaitech/ColossalAI/blob/main/colossalai/pipeline/schedule/zero_bubble_pp.py#L474",mdxType:"Title"}),(0,a.kt)(o.Pc,{mdxType:"Signature"},"model_chunk: typing.Union[torch.nn.modules.container.ModuleList, torch.nn.modules.module.Module], model_chunk_id: int, micro_batch: typing.Optional[dict], input_obj: typing.Optional[dict], criterion: typing.Callable, accum_loss: typing.Optional[torch.Tensor] = None, outputs: typing.Optional[typing.List[typing.Any]] = None"),(0,a.kt)(o.aE,{mdxType:"Parameters"},"- **model_chunk** (ModuleList or Module) -- Model Chunk to be run;\n- **model_chunk_id** (int) -- The current model chunk idx;\n- **input_obj** (Optional[dict]) -- x;\n- **criterion** (Callable) -- loss function;\n- **accum_loss** (Optional[torch.Tensor], optional) -- Accumulated loss. Defaults to None.\n- **outputs** (Optional[List[Any]], optional) -- List to store the output of the last stage (final output). Defaults to None."),(0,a.kt)(o.nT,{name:"Union[torch.Tensor, dict]",desc:"The intermediate output (dict) of the current stage. If it is the last stage, the output is the loss (Tensor).",mdxType:"Returns"})),(0,a.kt)("div",null,(0,a.kt)(o.iz,{name:"Description",mdxType:"Divider"}),"Forward one step of the pipeline")),(0,a.kt)(o.Cl,{mdxType:"DocStringContainer"},(0,a.kt)("div",null,(0,a.kt)(o.Dx,{type:"function",name:"get_model_chunk_id",source:"https://github.com/hpcaitech/ColossalAI/blob/main/colossalai/pipeline/schedule/zero_bubble_pp.py#L219",mdxType:"Title"}),(0,a.kt)(o.Pc,{mdxType:"Signature"},"microbatch_id: int, is_forward: bool"),(0,a.kt)(o.aE,{mdxType:"Parameters"},"- **microbatch_id** (int) -- the current microbatch idx\n- **forward** (bool) -- if is the forward process"),(0,a.kt)(o.nT,{name:"int",desc:"The model chunk idx of the input microbatch_id",mdxType:"Returns"})),(0,a.kt)("div",null,(0,a.kt)(o.iz,{name:"Description",mdxType:"Divider"}),"Helper method to get the model chunk ID given the iteration number.")),(0,a.kt)(o.Cl,{mdxType:"DocStringContainer"},(0,a.kt)("div",null,(0,a.kt)(o.Dx,{type:"function",name:"load_batch",source:"https://github.com/hpcaitech/ColossalAI/blob/main/colossalai/pipeline/schedule/zero_bubble_pp.py#L170",mdxType:"Title"}),(0,a.kt)(o.Pc,{mdxType:"Signature"},"data_iter: typing.Iterable, device: typing.Optional[torch.device] = None"),(0,a.kt)(o.aE,{mdxType:"Parameters"},"- **data_iter** (Iterable) -- Data iterator.\n- **device** (Optional[torch.device], optional) -- Target device. Defaults to None.")),(0,a.kt)("div",null,(0,a.kt)(o.iz,{name:"Description",mdxType:"Divider"}),"Load a batch from data iterator.")),(0,a.kt)(o.Cl,{mdxType:"DocStringContainer"},(0,a.kt)("div",null,(0,a.kt)(o.Dx,{type:"function",name:"load_micro_batch",source:"https://github.com/hpcaitech/ColossalAI/blob/main/colossalai/pipeline/schedule/zero_bubble_pp.py#L205",mdxType:"Title"}),(0,a.kt)(o.Pc,{mdxType:"Signature"},"model_chunk_id: int"),(0,a.kt)(o.aE,{mdxType:"Parameters"},"- **microbatch_id** (int) -- the current model chunk idx."),(0,a.kt)(o.nT,{name:"Any",desc:"Micro batch.",mdxType:"Returns"})),(0,a.kt)("div",null,(0,a.kt)(o.iz,{name:"Description",mdxType:"Divider"}),"Load a micro batch from the current batch.")),(0,a.kt)(o.Cl,{mdxType:"DocStringContainer"},(0,a.kt)("div",null,(0,a.kt)(o.Dx,{type:"function",name:"recv_backward",source:"https://github.com/hpcaitech/ColossalAI/blob/main/colossalai/pipeline/schedule/zero_bubble_pp.py#L297",mdxType:"Title"}),(0,a.kt)(o.Pc,{mdxType:"Signature"},"model_chunk_id: int, next_rank: int = None"),(0,a.kt)(o.aE,{mdxType:"Parameters"},"- **model_chunk_id** (int) -- The current model chunk idx.\n- **next_rank** (int, optional) -- The rank of the source of the tensor."),(0,a.kt)(o.nT,{name:"Any",desc:"The input gradient tensor or gradient tensor list.\nAny: The wait handles for the communication.",mdxType:"Returns"})),(0,a.kt)("div",null,(0,a.kt)(o.iz,{name:"Description",mdxType:"Divider"}),"Copy the gradient tensor from the next stage in pipeline as the input gradient of this stage. For ZBV.")),(0,a.kt)(o.Cl,{mdxType:"DocStringContainer"},(0,a.kt)("div",null,(0,a.kt)(o.Dx,{type:"function",name:"recv_forward",source:"https://github.com/hpcaitech/ColossalAI/blob/main/colossalai/pipeline/schedule/zero_bubble_pp.py#L239",mdxType:"Title"}),(0,a.kt)(o.Pc,{mdxType:"Signature"},"model_chunk_id: int, prev_rank: int = None"),(0,a.kt)(o.aE,{mdxType:"Parameters"},"- **model_chunk_id** (int) -- The current model chunk idx.\n- **prev_rank** (int, optional) -- The rank of the source of the tensor."),(0,a.kt)(o.nT,{name:"Any",desc:"The input tensor or input tensor list.\nAny: The wait handles for the communication.",mdxType:"Returns"})),(0,a.kt)("div",null,(0,a.kt)(o.iz,{name:"Description",mdxType:"Divider"}),"Copy the forward output from the previous stage in pipeline as the input tensor of this stage. For ZBV.")),(0,a.kt)(o.Cl,{mdxType:"DocStringContainer"},(0,a.kt)("div",null,(0,a.kt)(o.Dx,{type:"function",name:"run_forward_backward",source:"https://github.com/hpcaitech/ColossalAI/blob/main/colossalai/pipeline/schedule/zero_bubble_pp.py#L871",mdxType:"Title"}),(0,a.kt)(o.Pc,{mdxType:"Signature"},"model_chunk: typing.Union[torch.nn.modules.container.ModuleList, torch.nn.modules.module.Module], data_iter: typing.Iterable, criterion: typing.Callable[..., typing.Any], optimizer: typing.Optional[colossalai.interface.optimizer.OptimizerWrapper] = None, return_loss: bool = False, return_outputs: bool = False")),(0,a.kt)("div",null,(0,a.kt)(o.iz,{name:"Description",mdxType:"Divider"}),(0,a.kt)("p",null,"Runs Zerobubble schedule, with communication between pipeline stages."))),(0,a.kt)(o.Cl,{mdxType:"DocStringContainer"},(0,a.kt)("div",null,(0,a.kt)(o.Dx,{type:"function",name:"schedule_b",source:"https://github.com/hpcaitech/ColossalAI/blob/main/colossalai/pipeline/schedule/zero_bubble_pp.py#L741",mdxType:"Title"}),(0,a.kt)(o.Pc,{mdxType:"Signature"},"scheduled_node, model_chunk: typing.Union[torch.nn.modules.container.ModuleList, torch.nn.modules.module.Module], model_chunk_id: int, optimizer: OptimizerWrapper"),(0,a.kt)(o.aE,{mdxType:"Parameters"},"scheduled_node --\n- **model_chunk** (ModuleList or Module) -- Model Chunk to be run;\n- **model_chunk_id** (int) -- The current model chunk idx;"),(0,a.kt)(o.nT,{name:"",desc:"Nothing.",mdxType:"Returns"})),(0,a.kt)("div",null,(0,a.kt)(o.iz,{name:"Description",mdxType:"Divider"}),"A complete backward b schedule; Include recv bwd --\x3e cal bwd step --\x3e send bwd;")),(0,a.kt)(o.Cl,{mdxType:"DocStringContainer"},(0,a.kt)("div",null,(0,a.kt)(o.Dx,{type:"function",name:"schedule_f",source:"https://github.com/hpcaitech/ColossalAI/blob/main/colossalai/pipeline/schedule/zero_bubble_pp.py#L636",mdxType:"Title"}),(0,a.kt)(o.Pc,{mdxType:"Signature"},"scheduled_node, model_chunk: ModuleList, model_chunk_id: int, criterion: typing.Callable, accum_loss: typing.Optional[torch.Tensor] = None, outputs: typing.Optional[typing.List[typing.Any]] = None"),(0,a.kt)(o.aE,{mdxType:"Parameters"},"scheduled_node --\n- **model_chunk** (ModuleList or Module) -- Model Chunk to be run;\n- **model_chunk_id** (int) -- The current model chunk idx;\n- **criterion** (Callable) -- loss function;\n- **accum_loss** (Optional[torch.Tensor], optional) -- Accumulated loss. Defaults to None.\n- **outputs** (Optional[List[Any]], optional) -- List to store the output of the last stage (final output). Defaults to None."),(0,a.kt)(o.nT,{name:"",desc:"Nothing.",mdxType:"Returns"})),(0,a.kt)("div",null,(0,a.kt)(o.iz,{name:"Description",mdxType:"Divider"}),"A complete forward schedule; Include recv fwd --\x3e cal fwd --\x3e send fwd;")),(0,a.kt)(o.Cl,{mdxType:"DocStringContainer"},(0,a.kt)("div",null,(0,a.kt)(o.Dx,{type:"function",name:"schedule_w",source:"https://github.com/hpcaitech/ColossalAI/blob/main/colossalai/pipeline/schedule/zero_bubble_pp.py#L809",mdxType:"Title"}),(0,a.kt)(o.Pc,{mdxType:"Signature"},"scheduled_node, model_chunk: typing.Union[torch.nn.modules.container.ModuleList, torch.nn.modules.module.Module], model_chunk_id: int, optimizer: OptimizerWrapper"),(0,a.kt)(o.aE,{mdxType:"Parameters"},"scheduled_node --\n- **model_chunk** (ModuleList or Module) -- Model Chunk to be run;\n- **model_chunk_id** (int) -- The current model chunk idx;"),(0,a.kt)(o.nT,{name:"",desc:"Nothing.",mdxType:"Returns"})),(0,a.kt)("div",null,(0,a.kt)(o.iz,{name:"Description",mdxType:"Divider"}
1),"A complete backward w schedule; Include get y & dy from buffer --\x3e cal bwd w step(cal dw & update w);")),(0,a.kt)(o.Cl,{mdxType:"DocStringContainer"},(0,a.kt)("div",null,(0,a.kt)(o.Dx,{type:"function",name:"send_backward",source:"https://github.com/hpcaitech/ColossalAI/blob/main/colossalai/pipeline/schedule/zero_bubble_pp.py#L415",mdxType:"Title"}),(0,a.kt)(o.Pc,{mdxType:"Signature"},"model_chunk_id: int, prev_rank: int = None"),(0,a.kt)(o.aE,{mdxType:"Parameters"},"- **model_chunk_id** (int) -- The current model chunk idx.\n- **prev_rank** (int, optional) -- The rank of the recipient of the tensor"),(0,a.kt)(o.nT,{name:"Any",desc:"The wait handles for the communication.",mdxType:"Returns"})),(0,a.kt)("div",null,(0,a.kt)(o.iz,{name:"Description",mdxType:"Divider"}),"Sends the gradient tensor to the previous stage in pipeline. For ZBV.")),(0,a.kt)(o.Cl,{mdxType:"DocStringContainer"},(0,a.kt)("div",null,(0,a.kt)(o.Dx,{type:"function",name:"send_forward",source:"https://github.com/hpcaitech/ColossalAI/blob/main/colossalai/pipeline/schedule/zero_bubble_pp.py#L356",mdxType:"Title"}),(0,a.kt)(o.Pc,{mdxType:"Signature"},"model_chunk_id: int, next_rank: int = None"),(0,a.kt)(o.aE,{mdxType:"Parameters"},"- **model_chunk_id** (int) -- The current model chunk idx.\n- **next_rank** (int, optional) -- The rank of the recipient of the tensor."),(0,a.kt)(o.nT,{name:"Any",desc:"The wait handles for the communication.",mdxType:"Returns"})),(0,a.kt)("div",null,(0,a.kt)(o.iz,{name:"Description",mdxType:"Divider"}),"Sends the input tensor to the next stage in pipeline. For ZBV."))))}m.isMDXComponent=!0}}]);
Line numbers count LF bytes from the start of the resource, as the search results do. Vendor segments are library code the classifier recognised; they are stored but not indexed. Bytes are shown as Latin1 characters, one per byte.