無料のSageMakerStudioLabで Gemma 2B ITを動かす

# はじめに
タイトル通りSageMaker Studio Lab 上で Gemma 2B ITの推論実行を行う手順を紹介します。


# SageMaker Studio Labとは
機械学習のモデル開発やデータ解析を行うための統合開発環境であるSageMaker Studioの無料版です。
GPU付きのJupyterLab環境を4時間、GPU無し版なら8時間、使うことが出来ます。
15GBのストレージが付いており、作業を残しておくことが可能です。



# SageMaker Studio Labの起動

## アカウント登録
Studio Labの利用には登録が必要です。（AWSアカウントは不要）
アカウント作成がまだの方は、以下の「アカウント作成～ログイン」を進めてください。

https://github.com/aws-sagemaker-jp/awesome-studio-lab-jp/blob/main/README_usage.md

## GPU Notebookの起動
ログイン後の画面のCompute Typeで`GPU`を選択します。

![2024-03-27-21-31-23.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/3115707/c69e241c-c9dc-7c83-0e43-bfac3ff0bce8.png)

注意が表示されるので`Use GPU`
![2024-03-27-21-31-42.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/3115707/105e13e4-e586-ade8-1564-a438f376677e.png)

暫くすると起動状態になるので、`Open Project`をクリック
![2024-03-27-22-34-16.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/3115707/a2efe787-bc7f-dcea-9545-22f69e96e783.png)

Jupyter Labが開きます
![2024-03-27-22-36-00.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/3115707/4eaf65b9-3ad5-5bbf-bf65-d46d5495f62e.png)

# スペック確認
使い始める前にマシンスペックを確認しておきます。
スペック情報が不要な方は飛ばしてください。

### GPU
GPUはTesla T4、VRAM 15GB 程度です。
```bash
$ nvidia-smi
Wed Mar 27 13:37:20 2024       
+---------------------------------------------------------------------------------------+
| NVIDIA-SMI 535.129.03             Driver Version: 535.129.03   CUDA Version: 12.2     |
|-----------------------------------------+----------------------+----------------------+
| GPU  Name                 Persistence-M | Bus-Id        Disp.A | Volatile Uncorr. ECC |
| Fan  Temp   Perf          Pwr:Usage/Cap |         Memory-Usage | GPU-Util  Compute M. |
|                                         |                      |               MIG M. |
|=========================================+======================+======================|
|   0  Tesla T4                       Off | 00000000:00:1E.0 Off |                    0 |
| N/A   23C    P0              25W /  70W |      2MiB / 15360MiB |      5%      Default |
|                                         |                      |                  N/A |
+-----------------------------------------+----------------------+----------------------+
                                                                                         
+---------------------------------------------------------------------------------------+
| Processes:                                                                            |
|  GPU   GI   CI        PID   Type   Process name                            GPU Memory |
|        ID   ID                                                             Usage      |
|=======================================================================================|
|  No running processes found                                                           |
+---------------------------------------------------------------------------------------+
```

### CUDA
CUDAのバージョンは 11.2 です。
```bash
$ nvcc --version
nvcc: NVIDIA (R) Cuda compiler driver
Copyright (c) 2005-2021 NVIDIA Corporation
Built on Sun_Feb_14_21:12:58_PST_2021
Cuda compilation tools, release 11.2, V11.2.152
Build cuda_11.2.r11.2/compiler.29618528_0
```


### RAM
RAM容量は 15Gi です。
```bash
$ free -h
              total        used        free      shared  buff/cache   available
Mem:           15Gi       2.6Gi       4.6Gi        12Mi       8.1Gi        12Gi
Swap:            0B          0B          0B
```


### CPU
CPUは 4 Coreです。Intel(R) Xeon(R) Platinum 8259CL CPU @ 2.50GHzらしいです。


