06_昇腾流水线优化技术详解

昇腾流水线优化技术详解

1. pipeline多级流水

Pipeline多级流水的核心思想是将训练步骤拆分为多个连续的阶段（Stage），通过异步调度，使不同批次的处理过程在流水线上重叠执行。这种流水线并行机制有效掩盖了阶段间的通信与计算耗时，从而显著提升了训练吞吐量与硬件利用率。

如下图所示，一个训练步骤依次经过 Stage 1、2、3。当流水线稳定运行时，在某一时刻（如红色箭头标识）：Batch 0 正执行其 Stage 3，与此同时，Batch 1 正执行其 Stage 2，而 Batch 2 正执行其 Stage 1。这三个不同批次的阶段在同一时间点并行运行，后一批次的前期阶段被前一批次的后期阶段所“掩盖”，从而实现系统吞吐量的提升。

因为NPU和GPU硬件差异较大，所以两者在多级流水的底层实现逻辑上会有所差异。接下来将对NPU和GPU侧任务。

2. Rec SDK多级流水

本系统采用参数服务器（Parameter Server）架构，结合NPU多流并行技术，实现大规模稀疏特征的高效训练。通过计算与通信重叠、动态换入换出等优化策略，显著提升训练吞吐量。NPU中对于单个batch数据的处理流程如下：

从图中可以看到单个bath数据的数据训练流程可被切分成五个细粒度的阶段，代码对应如下：

NPU侧的多流通过Torchrec实现了EmbCacheTrainPipelineSparseDist，因涉及cpu和npu的交互，分为了5个stage。如下图所示为10个step之间，不同batch的流水线并行时不同stage之间的交互情况。

data H2D：dataloader返回batch拷贝到NPU，使用memcpy stream实现异步；
input_dist：data分桶并执行all2all，使用data_dist stream实现异步；
compute swap: 依据当前batch和已经在NPU上的key计算需要换入换出的key，使用线程池实现异步；
swap：涉及换入换出多个操作，npu emb gather/cpu emb update/cpu emb lookup/npu emb update；
fwd/bwd：包含emb查询，dense的前反向更新，emb的反向更新，使用default stream实现异步；

当前该方案已在gitcode上开源，代码实现可参考gitcode.com：

用户端接口调用

from torchrec_embcache.distributed.train_pipeline import EmbcacheTrainPipelineSparseDist

model = 
optimizer = 
dataloader_iter = 

EmbCacheTrainPipelineSparseDist(model, optimizer, cpu_device=cpu_device, npu_device="npu:0")

while True:
    try:
        pipeline.progress(dataloader_iter)

    except StopIteration:
        break

内部实现：

调用fill_pipeline，在真正的持续输出开始之前，需要先将前几个 batch 推进到各个 stage 中，形成“满载”状态。

fill_pipeline 方法：流水线初始化填充

首先，检查流水线是否已满（4个批次），避免重复填充

def fill_pipeline(self, dataloader_iter: Iterator[In]) -> None:
    # pipeline is already filled
    if len(self.batches) >= 4:
        return

batch_0的阶段性处理

初始化流水线模块，设置前向传播管道

self._init_pipelined_modules(self.batches[0], self.contexts[0], EmbCachePipelinedForward)

等待分布式所有阶段的稀疏数据同步完成，并将结果填充到训练管道上下文中。执行后输入分布处理，通过do_post_input_dist过程等待AlltoAll完成，并进行分桶去重的操作。

self.wait_sparse_data_dist(self.contexts[0])
self.do_post_input_dist(self.contexts[0])

同步等待目前批次batch_0的数据就绪，预先确定嵌入层的交换策略。

with record_function("## wait_for_batch ##"):
    _wait_for_batch(cast(In, self.batches[0]), self._data_dist_stream)

开始并等待交换信息计算（确定哪些嵌入需要换入换出）

self.start_compute_swap_info(self.contexts[0])  # 分布式训练中管理嵌入层的缓存交换策略
self.wait_and_get_swap_info(self.contexts[0]) # 等待并处理缓存交换信息

在host侧通过异步操作执行嵌入查找，执行异步操作从缓存中恢复需要的嵌入向量。

