44from collections import OrderedDict
55from contextlib import ExitStack
66from pathlib import Path
7- from typing import TYPE_CHECKING , Any , ContextManager , Dict , List , Mapping , Optional , TypeVar , Union , cast
7+ from typing import TYPE_CHECKING , Any , ContextManager , Dict , List , Literal , Mapping , Optional , TypeVar , Union , cast
88
99import pytorch_lightning as pl
1010import torch
1111import torch .distributed
1212from lightning_fabric .plugins import CheckpointIO , ClusterEnvironment
1313from lightning_fabric .utilities .optimizer import _optimizers_to_device
14+ from megatron .core .distributed import DistributedDataParallelConfig
1415from pytorch_lightning .accelerators import CPUAccelerator
1516from pytorch_lightning .callbacks .progress import TQDMProgressBar
1617from pytorch_lightning .loops import _AutomaticOptimization , evaluation_loop , fit_loop , prediction_loop
3839ConfigT = TypeVar ("ConfigT" )
3940
4041
42+ DDPLiteral = Literal ["megatron" , "pytorch" ]
43+
44+
4145class MegatronStrategy (DDPStrategy , io .IOMixin ):
4246 """Megatron plugin for Pytorch Lightning.
4347
@@ -58,11 +62,11 @@ def __init__(
5862 parallel_devices : Optional [List [torch .device ]] = None ,
5963 cluster_environment = None , # TODO: Add type-hint
6064 checkpoint_io = None , # TODO: Add type-hint
61- no_ddp_communication_hook : bool = True ,
6265 find_unused_parameters : bool = False ,
6366 enable_nemo_ckpt_io : bool = True ,
6467 ckpt_type : TrainerCkptProtocol = TrainerCheckpoint ,
6568 ckpt_include_optimizer : bool = False ,
69+ ddp : Union [DDPLiteral , DistributedDataParallelConfig ] = "megatron" ,
6670 lazy_init : bool = False ,
6771 ** kwargs ,
6872 ) -> None :
@@ -73,7 +77,7 @@ def __init__(
7377 find_unused_parameters = find_unused_parameters ,
7478 ** kwargs ,
7579 )
76- self . no_ddp_communication_hook = no_ddp_communication_hook
80+
7781 self .megatron_callbacks = CallbackConnector ()
7882 self .data_sampler : Optional ['DataSampler' ] = data_sampler
7983 self .tensor_model_parallel_size = tensor_model_parallel_size
@@ -85,6 +89,16 @@ def __init__(
8589 self .lazy_init = lazy_init
8690 self .ckpt_include_optimizer = ckpt_include_optimizer
8791
92+ if ddp == "megatron" :
93+ self .ddp_config = DistributedDataParallelConfig ()
94+ elif isinstance (ddp , DistributedDataParallelConfig ):
95+ self .ddp_config = ddp
96+ elif ddp == "pytorch" :
97+ self .ddp_config = None
98+ self .no_ddp_communication_hook = False
99+ else :
100+ raise ValueError (f"Invalid DDP type: { ddp } )
101+
88102 # used in NVIDIA NGC PyTorch containers
89103 _strategy_lib .enable_nvidia_optimizations ()
90104
@@ -153,6 +167,9 @@ def setup(self, trainer: pl.Trainer) -> None:
153167
154168 # set up optimizers after the wrapped module has been moved to the device
155169 self .setup_optimizers (trainer )
170+
171+ # TODO: Throw an execption if we have a mcore optimizer and no ddp_config
172+
156173 if hasattr (self .precision_plugin , "convert_optimizer" ):
157174 _optimizers = [* self .optimizers ]
158175 _optimizers [0 ] = self .precision_plugin .convert_optimizer (self .optimizers [0 ])
@@ -204,6 +221,7 @@ def setup_megatron_parallel(self, trainer: pl.Trainer) -> None:
204221 precision_plugin = self .precision_plugin ,
205222 vp_size = self .virtual_pipeline_model_parallel_size ,
206223 cpu = isinstance (trainer .accelerator , CPUAccelerator ),
224+ ddp_config = self .ddp_config ,
207225 )
208226 self .model = self .megatron_parallel
209227 self .model .trainer = trainer
@@ -212,6 +230,10 @@ def setup_megatron_parallel(self, trainer: pl.Trainer) -> None:
212230 self .model = self .precision_plugin .convert_module (self .model )
213231 self .model .callbacks .add (getattr (trainer , "callbacks" ))
214232
233+ if hasattr (self , "optimizers" ) and self .optimizers :
234+ for optimizer in self .optimizers :
235+ self .model .callbacks .add (optimizer )
236+
215237 if self .data_sampler :
216238 self .model .callbacks .add (self .data_sampler )
217239
@@ -223,10 +245,11 @@ def setup_megatron_parallel(self, trainer: pl.Trainer) -> None:
223245 def configure_ddp (self ) -> None :
224246 logging .debug (f"{ self .__class__ .__name__ } )
225247 self .model = self ._setup_model (self .model )
226- self ._register_ddp_hooks ()
248+ if self .ddp_config is None :
249+ self ._register_ddp_hooks ()
227250
228251 @override
229- def _setup_model (self , model : nn .Module ) -> DistributedDataParallel :
252+ def _setup_model (self , model : nn .Module ) -> nn . Module :
230253 """Only called when we need to wrap the model for pytorch's ddp."""
231254 from megatron .core import parallel_state
232255
@@ -236,16 +259,19 @@ def _setup_model(self, model: nn.Module) -> DistributedDataParallel:
236259 if app_state .model_parallel_size is not None :
237260 self ._ddp_kwargs ["process_group" ] = parallel_state .get_data_parallel_group ()
238261
239- dist_data_parallel : DistributedDataParallel = super ()._setup_model (model )
240- if self .no_ddp_communication_hook :
241- # When using custom gradient accumulation and allreduce, disable
242- # DDP communication hook that works on the gradient bucket.
243- # Instead, use the custom gradient function and communication hook,
244- # which is defined in the master optimizer wrapper.
245- dist_data_parallel .require_backward_grad_sync = False
246- dist_data_parallel .register_comm_hook (None , noop_hook )
262+ # Only wrap the model if we are not using Megatron's DDP
263+ if not self .ddp_config :
264+ dist_data_parallel : DistributedDataParallel = super ()._setup_model (model )
265+ if self .no_ddp_communication_hook :
266+ # When using custom gradient accumulation and allreduce, disable
267+ # DDP communication hook that works on the gradient bucket.
268+ # Instead, use the custom gradient function and communication hook,
269+ # which is defined in the master optimizer wrapper.
270+ dist_data_parallel .require_backward_grad_sync = False
271+ dist_data_parallel .register_comm_hook (None , noop_hook )
272+ model = dist_data_parallel
247273
248- return dist_data_parallel
274+ return model
249275
250276 def _setup_parallel_ranks (self ) -> None :
251277 self .set_world_ranks ()
@@ -260,7 +286,7 @@ def training_step(self, dataloader_iter, *args: Any, **kwargs: Any) -> STEP_OUTP
260286 kwargs = self ._update_step_kwargs (dataloader_iter , kwargs , "training" )
261287
262288 with self .precision_plugin .train_step_context (): # TODO: Do we need this?
263- return self .model (dataloader_iter , * args , ** kwargs )
289+ return self .model (dataloader_iter , forward_only = False , * args , ** kwargs )
264290
265291 @override
266292 def validation_step (self , dataloader_iter , * args : Any , ** kwargs : Any ) -> STEP_OUTPUT :
@@ -269,7 +295,7 @@ def validation_step(self, dataloader_iter, *args: Any, **kwargs: Any) -> STEP_OU
269295 kwargs = self ._update_step_kwargs (dataloader_iter , kwargs , "validation" )
270296
271297 with self .precision_plugin .val_step_context (): # TODO: Do we need this?
272- return self .model (dataloader_iter , * args , ** kwargs )
298+ return self .model (dataloader_iter , forward_only = True , * args , ** kwargs )
273299
274300 @override
275301 def test_step (self , dataloader_iter , * args : Any , ** kwargs : Any ) -> STEP_OUTPUT :
@@ -278,7 +304,7 @@ def test_step(self, dataloader_iter, *args: Any, **kwargs: Any) -> STEP_OUTPUT:
278304 kwargs = self ._update_step_kwargs (dataloader_iter , kwargs , "test" )
279305
280306 with self .precision_plugin .test_step_context (): # TODO: Do we need this?
281- return self .model (dataloader_iter , * args , ** kwargs )
307+ return self .model (dataloader_iter , forward_only = True , * args , ** kwargs )
282308
283309 @override
284310 def predict_step (self , dataloader_iter , * args : Any , ** kwargs : Any ) -> STEP_OUTPUT :
@@ -287,7 +313,7 @@ def predict_step(self, dataloader_iter, *args: Any, **kwargs: Any) -> STEP_OUTPU
287313 kwargs = self ._update_step_kwargs (dataloader_iter , kwargs , "predict" )
288314
289315 with self .precision_plugin .predict_step_context (): # TODO: Do we need this?
290- return self .model (dataloader_iter , * args , ** kwargs )
316+ return self .model (dataloader_iter , forward_only = True , * args , ** kwargs )
291317
292318 @override
293319 def teardown (self ) -> None :
![github-actions[bot]](https://avatars.githubusercontent.com/in/15368?v=4&size=40){kind=avatar}
0 commit comments