Merge branch 'cherry-pick-51648635' into 'main'

ckpt learning rate scheduler fix See merge request ADLR/megatron-lm!815
NVIDIA · Oct 6, 2023 · 0c9c4f3 · 0c9c4f3
2 parents 4c0daab + 809bd3c
commit 0c9c4f3
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/megatron/core/dist_checkpointing/dict_utils.py b/megatron/core/dist_checkpointing/dict_utils.py
@@ -20,7 +20,7 @@ def extract_matching_values(
                 match, nonmatch = extract_matching_values(v, predicate)
                 if match:
                     matching_vals[k] = match
-                if nonmatch:
+                if nonmatch or not v:
                     nonmatching_vals[k] = nonmatch
             elif predicate(v):
                 matching_vals[k] = v
@@ -35,7 +35,7 @@ def extract_matching_values(
                 match, nonmatch = extract_matching_values(v, predicate)
                 if match:
                     matching_vals.append(match)
-                if nonmatch:
+                if nonmatch or not v:
                     nonmatching_vals.append(nonmatch)
             elif predicate(v):
                 matching_vals.append(v)