automatically amend IDs in datasets

realhf/api/core/data_api.py

@@ -636,6 +636,13 @@ def load_shuffle_split_dataset(
         assert dataset_builder is not None
         data = dataset_builder()
 
+    if any("id" not in d for d in data):
+        logger.warning(
+            f'Key "id" not found in the dataset. Use indices as dataset IDs.'
+        )
+        for idx, d in enumerate(data):
+            d["id"] = idx
+
     datasize_per_rank = len(data) // util.world_size
     shuffle_indices = get_shuffle_indices(
         util.seed, datasize_per_rank * util.world_size

realhf/experiments/common/utils.py

@@ -83,7 +83,9 @@ def make_train_backend_config(
         if model_cfg.zero_stage == 3:
             raise ValueError("Zero stage 3 is not supported in Megatron backend.")
         if model_cfg.zero_stage == 2:
-            logger.warning("Megatron does not ZeRO stage 2. Degenerates to stage 1.")
+            logger.warning(
+                "Megatron does not support ZeRO stage 2. Degenerates to stage 1."
+            )
             model_cfg.zero_stage = 1
         return ModelBackendAbstraction(
             "megatron",