このモデルは事前学習モデルなので、そのままではChataGPTなどの普通の人が良く使うLLMと違ってうまく指示を聞くようになっていませんがSFTなどの事後学習を行えばいろいろなタスクをこなせるようにすることが可能です。

この記事では、このような事後学習のやり方の一つとしてGoogle Colabで比較的安く使えるL4というGPUを使って、LoRAという方法で学習するやり方を紹介します。

今回紹介するコードは以下のところにありますので、参考にしてください。

https://github.com/shu65/plamo-2-8b-lora-sft-example/blob/main/PLaMo_2_8B_LoRA_SFT.ipynb

LoRAとは

LoRAとは「Low-Rank Adaptation」の略で、大規模言語モデルのような大きなモデルの学習を、メモリが小さいGPU１枚など小規模な計算機環境で学習するために提案された手法です。

具体的にはモデルの一部のLayerに対して低ランクの行列を導入し、低ランクの行列のみ学習することでモデル全体を学習するのと比べて非常に少ないメモリで学習できるようにしています。

このLoRAを使った学習はライブラリがそろっていることもあり、簡単なものはかなり短いコードで書くことができます。

PLaMo 2 8B のLoRA

ここからPLaMo 2 8BでLoRAを使って学習する方法を説明していきます。

PLaMo 2 8B 利用規約への同意

先ほども説明した通りPLaMo 2 8Bは特殊なライセンスということもあり、事前に同意しておく必要があります。

これにはまず、Hugging FaceのサイトのPLaMo 2 8Bのページに行きます。URLは以下の通りです。

https://huggingface.co/pfnet/plamo-2-8b

このページに行くとまだライセンスに同意していない場合は以下のようにライセンスの一部が表示されていると思います。

この場合はライセンスを確認の上、同意してください。同意すると以下のような表示になります。

これでPLaMo 2 8Bを使う準備ができました。

Google ColabのランタイムでL4を使うようにする

次に、Google ColabでL4を使う準備をします。Google ColabでL4が使えるように課金が必要になりますので、まずは課金をします。

課金についてはこちらをご覧ください。

https://colab.research.google.com/signup?hl=ja

今回のコードを動かすだけであれば「Pay As You Go」で100 コンピューティング ユニットを購入すれば十分です。この記事を執筆時点では1200円に満たない程度で購入できます。

課金が済んだら、メニューバーから「ランタイム」→「ランタイムのタイプを変更」をクリックします。すると無料枠では選択できないL4 GPUが選択できるようになっていると思うので、L4 GPUを選択します。

これでGPUを使う準備ができました。

必要パッケージのインストール

次に今回の学習で必要なパッケージをインストールします。コマンドとしては以下の通りです。

!pip install torch==2.4.1 torchvision==0.19.1 torchaudio==2.4.1 --index-url https://download.pytorch.org/whl/cu124
!pip install trl numba>=0.60.0 mamba-ssm>=2.2.2 causal-conv1d>=1.4.0 transformers>=4.44.2

最初にPytorchのバージョンを少し下げていますが、これはPLaMo 2の中で使われているライブラリの一部が最新のPyTorchに対応させるのが結構大変なため、簡単に実行できるようにするために少し古いPyTorchを入れています。

上記のコマンドを実行したあとは、以下のバージョンになっていました。

causal-conv1d                      1.5.0.post8
mamba-ssm                          2.2.4
numba                              0.60.0
numba-cuda                         0.2.0
sentence-transformers              3.4.1
torch                              2.4.1+cu124
torchaudio                         2.4.1+cu124
torchsummary                       1.5.1
torchvision                        0.19.1+cu124
transformers                       4.50.0

LoRAのコード

パッケージをインストールしたら次は以下のようにHugging Faceにログインします。

from huggingface_hub import login

login()

これを実行するとHugging Faceのtoken を聞かれますのでHugging Faceのtokeを入力してください。

次に各パッケージをimportしておきます。

from transformers import AutoModelForCausalLM
from transformers import AutoTokenizer
import datasets
import string
from trl import DataCollatorForCompletionOnlyLM
from trl import SFTConfig, SFTTrainer
from peft import LoraConfig
import torch

そして、PLaMo 2 8Bのモデルとtokenizerを以下のようにロードします。

model_name = "pfnet/plamo-2-8b"
model = AutoModelForCausalLM.from_pretrained(model_name, trust_remote_code=True, torch_dtype=torch.bfloat16).to("cuda")
tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)

「PLaMo 2 8B 利用規約への同意」の部分の手順ができていなかったり、Hugging Faceのログインがうまくできていないと、この部分でエラーがでると思われます。その場合は利用規約の同意ができているかや、正しくHugging Faceのログインができているかなどを確認してください。

次に今回使うinstructionデータをダウウンロードして、前処理します。今回はkunishou/databricks-dolly-15k-ja のinputがないデータだけを取り出して利用します。

dataset = datasets.load_dataset("kunishou/databricks-dolly-15k-ja")
train_dataset = dataset["train"].filter(lambda data: "instruction" in data and "output" in data and data["input"] == "").select(range(2000))

data_collator = DataCollatorForCompletionOnlyLM(
    response_template=tokenizer.encode(" Answer:\n", add_special_tokens=False),
    tokenizer=tokenizer
)

次にLoRAとSFTの引数を指定します。PLaMo 2特有の部分としてLoraConfig でLoRAを使って学習するレイヤーを指定するtarget_modules という引数があります。ここでLLaMa系のようなTransformerの場合、AttensionのQeury, Keyを作るLinearレイヤーをLoRAで学習することが多い印象なので、PLaMo 2でも同じようにAttensionのQueryとKeyを作るLinearレイヤーをLoRAで学習するようにします。PLaMo 2の場合は、Qeury, Key、Valueを作るLinearレイヤーをすべてまとめたqkv_proj というLinearレイヤーがありますので、このqkv_proj をtarget_modules に指定しています。

peft_config = LoraConfig(
    task_type="CAUSAL_LM",
    target_modules=[
        "qkv_proj",
    ],
)

sft_args = SFTConfig(
    output_dir="./outputs",
    evaluation_strategy="no",
    per_device_train_batch_size=1,
    gradient_accumulation_steps=4,
    learning_rate=5e-5,
    num_train_epochs=1.0,
    lr_scheduler_type="cosine",
    warmup_ratio=0.3,
    logging_steps=10,
    save_strategy="epoch",
    report_to="tensorboard",
    bf16=True,
    max_seq_length=1024,
    gradient_checkpointing=True,
)

あとはデータのサンプルに対してフォーマットに合わせて１つのテキストを生成するformatting_func を定義します。今回は以下のようなものを使います。

INSTRUCTION_TEMPLATE = string.Template(
    """### Question:
{input} ### Answer:
{response}<|plamo:eos|>
"""
)

def formatting_func(example):
  text = INSTRUCTION_TEMPLATE.substitute(input=example["instruction"], response=example["output"])
  return text

最後に、ここまで用意したものをSFTTrainer に渡して学習を開始します。

trainer = SFTTrainer(
    model=model,
    args=sft_args,
    peft_config=peft_config,
    data_collator=data_collator,
    train_dataset=train_dataset,
    formatting_func=formatting_func,
)

trainer.train()
trainer.save_model()
     

今回の設定では500イテレーション回ることになるはずです。私が試した限りは15分程度で処理が完了しました。

できたモデルの出力が正しいかは以下のコードで確認できます。

import torch

# プロンプトの準備
prompt = "### Question:\n埼玉の県庁所在地は何市？\n\n### Answer:\n"

# 推論の実行
inputs = tokenizer(prompt, return_tensors="pt").to("cuda")
generated_tokens = trainer.model.generate(
    **inputs,
    max_new_tokens=64,
    pad_token_id=tokenizer.pad_token_id,
)[0]
generated_text = tokenizer.decode(generated_tokens)
print(generated_text)

おそらく以下のような出力がでるはずです

<|plamo:bos|>### Question:
埼玉の県庁所在地は何市？

### Answer:
さいたま市<|plamo:eos|>

終わりに

いかがだったでしょうか？今回はPLaMo 2 8Bに対してLoRAによって学習する方法を紹介しました。LoRAを使えばメモリが少ないL4のようなGPUでも8Bモデル程度で学習することができます。今回のコードを使えば事後学習が簡単にできると思われますので、みなさんもいろいろ試していただければと思っています。無料枠で使えるGPUのT4でも量子化などを頑張れば8Bモデルも学習できる気がしますが、うまくできそうであればそちらも記事にしようと思います。

この他にもPLaMoを含めたLLMの技術も紹介できればと思っています。

The post PFNの最新LLM PLaMo 2 8BをGoogle ColabでLoRAで学習してみる first appeared on まったり勉強ノート.

Google Colabの無料枠で推論を回す方法は前回記事にしましたので、そもそもPLaMo 2 1Bって何と思った方や推論を回してみたいという方はそちらをご覧ください。

また、今回説明に使うコードはこちらに置いてありますので、適宜参照してください。

https://github.com/shu65/plamo-2-1b-sft-example

Google Colabにおける一連の実行に関してはJupyter Notebookにまとめてありますので、細かい実行方法がわからないという方はこちらをご覧ください

https://github.com/shu65/plamo-2-1b-sft-example/blob/main/run_sft_google_colab.ipynb

Supervised Fine-Tuning(SFT)とは？

SFTを知らない方に簡単に説明すると、SFTは指示と想定されている回答のペアを用意し、LLMに対して学習を行い、指示に従いやすいモデルを作る方法になります。