self.host_embedding_lookup_async(self.contexts[0])  # 
self.do_restore_async(self.contexts[0])

记录可以执行换出操作的事件，异步执行换出操作，将NPU上不需要的嵌入向量和优化器状态换出到CPU内存。

self.contexts[0].event_can_swapout.record(self._default_stream)
self.swap_out(self.contexts[0])

swap_out是异步操作。与此同时，为了最大化硬件利用率和实现计算重叠，流水线立即转入Batch 1的预处理阶段。Batch 1将执行数据分布、后处理等操作，并开始异步计算交换信息，从而与Batch 0的换出操作形成时间上的重叠，构建出高效的流水线执行梯度。

batch_1的阶段性处理

在batch 0执行换出操作的同时，启动batch 1的预处理以实现计算重叠，形成流水线梯度。

# batch i+1
if not self.enqueue_batch(dataloader_iter):
    return
self.start_sparse_data_dist(self.batches[1], self.contexts[1])
_fuse_input_dist_splits(self.contexts[1])
self.wait_sparse_data_dist(self.contexts[1])
self.do_post_input_dist(self.contexts[1])
with record_function("## wait_for_batch##"):
    _wait_for_batch(cast(In, self.batches[1]), self._data_dist_stream)
self.start_compute_swap_info(self.contexts[1])

start_compute_swap_info是异步启动操作，不阻塞当前流。由于计算交换信息需要时间，此时可以并行处理Batch 2的数据分布，避免等待Batch 1的计算完成，充分利用计算资源

batch_2和batch_3的阶段性处理

每个后续批次比前一个少执行一些阶段，构建阶梯式流水线；batch_2完成数据准备之后，立即启动batch_3的数据操作，仅对batch_3启动数据分布。

# batch i+2
if not self.enqueue_batch(dataloader_iter):
    return
self.start_sparse_data_dist(self.batches[2], self.contexts[2])
_fuse_input_dist_splits(self.contexts[2])
self.wait_sparse_data_dist(self.contexts[2])
self.do_post_input_dist(self.contexts[2])
with record_function("## wait_for_batch##"):
    _wait_for_batch(cast(In, self.batches[2]), self._data_dist_stream)

# batch i+3, 仅启动数据分布，提前隐藏I/O延迟
if not self.enqueue_batch(dataloader_iter):
   return
self.start_sparse_data_dist(self.batches[3], self.contexts[3])
_fuse_input_dist_splits(self.contexts[3])

start_sparse_data_dist是异步操作，由于数据分布需要较长时间，提前启动可以隐藏延迟。

并行Batch的预加载机制

预加载更多批次，只执行最基本的数据分布，确保流水线永不"饿死"，减少后续等待时间。

# 额外预取多个batch
for i in range(self.local_unique_parallel_batch_num):
    if not self.enqueue_batch(dataloader_iter):
        return
    self.start_sparse_data_dist(self.batches[4 + i], self.contexts[4 + i])

`progress` 方法：流水线推进执行

首先进行流水线初始化准备，步数统计，模型检查，填充流水线，设置上下文，梯度清零

self._global_steps += 1
if not self._model_attached:
    self.attach(self._model)

self.fill_pipeline(dataloader_iter)
self._set_module_context(self.contexts[0])
if self._model.training:
    with record_function("## zero_grad ##"):
        self._optimizer.zero_grad()

流水线末端数据同步

确保流水线末端的batch_3数据准备就绪

# wait batch_ip2
if len(self.batches) >= 4:
    self.wait_sparse_data_dist(self.contexts[3])
    with record_function("## wait_for_batch ##"):
        _wait_for_batch(cast(In, self.batches[3]), self._data_dist_stream)

if len(self.batches) >= 5:
    _fuse_input_dist_splits(self.contexts[4])

加载新batch数据并管理流水线深度

加载新batch数据，如果流水线过长则开始新的数据分布

self.enqueue_batch(dataloader_iter)  # 加载新的批次
if len(self.batches) >= 5 + self.local_unique_parallel_batch_num:
    self.start_sparse_data_dist( # 防止流水线过长
        self.batches[4 + self.local_unique_parallel_batch_num],
        self.contexts[4 + self.local_unique_parallel_batch_num], )

