Distributed dataloader

chemlactica/config/create_train_config.py

@@ -20,6 +20,9 @@
 model_train_configs["125m"][
     "tokenizer_path"
 ] = "chemlactica/tokenizer/ChemLacticaTokenizer66"
+model_train_configs["small_opt"][
+    "tokenizer_path"
+] = "chemlactica/tokenizer/ChemLacticaTokenizer66"
 model_train_configs["1.3b"][
     "tokenizer_path"
 ] = "chemlactica/tokenizer/ChemLacticaTokenizer66"

chemlactica/jsonl_dataset.py

@@ -1,8 +1,12 @@
 from typing import List
-import torch
+
+# import torch
 
 # from io import StringIO
 import os
+from accelerate.state import PartialState
+
+distributed_state = PartialState()
 
 
 def generator_init_print(shared_jsonl_files, files):
@@ -23,51 +27,56 @@ def setup_generator(shared_jsonl_files, files):
     return file_states
 
 
-def get_batch(file, state, chunk_size):
-    with open(file) as f:
-        f.seek(state["position"])
-        batch = f.read(chunk_size)
-        if not batch:
-            raise StopIteration
+# def get_batch(file, state, chunk_size):
+#     with open(file) as f:
+#         f.seek(state["position"])
+#         batch = f.read(chunk_size)
+#         if not batch:
+#             raise StopIteration
+
+#         batch += f.readline()
+#         batch = batch.splitlines()
 
-        batch += f.readline()
-        batch = batch.splitlines()
+#         # batch = [line.rstrip("\n") for line in batch]
+#         state["position"] = f.tell()
+#         batch_len = len(batch)
+#         state["line_number"] += batch_len
+#     return batch, batch_len, state
 
-        # batch = [line.rstrip("\n") for line in batch]
-        state["position"] = f.tell()
-        batch_len = len(batch)
-        state["line_number"] += batch_len
-    return batch, batch_len, state
 
+def should_yield_on_current_rank(i, num_processes, process_index):
+    return i % num_processes == process_index
 
-def format_sample(sample, return_line_info, batch_len, file, state, i):
+
+def format_sample(line):
+    sample = line.strip()
     ret = {"text": sample}
-    if return_line_info:
-        ret["line_info"] = {
-            "file": file,
-            "line_number": state["line_number"] - batch_len + i,
-        }
     return ret
 
 
 def samples_generator(
     files: List[str], shared_jsonl_files, chunk_size=25000, return_line_info=False
 ):
-    if not torch.distributed.is_initialized() or torch.distributed.get_rank() == 0:
-        file_states = setup_generator(shared_jsonl_files, files)
-
-        returned = True
-        while returned:
-            returned = False
-            for file, state in file_states.items():
-                try:
-                    batch, batch_len, state = get_batch(file, state, chunk_size)
-                except StopIteration:
-                    break
-                for i, sample in enumerate(batch, start=1):
-                    returned = True
-                    ret = format_sample(
-                        sample, return_line_info, batch_len, file, state, i
-                    )
-                    yield ret
-                shared_jsonl_files[file] = state
+    file_states = setup_generator(shared_jsonl_files, files)
+
+    returned = True
+    while returned:
+        returned = False
+        for file, state in file_states.items():
+            with open(file) as f:
+                f.seek(state["position"])
+                line = f.readline()
+                counter = 0
+                while line:
+                    state["position"] = f.tell()
+                    if should_yield_on_current_rank(
+                        counter,
+                        distributed_state.num_processes,
+                        distributed_state.process_index,
+                    ):
+                        returned = True
+                        ret = format_sample(line)
+                        yield ret
+                    counter = counter + 1
+                    shared_jsonl_files[file] = state
+                    line = f.readline()

chemlactica/train.py

@@ -46,7 +46,6 @@
 logger = logging.get_logger("transformers")
 
 os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "caching_allocator"
-os.environ["TOKENIZERS_PARALLELISM"] = "true"
 
 signal.signal(signal.SIGINT, signal_handler)
 signal.signal(signal.SIGTERM, signal_handler)
@@ -188,9 +187,9 @@ def train(
 
     accelerator.wait_for_everyone()
 
-    with multiprocessing.Manager() if accelerator.is_main_process else nullcontext() as manager:
+    with multiprocessing.Manager() as manager:
         shared_jsonl_files = None
-        if accelerator.is_main_process and train_type == "pretrain":
+        if train_type == "pretrain":
             shared_jsonl_files = manager.dict()
             trainer_callback_dict[
                 "json_dataset_resume_callback"

chemlactica/utils/model_utils.py

@@ -99,7 +99,7 @@ def load_model(
                 ffn_dim=train_config["ffn_dim"],
                 max_position_embeddings=train_config["max_position_embeddings"],
                 num_attention_heads=train_config["num_attention_heads"],
-                word_embed_proj_dim=train_config["word_sembed_proj_dim"],
+                word_embed_proj_dim=train_config["word_embed_proj_dim"],
             )
         )
     if "galactica" in from_pretrained.lower():

test_status.yaml

@@ -1 +1 @@
-42fb5cdbfbb48de3934afa0bacba13c773c48473: PASS
+82ee7453c84ef0f5e3be607e7850a0db4d541eac: PASS

tests/_test_dataloader_speed.py

@@ -5,7 +5,7 @@
 
 from datasets import load_dataset
 from torch.utils.data import DataLoader
-from dataset_utils import process_dataset
+from chemlactica.utils.dataset_utils import process_dataset
 from transformers import TrainingArguments
 from transformers.data.data_collator import default_data_collator
 from transformers.trainer_pt_utils import IterableDatasetShard