特にPLaMo 2 1Bのような事前学習モデルでは、特に指示に従うように学習されていないケースもあり、そのまま利用した際、余計なことをだらだらと出力し続けたり、頓珍漢な回答が返ってきたりという問題が発生することがあります。

このため指示に適切にこたえてもらうための技術がいろいろあるのですが、そのうちの一つにSFTというものがあります。

Google ColabでPLaMo 2 1BをSFTする

それでは本題のGoogle ColabでPLaMo 2 1BをSFTする方法について説明します。今回はGPUメモリの関係上、おそらく無料で使えるT4だと無改造では実行できない気がするのでL4を使った説明をします。

L4 GPUの利用

まず、Google ColabでL4が使えるように、課金が必要になります。

課金についてはこちらをご覧ください。

https://colab.research.google.com/signup?hl=ja

今回のコードを動かすだけであれば「Pay As You Go」で100 コンピューティング ユニットを購入すれば十分です。この記事を執筆時点では1200円に満たない程度で購入できます。

課金が済んだら、メニューバーから「ランタイム」→「ランタイムのタイプを変更」をクリックします。すると無料枠では選択できないL4 GPUが選択できるようになっていると思うので、L4 GPUを選択します。

これでGPUを使う準備ができました。

実行環境準備

L4を利用するようにしたら、実行するコードのダウンロードやPythonパッケージのインストールを行います。

まずGithubよりコードをcloneしてきます

!git clone https://github.com/shu65/plamo-2-1b-sft-example.git

次に、PyTorchのバージョンを現在の最新版よりも前の以下のものに変更します。

!pip install torch==2.4.1 torchvision==0.19.1 torchaudio==2.4.1 --index-url https://download.pytorch.org/whl/cu124

この後は以下のようにPyTorch以外のPLaMo 2 1Bの実行に必要なパッケージやSFTに必要なパッケージなどをインストールします。

!pip install -r plamo-2-1b-sft-example/requirements.txt

ここまで実行すると2025/02/12現在以下のようなバージョンがインストールされました。

causal-conv1d                      1.5.0.post8
fastrlock                          0.8.3
mamba-ssm                          2.2.4
numba                              0.61.0
numba-cuda                         0.0.17.1
sentence-transformers              3.4.1
torch                              2.4.1+cu124
torchaudio                         2.4.1+cu124
torchsummary                       1.5.1
torchvision                        0.19.1+cu124
transformers                       4.48.2
trl                                0.14.0

これであとはSFTのコードを実行すれば、SFTをすることができます。このSFTの中身に関しては次で紹介していきます。

PLaMo 2 1BをSFTする

SFTをする部分はsft.py　というスクリプトにまとめてあります。このスクリプトの重要な部分について簡単にですが説明していきます。

まず、今回はすぐに実行が終わるように少量の質問と回答のペアのデータを用います。

今回は日本語の指示学習でよく使われるkunishou/databricks-dolly-15k-jaというデータセットのうち、input がなくinstruction とoutput のペアになっているデータのみを取り出しその一部だけを利用します。一つ例を見せると以下のようなデータを利用します。

{
  "output": "イコクエイラクブカ",
  "input": "",
  "index": "1",
  "category": "classification",
  "instruction": "魚の種類はどっち？イコクエイラクブカとロープ"
}

一部だけ取り出すコードは以下の通りです。

    dataset = datasets.load_dataset("kunishou/databricks-dolly-15k-ja")
    train_dataset = dataset["train"].filter(lambda data: data["input"] == "")

次にSFTConfig というSFTの実行の設定のクラスのインスタンスを用意します。具体的には以下の通りです。

    sft_args = SFTConfig(
        output_dir="./outputs",
        evaluation_strategy="no",
        per_device_train_batch_size=1,
        gradient_accumulation_steps=4,
        learning_rate=5e-5,
        num_train_epochs=0.1,
        lr_scheduler_type="cosine",
        warmup_ratio=0.3,
        logging_steps=10,
        save_strategy="epoch",
        report_to="tensorboard",
        bf16=True,
        max_seq_length=1024,
        gradient_checkpointing=True,
        deepspeed='./deepspeed_config.json',
    )

重要なこととして、今回はGPUのメモリが少ないため、DeepSpeedのStage 3という学習時に一部のデータをCPU側に置いておくモードを利用します。

これによりGPUメモリが少ない環境でもSFTを回すことができます。

DeepSpeed周りの設定はdeepspeed_config.json に書いてありますので気になる方はご覧ください。

また、今回は学習データの10%だけを利用するようにしています。これはこの学習を早く終わらせるためであり、本来はもっと回す必要があると考えられますので、本気でSFTをする場合は注意してください。

次にデータをどのようなフォーマットでLLMに入力するかを指定するformatting_func という関数を用意します。今回は以下のようにしました。

INSTRUCTION_TEMPLATE = string.Template(
    """### Question:
${input}

### Answer:
${response}<|plamo:eos|>
"""
)


def formatting_func(examples):
    output_texts = []
    for i in range(len(examples['instruction'])):
        text = INSTRUCTION_TEMPLATE.substitute(input=examples['instruction'][i], response=examples['output'][i])
        output_texts.append(text)
    return output_texts

INSTRUCTION_TEMPLATE が今回のフォーマットで、### Question:\n の後に指示、### Answer:\n のあとに回答が続き、最後にend of sequenceである<|plamo:eos|> が来るようになっています。

また、学習時には回答部分だけを学習してほしいので、どこからが回答かがわかるように‎DataCollatorForCompletionOnlyLM のインスタンスも用意します。これは以下の通りです。

    data_collator = DataCollatorForCompletionOnlyLM(
        response_template=tokenizer.encode(" Answer:\n", add_special_tokens=False),
        tokenizer=tokenizer
    )

response_template のところで回答前の部分がどのようなtoken idになるかを指定する部分があるので、上記のように指定します。前後の文字の影響で指定したtoken idが出現しないケースがあるので、その時はいろいろresponse_template に指定する文字列を調整してみてください。

最後にSFTを実行するためのクラスの‎SFTTrainer を以下のように用意します。

    trainer = SFTTrainer(
        model=model,
        args=sft_args,
        data_collator=data_collator,
        train_dataset=train_dataset,
        tokenizer=tokenizer,
        formatting_func=formatting_func,
    )

そして、以下のように実行し、結果を保存します。

    trainer.train()
    trainer.save_model()

これで学習が終わるとSFTConfig のoutput_dir で指定した./outputs に結果が出力されます。試しに私がGoogle Colabで実行した際は13分程度で学習が終わりました。コンピューティングユニットとしてはパッケージなどのインストールも含めて4だけ消費しました。

SFTされたモデルで推論してみる

最後にSFTされたモデルで推論するというのを行います。

これはPLaMo 2 1Bのexampleとほぼ同じでpromptだけ少し変えたものを例として用います。コードとしては以下の通りです。

from transformers import AutoTokenizer, AutoModelForCausalLM
import torch


model_name = "./plamo-2-1b-sft-example/outputs"

tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained(model_name, trust_remote_code=True)


# プロンプトの準備
prompt = "### Question:\n埼玉の県庁所在地は何市？\n\n### Answer:\n"

# 推論の実行
inputs = tokenizer(prompt, return_tensors="pt")
generated_tokens = model.generate(
    **inputs,
    max_new_tokens=64,
    pad_token_id=tokenizer.pad_token_id,
)[0]
generated_text = tokenizer.decode(generated_tokens)
print(generated_text)

出力結果は以下のようになります。

<|plamo:bos|>### Question:
埼玉の県庁所在地は何市？

### Answer:
埼玉県の県庁所在地はさいたま市です。<|plamo:eos|>

ちゃんと学習で指定されたように### Answer:\n の後に質問に対する回答をし、その後<|plamo:eos|> を出力するということができています。

ちなみにSFTしていないモデルではどうなるかというと、以下のように余計なことを出力するうえ、出力が止まらないという状態になっています。

<|plamo:bos|>### Question:
埼玉の県庁所在地は何市？

### Answer:
さいたま市

### 解説
「県庁所在地」とは、都道府県庁が置かれている都市のことです。
「さいたま市」は埼玉県の県庁所在地です。

### 関連記事
### 取り急ぎお知らせ
「埼玉の県庁所在地は何市？」の解説は以上です。
「埼玉の県庁所在地は何市？」の解説は以上です。

このため、SFTでうまくフォーマットに従うよう学習できたと考えられます。

終わりに

今回はPLaMo 2 1Bを使ってSFTをする例を示しました。今回示したように簡単なSFTなら十分Google Colabで実行することができます。みなさんもぜひいろいろ試していただければと思います。

The post 小型LLM PLaMo 2 1BをGoogle ColabでSFTしてみる first appeared on まったり勉強ノート.

HistogramはCUDA Samplesの中にもありますが、全然違う実装が、NVIDIAのA100の最適化に関する発表の資料の中で紹介されています。この資料ではatomicAddとL2キャッシュの「persistent data accesses」を利用してhistogramの実装をしています。このpersistent data accessesは発表当時気になっていて、あとで調べようと思って忘れてたのですが、最近調べたのでせっかくなので記事にしました。

今回調査するうえで特に知りたかった点としてpersistent data accessesを使う場合とshared memoryを使う場合だとどっちが速いのかというものがあります。A100の最適化の資料の中にはこれについて特に書いてなかったので、この二つの性能にどれくらい差があるのかの比較を行いました。

検証で使ったコードはこちらにあげてあります。

https://github.com/shu65/cuda-histogram

Histogramと今回対象とする部分に関して

Histogramに関して知っている方も多いと思いますが、どういうものか簡単に紹介します。Histogramはデータの範囲をいくつかのbinに区切り、データの中の各要素がどのbinに含まれるかを計算し、binごとに含まれる要素の個数をカウントするというものになります。

