ABEJAでデータサイエンティストをしている岩城です。
今回は株式会社ハイレゾ様のGPUクラウドサービス「GPUSOROBAN」を利用させていただき、NeMo 2.0を用いて大規模言語モデルの事前学習を行いました。 本検証では、あえてNVIDIA公式のDockerコンテナを使用せず環境構築してみたので、その詳細についてお伝えできればと思います。 その理由としては、NeMo関連のライブラリは絶賛開発が進められている中で、新たな昨日や新たなLLMへの対応などが日々進められています。 そうした新規機能が出たとき、柔軟に取り込んで試してみたいと考え、uvで関連ライブラリのバージョン管理をしたいなと思いました。 uvでの管理が可能か、現実的か、どれくらい大変かなどの所感を得るため、この検証に取り組みました。
概要
NeMoとは、生成AIモデルの構築、チューニングなどを行うためのNVIDIA製フレームワークです。
https://docs.nvidia.com/nemo-framework/user-guide/latest/_images/nemo-llm-mm-stack.pngdocs.nvidia.com
NeMo Frameworkを利用するためのAPIとして、現在NeMo 1.0とNeMo 2.0が提供されています。 NeMo 1.0では、学習などのタスクを設定したYAMLファイルを用意し、それを元にジョブを実行する形でした。 NeMo 2.0ではより柔軟に運用するため、Pythonスクリプトで設定を管理する仕様となりました。
弊社では、経済産業省とNEDOが実施する、国内の生成AIの開発力強化を目的としたプロジェクト「GENIAC」の2期でもNeMoを利用していましたが、このときは1.0でした。
NeMoの実行環境については、NVIDIA NGC上にDockerイメージが用意されています。 Megatron-CoreやTransformer Engineなどの複雑な依存関係が解決された環境となっているため、この公式のコンテナを使用することが推奨されています。
今回は、あえてNeMoの実行環境をuvで管理し、拡張性を持たせたNeMo Framework実行環境を構築してみたいと思います。
環境構築
実行環境
今回は、株式会社ハイレゾ様のGPUクラウドサービス「GPUSOROBAN」のGPUサーバを利用させていただきました。
弊社はハイレゾ様とパートナーシップを締結しています。
今回の検証では、NVIDIA H200 × 8基を搭載したGPUインスタンスを使用しました。 ハイレゾ様からいただいた手順書の通りに進めることで簡単にサーバへ接続でき、迅速に作業を開始できました。
サーバ環境については以下の通りです。
| GPUサーバ | H200 × 8基 × 1ノード |
| OS | Ubuntu 24.04 |
| Driver | NVIDIA Driver 570.133.20 |
| CUDA | 12.8 |
また、コンテナ環境については以下の通りです。
| Docker | 28.3.0 |
| Dockerイメージ | nvidia/cuda:12.8.1-cudnn-devel-ubuntu22.04 |
| Python | 3.10.12 |
| uv | 0.7.13 |
NeMo Framework
▶︎ 公式のDockerイメージを利用する場合
- NVIDIA NGCにユーザ登録する必要があります。
- ログインしたら、画面右上のユーザ情報の欄から
Setupを選択します。 Dashboard画面に移動するので、Keys/SecretsのGenerate API KeyからAPIキーを作成します。- 使用する計算機サーバで以下のコマンドを実行し、作成したAPIキーを入力します。
export NGC_API_KEY=<your api key>
echo "$NGC_API_KEY" | docker login nvcr.io --username '$oauthtoken' --password-stdin - NGC上のコンテナが使えるようになります。
docker pull nvcr.io/nvidia/nemo:25.04
NeMo 2.0での中心的なライブラリは nemo-toolkit と nemo-run です。
nemo-toolkit はNeMo Frameworkのライブラリです。LLM, VLM, Diffusionなどの様々なモデルに対応しています。それらのモデルの事前学習やファインチューニングのためのレシピも用意されており、手軽に大規模モデルの学習を実行できます。
nemo-run はNeMo Frameworkによる学習などの設定、実行、管理を効率的に行うためのツールです。
この他にもいくつかNeMo関連のライブラリがあります（データの前処理のためのnemo-curatorなど）。 NeMo Frameworkはmegatron-core, transformer-engine, mamba-ssm, flash-attnなどのライブラリに依存しており、これらがバージョンによっては異なるpytorchに依存していたりします。
uvで環境構築
こちらにNGC上に用意されているコンテナを構成するパッケージ情報が載っていたので、これを参考にしました。
flash-attnなどpytorchに依存しているライブラリは --no-build-isolation オプションでインストールする必要があります。
uvではこれを自動でやってくれる機能があり、 [tool.uv] の no-build-isolation-package で指定します。
pytorchが未インストールの状態でそれに依存したライブラリを入れようとするとエラーになってしまうため、インストールを2段階に分ける必要があります。
uvの optional-dependencies 機能を利用してこれを実現しています。
こうすることで、 uv sync で通常の dependencies で指定されたライブラリのみインストールし、その後 uv sync --extra compile で上のpytorch依存のライブラリをインストールできます。
transformer-engine はpytorch環境で使用する場合、 transformer-engine[pytorch] のようにextraを指定する必要があります。
しかし、 --no-build-isolation-package や tool.uv.dependency-metadata でextraを指定すると構文エラーになります。
なので、 transformer-engine[pytorch] でインストールされる3つのパッケージを明示的に指定しました。
▶︎ pyproject.toml
[project]
name = "project"
version = "0.1.0"
requires-python = "==3.10.*"
dependencies = [
"setuptools<80",
"cython>=3.1.2",
"packaging>=24.2",
"torch==2.7.0",
"pytorch-lightning==2.5.0",
"huggingface-hub[cli]>=0.33.2",
"megatron-core",
"nemo-toolkit[all]",
"nemo-run",
"nemo-curator",
"nvidia-resiliency-ext>=0.3.0",
"ipykernel>=6.27.1",
]
[project.optional-dependencies]
compile = [
"transformer-engine==2.2.0",
"transformer-engine-cu12==2.2.0",
"transformer-engine-torch==2.2.0",
"mamba-ssm",
"flash-attn==2.7.3",
]
[tool.uv]
default-groups = ["dev"]
no-build-isolation-package = [
"transformer-engine",
"transformer-engine-cu12",
"transformer-engine-torch",
"mamba-ssm",
"flash-attn",
]
[tool.uv.sources]
nemo-toolkit = {git = "https://github.com/NVIDIA/NeMo.git", tag = "v2.3.0" }
nemo-run = { git = "https://github.com/NVIDIA/NeMo-Run.git", tag = "v0.4.0" }
nemo-curator = { git = "https://github.com/NVIDIA-NeMo/Curator.git", tag = "v0.8.0" }
megatron-core = { git = "https://github.com/NVIDIA/Megatron-LM.git", tag = "core_v0.12.0" }
mamba-ssm = { git = "https://github.com/state-spaces/mamba.git", tag = "v2.2.2" }
[dependency-groups]
dev = [
"mypy~=1.15.0",
"ruff~=0.11.2",
]
[[tool.uv.dependency-metadata]]
name = "transformer-engine"
version = "2.2.0"
requires-dist = ["torch"]
[[tool.uv.dependency-metadata]]
name = "transformer-engine-cu12"
version = "2.2.0"
requires-dist = ["torch"]
[[tool.uv.dependency-metadata]]
name = "transformer-engine-torch"
version = "2.2.0"
requires-dist = ["torch"]
[[tool.uv.dependency-metadata]]
name = "mamba-ssm"
version = "2.2.2"
requires-dist = ["torch"]
[[tool.uv.dependency-metadata]]
name = "flash-attn"
version = "2.7.3"
requires-dist = ["torch"]
このTOMLファイルを使えば以下のコマンドで環境構築できると思います。 ※ mamba-ssm のインストールに15分以上かかる場合があります。
uv sync uv sync --extra compile
事前学習
上で作成した環境で実際にNeMo 2.0の事前学習を実施してみたいと思います。
とりあえず学習を回してみる
以下は、NeMoのチュートリアルを参考にして簡易的にQwen2.5-7Bの事前学習を実装したサンプルコードです。
▶︎ train.py
from datetime import datetime from nemo.collections import llm import nemo_run as nemo_run def configure_recipe( name: str, checkpoint_dir: str, num_nodes: int = 1, num_gpus_per_node: int = 8 ) -> nemo_run.Partial: recipe = llm.qwen25_7b.pretrain_recipe( name=name, dir=checkpoint_dir, num_nodes=num_nodes, num_gpus_per_node=num_gpus_per_node, ) recipe.model.config.num_layers = 1 recipe.model.config.hidden_size = 128 recipe.trainer.max_steps = 30 recipe.trainer.val_check_interval = 20 return recipe def run_pretrain(exp_name: str, exp_id: str) -> None: recipe = configure_recipe( name=exp_name, checkpoint_dir="./checkpoints", num_nodes=1, num_gpus_per_node=8, ) executor = nemo_run.LocalExecutor() with nemo_run.Experiment( exp_name, base_dir="./nemo_run", id=exp_id, log_level="INFO") as exp: exp.add(recipe, executor=executor, tail_logs=True, name="training") exp.run(detach=False) if __name__ == "__main__": exp_name = "qwen25_7b_tutorial_pretrain" exp_id = str(datetime.now().strftime('%Y%m%d_%H%M%S')) run_pretrain(exp_name, exp_id)
uv run train.py で学習を実行します。
以下のように SUCCEEDED で終了すれば正常に学習が完了したことになります。
[13:26:37] INFO Job training-rqh07q9pnstfl finished: SUCCEEDED launcher.py:160
# The experiment was run with the following tasks: ['training']
# You can inspect and reconstruct this experiment at a later point in time using:
experiment = run.Experiment.from_id("20250722_132550")
experiment.status() # Gets the overall status
experiment.logs("training") # Gets the log for the provided task
experiment.cancel("training") # Cancels the provided task if still running
# You can inspect this experiment at a later point in time using the CLI as well:
nemo experiment status 20250722_132550
nemo experiment logs 20250722_132550 0
nemo experiment cancel 20250722_132550 0
train.py と同じ階層に、学習の設定などが格納された nemo_run ディレクトリとモデルのチェックポイントやログが格納された checkpoints ディレクトリが作成されていると思います。
学習用データはどうなっているのかですが、上のように何も指定しない場合はランダムにトークンを生成する MockDataModule が自動的に適用されます。
リアルデータを使用する
以下のチュートリアルを参考にして、リアルデータで事前学習をします。
- NeMo/tutorials/llm/llama/slimpajama/data_pipeline.ipynb at main · NVIDIA/NeMo · GitHub
- NeMo/tutorials/llm/llama/slimpajama/pretraining.ipynb at main · NVIDIA/NeMo · GitHub
モデルのダウンロード
データの前処理を行うためのトークナイザが必要になります。
今回は、学習するモデルと合わせてQwen2.5-7Bを使用します。学習モデルはホストマシンにダウンロードしておく必要はありませんが、トークナイザとして使用するためには落としておく必要があります。
Hugging Faceからモデルをダウンロードする場合は、Hugging Face上でAccess Tokenを作成し、以下のコマンドで使用したいモデルをダウンロードします。
export HF_TOKEN=xxxxxxxxxxxxxxxxxxxx uv run huggingface-cli download [huggingface/model-paths] --local-dir path/to/models-directory/[model-name] # ex) uv run huggingface-cli download Qwen/Qwen2.5-7B --local-dir ./models/Qwen2.5-7B
※ この時、huggingface-cliの機能か何かで Qwen2.5-7B/.cache のようにキャッシュディレクトリが作成されることがあります。NeMo実行時にモデルを複製する際、モデルのディレクトリ直下にディレクトリがあるとエラーになってしまうため、削除しておく必要があります。
データの前処理
NeMoのGitHubから以下の2つのスクリプトを拝借してきます。
NeMo/tutorials/llm/llama/slimpajama/data/preprocess.py at ce30972017a9803fab4e805588d4358f7aab2854 · NVIDIA/NeMo · GitHub
こちらは以下のように少し手を加えました。
▶︎ preprocess.py
# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. import multiprocessing import os import subprocess import time from typing import Optional from pathlib import Path # from data.extract import _get_shard_list def _split_shards(dataset: list[str], w_size: int) -> list: shards = [] for shard in range(w_size): idx_start = (shard * len(dataset)) // w_size idx_end = ((shard + 1) * len(dataset)) // w_size shards.append(dataset[idx_start:idx_end]) return shards def _get_shard_list(data_dir: str, w_size: int, extension: str = "*zst") -> list: files = Path(data_dir).rglob(extension) files = sorted([str(f) for f in files]) return _split_shards(files, w_size) def _execute_cmd(cmd_tuple: tuple): cmd, task_id = cmd_tuple start_time = time.time() print(f" ****** Task ID {task_id:02d} starts to preprocess {os.path.basename(cmd[2])}...") subprocess.check_call(cmd) print(f" ****** Task ID {task_id:02d} finished preprocessing {os.path.basename(cmd[2])}...") print(f" ****** Task ID {task_id:02d} time elapsed {(time.time() - start_time) / 60:.2f} min.") def preprocess_data( data_dir: str, output_dir: str, dataset_impl: str = "", tokenizer_type: str = "", tokenizer_library: str = "sentencepiece", tokenizer_model: str = "", vocab_file_path: Optional[str] = None, merges_file_path: Optional[str] = None, num_tasks: Optional[int] = None, task_id: Optional[int] = None, extra_args: Optional[list[str]] = None, ): """ Preprocess data for Megatron Core using scripts/nlp_language_modeling/preprocess_data_for_megatron.py Args: data_dir: Path to the directory containing the data to preprocess. output_dir: Path to the directory where the preprocessed data will be saved. dataset_impl: Dataset implementation to use. tokenizer_type: Tokenizer type to use. tokenizer_library: Tokenizer library to use. tokenizer_model: Tokenizer model to use. vocab_file_path: Path to the vocabulary file. merges_file_path: Path to the merges file. num_tasks: Number of tasks to split the data into. task_id: Task ID of run. extra_args: Extra arguments to pass to the preprocess_data_for_megatron.py script. """ if not num_tasks: if "SLURM_ARRAY_TASK_COUNT" in os.environ: num_tasks = int(os.environ["SLURM_ARRAY_TASK_COUNT"]) task_id = int(os.environ["SLURM_ARRAY_TASK_ID"]) else: num_tasks = 1 task_id = 0 shards_to_extract = _get_shard_list(data_dir, num_tasks, extension="*.jsonl") shard_files = shards_to_extract[task_id] cmd = [ "uv", "run", "python", "preprocess_data_for_megatron.py", ] os.makedirs(output_dir, exist_ok=True) final_cmds = [] for split in shard_files: if not split: # Remove empty split continue output_arg = os.path.join(output_dir, os.path.basename(split)) flags = [ f"--input={split}", f"--output-prefix={output_arg}", f"--tokenizer-library={tokenizer_library}", f"--tokenizer-type={tokenizer_type}" if tokenizer_type else f"--tokenizer-model={tokenizer_model}", f"--workers={multiprocessing.cpu_count()}", "--log-interval=100000", "--apply-ftfy", ] if dataset_impl: flags += [f"--dataset-impl={dataset_impl}"] if vocab_file_path: flags += [ f"--vocab-file={vocab_file_path}", "--append-eod", ] if merges_file_path: flags += [f"--merge-file={merges_file_path}"] final_cmd = cmd + flags if extra_args: final_cmd += extra_args final_cmds.append((final_cmd, task_id)) for cmd in final_cmds: _execute_cmd(cmd)
今回は、Hugging Faceからサイズが小さめのcommon-crawlのデータセットをダウンロードして使用します。
keirp/common_crawl_sample · Datasets at Hugging Face
NeMoで自前のjsonlデータを使用するためには、トークナイザを用いてBINとIDXデータに変換する必要があります。
以下のコードでは、Hugging Faceからデータをダウンロードしてjsonlで保存し、それをBIN, IDXに変換するという処理を行っています。
モデルは ~/models/Qwen2.5-7B に格納しました。
▶︎ data_preprocess.py
必要に応じてモデルのパスを変更してください。
from datasets import load_dataset import json from pathlib import Path import nemo_run from preprocess import preprocess_data def import_dataset(name: str, output_dir: str) -> None: ds = load_dataset(name) file_name = f"{name.split('/')[-1]}.jsonl" data_dir = Path(output_dir) data_dir.mkdir(exist_ok=True) for _, split_data in ds.items(): output_file = data_dir / file_name with open(output_file, 'w', encoding='utf-8') as f: for example in split_data: json.dump(example, f, ensure_ascii=False) f.write('\n') def preprocess(data_dir: str, output_dir: str): with nemo_run.Experiment( "data_preprocess", base_dir="./data_preprocess", log_level="INFO") as exp: preprocess_task = nemo_run.Partial( preprocess_data, data_dir=data_dir, output_dir=output_dir, tokenizer_library="huggingface", tokenizer_type=str(Path.home() / "models" / "Qwen2.5-7B"), tokenizer_model=str(Path.home() / "models" / "Qwen2.5-7B"), project_root=str(Path.home() / "llm-train-nemo2"), ) executor = nemo_run.LocalExecutor() exp.add(preprocess_task, executor=executor, name="preprocess") exp.run(detach=False, tail_logs=True) if __name__ == "__main__": import_dataset( name="keirp/common_crawl_sample", output_dir="./data" ) preprocess( data_dir="./data", output_dir="./data", )
uv run data_preprocess.py でこのコードを実行すると、以下のデータが作成されると思います。
data/ ├── common_crawl_sample.jsonl ├── common_crawl_sample.jsonl_text_document.bin └── common_crawl_sample.jsonl_text_document.idx
また、実行時の設定が data_preprocess/ に保存されます。
学習の実行
最終的なディレクトリ構成は以下です。
. ├── checkpoints/ ├── data/ │ ├── common_crawl_sample.jsonl │ ├── common_crawl_sample.jsonl_text_document.bin │ └── common_crawl_sample.jsonl_text_document.idx ├── data_preprocess/ ├── nemo_run/ ├── data_preprocess.py ├── preprocess_for_megatron.py ├── preprocess.py ├── train_common_crawl.py └── train.py
そして、実行スクリプト train_common_crawl.py は以下のように作成しました。
from datetime import datetime from pathlib import Path import lightning.pytorch as pl from nemo.collections import llm from nemo.collections.common.tokenizers import AutoTokenizer import nemo_run as nemo_run def data_module( global_batch_size: int = 256, micro_batch_size: int = 4, seq_length: int = 8192, ) -> nemo_run.Config[pl.LightningDataModule]: return nemo_run.Config( llm.PreTrainingDataModule, paths=["./data/common_crawl_sample.jsonl_text_document"], seq_length=seq_length, global_batch_size=global_batch_size, micro_batch_size=micro_batch_size, tokenizer=nemo_run.Config( AutoTokenizer, pretrained_model_name=str(Path.home() / "models" / "Qwen2.5-7B"), ), split="99,8,2", num_workers=27, ) def configure_recipe( name: str, checkpoint_dir: str, num_nodes: int = 1, num_gpus_per_node: int = 8 ) -> nemo_run.Partial: recipe = llm.qwen25_7b.pretrain_recipe( name=name, dir=checkpoint_dir, num_nodes=num_nodes, num_gpus_per_node=num_gpus_per_node, ) recipe.model.config.num_layers = 1 recipe.model.config.hidden_size = 128 recipe.trainer.max_steps = 10 recipe.trainer.val_check_interval = 10 recipe.data = data_module( global_batch_size=32, micro_batch_size=1, ) return recipe def run_pretrain(exp_name: str, exp_id: str) -> None: recipe = configure_recipe( name=exp_name, checkpoint_dir="./checkpoints", num_nodes=1, num_gpus_per_node=8, ) executor = nemo_run.LocalExecutor() with nemo_run.Experiment( exp_name, base_dir="./nemo_run", id=exp_id, log_level="INFO") as exp: exp.add(recipe, executor=executor, tail_logs=True, name="training") exp.run(detach=False) if __name__ == "__main__": exp_name = "qwen25_7b_tutorial_pretrain_common_crawl" exp_id = str(datetime.now().strftime('%Y%m%d_%H%M%S')) run_pretrain(exp_name, exp_id)
uv run train_common_crawl.py で学習を実行します。
上と同じように SUCCEEDED が出ればOKです。
並列化
並列化は学習効率を高めるための重要な技術です。NeMo Frameworkでは複数の並列化技術をサポートしています。
今回の環境で、並列化処理も問題なく動作するか検証したいと思います。
Model Parallelism
モデルのパラメータを複数のGPUに分散させてGPUあたりの使用メモリを削減する手法です。
Tensor parallelism
https://docs.nvidia.com/nemo-framework/user-guide/24.09/_images/tp.gifdocs.nvidia.com
モデルの、個別レイヤのパラメータテンソルを複数のGPUに分散させる手法です。
GPUあたりの使用メモリを削減する代わりにGPUごとのカーネルの処理負荷が小さくなるため、CPUの間接的な使用量が増加します。
recipe.trainer.strategy.tensor_model_parallel_sizeを2以上にすることで利用できます。
GPU数が足りていれば、リアルデータでの事前学習のコードに上の変数指定を付け足すだけで問題なく学習が回るはずです。
train_common_crawl.py に以下の変更を加え、 train_tensor_parallelism.py を作成します。
def configure_recipe(
...
recipe.data = data_module(
global_batch_size=32,
micro_batch_size=1,
)
+ recipe.trainer.strategy.tensor_model_parallel_size = 2
...
if __name__ == "__main__":
- exp_name = "qwen25_7b_tutorial_pretrain_common_crawl"
+ exp_name = "qwen25_7b_tutorial_pretrain_tensor_parallelism"
...
uv run train_tensor_parallelism.py で学習を実行します。
Pipeline parallelism
https://docs.nvidia.com/nemo-framework/user-guide/24.09/_images/pp.gifdocs.nvidia.com
連続したレイヤやニューラルネットの一部を複数のGPUに割り当てる手法です。
各GPUは異なる層のネットワークをシーケンシャルに処理します。例えば、GPU0は0 ~ 8層まで、GPU1は9 ~ 17層まで… という具合です。
recipe.trainer.strategy.pipeline_model_parallel_sizeを2以上にすることで利用できます。- チュートリアルのコードでは
recipe.model.config.num_layer = 1となっています。Pipeline parallelismを利用するためには、num_layerがpipeline_model_parallel_sizeで割り切れる数値になっている必要があります。
同様に train_common_crawl.py に以下の変更を加え train_pipeline_parallelism.py を作成します。
def configure_recipe( ... - recipe.model.config.num_layers = 1 + recipe.model.config.num_layers = 2 ... recipe.data = data_module( global_batch_size=32, micro_batch_size=1, ) + recipe.trainer.strategy.pipeline_model_parallel_size = 2 ... if __name__ == "__main__": - exp_name = "qwen25_7b_tutorial_pretrain_common_crawl" + exp_name = "qwen25_7b_tutorial_pretrain_pipeline_parallelism" ...
uv run train_pipeline_parallelism.py で学習を実行します。
Interleave pipeline parallel schedule
各GPUの計算を単一の連続したブロックではなく、複数のレイヤのサブセット（モデルチャンク）に分割する手法です。
例えば、4つのGPUがそれぞれ4つの層を処理するとしたとき、さらにそれを2層ごとに分割させた場合次のようになります。GPU0は0, 1層と8, 9層、GPU1は2, 3層と10, 11層…という形です。
これは、層ごとの計算負荷に分散がある場合、GPUあたりの負荷をより分散させることができる技術だと考えられます。
以下の設定で利用できます。
recipe.model.config.num_layers *=* 4recipe.trainer.strategy.pipeline_model_parallel_size = 2recipe.trainer.strategy.virtual_pipeline_model_parallel_size = 2
以下の内容で train_interleave_pipeline_parallelism.py を作成します。
def configure_recipe( ... - recipe.model.config.num_layers = 1 + recipe.model.config.num_layers = 4 ... recipe.data = data_module( global_batch_size=32, micro_batch_size=1, ) + recipe.trainer.strategy.pipeline_model_parallel_size = 2 + recipe.trainer.strategy.virtual_pipeline_model_parallel_size = 2 ... if __name__ == "__main__": - exp_name = "qwen25_7b_tutorial_pretrain_common_crawl" + exp_name = "qwen25_7b_tutorial_interleave_pretrain_pipeline_parallelism" ...
uv run train_interleave_pipeline_parallelism.py で学習を実行します。
Activation Partitioning
効率的なactivationの分散手法で、シーケンス長が長い時やGPUあたりのマイクロバッチサイズが大きい設定で学習するときに有効な手法になります。
Sequence parallelism
https://docs.nvidia.com/nemo-framework/user-guide/24.09/_images/sp.gifdocs.nvidia.com
テンソルレベルのモデル並列化を、transformer層の系列的次元に沿って計算経路とactivationメモリを複数のGPUに分散させることで拡張します。つまり、tensor parallelismをもっと細かくするということです。
これは、特に上記まででは並列化されなかったレイヤの一部に対して有効な技術になります。
tensor_model_parallel_sizeを2以上にするsequence_parallel = Trueとする
以下の内容で train_sequence_parallelism.py を作成します。
def configure_recipe(
...
recipe.data = data_module(
global_batch_size=32,
micro_batch_size=1,
)
+ recipe.trainer.strategy.tensor_model_parallel_size = 2
+ recipe.trainer.strategy.sequence_parallel = True
...
if __name__ == "__main__":
- exp_name = "qwen25_7b_tutorial_pretrain_common_crawl"
+ exp_name = "qwen25_7b_tutorial_pretrain_sequence_parallelism"
...
uv run train_sequence_parallelism.py で学習を実行します。
Context parallelism
ニューラルネットのactivationsの処理を複数のGPUで分散する手法です。
系列的次元の入力テンソルを分割します。
Sequence parallelismが特定の層のactivationsを分割するのに対して、context parallelismは全ての層のactivationsを分割します。
ロングコンテキストでの学習時に利用することが推奨されています。
recipe.trainer.strategy.context_parallel_sizeを2以上にすることで利用できます
このままだと以下のようなエラーが出ると思います。
[DEBUG | DotProductAttention]: Running with config={'transformer_engine_version': '2.2.0', 'compute_capability': 'sm90', 'flash_attn_version': '2.7.3', 'flash_attn_3_version': 'not installed', 'cudnn_version': '9.8.0', 'qkv_type': <class 'torch.Tensor'>, 'qkv_dtype': torch.bfloat16, 'qkv_layout': 'sbhd_sbhd_sbhd', 'batch_size': 1, 'num_heads': 14, 'num_gqa_groups': 2, 'max_seqlen_q': 8192, 'max_seqlen_kv': 8192, 'head_dim_qk': 4, 'head_dim_v': 4, 'attn_mask_type': 'causal', 'window_size': (-1, 0), 'alibi_slopes_shape': None, 'core_attention_bias_type': 'no_bias', 'core_attention_bias_shape': None, 'core_attention_bias_requires_grad': False, 'pad_between_seqs': False, 'attention_dropout': 0.0, 'context_parallel': True, 'deterministic': False, 'is_training': False, 'fp8': False, 'fp8_meta': {'fp8_checkpoint': False, 'fp8_group': None}, 'inference_params': None}
DEBUG:DotProductAttention:Running with config={'transformer_engine_version': '2.2.0', 'compute_capability': 'sm90', 'flash_attn_version': '2.7.3', 'flash_attn_3_version': 'not installed', 'cudnn_version': '9.8.0', 'qkv_type': <class 'torch.Tensor'>, 'qkv_dtype': torch.bfloat16, 'qkv_layout': 'sbhd_sbhd_sbhd', 'batch_size': 1, 'num_heads': 14, 'num_gqa_groups': 2, 'max_seqlen_q': 8192, 'max_seqlen_kv': 8192, 'head_dim_qk': 4, 'head_dim_v': 4, 'attn_mask_type': 'causal', 'window_size': (-1, 0), 'alibi_slopes_shape': None, 'core_attention_bias_type': 'no_bias', 'core_attention_bias_shape': None, 'core_attention_bias_requires_grad': False, 'pad_between_seqs': False, 'attention_dropout': 0.0, 'context_parallel': True, 'deterministic': False, 'is_training': False, 'fp8': False, 'fp8_meta': {'fp8_checkpoint': False, 'fp8_group': None}, 'inference_params': None}
[DEBUG | DotProductAttention]: Disabling FlashAttention 2 due to unsupported head_dim_qk and head_dim_v. Supported: head_dim_qk = head_dim_v, head_dim_qk %8 = 0, head_dim_qk <= 256 (>192 requires sm80/90/100+). Found: head_dim_qk = 4, head_dim_v = 4, on sm9.0.
DEBUG:DotProductAttention:Disabling FlashAttention 2 due to unsupported head_dim_qk and head_dim_v. Supported: head_dim_qk = head_dim_v, head_dim_qk %8 = 0, head_dim_qk <= 256 (>192 requires sm80/90/100+). Found: head_dim_qk = 4, head_dim_v = 4, on sm9.0.
[DEBUG | DotProductAttention]: Disabling UnfusedDotProductAttention as it does not support context parallelism
DEBUG:DotProductAttention:Disabling UnfusedDotProductAttention as it does not support context parallelism
[DEBUG | DotProductAttention]: Disabling FusedAttention as no backend supports the provided input
DEBUG:DotProductAttention:Disabling FusedAttention as no backend supports the provided input
[DEBUG | DotProductAttention]: Available backends = {FlashAttention=False, FusedAttention=False, UnfusedDotProductAttention=False}
DEBUG:DotProductAttention:Available backends = {FlashAttention=False, FusedAttention=False, UnfusedDotProductAttention=False}
[DEBUG | DotProductAttention]: Selected backend = NoBackend
DEBUG:DotProductAttention:Selected backend = NoBackend
このエラーは、 hidden_size を以下のようにQwen2.5-7Bのデフォルト設定に合わせることで回避できます。
recipe.model.config.hidden_size = 3584
これについては事項の「DotProductAttention エラーについて」で解説します。
以下の内容で train_context_parallelism.py を作成します。
def configure_recipe(
...
recipe.model.config.num_layers = 1
- recipe.model.config.hidden_size = 128
+ recipe.model.config.hidden_size = 3584
recipe.trainer.max_steps = 10
recipe.trainer.val_check_interval = 10
recipe.data = data_module(
global_batch_size=32,
micro_batch_size=1,
)
+ recipe.trainer.strategy.context_parallel_size = 2
...
if __name__ == "__main__":
- exp_name = "qwen25_7b_tutorial_pretrain_common_crawl"
+ exp_name = "qwen25_7b_tutorial_pretrain_context_parallelism"
...
uv run train_context_parallelism.py で学習を実行します。
DotProductAttention エラーについて
これは、head dimの次元数がサポートされていないサイズになっていることが原因です。
head dimの次元数は以下の式で決まります。
このhead dimが8の倍数の数値になっている必要があるようです。
上で参考にしたNeMoの pretraining.ipynb では、 recipe.model.config.hidden_size = 128 となっています。 recipe.model.config.num_attention_heads = 28 がデフォルトのため、128 / 28 = 4.57… となり、エラーが発生しました。
そのため、head dimが8の倍数になるよう設定することで、このエラーを回避できます。hidden_size = 3584はQwen2.5-7Bのデフォルト設定です。
hidden_size = 224などでも動作しますが、この model.config 周りはNeMoの中でそのモデルに合わせて設定されているので下手に触らないほうが良いと思われます。
この問題は context parallelism だけでなく tensor parallelism などの他の並列化手法でも発生することがあります。
ということで、並列化の機能も動作することが確認できました！
最後に
NeMoの実行環境を公式のDockerイメージを使うことなく作成することのメリットの一つとして、NeMo関連のライブラリのリリースされていないバージョンをいち早く柔軟に取り入れられることが挙げられます。
実際、検証時点 (2025/07/14) のNeMo 2.0のリリースバージョン v2.3.0 ではQwen3の設定がサポートされていませんでした。
しかし、GitHubにはQwen3をサポートするコードが存在していたため、リポジトリのmainブランチを指定してnemo-toolkitをインストールすることでそうした機能を早く使うことができました。
※ 記事執筆時点 (2025/07/28) では v2.4.0 がリリースされ、Qwen3のサポートが含まれました。
他方で、最初は依存関係の解決にかなりの時間を要したので、特別な事情がない限り公式コンテナを使う方がいいと思いました。複数ノードでの学習を進めるための複雑な環境設定
謝辞
本検証において、GPUクラウドサービス「GPUSOROBAN」のGPUサーバをご提供いただきました株式会社ハイレゾ様に、心より御礼申し上げます。
We are hiring!
ABEJAは、テクノロジーの社会実装に取り組んでいます。 技術はもちろん、技術をどのようにして社会やビジネスに組み込んでいくかを考えるのが好きな方は、下記採用ページからエントリーください！ （新卒の方やインターンシップのエントリーもお待ちしております！）
