私自身、開発にかかわっているメンバーの一人です。このモデルは1Bという小さいサイズで手軽に動かすことができます。このため、今回はGoogle Colabの無料枠の範囲で簡単に使う方法を紹介します。

ちなみに今回紹介するGoogle Colabのコードはこちらにあげてあります。
https://github.com/shu65/plamo-2-1b-examples/blob/main/plamo_2_1b_inference_example_google_colab_t4.ipynb

コードだけ見たいという方はこちらをご覧ください。動作に関しては2025/02/11現在のGoogle Colabで動くことは確認してありますが、時間がたつと動かなくなる可能性があるので注意してください。

そもそもPLaMo 2 1Bとはどういうモデルか？

まず、このモデルがどういうモデルか知っておくと、いろいろトラブルに対処する心構えができると思ったので、簡単にこのモデルについて説明します。

このLLMは、LLaMaなどに代表されるような、よくあるオープンな他のLLMとは違い、独自路線を突っ走っているモデルです。一番際立って違う点として、状態空間モデル（Sate Space Model, SSM）とSliding Window Attensionを組み合わせたSambaで提案されたアーキテクチャベースのモデルになっています。知っている人からすると「マジで？」と思うかもしれませんが、マジです。詳しくはこちらをご覧ください。（アーキテクチャパートはこの辺いろいろ頑張ってくれてたPFEメンバーの力作の説明になっています）

これ以外にも際立って違う部分としてTokenizerがあります。こちらも特に既存のモデルと違う点として、トークン効率を上げるために英語でもスペース区切りをやめている点などがあります。私も最初「スペース区切りやめます！」って言われた時は「マジかよ」って思いましたが、マジで採用しました。これ以外の工夫も開発した方が熱い思いを記事にしていますので興味がある人はご覧ください。

さて、ここでいろいろ既存のモデルと違うということが伝わったかと思いますが、このようにいろいろ独自な部分があり、面白いモデルになっていると思いますが、一方で他のLLMの感覚で使っているとハマる部分があるので、ここからはPLaMo 2 1Bの推論するやり方に焦点を絞って説明していきます。

PLaMo 2 1BをGoogle Colabの無料枠で使う

Google Colabでの使い方に関してはすでにあげてくれている方もいます。記事にしていただきありがとうございます。

https://qiita.com/autotaker1984/items/32109944a6a058161eee

こちらを見ると、有料でしか使えないL4というGPUでのみ動作確認が取れたと報告をいただきました。私も同じように試したところ、L4での動作は確認できたのですが、いろいろな人に使ってもらうためにはやはり無料枠の範囲で試せるほうが良いだろうと思っています。

このため、ここからは無料で使えるT4というGPUでPLaMo 2 1Bを動かす手順を紹介します。

まず、Google ColabでT4が使えるようにメニューバーから「ランタイム」→「ランタイムのタイプを変更」をクリックして、T4 GPUを選択しておいてください。

そして、まずは最初にPyTorchのバージョンを以下のように2.4系に落とします。

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

T4でPLaMo 2 1Bを動かすにはこのPyTorchのバージョンを落とすということが重要でした。

この後は以下のように他に必要なパッケージをインストールするだけになります。

!pip install transformers>=4.44.2 numba>=0.60.0 causal-conv1d==1.4.0 mamba-ssm==2.2.2

2025/02/11現在Googel Colab上で上記のコマンドを叩くと以下のようなバージョンのパッケージが入りました。

causal-conv1d                      1.4.0
mamba-ssm                          2.2.2
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

インストールが終わればあとは簡単で、PLaMo 2 1BのREADMEにある通りに実行するだけになります。実行コードの例としては以下の通りです。

from transformers import AutoTokenizer, AutoModelForCausalLM