上記の説明の通り、histogramを計算するうえで、大まかに3つくらいのステップに分けることができます。

1. 各binの範囲を決める
2. データの各要素がどのbinに入るのかを計算する
3. bin毎に何個のデータの要素が含まれるかをカウントする

このHistogramは入力データによってどういう風にbinの範囲を決めればよいかが変わるため、データに応じて1,2あたりの処理はデータに応じて変化させる必要があります。また、GPU的にも難しいのは3のところなため、今回は3に注目して説明します。

3の部分は入力としては各要素がどのbinに入るかを表したbinのidの配列を受け取り、bin毎に何個要素があるかをカウントするという処理になります。どういう処理かイメージしやすいようにCPU版のコードを以下に示します。

#include <stdint.h>

void HistogramCPU(const int *data, const uint32_t n, const uint32_t n_bins, uint32_t *bin_counts)
{
  for (uint32_t i = 0; i < n_bins; ++i)
  {
    bin_counts[i] = 0;
  }

  for (uint32_t i = 0; i < n; ++i)
  {
    const int bin_i = data[i];
    ++bin_counts[bin_i];
  }
}

このあと、このコードと同じ結果になるようなGPUコードを紹介していきます。

GPUでHistogramが難しい理由

先にGPUでhistogramを計算する際、難しいポイントに関して触れておきます。GPUに限らず並列処理でhistogramを計算する際、bin毎のカウントをするところで複数のスレッドが同じメモリ領域にアクセスすることになるので、bin毎のカウント部分で排他制御が必要になります。

GPUで簡単に実装するなら、後ほど示す通りatomicAddを使えばいいのですが、atomicAddは遅いという問題があります。特にglobal memoryに対してのatomicAddはshared memoryに対するものよりも遅いです。このため、個人的にはatomicAdd、特にglobal memoryに対するものは注意が必要な計算という認識でした。

それがA100の最適化の資料でL2キャッシュのpersistent data accessesを使うとましになるよ、ということが書かれています。次にこのL2キャッシュのpersistent data accessesについて詳しく説明します。

L2キャッシュのpersistent data accessesについて

L2キャッシュのpersistent data accessesは、L2キャッシュのメモリ領域を分割してpersistent data用の領域を確保して、よくアクセスするものはpersistent data用の領域にキャッシュしてメモリアクセスを高速化するための機能です。

A100最適化の資料の18ページ目あたりからこの機能の紹介があります。

https://developer.download.nvidia.com/video/gputechconf/gtc/2020/presentations/s21819-optimizing-applications-for-nvidia-ampere-gpu-architecture.pdf

CUDAのprogramming guideでは以下の部分に説明があります。

https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html#device-memory-l2-access-management

これを使うと、よくデータアクセスする一部の領域とそれ以外の領域のキャッシュを分けることができます。結果として一部だけ何度もアクセスするという場合はこの機能を使うことで高速化が狙えます。

制限としてはL2キャッシュのすべてをpersistent dataにすることはできず、最大値が決まっています。最大値は以下のようにすると確認できます。

  cudaDeviceProp prop;
  CheckCudaErrors(cudaGetDeviceProperties(&prop, device_id));
  cout << "persistingL2CacheMaxSize:" <<   prop.persistingL2CacheMaxSize << endl;

A100で確認すると30MBが最大値になっています。

使い方としてはprogramming guideにある通り、以下の手順で使うことができます。

Persistent data accesses用の領域を確保

以下のようにpersistent data accesses用の領域として最大どれくらい使うかを設定します。コード中のsize にpersistent data accesses用の領域のサイズを入れてcudaDeviceSetLimitを呼ぶことで、使用するpersistent data accesses用の領域の最大値を設定します。

cudaDeviceSetLimit(cudaLimitPersistingL2CacheSize, size); 

Persistent data accessesの設定適用

次にstream、もしくはcuda graphのnodeに対してpersistent data accessesの設定を行います。ここではstreamに対しての設定方法を示します。programming guideにある通り、以下のように設定していきます。

cudaStreamAttrValue stream_attribute;                                         // Stream level attributes data structure
stream_attribute.accessPolicyWindow.base_ptr  = reinterpret_cast<void*>(ptr); // Global Memory data pointer
stream_attribute.accessPolicyWindow.num_bytes = num_bytes;                    // Number of bytes for persistence access.
                                                                              // (Must be less than cudaDeviceProp::accessPolicyMaxWindowSize)
stream_attribute.accessPolicyWindow.hitRatio  = 1.0;                          // Hint for cache hit ratio
stream_attribute.accessPolicyWindow.hitProp   = cudaAccessPropertyPersisting; // Type of access property on cache hit
stream_attribute.accessPolicyWindow.missProp  = cudaAccessPropertyStreaming;  // Type of access property on cache miss.

//Set the attributes to a CUDA stream of type cudaStream_t
cudaStreamSetAttribute(stream, cudaStreamAttributeAccessPolicyWindow, &stream_attribute);

注意する点としてはhitRatio の値です。hitRatio はアクセスするglobal memoryのサイズがnum_bytesよりも大きい場合は適切に指定しないとパフォーマンスが落ちることが以下の部分で示されています。

https://docs.nvidia.com/cuda/cuda-c-best-practices-guide/#tuning-the-access-window-hit-ratio

このため、hitRatio は自分のケースでどのくらいのサイズにすべきか？を考えて設定したほうがよさそうです。

ただ、今回のhistogramの例では後ほど示す通り、global memoryで最大20MB分の領域をpersistent data accessesに利用するので、persistent data accessesで指定できるサイズに収まります。このため、hitRatio は1.0でOKです。

HistogramのGPU実装

ここからは今回検証に使うhistogramのGPU実装に関してです。3つありますので、順番にどういうものかを説明していきます。

GPU実装のベースライン

まずはGPU実装のベースラインです。コードとしてはCPUをそのままCUDAで実装したような形になっています。

__global__ void HistogramGPUv1Kernel(const int *data, const uint32_t n, uint32_t *bin_counts)
{
  const int tid = blockIdx.x * blockDim.x + threadIdx.x;
  if (tid >= n)
  {
    return;
  }
  const int bin_i = data[tid];
  atomicAdd(bin_counts + bin_i, 1);
}

こちらのコードがA100の最適化の資料で示されているhistogramのコードとほぼ同じものになっています。bin_countsへのアクセスは全スレッドが同時に行うため、atomicAddを使って排他制御しながらカウントするようにしています。

GPU実装のshared memory版

Histogramの計算でbin_countsがshared memoryに収まる範囲であれば、shared memoryを使うという手があります。

先ほど説明した通り、shared memoryへのatomicAddはglobal memoryに比べて速いので、shared memoryを使ってblock毎に集計し、その後各blockの結果をatomicAddを使ってglobal memoryの領域に加算するという方法で計算します。こうすることでglobal memoryへのatomicAddの回数は減らすことができます。コードとしては以下の通りです。

__global__ void HistogramGPUv2Kernel(const int *data, const uint32_t n, const uint32_t n_bins, uint32_t *bin_counts)
{
  cg::thread_block cta = cg::this_thread_block();
  extern __shared__ uint32_t s_bin_counts[];
  const int tid = blockIdx.x * blockDim.x + threadIdx.x;
  const int stride = gridDim.x * blockDim.x;
  if (threadIdx.x < n_bins)
  {
    s_bin_counts[threadIdx.x] = 0;
  }
  cg::sync(cta);
  for (int i = tid; i < n; i += stride)
  {
    const int bin_i = data[i];
    atomicAdd(s_bin_counts + bin_i, 1);
  }
  cg::sync(cta);
  if (threadIdx.x < n_bins)
  {
    uint32_t sum = s_bin_counts[threadIdx.x];
    atomicAdd(bin_counts + threadIdx.x, sum);
  }
}

注意点としてはshared memoryのサイズは最大でA100の場合でも164KBらしいので、bin_counts に必要なサイズがこれ以上のときはこの戦略はそのまま使うことができません。

GPU実装のshared memory + reduction版

shared memory版では最後global memoryへの加算はatomicAddを利用しましたが、この部分もatomicAddなしで実行するようにします。具体的にはCUDAにおけるreductionのような戦略をとり、各blockが計算した結果をCUDAにおけるparallel reductionに似たアルゴリズムで集計します。コードとしては以下の通り。

__global__ void HistogramGPUv3Kernel(const int *data, const uint32_t n, const uint32_t n_bins, uint32_t *tmp_bin_counts)
{
  cg::thread_block cta = cg::this_thread_block();
  extern __shared__ uint32_t s_bin_counts[];
  const int tid = blockIdx.x * blockDim.x + threadIdx.x;
  const int tmp_bin_counts_offset = blockIdx.x * n_bins;
  const int stride = gridDim.x * blockDim.x;
  if (threadIdx.x < n_bins)
  {
    s_bin_counts[threadIdx.x] = 0;
  }
  cg::sync(cta);
  for (int i = tid; i < n; i += stride)
  {
    const int bin_i = data[i];
    atomicAdd(s_bin_counts + bin_i, 1);
  }
  cg::sync(cta);
  if (threadIdx.x < n_bins)
  {
    uint32_t sum = s_bin_counts[threadIdx.x];
    tmp_bin_counts[tmp_bin_counts_offset + threadIdx.x] = sum;
  }
}

