はじめに
Turingの基盤AIチームに業務委託として所属している東京科学大学(Institute of Science Tokyo)の藤井です。本記事では、PyTorch Lightningを利用してマルチノード(=2インスタンス以上の環境)で学習を行う際のTipsと、PyTorch Lightningと Lightningの併用により生じる問題の解決策について紹介を行います。
普段はSwallow Projectや横田研究室にて大規模モデルの分散並列学習や低精度学習について研究を行っていますので、そちらもご覧いただけますと幸いです。
2ノード以上で発生する問題
PyTorch Lightningを利用した学習を行う際に、後述する環境変数が正しく設定されていないと、以下のように、2ノードで学習するときにGLOBAL_RANKが重複して設定されてしまう問題が発生します。
このような問題が発生すると、本来あるべきGLOBAL_RANK=8から15のprocessのinitializeを待って、NCCL Timeout(defaultは10分)まで永遠に待ち続けることになります。
Initializing distributed: GLOBAL_RANK: 5, MEMBER: 6/16
Initializing distributed: GLOBAL_RANK: 6, MEMBER: 7/16
Initializing distributed: GLOBAL_RANK: 1, MEMBER: 2/16
Initializing distributed: GLOBAL_RANK: 3, MEMBER: 4/16
Initializing distributed: GLOBAL_RANK: 2, MEMBER: 3/16
Initializing distributed: GLOBAL_RANK: 4, MEMBER: 5/16
Initializing distributed: GLOBAL_RANK: 7, MEMBER: 8/16
Initializing distributed: GLOBAL_RANK: 7, MEMBER: 8/16
Initializing distributed: GLOBAL_RANK: 6, MEMBER: 7/16
Initializing distributed: GLOBAL_RANK: 5, MEMBER: 6/16
Initializing distributed: GLOBAL_RANK: 1, MEMBER: 2/16
Initializing distributed: GLOBAL_RANK: 2, MEMBER: 3/16
Initializing distributed: GLOBAL_RANK: 0, MEMBER: 1/16
Initializing distributed: GLOBAL_RANK: 4, MEMBER: 5/16
Initializing distributed: GLOBAL_RANK: 3, MEMBER: 4/16
では、このような問題はどうして発生しているのでしょうか？
分散学習のconnectionまわりを設定しているコードをまず確認してみましょう。
lightning/fabric/utilities/distributed.py
def _init_dist_connection(
cluster_environment: "ClusterEnvironment",
torch_distributed_backend: str,
global_rank: Optional[int] = None,
world_size: Optional[int] = None,
**kwargs: Any,
) -> None:
"""Utility function to initialize distributed connection by setting env variables and initializing the distributed
process group.
Args:
cluster_environment: ``ClusterEnvironment`` instance
torch_distributed_backend: Backend to use (includes `nccl` and `gloo`)
global_rank: Rank of the current process
world_size: Number of processes in the group
kwargs: Kwargs for ``init_process_group``
Raises:
RuntimeError:
If ``torch.distributed`` is not available
"""
if not torch.distributed.is_available():
raise RuntimeError("torch.distributed is not available. Cannot initialize distributed process group")
if torch.distributed.is_initialized():
log.debug("torch.distributed is already initialized. Exiting early")
return
global_rank = global_rank if global_rank is not None else cluster_environment.global_rank()
world_size = world_size if world_size is not None else cluster_environment.world_size()
os.environ["MASTER_ADDR"] = cluster_environment.main_address
os.environ["MASTER_PORT"] = str(cluster_environment.main_port)
log.info(f"Initializing distributed: GLOBAL_RANK: {global_rank}, MEMBER: {global_rank + 1}/{world_size}")
torch.distributed.init_process_group(torch_distributed_backend, rank=global_rank, world_size=world_size, **kwargs)
先程のログは、log.info(f"Initializing distributed: GLOBAL_RANK: {global_rank}, MEMBER: {global_rank + 1}/{world_size}")で出力されているので、この前の部分ですでにおかしい状態になっているはずです。
_init_dist_connectionを呼び出しているlightning/pytorch/strategies/ddp.pyの実装は以下のようになっています。
ここで怪しいのは、sefl.set_world_sizeです。OpenMPIなどできちんとWORLD_SIZE、RANKを設定しているにも関わらず、おかしな現象が発生するときは大体、ライブラリ側が余計な処理を入れているからのことが多いです。
def setup_distributed(self) -> None:
log.debug(f"{self.__class__.__name__}: setting up distributed...")
reset_seed()
self.set_world_ranks()
self._process_group_backend = self._get_process_group_backend()
assert self.cluster_environment is not None
_init_dist_connection(self.cluster_environment, self._process_group_backend, timeout=self._timeout)
以下に示すように、実装を確認すると、想定しない部分が発見されました。
global_rankは、環境変数から直接設定されるのではなく、self.node_rankとself.num_processとself.local_rankから設定されるという部分です。
def set_world_ranks(self) -> None:
if self.cluster_environment is not None:
self.cluster_environment.set_global_rank(self.node_rank * self.num_processes + self.local_rank)
self.cluster_environment.set_world_size(self.num_nodes * self.num_processes)
そのため、NODE_RANKの環境変数が正しく設定されていないと、default値の0が設定されていしまい、すべてのノードでGLOBAL_RANK = 0 * 8 + LOCAL_RANLとなってしまいます。
修正方法
上記を修正する方法は単純です。
以下のように、正しい環境変数を設定してあげるだけで問題ありません。
(下記はOpenMPIから環境変数を設定していますが、Slurmなどお使いの環境の環境変数に読み替えてください)
global_rank = int(os.getenv('OMPI_COMM_WORLD_RANK', 0))
local_rank = int(os.getenv('OMPI_COMM_WORLD_LOCAL_RANK', 0))
world_size = int(os.getenv('OMPI_COMM_WORLD_SIZE', 1))
local_world_size = int(os.getenv('OMPI_COMM_WORLD_LOCAL_SIZE', 8))
node_rank = global_rank // local_world_size
os.environ['RANK'] = str(global_rank)
os.environ['LOCAL_RANK'] = str(local_rank)
os.environ['WORLD_SIZE'] = str(world_size)
os.environ['LOCAL_WORLD_SIZE'] = str(local_world_size)
os.environ['NODE_RANK'] = str(node_rank)
これで以下のように正しいGLOBAL_RANKが設定できるようになりました。
Initializing distributed: GLOBAL_RANK: 8, MEMBER: 9/16
Initializing distributed: GLOBAL_RANK: 10, MEMBER: 11/16
Initializing distributed: GLOBAL_RANK: 12, MEMBER: 13/16
Initializing distributed: GLOBAL_RANK: 13, MEMBER: 14/16
Initializing distributed: GLOBAL_RANK: 15, MEMBER: 16/16
Initializing distributed: GLOBAL_RANK: 11, MEMBER: 12/16
Initializing distributed: GLOBAL_RANK: 9, MEMBER: 10/16
Initializing distributed: GLOBAL_RANK: 14, MEMBER: 15/16
Initializing distributed: GLOBAL_RANK: 5, MEMBER: 6/16
Initializing distributed: GLOBAL_RANK: 7, MEMBER: 8/16
Initializing distributed: GLOBAL_RANK: 4, MEMBER: 5/16
Initializing distributed: GLOBAL_RANK: 2, MEMBER: 3/16
Initializing distributed: GLOBAL_RANK: 1, MEMBER: 2/16
Initializing distributed: GLOBAL_RANK: 3, MEMBER: 4/16
Initializing distributed: GLOBAL_RANK: 6, MEMBER: 7/16
lightning.pytorch と pytorch_lightning
PyTorch Lightningを利用していて直面するよくある問題として紹介するもう1つの例は、lightning.pytorchとpytorch_lightningの併用時に生じるエラーになります。
このIsuueのように、PyTorch LightningとLightningのpytorchを併用していると発生するエラーが数多く存在します。
今回は中でも、NVIDIAが開発メンテナンスを行っていたNeMo-Alinerを利用する際によく発生しており、公式のIssueでも解決されずに放置されていることが多い下記のエラーを具体例に挙げて解説を行います。
pytorch_lightning/utilities/model_helpers.py", line 42, in is_overridden
[rank0]: raise ValueError("Expected a parent")
[rank0]: ValueError: Expected a parent
以下のように instanceが lightning.pytorchの指定のModuleかどうか判定する箇所があるのですが、ここでpytorch_lightningとlightning.pytorchを併用している弊害が出てきます。実装としてのの実体はどちらも同等なのですが、インスタンス判定ではそうもいきません。
そのため、以下の場所まで進んでしまい、エラーが発生するのです。
(= Trainerの初期化はpytorch_lightning、上述のmoduleチェックではlightning.pytorch側の実装が呼ばれている場合はエラーとなってしまいます)
if parent is None:
_check_mixed_imports(instance)
raise ValueError("Expected a parent")
これを解決する方法は簡単で、import lightning.pytorch as plとimport pytorch_lightning as plが混在しているコードの場合は、どちらかに統一することです。(注意: 関連するライブラリの依存先がどうなっているのかも確認する必要があります)
NeMo-Alignerで発生していたエラーは以下のような方法で修正可能です。
- from pytorch_lightning import TrainerAdd commentMore actions
+ from lightning.pytorch import Trainer
おわりに
本記事では、PyTorch Lightningでマルチノード学習をする際にハマりがちな点に関してなぜエラーが発生するのかと、その修正方法に焦点を当てて簡単に解説を行いました。
これらの問題は、多くの場合、ChatGPT、ClaudeCodeでは解決することが難しく、原因をきちんと考えながら解決する必要が依然としてあります。慣れていない場合ですと数時間かかるような場合もあるかと思いますので、Tipsという形で公開させていただきました。何かの役に立てば幸いです。
普段は、LLMの開発や、低精度学習による高速化を中心に研究開発を行っていますので、よろしければそちらもご覧ください。
Views: 0