tokenizer = AutoTokenizer.from_pretrained("pfnet/plamo-2-1b", trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained("pfnet/plamo-2-1b", trust_remote_code=True)

text = "これからの人工知能技術は"
input_ids = tokenizer(text, return_tensors="pt").input_ids
generated_tokens = model.generate(
    inputs=input_ids,
    max_new_tokens=32,
    do_sample=True,
    top_k=50,
    top_p=0.95,
    temperature=1.0,
)[0]
generated_text = tokenizer.decode(generated_tokens)
print(generated_text)

私が実行した際は以下のように出力されました。

<|plamo:bos|>これからの人工知能技術は人間の脳機能をコンピュータに取り入れ、私たちの生活に様々な影響を与えると言われています。人工知能技術の進化により、社会に混乱が生じる可能性があり、AIの規制に関して議論

入力で与えた文章の続きとして問題ない文章がちゃんと出力できていると思われます。

このようにT4でも問題なくPLaMo 2 1Bを動かすことができました。

終わりに

この記事ではGoogle ColabのPLaMo 2 1BをGoogle Colabで動かす手順を紹介しました。おそらくGoogle Colabじゃなくてもcausal-conv1dとmamba-ssm がサポートされている環境であれば動作すると思われます。逆に言えばこの二つがサポートしてない環境では使うのにいろいろ魔改造が必要な可能性があります。

この辺りはハマる人が多いと思われるので、知見がたまったらまた記事にしようと思います。

この記事を参考にみなさんもPLaMo 2 1Bで遊んでもらえればと思います。

PLaMo 2 1BをSupervised Fine-Tuning（SFT）するコードに関しても準備中で、動作確認は済んだので今週中にあげようと思いますのでお楽しみに！

The post 小型LLM PLaMo 2 1BをGoogle Colabの無料枠の範囲で使ってみる first appeared on まったり勉強ノート.

時間計測の部分で測定に使ったコードはここにあげてあります。

https://github.com/shu65/pytorch_2_compile_example/blob/main/torch_2_0_compile.ipynb

torch.compileとは？

torch.compileはPyTorch 2.0の新機能で、PyTorchの複数の機能を組み合わせて使い関数や深層学習のモデルを実行時に最適化して、その後の呼び出して高速に実行できるようにする機能です。

torch.compileの中身の詳しい説明はここにかかれています。

https://pytorch.org/get-started/pytorch-2.0/#technology-overview

簡単に説明するとtorch.compileの中身としては以下の３つで構成されています。

1. Graph acquisition: 計算グラフの構築
2. Graph lowering: PyTorchのオペレーションをバックエンドのデバイス（CPUやGPU）に特化した細かい命令に分解
3. Graph compilation: バックエンドのデバイス特化の命令を呼び出し

これらのステップを経ることで、より効率よく計算リソースを使えるようにし、高速化を実現しています。

また、この機能のすばらしいところは使い方も非常に簡単であるというものがあります。以下にデコレータで使う方法とtorch.compileの関数を呼び出して使う方法を示します。

デコレータで使うやり方

まずデコレータで使う方法です。これは以下のようになります (このチュートリアルの例：https://pytorch.org/tutorials/intermediate/torch_compile_tutorial.html#basic-usage)

@torch.compile
def opt_foo2(x, y):
    a = torch.sin(x)
    b = torch.cos(x)
    return a + b
opt_foo2(torch.randn(10, 10), torch.randn(10, 10))

torch.jit.scriptを使ったことがある方は、それと同じ感覚で使えるというと使い方がイメージしやすいかもしれません。

torch.compileの関数を呼び出して使うやり方

torch.compileの関数を呼び出してコンパイルする場合は以下のようにやります。(このチュートリアルの例：https://pytorch.org/tutorials/intermediate/torch_compile_tutorial.html#basic-usage)

class MyModule(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.lin = torch.nn.Linear(100, 10)

    def forward(self, x):
        return torch.nn.functional.relu(self.lin(x))

mod = MyModule()
opt_mod = torch.compile(mod)
opt_mod(torch.randn(10, 100))

こちらもtorch.jit.scriptのときと同じような使い方だと思います。

torch.compileによるパフォーマンスの評価

次にtorch.compileを実際に使ってみたときの計算時間を計測したので、その紹介です。今回は以下の二つのGPUで測定しました。

1. T4
2. V100

T4はTuringなので公式のドキュメントでtorch.compileのサポートが書かれてないものになっています。ただ、やってみたら少し早くなったので、測定結果を載せています。GitHubにあげたコードはT4で測定したほうです。

また、CUDAのバージョンはどちらのケースも12.0利用し、測定に使ったモデルはチュートリアルにあったtorchvisionのResNet18を使用しました。

また、torch.compileにはモードが以下の３つあります。

1. デフォルト
2. reduce-overhead
3. max-autotune

これらと何もしてない場合も含めて合計４つパターンの測定をしています。

具体的な測定方法が分かりやすいようにコードの一部を紹介します（torch.compleのデフォルトの場合）。

import time 

import torch
import torchvision.models as models
import torch._dynamo

batch_size = 64
n_warmup_iters = 10
n_iters = 500

x = torch.randn(batch_size, 3, 224, 224).cuda()

def get_mode():
    return models.resnet18()

torch._dynamo.reset()

model = get_mode().cuda()
optimizer = torch.optim.SGD(model.parameters(), lr=0.01)

# compile
compiled_model = torch.compile(model)
for i in range(n_warmup_iters):
    optimizer.zero_grad()
    torch.cuda.synchronize()
    start = time.time()
    out = compiled_model(x)
    torch.cuda.synchronize()
    forward_elapsed_time = time.time() - start
    torch.cuda.synchronize()
    start = time.time()
    out.sum().backward()
    backward_elapsed_time = time.time() - start
    print(f"with compile {i} iter forward: {forward_elapsed_time/1000:.3e} msec., backward: {backward_elapsed_time/1000:.3e} msec.")
    optimizer.step()

print("-"*10)

torch.cuda.synchronize()
start = time.time()
for i in range(n_iters):
    optimizer.zero_grad()
    out = compiled_model(x)
    out.sum().backward()
    optimizer.step()
torch.cuda.synchronize()
elapsed_time = time.time() - start

print(f"with compile total:{elapsed_time:.3e} sec. {batch_size*n_iters/elapsed_time:.3e} imgs/sec.")

最初に、モデルの入力とモデルを作ったあと、コンパイルする場合はtorch.compile(model)でコンパイルします。このときコンパイルのモードを変える場合は引数のmodeにモードの名前を渡します。

その後、最初の数回はforward、backwardの呼び出し時にコンパイルなどのオーバーヘッドが入って遅いので、あらかじめ何度か呼びます。そして最後に実際に時間を計測します。今回は10回あらかじめforwardとbackwardを呼んでおいて、その後500回イテレーションを回したときの時間を測定しています。バッチサイズに関しては変化させると高速化率が変化することはわかっていますが、今回固定で64で実行しています。

T4, V100ともに同様の方法でtorch.compileのありなし等を測定しています。

では、時間計測の結果です。500回イテレーションを回したときの実際の計算時間を順番に示していきます。まずはT4の場合です。

T4はtorch.compileのサポートが書かれてない世代のGPUなので、効果が全くでないのかと思ったのですが、そんなことはなかったです。ただ、10％は満たない高速化にとどまっているという印象です。ちなみにT4を使ったケースではtorch.compileのmodeをmax-autotuneに変えると以下のようにサポートされてないGPUであると警告がでてきます。

[2023-03-17 18:31:06,314] torch._inductor.utils: [WARNING] not enough cuda cores to use max_autotune mode

次にV100のResNet18の結果です。

V100のほうはtorch.compileのサポートされていると書かれているGPUです。実際、V100はtorch.compileのmodeをmax-autotuneに変えると確かにより速くなり、高速化率も最大値は10%台に入っています。

現状のtorch.compileの注意点

最後にtorch.compileの注意したほうがよさそうな点を書いておきます。

まず、公式で書かれいたものの紹介です。基本的な注意点はこのドキュメントに書いてあります。

https://pytorch.org/get-started/pytorch-2.0/#pytorch-2x-faster-more-pythonic-and-as-dynamic-as-ever

重要なものとして、現在提供されているtorch.compileの機能を最大限活かせるのはCPU、NVIDIAのVoltaとAmpere世代のGPUのみになっています。他のGPUでは使おうとすると警告が出てきます。ただ、私が試した範囲では警告がでるだけで現状では使えないわけではなさそうです。

また、私が使ったときに感じた注意点としては

1. おそらくforwardとbackwardで別々にコンパイルが走るので、forward、backwardの両方とも最初は遅い
2. 実行が遅いのは最初の１回目だけでなく、最初の数回の呼び出しが遅いケースがある
3. Google ColabなどでCellの実行を一度止めて再度実行しようとするとエラーがでて、ランタイムの再起動をしないと復帰できないケースがある

1と２は時間計測をしようとしたときにはまったポイントです。まず、１に関してです。torch.compileの直後の呼び出しはコンパイルが走るので、遅いというのはドキュメントにも書かれています。ただ、forwadだけがおそいのかな？と思ってました。ただ、torch.compileの説明をちゃんと読めば想像できると思いますが、backwardも最初の実行のときは遅いです。なので、時間を計測するときは、forwardとbackwardの両方が遅いことを考慮して測定する必要があります。

次に２です。これに関しては私が見逃してなければドキュメントに明示的に説明が書いてあるわけではないのですが、チュートリアルの時間計測の結果や実際に測定してみるとどうやら遅いのは最初の１回目の呼び出しだけではなく、そのあと数回遅いケースが存在しているようです。このため、計算時間の測定の際、最初に数回呼び出してから測定しないとtorch.compileを使ったときよりも遅いみたいな誤った結果になるので注意してください。

最後に３です。これは何度かはまったのですが、どこかにキャッシュか何か残っているのか変なところで止めるとコード的には問題ないはずなのに、エラーがでるようになるときがあります。調べても解決方法が分からなかったので、エラーがでるようになったらランタイムごと再起動するということを何度かやりました。Google Colabでやるときは注意してください。

終わりに

今回はtorch.compileについて使ってみたのでまとめを書きました。去年発表があったときから楽しみにしていましたが、期待通りのものとなっていました。なにより使い方が非常に簡単なことには驚きました。

今回はT4とV100の測定結果でしたが、A100だとどうなるのかも今度測定しようかなと思っています。

この記事がみなさんのお役に立てば幸いです。

The post PyTorch 2.0の新機能「torch.compile」使ってみた 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 まったり勉強ノート.

ちなみに、測定するときにいろいろ気を付けないといけないポイントがあったので、Fuse pointwise operationsのために利用したtorch.jit.script の謎現象で困る人が減るように、それについても後半で説明していきます。

第一弾の「parameter.grad = Noneを使う」というのもありますので、PyTorchの高速化に興味がある方はそちらも合わせてご覧ください。

Fuse pointwise operationsとは？

elementwiseの加算や乗算、sin(), cos(), sigmoid() などなど、行列やベクトルの要素単位で実行される演算をまとめてpointwise operationsと呼ぶときがあります。これらの演算は一つの演算にかかる時間は非常に短いため、GPUのような関数1回の実行のオーバーヘッドやメモリアクセスのオーバーヘッドが大きい演算器では計算量のわりに長い計算時間がかかってしまいます。

このようなメモリアクセスや関数の実行のオーバーヘッドを削減する工夫として、複数の独立した演算を一つの関数にまとめる(fuse)という方法が良く用いられます。

PyTorchでも演算をまとめる仕組みがあります。その中でもpointwise operationsをfuseする仕組みとしてよく例で用いられるのが torch.jit.script です。

今回はこのtorch.jit.script によって、どれくらいfuseしたpointwise operationsが速くなるのかを確認していきます。

実際に効果を測定してみる

torch.jit.script でfuseするとどれくらい速くなるのか？を測定するための環境と実際に用いたコードは以下の通りです。

• 実行環境：Google Colab
• GPU: Tesla T4
• PyTorch: 1.8.1
• torchvision: 0.9.1
• 測定に使ったnotebook: https://github.com/shu65/pyorch_performance_tuning_guide_examples/blob/main/Fuse_pointwise_operations.ipynb

また、今回測定に利用した関数は「PERFORMANCE TUNING GUIDE」で示されていたGELUです。実装自体は単純で、以下の通りです。

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

また、今回はGPUのだけでなく念のためCPUも測定しました。

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

CPUのほうはあまり期待してなかったですが、予想通りほぼ変わらずという結果でした。一方、GPUのほうは劇的に速度が変化し、今回のGULEの例では4.5倍速くなることが確認できました。個人的にはtorch.jit.script で速くなることはあまりないようなイメージだったので、シンプルなFuse pointwise operationsならちゃんと速くなるというのがわかって少し感動してます。

torch.jit.script を使った実行時間測定の注意点

さて、この記事を書くにあたってかなり苦労したので、その苦労話もちゃんと書いておこうと思います。この分量の内容の記事なら数時間で実験して書けるだろうと当初は思っていたのですが、torch.jit.scriptの謎現象に悩まされて実験がちゃんと安定して取れるようになるまで、実は数日かかりました。なので、torch.jit.scriptを使った計算時間測定の注意点をまとめておきます。

1. GPUの計算時間を正しく測定する

以前自分で「PyTorchでGPUの計算時間を正しく計測する」という記事を書きましたが、恥ずかしながら最初は正しく測定するのを忘れていました。なので、自戒も込めて何度も書きますが、GPUの計算時間を測定するときは注意してください。

2. torch.jit.scriptの1回目の実行はオーバーヘッドが大きいので無視する

torch.jit.scriptは名前の通りJITなので、1回目の実行時はオーバーヘッドが大きいです。このため、1回も含めていて、かつ、少ない実行回数で平均を取るとtorch.jit.scriptを使っているのに速くなっていないというような状態になります。このため、ちゃんと測定する場合は1回目の実行は別にするようにするとよいかと思います。

3. 入力のTensorのshapeやdeviceが違う場合はtorch.jit.script()の実行前にキャッシュをクリアする

今回の測定で気が付くのに苦労した点がこれです。PyTorch 1.8.1現在、torch.jit.script()は一度関数オブジェクトをtorch.jit.script化したあと、2回目以降はこの部分をスキップするためにキャッシュしています。このため、全く別のshapeやdeviceのTensorを入力に使う場合はtorch.jit.script() を実行する前にキャッシュをクリアしておかないと、本来はJITを使って速くなるはずなのにキャッシュに残ったものがそのまま使われて全然速くならないという現象が発生します。

今回の測定に用いたnotebookではCPUを測定したあとGPUの測定をしています。このため、何もしていないとCPUでJITが走っているので、その後、いくら入力をGPUにしていてもGPU用のJITが走らず、torch.jit.scriptをGPUで使っているのに全然速くならないという状態になります。

これを回避するために以下のようにキャッシュのクリアしてからtorch.jit.script化して測定を行うようにしています。

torch.jit._state._jit_function_overload_caching.clear()
torch.jit._state._jit_caching_layer.clear()
scripted_gelu = torch.jit.script(gelu)

ちなみにちゃんと最適化が走っているか確認する際はtorch.jit.last_executed_optimized_graph() で直前の関数の実行時のグラフが出力できるので、JITが走っているはずの1回目の実行で「prim::profile」というものが出てきているか確認してください。現状のPyTorchのデフォルトだと最初の1回目はプロファイル測定のためにこのようなIRが挿入されるようになっています。

最後に

PyTorchには「PERFORMANCE TUNING GUIDE」の「Fuse pointwise operations」を試したときのまとめを書きました。個人的にはtorch.jit.script() を使う際の注意点がいろいろわかってかなり勉強になりました。他にもまだまだ試したい高速化テクニックがあるので、試した際はまたこうしたまとめ記事を書こうと思います。

The post PyTorchのPERFORMANCE TUNING GUIDEの効果を確認してみる その2 「Fuse pointwise operations」 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 まったり勉強ノート.