__global__ void HistogramGPUv3MergeKernel(const uint32_t *tmp_bin_counts, const int n, uint32_t *bin_counts)
{
  cg::thread_block cta = cg::this_thread_block();
  extern __shared__ uint32_t s_data[];

  uint32_t sum = 0;
  for (int i = threadIdx.x; i < n; i += blockDim.x)
  {
    sum += tmp_bin_counts[blockIdx.x + i * blockDim.x];
  }
  s_data[threadIdx.x] = sum;
  for (uint stride = blockDim.x / 2; stride > 0; stride >>= 1)
  {
    cg::sync(cta);
    if (threadIdx.x < stride)
    {
      s_data[threadIdx.x] += s_data[threadIdx.x + stride];
    }
  }

  if (threadIdx.x == 0)
  {
    bin_counts[blockIdx.x] = s_data[0];
  }
}

ちなみに、アルゴリズム的にこれはほぼCUDA Samplesに含まれるhistogramと同じになります。

評価

今回、binの数で傾向が変わったので、以下の2種類のデータで比較します。

1. 入力データ数は256M個、binの数が256
2. 入力データ数は256M個、binの数が5M個

1つ目のほうがCUDA Samplesに含まれるhistogramの条件に近いもので、2つ目がA100の最適化の資料に書かれている条件になります。また、1のほうは先ほど紹介したアルゴリズムすべてが実行できますが、2つ目のほうはshared memoryが足りないのでベースラインのみとなっています。

また、persistent data accessesのありなしでどれくらい計算結果が変化するのかも知りたいので、各アルゴリズムでbin_countsの部分にpersistent data accessesを使う場合と使わなかった場合も比較します。

時間の計測方法としては10回の平均時間を算出して比較します。実行環境としてはCUDA 12.0、A100を利用しています。

計測した計算時間はそれぞれ以下の通りです。

表からわかる通り、1のデータではglobal memoryへのatomicAddが少なければ少ないほど高速化できていることがわかります。また、persistent data accessesは1のデータでは効果がありませんでした。

また、2のデータに関してはpersistent data accessesありなしで若干差がありますが、今回は約7%の向上と効果は小さいという結果になりました。A100の最適化のほうの資料では43%向上とあるのでどこか設定を間違えているのかもしれません。（いろいろ試しましたがわからなかったのでご存じの方いたら教えていただけるとありがたいです。）

ただ、結果からshared memoryが使える状況下であればL2キャッシュのチューニングするよりもshared memoryを使ったほうが速くなりそうという印象を持ちました。

終わりに

今回、気になっていたatomicAddとL2キャッシュのpersistent data accessesのパフォーマンスについて調べました。結果としてやっぱりglobal memoryへのatomicAddはできるだけ避けたほうがいいということが確認できてよかったです。

これまでも最近のGPU、CUDAを使ってreduction、vectorized memory accessに関しても調査したまとめを書いたのでもしよろしければそちらもご覧ください。

The post CUDAの高速化の復習2023年版 Histogram（主にatomicAdd）編 first appeared on まったり勉強ノート.

このvectorized memory accessは昔からあるテクニックです。ただ、最近CUDAの高速化をしようとして、vectorized memory accessを試してみるのですが、いまいち効果がなさそうな気配があったので、ちゃんと調べようと思い今回記事をかきました。ちなみに結論からいうと今でもちょっとは効果ありそうでした。

検証に利用したコードはこちらにあげてあります。

https://github.com/shu65/cuda-vectorized-memory-access

検証環境はCUDA 12.0、GPUはA100を使っています。

今回のVectorized Memory Accessは少しマニアックなテクニックなので、CUDAの高速化全般に関して簡単に知りたいという方はReductionの記事のほうがおすすめです。リンクは以下の通りです。

Vectorized Memory Accessとは？

Vectorized Memory AccessとはCUDAにおいて連続するグローバルメモリへのアクセスを高速化するテクニックの一つです。このテクニックは結構昔から知られていてNVIDIAのblogでも2013年に紹介されています。

https://developer.nvidia.com/blog/cuda-pro-tip-increase-performance-with-vectorized-memory-access/

詳細はこちらのNVIDIAの記事を見ていただきたいと思いますが、ざっくり簡単に説明すると、連続したグローバルメモリにアクセスする際にintなど32 bit単位でアクセスするよりもint2やint4でアクセスするほうが速いよ、というものです。int2、int4はCUDAで定義されている構造体でintを2つ、または4つもった構造体です。なので、普通のintが32 bitなのにたいしてint2だと64 bit, int4だと128 bitのサイズになってintよりも大きいデータに一気にアクセスすることになります。

Vectorized Memory Accessの検証コード

今回、配列の要素数を1Kから1Gまで増加させたとき、配列の全要素を別の配列にコピーする単純なカーネルで測定します。Vectorized memory accessを使わない場合は以下のようなコードになります。

template <typename T>
__global__ void CopyScalarKernel(T *d_in, T *d_out, const size_t n)
{
  const int idx = blockIdx.x * blockDim.x + threadIdx.x;
  const int stride = blockDim.x * gridDim.x;
  for (int i = idx; i < n; i += stride)
  {
    d_out[i] = d_in[i];
  }
}

template <typename T>
void CopyScalar(T *d_in, T *d_out, size_t n)
{
  int max_blocks = 4096;
  int threads = 1024;
  int blocks = min((int)(n + threads - 1) / threads, max_blocks);
  CopyScalarKernel<T><<<blocks, threads>>>(d_in, d_out, n);
}

thread数やblock数を変化させるとパフォーマンスが若干変化するのですが、ここのチューニングするのは大変なので、すべて同じ方法で決めて使います。

このコードでvectorized memory accessを使って64 bit, 128 bitでアクセスするときはこのようなカーネルになります。

template <typename T>
__global__ void CopyVector2Kernel(T *d_in, T *d_out, const size_t n)
{
  const float ratio = ((float)sizeof(int2)) / sizeof(T);
  const int idx = blockIdx.x * blockDim.x + threadIdx.x;
  const int stride = blockDim.x * gridDim.x;
  const int m = n / ratio;
  for (int i = idx; i < m; i += stride)
  {
    reinterpret_cast<int2 *>(d_out)[i] = reinterpret_cast<int2 *>(d_in)[i];
  }
}

template <typename T>
__global__ void CopyVector4Kernel(T *d_in, T *d_out, const size_t n)
{
  const float ratio = ((float)sizeof(int4)) / sizeof(T);
  const int idx = blockIdx.x * blockDim.x + threadIdx.x;
  const int stride = blockDim.x * gridDim.x;
  const int m = n / ratio;
  for (int i = idx; i < m; i += stride)
  {
    reinterpret_cast<int4 *>(d_out)[i] = reinterpret_cast<int4 *>(d_in)[i];
  }
}

重要な点として、グローバルメモリからデータと読み込むときと書き込むときでint2やint4など大きい型のポインタにキャストしてからアクセスするということをしています。

Vectorized Memory Accessの結果

では、さきほどのコードを動かして実際にどのくらいのスループットになるかを示します。計測する際は10回の平均時間を出してスループットを算出しました。比較には最近よく使う、halfの配列とfloatの配列の２種類を使います。そしてデータにアクセスするときは、何もしないscalerのまま、32 bit, 64 bit, 128 bitでアクセスする場合の合計4つを示します。

halfとfloatの結果のグラフを以下に示します。

結果を見るとコピーするサイズが小さいときはvectorized memory accessなし、ありでそれほど差がなく、数MBくらいでちょっとずつ差がでるという感じの結果でした。ちなみにfloatで32bitでアクセスするとscalerよりも遅くなっていますが、これはキャストのオーバーヘッドがあるためだと思われます。

やってみた感想としては今も多少は効果があるけど、そこまで劇的に変化するわけではなさそうという印象です。なので、最適化をできるだけ頑張って、もう次やることがないってなったときに試してみるくらいでよいかなということを思いました。

終わりに

今回は昔からあるCUDAの高速化テクニックの一つのvectorized memory accessが今でも有効なのか確認したので、そのまとめを書きました。CUDAのコンパイラやGPUのアーキテクチャもどんどん変化しているので、昔は効果あったけど今はない、ってものも少なからずあるので、今後もこういう高速化テクニックの確認をしていければと思います。

The post CUDAの高速化の復習2023年版 Vectorized Memory Access編 first appeared on まったり勉強ノート.

私はCUDAを2009年のころから研究で使っていました。当時は頑張って勉強していたので自分の研究分野以外のCUDA応用の論文などを読んで高速化テクニックを勉強していました。ただ、2015年に企業に就職して以降、GPUを使うことはあってもCUDAのコードを直接自分で書くことはほとんどなくなってしまいました。このため、最近CUDAのコードを速くするためにどうすればいいのか？みたいな議論のときに、「昔は～」みたいな老害なコメントしかできない状態になってしまっていました。

この状態はさすがにまずいということでCUDAの勉強をし直そうと思い、この記事はそのまとめの第一弾でreductionをテーマで勉強した内容のまとめになります。私のように昔（2009年ごろ）CUDAを勉強したけど最近のCUDAわからんって人向けに記事は書いています。

今回測定のために書いたコードはこちらです。

https://github.com/shu65/cuda_reduction

この記事にはほとんど登場しないですが、昔風に書いたコードも一緒に含めています。

また元にしたNVIDIAのCUDA Samplesはこちらです。

https://github.com/NVIDIA/cuda-samples/tree/v11.8/Samples/2_Concepts_and_Techniques/reduction

計算時間に関してはCUDA Versionは12.0、GPUはV100を使った結果になります。計測はwarmupとしてreductionを実行したあと100回測定して平均計算時間を示しています。reductionの要素数は32M個に固定して測定します。また、reductionの計算では最初の段階はGPUを使いますが、CUDAの各thread blockの結果をさらにreductionする段階では簡単のためにCPUを使うことにします。

