Profile工作判斷模型的計(jì)算以及內(nèi)存瓶頸

Profile模型訓(xùn)練的循環(huán)

標(biāo)記任意代碼范圍

Profile CPU/GPU的活動(dòng)

Profile 內(nèi)存消耗

總覽

Flops 測(cè)量

多GPU，多節(jié)點(diǎn)，數(shù)據(jù)并行和模型并行

例子

和DeepSpeed運(yùn)行時(shí)一起使用

在Megatron-LM中使用

在 DeepSpeed 運(yùn)行環(huán)境之外的使用方法

訓(xùn)練工作流例子

0x0. 前言

這篇翻譯是對(duì) https://www.deepspeed.ai/tutorials/pytorch-profiler/ 和 https://www.deepspeed.ai/tutorials/flops-profiler/ 兩篇教程做的，使用DeepSpeed訓(xùn)練模型可以基于這兩個(gè)教程做一下Profile工作判斷模型的計(jì)算以及內(nèi)存瓶頸在哪個(gè)地方。

0x1. 在 DeepSpeed 中使用 PyTorch Profiler做性能調(diào)試

對(duì)應(yīng)原始的教程：https://www.deepspeed.ai/tutorials/pytorch-profiler/

這個(gè)教程描述的是如何在DeepSpeed中使用PyTorch Profiler工具（https://pytorch.org/blog/introducing-pytorch-profiler-the-new-and-improved-performance-tool/）。