```bash
$ lscpu
Architecture:                       x86_64
CPU op-mode(s):                     32-bit, 64-bit
Byte Order:                         Little Endian
Address sizes:                      46 bits physical, 48 bits virtual
CPU(s):                             4
On-line CPU(s) list:                0-3
Thread(s) per core:                 2
Core(s) per socket:                 2
Socket(s):                          1
NUMA node(s):                       1
Vendor ID:                          GenuineIntel
CPU family:                         6
Model:                              85
Model name:                         Intel(R) Xeon(R) Platinum 8259CL CPU @ 2.50GHz
Stepping:                           7
CPU MHz:                            3203.138
BogoMIPS:                           4999.99
Hypervisor vendor:                  KVM
Virtualization type:                full
L1d cache:                          64 KiB
L1i cache:                          64 KiB
L2 cache:                           2 MiB
L3 cache:                           35.8 MiB
NUMA node0 CPU(s):                  0-3
Vulnerability Gather data sampling: Unknown: Dependent on hypervisor status
Vulnerability Itlb multihit:        KVM: Vulnerable
Vulnerability L1tf:                 Mitigation; PTE Inversion
Vulnerability Mds:                  Vulnerable: Clear CPU buffers attempted, no microcode; SMT Host state unknown
Vulnerability Meltdown:             Mitigation; PTI
Vulnerability Mmio stale data:      Vulnerable: Clear CPU buffers attempted, no microcode; SMT Host state unknown
Vulnerability Retbleed:             Vulnerable
Vulnerability Spec store bypass:    Vulnerable
Vulnerability Spectre v1:           Mitigation; usercopy/swapgs barriers and __user pointer sanitization
Vulnerability Spectre v2:           Mitigation; Retpolines, STIBP disabled, RSB filling, PBRSB-eIBRS Not affected
Vulnerability Srbds:                Not affected
Vulnerability Tsx async abort:      Not affected
Flags:                              fpu vme de pse tsc msr pae mce cx8 apic sep mtrr pge mca cmov pat pse36 clflush mmx fxsr
                                     sse sse2 ss ht syscall nx pdpe1gb rdtscp lm constant_tsc rep_good nopl xtopology nonsto
                                    p_tsc cpuid aperfmperf tsc_known_freq pni pclmulqdq ssse3 fma cx16 pcid sse4_1 sse4_2 x2
                                    apic movbe popcnt tsc_deadline_timer aes xsave avx f16c rdrand hypervisor lahf_lm abm 3d
                                    nowprefetch invpcid_single pti fsgsbase tsc_adjust bmi1 avx2 smep bmi2 erms invpcid mpx 
                                    avx512f avx512dq rdseed adx smap clflushopt clwb avx512cd avx512bw avx512vl xsaveopt xsa
                                    vec xgetbv1 xsaves ida arat pku ospke avx512_vnni
```


# Gemma 2B ITの起動

NotebookのPythonアイコンをクリックしてNotebookを開きます。
![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/3115707/63b0ccb4-a995-221f-8768-bc0b7201a085.png)

以降はNotebook上での実行コマンドです。

<br>

必要なライブラリをインストールします。
```pythoon
!pip install -U transformers
!pip install accelerate
```

Hugging Faceのトークンをセットします。
```python
from huggingface_hub import HfFolder
hf = HfFolder()
hf.save_token("your token")
```

公式のコードを参考に、モデルをロードして推論実行します。
（初回はダウンロードに時間がかかると思います）
```python
from transformers import AutoTokenizer, AutoModelForCausalLM

tokenizer = AutoTokenizer.from_pretrained("google/gemma-2b-it")
model = AutoModelForCausalLM.from_pretrained("google/gemma-2b-it", device_map="auto")

input_text = "リンゴの色は"
input_ids = tokenizer(input_text, return_tensors="pt").to("cuda")

outputs = model.generate(**input_ids)
print(tokenizer.decode(outputs[0]))
```

推論の結果は全然いけてないですが、推論実行自体は無事に出来ました。
![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/3115707/022eda98-9713-e8de-db8d-613479771c39.png)

<br>


次にこちらの方の記事にならって推論パイプラインを使います。

https://note.com/owlet_notes/n/nd41d2a6d24ce


パイプライン作成
```python
from transformers import AutoTokenizer, pipeline
import torch

model = "google/gemma-2b-it"

tokenizer = AutoTokenizer.from_pretrained(model)
pipeline = pipeline(
    "text-generation",
    model=model,
    model_kwargs={"torch_dtype": torch.bfloat16},
    device="cuda",
)
```

推論実行
```python
messages = [
    {"role": "user", "content": "眠いときはどうする?"}
]
prompt = pipeline.tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
outputs = pipeline(
    prompt,
    max_new_tokens=512,
    add_special_tokens=True,
    do_sample=True,
    temperature=0.7,
    top_k=50,
    top_p=0.95
)
print(outputs[0]["generated_text"][len(prompt):])
```

```text: 出力
1. **睡眠の必要性とメリットを理解する**
2. **睡眠時間を決める**
3. **睡眠機会を確保する**
4. **睡眠環境を改善する**
5. **ストレスを軽減するためのヒントを学ぶ**
6. **睡眠を確保するための習慣を始める**
7. **睡眠の維持を続ける**
```

見慣れた形で推論実行できました。
以上、簡単ですがGemma 2B ITを動かすことはできました。


## ちなみに
モデルロード時のGPUのVRAM使用量は9.5 GBほどです。2Bモデルでも結構消費するようですね。7Bは難しいかもしれません。
![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/3115707/7aa6df9a-4854-e824-7706-a7ba03c03b42.png)



# おわりに
今回は無料で利用できるSageMaker Studio Labを使って、軽量LLMのGemma 2B ITによる推論を試してみました。今回の結果を見るに2Bモデルを動かすのがやっとで、7Bモデルは厳しそうな感じがします。
SageMaker Studioの本家でもGemma 2Bを簡単にデプロイできますが、最小インスタンスでもg5.xlargeなので個人利用には地味に財布に効いてきます。4時間の制限はあるものの無料環境を提供してくれるのは有難い限りです。今後も継続して提供してくれますように。。。


# 参考文献

https://huggingface.co/google/gemma-2b-it

https://note.com/owlet_notes/n/nd41d2a6d24ce