Apply packed sequence params change for fused rope compatibility

nemo/collections/llm/api.py

@@ -894,16 +894,17 @@ def _validate_config(
 ) -> None:
 
     ## Model validation
-    assert getattr(model.config, "seq_length", 1) > 0
-    assert getattr(model.config, "max_position_embeddings", 1) > 0
-    assert model.config.num_layers > 0
-    assert model.config.hidden_size > 0
-    assert model.config.num_attention_heads > 0
-    assert model.config.ffn_hidden_size > 0
-
-    if hasattr(model.config, "seq_length"):
-        if getattr(model.config, "max_position_embeddings", None) is not None:
-            assert model.config.seq_length <= model.config.max_position_embeddings
+    if hasattr(model, "config"):
+        assert getattr(model.config, "seq_length", 1) > 0
+        assert getattr(model.config, "max_position_embeddings", 1) > 0
+        assert model.config.num_layers > 0
+        assert model.config.hidden_size > 0
+        assert model.config.num_attention_heads > 0
+        assert model.config.ffn_hidden_size > 0
+
+        if hasattr(model.config, "seq_length"):
+            if getattr(model.config, "max_position_embeddings", None) is not None:
+                assert model.config.seq_length <= model.config.max_position_embeddings
 
     ## Data validation
     assert data.micro_batch_size > 0

nemo/collections/nlp/modules/common/megatron/adapters/mcore_mixins.py

@@ -267,10 +267,17 @@ def forward(
             q_pos_emb, k_pos_emb = rotary_pos_emb
 
             if packed_seq_params is not None:
-                cu_seqlens_q = packed_seq_params.cu_seqlens_q
-                cu_seqlens_kv = packed_seq_params.cu_seqlens_kv
+                if packed_seq_params.cu_seqlens_q_padded is not None:
+                    cu_seqlens_q = packed_seq_params.cu_seqlens_q_padded
+                else:
+                    cu_seqlens_q = packed_seq_params.cu_seqlens_q
+                if packed_seq_params.cu_seqlens_kv_padded is not None:
+                    cu_seqlens_kv = packed_seq_params.cu_seqlens_kv_padded
+                else:
+                    cu_seqlens_kv = packed_seq_params.cu_seqlens_kv
             else:
                 cu_seqlens_q = cu_seqlens_kv = None
+
             query = apply_rotary_pos_emb(query, q_pos_emb, config=self.config, cu_seqlens=cu_seqlens_q)
             key = apply_rotary_pos_emb(key, k_pos_emb, config=self.config, cu_seqlens=cu_seqlens_kv)
             # TODO, can apply positional embedding to value_layer so it has

tests/collections/llm/gpt_finetuning.py

@@ -15,6 +15,7 @@
 import os
 from dataclasses import dataclass
 
+import torch
 from megatron.core.optimizer import OptimizerConfig
 
 from nemo import lightning as nl
@@ -55,6 +56,8 @@ def get_args():
     strategy = nl.MegatronStrategy(
         tensor_model_parallel_size=args.tp_size,
         pipeline_model_parallel_size=args.pp_size,
+        # Pipeline dtype is coupled with the bf16 mixed precision plugin
+        pipeline_dtype=torch.bfloat16,
     )
 
     trainer = nl.Trainer(