Some fixes to bump mcore (#11600)

* chore(beep boop 🤖): Bump `MCORE_TAG=71c394b...` (2024-12-15) Signed-off-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com> * ci: Add `no-fail-fast` mode Signed-off-by: Oliver Koenig <[email protected]> * fix _get_layer_offset api for mllama Signed-off-by: yaoyu-33 <[email protected]> * bump Signed-off-by: Oliver Koenig <[email protected]> --------- Signed-off-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com> Signed-off-by: Oliver Koenig <[email protected]> Signed-off-by: yaoyu-33 <[email protected]> Signed-off-by: oliver könig <[email protected]> Co-authored-by: pablo-garay <[email protected]> Co-authored-by: yaoyu-33 <[email protected]>
NVIDIA · Jan 3, 2025 · 778ea30 · 778ea30
1 parent 2797fc1
commit 778ea30
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/nemo/collections/vlm/mllama/model/language.py b/nemo/collections/vlm/mllama/model/language.py
@@ -390,7 +390,7 @@ def sharded_state_dict(
         layer_prefix = f'{prefix}layers.'
         num_layers = self.config.num_layers
         for layer in self.layers:
-            offset = layer._get_layer_offset()
+            offset = layer._get_layer_offset(layer.config)
             global_layer_offset = layer.layer_number - 1  # self.layer_number starts at 1
             state_dict_prefix = f'{layer_prefix}{global_layer_offset - offset}.'  # module list index in TransformerBlock # pylint: disable=line-too-long
             sharded_prefix = layer_prefix
@@ -403,7 +403,7 @@ def sharded_state_dict(
         for xlayer in self.xattn_layers:
             if isinstance(xlayer, DummyCrossAttentionTransformerLayer):
                 continue
-            offset = xlayer._get_layer_offset()
+            offset = xlayer._get_layer_offset(xlayer.config)
             global_layer_offset = xlayer.layer_number - 1
             state_dict_prefix = f'{xlayer_prefix}{global_layer_offset - offset}.'  # module list index in TransformerBlock # pylint: disable=line-too-long
             sharded_prefix = f'{xlayer_prefix}{global_layer_offset}.'