しの (@shino1609) です。8月10日から9月29日まで、株式会社 Preferred Networks の夏季インターンシップに「CuPy の開発」というテーマで参加しました。

この記事では、業務の大まかな内容や、選考課題への取り組み、感想や手応えを書きます。インターンに行ったことがない人にとっても有益な記事にしたいと思います。

この企業を選んだ理由

Preferred Networks (PFN) は強い競プロerが集まる機械学習系の企業として昔から認識していましたが、機械学習は全くの初心者なので、縁の遠い存在だと思っていました。しかし、冬に別の企業で高速化関連のインターンをしたときに「PFN もこういう高速化をやってるよ」という情報を得て、もしかしたら行けるかもと考えるようになりました。

他の企業からもインターンシップのオファーを受けましたが、内容がとても興味に合っていたので PFN を選びました。大規模な web サービスに関われるインターンもあり悩みましたが、PFN でのテーマに直感的な魅力を感じました。

テーマ「CuPy の開発」と選考課題

CuPy は NumPy/SciPy の GPU 版ライブラリです。NumPy を使う既存コードに対して、基本的には numpy.XXX を全て cupy.XXX に置き換えるだけで、データを GPU 上で処理するようにできます。

例：https://cupy.dev より

1
2
3
4
5
6
7

>>> import cupy as cp
>>> x = cp.arange(6).reshape(2, 3).astype('f')
>>> x
array([[ 0.,  1.,  2.],
       [ 3.,  4.,  5.]], dtype=float32)
>>> x.sum(axis=1)
array([  3.,  12.], dtype=float32)

この CuPy に新機能の追加を行うのがこのテーマの趣旨でした。

募集テーマ一覧では、必須要件として Python, NumPy, CUDA が、歓迎要件として NCCL, MPI, Metal が挙げられていました。この時点で CUDA については概要しか知りませんでしたが、とにかく応募を出し、テーマ別選考課題に沿って CUDA をがっつり勉強しました。

​

今年、いくつかの選考課題付きインターンに応募しましたが、背伸びしたところに挑戦したことが非常に大きな学びの機会になりました。「どこどこのインターンに行きたい」という気持ちを原動力に、新たに言語やフレームワークを学んだり、それらを使って何かを製作する経験ができました。仮に不合格だったとしても、これだけのことを学べただけでも価値があったかなと思います。

業務内容

要約：CuPy の GitHub レポジトリに PR を出しました。この PR で、CuPy が提供する N-dimensional array (ndarray) の複数 GPU 分散バージョン、DistributedArray を追加しました。

基本の API

単一のホスト (CPU) に複数のデバイス (GPU) が接続されている環境を仮定します。

1
2

a = numpy.arange(10)
A = distributed_array(a, {0: slice(5), 1: slice(5, 10)})

このように記述すると、a[:5] の部分が 0 番目のデバイスに、a[5:10] の部分が 1 番目のデバイスに配置されます。2 次元以上の場合も同様です。

（注釈：slice(x, y, z) は a[x:y:z] というインデクシングに対応します。参照）

DistributedArray は ufunc に対応しており、A+B や sin(A) といった要素ごとの演算が行えます。A.sum(axis=0) などの reduction や行列積 A @ B も、基本的なところをサポートしました。

これに付随して DistributedArray 特有の概念がいくつかありますが、特に resharding という機能に苦労しました。業務の雰囲気が分かると思うので、これがどのようなものだったか簡単に説明します。

Resharding (再振り分け)

以下の状況を考えます。

1
2
3

A = distributed_array(a, {0: slice(5), 1: slice(5, 10)})
B = distributed_array(b, {0: slice(5), 1: slice(5, 10)})
C = distributed_array(b, {0: slice(7), 1: slice(7, 10)})

A+B の計算は簡単で、デバイス 0 とデバイス 1 がそれぞれ a[:5] + b[:5] と a[5:10] + b[5:10] にあたる部分を計算すればよいです。しかし A+C のように、各デバイスが持つ領域が一致していない場合は簡単にはいきません。そこで、次のようにします。

1

C2 = C.reshard(A.index_map)

こうすると、C のデータが A に合わせて再配置されます。この例では、c[5:7] の部分をデバイス 0 からデバイス 1 に転送します。こうして得られた C2 は A と同様に、c[0:5] の部分をデバイス 0 に、c[5:10] をデバイス 1 に持つようになります。A+C2 は先ほどと同様で容易に行えます。

reshard は実装をかなり工夫しました。NCCL というライブラリを使ってデータを非同期的に転送し、データ転送とその直後の演算を並列に実行できるようにしています。例えば、次のコードを考えます。

1

A + C.reshard(A.index_map)

このコードは以下のことを実行します。

  1. [転送] デバイス 0 からデバイス 1 に c[5:7] を転送する

  2. [演算] デバイス 0 が a[0:5] + c[0:5] を計算する

  3. [演算] デバイス 1 が a[5:7] + c[5:7] を計算する

  4. [演算] デバイス 1 が a[7:10] + c[7:10] を計算する

このうち、reshard が 1. を実行し、それと並列に 2. と 4. が実行されます。1. が完了すると、続けて 3. も実行されるという流れです。これにより、データ転送中の暇な時間を他の演算で有効に活用できます。

振り返り

自分の取り組みについて

試行錯誤と改善を繰り返すたびに CUDA 関連の理解が深まり、やりたいことが成果として現れることにやりがいを感じました。できることならもっと長い時間をかけて作り込みたかったです。

また、専攻や研究活動がまともに無いことを今まで引け目に感じていましたが、それでもインターンとして採用され、期間中の取り組みが評価されたことで、自分の強みを認識でき自信を持てました。実装が得意なことは誇っていきたいです。

周りからの刺激について

他のインターン生や正社員の方々と交流する機会を提供してくださり、さまざまな話を聞くことができました。周りの方々の優秀さを目の当たりにして身が引き締まりましたし、仕事や学びに積極的な姿勢が魅力的でした。自分が所属したチームのメンバーは親切で話しやすく非常に優秀で、交流がとても楽しかったです。本当にお世話になりました。

今後

GPU 関連の技術に対する興味が高まりました。このようなパフォーマンスを引き出す技術を扱っていくのは面白いですし、機械学習分野でもこの方面の需要は高いようなので、これからもモチベーションを保ちながら続けていきたいと思います。

半年前には CUDA について何も知らなかったことを考えると、短期間でこれだけ多くのことを学べたことに驚きます。新たな興味を発掘して実践的に身に付けることができる、とてもありがたい機会でした。今後も積極的にインターンに応募し、さまざまな技術を学ぶつもりです。