# Copyright (c) 2022, 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. # # flake8: noqa # pylint: skip-file import itertools from functools import partial from typing import Any, Optional import torch from lightning.pytorch.accelerators import CPUAccelerator from lightning.pytorch.trainer.trainer import Trainer from omegaconf.dictconfig import DictConfig from nemo.collections.nlp.data.language_modeling.megatron.data_samplers import MegatronPretrainingSampler from nemo.collections.nlp.models.language_modeling.megatron_base_model import MegatronBaseModel from nemo.collections.nlp.modules.common.megatron.build_model import build_model from nemo.collections.nlp.modules.common.megatron.module import Float16Module, MegatronModule from nemo.collections.nlp.modules.common.megatron.utils import ( average_losses_across_data_parallel_group, get_all_params_for_weight_decay_optimization, get_linear_layer, get_params_for_weight_decay_optimization, init_method_normal, scaled_init_method_normal, ) from nemo.collections.nlp.parts.utils_funcs import get_last_rank, torch_dtype_from_precision from nemo.collections.vision.data.megatron.data_samplers import MegatronVisionPretrainingRandomSampler from nemo.collections.vision.data.megatron.vit_dataset import build_train_valid_datasets from nemo.collections.vision.modules.vit.vit_backbone import VitBackbone, VitMlpHead from nemo.core.classes.common import PretrainedModelInfo from nemo.utils import logging try: from megatron.core import parallel_state from megatron.core.pipeline_parallel.schedules import get_forward_backward_func HAVE_MEGATRON_CORE = True except (ImportError, ModuleNotFoundError): HAVE_MEGATRON_CORE = False try: from megatron.core.num_microbatches_calculator import get_num_microbatches except (ImportError, ModuleNotFoundError): logging.warning("Megatron num_microbatches_calculator not found, using Apex version.") from apex.transformer.pipeline_parallel.utils import get_num_microbatches class VitClassificationModel(MegatronModule): """Vision Transformer Model.""" def __init__( self, model_cfg, model_parallel_config, num_classes, finetune=False, pre_process=True, post_process=True ): super(VitClassificationModel, self).__init__() scaled_init_method = ( scaled_init_method_normal(model_cfg.init_method_std, model_cfg.num_layers) if model_cfg.use_scaled_init_method else init_method_normal(model_cfg.init_method_std) ) self.config = model_parallel_config self.hidden_size = model_cfg.hidden_size self.num_classes = num_classes self.finetune = finetune self.pre_process = pre_process self.post_process = post_process self.backbone = VitBackbone( model_cfg, model_parallel_config, init_method=init_method_normal(model_cfg.init_method_std), scaled_init_method=scaled_init_method, pre_process=self.pre_process, post_process=self.post_process, single_token_output=True, ) if self.post_process: if not self.finetune: self.head = VitMlpHead(self.hidden_size, self.num_classes) else: self.head = get_linear_layer(self.hidden_size, self.num_classes, torch.nn.init.zeros_) def set_input_tensor(self, input_tensor): """See megatron.model.transformer.set_input_tensor()""" self.backbone.set_input_tensor(input_tensor) def forward(self, input): hidden_states = self.backbone(input) if self.post_process: hidden_states = self.head(hidden_states) hidden_states = hidden_states.contiguous() return hidden_states class MegatronVitClassificationModel(MegatronBaseModel): """Megatron Vision Transformer Model.""" def __init__(self, cfg: DictConfig, trainer: Trainer): if not HAVE_MEGATRON_CORE: raise ImportError( "megatron-core was not found. Please see the NeMo README for installation instructions: https://github.com/NVIDIA/NeMo#megatron-gpt." ) super().__init__(cfg, trainer=trainer) self._validate_trainer() # TODO(yuya): clean up all default values self.megatron_amp_O2 = cfg.get('megatron_amp_O2', False) if not self.megatron_amp_O2 and self.cfg.get('virtual_pipeline_model_parallel_size', None): raise ValueError('Virtual pipeline model parallel is only supported when using megatron_amp_O2') # build_model returns a list of modules which are used for interleaved pipeline parallelism if isinstance(self.trainer.accelerator, CPUAccelerator): self.model = build_model( model_provider_func=self.model_provider_func, wrap_with_ddp=False, on_cpu=True, virtual_pipeline_model_parallel_size=self.cfg.get('virtual_pipeline_model_parallel_size', None), ) else: self.model = build_model( model_provider_func=self.model_provider_func, wrap_with_ddp=False, virtual_pipeline_model_parallel_size=self.cfg.get('virtual_pipeline_model_parallel_size', None), ) # if we're not using interleaved, then self.model is a module. if self.cfg.get('virtual_pipeline_model_parallel_size', None) is None: self.model = self.model[0] if self.megatron_amp_O2: if not self.with_distributed_adam: # Pre-allocate the model on GPU to have master parameters allocated on the same device with matching data type if isinstance(self.model, list): for module in self.model: module.cuda(torch.cuda.current_device()) else: self.model.cuda(torch.cuda.current_device()) # Model wrapper to convert both model and inputs to half precision if isinstance(self.model, list): converted_model = [] for module in self.model: converted_model.append( Float16Module(config=self.model_parallel_config, module=module, precision=cfg.precision) ) self.model = converted_model else: self.model = Float16Module( config=self.model_parallel_config, module=self.model, precision=cfg.precision ) self.autocast_dtype = torch_dtype_from_precision(self.trainer.precision) self.enable_autocast = ( True if (not self.megatron_amp_O2) and (self.autocast_dtype in [torch.float16, torch.bfloat16]) else False ) self.transformer_engine = cfg.get('transformer_engine', False) # Convert the global-batch-based profile index to micro-batch index if hasattr(self, '_nsys_profile_enabled') or hasattr(self, '_memory_profile_enabled'): mp_size = cfg.get('tensor_model_parallel_size', 1) * cfg.get('pipeline_model_parallel_size', 1) data_parallel_world_size = trainer.world_size // mp_size grad_accum_steps = cfg.get('global_batch_size') // (cfg.get('micro_batch_size') * data_parallel_world_size) if hasattr(self, '_nsys_profile_enabled'): self._nsys_profile_start_step *= grad_accum_steps self._nsys_profile_end_step *= grad_accum_steps if hasattr(self, '_memory_profile_enabled'): self._memory_profile_start_step *= grad_accum_steps self._memory_profile_end_step *= grad_accum_steps self.get_attention_mask_from_fusion = self.cfg.get('get_attention_mask_from_fusion', True) self.initialize_ub = self.cfg.get('ub_tp_comm_overlap', False) def get_module_list(self): if isinstance(self.model, list): return [model.module if isinstance(model, Float16Module) else model for model in self.model] elif isinstance(self.model, Float16Module): return [self.model.module] else: return [self.model] def model_provider_func(self, pre_process, post_process): """Model depends on pipeline paralellism.""" model = VitClassificationModel( model_cfg=self.cfg, model_parallel_config=self.model_parallel_config, num_classes=self.cfg.get("num_classes"), # TODO(yuya): clean this up finetune=self.cfg.get("finetune", False), pre_process=pre_process, post_process=post_process, ) return model def setup_optimizer_param_groups(self): """ModelPT override. Optimizer will get self._optimizer_param_groups""" if self.cfg.get('do_layer_norm_weight_decay', False): if isinstance(self.model, list): self._optimizer_param_groups = get_all_params_for_weight_decay_optimization(self.model) else: self._optimizer_param_groups = get_all_params_for_weight_decay_optimization([self.model]) else: self._optimizer_param_groups = get_params_for_weight_decay_optimization(self.model) def configure_optimizers(self): if self.with_distributed_adam: # Disable overlapped grad sync for layer norm grads when # sequence parallelism is enabled for param in self.parameters(): if getattr(param, 'sequence_parallel_enabled', False): param._disable_greedy_grad_copy = not self.megatron_amp_O2 param._disable_overlap_grad_sync = True # KJJ - Copied this entire block, up to "return" here blindly from megatron_gpt_model.py # Initialize parameter buckets for overlapped grad and param syncs # Note: Params with disabled overlapping are put in the # last param bucket buckets = [] if self.cfg.get('virtual_pipeline_model_parallel_size', None) is not None: # Initialize a bucket for each virtual pipeline stage for module in self.model: if isinstance(module, Float16Module): module = module.module stage_bucket = [] # for layer in module.language_model.encoder.layers: for layer in module.backbone.transformer.layers: stage_bucket.extend( p for p in layer.parameters() if not getattr(p, '_disable_overlap_grad_sync', False) ) buckets.append(stage_bucket) else: # Initialize a bucket for each Transformer layer modules = self.model if isinstance(self.model, list) else [self.model] for module in modules: if isinstance(module, Float16Module): module = module.module # for layer in module.language_model.encoder.layers: for layer in module.backbone.transformer.layers: buckets.append( [p for p in layer.parameters() if not getattr(p, '_disable_overlap_grad_sync', False)] ) buckets.reverse() used_params = set() for bucket in buckets: used_params.update(bucket) buckets[-1].extend(p for p in self.parameters() if p not in used_params) self.distributed_adam_buckets = buckets return super().configure_optimizers() def forward(self, tokens): output_tensor = self.model(tokens) return output_tensor def fwd_bwd_step(self, dataloader_iter, forward_only): # handle asynchronous grad reduction no_sync_func = None grad_sync_func = None param_sync_func = None if not forward_only and self.with_distributed_adam: no_sync_func = partial( self._optimizer.no_sync, greedy_grad_copy=self.megatron_amp_O2, ) grad_sync_func = self.reduce_overlap_gradients param_sync_func = self.sync_overlap_parameters # pipeline schedules will get these from self.model.config for module in self.get_module_list(): module.config.no_sync_func = no_sync_func module.config.grad_sync_func = grad_sync_func module.config.param_sync_func = param_sync_func # run forward and backwards passes for an entire global batch # we do this inside training_step to support pipeline parallelism fwd_bwd_function = get_forward_backward_func() # TODO @akhattar: add num_micro_batches_with_partial_activation_checkpoints when ready losses_reduced_per_micro_batch = fwd_bwd_function( forward_step_func=self.get_forward_output_and_loss_func(), data_iterator=dataloader_iter, model=[self.model], num_microbatches=get_num_microbatches(), forward_only=forward_only, seq_length=self.cfg.encoder_seq_length, micro_batch_size=self.cfg.micro_batch_size, ) # only the last stages of the pipeline return losses if losses_reduced_per_micro_batch: if (not forward_only) or self.cfg.data.get('validation_drop_last', True): # average loss across micro batches loss_tensors_list = [loss_reduced['loss'] for loss_reduced in losses_reduced_per_micro_batch] loss_tensor = torch.stack(loss_tensors_list) loss_mean = loss_tensor.mean() acc_tensors_list = [loss_reduced['accuracy'] for loss_reduced in losses_reduced_per_micro_batch] acc_tensor = torch.stack(acc_tensors_list) accuracy_mean = acc_tensor.mean() else: # Get the total loss since micro batches sizes are not uniform raise NotImplementedError("Losses of micro batches sizes must be uniform!") else: # we're not on the last pipeline stage so no losses if forward_only: loss_mean = [] accuracy_mean = [] else: loss_mean = torch.tensor(0.0).cuda() accuracy_mean = loss_mean.copy() return loss_mean, accuracy_mean def initialize_ub_func(self): ub_cfgs = self.cfg.get('ub_tp_comm_overlap_cfg', None) if ub_cfgs is None: warnings.warn( "Couldn't find TP config. Please check the path correctness. Initializing TP comm overlap with the default config." ) input_shape = [ self.cfg.get('encoder_seq_length') * self.cfg.get('micro_batch_size'), self.cfg.get('hidden_size'), ] te_module.base.initialize_ub( shape=input_shape, tp_size=self.cfg.get('tensor_model_parallel_size'), use_fp8=self.cfg.get('fp8'), ub_cfgs=ub_cfgs, ) self.initialize_ub = False def training_step(self, dataloader_iter): """ Our dataloaders produce a micro-batch and then we fetch a number of microbatches depending on the global batch size and model parallel size from the dataloader to produce a list of microbatches. Batch should be a list of microbatches and those microbatches should on CPU. Microbatches are then moved to GPU during the pipeline. The list of microbatches is then piped through the pipeline using Apex fwd/bwd functions. """ # Initialize userbuffer communicators. if self.initialize_ub: self.initialize_ub_func() # we zero grads here because we also call backward in the megatron-core fwd/bwd functions self._optimizer.zero_grad() loss_mean, _ = self.fwd_bwd_step(dataloader_iter, False) # when using sequence parallelism, the sequence parallel layernorm grads must be all-reduced if self.cfg.get('tensor_model_parallel_size', 1) > 1 and self.cfg.get('sequence_parallel', False): self.allreduce_sequence_parallel_gradients() if self.with_distributed_adam: # KJJ - Added this block from megatron_gpt_model. It says it's not necessary # and it's not clear if the remaining "if not" logic is still needed. # keeping it for now, but might need to delete one or both of these. # synchronize asynchronous grad reductions # note: not necessary, but reduces performance degradation # from multiple simultaneous NCCL calls self._optimizer._finish_bucket_grad_sync() # launch grad reductions # Note: grads in first pipeline stage have already been # reduced assert ( self.cfg.get("virtual_pipeline_model_parallel_size", None) is None ), "vpp is not supported yet in MegatronVitClassificationModel" if not parallel_state.is_pipeline_first_stage(): self.reduce_overlap_gradients() elif self.megatron_amp_O2: # # when using pipeline parallelism grads must be all-reduced after the pipeline (not asynchronously) self._optimizer.allreduce_main_grads() else: # async grad allreduce is not currently implemented for O1/autocasting mixed precision training # so we all-reduce gradients after the pipeline self.allreduce_gradients() # @sangkug we think this is causing memory to blow up (hurts perf) ## logging # we can only log on one rank if it is rank zero so we broadcast from last rank # we can avoid this broadcast by updating the PTL log function to accept specific ranks torch.distributed.broadcast(loss_mean, get_last_rank()) if self.cfg.precision in [16, '16', '16-mixed']: loss_scale = self.trainer.precision_plugin.scaler._scale if loss_scale is not None: self.log('loss_scale', loss_scale, batch_size=1) self.log('reduced_train_loss', loss_mean, prog_bar=True, rank_zero_only=True, batch_size=1) lr = self._optimizer.param_groups[0]['lr'] self.log('lr', lr, rank_zero_only=True, batch_size=1) self.log('global_step', self.trainer.global_step + 1, prog_bar=True, rank_zero_only=True, batch_size=1) self.log( 'consumed_samples', self.compute_consumed_samples(self.trainer.global_step + 1 - self.init_global_step), prog_bar=True, rank_zero_only=True, batch_size=1, ) return loss_mean def backward(self, *args, **kwargs): """LightningModule hook to do backward. We want this to do nothing since we run backward in the fwd/bwd functions from apex. No need to call it here. """ pass def optimizer_zero_grad(self, *args, **kwargs): """LightningModule hook to zero grad. We want this to do nothing as we are zeroing grads during the training_step. """ pass def _append_sequence_parallel_module_grads(self, module, grads): """Helper method for allreduce_sequence_parallel_gradients""" for param in module.parameters(): sequence_parallel_param = getattr(param, 'sequence_parallel', False) if sequence_parallel_param and param.requires_grad: if self.megatron_amp_O2: grad = param.main_grad else: grad = param.grad grads.append(grad.data) def allreduce_sequence_parallel_gradients(self): """All-reduce layernorm parameters across model parallel nodes when sequence parallelism is used. Modified from megatron-lm: https://gitlab-master.nvidia.com/ADLR/megatron-lm/-/blob/3f91f09bb2ab32f9904b47f46f19d2fc3f518ed8/megatron/training.py#L425 """ grads = [] if isinstance(self.model, list): for module in self.model: self._append_sequence_parallel_module_grads(module, grads) else: self._append_sequence_parallel_module_grads(self.model, grads) coalesced = torch._utils._flatten_dense_tensors(grads) torch.distributed.all_reduce(coalesced, group=parallel_state.get_tensor_model_parallel_group()) for buf, synced in zip(grads, torch._utils._unflatten_dense_tensors(coalesced, grads)): buf.copy_(synced) def get_forward_output_and_loss_func(self, validation_step=False): def loss_func(labels, output_tensor): logits = output_tensor.contiguous().float() loss = torch.nn.functional.cross_entropy(logits, labels) outputs = torch.argmax(logits, -1) correct = (outputs == labels).float() accuracy = torch.mean(correct) averaged_loss = average_losses_across_data_parallel_group([loss, accuracy]) return loss, {"loss": averaged_loss[0], "accuracy": averaged_loss[1]} def fwd_output_and_loss_func(dataloader_iter, model): batch, _, _ = next(dataloader_iter) if parallel_state.get_pipeline_model_parallel_world_size() == 1: batch = [x.cuda(non_blocking=True) for x in batch] tokens, labels = batch else: # Vision transformer doesn't need attention mask assert ( self.cfg.get("virtual_pipeline_model_parallel_size", None) is None ), "vpp is not supported yet in MegatronVitClassificationModel" if parallel_state.is_pipeline_first_stage(): # First pipeline stage needs only the tokens and position_ids tokens = batch[0].cuda(non_blocking=True) labels = None elif parallel_state.is_pipeline_last_stage(): # Last pipeline stage needs only the labels and loss_mask labels = batch[1].cuda(non_blocking=True) tokens = None else: # Intermediate pipeline stage doesn't need any inputs tokens, labels = None, None output_tensor = model(tokens) return output_tensor, partial(loss_func, labels) return fwd_output_and_loss_func def get_forward_output_only_func(self): def fwd_output_only_func(batch, model): raise NotImplementedError return fwd_output_only_func def validation_step(self, dataloader_iter): """ Our dataloaders produce a micro-batch and then we fetch a number of microbatches depending on the global batch size and model parallel size from the dataloader to produce a list of microbatches. The list of microbatches is then piped through the pipeline using megatron-core fwd/bwd functions. """ mode = 'test' if self.trainer.testing else 'val' # Initialize userbuffer communicators. if self.initialize_ub: self.initialize_ub_func() loss, accuracy = self.fwd_bwd_step(dataloader_iter, True) ( self.validation_step_outputs.append((loss, accuracy)) if mode == 'val' else self.test_step_outputs.append((loss, accuracy)) ) return loss, accuracy def on_validation_epoch_end(self): # TODO (yuya): need fix later, check with Sean if not self.validation_step_outputs: return None assert ( self.cfg.get("virtual_pipeline_model_parallel_size", None) is None ), "vpp is not supported yet in MegatronVitClassificationModel" if parallel_state.is_pipeline_last_stage(): loss_outputs = [output[0] for output in self.validation_step_outputs] acc_outputs = [output[1] for output in self.validation_step_outputs] averaged_metrics = torch.tensor( [torch.stack(loss_outputs).mean(), torch.stack(acc_outputs).mean()], dtype=torch.float32, device='cuda' ) else: averaged_metrics = torch.tensor([0.0, 0.0], dtype=torch.float32, device='cuda') # we can only log on one rank if it is rank zero so we broadcast from last rank torch.distributed.broadcast(averaged_metrics, get_last_rank()) averaged_loss, averaged_acc = averaged_metrics self.log('global_step', self.trainer.global_step, prog_bar=True, rank_zero_only=True, batch_size=1) self.log('val_loss', averaged_loss, prog_bar=True, rank_zero_only=True, batch_size=1) self.log('val_accuracy', averaged_acc, prog_bar=True, rank_zero_only=True, batch_size=1) self.validation_step_outputs.clear() # free memory return averaged_loss def test_step(self, batch, batch_idx): return self.validation_step(batch) def on_test_epoch_end(self): pass def build_train_valid_test_datasets(self): logging.info('Building datasets for ViT...') if self.trainer.limit_val_batches > 1.0 and isinstance(self.trainer.limit_val_batches, float): raise ValueError("limit_val_batches must be an integer or float less than or equal to 1.0.") self._train_ds, self._validation_ds = build_train_valid_datasets( model_cfg=self.cfg, data_path=self.cfg.data.data_path, image_size=(self.cfg.img_h, self.cfg.img_w), ) self._test_ds = None if self._train_ds is not None: logging.info(f'Length of train dataset: {len(self._train_ds)}') if self._validation_ds is not None: logging.info(f'Length of val dataset: {len(self._validation_ds)}') if self._test_ds is not None: logging.info(f'Length of test dataset: {len(self._test_ds)}') logging.info(f'Finished building datasets for ViT.') return self._train_ds, self._validation_ds, self._test_ds def build_pretraining_data_loader(self, dataset, consumed_samples, drop_last=True): """Buld dataloader given an input dataset.""" logging.info(f'Building dataloader with consumed samples: {consumed_samples}') # Megatron sampler if hasattr(self.cfg.data, 'dataloader_type') and self.cfg.data.dataloader_type is not None: if self.cfg.data.dataloader_type == 'single': batch_sampler = MegatronPretrainingSampler( total_samples=len(dataset), consumed_samples=consumed_samples, micro_batch_size=self.cfg.micro_batch_size, global_batch_size=self.cfg.global_batch_size, data_parallel_rank=parallel_state.get_data_parallel_rank(), data_parallel_size=parallel_state.get_data_parallel_world_size(), drop_last=drop_last, ) elif self.cfg.data.dataloader_type == 'cyclic': batch_sampler = MegatronVisionPretrainingRandomSampler( dataset=dataset, total_samples=len(dataset), consumed_samples=consumed_samples, micro_batch_size=self.cfg.micro_batch_size, global_batch_size=self.cfg.global_batch_size, data_parallel_rank=parallel_state.get_data_parallel_rank(), data_parallel_size=parallel_state.get_data_parallel_world_size(), drop_last=drop_last, data_sharding=self.cfg.data.get("data_sharding", True), ) else: raise ValueError('cfg.data.dataloader_type must be "single" or "cyclic"') else: raise ValueError('cfg.data.dataloader_type not found. Must be "single" or "cyclic"') return torch.utils.data.DataLoader( dataset, batch_sampler=batch_sampler, num_workers=self.cfg.data.num_workers, pin_memory=True, persistent_workers=True if self.cfg.data.num_workers > 0 else False, ) def setup(self, stage=None): """ PTL hook that is executed after DDP spawns. We setup datasets here as megatron datasets require DDP to instantiate. See https://pytorch-lightning.readthedocs.io/en/latest/common/lightning_module.html#setup for more information. Args: stage (str, optional): Can be 'fit', 'validate', 'test' or 'predict'. Defaults to None. """ # log number of parameters if isinstance(self.model, list): num_parameters_on_device = sum( [sum([p.nelement() for p in model_module.parameters()]) for model_module in self.model] ) else: num_parameters_on_device = sum([p.nelement() for p in self.model.parameters()]) # to be summed across data parallel group total_num_parameters = torch.tensor(num_parameters_on_device).cuda() torch.distributed.all_reduce(total_num_parameters, group=parallel_state.get_model_parallel_group()) logging.info( f'Pipeline model parallel rank: {parallel_state.get_pipeline_model_parallel_rank()}, ' f'Tensor model parallel rank: {parallel_state.get_tensor_model_parallel_rank()}, ' f'Number of model parameters on device: {num_parameters_on_device:.2e}. ' f'Total number of model parameters: {total_num_parameters:.2e}.' ) resume_checkpoint_path = self.trainer.ckpt_path if resume_checkpoint_path: init_consumed_samples = self._extract_consumed_samples_from_ckpt(resume_checkpoint_path) else: init_consumed_samples = 0 self.init_consumed_samples = init_consumed_samples self.init_global_step = self.trainer.global_step # allowing restored models to optionally setup datasets self.build_train_valid_test_datasets() self.setup_training_data(self.cfg.data) self.setup_validation_data(self.cfg.data) self.setup_test_data(self.cfg.data) def setup_training_data(self, cfg): if hasattr(self, '_train_ds') and self._train_ds is not None: consumed_samples = self.compute_consumed_samples(0) logging.info( f'Setting up train dataloader with len(len(self._train_ds)): {len(self._train_ds)} and consumed samples: {consumed_samples}' ) self._train_dl = self.build_pretraining_data_loader(self._train_ds, consumed_samples) def setup_validation_data(self, cfg): if hasattr(self, '_validation_ds') and self._validation_ds is not None: consumed_samples = 0 logging.info( f'Setting up validation dataloader with len(len(self._validation_ds)): {len(self._validation_ds)} and consumed samples: {consumed_samples}' ) drop_last = True if not self.cfg.data.get('validation_drop_last', True): logging.info(f'Drop last in validation dataset is set to False') drop_last = False self._validation_dl = self.build_pretraining_data_loader( self._validation_ds, consumed_samples, drop_last=drop_last ) def setup_test_data(self, cfg): if hasattr(self, '_test_ds') and self._test_ds is not None: consumed_samples = 0 logging.info( f'Setting up test dataloader with len(len(self._test_ds)): {len(self._test_ds)} and consumed samples: {consumed_samples}' ) self._test_dl = self.build_pretraining_data_loader(self._test_ds, consumed_samples) def predict_step(self, batch: Any, batch_idx: int, dataloader_idx: Optional[int] = None) -> Any: raise NotImplementedError def transfer_batch_to_device(self, batch: Any, device: torch.device, dataloader_idx: int) -> Any: """PTL hook: https://pytorch-lightning.readthedocs.io/en/latest/common/lightning_module.html#transfer-batch-to-device When using pipeline parallelism, we need the global batch to remain on the CPU, since the memory overhead will be too high when using a large number of microbatches. Microbatches are transferred from CPU to GPU inside the pipeline. """ return batch def _validate_trainer(self): """Certain trainer configurations can break training. Here we try to catch them and raise an error. """ if self.trainer.accumulate_grad_batches > 1: raise ValueError( f'Gradient accumulation is done within training_step. trainer.accumulate_grad_batches must equal 1' ) @classmethod def list_available_models(cls) -> Optional[PretrainedModelInfo]: return None def on_save_checkpoint(self, checkpoint) -> None: """LightningModule hook: https://pytorch-lightning.readthedocs.io/en/stable/common/lightning_module.html#on-save-checkpoint """ if isinstance(self.model, list): for i in range(len(self.model)): checkpoint[f'model{i}'] = self.model[i].module.state_dict_for_save_checkpoint() def on_load_checkpoint(self, checkpoint) -> None: """LightningModule hook: https://pytorch-lightning.readthedocs.io/en/stable/common/lightning_module.html#on-load-checkpoint """ if isinstance(self.model, list): for i in range(len(self.model)): self.model[i].module.load_state_dict(checkpoint[f'model{i}'], strict=True) def parameters(self): if isinstance(self.model, list): return itertools.chain.from_iterable(module.parameters() for module in self.model) else: return self.model.parameters()