フロート64で100億パラメータ(10万×10万の行列)を持つモデルのメモリー使用量は、約74.5GBになります。このクラスの重み行列を持つニューラルネットワークモデルはとても強力であることが予想されます。
（FP16だと18.75GBとなり、メモリに乗らない。）
行列のドット積は、行列Aの行ベクトルと行列Bの列ベクトルで内積計算を行うことで計算される。しかし、行列サイズが大規模になると、計算処理に必要なメモリが膨大となり、メモリ不足が問題となる。行列のサイズを小さくするためには、計算対象の行列を小さなチャンク（ブロック）に分割し、部分的に計算する手法が考えられる。しかし、行列を単純にチャンクに分割すると、チャンク内の局所的な範囲でしか内積計算が行えず、グローバルな行列全体の情報を捉えることができないという問題が生じる。
本稿では、大規模行列のドット積計算において、2つの行列をテキストファイルからストリーム処理し、対応する行と列を部分的に読み込んで計算する手法を用いています。この手法の目的は、メモリ消費量を最小限に抑えつつ、行列全体の情報をグローバルに捉えた計算を行うことです。
つまり10万次元ベクトル（ベクトル1つでも巨大です。）の内積計算を逐次行うということです。
例 FP64
ベクトル1（10万次元ベクトル）のメモリ使用量: 0.000745 GB
ベクトル2（10万次元ベクトル）のメモリ使用量: 0.000745 GB
合計メモリ使用量: 0.001490 GB
-
行列データのストリーム処理
大規模行列を直接メモリに保持すると、メモリ不足が発生する可能性があるため、計算対象の行列をディスク上のテキストファイルに保存し、必要な行と列データを逐次的に読み込みながら処理を行います。具体的には、行列Aの行ベクトルと行列Bの列ベクトルをそれぞれチャンクごとにテキストファイルから読み込み、そのドット積を計算します。（SSD上のテキストファイルを仮想メモリとして利用します。） -
行と列の対応関係
この手法では、行列Aの特定の行と、行列Bの対応する列を一度にメモリに読み込みます。行ベクトルと列ベクトルを読み込むことで、局所的な演算により、グローバルな行列全体のドット積を少しずつ計算し、最終的な結果行列を構築します。 -
メモリ効率の向上
行列のデータを一度にすべて読み込むのではなく、必要な行と列を小さなチャンク単位で読み込むことで、使用するメモリ量を大幅に削減できます。チャンクサイズは柔軟に設定でき、計算精度を保ちながら、システムのメモリ制約に対応することが可能です。
このストリーム処理方式により、大規模な行列のドット積演算を効率的かつ正確に行うことができ、行列全体のグローバルな情報を保持した計算が実現されます。この手法は、大規模データ処理の一つの解決策となります。
実行結果。マトリックス サイズ1000の場合。
ランダムなFP16行列がテキストファイルに保存されました。
matrix1.txtのサイズ: 0.023284 GB
matrix2.txtのサイズ: 0.023284 GB
Progress: 1.00%
Progress: 2.00%
Progress: 3.00%
Progress: 4.00%
Progress: 5.00%
Progress: 6.00%
Progress: 98.00%
Progress: 99.00%
Progress: 100.00%
ドット積の計算が完了しました。
テキストファイルが削除されました。
import numpy as np
import os
# 行列のサイズを定義
matrix_size = 1000 # 1000x1000の行列サイズ
block_size = 100 # 一度に読み込むブロックサイズ（部分行列のサイズ）
# ステップ1: ランダムなFP16行列を生成してファイルに保存
matrix1 = np.random.rand(matrix_size, matrix_size).astype(np.float16)
matrix2 = np.random.rand(matrix_size, matrix_size).astype(np.float16)
# 行列1と行列2をテキストファイルとして保存
np.savetxt('matrix1.txt', matrix1)
np.savetxt('matrix2.txt', matrix2)
print("ランダムなFP16行列がテキストファイルに保存されました。")
# テキストファイルのサイズをGB単位で表示
matrix1_size = os.path.getsize('matrix1.txt') / (1024 ** 3) # バイトをGBに変換
matrix2_size = os.path.getsize('matrix2.txt') / (1024 ** 3) # バイトをGBに変換
print(f"matrix1.txtのサイズ: {matrix1_size:.6f} GB")
print(f"matrix2.txtのサイズ: {matrix2_size:.6f} GB")
# ステップ2: 部分的に行列を読み込み、ドット積を計算
result = np.zeros((matrix_size, matrix_size), dtype=np.float16) # 結果を保存するための行列
total_steps = (matrix_size // block_size) ** 2 # 進捗表示のための総ステップ数
step = 0 # 現在のステップ数
# "matrix1.txt"と"matrix2.txt"を読み込みながら計算
with open('matrix1.txt', 'r') as f1:
for i in range(0, matrix_size, block_size):
# matrix1の部分行列（i番目からblock_size行分）を読み込む
rows_matrix1 = np.loadtxt(f1, max_rows=block_size, dtype=np.float16)
with open('matrix2.txt', 'r') as f2:
for j in range(0, matrix_size, block_size):
# matrix2のj番目の列からblock_size列分を読み込む
f2.seek(0) # ファイルの先頭に戻る
cols_matrix2 = np.loadtxt(f2, usecols=range(j, j + block_size), max_rows=matrix_size, dtype=np.float16)
# 部分行列のドット積を計算して結果行列に保存
result[i:i+block_size, j:j+block_size] = np.dot(rows_matrix1, cols_matrix2)
# 進捗表示
step += 1
progress = (step / total_steps) * 100
print(f"Progress: {progress:.2f}%")
print("ドット積の計算が完了しました。")
# ステップ3: テキストファイルを削除
os.remove('matrix1.txt')
os.remove('matrix2.txt')
print("テキストファイルが削除されました。")
chank 計算のイメージ
1024×1024の行列の類似度計算を行い、512×512のチャンクに分けてブロック単位で計算するように変更します。
具体的なポイント:
ランダムなテクスチャ生成関数を1024×1024に対応させる。
512×512のチャンクごとにシェーダーの実行を行う
<!DOCTYPE html>
<html>
<head>
<title>Matrix Similarity Calculation with Three.js</title>
<style>
/* キャンバスのスタイル設定 */
canvas {
width: 256px;
height: 256px;
display: inline-block;
margin: 10px;
}
</style>
</head>
<body>
<div id="container"></div>
<!-- Three.jsのライブラリを読み込み -->
<script src="https://cdnjs.cloudflare.com/ajax/libs/three.js/r128/three.min.js"></script>
<!-- フラグメントシェーダーのコード -->
<script type="x-shader/x-fragment" id="fragment-shader">
uniform sampler2D u_texture1;
uniform sampler2D u_texture2;
varying vec2 vUv;
void main() {
// テクスチャ1とテクスチャ2から色の値を取得
vec4 color1 = texture2D(u_texture1, vUv);
vec4 color2 = texture2D(u_texture2, vUv);
// ドット積を計算（対応するRGB要素の積とその合計）
float similarity = dot(color1.rgb, color2.rgb);
// 類似度をグレースケールの色として出力
gl_FragColor = vec4(vec3(similarity), 1.0);
}
</script>
<script>
// ランダムなテクスチャを生成する関数（1024×1024に対応）
function createRandomTexture(size) {
const data = new Float32Array(size * size * 4); // RGBAの4チャンネル
for (let i = 0; i < data.length; i++) {
data[i] = Math.random(); // ランダムな値を設定
}
const texture = new THREE.DataTexture(data, size, size, THREE.RGBAFormat, THREE.FloatType);
texture.needsUpdate = true; // テクスチャの更新を指示
return texture;
}
// チャンクごとのシーンを作成し、テクスチャを描画する関数
function createScene(container, texture1, texture2, isSimilarity, chunkX, chunkY) {
// シーンとカメラを作成
const scene = new THREE.Scene();
const camera = new THREE.OrthographicCamera(-1, 1, 1, -1, 0.1, 10);
camera.position.z = 1;
// 平面ジオメトリとシェーダーマテリアルの作成
const geometry = new THREE.PlaneBufferGeometry(2, 2);
let material;
if (isSimilarity) {
// 類似度を計算するシェーダーマテリアルを設定
material = new THREE.ShaderMaterial({
uniforms: {
u_texture1: { type: 't', value: texture1 },
u_texture2: { type: 't', value: texture2 }
},
vertexShader: `
varying vec2 vUv;
void main() {
vUv = uv;
gl_Position = projectionMatrix * modelViewMatrix * vec4(position, 1.0);
}
`,
fragmentShader: document.getElementById('fragment-shader').textContent
});
} else {
// 単純なテクスチャを表示するためのマテリアル
material = new THREE.MeshBasicMaterial({ map: texture1 });
}
const mesh = new THREE.Mesh(geometry, material);
scene.add(mesh);
// レンダラーの作成
const renderer = new THREE.WebGLRenderer();
renderer.setSize(256, 256); // 各チャンクのサイズ
container.appendChild(renderer.domElement);
// チャンクごとのオフセットを設定
const chunkOffset = 512; // 各チャンクのサイズ（512×512）
renderer.setViewport(chunkX * chunkOffset, chunkY * chunkOffset, chunkOffset, chunkOffset);
// シーンを描画
renderer.render(scene, camera);
}
function main() {
const container = document.getElementById('container');
// 2つのランダムなテクスチャを生成（1024×1024）
const texture1 = createRandomTexture(1024);
const texture2 = createRandomTexture(1024);
// チャンクごとの表示処理（512×512のチャンクに分割）
for (let i = 0; i < 2; i++) { // 2x2のチャンク
for (let j = 0; j < 2; j++) {
// 類似度を計算して表示
createScene(container, texture1, texture2, true, i, j);
}
}
}
main();
</script>
</body>
</html>
まとめ
チャンクサイズが小さすぎると、全体の傾向を捉えられず、計算結果が全く異なるものになる可能性が高いです。
したがって、近似的な計算を目指す場合には、できるだけチャンクサイズを大きくし、全体の行列に対する類似度計算が行えるようにすることが好ましいです。
計算資源に限りがある場合でも、チャンクサイズを可能な限り大きく設定することで、無駄を減らしつつ精度を高めることが可能です。
もし、計算リソースに問題がないのであれば、1024×1024の行列全体での計算を目指すのが理想的だと言えます。