Reductionとは？

Reductionとは配列を入力としてとり、配列の全要素の和などを計算する処理になります。C++のコードを見るのが一番わかりやすいと思うので、以下にC++のコードを示します。

int reduce(const int *in, size_t n)
{
    int ret = 0;
    for (size_t i = 0; i < n; ++i)
    {
        ret += in[i];
    }
    return ret;
}

「和など」と書きましたが、各要素に対して使える演算は加算以外にも積、max, minなどいくつかあります。今回は説明を簡単にするために、加算で説明します。

そもそもなんでreductionをテーマにしようとしたかというと、私が勉強した当時、CUDAの高速化のテクニックがいろいろ詰まった題材としてreductionがよく紹介されていて、私もreductionを写経して勉強した経験がありました。このため、昔と現在の違いが分かりやすいということで選びました。

ちなみに昔の私が読んでたNVIDIAのreductionの資料はまだ公開されているようです。

https://developer.download.nvidia.com/assets/cuda/files/reduction.pdf

現在のCUDA Samplesで公開されているコードも基本的にはこの高速化テクニックに沿って実装されているようなので、この記事でも同じように沿って説明します。

Reductionを高速化していく

Reduce1: Baseline

まずはできるだけシンプルなreducetionのコードを最近のCUDA Sampleのコードを参考にしながら書いたコードを示します。

__global__ void reduce_gpu_v1_kernel(const int *g_in, size_t n, int *g_out)
{
    cg::thread_block cta = cg::this_thread_block();
    extern __shared__ int sdata[];
    int tid = threadIdx.x;
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i < n)
    {
        sdata[tid] = g_in[i];
    }
    else
    {
        sdata[tid] = 0;
    }

    cg::sync(cta);

    for (int s = 1; s < blockDim.x; s *= 2)
    {
        if (tid % (2 * s) == 0)
        {
            sdata[tid] += sdata[tid + s];
        }
        cg::sync(cta);
    }
    if (tid == 0)
    {
        g_out[blockIdx.x] = sdata[0];
    }
}

このコードを見てblock内のスレッドの同期をする関数として__syncthreads() を使うんじゃないの？と思った方、いますよね？ちなみに私は思いました。昔の資料を参考にすると昔のreductionのコードは以下のような感じでした。

__global__ void reduce_gpu_old_v1_kernel(const int *g_in, size_t n, int *g_out)
{
    extern __shared__ int sdata[];
    int tid = threadIdx.x;
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i < n)
    {
        sdata[tid] = g_in[i];
    }
    else
    {
        sdata[tid] = 0;
    }

    __syncthreads();

    for (int s = 1; s < blockDim.x; s *= 2)
    {
        if (tid % (2 * s) == 0)
        {
            sdata[tid] += sdata[tid + s];
        }
        __syncthreads();
    }
    if (tid == 0)
    {
        g_out[blockIdx.x] = sdata[0];
    }
}

ほとんど同じですが、同期の部分で昔のコードでは__syncthreads() を使っているのに対して最近のコードではcooperative_groups:sync() を使うようになっています。

最近のCUDAを追ってない方はCooperative Groupsって何？と思った方もいると思うので簡単に説明します。Cooperative GroupsはCUDA 9から導入されたもので、様々な単位でスレッドの同期などを行うための仕組みになります。

CUDAでは通常スレッドのまとまりとしてblock、gridなどの単位があります。ただ、同期ができるスレッドの単位はこれまでblockくらいで、CUDAのカーネル関数内で他の単位で同期するのは結構面倒でした。

この問題を解決して、様々な単位、例えばblockよりも少ないスレッド数や、grid単位で同期したりできる仕組みがCooperative Groupsです。より詳しく知りたい方はこれらの資料に詳しく書かれていますのでご覧ください。

• https://developer.nvidia.com/blog/cooperative-groups/
• https://www.nvidia.com/content/apac/gtc/ja/pdf/2017/1041.pdf

このreduce1のほうの平均時間は以下の通りです。

これをベースにして他の改良をしたらどうなるか？を示していきます。

Reduce2: Branch Divergenceの削減

Reduce1のコードの問題点の1つ目として以下のif文で実行されるスレッドが飛び飛びになってしまっているという問題点があります。

    for (int s = 1; s < blockDim.x; s *= 2)
    {
        if (tid % (2 * s) == 0)
        {
            sdata[tid] += sdata[tid + s];
        }
        cg::sync(cta);
    }

この結果、1つのwarpでバラバラの処理が実行されることになり、branch divergenceが発生することになります。これはCUDAのコードの高速化をする際には注意する点の一つとなっています。詳しく知りたい方はこちらをご覧ください。

https://docs.nvidia.com/cuda/cuda-c-best-practices-guide/#branching-and-divergence

これを解決するためにblock内のスレッドの0番から連続したスレッドがif文の中を実行するようにします。コードとしては以下の通りです。

    for (uint32_t s = 1; s < blockDim.x; s *= 2)
    {
        uint32_t index = 2 * s * tid;
        if (index < blockDim.x)
        {
            sdata[index] += sdata[index + s];
        }
        cg::sync(cta);
    }

コード全体としては以下のようになります。

https://github.com/shu65/cuda_reduction/blob/main/src/reduction_gpu_old.cu#L90-L119

このwarpのdivergenceの削減は昔から重要な高速化ポイントの一つで、昔の資料でも2.33倍高速化すると書かれています。ではこれを今回の環境で測定すると以下の通りです。

ご覧の通り、現在でもwap divergenceはちゃんと削減すると効果があることがわかりました。

Reduce3: shared memoryのbank conflictの削減

次はshared memoryのbank conflictの削減のための工夫です。shared memoryは高速アクセスできるのですがbank confictに注意する必要があります。memory bankは連続アドレスに割り当てられ、複数のスレッドが同じbankを使う場合は遅延が発生します。V100の場合、以下の資料にいくつか例でどういうときにbank conflictが起きるかが書かれているので、詳しく知りたい方は参考にしてください。

https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html#shared-memory-5-x

注意点として昔のGPUとはmemory bankが16でしたが、最近のGPUはmemory bankが32らしいので注意してください。

bank conflictを削減するためには、飛び飛びのアドレスにアクセスするのではなく、各スレッドが連続したアドレスにアクセスするように改良します。先ほどしめしましたが、元々のコードは以下の通りです。

    for (uint32_t s = 1; s < blockDim.x; s *= 2)
    {
        uint32_t index = 2 * s * tid;
        if (index < blockDim.x)
        {
            sdata[index] += sdata[index + s];
        }
        cg::sync(cta);
    }

このコードでは例えば0番と1番のスレッドは5回目のイテレーションで32個隣をみるようになるので、このタイミングでmemory bankが衝突するようになります。

このコードに対してshared memoryのbank confilict削減するために以下のようにforループを工夫します。

    for (unsigned int s = blockDim.x / 2; s > 0; s >>= 1)
    {
        if (tid < s)
        {
            sdata[tid] += sdata[tid + s];
        }
        cg::sync(cta);
    }

こうすることでmemory bankが衝突することはなくなります。本当にbank conflictが削減しているのか？を確認したい場合、Nsight Computeというプロファイラを使うと確認できます。ちなみに昔からあったプロファイラのnvprofとnvvpはAmpereからサポートされなくなったので注意してください。

このコードを測定すると結果は以下の通りです。

reduce3では1.2倍なので若干速くなっていますが、昔の資料をみると2倍速くなってたらしいのでだいぶ効果がうすれたなかという印象があります。

Reduce4: スレッドの実行効率向上

次にスレッドの実行効率の向上を図ります。

    for (unsigned int s = blockDim.x / 2; s > 0; s >>= 1)
    {
        if (tid < s)
        {
            sdata[tid] += sdata[tid + s];
        }
        cg::sync(cta);
    }

このループはよく見ると最初のイテレーションでblockの半分のスレッドはif文に入らないことが分かります。この結果半分のスレッドは一度もreductionの加算を実行しないことになります。これではせっかくスレッドを立ち上げたのにもったいないことになります。このため、効率をもう少しあげるために、以下のように最初にshared memoryに代入する部分を工夫します。

    int tid = threadIdx.x;
    int i = blockIdx.x * (blockDim.x * 2) + threadIdx.x;
    int sum_value = 0;
    if (i < n)
    {
        sum_value = g_in[i];
    }
    if ((i + blockDim.x) < n)
    {
        sum_value += g_in[i + blockDim.x];
    }
    sdata[tid] = sum_value;

このように最初shared memoryに足す前に各スレッドが2か所データを読み込んで加算してshared memoryに足すようにします。これで少なくとも1回は各スレッドがreductionの加算を実行することになります。またこの結果、カーネル実行時に起動するblock数を半分にできます。この結果計算時間は以下のようになります。

結果からわかる通りやってみると1.8倍速くなっていて、昔も1.7倍の高速化があったらしいので、この改良は今でも効果的なことがわかります。

Reduce6: 完全なloop unroll

さて、次は1つ飛んでReduce6のloop unrollについて説明します。

昔のReduce5はwarp周りの改良なのですが、このwarp周りの事情が昔と今で変わっているのと、同じ部分に別の最適化の話が新しくCUDA Sampleに追加されているので最後にまとめて説明します。

では、reduce6の完全なloop unrollについてです。loop unrollはCUDAをやっていれば最後の手段的に出てくるテクニックで、ループの回数が分かっているならfor文やwhile文を使わずに直書きするというテクニックです。こうすることで、ループを抜けるかの条件判定などをなくすことができます。この結果、速くなるというものです。