PyTorch Profiler 是一個(gè)開源工具，能夠?yàn)榇笠?guī)模深度學(xué)習(xí)模型提供精確且高效的性能分析和故障排查。分析結(jié)果可以輸出為 .json 追蹤文件，并在 Google Chrome 的追蹤查看器 (chrome://tracing) 中查看。Microsoft Visual Studio Code 的 Python 擴(kuò)展將 TensorBoard 集成到代碼編輯器中，包括對(duì) PyTorch Profiler 的支持。更多的細(xì)節(jié)可以參考（https://pytorch.org/tutorials/recipes/recipes/profiler_recipe.html#pytorch-profiler）

Profile模型訓(xùn)練的循環(huán)

下面展示了如何通過在 Profiler 上下文管理器中封裝代碼來分析訓(xùn)練循環(huán)。Profiler 假設(shè)訓(xùn)練過程由steps（從零開始編號(hào)）組成。PyTorch Profiler 接受許多參數(shù)，例如 schedule, on_trace_ready, with_stack 等。

在下面的例子中，分析器將跳過前5步，使用接下來的2步作為預(yù)熱，并記錄接下來的6步。由于repeat設(shè)為2，所以分析器將在兩個(gè)周期后停止記錄（這里的周期的意思是將active的step數(shù)重復(fù)repeat次）。關(guān)于 schedule 的詳細(xì)使用方法，請(qǐng)參考使用Profiler分析長(zhǎng)時(shí)間運(yùn)行的任務(wù)（https://pytorch.org/tutorials/recipes/recipes/profiler_recipe.html#using-profiler-to-analyze-long-running-jobs）。

fromtorch.profilerimportprofile,record_function,ProfilerActivity

withtorch.profiler.profile(
schedule=torch.profiler.schedule(
wait=5,#Duringthisphaseprofilerisnotactive.
warmup=2,#Duringthisphaseprofilerstartstracing,buttheresultsarediscarded.
active=6,#Duringthisphaseprofilertracesandrecordsdata.
repeat=2),#Specifiesanupperboundonthenumberofcycles.
on_trace_ready=tensorboard_trace_handler,
with_stack=True#Enablestacktracing,addsextraprofilingoverhead.
)asprofiler:
forstep,batchinenumerate(data_loader):
print("step:{}".format(step))

#forward()method
loss=model_engine(batch)

#runsbackpropagation
model_engine.backward(loss)

#weightupdate
model_engine.step()
profiler.step()#Sendthesignaltotheprofilerthatthenextstephasstarted.

標(biāo)記任意代碼范圍

可以使用 record_function 上下文管理器標(biāo)記任意用戶指定的代碼范圍。例如，以下代碼將 "model_forward" 標(biāo)記為一個(gè)label：

withprofile(record_shapes=True)asprof:#record_shapesindicateswhethertorecordshapesoftheoperatorinputs.
withrecord_function("model_forward"):"
model_engine(inputs)

后續(xù)在profile結(jié)果里面就可以看到標(biāo)記的這個(gè)"model_forward"的耗時(shí)情況了。

Profile CPU/GPU的活動(dòng)

傳遞給 Profiler 的 activities 參數(shù)指定了在使用 profiler 上下文管理器包裝的代碼范圍執(zhí)行期間要進(jìn)行性能分析的活動(dòng)列表：

ProfilerActivity.CPU - PyTorch 操作符、TorchScript 函數(shù)和用戶定義的代碼標(biāo)簽（record_function）。

ProfilerActivity.CUDA - 在設(shè)備上的 CUDA 核函數(shù)。請(qǐng)注意，CUDA 性能分析會(huì)帶來不可忽視的開銷。下面的例子在模型的前向傳播中對(duì) CPU 和 GPU 的活動(dòng)進(jìn)行了性能分析，并按總 CUDA 時(shí)間排序打印了總結(jié)表。

withprofile(activities=[
ProfilerActivity.CPU,ProfilerActivity.CUDA],record_shapes=True)asprof:
withrecord_function("model_forward"):
model_engine(inputs)

print(prof.key_averages().table(sort_by="cuda_time_total",row_limit=10))

Profile 內(nèi)存消耗

通過向 PyTorch Profiler傳遞 profile_memory=True，我們啟用了內(nèi)存分析功能，該功能記錄在模型OP的執(zhí)行過程中分配（或釋放）的模型張量所使用的內(nèi)存量。例如：

withprofile(activities=[ProfilerActivity.CUDA],
profile_memory=True,record_shapes=True)asprof:
model(inputs)

print(prof.key_averages().table(sort_by="self_cuda_memory_usage",row_limit=10))

0x2. Flops Profiler

對(duì)應(yīng)原始的教程：https://www.deepspeed.ai/tutorials/flops-profiler/

在這個(gè)教程中，我們將介紹 DeepSpeed Flops Profiler，并提供其使用的示例。

總覽

有效利用硬件資源對(duì)于良好的性能至關(guān)重要，但在現(xiàn)有的大規(guī)模模型訓(xùn)練和推理實(shí)現(xiàn)中，性能低效往往難以察覺并歸因于特定的模塊組件。DeepSpeed Flops Profiler幫助用戶輕松測(cè)量模型及其子模塊的訓(xùn)練/推理速度（延遲，吞吐量）和效率（每秒浮點(diǎn)運(yùn)算次數(shù)，即FLOPS），旨在消除現(xiàn)有實(shí)現(xiàn)中的效率低下問題。

以下是在A100 GPU上，批量大小為80的BERT-Large（NVIDIA）的示例輸出：

-------------------------- DeepSpeed Flops Profiler --------------------------
Profile Summary at step 10:
Notations:
data parallel size (dp_size), model parallel size(mp_size),
number of parameters (params), number of multiply-accumulate operations(MACs),
number of floating-point operations (flops), floating-point operations per second (FLOPS),
fwd latency (forward propagation latency), bwd latency (backward propagation latency),
step (weights update latency), iter latency (sum of fwd, bwd and step latency)

world size:                                                   1
data parallel size:                                           1
model parallel size:                                          1
batch size per GPU:                                           80
params per gpu:                                               336.23 M
params of model = params per GPU * mp_size:                   336.23 M
fwd MACs per GPU:                                             3139.93 G
fwd flops per GPU:                                            6279.86 G
fwd flops of model = fwd flops per GPU * mp_size:             6279.86 G
fwd latency:                                                  76.67 ms
bwd latency:                                                  108.02 ms
fwd FLOPS per GPU = fwd flops per GPU / fwd latency:          81.9 TFLOPS
bwd FLOPS per GPU = 2 * fwd flops per GPU / bwd latency:      116.27 TFLOPS
fwd+bwd FLOPS per GPU = 3 * fwd flops per GPU / (fwd+bwd latency):   102.0 TFLOPS
step latency:                                                 34.09 us
iter latency:                                                 184.73 ms
samples/second:                                               433.07

----------------------------- Aggregated Profile per GPU -----------------------------
Top modules in terms of params, MACs or fwd latency at different model depths:
depth 0:
    params      - {'BertForPreTrainingPreLN': '336.23 M'}
    MACs        - {'BertForPreTrainingPreLN': '3139.93 GMACs'}
    fwd latency - {'BertForPreTrainingPreLN': '76.39 ms'}
depth 1:
    params      - {'BertModel': '335.15 M', 'BertPreTrainingHeads': '32.34 M'}
    MACs        - {'BertModel': '3092.96 GMACs', 'BertPreTrainingHeads': '46.97 GMACs'}
    fwd latency - {'BertModel': '34.29 ms', 'BertPreTrainingHeads': '3.23 ms'}
depth 2:
    params      - {'BertEncoder': '302.31 M', 'BertLMPredictionHead': '32.34 M'}
    MACs        - {'BertEncoder': '3092.88 GMACs', 'BertLMPredictionHead': '46.97 GMACs'}
    fwd latency - {'BertEncoder': '33.45 ms', 'BertLMPredictionHead': '2.61 ms'}
depth 3:
    params      - {'ModuleList': '302.31 M', 'Embedding': '31.79 M', 'Linear': '31.26 M'}
    MACs        - {'ModuleList': '3092.88 GMACs', 'Linear': '36.23 GMACs'}
    fwd latency - {'ModuleList': '33.11 ms', 'BertPredictionHeadTransform': '1.83 ms''}
depth 4:
    params      - {'BertLayer': '302.31 M', 'LinearActivation': '1.05 M''}
    MACs        - {'BertLayer': '3092.88 GMACs', 'LinearActivation': '10.74 GMACs'}
    fwd latency - {'BertLayer': '33.11 ms', 'LinearActivation': '1.43 ms'}
depth 5:
    params      - {'BertAttention': '100.76 M', 'BertIntermediate': '100.76 M'}
    MACs        - {'BertAttention': '1031.3 GMACs', 'BertIntermediate': '1030.79 GMACs'}
    fwd latency - {'BertAttention': '19.83 ms', 'BertOutput': '4.38 ms'}
depth 6:
    params      - {'LinearActivation': '100.76 M', 'Linear': '100.69 M'}
    MACs        - {'LinearActivation': '1030.79 GMACs', 'Linear': '1030.79 GMACs'}
    fwd latency - {'BertSelfAttention': '16.29 ms', 'LinearActivation': '3.48 ms'}

------------------------------ Detailed Profile per GPU ------------------------------
Each module profile is listed after its name in the following order:
params, percentage of total params, MACs, percentage of total MACs, fwd latency, percentage of total fwd latency, fwd FLOPS

BertForPreTrainingPreLN(
  336.23 M, 100.00% Params, 3139.93 GMACs, 100.00% MACs, 76.39 ms, 100.00% latency, 82.21 TFLOPS,
  (bert): BertModel(
    335.15 M, 99.68% Params, 3092.96 GMACs, 98.50% MACs, 34.29 ms, 44.89% latency, 180.4 TFLOPS,
    (embeddings): BertEmbeddings(...)
    (encoder): BertEncoder(
      302.31 M, 89.91% Params, 3092.88 GMACs, 98.50% MACs, 33.45 ms, 43.79% latency, 184.93 TFLOPS,
      (FinalLayerNorm): FusedLayerNorm(...)
      (layer): ModuleList(
        302.31 M, 89.91% Params, 3092.88 GMACs, 98.50% MACs, 33.11 ms, 43.35% latency, 186.8 TFLOPS,
        (0): BertLayer(
          12.6 M, 3.75% Params, 128.87 GMACs, 4.10% MACs, 1.29 ms, 1.69% latency, 199.49 TFLOPS,
          (attention): BertAttention(
            4.2 M, 1.25% Params, 42.97 GMACs, 1.37% MACs, 833.75 us, 1.09% latency, 103.08 TFLOPS,
            (self): BertSelfAttention(
              3.15 M, 0.94% Params, 32.23 GMACs, 1.03% MACs, 699.04 us, 0.92% latency, 92.22 TFLOPS,
              (query): Linear(1.05 M, 0.31% Params, 10.74 GMACs, 0.34% MACs, 182.39 us, 0.24% latency, 117.74 TFLOPS,...)
              (key): Linear(1.05 M, 0.31% Params, 10.74 GMACs, 0.34% MACs, 57.22 us, 0.07% latency, 375.3 TFLOPS,...)
              (value): Linear(1.05 M, 0.31% Params, 10.74 GMACs, 0.34% MACs, 53.17 us, 0.07% latency, 403.91 TFLOPS,...)
              (dropout): Dropout(...)
              (softmax): Softmax(...)
            )
            (output): BertSelfOutput(
              1.05 M, 0.31% Params, 10.74 GMACs, 0.34% MACs, 114.68 us, 0.15% latency, 187.26 TFLOPS,
              (dense): Linear(1.05 M, 0.31% Params, 10.74 GMACs, 0.34% MACs, 64.13 us, 0.08% latency, 334.84 TFLOPS, ...)
              (dropout): Dropout(...)
            )
          )
          (PreAttentionLayerNorm): FusedLayerNorm(...)
          (PostAttentionLayerNorm): FusedLayerNorm(...)
          (intermediate): BertIntermediate(
            4.2 M, 1.25% Params, 42.95 GMACs, 1.37% MACs, 186.68 us, 0.24% latency, 460.14 TFLOPS,
            (dense_act): LinearActivation(4.2 M, 1.25% Params, 42.95 GMACs, 1.37% MACs, 175.0 us, 0.23% latency, 490.86 TFLOPS,...)
          )
          (output): BertOutput(
            4.2 M, 1.25% Params, 42.95 GMACs, 1.37% MACs, 116.83 us, 0.15% latency, 735.28 TFLOPS,
            (dense): Linear(4.2 M, 1.25% Params, 42.95 GMACs, 1.37% MACs, 65.57 us, 0.09% latency, 1310.14 TFLOPS,...)
            (dropout): Dropout(...)
          )
        )
        ...
        (23): BertLayer(...)
      )
    )
    (pooler): BertPooler(...)
  )
  (cls): BertPreTrainingHeads(...)
)
------------------------------------------------------------------------------

在 profile 總結(jié)中，DeepSpeed Flops Profiler輸出了模型的參數(shù)量，浮點(diǎn)運(yùn)算數(shù)（flops），F(xiàn)LOPS，延遲，以及樣本/秒的吞吐量。此概況顯示了當(dāng)前模型執(zhí)行與硬件峰值性能之間的性能差距，并幫助用戶調(diào)整訓(xùn)練或推理設(shè)置（例如，超參數(shù)，數(shù)據(jù)并行性，模型并行性，系統(tǒng)配置等）以獲得更好的性能。

DeepSpeed Flops Profiler還可以在不同的模型深度（聚合profile）和模型架構(gòu)中的特定模塊（詳細(xì)profile）Profile重要模塊。通過這些Profile，DeepSpeed用戶可以理解每個(gè)層或子模塊對(duì)整個(gè)模型復(fù)雜性/性能的貢獻(xiàn)。然后，用戶可以調(diào)整或重構(gòu)模型設(shè)計(jì)以提高性能。例如，使用Profiler，DeepSpeed用戶可以量化地判斷堆疊較小的層是否比擁有較大的層更輕量或性能更好。聚合和詳細(xì) Profile 還允許用戶快速識(shí)別瓶頸模塊。在上面的BERT-Large示例中，使用DeepSpeed Flops Profiler，我們發(fā)現(xiàn)BertLayer是最重要的層，并且包含了很多dropout，softmax和layer norm以及線性層模塊。這些模塊在flops中并不heavy，但會(huì)觸發(fā)許多GPU Kernel調(diào)用并創(chuàng)建過多的內(nèi)存讀/寫請(qǐng)求。詳細(xì)Profile中顯示的Pattern表明這是Kernel融合的完美匹配，我們開發(fā)了fused transformer-kernels來減少數(shù)據(jù)移動(dòng)（參見 https://www.deepspeed.ai/tutorials/bert-pretraining/）。在應(yīng)用我們的優(yōu)化之后，我們?cè)贒eepSpeed Flops Profiler的輸出中看到每GPU的FLOPS和總體訓(xùn)練樣本/秒提高了25%。

DeepSpeed Flops Profiler可以與DeepSpeed運(yùn)行時(shí)一起使用，并無需任何用戶代碼更改，也可以獨(dú)立于DeepSpeed作為一個(gè)獨(dú)立的包使用。在使用DeepSpeed進(jìn)行模型訓(xùn)練時(shí)，可以在DeepSpeed配置文件（https://www.deepspeed.ai/docs/config-json/#flops-profiler）中啟用分析器。作為一個(gè)獨(dú)立的包，分析器API可以在訓(xùn)練和推理代碼中使用。DeepSpeed分析器仍在積極開發(fā)中，目前僅包含初始功能。請(qǐng)保持關(guān)注，很快會(huì)添加更多激動(dòng)人心的功能。

Flops 測(cè)量

與現(xiàn)有的flops計(jì)算工具或方法類似，DeepSpeed Flops分析器測(cè)量Module前向傳播的flops，而反向傳播的flops則被估計(jì)為前向傳播flops的兩倍。與計(jì)算PyTorch Op的flops的PyTorch分析器不同，DeepSpeed Flops分析器測(cè)量模型中模塊內(nèi)部的flops，并為用戶提供關(guān)于模型執(zhí)行的更多洞察。flops估計(jì)部分受到ptflops（https://github.com/sovrasov/flops-counter.pytorch）的啟發(fā)，主要區(qū)別在于，DeepSpeed Flops分析器不僅支持直接在模塊級(jí)別進(jìn)行FLOPS計(jì)算，還可以捕獲在模塊中調(diào)用的torch.nn.functional來估計(jì)flops。因此，DeepSpeed Flops分析器允許在模型中使用自定義模塊，例如Megatron-LM中的ParallelTransformerLayerworks、ParallelSelfAttention、RowParallelLinear等。這與ptflops形成對(duì)比，ptflops要求用戶為每個(gè)自定義模塊編寫自定義的flops計(jì)算函數(shù)。

多GPU，多節(jié)點(diǎn)，數(shù)據(jù)并行和模型并行

DeepSpeed Flops 分析器輸出每個(gè) GPU 的分析結(jié)果以及world size，數(shù)據(jù)并行大小和模型并行大小。

對(duì)于在多 GPU 或多節(jié)點(diǎn)上運(yùn)行的模型，只有模型并行（例如，Megatron-LM 中的 --model-parallel-size）的改變會(huì)影響浮點(diǎn)操作數(shù)和Paramater的分析結(jié)果，即，model_parallel_size * flops = total_flops 和 model_parallel_size * parameters = total_parameters。數(shù)據(jù)并行大小或world size（與 GPU 或節(jié)點(diǎn)的數(shù)量相關(guān)）不會(huì)影響每個(gè) GPU 的分析結(jié)果。

例子

DeepSpeed Flops 分析器可以與 DeepSpeed 運(yùn)行時(shí)一起使用，也可以作為一個(gè)獨(dú)立的包使用。當(dāng)使用 DeepSpeed 進(jìn)行模型訓(xùn)練時(shí)，用戶無需更改代碼，就可以在 deepspeed 配置文件（https://www.deepspeed.ai/docs/config-json/#flops-profiler）中配置分析器。要在 DeepSpeed 運(yùn)行時(shí)之外使用 flops 分析器，安裝 DeepSpeed 并導(dǎo)入 flops_profiler 包直接使用 API。下面給出了每種使用方法的示例。

和DeepSpeed運(yùn)行時(shí)一起使用

當(dāng)使用 DeepSpeed 進(jìn)行模型訓(xùn)練時(shí)，可以在 deepspeed 配置文件中配置分析器。使用分析器不需要明確的 API 調(diào)用?？梢酝ㄟ^在 deepspeed 的配置 json 文件中添加以下字段來啟用分析器。具體詳情請(qǐng)參考 flops profiler(https://www.deepspeed.ai/docs/config-json/#flops-profiler)。

{
"flops_profiler":{
"enabled":true,
"profile_step":1,
"module_depth":-1,
"top_modules":1,
"detailed":true,
"output_file":null
}
}

在Megatron-LM中使用

關(guān)于使用 DeepSpeed 運(yùn)行 Megatron-LM 的信息，請(qǐng)參考我們的教程 Megatron-LM。

下面展示了一個(gè) 12 層 Megatron-LM 模型的示例輸出（hidden_size = 8192，num_attention_heads = 32，batch_size = 1024，seq_length = 1024）。

-------------------------- DeepSpeed Flops Profiler --------------------------
Profile Summary at step 10:
Notations:
data parallel size (dp_size), model parallel size(mp_size),
number of parameters (params), number of multiply-accumulate operations(MACs),
number of floating-point operations (flops), floating-point operations per second (FLOPS),
fwd latency (forward propagation latency), bwd latency (backward propagation latency),
step (weights update latency), iter latency (sum of fwd, bwd and step latency)

world size:                                                   1
data parallel size:                                           1
model parallel size:                                          1
batch size per GPU:                                           1024
params per gpu:                                               1.29 M
params of model = params per GPU * mp_size:                   1.29 M
fwd MACs per GPU:                                             41271.95 G
fwd flops per GPU:                                            82543.9 G
fwd flops of model = fwd flops per GPU * mp_size:             82543.9 G
fwd latency:                                                  1.89 s
bwd latency:                                                  5.38 s
fwd FLOPS per GPU = fwd flops per GPU / fwd latency:          43.68 TFLOPS
bwd FLOPS per GPU = 2 * fwd flops per GPU / bwd latency:      30.7 TFLOPS
fwd+bwd FLOPS per GPU = 3 * fwd flops per GPU / (fwd+bwd latency):   34.07 TFLOPS
step latency:                                                 34.12 s
iter latency:                                                 41.39 s
samples/second:                                               24.74

----------------------------- Aggregated Profile per GPU -----------------------------
Top 1 modules in terms of params, MACs or fwd latency at different model depths:
depth 0:
    params      - {'GPT2Model': '1.29 M'}
    MACs        - {'GPT2Model': '41271.95 GMACs'}
    fwd latency - {'GPT2Model': '1.84 s'}
depth 1:
    params      - {'TransformerLanguageModel': '1.29 M'}
    MACs        - {'TransformerLanguageModel': '39584.03 GMACs'}
    fwd latency - {'TransformerLanguageModel': '1.83 s'}
depth 2:
    params      - {'ParallelTransformer': '1.29 M'}
    MACs        - {'ParallelTransformer': '39584.03 GMACs'}
    fwd latency - {'ParallelTransformer': '1.81 s'}
depth 3:
    params      - {'ModuleList': '1.28 M'}
    MACs        - {'ModuleList': '39584.03 GMACs'}
    fwd latency - {'ModuleList': '1.3 s'}
depth 4:
    params      - {'ParallelTransformerLayerPart2': '688.15 k'}
    MACs        - {'ParallelTransformerLayerPart2': '26388.28 GMACs'}
    fwd latency - {'ParallelTransformerLayerPart2': '865.73 ms'}
depth 5:
    params      - {'ParallelMLP': '491.54 k'}
    MACs        - {'ParallelMLP': '26388.28 GMACs'}
    fwd latency - {'ParallelMLP': '849.4 ms'}

------------------------------ Detailed Profile per GPU ------------------------------
Each module profile is listed after its name in the following order:
params, percentage of total params, MACs, percentage of total MACs, fwd latency, percentage of total fwd latency, fwd FLOPS

Note: 1. A module can have torch.nn.module or torch.nn.functional to compute logits (e.g. CrossEntropyLoss). They are not counted as submodules, thus not to be printed out. However they make up the difference between a parent's MACs(or latency) and the sum of its submodules'.
1. Number of floating-point operations is a theoretical estimation, thus FLOPS computed using that could be larger than the maximum system throughput.
2. The fwd latency listed in the top module's profile is directly captured at the module forward function in PyTorch, thus it's less than the fwd latency shown above which is captured in DeepSpeed.

GPT2Model(
  1.29 M, 100.00% Params, 41271.95 GMACs, 100.00% MACs, 1.84 s, 100.00% latency, 44.78 TFLOPS,
  (language_model): TransformerLanguageModel(
    1.29 M, 100.00% Params, 39584.03 GMACs, 95.91% MACs, 1.83 s, 99.11% latency, 43.34 TFLOPS,
    (embedding): Embedding(
      2, 0.00% Params, 0 MACs, 0.00% MACs, 18.1 ms, 0.98% latency, 0.0 FLOPS,
      (word_embeddings): VocabParallelEmbedding(1, 0.00% Params, 0 MACs, 0.00% MACs, 164.75 us, 0.01% latency, 0.0 FLOPS, )
      (position_embeddings): Embedding(1, 0.00% Params, 0 MACs, 0.00% MACs, 489.23 us, 0.03% latency, 0.0 FLOPS, 1024, 8192)
      (embedding_dropout): Dropout(0, 0.00% Params, 0 MACs, 0.00% MACs, 93.94 us, 0.01% latency, 0.0 FLOPS, p=0.1, inplace=False)
    )
    (transformer): ParallelTransformer(
      1.29 M, 100.00% Params, 39584.03 GMACs, 95.91% MACs, 1.81 s, 98.11% latency, 43.78 TFLOPS,
      (layers): ModuleList(
        1.28 M, 98.73% Params, 39584.03 GMACs, 95.91% MACs, 1.3 s, 70.66% latency, 60.79 TFLOPS,
        (0): ParallelTransformerLayerPart1(
          49.15 k, 3.80% Params, 1099.65 GMACs, 2.66% MACs, 23.5 ms, 1.27% latency, 93.6 TFLOPS,
          (input_layernorm): FusedLayerNorm(16.38 k, 1.27% Params, 0 MACs, 0.00% MACs, 128.75 us, 0.01% latency, 0.0 FLOPS, torch.Size([8192]), eps=1e-05, elementwise_affine=True)
          (attention): ParallelSelfAttention(
            32.77 k, 2.53% Params, 1099.65 GMACs, 2.66% MACs, 22.8 ms, 1.24% latency, 96.46 TFLOPS,
            (query_key_value): ColumnParallelLinear(24.58 k, 1.90% Params, 824.63 GMACs, 2.00% MACs, 8.93 ms, 0.48% latency, 184.7 TFLOPS, )
            (scale_mask_softmax): FusedScaleMaskSoftmax(0, 0.00% Params, 134.22 MMACs, 0.00% MACs, 151.16 us, 0.01% latency, 1.78 TFLOPS, )
            (attention_dropout): Dropout(0, 0.00% Params, 0 MACs, 0.00% MACs, 79.63 us, 0.00% latency, 0.0 FLOPS, p=0.1, inplace=False)
            (dense): RowParallelLinear(8.19 k, 0.63% Params, 274.88 GMACs, 0.67% MACs, 2.67 ms, 0.14% latency, 205.81 TFLOPS, )
          )
        )
        (1): ParallelTransformerLayerPart2(
          57.35 k, 4.43% Params, 2199.02 GMACs, 5.33% MACs, 77.53 ms, 4.21% latency, 56.73 TFLOPS,
          (post_attention_layernorm): FusedLayerNorm(16.38 k, 1.27% Params, 0 MACs, 0.00% MACs, 116.11 us, 0.01% latency, 0.0 FLOPS, torch.Size([8192]), eps=1e-05, elementwise_affine=True)
          (mlp): ParallelMLP(
            40.96 k, 3.16% Params, 2199.02 GMACs, 5.33% MACs, 76.19 ms, 4.13% latency, 57.72 TFLOPS,
            (dense_h_to_4h): ColumnParallelLinear(32.77 k, 2.53% Params, 1099.51 GMACs, 2.66% MACs, 10.79 ms, 0.59% latency, 203.81 TFLOPS, )
            (dense_4h_to_h): RowParallelLinear(8.19 k, 0.63% Params, 1099.51 GMACs, 2.66% MACs, 14.38 ms, 0.78% latency, 152.95 TFLOPS, )
          )
        )
        ...
        (23): ParallelTransformerLayerPart2(...)
      )
      (final_layernorm): FusedLayerNorm(16.38 k, 1.27% Params, 0 MACs, 0.00% MACs, 110.86 us, 0.01% latency, 0.0 FLOPS, torch.Size([8192]), eps=1e-05, elementwise_affine=True)
    )
  )
)
------------------------------------------------------------------------------

可以參考最新的DeepSpeed-Megatron倉(cāng)庫(kù)，然后在訓(xùn)練模型時(shí)將DeepSpeed的config文件配置DeepSpeed Profiler。

在 DeepSpeed 運(yùn)行環(huán)境之外的使用方法

profiler 可以在 DeepSpeed 運(yùn)行時(shí)環(huán)境之外作為一個(gè)獨(dú)立的包來使用。你只需要簡(jiǎn)單地安裝 DeepSpeed 并導(dǎo)入 flops_profiler 包來直接使用 API。關(guān)于如何安裝 DeepSpeed，請(qǐng)參考 DeepSpeed 的安裝指南。

在模型推理中

要對(duì)推理狀態(tài)的訓(xùn)練模型進(jìn)行性能分析，請(qǐng)使用 get_model_profile 函數(shù)。下面給出了一些示例。

AlexNet例子

以下示例展示了如何使用 DeepSpeed flops 分析器對(duì) AlexNet 進(jìn)行性能分析。

import torchvision.models as models
import torch
from deepspeed.profiling.flops_profiler import get_model_profile
from deepspeed.accelerator import get_accelerator

with get_accelerator().device(0):
    model = models.alexnet()
    batch_size = 256
    flops, macs, params = get_model_profile(model=model, # model
                                    input_shape=(batch_size, 3, 224, 224), # input shape to the model. If specified, the model takes a tensor with this shape as the only positional argument.
                                    args=None, # list of positional arguments to the model.
                                    kwargs=None, # dictionary of keyword arguments to the model.
                                    print_profile=True, # prints the model graph with the measured profile attached to each module
                                    detailed=True, # print the detailed profile
                                    module_depth=-1, # depth into the nested modules, with -1 being the inner most modules
                                    top_modules=1, # the number of top modules to print aggregated profile
                                    warm_up=10, # the number of warm-ups before measuring the time of each module
                                    as_string=True, # print raw numbers (e.g. 1000) or as human-readable strings (e.g. 1k)
                                    output_file=None, # path to the output file. If None, the profiler prints to stdout.
                                    ignore_modules=None) # the list of modules to ignore in the profiling

BERT例子

from functools import partial
import torch
from transformers import BertForSequenceClassification, BertTokenizer
from deepspeed.profiling.flops_profiler import get_model_profile
from deepspeed.accelerator import get_accelerator


def bert_input_constructor(batch_size, seq_len, tokenizer):
    fake_seq = ""
    for _ in range(seq_len - 2):  # ignore the two special tokens [CLS] and [SEP]
      fake_seq += tokenizer.pad_token
    inputs = tokenizer([fake_seq] * batch_size,
                       padding=True,
                       truncation=True,
                       return_tensors="pt")
    labels = torch.tensor([1] * batch_size)
    inputs = dict(inputs)
    inputs.update({"labels": labels})
    return inputs


with get_accelerator().device(0):
    tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
    model = BertForSequenceClassification.from_pretrained('bert-base-uncased')
    batch_size = 4
    seq_len = 128
    enable_profile = True
    if enable_profile:
      flops, macs, params = get_model_profile(
          model,
          kwargs=bert_input_constructor(batch_size, seq_len, tokenizer),
          print_profile=True,
          detailed=True,
      )
    else:
      inputs = bert_input_constructor((batch_size, seq_len), tokenizer)
      outputs = model(inputs)

在模型訓(xùn)練工作流中

要在訓(xùn)練工作流中對(duì)模型的前向過程進(jìn)行性能分析，請(qǐng)使用 FlopsProfiler 類。FlopsProfiler 類提供了以下方法：

start_profile() - 開始profiling。

get_total_flops(as_string=False)- 返回模型中的浮點(diǎn)操作的總數(shù)。

get_total_macs(as_string=False- 返回模型中的macs的總數(shù)。

get_total_params(as_string=False)- 返回模型中參數(shù)的總數(shù)。

print_model_profile(profile_step=1, module_depth=-1, top_modules=3, detailed=True, output_file=None)-打印模型profile。

stop_profile()-停止性能分析。這將停止模型中的浮點(diǎn)運(yùn)算計(jì)數(shù)。

end_profile()-進(jìn)行清理。這將清理在性能分析過程中添加到模型的性能分析屬性。這應(yīng)該在性能分析結(jié)束并且在調(diào)用get_total_flops、get_total_params或print_model_profile之后進(jìn)行。

訓(xùn)練工作流例子

以下是一個(gè)典型的訓(xùn)練工作流中使用該方法的示例。

from deepspeed.profiling.flops_profiler import FlopsProfiler

model = Model()
prof = FlopsProfiler(model)

profile_step = 5
print_profile= True

for step, batch in enumerate(data_loader):
  # start profiling at training step "profile_step"
  if step == profile_step:
    prof.start_profile()

  # forward() method
  loss = model(batch)

  # end profiling and print output
  if step == profile_step: # if using multi nodes, check global_rank == 0 as well
    prof.stop_profile()
    flops = prof.get_total_flops()
    macs = prof.get_total_macs()
    params = prof.get_total_params()
    if print_profile:
        prof.print_model_profile(profile_step=profile_step)
    prof.end_profile()

  # runs backpropagation
  loss.backward()

  # weight update
  optimizer.step()