Neva ETP EPP support

nemo/collections/llm/api.py

@@ -908,10 +908,17 @@ def _validate_config(
         assert trainer.strategy.pipeline_model_parallel_size > 0
         assert trainer.strategy.context_parallel_size > 0
 
+        encoder_tensor_model_parallel_size = trainer.strategy.encoder_tensor_model_parallel_size
+        # By default, encoder has the same TP size as decoder
+        if encoder_tensor_model_parallel_size == 0:
+            encoder_tensor_model_parallel_size = trainer.strategy.tensor_model_parallel_size
+
         # DP validation
         assert (trainer.num_devices * trainer.num_nodes) % (
-            trainer.strategy.tensor_model_parallel_size
-            * trainer.strategy.pipeline_model_parallel_size
+            (
+                trainer.strategy.tensor_model_parallel_size * trainer.strategy.pipeline_model_parallel_size
+                + encoder_tensor_model_parallel_size * trainer.strategy.encoder_pipeline_model_parallel_size
+            )
             * trainer.strategy.context_parallel_size
         ) == 0, "Number of GPUs must be divisible by the product of all parallelism sizes for data parallel."
 
@@ -922,8 +929,11 @@ def _validate_config(
                 * (
                     (trainer.num_devices * trainer.num_nodes)
                     / (
-                        trainer.strategy.tensor_model_parallel_size
-                        * trainer.strategy.pipeline_model_parallel_size
+                        (
+                            trainer.strategy.tensor_model_parallel_size * trainer.strategy.pipeline_model_parallel_size
+                            + encoder_tensor_model_parallel_size
+                            * trainer.strategy.encoder_pipeline_model_parallel_size
+                        )
                         * trainer.strategy.context_parallel_size
                     )
                 )

nemo/collections/multimodal/data/energon/base.py

@@ -413,20 +413,3 @@ def transform_dataloader(self, dataloader: DataLoader) -> DataLoader:
         DataLoader: The transformed DataLoader.
         """
         return dataloader
-
-    @property
-    def megatron_data_kwargs(self) -> Dict[str, Any]:
-        """
-        Return the keyword arguments required for Megatron data handling.
-
-        This property provides the necessary arguments that Megatron uses to handle data, including sequence length,
-        micro-batch size, and the number of micro-batches.
-
-        Returns:
-        Dict[str, Any]: A dictionary containing the Megatron data handling arguments.
-        """
-        return {
-            "seq_length": self.seq_len,
-            "micro_batch_size": self.micro_batch_size,
-            "num_microbatches": self.num_microbatches,
-        }

nemo/collections/multimodal/models/multimodal_llm/neva/neva_model.py

@@ -943,7 +943,6 @@ def fwd_bwd_step(self, dataloader_iter, forward_only, first_val_step=None):
             # 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()
-            # print(f"{torch.distributed.get_rank()}: {parallel_state.is_pipeline_last_stage()} {fwd_bwd_function}")
 
             # TODO @akhattar: add num_micro_batches_with_partial_activation_checkpoints when ready
             losses_reduced_per_micro_batch = fwd_bwd_function(

nemo/collections/vlm/neva/data/mock.py

@@ -74,7 +74,7 @@ def __init__(
         )
 
     def setup(self, stage: str = "") -> None:
-        seq_length = self.seq_length
+        seq_length = self.decoder_seq_len or self.seq_length
         if self.packed_sequence and self.micro_batch_size > 1:
             seq_length = seq_length // self.micro_batch_size
             logging.warning(

nemo/collections/vlm/neva/model/base.py

@@ -225,7 +225,8 @@ def configure_model(self, tokenizer) -> "MCoreNevaModel":
         model = MCoreNevaModel(
             config=self,
             tokenizer=tokenizer,
-            pre_process=ps.is_pipeline_first_stage(),
+            pre_process=ps.is_pipeline_first_stage()
+            or ps.get_pipeline_model_parallel_rank() == self.encoder_pipeline_model_parallel_size,
             post_process=ps.is_pipeline_last_stage(),
             add_encoder=ps.is_pipeline_first_stage(),
             add_decoder=ps.is_pipeline_last_stage()
@@ -361,14 +362,15 @@ def __init__(
             freeze_vision_projection=config.freeze_vision_projection,
         )
 
-        self.model_type = ModelType.encoder_or_decoder
+        self.model_type = ModelType.encoder_and_decoder
         # This attribute is needed to check if an all-reduce is required
         # on the word embeddings inside `finalize_model_grads._allreduce_word_embedding_grads`.
 
         self.vision_model_from_hf = hasattr(vision_transformer_config, "image_size")
         self._img_seq_len = vision_transformer_config.num_image_embeddings_per_tile
         if drop_vision_class_token and vision_transformer_config.add_class_token:
             self._img_seq_len -= vision_transformer_config.class_token_len
+        self._language_hidden_size = language_transformer_config.hidden_size
 
     def forward(
         self,
@@ -426,7 +428,9 @@ def forward(
         elif self.add_encoder and not has_images:
             vision_param = next(self.vision_model.parameters())
             # If no images provided, use an empty image embeddings tensor.
-            image_embeddings = torch.tensor([], dtype=vision_param.dtype, device=vision_param.device).reshape(0, 0, 0)
+            image_embeddings = torch.tensor([], dtype=vision_param.dtype, device=vision_param.device).reshape(
+                self._img_seq_len, 0, self._language_hidden_size
+            )
         elif self.add_encoder and has_images:
             # images is in shape of (num_images_in_mbs, c, h, w)
             # note num_images_in_mbs is not mbs but total images in this mbs.
@@ -459,8 +463,23 @@ def forward(
                 )
         else:
             image_embeddings = self.encoder_hidden_state
+            if self.config.encoder_pipeline_model_parallel_size > 0:
+                num_images = len(images) if images is not None else 0
+                image_embeddings = image_embeddings[:, :num_images]
 
         if not self.add_decoder:
+            if self.config.encoder_pipeline_model_parallel_size > 0:
+                _, num_images, _ = image_embeddings.shape
+                pad_amount = max(input_ids.size(0) - num_images, 0)
+                if pad_amount > 0:
+                    pad_tensor = torch.zeros(
+                        self._img_seq_len,
+                        pad_amount,
+                        self._language_hidden_size,
+                        dtype=image_embeddings.dtype,
+                        device=image_embeddings.device,
+                    )
+                    image_embeddings = torch.cat([image_embeddings, pad_tensor], dim=1)
             return image_embeddings
 
         language_embeddings = None

nemo/collections/vlm/vision/base.py

@@ -85,6 +85,7 @@ def configure_model(self) -> "MCoreMultimodalProjector":
         if self.projector_type.startswith("mcore") and self.layer_spec is None:
             if self.projector_type == "mcore_mlp":
                 self.projector_type = "mlp"  # strip "mcore_" for mcore init
+                self.add_bias_linear = self.bias
                 self.layer_spec = ModuleSpec(
                     module=MLP,
                     submodules=MLPSubmodules(

nemo/collections/vlm/vision/intern_vit.py

@@ -168,6 +168,7 @@ def get_mlp_module_spec(use_te: bool = True) -> ModuleSpec:
     )
 
 
+# Handle InternViT's layer scaling.
 def _bias_dropout_add_func_internvit(ls, x_with_bias, residual, prob, training):
     """Handle InternViT's layer scaling."""
     x, bias = x_with_bias  # unpack
@@ -199,6 +200,7 @@ def get_bias_dropout_add_internvit(ls, training, fused):
     return bias_dropout_add_unfused_internvit(ls, training)
 
 
+# Add InternViT specialties to our default TransformerLayer.
 class InternViTTransformerLayer(TransformerLayer):
     """Add InternViT specialties to our default TransformerLayer."""
 
@@ -212,6 +214,7 @@ def __init__(self, *args, **kwargs):
         self.mlp_bda = partial(self.mlp_bda, self.ls2)
 
 
+# Override a few things that are special in InternViT and not supported by the SelfAttention class.
 class InternViTSelfAttention(SelfAttention):
     """Override a few things that are special in InternViT and not supported by the SelfAttention class."""
 

nemo/lightning/_strategy_lib.py

@@ -620,12 +620,24 @@ def _sync_from_last_pipeline_stage(value: torch.Tensor, broadcast: bool = False)
     if parallel_state.get_pipeline_model_parallel_world_size() > 1:
         src_rank = parallel_state.get_pipeline_model_parallel_last_rank()
 
+        if not isinstance(src_rank, list):
+            src_rank = [src_rank]
+
         if not broadcast:
-            pp_ranks = torch.distributed.get_process_group_ranks(parallel_state.get_pipeline_model_parallel_group())
-            if torch.distributed.get_rank() == src_rank and 0 in pp_ranks:
-                torch.distributed.send(value, 0)
-            elif torch.distributed.get_rank() == 0:
-                torch.distributed.recv(value, src_rank)
+            group = parallel_state.get_pipeline_model_parallel_group()
+            if isinstance(group, list):
+                pp_ranks = []
+                for g in group:
+                    pp_ranks.append(torch.distributed.get_process_group_ranks(g))
+            else:
+                pp_ranks = torch.distributed.get_process_group_ranks(group)
+
+            for src_rank_idx in src_rank:
+                if torch.distributed.get_rank() == 0:
+                    torch.distributed.recv(value, src_rank_idx)
+                elif torch.distributed.get_rank() == src_rank_idx and 0 in pp_ranks:
+                    torch.distributed.send(value, 0)
+
         else:
             torch.distributed.broadcast(
                 value,

nemo/lightning/pytorch/callbacks/megatron_comm_overlap.py

@@ -184,8 +184,8 @@ def _get_optimizer_overlap_cfgs(self, parallelism_cfg: ParallelismConfig) -> _Co
 
         if data_parallel_size > 1:
             comm_overlap_cfg.bucket_size = 128 * 1024 * 1024
-            comm_overlap_cfg.overlap_grad_reduce = True
-            comm_overlap_cfg.overlap_param_gather = True
+            comm_overlap_cfg.overlap_grad_reduce = False
+            comm_overlap_cfg.overlap_param_gather = False
             if parallelism_cfg.pipeline_model_parallel_size > 1 and vp_size > 1:
                 # Currently disabled due to an issue with checkpointing
                 # comm_overlap_cfg.overlap_param_gather_with_optimizer_step = True

scripts/vlm/neva_finetune.py

@@ -43,6 +43,7 @@
 from nemo.collections import llm, vlm
 from nemo.collections.multimodal.data.energon.task_encoder import MultiModalTaskEncoder
 from nemo.collections.vlm import ImageDataConfig
+from nemo.lightning.pytorch.callbacks import NsysCallback
 from nemo.lightning.pytorch.callbacks.megatron_comm_overlap import MegatronCommOverlapCallback
 from nemo.lightning.pytorch.optim import CosineAnnealingScheduler
 from nemo.lightning.pytorch.optim.megatron import MegatronOptimizerModule
@@ -73,6 +74,8 @@
         input_size=vision_transformer_config.hidden_size,
         hidden_size=language_transformer_config.hidden_size,
         ffn_hidden_size=language_transformer_config.hidden_size,
+        bias=False,
+        bias_activation_fusion=False,
     )
 
     # NEVA model configuration
@@ -84,7 +87,11 @@
         freeze_language_model=False,
         freeze_vision_model=True,
     )
-    num_image_embeddings_per_tile = vision_transformer_config.num_image_embeddings_per_tile
+    num_image_embeddings_per_tile = vision_transformer_config.num_image_embeddings_per_tile - int(
+        neva_config.drop_vision_class_token and vision_transformer_config.add_class_token
+    )
+
+    seq_length = num_image_embeddings_per_tile
 
     if args.data_type == "llava":
         # Data configuration
@@ -97,8 +104,8 @@
         data = vlm.NevaPreloadedDataModule(
             paths=args.data_path,
             data_config=data_config,
-            seq_length=decoder_seq_length,
-            decoder_seq_length=None,
+            seq_length=seq_length,
+            decoder_seq_length=decoder_seq_length,
             global_batch_size=gbs,
             micro_batch_size=mbs,
             tokenizer=None,
@@ -133,7 +140,8 @@
             path=args.data_path,
             tokenizer=tokenizer,
             image_processor=image_processor,
-            seq_length=decoder_seq_length,
+            seq_length=seq_length,
+            decoder_seq_length=decoder_seq_length,
             micro_batch_size=mbs,
             global_batch_size=gbs,
             num_workers=0,
@@ -143,14 +151,15 @@
                 image_processor=image_processor,
                 multimodal_sample_config=config,
                 packed_sequence=args.use_packed_sequence,
-                packed_sequence_size=decoder_seq_length,
+                packed_sequence_size=seq_length,
                 num_image_embeddings_per_tile=num_image_embeddings_per_tile,
             ),
             packing_buffer_size=200 if args.use_packed_sequence else None,
         )
     elif args.data_type == "mock":
         data = vlm.NevaMockDataModule(
-            seq_length=decoder_seq_length,
+            seq_length=seq_length,
+            decoder_seq_length=decoder_seq_length,
             global_batch_size=gbs,
             micro_batch_size=mbs,
             tokenizer=None,
@@ -168,9 +177,10 @@
         tensor_model_parallel_size=args.tp_size,
         pipeline_model_parallel_size=args.pp_size,
         encoder_pipeline_model_parallel_size=args.encoder_pp_size,
+        encoder_tensor_model_parallel_size=args.encoder_tp_size,
         context_parallel_size=args.cp_size,
         pipeline_dtype=torch.bfloat16,
-        sequence_parallel=True,
+        sequence_parallel=False,
         ddp=DistributedDataParallelConfig(
             check_for_nan_in_grad=True,
             grad_reduce_in_fp32=True,
@@ -290,6 +300,7 @@
     parser.add_argument("--pp_size", type=int, required=False, default=1)
     parser.add_argument("--cp_size", type=int, required=False, default=1)
     parser.add_argument("--encoder_pp_size", type=int, required=False, default=0)
+    parser.add_argument("--encoder_tp_size", type=int, required=False, default=0)
     parser.add_argument("--projector_type", type=str, required=False, default="mcore_mlp")
     parser.add_argument("--name", type=str, required=False, default="neva_pretrain")
     parser.add_argument("--peft", type=str, default='none', help="none | lora")

scripts/vlm/neva_generate.py

@@ -78,16 +78,6 @@ def main(args) -> None:
     if raw_image is None:
         return  # Exit if the image can't be loaded
 
-    inputs = processor(prompt, raw_image, return_tensors='pt').to(0, torch.float16)
-    input_ids = hf_tokenizer(prompt, return_tensors='pt')['input_ids'].cuda()
-    input_ids[input_ids == 32000] = -200
-    media = inputs['pixel_values'].cuda()
-    media = media.reshape(media.size(0), 3, 336, 336)
-
-    position_ids = (
-        torch.arange(input_ids.size(1), dtype=torch.long, device=input_ids.device).unsqueeze(0).expand_as(input_ids)
-    )
-
     fabric = trainer.to_fabric()
 
     # Decide whether to import or load the model based on the input arguments
@@ -99,13 +89,23 @@ def main(args) -> None:
 
     model = model.module.cuda()
     model.eval()
+
+    inputs = processor(prompt, raw_image, return_tensors='pt').to(0, torch.float16)
+    input_ids = hf_tokenizer(prompt, return_tensors='pt')['input_ids'].to(model.device)
+    input_ids[input_ids == 32000] = -200
+    images = inputs['pixel_values'].to(model.device)
+    images = images.reshape(images.size(0), 3, 336, 336)
+
+    position_ids = (
+        torch.arange(input_ids.size(1), dtype=torch.long, device=input_ids.device).unsqueeze(0).expand_as(input_ids)
+    )
     generated_ids = input_ids.clone()
 
     # Greedy generation loop
     for _ in range(20):
         with torch.no_grad():
             output = model(
-                media=media,
+                images=images,
                 input_ids=input_ids,
                 position_ids=position_ids,
                 attention_mask=None,