espnet2.speechlm.model.speechlm.parallel_utils.pipeline.build_pipeline

source

espnet2.speechlm.model.speechlm.parallel_utils.pipeline.build_pipeline(model: Module | ModuleList, parallel_dims: ParallelDims, titan_config: Dict[str, Any], n_microbatches: int) → Tuple[Any, bool]

Build PipelineStage(s) and PP schedule for stage model(s).

For single-stage schedules (e.g. 1F1B), model is a single nn.Module. For multi-stage schedules (e.g. Interleaved1F1B), model is an nn.ModuleList of model chunks (one per virtual stage on this rank).

• Parameters:
  • model – Single PP stage model, or nn.ModuleList of model chunks for interleaved schedules. Each model must have stage_idx, num_virtual_stages, pp_rank, pp_degree, and is_last_stage attributes.
  • parallel_dims – TorchTitan ParallelDims with PP mesh built.
  • titan_config – Config dict. Reads pp_schedule (default “1F1B”).
  • n_microbatches – Number of microbatches per optimizer step.
• Returns:(schedule, pp_has_last_stage) tuple.