Support stage1-2 MoE resume on g0050

training/megatron_bridge/laoyao_2b_moe_pretrain.py

@@ -44,7 +44,7 @@ def _token_weights_from_manifest(prefix_list: list[str], manifest: dict) -> list
     token_by_prefix = {}
     for result in manifest.get("results") or []:
         output_prefix = result.get("output_prefix")
-        tokens = result.get("tokens_estimated_from_int32_bin")
+        tokens = result.get("tokens_estimated_from_int32_bin") or result.get("tokens")
         if output_prefix and tokens:
             token_by_prefix[str(output_prefix) + "_text_document"] = float(tokens)
 
@@ -55,6 +55,42 @@ def _token_weights_from_manifest(prefix_list: list[str], manifest: dict) -> list
     return _normalize_weights(weights)
 
 
+def _target_mix_weights_from_manifest(prefix_list: list[str], manifest: dict) -> list[float] | None:
+    target_mix = manifest.get("target_mix") or {}
+    if not target_mix:
+        return None
+
+    prefix_to_category = {}
+    prefix_to_tokens = {}
+    category_totals = {}
+    for result in manifest.get("results") or []:
+        output_prefix = result.get("output_prefix")
+        category = result.get("mixture_category") or result.get("category")
+        tokens = result.get("tokens")
+        if not output_prefix or not category or not tokens:
+            continue
+
+        prefix = str(output_prefix) + "_text_document"
+        prefix_to_category[prefix] = str(category)
+        prefix_to_tokens[prefix] = float(tokens)
+        category_totals[str(category)] = category_totals.get(str(category), 0.0) + float(tokens)
+
+    if not prefix_to_category:
+        return None
+
+    raw_weights = []
+    for prefix in prefix_list:
+        category = prefix_to_category.get(prefix)
+        tokens = prefix_to_tokens.get(prefix)
+        category_total = category_totals.get(category or "", 0.0)
+        target_ratio = float(target_mix.get(category or "", 0.0))
+        if not category or not tokens or category_total <= 0 or target_ratio <= 0:
+            return None
+        raw_weights.append(target_ratio * (tokens / category_total))
+
+    return _normalize_weights(raw_weights)
+
+
 def _target_mix_weights_from_category_stats(
     prefix_list: list[str],
     category_stats_path: Path,
@@ -91,7 +127,9 @@ def load_data_blend(args: argparse.Namespace) -> tuple[list[str], list[float] |
             raise ValueError(f"{manifest_path} has no ok_prefixes")
         prefix_list = [str(prefix) for prefix in prefixes]
         category_stats_path = manifest_path.with_name("prefix_category_stats.json")
-        weights = _target_mix_weights_from_category_stats(prefix_list, category_stats_path)
+        weights = _target_mix_weights_from_manifest(prefix_list, manifest)
+        if weights is None:
+            weights = _target_mix_weights_from_category_stats(prefix_list, category_stats_path)
         if weights is None:
             weights = _token_weights_from_manifest(prefix_list, manifest)
         return (prefix_list, weights)