对batch_0数据进行嵌入缓存管理

对fill_pipeline中的batch_0的数据同步换出操作，异步更新主机端嵌入，执行换入操作

# 更新 host 侧的 Embedding 和 优化器参数
self.contexts[0].event_gather_swapouted.synchronize()
self.host_embedding_update_async(self.contexts[0])

# 等待 swapin tensor 并 swapin
self.swapin_tensors_to_npu(self.contexts[0])
self.wait_and_swapin(self.contexts[0])

核心计算：前向传播

batch_0的全部数据处理操作已经完成，对目前host侧的embeddding数据执行模型前向传播，计算损失和输出

# forward
with record_function("## forward ##"):
    losses, output = self._model_fwd(self.batches[0])

流水线协调与预计算

协调批次间依赖，定期淘汰不常用特征，预计算后续批次信息

# 等待 i+1 swap 参数计算完成
if len(self.batches) >= 2:
    self.wait_and_get_swap_info(self.contexts[1])

# 处理淘汰信息
if (self._evict_step_interval and (self._global_steps + 1) % self._evict_step_interval == 0 ):
    with record_function("## feature_evict ##"):
        self._start_feature_evict()

# 后续batch的swap计算，计算 i+2 batch 的 swap pairs 和 key2offset
if len(self.batches) >= 3:
    self.start_compute_swap_info(self.contexts[2])

batch_1数据预处理

在当前批次计算时，异步预处理下一个批次。保险起见，i+1轮的swapout等i轮的swapin做完之后做

# batches[i+1] 在host侧预查
if len(self.batches) >= 2:
    self.contexts[1].event_can_swapout.record(self._default_stream)
    self.host_embedding_lookup_async(self.contexts[1])
    self.do_restore_async(self.contexts[1])
    
if self._model.training:
    # backward
    with record_function("## backward ##"):
        torch.sum(losses, dim=0).backward()
    # update
    with record_function("## optimizer ##"):
        self._optimizer.step()

流水线清理和推进

完成远端批次处理，执行换出操作，等待主机更新，出队当前批次

if len(self.batches) >= 4:
    self.do_post_input_dist(self.contexts[3])

# swapout
if len(self.batches) >= 2:
    self.swap_out(self.contexts[1])

self.wait_host_update(self.contexts[0])
self.dequeue_batch()
return output, losses

这个设计确保了在大规模嵌入表训练场景下，系统能够保持高吞吐量和低延迟。

3. GPU侧多级流水

GPU侧采用自定义的StagedTrainPipeline，通过CUDA_Stream实现异步机制。当前stage负责提交任务，下个stage开始前需要wait上一个stage的结果。图中的prefecth就是换入（emb的H2D）或换出操作（emb的D2H）。图中data H2D、input_dist操作与NPU侧含义一致。

batch_1的prefecth会prefecth当前batch所有的embedding到HBM；
batch_0的bwd（update）结束后，需要将batch_0的所有embedding刷回到MEM，同时需要把 batch_0和batch_1交集的embedding进行更新，这是因为batch_1训练的数据需要从batch_0中获取最新取值；
batch_0的bwd（update）中（求交&更新）的执行需要等batch_1的prefecth完成后开始；
batch_2的prefect启动需要依赖batch_0的所有embedding刷回HBM；

外部接口调用：

import torch

model = ...
optimizer = ...
dataloader_iter = ...
pipeline_stages = [
    PipelineStage(
        name="data H2D",
        runnable=partial(model.stage1, device=self.device),
    ),
    PipelineStage(
        name="input_dist",
        runnable=model.stage2,
    ),
    PipelineStage(
        name="prefetch",
        runnable=model.stage3,
    ),
]
emb = GpuPpipeEmbedding(
    dataloader_iter
)
def foward():
    sync & swicth # 流水对齐，等待backward完成之后再去做下一个batch的换入换出
    data_h2d     # n+3
    input_dist   # n+2
    prefetch(dump and load) # n+1
   return batch