今回のreductionでいえば、blockのスレッド数に依存してfor文のループ回数が決まっているので、blockのスレッド数を決め打ちしてあげればfor文なしで書けることになります。この時templateを使えば条件分岐で指定のblockのスレッド数の関数を呼び出すということもでき、ある程度柔軟にスレッド数を指定できるカーネル関数もできます。reductionの加算部分のfor文をunrollするとコードになります(長いので一部省略してます)。

    if (kBlockSize >= 512)
    {
        if (tid < 256)
        {
            sdata[tid] += sdata[tid + 256];
        }
        cg::sync(cta);
    }
    if (kBlockSize >= 256)
    {
        if (tid < 128)
        {
            sdata[tid] += sdata[tid + 128];
        }
        cg::sync(cta);
    }
　　...
    if (kBlockSize >= 2)
    {
        if (tid < 1)
        {
            sdata[tid] += sdata[tid + 1];
        }
        cg::sync(cta);
    }

今回私の書いたコードは簡単のためにカーネル関数の呼び出しもとでblock数を512で決め打ちで書いています。

https://github.com/shu65/cuda_reduction/blob/main/src/reduction_gpu.cu#L284-L285

templateを使ったスレッド数の分岐に関して詳しく知りたい場合はCUDAのSampleの以下のコードをご覧ください。

https://github.com/NVIDIA/cuda-samples/blob/v11.8/Samples/2_Concepts_and_Techniques/reduction/reduction_kernel.cu#L650-L703

この変更を加えた実行時間は以下のようになります。

この変更も昔は1.41倍速くなっていたらしいのですが、ほとんど効果がなくなっているような印象です。個人的にはloop をunrollするとコードのメンテナンス性が非常に悪くなって嫌いなので、これの効果が小さくなっていることは私としてはちょっとうれしいです。

Reduce7: 1スレッドあたりの仕事を増やす

Reduce7では、昔の資料の最後の最適化で起動するスレッド数を減らしつつ、1スレッドあたりのreductionの加算の回数を増やすということをします。これを実現するためにreductionの最初はシーケンシャルに加算を実行していき、その後、いままでの加算する担当のスレッドを半分ずつ減らすというアルゴリズムにします。まず変更前のコードは以下のとおりです。

    int tid = threadIdx.x;
    int i = blockIdx.x * (blockDim.x * 2) + threadIdx.x;
    int sum_value = 0;
    if (i < n)
    {
        sum_value = g_in[i];
    }
    if ((i + blockDim.x) < n)
    {
        sum_value += g_in[i + blockDim.x];
    }
    sdata[tid] = sum_value;
    cg::sync(cta);

この部分、以下のようにwhileループを追加します。

    int tid = threadIdx.x;
    int i = blockIdx.x * (blockDim.x * 2) + threadIdx.x;
    int grid_size = 2 * blockDim.x * gridDim.x;
    int sum_value = 0;

    while (i < n)
    {
        sum_value += g_in[i];
        if ((i + kBlockSize) < n)
        {
            sum_value += g_in[i + kBlockSize];
        }
        i += grid_size;
    }
    sdata[tid] = sum_value;
    cg::sync(cta);

これに加えてカーネル関数の起動時のblock数を減らします。今回は決め打ちで128に固定しています。結果は以下の通りです。

この最適化に関しては昔も1.42倍と効果が小さかったですが、さらに効果が小さくなっている印象です。

Reduce8 : スレッド数がwarpサイズになった以降の同期の最適化(Reduce5の部分)

昔の資料ではreduce7までの最適化ですが、ここからCUDAの新しい機能による最適化に関して紹介していきます。まずはスレッド数がwarpサイズになったところからの同期に関してです。昔、勉強してた人は「え、reduce5がそれでは？」と思った方もいるかもしれませんが、reduce5のコードそのままだと実は危険なので、改良したのものを紹介します。

昔、warpサイズは32で、連続したwarpサイズ分のスレッドは同時に実行されるので同期が必要ないという話がありました。これに基づいて同期を削除して高速化したのがreduce5です。ただ、CUDA 9からwarpに関するいろいろな機能追加があり、その中でwarpサイズ分のスレッドがすべて同時に実行されるという保証がなくなってしまいました。このためwarpサイズ分以下のスレッド数しかないとはいえ、同期なしだと何等かのエラーが発生する可能性があります。

このため、reduce5に同期を加える形にします。コードとしては以下の通りです(長いので一部省略してます)。

    cg::thread_block_tile<32> tile32 = cg::tiled_partition<32>(cta);
    if (cta.thread_rank() < 32)
    {
        if (kBlockSize >= 64)
        {
            if (tid < 32)
            {
                sdata[tid] += sdata[tid + 32];
            }
            tile32.sync();
        }
        if (kBlockSize >= 32)
        {
            if (tid < 16)
            {
                sdata[tid] += sdata[tid + 16];
            }
            tile32.sync();
        }
        ...
        if (kBlockSize >= 2)
        {
            if (tid < 1)
            {
                sdata[tid] += sdata[tid + 1];
            }
            tile32.sync();
        }
    }

こちらのコードではまず、warpサイズの32でスレッド数が十分になったタイミングでtiled_partitionを使って32区切りのthread_block_tile を作り、そのうち0番から31番までのスレッドだけその後の処理をするようにします。さらに同期はこのタイル単位で行うことでwarpサイズ内での同期にしています。thread_block_tile を使わずに__syncwarp() を呼んで同期するという手もありますが、速度的にはそこまで変化しなかったので、ここではthread_block_tile を使うバージョンだけ測定します。

測定結果は以下の通りです。

やってみるとこれに関しては全然速くならないという結果でした。ただ、ループのunrollはreduce6ですでにやってあるし、同期なしにはできなくなってしまっているのでこんなものかもしれないと思っています。

Reduce9 : shfl_down()の使用

最後にshfl_down() の使用による高速化です。先ほど紹介した通りCUDA 9のタイミングでwarp levelの関数がいくつか追加されました。warpレベルの同期もその一つですが、それに加えてデータ交換の関数もいくつか追加されています。これらの関数はCooperative Groupsにも同様の関数があり、今回はそちらを使うようにします。reduce8の段階のコードは以下の通りです。

    cg::thread_block_tile<32> tile32 = cg::tiled_partition<32>(cta);
    if (cta.thread_rank() < 32)
    {
        if (kBlockSize >= 64)
        {
            if (tid < 32)
            {
                sdata[tid] += sdata[tid + 32];
            }
            tile32.sync();
        }
        if (kBlockSize >= 32)
        {
            if (tid < 16)
            {
                sdata[tid] += sdata[tid + 16];
            }
            tile32.sync();
        }
        ...
        if (kBlockSize >= 2)
        {
            if (tid < 1)
            {
                sdata[tid] += sdata[tid + 1];
            }
            tile32.sync();
        }
    }

shfl_down() を使う場合は以下のようになります。

    cg::thread_block_tile<32> tile32 = cg::tiled_partition<32>(cta);
    if (cta.thread_rank() < 32)
    {
        sum_value = sdata[tid];
        if (kBlockSize >= 64)
        {
            sum_value += sdata[tid + 32];
        }
        for (int offset = tile32.size() / 2; offset > 0; offset /= 2)
        {
            sum_value += tile32.shfl_down(sum_value, offset);
        }
    }

このときCUDA Sampleのコードに合わせてループのunrollをやめているので注意してください。またshfl_down()を使うケースではshared memoryを使わなくなっているので同期が不要になります。それ以外は基本処理の流れは同じです。

このコードの実行結果は以下の通りです。

結果をみると少しだけ速くなっていることがわかります。

終わりに

今回、改めてCUDAの復習ということでreductionについての今時のコードの解説を書きました。やってみると結構昔と結果が変わって驚きました。なので、ちゃんと改めて復習してよかったと思っています。とりあえず、これで少しは老害発言をへらせるんじゃないかと思っています。

CUDA Sampleには今回紹介しなかった他のバージョンのreductionも書かれているのでもし気になる方は見てみてください。

The post CUDAの高速化の復習2023年版 Reduction編 first appeared on まったり勉強ノート.

最近、久しぶりにCUDAでコードを書きそうな状況になってきました。このため、ここ数日CUDAの開発環境の準備をしているところです。その際、周りの人が結構Visual Studio Code (VSCode) を使っていろいろな言語のコードを書いているという話を聞いたことを思い出し、私自身、まだVSCodeを使ったことなかったので、せっかくならCUDAの開発をVSCodeを使ってやってみるかーと思ってCUDAをVSCodeを使って書くための準備を開始しした次第です。

そして、VSCodeが何もわからない状態からCUDAのコードをある程度VSCodeを使って書くところまでできたので、今回調べたことをまとめたものを記事にしました。

前提として、CUDAのセットアップは終わっている状態であると仮定して説明をします。また、開発環境はWindows11のマシンでVSCodeを起動して、VSCodeからリモートのUbuntusマシンに接続して開発することを想定して説明します。

今回のVSCodeの設定をして、CUDAの簡単なコード作成したものをGitHubに以下のところにあげてあります。コード全体を見たいという方は見てみてください。

https://github.com/shu65/cuda_vscode_sample

Visual Studio Codeインストール

まずは以下のURLからVSCodeの本家のサイトからVSCodeのインストーラーを自分の環境に合わせてダウンロードしてください。

https://code.visualstudio.com/

その後、インストーラーを実行してVSCodeをインストールしてください。

Visual Studio CodeのExtensionsインストール

CUDAを開発するにあたり最低限必要なextensionを入れます。まず、VSCodeのExtensionsのビューを開きます。ビューは以下のボタンをクリックするか Ctrl + Shift + x で開くことができます。

