Even more detailed logging.

mlfoundations · May 14, 2024 · e10d61d · e10d61d
1 parent 236817a
commit e10d61d
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/open_lm/main.py b/open_lm/main.py
@@ -59,6 +59,7 @@
     check_exists,
     start_sync_process,
     remote_sync_with_expon_backoff,
+    get_metadata_file,
     get_string_for_epoch,
     log_num_checkpoints,
     terminate_sync_process,
@@ -875,6 +876,9 @@ def main(args):
             if args.dataset_manifest is not None:
                 for i in range(len(next_shard_per_source)):
                     end_of_epoch_log[f"next_shard_{i}"] = next_shard_per_source[i]
+                    end_of_epoch_log[f"dataset_pass_{i}"] = next_shard_per_source[i] // len(
+                        get_metadata_file(args.dataset_manifest[i])
+                    )
 
             for name, val in end_of_epoch_log.items():
                 name = "train/" + name