model = BuildMode(emb)
engine = unrec_engine(model)
while True:     # 模型训练
    try:
        optimizer.zero_grad()
        logits = model()
        loss = Loss(logits)
        loss.backward()
        optimizer.step()
        emb.save()     # embedding保存
        torch.save()   # dense层保存
    except StopIteration:
        break

内部实现:

初始化阶段

在真正的持续输出开始之前，需要先将前几个 batch 推进到各个 stage 中，形成“满载”状态。

def _fill_pipeline(self, dataloader_iter: Iterator[In]) -> None:
        for batch_offset in range(self.num_stages):
            stages_to_run = self.num_stages - batch_offset
            for stage_idx in range(stages_to_run):
                self._run_stage(
                    batch_offset=batch_offset,
                    stage_idx=stage_idx,
                    dataloader_iter=dataloader_iter,
                    fill=True,)
        self._initialized = True

循环执行模块

每调用一次 progress() 输出一个完成 batch，该方法构成了训练 loop 的核心接口：

def progress(self,dataloader_iter: Iterator[In],) -> Optional[StageOut]:
    if not self._initialized:
        self._fill_pipeline(dataloader_iter)

    output = self._advance()
    if output is None:
        raise StopIteration

    self._num_steps += 1

    for stage_idx in range(self.num_stages):
        stage_output_idx = self.num_stages - 1 - stage_idx
        self._run_stage(
            batch_offset=stage_output_idx,
            stage_idx=stage_idx,
            dataloader_iter=dataloader_iter,
        )
    return output

流水线stage执行模块

根据输入stage的值，执行单个 stage 的逻辑。

def _run_stage(self,batch_offset: int,stage_idx: int,dataloader_iter: Iterator[In],fill: bool = False,) -> 					StageOutputWithEvent:
        stage = self._pipeline_stages[stage_idx]
        with record_function(
            f"## Pipeline Stage {stage_idx} : {stage.name} for batch {batch_offset + self._num_steps} ##"):
            if stage_idx == 0:
                batch_to_wait = self._next_batch(dataloader_iter)
            else:
                batch_to_wait = self._stage_outputs[batch_offset]
                assert batch_to_wait is not None
            new_result = self._run_with_event(
                runnable=stage.runnable,
                inputs=batch_to_wait,
            )
        self._stage_outputs[batch_offset] = new_result
        if fill and (fill_callback := stage.fill_callback) is not None:
            if self._debug_mode:
                logger.info(f"Finished callback for {stage.name}")
            fill_callback()
        return new_result

posted @ 2025-12-27 14:45 永荣带你玩转昇腾阅读(0) 评论(0) 收藏举报

刷新页面返回顶部

pengyongrong

06_昇腾流水线优化技术详解

昇腾流水线优化技术详解

1. pipeline多级流水

2. Rec SDK多级流水

用户端接口调用

内部实现：

fill_pipeline 方法：流水线初始化填充

batch_0的阶段性处理

batch_1的阶段性处理

batch_2和batch_3的阶段性处理

并行Batch的预加载机制

`progress` 方法：流水线推进执行

流水线末端数据同步

加载新batch数据并管理流水线深度

对batch_0数据进行嵌入缓存管理

核心计算：前向传播

流水线协调与预计算

batch_1数据预处理

流水线清理和推进

3. GPU侧多级流水

外部接口调用：

内部实现:

初始化阶段

循环执行模块

流水线stage执行模块

公告

pengyongrong

06_昇腾流水线优化技术详解

昇腾流水线优化技术详解

1. pipeline多级流水

2. Rec SDK多级流水

用户端接口调用

内部实现：

fill_pipeline 方法：流水线初始化填充

batch_0的阶段性处理

batch_1的阶段性处理

batch_2和batch_3的阶段性处理

并行Batch的预加载机制

progress 方法：流水线推进执行

流水线末端数据同步

加载新batch数据并管理流水线深度

对batch_0数据进行嵌入缓存管理

核心计算：前向传播

流水线协调与预计算

batch_1数据预处理

流水线清理和推进

3. GPU侧多级流水

外部接口调用：

内部实现:

初始化阶段

循环执行模块

流水线stage执行模块

公告

`progress` 方法：流水线推进执行