Extensionsのビューを開いたら以下の二つのものを検索して[Install]ボタンを押します。

C/C++ のほうは名前の通り、C/C++のコードを書くためのextensionで、Nsight Visual Studio Code EditionはCUDAのdebuggerをVSCodeを使ってGUIで操作するためのものになっています。

CUDA開発で必要なVisual Studio Codeの設定ファイルを作る

c_cpp_properties.jsonの作成と編集

c_cpp_properties.json はコード補完や各種C++に関する設定を書くファイルです。VSCodeのウインドウの上部にある[View]→[Command Pallet]を選択するかCtrl + Shift + p というショートカットキーで以下のような入力がでてきます。

ここで[C/C++: Edit Configurations]を選択するとデフォルトの c_cpp_properties.json を.vscode というディレクトリ下に作ってくれます。

C++のコードであれば基本はデフォルトのままでもよいですが、CUDAのコードの開発の場合はインクルードパスにCUDAのディレクトリを追記します。具体的には以下のようにします。

            "includePath": [
                "${workspaceFolder}/**",
                "/usr/local/cuda/include"
            ],

ただし、WindowsにCUDAを入れている方やLinux系でもインストール時のパスの指定などが違うと上記のパスにincludeディレクトリがない可能性があるので注意してください。

tasks.jsonの作成と編集

次にビルドを簡単に実行できるように tasks.json を準備します。VSCodeの場合、taskでコンパイラを叩いて実行する例をよくみます。ただ、この形式だとVSCodeを使わないでビルドする方法が謎だったので、今回はよくあるMakefileを用意してmakeコマンドでビルドすることにします。

Makefileは以下のものを予めディレクトリに追加しておきます。

https://github.com/shu65/cuda_vscode_sample/blob/main/Makefile

c_cpp_properties.json と同様にVSCodeのウインドウの上部にある[View]→[Command Pallet]を選択するかCtrl + Shift + p を叩きます。

そして[Tasks: Configure Task]を選択します。

次に以下のように[Create tasksjsonfile from template]を選択します。

その後、[Others]を選択してファイルを作ります。

これでtasks.jsonを.vscode 下に作ってくれます。あとは以下のようにmakeを叩くtaskを追加します。

        {
            "label": "make",
            "type": "shell",
            "command": "make",
            "args": [],
            "problemMatcher": [],
            "group": {
                "kind": "build",
                "isDefault": true
            }
        }

これでCtrl + Shift + b で上記のmakeを実行するtaskを実行できます。

launch.jsonの作成と編集

最後にlaunch.json の作成と必要な部分の編集についてです。

launch.json を作成する際は、ウインドウ上部の[Run]→[Add Configuration…]を選択するかCtrl + Shift + p を叩いて、以下のように[Debug: Add Configuration…]を選択します。

次に[CUDA C++ (CUDA-GDB)]を選択します。

そしてprogramの部分を以下のように変更します。

 "program": "${workspaceFolder}/main"

実行ファイル名はMakefileでmain という名前を指定していたので、それをそのまま書いています。この部分はプロジェクトごとに変わると思うので注意してください。

launch.jsonの修正が完了したらF5 でcuda-gdb を使ったデバッグを実行できます。

Visual Studio CodeでCUDAのDebuggerを試す

ここまででCUDA開発する最低限の設定ができました。ここからはcuda-gdb の使い方に関して少し説明します。基本的な使い方はC++ のgdb を使ったデバッグと同じだと思いますが、cuda-gdbを使うとCUDAはカーネル関数内の各スレッドの変数の状況も見ることができます。ただ、デバッガーで注目しているスレッドを切り替える方法を見つけるのに苦労したので、同じように苦労する人が減ることを願ってその部分だけここでは説明します。

デバッグするコードはこちらのリポジトリにあるもので説明します。

https://github.com/shu65/cuda_vscode_sample

ちなみにドキュメントのほうにはここに詳しく書いてあります。

https://docs.nvidia.com/nsight-visual-studio-code-edition/cuda-inspect-state/index.html#abstract

まず、CUDAカーネルの関数のデバッグをするときはコンパイル時に以下のように-G を付けてコンパイルしておきます。

nvcc -c -g -G src/main.cu -o src/main.o 

そして、ブレイクポイントを指定したい行番号の左隣をクリックしてブレイクポイントを指定します。

これで指定された位置でプログラムの実行が止まるようになります。

準備ができたらF5 でデバッガーを使って実行します。すると以下のようにブレイクポイントの地点でとまります。

ちなみにデバッガーを起動したタイミングではblockとスレッドのインデックスがそれぞれ（0, 0, 0）(0, 0, 0)のスレッドの情報が表示されています。現在どのスレッドが見えているかは右下のこの部分に表示されています。

見たいスレッドを切り替える場合はCtrl + Shift + p でコマンドパレットを開いて[CUDA: Change CUDA debug focus]を選択するか、上の図のスレッドのインデックス部分をクリックすると以下のような入力が出てきます。

ここに以下のようにblock (X, Y, Z) thread (X, Y, Z) の形式で見たいスレッドのインデックスを入力してエンターを押します。

上の図の例ではblock threadのインデックスが（1, 0, 0）(14, 0, 0)のものを指定しています。こうすると表示するスレッドを切り替えることができます。

他にもデバッガーの使い方は基本的にはドキュメントに書いてありますので、詳しく知りたい方は参考資料のリンクからドキュメントに飛んでご確認ください。

終わりに

今回はVSCodeがそもそも初めてだったこともあり、VSCodeの使い方を調べるのに苦労した感があります。また、デバッガーの使い方も昔使っていたeclipse版と結構違う部分があって苦労しました。同じようにどうすればいいかわからない人の参考になれば幸いです。

参考資料

• https://developer.nvidia.com/nsight-visual-studio-code-edition
• https://docs.nvidia.com/nsight-visual-studio-code-edition/
• https://www.youtube.com/watch?v=gN3XeFwZ4ng

The post Visual Studio Codeを使ってCUDAのコードを書く first appeared on まったり勉強ノート.

10/21にPyTorch 1.10がリリースされ、今回も面白そうな機能が追加されました。個人的には楽しみにしていた「CUDA Graphs」のAPIのベータ版が追加されたということで早速試してみました。今回はこの試した結果の記事になります。

CUDA Graphsとは？

CUDA GraphsはCUDA 10で追加されたCUDAの機能の一つで、複数のCUDA Kernelの実行にかかるオーバーヘッドを減らすための機能です。

基本的には依存関係表すことができるグラフにCUDA Kernelを登録して、依存関係を考慮して順番にCUDA Kernelを実行するという仕組みです。このCUDA Graphsを通して実行すると普通にCUDA Knernelを実行するのに比べてCUDA Kernelの実行オーバーヘッドを減らすことができます。

詳しくはNVIDIA Developer Blogに記事があるのでご覧ください。

PyTorchでCUDA Graphsを使う

PyTorchでCUDA Graphsを使うには主に以下の2つのステップを踏みます。

1. CUDA GraphsのStream Captureの機能を使ってグラフを構築
2. 構築したグラフを実行

それぞれについて順番に説明します。

また、ディープラーニングにおいてすべてのレイヤーがグラフに登録できるものでなかった場合、ネットワークの一部部分だけグラフを構築する方法も用意されています。こちらは今回は触れません。詳しく知りたい方は以下のドキュメントをご覧ください。

https://pytorch.org/docs/master/notes/cuda.html#partial-network-capture

CUDA GraphsのStream Captureの機能を使ってグラフを構築

PyTorchではCUDA Graphsのグラフ構築の一つにStream Captureベースの方法が提供されています。これはtorch.cuda.graph() 以下の実行された関数を自動的にグラフに登録するというものです。
注意点としてはグラフ構築の前のwarmupでは別streamで実行したほうが良いらしいです。詳しくは参考資料の公式ドキュメントをご覧ください。

warmupも含めたグラフ構築は以下の通りです。

static_input = torch.empty((5,), device="cuda")
# Warmup before capture
s = torch.cuda.Stream()
s.wait_stream(torch.cuda.current_stream())
with torch.cuda.stream(s):
    for _ in range(3):
        static_output = static_input * 2
torch.cuda.current_stream().wait_stream(s)

# Captures the graph
g = torch.cuda.CUDAGraph()
with torch.cuda.graph(g):
    static_output = static_input * 2

これで入力をstatic_input、出力をstatic_outputとし、入力を2倍にする計算のグラフgが準備できました。

構築したグラフを実行

構築されたグラフg を実行する際には入力データをstatic_input に上書きして、replay()を実行します。

static_input.copy_(torch.full((5,), 3, device="cuda"))
print("input of cuda graph", static_input)
g.replay()
# static_output holds the results
print("output of cuda graph", static_output) 

出力は以下の通りです。

input of cuda graph tensor([3., 3., 3., 3., 3.], device='cuda:0')
output of cuda graph tensor([6., 6., 6., 6., 6.], device='cuda:0')

注意事項

CUDA Graphsは簡単に使えそうですが、入力のtensorのshapeが変えられないなど制約がいくつかあります。詳しくはこちらをご覧ください。

https://pytorch.org/docs/master/notes/cuda.html#constraints

パフォーマンスの評価

使い方がわかったところで、どれくらい速くなるのか？ということが気になったので測定してみました。測定したときのnotebookは以下のところに置いておきます。

https://github.com/shu65/blog-pytorch-notebooks/blob/main/pytorch_CUDA_Graphs.ipynb

今回は気になった2つのパターンで評価しました。

• GELU
• シンプルなLinearとDropoutのモデルの学習

評価環境は以下の通り。

