big, breaking change but large upside: swap previous FineWeb-EDU dataset to NVIDIA ClimbMix dataset. Requires people to download the data shards. The upside is that training GPT-2 capablity model now only takes ~2 hours, down from 2.76 hours, so this is a huge win data-wise

2026-03-04 19:47:12 +00:00
parent b07604ebaa
commit 324e69c45d
5 changed files with 73 additions and 17 deletions
@@ -32,7 +32,8 @@ def _document_batches(split, resume_state_dict, tokenizer_batch_size):
    """
    ddp, ddp_rank, ddp_local_rank, ddp_world_size = get_dist_info()

-    parquet_paths = list_parquet_files()
+    warn_on_legacy = ddp_rank == 0 and split == "train" # rank 0 on train split will warn on legacy
+    parquet_paths = list_parquet_files(warn_on_legacy=warn_on_legacy)
    assert len(parquet_paths) != 0, "No dataset parquet files found, did you run dataset.py?"
    parquet_paths = parquet_paths[:-1] if split == "train" else parquet_paths[-1:]