• 実行環境：Google Colab
  • PyTorch: 1.10.0
  • CUDA: 11.1
  • GPU: K80 (たまたま取れた)

GELU

簡単な例として以下ようなGELUをCUDA Graphsで実行してみます。

def gelu(x):
    return x * 0.5 * (1.0 + torch.erf(x / 1.41421))

また、この際、入力のtensorで小さい例と大きい例の2種類を使って測定してみます。

それぞれのtensorのshapeとしては以下の通りです。

• 小さいtensor: (1, 3, 224, 224)
• 大きいtensor: (32, 3, 224, 224)

上記のサイズのtensorそれぞれを10000回実行して平均計算時間を測定しました。結果は以下の通りです。

評価結果としては個人的には思った通りの結果という印象で、CUDA Kernelのオーバーヘッドの割合が大きい、小さいtensorの時は効果がある程度出ているが、大きいtensorの時はオーバーヘッドの割合が小さいため、ほぼ変わらないという結果になりました。

シンプルなLinearとDropoutのモデルの学習

PyTorchでCUDA Graphsの真価を発揮するのは学習のタイミングかと思いますので、公式ドキュメントにあった例の評価をしてみます。CUDA Graphsに登録する関数train_step()とモデル、各種入力は以下の通りです。

def training_step(model, loss_fn, optimizer, data, target):
    y_pred = model(data)
    loss = loss_fn(y_pred, target)
    loss.backward()
    optimizer.step()

N, D_in, H, D_out = 32, 128, 256, 16
model = torch.nn.Sequential(
    torch.nn.Linear(D_in, H),
    torch.nn.Dropout(p=0.2),
    torch.nn.Linear(H, D_out),
    torch.nn.Dropout(p=0.1)
).cuda()
loss_fn = torch.nn.MSELoss()
optimizer = torch.optim.SGD(model.parameters(), lr=0.1)

# Placeholders used for capture
static_input = torch.randn(N, D_in, device='cuda')
static_target = torch.randn(N, D_out, device='cuda')

ちなみにGoogle Colabで実行しようとしたとき、公式ドキュメントの入力サイズそのままだとcuBLASの内部でエラーが発生して実行できなかったため、サイズを小さくしてあります。

これらを10イテレーション分実行したときの評価結果は以下の通りです。

こちらは思ったよりも速度に差がでました。CUDA Graphsを利用できる場合は使うと効果的かもしれません。

おまけ

CUDA Graphsの制限を見ていて思いましたが、これならtorch.jit.traceやtorch.jit.scriptも併用できるのでは？と思ってやってみました。以前、以下の記事で行ったように torch.jit.script + GELUを使用して評価しました。

評価結果は以下の通りです。

torch.jit.script の効果が大きいですが、CUDA Graphsを使うことでさらに速くなることが確認できました。個人的には CUDA Graphs が使える状況なら torch.jit.trace や torch.jit.script も使えると思われるので併用してよいのではないかと思います。

終わりに

楽しみにしていたCUDA GraphsがPyTorchで使えるようになったということで、評価してみました。一部思った以上の効果を発揮したところもあるので、仕事でも使ってみてノウハウを貯めていこうと思います。

参考資料

• PyTorchの公式ドキュメントのCUDA Graphsの説明部分: https://pytorch.org/docs/master/notes/cuda.html#cuda-graphs

The post PyTorch 1.10の新機能「CUDA Graphs」のパフォーマンスを測定してみる first appeared on まったり勉強ノート.

はじめに

仕事がらPyTorchで高速な学習方法をいろいろ調べることがよくあります。
その際、blog記事などで、Pythonの time() を利用して計算時間を測定して「こんなに速くなりました！」という紹介記事を見かけることがあります。ただ、そこに載っているコードがGPU用の測定方法を用いていないため、正しく測定できていなくて数値が参考にならないということが本当によくあります。

せっかく、いいまとめなのにもったいない・・・ということが多いため、この記事では少しでもそういうものが減ってくれればと思い、PyTorchのGPUの処理の正しい計算時間測定方法についてまとめました。ちなみに、profilerを使ったやり方は別の記事にしようとかと思うので今回は言及しません。

サンプルコードはこちらにありますので、わからないところがあれば実際に動かして確かめてみてください。

https://github.com/shu65/pytorch-cuda-time-measurement/blob/main/Pytorch_GPU_Time_Measurement.ipynb

具体的なGPUの計算時間の測定方法

今回は torch.cuda.synchronize() と torch.cuda.Event を利用した2種類の方法を紹介します。

torch.cuda.synchronize() を利用した方法

torch.cuda.synchronize() を利用する場合は以下のように time()の前に torch.cuda.synchronize() を実行するようにします。

torch.cuda.synchronize()
start = time.time()
# 測定したい部分開始
with torch.no_grad():
  out = model_gpu(input_batch_gpu)  
# 測定したい部分終了
torch.cuda.synchronize()
elapsed_time = time.time() - start

print(elapsed_time, 'sec.')

torch.cuda.Eventを利用した方法

torch.cuda.Event を利用する場合は開始用と終了用のtorch.cuda.Event を作り、測定したい関数の前後でrecord()を呼びます。その後、GPUの処理が終わるまで待つために torch.cuda.synchronize() を呼び、elapsed_time() で計算時間を取得するという流れになります。

start = torch.cuda.Event(enable_timing=True)
end = torch.cuda.Event(enable_timing=True)

start.record()
# 測定したい部分開始
with torch.no_grad():
  out = model_gpu(input_batch_gpu)  
# 測定したい部分終了
end.record()
torch.cuda.synchronize()
elapsed_time = start.elapsed_time(end)

print(elapsed_time / 1000, 'sec.')

詳しい説明

まず、よくある間違えがどんなコードか？を説明します。よくある間違えているコードのサンプルとしては以下の通りです。

start = time.time()
# 測定したい部分開始
with torch.no_grad():
  out = model_gpu(input_batch_gpu)
# 測定したい部分終了
elapsed_time = time.time() - start

print(elapsed_time, 'sec.')

このように普通に time() を呼ぶという間違いが多い印象です。CPUではこれで問題ないのですが、GPUを使った場合はこれでは正しく計算時間が測定できていません。その理由はCPUとGPUの処理が非同期で行われているからです。

CPUとGPUの処理が非同期とは？

PyTorchではGPUの処理を実行する際にCUDAを利用しています。このCUDAではGPUで処理する関数をkernel関数、または単にkernelと呼びます。このkernel関数はCPUとGPUのリソースを最大限に活用できるように、基本的にはCPUがkernel関数の実行を依頼するGPUのタスクキューに積むところまで行い、kernel関数の処理が終わるのを待たないで返ってくるということをしています。

この結果、GPUが処理している最中もCPUが別の処理を実行でき、計算リソースを有効活用することができます。

GPUを活用するという意味ではこの非同期処理という仕組みは非常に有用なのですが、GPUで行っている処理の計算時間を測定する場合には注意が必要になります。というのも、kernel関数を呼んで返ってきたタイミングではGPUの処理が終わってないためです。このため、処理の時間を測定する場合はCPUとGPUとの同期をしたり、CUDA Eventなどの特別な方法で測定する必要があります。

正しく測定してない場合と正しく測定した場合どれくらい差がでるのか？

間違っていたらどれくらいひどいことになるか？を実感してもらうために、Google ColabでResNet50というモデルの推論を行ったときの結果を紹介します。コードはこちらになります

https://github.com/shu65/pytorch-cuda-time-measurement/blob/main/Pytorch_GPU_Time_Measurement.ipynb

実行環境は以下の通りです。

• GPU: T4
• CUDA: 10.1
• PyTorch: 1.8.0
• torchvision: 0.9.0

また、ResNet50の入力バッチサイズは128として実行します。

この条件で実行した結果は以下の通りです

結果として、間違った測定方法だとCPUとGPUを比較すると「1883倍速くなりました！」という主張をしてしまうことになります。ちなみにGPUで1000倍なんて数字が出てきたら確実にどこか間違えています。実際、今回のケースでは本当は「約59倍速くなりました！」というのが正しい結果になります。

torch.cuda.synchronize()とtorch.cuda.Eventを使った場合の違い

今回torch.cuda.synchronize()とtorch.cuda.Event の2種類を紹介しました。場合によっては使い分けをしたほうがいいのでこの二つの違いを説明していきます。

torch.cuda.synchronize() を利用した場合、簡単なので測定しやすいのでいいので、ぱっと測定したい場合はこちらの方法が楽でよいかと思います。ただ、こちらの方法はkernel関数の発行と測定終了のtorch.cuda.synchronize() の終了までの時間も含むことになります。kernel関数の発行もtorch.cuda.synchronize() も時間としては十分短いことが多いので、ほとんどの場合は無視できると思います。ただ、常時監視する目的で測定する際には、torch.cuda.synchronize() を測定したい部分の終了時に呼ぶため、CPUの処理がtorch.cuda.synchronize()のところで止まってしまうのでオーバーヘッドが大きすぎるという問題があります。
一方、torch.cuda.Event を利用した場合はelapsed_time() を呼ぶ直前に何等かの方法でCPUと同期すればいいので、学習のイテレーションの最後に同期するなど工夫することができ、この結果、オーバーヘッドを小さくすることができます。このため、常時監視する目的で測定する際はtorch.cuda.Event の利用をお勧めします。

終わりに

今回はPyTorchのGPUの計算時間を正しく計測する方法について紹介しました。この記事でGPUの計算時間の測定方法を間違えておかしなことを主張する記事が少しでも減ってくれれば幸いです。
profilerについても今度調べて記事にできればと思っています。

The post PyTorchでGPUの計算時間を正しく計測する first appeared on まったり勉強ノート.